Refactor to make packet handler an association handler.

caddy/shadowsocks_handler.go

@@ -51,9 +51,8 @@ type ShadowsocksHandler struct {
 	Keys []KeyConfig `json:"keys,omitempty"`
 
 	streamHandler outline.StreamHandler
-	packetHandler outline.PacketHandler
+	associationHandler outline.AssociationHandler
 	metrics       outline.ServiceMetrics
-	tgtListener   transport.PacketListener
 	logger        *slog.Logger
 }
 
@@ -106,13 +105,12 @@ func (h *ShadowsocksHandler) Provision(ctx caddy.Context) error {
 	ciphers := outline.NewCipherList()
 	ciphers.Update(cipherList)
 
-	h.streamHandler, h.packetHandler = outline.NewShadowsocksHandlers(
+	h.streamHandler, h.associationHandler = outline.NewShadowsocksHandlers(
 		outline.WithLogger(h.logger),
 		outline.WithCiphers(ciphers),
 		outline.WithMetrics(h.metrics),
 		outline.WithReplayCache(&app.ReplayCache),
 	)
-	h.tgtListener = outline.MakeTargetUDPListener(defaultNatTimeout, 0)
 	return nil
 }
 
@@ -122,11 +120,7 @@ func (h *ShadowsocksHandler) Handle(cx *layer4.Connection, _ layer4.Handler) err
 	case transport.StreamConn:
 		h.streamHandler.HandleStream(cx.Context, conn, h.metrics.AddOpenTCPConnection(conn))
 	case net.Conn:
-		assoc, err := outline.NewPacketAssociation(conn, h.tgtListener, h.metrics.AddOpenUDPAssociation(conn))
-		if err != nil {
-			return fmt.Errorf("failed to handle association: %v", err)
-		}
-		outline.HandleAssociation(assoc, h.packetHandler.HandlePacket)
+		h.associationHandler.HandleAssociation(cx.Context, conn, h.metrics.AddOpenUDPAssociation(conn))
 	default:
 		return fmt.Errorf("failed to handle unknown connection type: %t", conn)
 	}

cmd/outline-ss-server/main.go

@@ -225,10 +225,11 @@ func (s *OutlineServer) runConfig(config Config) (func() error, error) {
 				ciphers := service.NewCipherList()
 				ciphers.Update(cipherList)
 
-				streamHandler, packetHandler := service.NewShadowsocksHandlers(
+				streamHandler, associationHandler := service.NewShadowsocksHandlers(
 					service.WithCiphers(ciphers),
 					service.WithMetrics(s.serviceMetrics),
 					service.WithReplayCache(&s.replayCache),
+					service.WithPacketListener(service.MakeTargetUDPListener(s.natTimeout, 0)),
 					service.WithLogger(slog.Default()),
 				)
 				ln, err := lnSet.ListenStream(addr)
@@ -245,27 +246,22 @@ func (s *OutlineServer) runConfig(config Config) (func() error, error) {
 					return err
 				}
 				slog.Info("UDP service started.", "address", pc.LocalAddr().String())
-				tgtListener := service.MakeTargetUDPListener(s.natTimeout, 0)
-				go service.PacketServe(pc, func(conn net.Conn) (service.PacketAssociation, error) {
-					m := s.serviceMetrics.AddOpenUDPAssociation(conn)
-					assoc, err := service.NewPacketAssociation(conn, tgtListener, m)
-					if err != nil {
-						return nil, fmt.Errorf("failed to handle association: %v", err)
-					}
-					return assoc, nil
-				}, packetHandler.HandlePacket, s.serverMetrics)
+				go service.PacketServe(pc, func(ctx context.Context, conn net.Conn) {
+					associationHandler.HandleAssociation(ctx, conn, s.serviceMetrics.AddOpenUDPAssociation(conn))
+				}, s.serverMetrics)
 			}
 
 			for _, serviceConfig := range config.Services {
 				ciphers, err := newCipherListFromConfig(serviceConfig)
 				if err != nil {
 					return fmt.Errorf("failed to create cipher list from config: %v", err)
 				}
-				streamHandler, packetHandler := service.NewShadowsocksHandlers(
+				streamHandler, associationHandler := service.NewShadowsocksHandlers(
 					service.WithCiphers(ciphers),
 					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.WithLogger(slog.Default()),
 				)
 				if err != nil {
@@ -298,15 +294,9 @@ func (s *OutlineServer) runConfig(config Config) (func() error, error) {
 							}
 							return serviceConfig.Dialer.Fwmark
 						}())
-						tgtListener := service.MakeTargetUDPListener(s.natTimeout, serviceConfig.Dialer.Fwmark)
-						go service.PacketServe(pc, func(conn net.Conn) (service.PacketAssociation, error) {
-							m := s.serviceMetrics.AddOpenUDPAssociation(conn)
-							assoc, err := service.NewPacketAssociation(conn, tgtListener, m)
-							if err != nil {
-								return nil, fmt.Errorf("failed to handle association: %v", err)
-							}
-							return assoc, nil
-						}, packetHandler.HandlePacket, s.serverMetrics)
+						go service.PacketServe(pc, func(ctx context.Context, conn net.Conn) {
+							associationHandler.HandleAssociation(ctx, conn, s.serviceMetrics.AddOpenUDPAssociation(conn))
+						}, s.serverMetrics)
 					}
 				}
 				totalCipherCount += len(serviceConfig.Keys)

