Merge remote-tracking branch 'origin/pr/96'

handler/ss/handler.go

@@ -5,6 +5,7 @@ import (
 	"bytes"
 	"context"
 	"crypto/tls"
+	"errors"
 	"net"
 	"time"
 
@@ -14,6 +15,7 @@ import (
 	"github.com/go-gost/core/observer/stats"
 	"github.com/go-gost/core/recorder"
 	"github.com/go-gost/go-shadowsocks2/core"
+	"github.com/go-gost/go-shadowsocks2/socks"
 	"github.com/go-gost/go-shadowsocks2/utils"
 	xctx "github.com/go-gost/x/ctx"
 	ictx "github.com/go-gost/x/internal/ctx"
@@ -54,19 +56,22 @@ func (h *ssHandler) Init(md md.Metadata) (err error) {
 	if err = h.parseMetadata(md); err != nil {
 		return
 	}
-	if h.options.Auth != nil {
-		method := h.options.Auth.Username()
-		password, _ := h.options.Auth.Password()
 
-		serverConfig, err := utils.NewServerConfig(method, password, h.md.users)
-		if err != nil {
-			return err
-		}
-		h.server = core.NewTCPServer(serverConfig)
-		err = h.server.Init()
-		if err != nil {
-			return err
-		}
+	if h.options.Auth == nil {
+		return errors.New("ss: auth is required")
+	}
+
+	method := h.options.Auth.Username()
+	password, _ := h.options.Auth.Password()
+
+	serverConfig, err := utils.NewServerConfig(method, password, h.md.users)
+	if err != nil {
+		return err
+	}
+	h.server = core.NewTCPServer(serverConfig)
+	err = h.server.Init()
+	if err != nil {
+		return err
 	}
 
 	for _, ro := range h.options.Recorders {
@@ -139,7 +144,13 @@ func (h *ssHandler) Handle(ctx context.Context, conn net.Conn, opts ...handler.H
 	if err != nil {
 		return
 	}
-	target := wrappedConn.Target()
+	var target socks.Addr
+	if tc, ok := wrappedConn.(interface{ Target() socks.Addr }); ok {
+		target = tc.Target()
+	}
+	if len(target) == 0 {
+		return
+	}
 	conn = wrappedConn
 
 	conn.SetReadDeadline(time.Now().Add(h.md.readTimeout))

handler/ss/udp/handler.go

@@ -5,12 +5,10 @@ import (
 	"context"
 	"errors"
 	"net"
-	"net/netip"
 	"sync"
 	"time"
 
 	"github.com/go-gost/core/bypass"
-	"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"
@@ -19,6 +17,7 @@ import (
 	"github.com/go-gost/go-shadowsocks2/utils"
 	xctx "github.com/go-gost/x/ctx"
 	ictx "github.com/go-gost/x/internal/ctx"
+	"github.com/go-gost/x/internal/net/udp"
 	"github.com/go-gost/x/internal/util/relay"
 	rate_limiter "github.com/go-gost/x/limiter/rate"
 	xrecorder "github.com/go-gost/x/recorder"
@@ -54,25 +53,28 @@ func (h *ssuHandler) Init(md md.Metadata) (err error) {
 		return
 	}
 
-	if h.options.Auth != nil {
-		method := h.options.Auth.Username()
-		password, _ := h.options.Auth.Password()
+	if h.options.Auth == nil {
+		return errors.New("ss: auth is required")
+	}
 
-		serverConfig, err := utils.NewServerConfig(method, password, h.md.users)
-		if err != nil {
-			return err
-		}
-		h.udpServer = core.NewUDPServer(serverConfig, time.Duration(60*time.Second))
-		err = h.udpServer.Init()
-		if err != nil {
-			return err
-		}
+	method := h.options.Auth.Username()
+	password, _ := h.options.Auth.Password()
 
-		h.tcpServer = core.NewTCPServer(serverConfig)
-		err = h.udpServer.Init()
-		if err != nil {
-			return err
-		}
+	serverConfig, err := utils.NewServerConfig(method, password, h.md.users)
+	if err != nil {
+		return err
+	}
+	serverConfig.UDPTimeout = time.Minute
+	h.udpServer = core.NewUDPServer(serverConfig)
+	err = h.udpServer.Init()
+	if err != nil {
+		return err
+	}
+
+	h.tcpServer = core.NewTCPServer(serverConfig)
+	err = h.tcpServer.Init()
+	if err != nil {
+		return err
 	}
 
 	for _, ro := range h.options.Recorders {
@@ -141,198 +143,148 @@ func (h *ssuHandler) Handle(ctx context.Context, conn net.Conn, opts ...handler.
 		// WriteTo will write data to original remote addr instead of the addr parameter passed to
 		pc = relay.UDPTunServerConn(conn)
 
-		return h.relayPacketTCP(ctx, pc, ro, log)
-	} else {
-		return h.relayPacketUDP(ctx, pc, ro, log)
-	}
-}
-
-// UDP over TCP
-func (h *ssuHandler) relayPacketTCP(ctx context.Context, src net.PacketConn, ro *xrecorder.HandlerRecorderObject, log logger.Logger) error {
-	bufferSize := h.md.udpBufferSize
-	if bufferSize <= 0 {
-		bufferSize = defaultBufferSize
-	}
-
-	b := bufpool.Get(bufferSize)
-	defer bufpool.Put(b)
-
-	var dstConn net.Conn
-	defer func() {
-		if dstConn != nil {
-			dstConn.Close()
-		}
-	}()
-
-	for {
-		n, targetAddr, err := src.ReadFrom(b)
+		// UDP over TCP still behaves like a regular packet tunnel.
+		var buf bytes.Buffer
+		c, err := h.options.Router.Dial(ictx.ContextWithBuffer(ctx, &buf), "udp", "")
+		ro.Route = buf.String()
 		if err != nil {
+			log.Error(err)
+			return err
+		}
+		defer c.Close()
+
+		cc, ok := c.(net.PacketConn)
+		if !ok {
+			err := errors.New("ss: wrong connection type")
+			log.Error(err)
 			return err
 		}
 
-		if ro.Host == "" {
-			ro.Host = targetAddr.String()
-		}
+		r := udp.NewRelay(pc, cc).
+			WithService(h.options.Service).
+			WithBypass(h.options.Bypass).
+			WithBufferSize(h.md.udpBufferSize).
+			WithLogger(log)
 
-		if h.options.Bypass != nil && h.options.Bypass.Contains(ctx, targetAddr.Network(), targetAddr.String(), bypass.WithService(h.options.Service)) {
-			log.Warn("bypass: ", targetAddr)
-			return nil
-		}
-
-		if dstConn == nil {
-			// obtain a udp connection
-			var buf bytes.Buffer
-			dstConn, err := h.options.Router.Dial(ictx.ContextWithBuffer(ctx, &buf), "udp", "") // UDP association
-			ro.Route = buf.String()
-			if err != nil {
-				log.Error(err)
-				return err
-			}
-
-			cc, ok := dstConn.(net.PacketConn)
-			if !ok {
-				err := errors.New("ss: wrong connection type")
-				log.Error(err)
-				return err
-			}
-
-			log.Infof("%s <-> %s", src.LocalAddr(), cc.LocalAddr())
-
-			go func() {
-				defer dstConn.Close()
-				for {
-					n, raddr, err := cc.ReadFrom(b)
-					if err != nil {
-						log.Warnf("failed to read response from %v: %v", raddr, err)
-						return
-					}
-
-					if h.options.Bypass != nil && h.options.Bypass.Contains(ctx, raddr.Network(), raddr.String(), bypass.WithService(h.options.Service)) {
-						log.Warn("bypass: ", raddr)
-						return
-					}
-
-					if _, err = src.WriteTo(b[:n], targetAddr); err != nil {
-						log.Warnf("failed to write response from %v: %v", targetAddr, err)
-						return
-					}
-
-					log.Tracef("%s <<< %s data: %d",
-						cc.LocalAddr(), raddr, n)
-				}
-			}()
-		}
-
-		if _, err = dstConn.(net.PacketConn).WriteTo(b[:n], targetAddr); err != nil {
-			return err
-		}
-
-		log.Tracef("%s >>> %s data: %d",
-			dstConn.(net.PacketConn).LocalAddr(), targetAddr, n)
+		return r.Run(ctx)
 	}
+
+	return h.relayPacketUDP(ctx, h.udpServer.WrapConn(pc), ro, log)
 }
 
-// standard udp relay
-// Dataflow: src (encrypted data) <-> dst (plaintext data)
 func (h *ssuHandler) relayPacketUDP(ctx context.Context, src net.PacketConn, ro *xrecorder.HandlerRecorderObject, log logger.Logger) error {
 	bufferSize := h.md.udpBufferSize
 	if bufferSize <= 0 {
 		bufferSize = defaultBufferSize
 	}
 
-	b := bufpool.Get(bufferSize)
-	defer bufpool.Put(b)
-
+	b := make([]byte, bufferSize)
 	for {
 		n, addr, err := src.ReadFrom(b)
 		if err != nil {
 			return err
 		}
 
-		clientAddr, err := netip.ParseAddrPort(addr.String())
-		if err != nil {
-			return err
-		}
-		session, payload, err := h.udpServer.Inbound(b[:n], clientAddr)
-		if err != nil {
-			return err
+		ctxAddr, ok := addr.(interface {
+			Session() core.UDPSession
+			ClientAddr() net.Addr
+			TargetAddr() net.Addr
+		})
+		if !ok {
+			return errors.New("ss: missing udp session context")
 		}
 
-		targetAddr, err := net.ResolveUDPAddr("udp", session.Target().String())
-		if ro.Host == "" {
+		targetAddr := ctxAddr.TargetAddr()
+		if ro.Host == "" && targetAddr != nil {
 			ro.Host = targetAddr.String()
 		}
 
-		if h.options.Bypass != nil && h.options.Bypass.Contains(ctx, targetAddr.Network(), targetAddr.String(), bypass.WithService(h.options.Service)) {
-			log.Warn("bypass: ", addr)
+		if h.options.Bypass != nil && targetAddr != nil && h.options.Bypass.Contains(ctx, "udp", targetAddr.String(), bypass.WithService(h.options.Service)) {
+			log.Warn("bypass: ", targetAddr)
 			return nil
 		}
+
+		session := ctxAddr.Session()
+		if session == nil {
+			return errors.New("ss: udp session not found")
+		}
+
+		dstConn, err := h.packetConnForSession(ctx, src, session, ro, log)
 		if err != nil {
 			return err
 		}
 
-		dstConn, ok := h.connMap.Load(session.Hash())
-		if !ok {
-			// obtain a udp connection
-			var buf bytes.Buffer
-			c, err := h.options.Router.Dial(ictx.ContextWithBuffer(ctx, &buf), "udp", "") // UDP association
-			ro.Route = buf.String()
-			if err != nil {
-				log.Error(err)
-				return err
-			}
-
-			cc, ok := c.(net.PacketConn)
-			if !ok {
-				err := errors.New("ss: wrong connection type")
-				log.Error(err)
-				return err
-			}
-			h.connMap.Store(session.Hash(), c)
-
-			dstConn = cc
-
-			log.Infof("%s <-> %s", src.LocalAddr(), cc.LocalAddr())
-
-			go func() {
-				defer func() {
-					c.Close()
-					h.connMap.Delete(session.Hash())
-				}()
-
-				for {
-					n, raddr, err := dstConn.(net.PacketConn).ReadFrom(b)
-					if err != nil {
-						log.Warnf("failed to read response: %v", err)
-						return
-					}
-
-					clientAddr := session.ClientAddr()
-					if h.options.Bypass != nil && h.options.Bypass.Contains(ctx, raddr.Network(), raddr.String(), bypass.WithService(h.options.Service)) {
-						log.Warn("bypass: ", raddr)
-						return
-					}
-					encrypted, err := h.udpServer.Outbound(b[:n], session)
-					if _, err = src.WriteTo(encrypted, net.UDPAddrFromAddrPort(clientAddr)); err != nil {
-						log.Warnf("failed to write response to %v: %v", clientAddr, err)
-						return
-					}
-
-					log.Tracef("%s <<< %s data: %d",
-						dstConn.(net.PacketConn).LocalAddr(), raddr, n)
-				}
-			}()
-		}
-
-		if n, err = dstConn.(net.PacketConn).WriteTo(payload, targetAddr); err != nil {
+		if _, err = dstConn.WriteTo(b[:n], targetAddr); err != nil {
 			return err
 		}
 
-		log.Tracef("%s >>> %s data: %d",
-			dstConn.(net.PacketConn).LocalAddr(), targetAddr, n)
+		log.Tracef("%s >>> %s data: %d", dstConn.LocalAddr(), targetAddr, n)
 	}
 }
 
+func (h *ssuHandler) packetConnForSession(ctx context.Context, src net.PacketConn, session core.UDPSession, ro *xrecorder.HandlerRecorderObject, log logger.Logger) (net.PacketConn, error) {
+	if cc, ok := h.connMap.Load(session.Hash()); ok {
+		return cc.(net.PacketConn), nil
+	}
+
+	var buf bytes.Buffer
+	c, err := h.options.Router.Dial(ictx.ContextWithBuffer(ctx, &buf), "udp", "")
+	ro.Route = buf.String()
+	if err != nil {
+		log.Error(err)
+		return nil, err
+	}
+
+	cc, ok := c.(net.PacketConn)
+	if !ok {
+		c.Close()
+		err := errors.New("ss: wrong connection type")
+		log.Error(err)
+		return nil, err
+	}
+
+	actual, loaded := h.connMap.LoadOrStore(session.Hash(), cc)
+	if loaded {
+		c.Close()
+		return actual.(net.PacketConn), nil
+	}
+
+	go func() {
+		defer func() {
+			cc.Close()
+			h.connMap.Delete(session.Hash())
+		}()
+
+		b := make([]byte, h.md.udpBufferSize)
+		if len(b) == 0 {
+			b = make([]byte, defaultBufferSize)
+		}
+
+		for {
+			n, raddr, err := cc.ReadFrom(b)
+			if err != nil {
+				log.Warnf("failed to read response: %v", err)
+				return
+			}
+
+			if h.options.Bypass != nil && h.options.Bypass.Contains(ctx, "udp", raddr.String(), bypass.WithService(h.options.Service)) {
+				log.Warn("bypass: ", raddr)
+				return
+			}
+
+			addr := core.NewUDPServerPacketAddr(raddr, net.UDPAddrFromAddrPort(session.ClientAddr()), session)
+			if _, err = src.WriteTo(b[:n], addr); err != nil {
+				log.Warnf("failed to write response to %v: %v", session.ClientAddr(), err)
+				return
+			}
+
+			log.Tracef("%s <<< %s data: %d", cc.LocalAddr(), raddr, n)
+		}
+	}()
+
+	return cc, nil
+}
+
 func (h *ssuHandler) checkRateLimit(addr net.Addr) bool {
 	if h.options.RateLimiter == nil {
 		return true

handler/ss/udp/metadata.go

@@ -24,6 +24,9 @@ func (h *ssuHandler) parseMetadata(md mdata.Metadata) (err error) {
 
 	h.md.key = mdutil.GetString(md, "key")
 	h.md.readTimeout = mdutil.GetDuration(md, "readTimeout")
+	if h.md.readTimeout == 0 {
+		h.md.readTimeout = time.Minute // Default read timeout
+	}
 	h.md.udpBufferSize = mdutil.GetInt(md, "udpBufferSize", "udp.bufferSize")
 
 	usersMap := mdutil.GetStringMapString(md, "users")