From a5309eee9792dbe3042864e376531b4e3fd9abb4 Mon Sep 17 00:00:00 2001 From: ginuerzh Date: Wed, 30 Jul 2025 21:40:53 +0800 Subject: [PATCH] add buffer size option for udp connection --- connector/socks/v5/conn.go | 22 +++++---- connector/socks/v5/connector.go | 9 ++-- connector/socks/v5/metadata.go | 9 ++++ connector/ss/udp/connector.go | 2 +- connector/ss/udp/metadata.go | 11 ++--- handler/http/metadata.go | 2 + handler/http/udp.go | 1 + handler/relay/bind.go | 1 + handler/relay/metadata.go | 12 +++-- handler/router/associate.go | 6 ++- handler/router/entrypoint.go | 6 ++- handler/router/metadata.go | 8 ++- handler/socks/v5/handler.go | 3 +- handler/socks/v5/metadata.go | 2 + handler/socks/v5/udp.go | 14 +++--- handler/socks/v5/udp_tun.go | 8 +-- handler/ss/udp/handler.go | 20 ++++++-- handler/ss/udp/metadata.go | 17 +++---- handler/tungo/metadata.go | 8 ++- handler/tungo/stack.go | 30 ++++++++--- internal/net/udp/listener.go | 6 +++ internal/net/udp/relay.go | 38 +++++++++----- internal/util/relay/conn.go | 88 --------------------------------- internal/util/socks/conn.go | 35 +++++++------ internal/util/ss/conn.go | 43 ++++++++++------ listener/udp/metadata.go | 36 ++++++-------- 26 files changed, 218 insertions(+), 219 deletions(-) diff --git a/connector/socks/v5/conn.go b/connector/socks/v5/conn.go index c5562d82..933a1f23 100644 --- a/connector/socks/v5/conn.go +++ b/connector/socks/v5/conn.go @@ -2,10 +2,10 @@ package v5 import ( "bytes" - "math" "net" "time" + "github.com/go-gost/core/common/bufpool" "github.com/go-gost/core/logger" "github.com/go-gost/gosocks5" ) @@ -25,17 +25,18 @@ func (c *bindConn) RemoteAddr() net.Addr { } type udpRelayConn struct { - udpConn *net.UDPConn - tcpConn net.Conn - taddr net.Addr - rbuf [math.MaxUint16]byte - wbuf [math.MaxInt16]byte - logger logger.Logger + udpConn *net.UDPConn + tcpConn net.Conn + taddr net.Addr + bufferSize int + logger logger.Logger } func (c *udpRelayConn) ReadFrom(b []byte) (n int, addr net.Addr, err error) { - buf := c.rbuf[:] - nn, err := c.udpConn.Read(buf[:]) + buf := bufpool.Get(c.bufferSize) + defer bufpool.Put(buf) + + nn, err := c.udpConn.Read(buf) if err != nil { return } @@ -81,7 +82,8 @@ func (c *udpRelayConn) WriteTo(b []byte, addr net.Addr) (n int, err error) { Data: b, } - buf := c.wbuf[:] + buf := bufpool.Get(c.bufferSize) + defer bufpool.Put(buf) nn, err := dgram.WriteTo(bytes.NewBuffer(buf[:0])) if err != nil { diff --git a/connector/socks/v5/connector.go b/connector/socks/v5/connector.go index 5975a316..4d09591a 100644 --- a/connector/socks/v5/connector.go +++ b/connector/socks/v5/connector.go @@ -225,9 +225,10 @@ func (c *socks5Connector) relayUDP(ctx context.Context, conn net.Conn, addr net. } return &udpRelayConn{ - udpConn: cc.(*net.UDPConn), - tcpConn: conn, - taddr: addr, - logger: log, + udpConn: cc.(*net.UDPConn), + tcpConn: conn, + taddr: addr, + bufferSize: c.md.udpBufferSize, + logger: log, }, nil } diff --git a/connector/socks/v5/metadata.go b/connector/socks/v5/metadata.go index 50556347..f908660f 100644 --- a/connector/socks/v5/metadata.go +++ b/connector/socks/v5/metadata.go @@ -8,10 +8,15 @@ import ( mdutil "github.com/go-gost/x/metadata/util" ) +const ( + defaultUDPBufferSize = 4096 +) + type metadata struct { connectTimeout time.Duration noTLS bool relay string + udpBufferSize int udpTimeout time.Duration muxCfg *mux.Config } @@ -20,6 +25,10 @@ func (c *socks5Connector) parseMetadata(md mdata.Metadata) (err error) { c.md.connectTimeout = mdutil.GetDuration(md, "timeout") c.md.noTLS = mdutil.GetBool(md, "notls") c.md.relay = mdutil.GetString(md, "relay") + c.md.udpBufferSize = mdutil.GetInt(md, "udp.bufferSize", "udpBufferSize") + if c.md.udpBufferSize <= 0 { + c.md.udpBufferSize = defaultUDPBufferSize + } c.md.udpTimeout = mdutil.GetDuration(md, "udp.timeout") c.md.muxCfg = &mux.Config{ diff --git a/connector/ss/udp/connector.go b/connector/ss/udp/connector.go index ae98b999..4e6a74b4 100644 --- a/connector/ss/udp/connector.go +++ b/connector/ss/udp/connector.go @@ -85,7 +85,7 @@ func (c *ssuConnector) Connect(ctx context.Context, conn net.Conn, network, addr } // standard UDP relay - return ss.UDPClientConn(pc, conn.RemoteAddr(), taddr), nil + return ss.UDPClientConn(pc, conn.RemoteAddr(), taddr, c.md.udpBufferSize), nil } if c.cipher != nil { diff --git a/connector/ss/udp/metadata.go b/connector/ss/udp/metadata.go index 5227b1a5..d8ec3061 100644 --- a/connector/ss/udp/metadata.go +++ b/connector/ss/udp/metadata.go @@ -10,16 +10,13 @@ import ( type metadata struct { key string connectTimeout time.Duration + udpBufferSize int } func (c *ssuConnector) parseMetadata(md mdata.Metadata) (err error) { - const ( - key = "key" - connectTimeout = "timeout" - ) - - c.md.key = mdutil.GetString(md, key) - c.md.connectTimeout = mdutil.GetDuration(md, connectTimeout) + c.md.key = mdutil.GetString(md, "key") + c.md.connectTimeout = mdutil.GetDuration(md, "timeout", "connectTimeout") + c.md.udpBufferSize = mdutil.GetInt(md, "udpBufferSize", "udp.bufferSize") return } diff --git a/handler/http/metadata.go b/handler/http/metadata.go index e704e07d..b15a83a8 100644 --- a/handler/http/metadata.go +++ b/handler/http/metadata.go @@ -25,6 +25,7 @@ type metadata struct { compression bool probeResistance *probeResistance enableUDP bool + udpBufferSize int header http.Header hash string authBasicRealm string @@ -77,6 +78,7 @@ func (h *httpHandler) parseMetadata(md mdata.Metadata) error { } } h.md.enableUDP = mdutil.GetBool(md, "udp") + h.md.udpBufferSize = mdutil.GetInt(md, "udp.bufferSize", "udpBufferSize") h.md.hash = mdutil.GetString(md, "hash") h.md.authBasicRealm = mdutil.GetString(md, "authBasicRealm") diff --git a/handler/http/udp.go b/handler/http/udp.go index 36cd4bbe..dc8aecfb 100644 --- a/handler/http/udp.go +++ b/handler/http/udp.go @@ -75,6 +75,7 @@ func (h *httpHandler) handleUDP(ctx context.Context, conn net.Conn, ro *xrecorde relay := udp.NewRelay(socks.UDPTunServerConn(conn), pc). WithBypass(h.options.Bypass). + WithBufferSize(h.md.udpBufferSize). WithLogger(log) t := time.Now() diff --git a/handler/relay/bind.go b/handler/relay/bind.go index 2ff7aa0a..fab19d34 100644 --- a/handler/relay/bind.go +++ b/handler/relay/bind.go @@ -205,6 +205,7 @@ func (h *relayHandler) bindUDP(ctx context.Context, conn net.Conn, network, addr r := udp.NewRelay(relay_util.UDPTunServerConn(conn), pc). WithBypass(h.options.Bypass). + WithBufferSize(h.md.udpBufferSize). WithLogger(log) t := time.Now() diff --git a/handler/relay/metadata.go b/handler/relay/metadata.go index d4e84671..c7d09838 100644 --- a/handler/relay/metadata.go +++ b/handler/relay/metadata.go @@ -14,11 +14,12 @@ import ( ) type metadata struct { - readTimeout time.Duration - enableBind bool - noDelay bool - hash string - muxCfg *mux.Config + readTimeout time.Duration + udpBufferSize int + enableBind bool + noDelay bool + hash string + muxCfg *mux.Config observerPeriod time.Duration observerResetTraffic bool @@ -43,6 +44,7 @@ func (h *relayHandler) parseMetadata(md mdata.Metadata) (err error) { h.md.readTimeout = 15 * time.Second } + h.md.udpBufferSize = mdutil.GetInt(md, "udp.bufferSize", "udpBufferSize") h.md.enableBind = mdutil.GetBool(md, "bind") h.md.noDelay = mdutil.GetBool(md, "nodelay") h.md.hash = mdutil.GetString(md, "hash") diff --git a/handler/router/associate.go b/handler/router/associate.go index cad7e039..d6f20e7e 100644 --- a/handler/router/associate.go +++ b/handler/router/associate.go @@ -121,9 +121,11 @@ func (h *routerHandler) handleAssociate(ctx context.Context, conn net.Conn, netw log.Debugf("%s/%s: router=%s, connector=%s, weight=%d established", host, network, routerID, connectorID, connectorID.Weight()) - var b [MaxMessageSize]byte + b := bufpool.Get(h.md.bufferSize) + defer bufpool.Put(b) + for { - n, err := conn.Read(b[:]) + n, err := conn.Read(b) if err != nil { if err == io.EOF { return nil diff --git a/handler/router/entrypoint.go b/handler/router/entrypoint.go index 0f952dfd..8b794a4b 100644 --- a/handler/router/entrypoint.go +++ b/handler/router/entrypoint.go @@ -3,16 +3,18 @@ package router import ( "bytes" + "github.com/go-gost/core/common/bufpool" "github.com/go-gost/core/logger" "github.com/go-gost/relay" ) func (h *routerHandler) handleEntrypoint(log logger.Logger) error { - var buf [MaxMessageSize]byte + buf := bufpool.Get(h.md.bufferSize) + defer bufpool.Put(buf) for { err := func() error { - n, addr, err := h.epConn.ReadFrom(buf[:]) + n, addr, err := h.epConn.ReadFrom(buf) if err != nil { return err } diff --git a/handler/router/metadata.go b/handler/router/metadata.go index 923b3561..e3c125b0 100644 --- a/handler/router/metadata.go +++ b/handler/router/metadata.go @@ -1,7 +1,6 @@ package router import ( - "math" "time" "github.com/go-gost/core/ingress" @@ -13,13 +12,14 @@ import ( ) const ( - MaxMessageSize = math.MaxUint16 defaultTTL = 15 * time.Second + defaultBufferSize = 4096 defaultCacheExpiration = time.Second ) type metadata struct { readTimeout time.Duration + bufferSize int entryPoint string ingress ingress.Ingress @@ -40,6 +40,10 @@ type metadata struct { func (h *routerHandler) parseMetadata(md mdata.Metadata) (err error) { h.md.readTimeout = mdutil.GetDuration(md, "readTimeout") + h.md.bufferSize = mdutil.GetInt(md, "router.bufferSize", "bufferSize") + if h.md.bufferSize <= 0 { + h.md.bufferSize = defaultBufferSize + } h.md.entryPoint = mdutil.GetString(md, "entrypoint") h.md.ingress = registry.IngressRegistry().Get(mdutil.GetString(md, "ingress")) diff --git a/handler/socks/v5/handler.go b/handler/socks/v5/handler.go index 76bf9ab8..7f13e3d9 100644 --- a/handler/socks/v5/handler.go +++ b/handler/socks/v5/handler.go @@ -145,6 +145,7 @@ func (h *socks5Handler) Handle(ctx context.Context, conn net.Conn, opts ...handl } log.WithFields(map[string]any{ + "network": ro.Network, "duration": time.Since(start), "inputBytes": ro.InputBytes, "outputBytes": ro.OutputBytes, @@ -186,7 +187,7 @@ func (h *socks5Handler) Handle(ctx context.Context, conn net.Conn, opts ...handl return h.handleMuxBind(ctx, conn, "tcp", address, ro, log) case gosocks5.CmdUdp: ro.Network = "udp" - return h.handleUDP(ctx, conn, ro, log) + return h.handleUDP(ctx, conn, "udp", ro, log) case socks.CmdUDPTun: ro.Network = "udp" return h.handleUDPTun(ctx, conn, "udp", address, ro, log) diff --git a/handler/socks/v5/metadata.go b/handler/socks/v5/metadata.go index 34b6d994..2c75645c 100644 --- a/handler/socks/v5/metadata.go +++ b/handler/socks/v5/metadata.go @@ -19,6 +19,7 @@ type metadata struct { noTLS bool enableBind bool enableUDP bool + udpBufferSize int compatibilityMode bool hash string muxCfg *mux.Config @@ -50,6 +51,7 @@ func (h *socks5Handler) parseMetadata(md mdata.Metadata) (err error) { h.md.noTLS = mdutil.GetBool(md, "notls") h.md.enableBind = mdutil.GetBool(md, "bind") h.md.enableUDP = mdutil.GetBool(md, "udp") + h.md.udpBufferSize = mdutil.GetInt(md, "udp.bufferSize", "udpBufferSize") 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 d46d058f..ea155a0c 100644 --- a/handler/socks/v5/udp.go +++ b/handler/socks/v5/udp.go @@ -23,9 +23,10 @@ import ( xrecorder "github.com/go-gost/x/recorder" ) -func (h *socks5Handler) handleUDP(ctx context.Context, conn net.Conn, ro *xrecorder.HandlerRecorderObject, log logger.Logger) error { +func (h *socks5Handler) handleUDP(ctx context.Context, conn net.Conn, network string, ro *xrecorder.HandlerRecorderObject, log logger.Logger) error { log = log.WithFields(map[string]any{ - "cmd": "udp", + "network": network, + "cmd": network, }) if !h.md.enableUDP { @@ -39,7 +40,7 @@ func (h *socks5Handler) handleUDP(ctx context.Context, conn net.Conn, ro *xrecor Netns: h.options.Netns, } - cc, err := lc.ListenPacket(ctx, "udp", "") + cc, err := lc.ListenPacket(ctx, network, "") if err != nil { log.Error(err) reply := gosocks5.NewReply(gosocks5.Failure, nil) @@ -74,7 +75,7 @@ func (h *socks5Handler) handleUDP(ctx context.Context, conn net.Conn, ro *xrecor // obtain a udp connection var buf bytes.Buffer - c, err := h.options.Router.Dial(ctxvalue.ContextWithBuffer(ctx, &buf), "udp", "") // UDP association + c, err := h.options.Router.Dial(ctxvalue.ContextWithBuffer(ctx, &buf), network, "") // UDP association ro.Route = buf.String() if err != nil { log.Error(err) @@ -106,7 +107,7 @@ func (h *socks5Handler) handleUDP(ctx context.Context, conn net.Conn, ro *xrecor string(clientID), limiter.ServiceOption(h.options.Service), limiter.ScopeOption(limiter.ScopeClient), - limiter.NetworkOption("udp"), + limiter.NetworkOption(network), limiter.ClientOption(string(clientID)), limiter.SrcOption(conn.RemoteAddr().String()), ) @@ -119,8 +120,9 @@ func (h *socks5Handler) handleUDP(ctx context.Context, conn net.Conn, ro *xrecor } } - r := udp.NewRelay(socks.UDPConn(cc), pc). + r := udp.NewRelay(socks.UDPConn(cc, h.md.udpBufferSize), pc). WithBypass(h.options.Bypass). + WithBufferSize(h.md.udpBufferSize). WithLogger(log) go r.Run(ctx) diff --git a/handler/socks/v5/udp_tun.go b/handler/socks/v5/udp_tun.go index 4f4b067c..c7b1443c 100644 --- a/handler/socks/v5/udp_tun.go +++ b/handler/socks/v5/udp_tun.go @@ -22,7 +22,8 @@ import ( func (h *socks5Handler) handleUDPTun(ctx context.Context, conn net.Conn, network, address string, ro *xrecorder.HandlerRecorderObject, log logger.Logger) error { log = log.WithFields(map[string]any{ - "cmd": "udp-tun", + "network": network, + "cmd": "udp-tun", }) { @@ -66,7 +67,7 @@ func (h *socks5Handler) handleUDPTun(ctx context.Context, conn net.Conn, network // obtain a udp connection var buf bytes.Buffer - c, err := h.options.Router.Dial(ctxvalue.ContextWithBuffer(ctx, &buf), "udp", "") // UDP association + c, err := h.options.Router.Dial(ctxvalue.ContextWithBuffer(ctx, &buf), network, "") // UDP association ro.Route = buf.String() if err != nil { log.Error(err) @@ -94,7 +95,7 @@ func (h *socks5Handler) handleUDPTun(ctx context.Context, conn net.Conn, network Netns: h.options.Netns, } var err error - pc, err = lc.ListenPacket(ctx, "udp", bindAddr.String()) + pc, err = lc.ListenPacket(ctx, network, bindAddr.String()) if err != nil { log.Error(err) reply := gosocks5.NewReply(gosocks5.Failure, nil) @@ -133,6 +134,7 @@ func (h *socks5Handler) handleUDPTun(ctx context.Context, conn net.Conn, network r := udp.NewRelay(socks.UDPTunServerConn(conn), pc). WithBypass(h.options.Bypass). + WithBufferSize(h.md.udpBufferSize). WithLogger(log) t := time.Now() diff --git a/handler/ss/udp/handler.go b/handler/ss/udp/handler.go index 5cefa7ab..f135104f 100644 --- a/handler/ss/udp/handler.go +++ b/handler/ss/udp/handler.go @@ -7,6 +7,7 @@ import ( "net" "time" + "github.com/go-gost/core/common/bufpool" "github.com/go-gost/core/handler" "github.com/go-gost/core/logger" md "github.com/go-gost/core/metadata" @@ -112,7 +113,7 @@ func (h *ssuHandler) Handle(ctx context.Context, conn net.Conn, opts ...handler. pc = h.cipher.PacketConn(pc) } // standard UDP relay. - pc = ss.UDPServerConn(pc, conn.RemoteAddr()) + pc = ss.UDPServerConn(pc, conn.RemoteAddr(), h.md.udpBufferSize) } else { if h.cipher != nil { conn = ss.ShadowConn(h.cipher.StreamConn(conn), nil) @@ -150,11 +151,18 @@ func (h *ssuHandler) Handle(ctx context.Context, conn net.Conn, opts ...handler. func (h *ssuHandler) relayPacket(ctx context.Context, pc1, pc2 net.PacketConn, ro *xrecorder.HandlerRecorderObject, log logger.Logger) (err error) { errc := make(chan error, 2) + bufferSize := h.md.udpBufferSize + if bufferSize <= 0 { + bufferSize = defaultBufferSize + } + go func() { - var b [MaxMessageSize]byte + b := bufpool.Get(bufferSize) + defer bufpool.Put(b) + for { err := func() error { - n, addr, err := pc1.ReadFrom(b[:]) + n, addr, err := pc1.ReadFrom(b) if err != nil { return err } @@ -185,10 +193,12 @@ func (h *ssuHandler) relayPacket(ctx context.Context, pc1, pc2 net.PacketConn, r }() go func() { - var b [MaxMessageSize]byte + b := bufpool.Get(bufferSize) + defer bufpool.Put(b) + for { err := func() error { - n, raddr, err := pc2.ReadFrom(b[:]) + n, raddr, err := pc2.ReadFrom(b) if err != nil { return err } diff --git a/handler/ss/udp/metadata.go b/handler/ss/udp/metadata.go index 33145e84..b6da0734 100644 --- a/handler/ss/udp/metadata.go +++ b/handler/ss/udp/metadata.go @@ -1,7 +1,6 @@ package ss import ( - "math" "time" mdata "github.com/go-gost/core/metadata" @@ -9,22 +8,20 @@ import ( ) const ( - MaxMessageSize = math.MaxUint16 + defaultBufferSize = 4096 ) type metadata struct { - key string - readTimeout time.Duration + key string + readTimeout time.Duration + udpBufferSize int } func (h *ssuHandler) parseMetadata(md mdata.Metadata) (err error) { - const ( - key = "key" - readTimeout = "readTimeout" - ) - h.md.key = mdutil.GetString(md, key) - h.md.readTimeout = mdutil.GetDuration(md, readTimeout) + h.md.key = mdutil.GetString(md, "key") + h.md.readTimeout = mdutil.GetDuration(md, "readTimeout") + h.md.udpBufferSize = mdutil.GetInt(md, "udpBufferSize", "udp.bufferSize") return } diff --git a/handler/tungo/metadata.go b/handler/tungo/metadata.go index aa0b058c..aaab9114 100644 --- a/handler/tungo/metadata.go +++ b/handler/tungo/metadata.go @@ -10,8 +10,13 @@ import ( mdutil "github.com/go-gost/x/metadata/util" ) +const ( + defaultBufferSize = 4096 +) + type metadata struct { - udpTimeout time.Duration + udpTimeout time.Duration + udpBufferSize int sniffing bool sniffingUDP bool @@ -34,6 +39,7 @@ type metadata struct { func (h *tungoHandler) parseMetadata(md mdata.Metadata) (err error) { h.md.udpTimeout = mdutil.GetDuration(md, "udpTimeout", "tungo.udpTimeout") + h.md.udpBufferSize = mdutil.GetInt(md, "udp.bufferSize", "udpBufferSize") h.md.sniffing = mdutil.GetBool(md, "sniffing") h.md.sniffingUDP = mdutil.GetBool(md, "sniffing.udp") diff --git a/handler/tungo/stack.go b/handler/tungo/stack.go index 0948ccef..0943dca6 100644 --- a/handler/tungo/stack.go +++ b/handler/tungo/stack.go @@ -6,13 +6,13 @@ import ( "context" "crypto/tls" "io" - "math" "net" "net/netip" "strings" "sync" "time" + "github.com/go-gost/core/common/bufpool" "github.com/go-gost/core/handler" "github.com/go-gost/core/observer/stats" "github.com/go-gost/core/recorder" @@ -44,7 +44,8 @@ type transportHandler struct { procCancel context.CancelFunc // UDP session timeout. - udpTimeout time.Duration + udpTimeout time.Duration + udpBufferSize int sniffing bool sniffingUDP bool @@ -344,33 +345,46 @@ func (h *transportHandler) handleUDPConn(uc adapter.UDPConn) { t := time.Now() log.Infof("%s <-> %s", remoteAddr, dstAddr) - pipePacketData(conn, cc, h.sniffingUDP, ro, h.udpTimeout) + h.pipePacketData(conn, cc, ro) log.WithFields(map[string]any{ "duration": time.Since(t), }).Infof("%s >-< %s", remoteAddr, dstAddr) } -func pipePacketData(conn1, conn2 net.Conn, sniffing bool, ro *xrecorder.HandlerRecorderObject, timeout time.Duration) { +func (h *transportHandler) pipePacketData(conn1, conn2 net.Conn, ro *xrecorder.HandlerRecorderObject) { + timeout := h.udpTimeout if timeout <= 0 { timeout = udpSessionTimeout } + bufferSize := h.udpBufferSize + if bufferSize <= 0 { + bufferSize = defaultBufferSize + } + wg := sync.WaitGroup{} wg.Add(2) go func() { defer wg.Done() - copyPacketData(conn1, conn2, sniffing, false, ro, timeout) + + buf := bufpool.Get(bufferSize) + defer bufpool.Put(buf) + + copyPacketData(conn1, conn2, buf, h.sniffing, false, ro, timeout) }() go func() { defer wg.Done() - copyPacketData(conn2, conn1, sniffing, true, ro, timeout) + + buf := bufpool.Get(bufferSize) + defer bufpool.Put(buf) + + copyPacketData(conn2, conn1, buf, h.sniffing, true, ro, timeout) }() wg.Wait() } -func copyPacketData(dst, src net.Conn, sniffing bool, c2s bool, ro *xrecorder.HandlerRecorderObject, timeout time.Duration) error { - buf := make([]byte, math.MaxUint16/2) +func copyPacketData(dst, src net.Conn, buf []byte, sniffing bool, c2s bool, ro *xrecorder.HandlerRecorderObject, timeout time.Duration) error { isDNS := false for { diff --git a/internal/net/udp/listener.go b/internal/net/udp/listener.go index b7c34f49..0ad0b10e 100644 --- a/internal/net/udp/listener.go +++ b/internal/net/udp/listener.go @@ -176,6 +176,8 @@ func (c *conn) WriteTo(b []byte, addr net.Addr) (n int, err error) { if !c.keepalive { defer c.Close() } + c.SetIdle(false) + return c.PacketConn.WriteTo(b, addr) } @@ -212,6 +214,10 @@ func (c *conn) RemoteAddr() net.Addr { return c.remoteAddr } +func (c *conn) SetReadDeadline(t time.Time) error { + return nil +} + func (c *conn) IsIdle() bool { return atomic.LoadInt32(&c.idle) > 0 } diff --git a/internal/net/udp/relay.go b/internal/net/udp/relay.go index c14124ab..7ca268c3 100644 --- a/internal/net/udp/relay.go +++ b/internal/net/udp/relay.go @@ -2,23 +2,23 @@ package udp import ( "context" - "math" "net" "github.com/go-gost/core/bypass" + "github.com/go-gost/core/common/bufpool" "github.com/go-gost/core/logger" ) const ( - MaxMessageSize = math.MaxUint16 + defaultBufferSize = 4096 ) type Relay struct { - pc1 net.PacketConn - pc2 net.PacketConn - - bypass bypass.Bypass - logger logger.Logger + pc1 net.PacketConn + pc2 net.PacketConn + bufferSize int + bypass bypass.Bypass + logger logger.Logger } func NewRelay(pc1, pc2 net.PacketConn) *Relay { @@ -38,14 +38,26 @@ func (r *Relay) WithLogger(logger logger.Logger) *Relay { return r } +func (r *Relay) WithBufferSize(n int) *Relay { + r.bufferSize = n + return r +} + func (r *Relay) Run(ctx context.Context) (err error) { errc := make(chan error, 2) + bufferSize := r.bufferSize + if bufferSize <= 0 { + bufferSize = defaultBufferSize + } + go func() { - var b [MaxMessageSize]byte + b := bufpool.Get(bufferSize) + defer bufpool.Put(b) + for { err := func() error { - n, raddr, err := r.pc1.ReadFrom(b[:]) + n, raddr, err := r.pc1.ReadFrom(b) if err != nil { return err } @@ -64,7 +76,6 @@ func (r *Relay) Run(ctx context.Context) (err error) { if r.logger != nil { r.logger.Tracef("%s >>> %s data: %d", r.pc2.LocalAddr(), raddr, n) - } return nil @@ -78,10 +89,12 @@ func (r *Relay) Run(ctx context.Context) (err error) { }() go func() { - var b [MaxMessageSize]byte + b := bufpool.Get(bufferSize) + defer bufpool.Put(b) + for { err := func() error { - n, raddr, err := r.pc2.ReadFrom(b[:]) + n, raddr, err := r.pc2.ReadFrom(b) if err != nil { return err } @@ -100,7 +113,6 @@ func (r *Relay) Run(ctx context.Context) (err error) { if r.logger != nil { r.logger.Tracef("%s <<< %s data: %d", r.pc2.LocalAddr(), raddr, n) - } return nil diff --git a/internal/util/relay/conn.go b/internal/util/relay/conn.go index 865094d4..abe3c3df 100644 --- a/internal/util/relay/conn.go +++ b/internal/util/relay/conn.go @@ -1,10 +1,8 @@ package relay import ( - "bytes" "net" - "github.com/go-gost/core/common/bufpool" "github.com/go-gost/gosocks5" "github.com/go-gost/relay" ) @@ -108,89 +106,3 @@ func (c *udpTunConn) WriteTo(b []byte, addr net.Addr) (n int, err error) { func (c *udpTunConn) Write(b []byte) (n int, err error) { return c.WriteTo(b, c.taddr) } - -var ( - DefaultBufferSize = 4096 -) - -type udpConn struct { - net.PacketConn - raddr net.Addr - taddr net.Addr - bufferSize int -} - -func UDPConn(c net.PacketConn, bufferSize int) net.PacketConn { - return &udpConn{ - PacketConn: c, - bufferSize: bufferSize, - } -} - -// ReadFrom reads an UDP datagram. -// NOTE: for server side, -// the returned addr is the target address the client want to relay to. -func (c *udpConn) ReadFrom(b []byte) (n int, addr net.Addr, err error) { - rbuf := bufpool.Get(c.bufferSize) - defer bufpool.Put(rbuf) - - n, c.raddr, err = c.PacketConn.ReadFrom(rbuf) - if err != nil { - return - } - - socksAddr := gosocks5.Addr{} - header := gosocks5.UDPHeader{ - Addr: &socksAddr, - } - hlen, err := header.ReadFrom(bytes.NewReader(rbuf[:n])) - if err != nil { - return - } - n = copy(b, rbuf[hlen:n]) - - addr, err = net.ResolveUDPAddr("udp", socksAddr.String()) - return -} - -func (c *udpConn) Read(b []byte) (n int, err error) { - n, _, err = c.ReadFrom(b) - return -} - -func (c *udpConn) WriteTo(b []byte, addr net.Addr) (n int, err error) { - wbuf := bufpool.Get(c.bufferSize) - defer bufpool.Put(wbuf) - - socksAddr := gosocks5.Addr{} - if err = socksAddr.ParseFrom(addr.String()); err != nil { - return - } - - header := gosocks5.UDPHeader{ - Addr: &socksAddr, - } - dgram := gosocks5.UDPDatagram{ - Header: &header, - Data: b, - } - - buf := bytes.NewBuffer(wbuf[:0]) - _, err = dgram.WriteTo(buf) - if err != nil { - return - } - - _, err = c.PacketConn.WriteTo(buf.Bytes(), c.raddr) - n = len(b) - - return -} - -func (c *udpConn) Write(b []byte) (n int, err error) { - return c.WriteTo(b, c.taddr) -} - -func (c *udpConn) RemoteAddr() net.Addr { - return c.raddr -} diff --git a/internal/util/socks/conn.go b/internal/util/socks/conn.go index 49a91702..d0c888e7 100644 --- a/internal/util/socks/conn.go +++ b/internal/util/socks/conn.go @@ -2,9 +2,9 @@ package socks import ( "bytes" - "math" "net" + "github.com/go-gost/core/common/bufpool" "github.com/go-gost/gosocks5" ) @@ -86,18 +86,23 @@ func (c *udpTunConn) Write(b []byte) (n int, err error) { } const ( - MaxMessageSize = math.MaxUint16 + defaultBufferSize = 4096 ) type udpConn struct { net.PacketConn - raddr net.Addr - taddr net.Addr + raddr net.Addr + taddr net.Addr + bufferSize int } -func UDPConn(c net.PacketConn) net.PacketConn { +func UDPConn(c net.PacketConn, bufferSize int) net.PacketConn { + if bufferSize <= 0 { + bufferSize = defaultBufferSize + } return &udpConn{ PacketConn: c, + bufferSize: bufferSize, } } @@ -105,9 +110,10 @@ func UDPConn(c net.PacketConn) net.PacketConn { // NOTE: for server side, // the returned addr is the target address the client want to relay to. func (c *udpConn) ReadFrom(b []byte) (n int, addr net.Addr, err error) { - var rbuf [MaxMessageSize]byte + buf := bufpool.Get(c.bufferSize) + defer bufpool.Put(buf) - n, c.raddr, err = c.PacketConn.ReadFrom(rbuf[:]) + n, c.raddr, err = c.PacketConn.ReadFrom(buf) if err != nil { return } @@ -116,11 +122,11 @@ func (c *udpConn) ReadFrom(b []byte) (n int, addr net.Addr, err error) { header := gosocks5.UDPHeader{ Addr: &socksAddr, } - hlen, err := header.ReadFrom(bytes.NewReader(rbuf[:n])) + hlen, err := header.ReadFrom(bytes.NewReader(buf[:n])) if err != nil { return } - n = copy(b, rbuf[hlen:n]) + n = copy(b, buf[hlen:n]) addr, err = net.ResolveUDPAddr("udp", socksAddr.String()) return @@ -132,7 +138,6 @@ func (c *udpConn) Read(b []byte) (n int, err error) { } func (c *udpConn) WriteTo(b []byte, addr net.Addr) (n int, err error) { - socksAddr := gosocks5.Addr{} if err = socksAddr.ParseFrom(addr.String()); err != nil { return @@ -146,14 +151,16 @@ func (c *udpConn) WriteTo(b []byte, addr net.Addr) (n int, err error) { Data: b, } - var wbuf [MaxMessageSize]byte - buf := bytes.NewBuffer(wbuf[:][:0]) - _, err = dgram.WriteTo(buf) + buf := bufpool.Get(c.bufferSize) + defer bufpool.Put(buf) + + bw := bytes.NewBuffer(buf[:0]) + _, err = dgram.WriteTo(bw) if err != nil { return } - _, err = c.PacketConn.WriteTo(buf.Bytes(), c.raddr) + _, err = c.PacketConn.WriteTo(bw.Bytes(), c.raddr) n = len(b) return diff --git a/internal/util/ss/conn.go b/internal/util/ss/conn.go index 8627590c..c712aee6 100644 --- a/internal/util/ss/conn.go +++ b/internal/util/ss/conn.go @@ -2,14 +2,14 @@ package ss import ( "bytes" - "math" "net" + "github.com/go-gost/core/common/bufpool" "github.com/go-gost/gosocks5" ) const ( - MaxMessageSize = math.MaxUint16 + defaultBufferSize = 4096 ) var ( @@ -19,40 +19,52 @@ var ( type UDPConn struct { net.PacketConn - raddr net.Addr - taddr net.Addr + raddr net.Addr + taddr net.Addr + bufferSize int } -func UDPClientConn(c net.PacketConn, remoteAddr, targetAddr net.Addr) *UDPConn { +func UDPClientConn(c net.PacketConn, remoteAddr, targetAddr net.Addr, bufferSize int) *UDPConn { + if bufferSize <= 0 { + bufferSize = defaultBufferSize + } + return &UDPConn{ PacketConn: c, raddr: remoteAddr, taddr: targetAddr, + bufferSize: bufferSize, } } -func UDPServerConn(c net.PacketConn, remoteAddr net.Addr) *UDPConn { +func UDPServerConn(c net.PacketConn, remoteAddr net.Addr, bufferSize int) *UDPConn { + if bufferSize <= 0 { + bufferSize = defaultBufferSize + } + return &UDPConn{ PacketConn: c, raddr: remoteAddr, + bufferSize: bufferSize, } } func (c *UDPConn) ReadFrom(b []byte) (n int, addr net.Addr, err error) { - var rbuf [MaxMessageSize]byte + buf := bufpool.Get(c.bufferSize) + defer bufpool.Put(buf) - n, _, err = c.PacketConn.ReadFrom(rbuf[:]) + n, _, err = c.PacketConn.ReadFrom(buf) if err != nil { return } saddr := gosocks5.Addr{} - addrLen, err := saddr.ReadFrom(bytes.NewReader(rbuf[:n])) + addrLen, err := saddr.ReadFrom(bytes.NewReader(buf[:n])) if err != nil { return } - n = copy(b, rbuf[addrLen:n]) + n = copy(b, buf[addrLen:n]) addr, err = net.ResolveUDPAddr("udp", saddr.String()) return @@ -64,20 +76,21 @@ func (c *UDPConn) Read(b []byte) (n int, err error) { } func (c *UDPConn) WriteTo(b []byte, addr net.Addr) (n int, err error) { - var wbuf [MaxMessageSize]byte - socksAddr := gosocks5.Addr{} if err = socksAddr.ParseFrom(addr.String()); err != nil { return } - addrLen, err := socksAddr.Encode(wbuf[:]) + buf := bufpool.Get(c.bufferSize) + defer bufpool.Put(buf) + + addrLen, err := socksAddr.Encode(buf) if err != nil { return } - n = copy(wbuf[addrLen:], b) - _, err = c.PacketConn.WriteTo(wbuf[:addrLen+n], c.raddr) + n = copy(buf[addrLen:], b) + _, err = c.PacketConn.WriteTo(buf[:addrLen+n], c.raddr) return } diff --git a/listener/udp/metadata.go b/listener/udp/metadata.go index db8a5e13..1e5ec75e 100644 --- a/listener/udp/metadata.go +++ b/listener/udp/metadata.go @@ -9,47 +9,39 @@ import ( const ( defaultTTL = 5 * time.Second - defaultReadBufferSize = 8192 + defaultReadBufferSize = 4096 defaultReadQueueSize = 128 defaultBacklog = 128 ) type metadata struct { - readBufferSize int - readQueueSize int - backlog int - keepalive bool - ttl time.Duration + readBufferSize int + readQueueSize int + backlog int + keepalive bool + ttl time.Duration } func (l *udpListener) parseMetadata(md mdata.Metadata) (err error) { - const ( - readBufferSize = "readBufferSize" - readQueueSize = "readQueueSize" - backlog = "backlog" - keepalive = "keepalive" - ttl = "ttl" - ) + l.md.backlog = mdutil.GetInt(md, "backlog") + if l.md.backlog <= 0 { + l.md.backlog = defaultBacklog + } - l.md.ttl = mdutil.GetDuration(md, ttl) + l.md.keepalive = mdutil.GetBool(md, "keepalive") + l.md.ttl = mdutil.GetDuration(md, "ttl", "keepalive.ttl") if l.md.ttl <= 0 { l.md.ttl = defaultTTL } - l.md.readBufferSize = mdutil.GetInt(md, readBufferSize) + l.md.readBufferSize = mdutil.GetInt(md, "readBufferSize", "udp.bufferSize") if l.md.readBufferSize <= 0 { l.md.readBufferSize = defaultReadBufferSize } - l.md.readQueueSize = mdutil.GetInt(md, readQueueSize) + l.md.readQueueSize = mdutil.GetInt(md, "readQueueSize", "recvQueueSize") if l.md.readQueueSize <= 0 { l.md.readQueueSize = defaultReadQueueSize } - l.md.backlog = mdutil.GetInt(md, backlog) - if l.md.backlog <= 0 { - l.md.backlog = defaultBacklog - } - l.md.keepalive = mdutil.GetBool(md, keepalive) - return }