internal/integration_test/integration_test.go

@@ -317,15 +317,14 @@ func TestUDPEcho(t *testing.T) {
 	if err != nil {
 		t.Fatal(err)
 	}
-	proxy := service.NewPacketHandler(cipherList, &fakeShadowsocksMetrics{})
+	proxy := service.NewAssociationHandler(cipherList, &fakeShadowsocksMetrics{})
 
 	proxy.SetTargetIPValidator(allowAll)
 	natMetrics := &natTestMetrics{}
 	associationMetrics := &fakeUDPAssociationMetrics{}
-	go service.PacketServe(proxyConn, func(conn net.Conn) (service.PacketAssociation, error) {
-		assoc, _ := service.NewPacketAssociation(conn, &transport.UDPListener{Address: ""}, associationMetrics)
-		return assoc, nil
-	}, proxy.Handle, natMetrics)
+	go service.PacketServe(proxyConn, func(ctx context.Context, conn net.Conn) {
+		proxy.HandleAssociation(ctx, conn, associationMetrics)
+	}, natMetrics)
 
 	cryptoKey, err := shadowsocks.NewEncryptionKey(shadowsocks.CHACHA20IETFPOLY1305, secrets[0])
 	require.NoError(t, err)
@@ -546,14 +545,13 @@ func BenchmarkUDPEcho(b *testing.B) {
 	if err != nil {
 		b.Fatal(err)
 	}
-	proxy := service.NewPacketHandler(cipherList, &fakeShadowsocksMetrics{})
+	proxy := service.NewAssociationHandler(cipherList, &fakeShadowsocksMetrics{})
 	proxy.SetTargetIPValidator(allowAll)
 	done := make(chan struct{})
 	go func() {
-		service.PacketServe(server, func(conn net.Conn) (service.PacketAssociation, error) {
-			assoc, _ := service.NewPacketAssociation(conn, &transport.UDPListener{Address: ""}, nil)
-			return assoc, nil
-		}, proxy.Handle, &natTestMetrics{})
+		service.PacketServe(server, func(ctx context.Context, conn net.Conn) {
+			proxy.HandleAssociation(ctx, conn, &fakeUDPAssociationMetrics{})
+		}, &natTestMetrics{})
 		done <- struct{}{}
 	}()
 
@@ -593,14 +591,13 @@ func BenchmarkUDPManyKeys(b *testing.B) {
 	if err != nil {
 		b.Fatal(err)
 	}
-	proxy := service.NewPacketHandler(cipherList, &fakeShadowsocksMetrics{})
+	proxy := service.NewAssociationHandler(cipherList, &fakeShadowsocksMetrics{})
 	proxy.SetTargetIPValidator(allowAll)
 	done := make(chan struct{})
 	go func() {
-		service.PacketServe(proxyConn, func(conn net.Conn) (service.PacketAssociation, error) {
-			assoc, _ := service.NewPacketAssociation(conn, &transport.UDPListener{Address: ""}, nil)
-			return assoc, nil
-		}, proxy.Handle, &natTestMetrics{})
+		service.PacketServe(proxyConn, func(ctx context.Context, conn net.Conn) {
+			proxy.HandleAssociation(ctx, conn, &fakeUDPAssociationMetrics{})
+		}, &natTestMetrics{})
 		done <- struct{}{}
 	}()
 

service/shadowsocks.go

@@ -24,10 +24,8 @@ import (
 	onet "github.com/Jigsaw-Code/outline-ss-server/net"
 )
 
-const (
-	// 59 seconds is most common timeout for servers that do not respond to invalid requests
-	tcpReadTimeout time.Duration = 59 * time.Second
-)
+// 59 seconds is most common timeout for servers that do not respond to invalid requests
+const tcpReadTimeout time.Duration = 59 * time.Second
 
 // ShadowsocksConnMetrics is used to report Shadowsocks related metrics on connections.
 type ShadowsocksConnMetrics interface {
@@ -51,11 +49,12 @@ type ssService struct {
 	targetIPValidator onet.TargetIPValidator
 	replayCache       *ReplayCache
 
-	streamDialer transport.StreamDialer
+	streamDialer   transport.StreamDialer
+	packetListener transport.PacketListener
 }
 
 // NewShadowsocksHandlers creates new Shadowsocks stream and packet handlers.
-func NewShadowsocksHandlers(opts ...Option) (StreamHandler, PacketHandler) {
+func NewShadowsocksHandlers(opts ...Option) (StreamHandler, AssociationHandler) {
 	s := &ssService{
 		logger: noopLogger(),
 	}
@@ -74,10 +73,13 @@ func NewShadowsocksHandlers(opts ...Option) (StreamHandler, PacketHandler) {
 	}
 	sh.SetLogger(s.logger)
 
-	ph := NewPacketHandler(s.ciphers, &ssConnMetrics{s.metrics.AddUDPCipherSearch})
-	ph.SetLogger(s.logger)
+	ah := NewAssociationHandler(s.ciphers, &ssConnMetrics{s.metrics.AddUDPCipherSearch})
+	if s.packetListener != nil {
+		ah.SetTargetPacketListener(s.packetListener)
+	}
+	ah.SetLogger(s.logger)
 
-	return sh, ph
+	return sh, ah
 }
 
 // WithLogger can be used to provide a custom log target. If not provided,
@@ -115,6 +117,13 @@ func WithStreamDialer(dialer transport.StreamDialer) Option {
 	}
 }
 
+// WithPacketListener option function.
+func WithPacketListener(listener transport.PacketListener) Option {
+	return func(s *ssService) {
+		s.packetListener = listener
+	}
+}
+
 type ssConnMetrics struct {
 	metricFunc func(accessKeyFound bool, timeToCipher time.Duration)
 }

service/tcp_test.go

@@ -368,7 +368,7 @@ func TestProbeClientBytesBasicTruncated(t *testing.T) {
 	testMetrics := &probeTestMetrics{}
 	authFunc := NewShadowsocksStreamAuthenticator(cipherList, nil, &fakeShadowsocksMetrics{}, nil)
 	handler := NewStreamHandler(authFunc, 200*time.Millisecond)
-	handler.SetTargetDialerStream(MakeValidatingTCPStreamDialer(allowAll, 0))
+	handler.SetTargetDialer(MakeValidatingTCPStreamDialer(allowAll, 0))
 	done := make(chan struct{})
 	go func() {
 		StreamServe(
@@ -406,7 +406,7 @@ func TestProbeClientBytesBasicModified(t *testing.T) {
 	testMetrics := &probeTestMetrics{}
 	authFunc := NewShadowsocksStreamAuthenticator(cipherList, nil, &fakeShadowsocksMetrics{}, nil)
 	handler := NewStreamHandler(authFunc, 200*time.Millisecond)
-	handler.SetTargetDialerStream(MakeValidatingTCPStreamDialer(allowAll, 0))
+	handler.SetTargetDialer(MakeValidatingTCPStreamDialer(allowAll, 0))
 	done := make(chan struct{})
 	go func() {
 		StreamServe(
@@ -445,7 +445,7 @@ func TestProbeClientBytesCoalescedModified(t *testing.T) {
 	testMetrics := &probeTestMetrics{}
 	authFunc := NewShadowsocksStreamAuthenticator(cipherList, nil, &fakeShadowsocksMetrics{}, nil)
 	handler := NewStreamHandler(authFunc, 200*time.Millisecond)
-	handler.SetTargetDialerStream(MakeValidatingTCPStreamDialer(allowAll, 0))
+	handler.SetTargetDialer(MakeValidatingTCPStreamDialer(allowAll, 0))
 	done := make(chan struct{})
 	go func() {
 		StreamServe(
@@ -747,7 +747,7 @@ func TestStreamServeEarlyClose(t *testing.T) {
 	err = tcpListener.Close()
 	require.NoError(t, err)
 	// This should return quickly, without timing out or calling the handler.
-	StreamServeStream(WrapStreamAcceptFunc(tcpListener.AcceptTCP), nil)
+	StreamServe(WrapStreamAcceptFunc(tcpListener.AcceptTCP), nil)
 }
 
 // Makes sure the TCP listener returns [io.ErrClosed] on Close().