add traffic limiter for relay bind

handler/relay/bind.go

@@ -7,14 +7,19 @@ import (
 	"time"
 
 	"github.com/go-gost/core/handler"
+	"github.com/go-gost/core/limiter"
 	"github.com/go-gost/core/listener"
 	"github.com/go-gost/core/logger"
+	"github.com/go-gost/core/observer/stats"
 	"github.com/go-gost/relay"
+	ctxvalue "github.com/go-gost/x/ctx"
 	xnet "github.com/go-gost/x/internal/net"
 	"github.com/go-gost/x/internal/net/udp"
 	"github.com/go-gost/x/internal/util/mux"
 	relay_util "github.com/go-gost/x/internal/util/relay"
+	traffic_wrapper "github.com/go-gost/x/limiter/traffic/wrapper"
 	metrics "github.com/go-gost/x/metrics/wrapper"
+	stats_wrapper "github.com/go-gost/x/observer/stats/wrapper"
 	xservice "github.com/go-gost/x/service"
 )
 
@@ -26,6 +31,30 @@ func (h *relayHandler) handleBind(ctx context.Context, conn net.Conn, network, a
 
 	log.Debugf("%s >> %s", conn.RemoteAddr(), address)
 
+	{
+		clientID := ctxvalue.ClientIDFromContext(ctx)
+		rw := traffic_wrapper.WrapReadWriter(
+			h.limiter,
+			conn,
+			string(clientID),
+			limiter.ScopeOption(limiter.ScopeClient),
+			limiter.ServiceOption(h.options.Service),
+			limiter.NetworkOption(network),
+			limiter.AddrOption(address),
+			limiter.ClientOption(string(clientID)),
+			limiter.SrcOption(conn.RemoteAddr().String()),
+		)
+		if h.options.Observer != nil {
+			pstats := h.stats.Stats(string(clientID))
+			pstats.Add(stats.KindTotalConns, 1)
+			pstats.Add(stats.KindCurrentConns, 1)
+			defer pstats.Add(stats.KindCurrentConns, -1)
+			rw = stats_wrapper.WrapReadWriter(rw, pstats)
+		}
+
+		conn = xnet.NewReadWriteConn(rw, rw, conn)
+	}
+
 	resp := relay.Response{
 		Version: relay.Version1,
 		Status:  relay.StatusOK,
@@ -174,7 +203,6 @@ func (h *relayHandler) bindUDP(ctx context.Context, conn net.Conn, network, addr
 	r := udp.NewRelay(relay_util.UDPTunServerConn(conn), pc).
 		WithBypass(h.options.Bypass).
 		WithLogger(log)
-	r.SetBufferSize(h.md.udpBufferSize)
 
 	t := time.Now()
 	log.Debugf("%s <-> %s", conn.RemoteAddr(), pc.LocalAddr())

handler/relay/connect.go

@@ -40,6 +40,30 @@ func (h *relayHandler) handleConnect(ctx context.Context, conn net.Conn, network
 
 	log.Debugf("%s >> %s/%s", conn.RemoteAddr(), address, network)
 
+	{
+		clientID := ctxvalue.ClientIDFromContext(ctx)
+		rw := traffic_wrapper.WrapReadWriter(
+			h.limiter,
+			conn,
+			string(clientID),
+			limiter.ScopeOption(limiter.ScopeClient),
+			limiter.ServiceOption(h.options.Service),
+			limiter.NetworkOption(network),
+			limiter.AddrOption(address),
+			limiter.ClientOption(string(clientID)),
+			limiter.SrcOption(conn.RemoteAddr().String()),
+		)
+		if h.options.Observer != nil {
+			pstats := h.stats.Stats(string(clientID))
+			pstats.Add(stats.KindTotalConns, 1)
+			pstats.Add(stats.KindCurrentConns, 1)
+			defer pstats.Add(stats.KindCurrentConns, -1)
+			rw = stats_wrapper.WrapReadWriter(rw, pstats)
+		}
+
+		conn = xnet.NewReadWriteConn(rw, rw, conn)
+	}
+
 	resp := relay.Response{
 		Version: relay.Version1,
 		Status:  relay.StatusOK,
@@ -122,30 +146,6 @@ func (h *relayHandler) handleConnect(ctx context.Context, conn net.Conn, network
 		}
 	}
 
-	{
-		clientID := ctxvalue.ClientIDFromContext(ctx)
-		rw := traffic_wrapper.WrapReadWriter(
-			h.limiter,
-			conn,
-			string(clientID),
-			limiter.ScopeOption(limiter.ScopeClient),
-			limiter.ServiceOption(h.options.Service),
-			limiter.NetworkOption(network),
-			limiter.AddrOption(address),
-			limiter.ClientOption(string(clientID)),
-			limiter.SrcOption(conn.RemoteAddr().String()),
-		)
-		if h.options.Observer != nil {
-			pstats := h.stats.Stats(string(clientID))
-			pstats.Add(stats.KindTotalConns, 1)
-			pstats.Add(stats.KindCurrentConns, 1)
-			defer pstats.Add(stats.KindCurrentConns, -1)
-			rw = stats_wrapper.WrapReadWriter(rw, pstats)
-		}
-
-		conn = xnet.NewReadWriteConn(rw, rw, conn)
-	}
-
 	if h.md.sniffing {
 		if h.md.sniffingTimeout > 0 {
 			conn.SetReadDeadline(time.Now().Add(h.md.sniffingTimeout))

handler/relay/forward.go

@@ -38,6 +38,29 @@ func (h *relayHandler) handleForward(ctx context.Context, conn net.Conn, network
 
 	log.Debugf("%s >> %s", conn.RemoteAddr(), target.Addr)
 
+	{
+		clientID := ctxvalue.ClientIDFromContext(ctx)
+		rw := wrapper.WrapReadWriter(
+			h.limiter,
+			conn,
+			string(clientID),
+			limiter.ServiceOption(h.options.Service),
+			limiter.ScopeOption(limiter.ScopeClient),
+			limiter.NetworkOption(network),
+			limiter.AddrOption(target.Addr),
+			limiter.ClientOption(string(clientID)),
+			limiter.SrcOption(conn.RemoteAddr().String()),
+		)
+		if h.options.Observer != nil {
+			pstats := h.stats.Stats(string(clientID))
+			pstats.Add(stats.KindTotalConns, 1)
+			pstats.Add(stats.KindCurrentConns, 1)
+			defer pstats.Add(stats.KindCurrentConns, -1)
+			rw = stats_wrapper.WrapReadWriter(rw, pstats)
+		}
+		conn = xnet.NewReadWriteConn(rw, rw, conn)
+	}
+
 	cc, err := h.options.Router.Dial(ctx, network, target.Addr)
 	if err != nil {
 		// TODO: the router itself may be failed due to the failed node in the router,
@@ -89,29 +112,9 @@ func (h *relayHandler) handleForward(ctx context.Context, conn net.Conn, network
 		conn = rc
 	}
 
-	clientID := ctxvalue.ClientIDFromContext(ctx)
-	rw := wrapper.WrapReadWriter(
-		h.limiter,
-		conn,
-		string(clientID),
-		limiter.ServiceOption(h.options.Service),
-		limiter.ScopeOption(limiter.ScopeClient),
-		limiter.NetworkOption(network),
-		limiter.AddrOption(target.Addr),
-		limiter.ClientOption(string(clientID)),
-		limiter.SrcOption(conn.RemoteAddr().String()),
-	)
-	if h.options.Observer != nil {
-		pstats := h.stats.Stats(string(clientID))
-		pstats.Add(stats.KindTotalConns, 1)
-		pstats.Add(stats.KindCurrentConns, 1)
-		defer pstats.Add(stats.KindCurrentConns, -1)
-		rw = stats_wrapper.WrapReadWriter(rw, pstats)
-	}
-
 	t := time.Now()
 	log.Debugf("%s <-> %s", conn.RemoteAddr(), target.Addr)
-	xnet.Transport(rw, cc)
+	xnet.Transport(conn, cc)
 	log.WithFields(map[string]any{
 		"duration": time.Since(t),
 	}).Debugf("%s >-< %s", conn.RemoteAddr(), target.Addr)

handler/relay/metadata.go

@@ -4,7 +4,6 @@ import (
 	"crypto"
 	"crypto/tls"
 	"crypto/x509"
-	"math"
 	"time"
 
 	"github.com/go-gost/core/bypass"
@@ -15,12 +14,11 @@ import (
 )
 
 type metadata struct {
-	readTimeout            time.Duration
+	readTimeout time.Duration
-	enableBind             bool
+	enableBind  bool
-	udpBufferSize          int
+	noDelay     bool
-	noDelay                bool
+	hash        string
-	hash                   string
+	muxCfg      *mux.Config
-	muxCfg                 *mux.Config
 
 	observerPeriod       time.Duration
 	observerResetTraffic bool
@@ -47,13 +45,6 @@ func (h *relayHandler) parseMetadata(md mdata.Metadata) (err error) {
 
 	h.md.enableBind = mdutil.GetBool(md, "bind")
 	h.md.noDelay = mdutil.GetBool(md, "nodelay")
-
-	if bs := mdutil.GetInt(md, "udpBufferSize"); bs > 0 {
-		h.md.udpBufferSize = int(math.Min(math.Max(float64(bs), 512), 64*1024))
-	} else {
-		h.md.udpBufferSize = 4096
-	}
-
 	h.md.hash = mdutil.GetString(md, "hash")
 
 	h.md.muxCfg = &mux.Config{