From 2495b13053a5a78c8bf5c2e4ed8246e0a9bc310f Mon Sep 17 00:00:00 2001 From: ginuerzh Date: Tue, 18 Feb 2025 17:30:36 +0800 Subject: [PATCH] fix traffic logging for socks5 udp --- handler/socks/v4/handler.go | 48 ++++++++++++++++++------------------ handler/socks/v5/bind.go | 29 ++++++++++++++++++++++ handler/socks/v5/connect.go | 48 ++++++++++++++++++------------------ handler/socks/v5/handler.go | 4 +-- handler/socks/v5/mbind.go | 29 ++++++++++++++++++++++ handler/socks/v5/metadata.go | 8 ------ handler/socks/v5/udp.go | 40 +++++++++++++++++++++++------- handler/socks/v5/udp_tun.go | 27 +++++++++++++++++++- 8 files changed, 165 insertions(+), 68 deletions(-) diff --git a/handler/socks/v4/handler.go b/handler/socks/v4/handler.go index 8cc9e8c4..e4c548f7 100644 --- a/handler/socks/v4/handler.go +++ b/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)) diff --git a/handler/socks/v5/bind.go b/handler/socks/v5/bind.go index b9ace68c..1b3e2180 100644 --- a/handler/socks/v5/bind.go +++ b/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) } diff --git a/handler/socks/v5/connect.go b/handler/socks/v5/connect.go index 1e83c344..e71fed86 100644 --- a/handler/socks/v5/connect.go +++ b/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)) diff --git a/handler/socks/v5/handler.go b/handler/socks/v5/handler.go index 8d7f6955..2cdd5173 100644 --- a/handler/socks/v5/handler.go +++ b/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.OutputBytes = pStats.Get(stats.KindOutputBytes) + ro.InputBytes += pStats.Get(stats.KindInputBytes) + 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) diff --git a/handler/socks/v5/mbind.go b/handler/socks/v5/mbind.go index 3c25bd2b..6a1b2c79 100644 --- a/handler/socks/v5/mbind.go +++ b/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) } diff --git a/handler/socks/v5/metadata.go b/handler/socks/v5/metadata.go index 2fb17111..a8c264ef 100644 --- a/handler/socks/v5/metadata.go +++ b/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") diff --git a/handler/socks/v5/udp.go b/handler/socks/v5/udp.go index 9221b6d6..2011187b 100644 --- a/handler/socks/v5/udp.go +++ b/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) diff --git a/handler/socks/v5/udp_tun.go b/handler/socks/v5/udp_tun.go index 81df6d35..5f416ba8 100644 --- a/handler/socks/v5/udp_tun.go +++ b/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())