add traffic limiter for proxy handler

This commit is contained in:
ginuerzh
2023-11-18 18:28:09 +08:00
parent 330631fd79
commit 88cc6ff4d5
38 changed files with 633 additions and 200 deletions

View File

@ -8,11 +8,13 @@ import (
"net"
"time"
"github.com/go-gost/core/limiter/traffic"
"github.com/go-gost/core/logger"
"github.com/go-gost/relay"
ctxvalue "github.com/go-gost/x/internal/ctx"
xnet "github.com/go-gost/x/internal/net"
sx "github.com/go-gost/x/internal/util/selector"
serial "github.com/go-gost/x/internal/util/serial"
"github.com/go-gost/x/limiter/traffic/wrapper"
)
func (h *relayHandler) handleConnect(ctx context.Context, conn net.Conn, network, address string, log logger.Logger) (err error) {
@ -51,7 +53,7 @@ func (h *relayHandler) handleConnect(ctx context.Context, conn net.Conn, network
switch h.md.hash {
case "host":
ctx = sx.ContextWithHash(ctx, &sx.Hash{Source: address})
ctx = ctxvalue.ContextWithHash(ctx, &ctxvalue.Hash{Source: address})
}
var cc io.ReadWriteCloser
@ -103,9 +105,16 @@ func (h *relayHandler) handleConnect(ctx context.Context, conn net.Conn, network
}
}
rw := wrapper.WrapReadWriter(h.options.Limiter, conn, conn.RemoteAddr().String(),
traffic.NetworkOption(network),
traffic.AddrOption(address),
traffic.ClientOption(string(ctxvalue.ClientIDFromContext(ctx))),
traffic.SrcOption(conn.RemoteAddr().String()),
)
t := time.Now()
log.Infof("%s <-> %s", conn.RemoteAddr(), address)
xnet.Transport(conn, cc)
xnet.Transport(rw, cc)
log.WithFields(map[string]any{
"duration": time.Since(t),
}).Infof("%s >-< %s", conn.RemoteAddr(), address)