refactor: move IP validation to the PacketListener (#229)

* refactor: move IP validation to the PacketListener

* Rename `validatePacket()` to reflect what it currently does.

* Reorder for consistency.

* Handle the `EOF` case and stop reading from the connection.

* Pass through the connection errors from the IP validator.

* Address review comments.

* Rename `timeout` to `natTimeout` in test.

* Only resolve the `net.Addr` if it's not already a `UDPAddr`.

* Use `MakeNetAddr()` instead of resolving when extracting the addr.

Author: sbruens

cmd/outline-ss-server/main.go

@@ -268,7 +268,7 @@ func (s *OutlineServer) runConfig(config Config) (func() error, error) {
 					service.WithCiphers(ciphers),
 					service.WithMetrics(s.serviceMetrics),
 					service.WithReplayCache(&s.replayCache),
-					service.WithPacketListener(service.MakeTargetUDPListener(s.natTimeout, 0)),
+					service.WithPacketListener(service.MakeTargetUDPListener(onet.RequirePublicIP, s.natTimeout, 0)),
 					service.WithLogger(slog.Default()),
 				)
 				ln, err := lnSet.ListenStream(addr)
@@ -301,7 +301,7 @@ func (s *OutlineServer) runConfig(config Config) (func() error, error) {
 					service.WithMetrics(s.serviceMetrics),
 					service.WithReplayCache(&s.replayCache),
 					service.WithStreamDialer(service.MakeValidatingTCPStreamDialer(onet.RequirePublicIP, serviceConfig.Dialer.Fwmark)),
-					service.WithPacketListener(service.MakeTargetUDPListener(s.natTimeout, serviceConfig.Dialer.Fwmark)),
+					service.WithPacketListener(service.MakeTargetUDPListener(onet.RequirePublicIP, s.natTimeout, serviceConfig.Dialer.Fwmark)),
 					service.WithLogger(slog.Default()),
 				)
 				if err != nil {

internal/integration_test/integration_test.go

@@ -35,7 +35,10 @@ import (
 	"github.com/stretchr/testify/require"
 )
 
-const maxUDPPacketSize = 64 * 1024
+const (
+	maxUDPPacketSize = 64 * 1024
+	natTimeout       = 5 * time.Minute
+)
 
 func init() {
 	logging.SetLevel(logging.INFO, "")
@@ -321,7 +324,7 @@ func TestUDPEcho(t *testing.T) {
 	}
 	proxy := service.NewAssociationHandler(cipherList, &fakeShadowsocksMetrics{})
 
-	proxy.SetTargetIPValidator(allowAll)
+	proxy.SetTargetPacketListener(service.MakeTargetUDPListener(allowAll, natTimeout, 0))
 	natMetrics := &natTestMetrics{}
 	associationMetrics := &fakeUDPAssociationMetrics{}
 	go service.PacketServe(proxyConn, func(ctx context.Context, conn net.Conn) {
@@ -548,7 +551,7 @@ func BenchmarkUDPEcho(b *testing.B) {
 		b.Fatal(err)
 	}
 	proxy := service.NewAssociationHandler(cipherList, &fakeShadowsocksMetrics{})
-	proxy.SetTargetIPValidator(allowAll)
+	proxy.SetTargetPacketListener(service.MakeTargetUDPListener(allowAll, natTimeout, 0))
 	done := make(chan struct{})
 	go func() {
 		service.PacketServe(server, func(ctx context.Context, conn net.Conn) {
@@ -594,7 +597,7 @@ func BenchmarkUDPManyKeys(b *testing.B) {
 		b.Fatal(err)
 	}
 	proxy := service.NewAssociationHandler(cipherList, &fakeShadowsocksMetrics{})
-	proxy.SetTargetIPValidator(allowAll)
+	proxy.SetTargetPacketListener(service.MakeTargetUDPListener(allowAll, natTimeout, 0))
 	done := make(chan struct{})
 	go func() {
 		service.PacketServe(proxyConn, func(ctx context.Context, conn net.Conn) {

service/udp.go

@@ -109,7 +109,7 @@ func NewAssociationHandler(cipherList CipherList, ssMetrics ShadowsocksConnMetri
 		ciphers:           cipherList,
 		ssm:               ssMetrics,
 		targetIPValidator: onet.RequirePublicIP,
-		targetListener:    MakeTargetUDPListener(defaultNatTimeout, 0),
+		targetListener:    MakeTargetUDPListener(onet.RequirePublicIP, defaultNatTimeout, 0),
 	}
 }
 
@@ -118,8 +118,6 @@ type AssociationHandler interface {
 	HandleAssociation(ctx context.Context, conn net.Conn, assocMetrics UDPAssociationMetrics)
 	// SetLogger sets the logger used to log messages. Uses a no-op logger if nil.
 	SetLogger(l *slog.Logger)
-	// SetTargetIPValidator sets the function to be used to validate the target IP addresses.
-	SetTargetIPValidator(targetIPValidator onet.TargetIPValidator)
 	// SetTargetPacketListener sets the packet listener to use for target connections.
 	SetTargetPacketListener(targetListener transport.PacketListener)
 }
@@ -131,10 +129,6 @@ func (h *associationHandler) SetLogger(l *slog.Logger) {
 	h.logger = l
 }
 
-func (h *associationHandler) SetTargetIPValidator(targetIPValidator onet.TargetIPValidator) {
-	h.targetIPValidator = targetIPValidator
-}
-
 func (h *associationHandler) SetTargetPacketListener(targetListener transport.PacketListener) {
 	h.targetListener = targetListener
 }
@@ -176,7 +170,7 @@ func (h *associationHandler) HandleAssociation(ctx context.Context, clientConn n
 			}
 
 			var payload []byte
-			var tgtUDPAddr *net.UDPAddr
+			var tgtAddr net.Addr
 			if targetConn == nil {
 				ip := clientConn.RemoteAddr().(*net.UDPAddr).AddrPort().Addr()
 				var textData []byte
@@ -194,7 +188,7 @@ func (h *associationHandler) HandleAssociation(ctx context.Context, clientConn n
 				assocMetrics.AddAuthentication(keyID)
 
 				var onetErr *onet.ConnectionError
-				if payload, tgtUDPAddr, onetErr = h.validatePacket(textData); onetErr != nil {
+				if payload, tgtAddr, onetErr = h.extractPayloadAndDestination(textData); onetErr != nil {
 					return onetErr
 				}
 
@@ -219,15 +213,15 @@ func (h *associationHandler) HandleAssociation(ctx context.Context, clientConn n
 				}
 
 				var onetErr *onet.ConnectionError
-				if payload, tgtUDPAddr, onetErr = h.validatePacket(textData); onetErr != nil {
+				if payload, tgtAddr, onetErr = h.extractPayloadAndDestination(textData); onetErr != nil {
 					return onetErr
 				}
 			}
 
 			debugUDP(l, "Proxy exit.")
-			proxyTargetBytes, err = targetConn.WriteTo(payload, tgtUDPAddr) // accept only UDPAddr despite the signature
+			proxyTargetBytes, err = targetConn.WriteTo(payload, tgtAddr)
 			if err != nil {
-				return onet.NewConnectionError("ERR_WRITE", "Failed to write to target", err)
+				return ensureConnectionError(err, "ERR_WRITE", "Failed to write to target")
 			}
 			return nil
 		}()
@@ -246,21 +240,17 @@ func (h *associationHandler) HandleAssociation(ctx context.Context, clientConn n
 	}
 }
 
-// Given the decrypted contents of a UDP packet, return
-// the payload and the destination address, or an error if
-// this packet cannot or should not be forwarded.
-func (h *associationHandler) validatePacket(textData []byte) ([]byte, *net.UDPAddr, *onet.ConnectionError) {
+// extractPayloadAndDestination processes a decrypted Shadowsocks UDP packet and
+// extracts the payload data and destination address.
+func (h *associationHandler) extractPayloadAndDestination(textData []byte) ([]byte, net.Addr, *onet.ConnectionError) {
 	tgtAddr := socks.SplitAddr(textData)
 	if tgtAddr == nil {
 		return nil, nil, onet.NewConnectionError("ERR_READ_ADDRESS", "Failed to get target address", nil)
 	}
 
-	tgtUDPAddr, err := net.ResolveUDPAddr("udp", tgtAddr.String())
+	tgtUDPAddr, err := transport.MakeNetAddr("udp", tgtAddr.String())
 	if err != nil {
-		return nil, nil, onet.NewConnectionError("ERR_RESOLVE_ADDRESS", fmt.Sprintf("Failed to resolve target address %v", tgtAddr), err)
-	}
-	if err := h.targetIPValidator(tgtUDPAddr.IP); err != nil {
-		return nil, nil, ensureConnectionError(err, "ERR_ADDRESS_INVALID", "invalid address")
+		return nil, nil, onet.NewConnectionError("ERR_CONVERT_ADDRESS", fmt.Sprintf("Failed to convert target address %v", tgtAddr), err)
 	}
 
 	payload := textData[len(tgtAddr):]
@@ -408,6 +398,23 @@ func isDNS(addr net.Addr) bool {
 	return port == "53"
 }
 
+type validatingPacketConn struct {
+	net.PacketConn
+	targetIPValidator onet.TargetIPValidator
+}
+
+func (vpc *validatingPacketConn) WriteTo(p []byte, addr net.Addr) (int, error) {
+	udpAddr, err := net.ResolveUDPAddr("udp", addr.String())
+	if err != nil {
+		return 0, onet.NewConnectionError("ERR_RESOLVE_ADDRESS", fmt.Sprintf("Failed to resolve target address %v", udpAddr), err)
+	}
+	if err := vpc.targetIPValidator(udpAddr.IP); err != nil {
+		return 0, ensureConnectionError(err, "ERR_ADDRESS_INVALID", "invalid address")
+	}
+
+	return vpc.PacketConn.WriteTo(p, udpAddr) // accept only `net.UDPAddr` despite the signature
+}
+
 type timedPacketConn struct {
 	net.PacketConn
 	// Connection timeout to apply for non-DNS packets.

service/udp_linux.go

@@ -23,12 +23,18 @@ import (
 	"time"
 
 	"github.com/Jigsaw-Code/outline-sdk/transport"
+
+	onet "github.com/Jigsaw-Code/outline-ss-server/net"
 )
 
 type udpListener struct {
+	// The validator to be used to validate target IP addresses.
+	targetIPValidator onet.TargetIPValidator
+
 	// NAT mapping timeout is the default time a mapping will stay active
 	// without packets traversing the NAT, applied to non-DNS packets.
 	timeout time.Duration
+
 	// fwmark can be used in conjunction with other Linux networking features like cgroups, network
 	// namespaces, and TC (Traffic Control) for sophisticated network management.
 	// Value of 0 disables fwmark (SO_MARK) (Linux only)
@@ -37,14 +43,14 @@ type udpListener struct {
 
 // NewPacketListener creates a new PacketListener that listens on UDP
 // and optionally sets a firewall mark on the socket (Linux only).
-func MakeTargetUDPListener(timeout time.Duration, fwmark uint) transport.PacketListener {
-	return &udpListener{timeout: timeout, fwmark: fwmark}
+func MakeTargetUDPListener(targetIPValidator onet.TargetIPValidator, timeout time.Duration, fwmark uint) transport.PacketListener {
+	return &udpListener{timeout: timeout, targetIPValidator: targetIPValidator, fwmark: fwmark}
 }
 
 func (ln *udpListener) ListenPacket(ctx context.Context) (net.PacketConn, error) {
 	conn, err := net.ListenUDP("udp", nil)
 	if err != nil {
-		return nil, fmt.Errorf("Failed to create UDP socket: %w", err)
+		return nil, fmt.Errorf("failed to create UDP socket: %w", err)
 	}
 
 	if ln.fwmark > 0 {
@@ -57,9 +63,12 @@ func (ln *udpListener) ListenPacket(ctx context.Context) (net.PacketConn, error)
 		err = SetFwmark(rawConn, ln.fwmark)
 		if err != nil {
 			conn.Close()
-			return nil, fmt.Errorf("Failed to set `fwmark`: %w", err)
+			return nil, fmt.Errorf("failed to set `fwmark`: %w", err)
 
 		}
 	}
-	return &timedPacketConn{PacketConn: conn, defaultTimeout: ln.timeout}, nil
+	return &validatingPacketConn{
+		PacketConn:        &timedPacketConn{PacketConn: conn, defaultTimeout: ln.timeout},
+		targetIPValidator: ln.targetIPValidator,
+	}, nil
 }

service/udp_other.go

@@ -22,29 +22,41 @@ import (
 	"time"
 
 	"github.com/Jigsaw-Code/outline-sdk/transport"
+
+	onet "github.com/Jigsaw-Code/outline-ss-server/net"
 )
 
 type udpListener struct {
 	*transport.UDPListener
 
+	// The validator to be used to validate target IP addresses.
+	targetIPValidator onet.TargetIPValidator
+
 	// NAT mapping timeout is the default time a mapping will stay active
 	// without packets traversing the NAT, applied to non-DNS packets.
 	timeout time.Duration
 }
 
 // fwmark can be used in conjunction with other Linux networking features like cgroups, network namespaces, and TC (Traffic Control) for sophisticated network management.
 // Value of 0 disables fwmark (SO_MARK)
-func MakeTargetUDPListener(timeout time.Duration, fwmark uint) transport.PacketListener {
+func MakeTargetUDPListener(targetIPValidator onet.TargetIPValidator, timeout time.Duration, fwmark uint) transport.PacketListener {
 	if fwmark != 0 {
 		panic("fwmark is linux-specific feature and should be 0")
 	}
-	return &udpListener{UDPListener: &transport.UDPListener{Address: ""}}
+	return &udpListener{
+		targetIPValidator: targetIPValidator,
+		timeout:           timeout,
+		UDPListener:       &transport.UDPListener{Address: ""},
+	}
 }
 
 func (ln *udpListener) ListenPacket(ctx context.Context) (net.PacketConn, error) {
 	conn, err := ln.UDPListener.ListenPacket(ctx)
 	if err != nil {
 		return nil, err
 	}
-	return &timedPacketConn{PacketConn: conn, defaultTimeout: ln.timeout}, nil
+	return &validatingPacketConn{
+		PacketConn:        &timedPacketConn{PacketConn: conn, defaultTimeout: ln.timeout},
+		targetIPValidator: ln.targetIPValidator,
+	}, nil
 }

service/udp_test.go

@@ -24,6 +24,7 @@ import (
 	"testing"
 	"time"
 
+	"github.com/Jigsaw-Code/outline-sdk/transport"
 	"github.com/Jigsaw-Code/outline-sdk/transport/shadowsocks"
 	logging "github.com/op/go-logging"
 	"github.com/shadowsocks/go-shadowsocks2/socks"
@@ -62,6 +63,15 @@ func (ln *packetListener) ListenPacket(ctx context.Context) (net.PacketConn, err
 	return ln.conn, nil
 }
 
+func WrapWithValidatingPacketListener(conn net.PacketConn, targetIPValidator onet.TargetIPValidator) transport.PacketListener {
+	return &packetListener{
+		&validatingPacketConn{
+			PacketConn:        conn,
+			targetIPValidator: targetIPValidator,
+		},
+	}
+}
+
 type fakePacketConn struct {
 	net.PacketConn
 	send     chan fakePacket
@@ -225,7 +235,7 @@ func TestAssociationCloseWhileReading(t *testing.T) {
 func TestAssociationHandler_Handle_IPFilter(t *testing.T) {
 	t.Run("RequirePublicIP blocks localhost", func(t *testing.T) {
 		handler, sendPayload, targetConn := startTestHandler()
-		handler.SetTargetIPValidator(onet.RequirePublicIP)
+		handler.SetTargetPacketListener(WrapWithValidatingPacketListener(targetConn, onet.RequirePublicIP))
 
 		sendPayload(&localAddr, []byte{1, 2, 3})
 
@@ -239,7 +249,7 @@ func TestAssociationHandler_Handle_IPFilter(t *testing.T) {
 
 	t.Run("allowAll allows localhost", func(t *testing.T) {
 		handler, sendPayload, targetConn := startTestHandler()
-		handler.SetTargetIPValidator(allowAll)
+		handler.SetTargetPacketListener(WrapWithValidatingPacketListener(targetConn, allowAll))
 
 		sendPayload(&localAddr, []byte{1, 2, 3})