add traffic limiter for proxy handler
This commit is contained in:
@ -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)
|
||||
|
Reference in New Issue
Block a user