fix traffic logging for socks5 udp

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,19 +80,38 @@ func (h *socks5Handler) handleUDP(ctx context.Context, conn net.Conn, ro *xrecor
 		return err
 	}
 
-	clientID := ctxvalue.ClientIDFromContext(ctx)
-	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)
-		cc = stats_wrapper.WrapPacketConn(cc, pstats)
+	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)
+			pstats.Add(stats.KindCurrentConns, 1)
+			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)