fix traffic logging for socks5 udp

handler/socks/v4/handler.go

@@ -207,6 +207,30 @@ func (h *socks4Handler) handleConnect(ctx context.Context, conn net.Conn, req *g
 	})
 	log.Debugf("%s >> %s", conn.RemoteAddr(), addr)
 
+	{
+		clientID := ctxvalue.ClientIDFromContext(ctx)
+		rw := traffic_wrapper.WrapReadWriter(
+			h.limiter,
+			conn,
+			string(clientID),
+			limiter.ScopeOption(limiter.ScopeClient),
+			limiter.ServiceOption(h.options.Service),
+			limiter.NetworkOption("tcp"),
+			limiter.AddrOption(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)
+	}
+
 	if h.options.Bypass != nil && h.options.Bypass.Contains(ctx, "tcp", addr) {
 		resp := gosocks4.NewReply(gosocks4.Rejected, nil)
 		log.Trace(resp)
@@ -237,30 +261,6 @@ func (h *socks4Handler) handleConnect(ctx context.Context, conn net.Conn, req *g
 		return err
 	}
 
-	{
-		clientID := ctxvalue.ClientIDFromContext(ctx)
-		rw := traffic_wrapper.WrapReadWriter(
-			h.limiter,
-			conn,
-			string(clientID),
-			limiter.ScopeOption(limiter.ScopeClient),
-			limiter.ServiceOption(h.options.Service),
-			limiter.NetworkOption("tcp"),
-			limiter.AddrOption(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)
-	}
-
 	if h.md.sniffing {
 		if h.md.sniffingTimeout > 0 {
 			conn.SetReadDeadline(time.Now().Add(h.md.sniffingTimeout))

handler/socks/v5/bind.go

@@ -6,9 +6,14 @@ import (
 	"net"
 	"time"
 
+	"github.com/go-gost/core/limiter"
 	"github.com/go-gost/core/logger"
+	"github.com/go-gost/core/observer/stats"
 	"github.com/go-gost/gosocks5"
+	ctxvalue "github.com/go-gost/x/ctx"
 	xnet "github.com/go-gost/x/internal/net"
+	traffic_wrapper "github.com/go-gost/x/limiter/traffic/wrapper"
+	stats_wrapper "github.com/go-gost/x/observer/stats/wrapper"
 )
 
 func (h *socks5Handler) handleBind(ctx context.Context, conn net.Conn, network, address string, log logger.Logger) error {
@@ -26,6 +31,30 @@ func (h *socks5Handler) handleBind(ctx context.Context, conn net.Conn, network,
 		return reply.Write(conn)
 	}
 
+	{
+		clientID := ctxvalue.ClientIDFromContext(ctx)
+		rw := traffic_wrapper.WrapReadWriter(
+			h.limiter,
+			conn,
+			string(clientID),
+			limiter.ServiceOption(h.options.Service),
+			limiter.ScopeOption(limiter.ScopeClient),
+			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)
+	}
+
 	// BIND does not support chain.
 	return h.bindLocal(ctx, conn, network, address, log)
 }

handler/socks/v5/connect.go

@@ -29,6 +29,30 @@ func (h *socks5Handler) handleConnect(ctx context.Context, conn net.Conn, networ
 	})
 	log.Debugf("%s >> %s", conn.RemoteAddr(), address)
 
+	{
+		clientID := ctxvalue.ClientIDFromContext(ctx)
+		rw := traffic_wrapper.WrapReadWriter(
+			h.limiter,
+			conn,
+			string(clientID),
+			limiter.ServiceOption(h.options.Service),
+			limiter.ScopeOption(limiter.ScopeClient),
+			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.options.Bypass != nil && h.options.Bypass.Contains(ctx, network, address) {
 		resp := gosocks5.NewReply(gosocks5.NotAllowed, nil)
 		log.Trace(resp)
@@ -59,30 +83,6 @@ func (h *socks5Handler) handleConnect(ctx context.Context, conn net.Conn, networ
 		return err
 	}
 
-	{
-		clientID := ctxvalue.ClientIDFromContext(ctx)
-		rw := traffic_wrapper.WrapReadWriter(
-			h.limiter,
-			conn,
-			string(clientID),
-			limiter.ServiceOption(h.options.Service),
-			limiter.ScopeOption(limiter.ScopeClient),
-			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/socks/v5/handler.go

@@ -136,8 +136,8 @@ func (h *socks5Handler) Handle(ctx context.Context, conn net.Conn, opts ...handl
 		if err != nil {
 			ro.Err = err.Error()
 		}
-		ro.InputBytes = pStats.Get(stats.KindInputBytes)
+		ro.InputBytes += pStats.Get(stats.KindInputBytes)
-		ro.OutputBytes = pStats.Get(stats.KindOutputBytes)
+		ro.OutputBytes += pStats.Get(stats.KindOutputBytes)
 		ro.Duration = time.Since(start)
 		if err := ro.Record(ctx, h.recorder.Recorder); err != nil {
 			log.Errorf("record: %v", err)

handler/socks/v5/mbind.go

@@ -6,10 +6,15 @@ import (
 	"net"
 	"time"
 
+	"github.com/go-gost/core/limiter"
 	"github.com/go-gost/core/logger"
+	"github.com/go-gost/core/observer/stats"
 	"github.com/go-gost/gosocks5"
+	ctxvalue "github.com/go-gost/x/ctx"
 	xnet "github.com/go-gost/x/internal/net"
 	"github.com/go-gost/x/internal/util/mux"
+	traffic_wrapper "github.com/go-gost/x/limiter/traffic/wrapper"
+	stats_wrapper "github.com/go-gost/x/observer/stats/wrapper"
 )
 
 func (h *socks5Handler) handleMuxBind(ctx context.Context, conn net.Conn, network, address string, log logger.Logger) error {
@@ -27,6 +32,30 @@ func (h *socks5Handler) handleMuxBind(ctx context.Context, conn net.Conn, networ
 		return reply.Write(conn)
 	}
 
+	{
+		clientID := ctxvalue.ClientIDFromContext(ctx)
+		rw := traffic_wrapper.WrapReadWriter(
+			h.limiter,
+			conn,
+			string(clientID),
+			limiter.ServiceOption(h.options.Service),
+			limiter.ScopeOption(limiter.ScopeClient),
+			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)
+	}
+
 	return h.muxBindLocal(ctx, conn, network, address, log)
 }
 

handler/socks/v5/metadata.go

@@ -4,7 +4,6 @@ import (
 	"crypto"
 	"crypto/tls"
 	"crypto/x509"
-	"math"
 	"time"
 
 	"github.com/go-gost/core/bypass"
@@ -19,7 +18,6 @@ type metadata struct {
 	noTLS             bool
 	enableBind        bool
 	enableUDP         bool
-	udpBufferSize     int
 	compatibilityMode bool
 	hash              string
 	muxCfg            *mux.Config
@@ -51,12 +49,6 @@ func (h *socks5Handler) parseMetadata(md mdata.Metadata) (err error) {
 	h.md.enableBind = mdutil.GetBool(md, "bind")
 	h.md.enableUDP = mdutil.GetBool(md, "udp")
 
-	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.compatibilityMode = mdutil.GetBool(md, "comp")
 	h.md.hash = mdutil.GetString(md, "hash")
 

handler/socks/v5/udp.go

@@ -9,6 +9,7 @@ import (
 	"net"
 	"time"
 
+	"github.com/go-gost/core/limiter"
 	"github.com/go-gost/core/logger"
 	"github.com/go-gost/core/observer/stats"
 	"github.com/go-gost/gosocks5"
@@ -16,6 +17,8 @@ import (
 	xnet "github.com/go-gost/x/internal/net"
 	"github.com/go-gost/x/internal/net/udp"
 	"github.com/go-gost/x/internal/util/socks"
+	traffic_wrapper "github.com/go-gost/x/limiter/traffic/wrapper"
+	xstats "github.com/go-gost/x/observer/stats"
 	stats_wrapper "github.com/go-gost/x/observer/stats/wrapper"
 	xrecorder "github.com/go-gost/x/recorder"
 )
@@ -77,7 +80,26 @@ func (h *socks5Handler) handleUDP(ctx context.Context, conn net.Conn, ro *xrecor
 		return err
 	}
 
+	pStats := xstats.Stats{}
+	cc = stats_wrapper.WrapPacketConn(cc, &pStats)
+
+	defer func() {
+		ro.InputBytes = pStats.Get(stats.KindInputBytes)
+		ro.OutputBytes = pStats.Get(stats.KindOutputBytes)
+	}()
+
+	{
 		clientID := ctxvalue.ClientIDFromContext(ctx)
+		cc = traffic_wrapper.WrapPacketConn(
+			cc,
+			h.limiter,
+			string(clientID),
+			limiter.ServiceOption(h.options.Service),
+			limiter.ScopeOption(limiter.ScopeClient),
+			limiter.NetworkOption("udp"),
+			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)
@@ -85,11 +107,11 @@ func (h *socks5Handler) handleUDP(ctx context.Context, conn net.Conn, ro *xrecor
 			defer pstats.Add(stats.KindCurrentConns, -1)
 			cc = stats_wrapper.WrapPacketConn(cc, pstats)
 		}
+	}
 
-	r := udp.NewRelay(socks.UDPConn(cc, h.md.udpBufferSize), pc).
+	r := udp.NewRelay(socks.UDPConn(cc), pc).
 		WithBypass(h.options.Bypass).
 		WithLogger(log)
-	r.SetBufferSize(h.md.udpBufferSize)
 
 	go r.Run(ctx)
 

handler/socks/v5/udp_tun.go

@@ -7,6 +7,7 @@ import (
 	"net"
 	"time"
 
+	"github.com/go-gost/core/limiter"
 	"github.com/go-gost/core/logger"
 	"github.com/go-gost/core/observer/stats"
 	"github.com/go-gost/gosocks5"
@@ -14,6 +15,7 @@ import (
 	xnet "github.com/go-gost/x/internal/net"
 	"github.com/go-gost/x/internal/net/udp"
 	"github.com/go-gost/x/internal/util/socks"
+	traffic_wrapper "github.com/go-gost/x/limiter/traffic/wrapper"
 	stats_wrapper "github.com/go-gost/x/observer/stats/wrapper"
 	xrecorder "github.com/go-gost/x/recorder"
 )
@@ -23,6 +25,30 @@ func (h *socks5Handler) handleUDPTun(ctx context.Context, conn net.Conn, network
 		"cmd": "udp-tun",
 	})
 
+	{
+		clientID := ctxvalue.ClientIDFromContext(ctx)
+		rw := traffic_wrapper.WrapReadWriter(
+			h.limiter,
+			conn,
+			string(clientID),
+			limiter.ServiceOption(h.options.Service),
+			limiter.ScopeOption(limiter.ScopeClient),
+			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)
+	}
+
 	bindAddr, _ := net.ResolveUDPAddr(network, address)
 	if bindAddr == nil {
 		bindAddr = &net.UDPAddr{}
@@ -101,7 +127,6 @@ func (h *socks5Handler) handleUDPTun(ctx context.Context, conn net.Conn, network
 	r := udp.NewRelay(socks.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())