refactor: split out UDP association handling from serving

caddy/shadowsocks_handler.go

@@ -19,16 +19,21 @@ import (
 	"fmt"
 	"log/slog"
 	"net"
+	"time"
 
 	"github.com/Jigsaw-Code/outline-sdk/transport"
 	"github.com/Jigsaw-Code/outline-sdk/transport/shadowsocks"
-	outline "github.com/Jigsaw-Code/outline-ss-server/service"
 	"github.com/caddyserver/caddy/v2"
 	"github.com/mholt/caddy-l4/layer4"
+
+	outline "github.com/Jigsaw-Code/outline-ss-server/service"
 )
 
 const ssModuleName = "layer4.handlers.shadowsocks"
 
+// A UDP NAT timeout of at least 5 minutes is recommended in RFC 4787 Section 4.3.
+const defaultNatTimeout time.Duration = 5 * time.Minute
+
 func init() {
 	caddy.RegisterModule(ModuleRegistration{
 		ID:  ssModuleName,
@@ -45,8 +50,10 @@ type KeyConfig struct {
 type ShadowsocksHandler struct {
 	Keys []KeyConfig `json:"keys,omitempty"`
 
-	service outline.Service
-	logger  *slog.Logger
+	streamHandler      outline.StreamHandler
+	associationHandler outline.AssociationHandler
+	metrics            outline.ServiceMetrics
+	logger             *slog.Logger
 }
 
 var (
@@ -70,6 +77,7 @@ func (h *ShadowsocksHandler) Provision(ctx caddy.Context) error {
 	if !ok {
 		return fmt.Errorf("module `%s` is of type `%T`, expected `OutlineApp`", outlineModuleName, app)
 	}
+	h.metrics = app.Metrics
 
 	if len(h.Keys) == 0 {
 		h.logger.Warn("no keys configured")
@@ -97,26 +105,22 @@ func (h *ShadowsocksHandler) Provision(ctx caddy.Context) error {
 	ciphers := outline.NewCipherList()
 	ciphers.Update(cipherList)
 
-	service, err := outline.NewShadowsocksService(
+	h.streamHandler, h.associationHandler = outline.NewShadowsocksHandlers(
 		outline.WithLogger(h.logger),
 		outline.WithCiphers(ciphers),
-		outline.WithMetrics(app.Metrics),
+		outline.WithMetrics(h.metrics),
 		outline.WithReplayCache(&app.ReplayCache),
 	)
-	if err != nil {
-		return err
-	}
-	h.service = service
 	return nil
 }
 
 // Handle implements layer4.NextHandler.
 func (h *ShadowsocksHandler) Handle(cx *layer4.Connection, _ layer4.Handler) error {
 	switch conn := cx.Conn.(type) {
 	case transport.StreamConn:
-		h.service.HandleStream(cx.Context, conn)
-	case net.PacketConn:
-		h.service.HandlePacket(conn)
+		h.streamHandler.HandleStream(cx.Context, conn, h.metrics.AddOpenTCPConnection(conn))
+	case net.Conn:
+		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

@@ -16,6 +16,7 @@ package main
 
 import (
 	"container/list"
+	"context"
 	"flag"
 	"fmt"
 	"log/slog"
@@ -27,6 +28,7 @@ import (
 	"syscall"
 	"time"
 
+	"github.com/Jigsaw-Code/outline-sdk/transport"
 	"github.com/Jigsaw-Code/outline-sdk/transport/shadowsocks"
 	"github.com/lmittmann/tint"
 	"github.com/prometheus/client_golang/prometheus"
@@ -223,40 +225,43 @@ func (s *OutlineServer) runConfig(config Config) (func() error, error) {
 				ciphers := service.NewCipherList()
 				ciphers.Update(cipherList)
 
-				ssService, err := service.NewShadowsocksService(
+				streamHandler, associationHandler := service.NewShadowsocksHandlers(
 					service.WithCiphers(ciphers),
-					service.WithNatTimeout(s.natTimeout),
 					service.WithMetrics(s.serviceMetrics),
 					service.WithReplayCache(&s.replayCache),
+					service.WithPacketListener(service.MakeTargetUDPListener(s.natTimeout, 0)),
 					service.WithLogger(slog.Default()),
 				)
 				ln, err := lnSet.ListenStream(addr)
 				if err != nil {
 					return err
 				}
 				slog.Info("TCP service started.", "address", ln.Addr().String())
-				go service.StreamServe(ln.AcceptStream, ssService.HandleStream)
+				go service.StreamServe(ln.AcceptStream, func(ctx context.Context, conn transport.StreamConn) {
+					streamHandler.HandleStream(ctx, conn, s.serviceMetrics.AddOpenTCPConnection(conn))
+				})
 
 				pc, err := lnSet.ListenPacket(addr)
 				if err != nil {
 					return err
 				}
 				slog.Info("UDP service started.", "address", pc.LocalAddr().String())
-				go ssService.HandlePacket(pc)
+				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)
 				}
-				ssService, err := service.NewShadowsocksService(
+				streamHandler, associationHandler := service.NewShadowsocksHandlers(
 					service.WithCiphers(ciphers),
-					service.WithNatTimeout(s.natTimeout),
 					service.WithMetrics(s.serviceMetrics),
 					service.WithReplayCache(&s.replayCache),
 					service.WithStreamDialer(service.MakeValidatingTCPStreamDialer(onet.RequirePublicIP, serviceConfig.Dialer.Fwmark)),
-					service.WithPacketListener(service.MakeTargetUDPListener(serviceConfig.Dialer.Fwmark)),
+					service.WithPacketListener(service.MakeTargetUDPListener(s.natTimeout, serviceConfig.Dialer.Fwmark)),
 					service.WithLogger(slog.Default()),
 				)
 				if err != nil {
@@ -275,7 +280,9 @@ func (s *OutlineServer) runConfig(config Config) (func() error, error) {
 							}
 							return serviceConfig.Dialer.Fwmark
 						}())
-						go service.StreamServe(ln.AcceptStream, ssService.HandleStream)
+						go service.StreamServe(ln.AcceptStream, func(ctx context.Context, conn transport.StreamConn) {
+							streamHandler.HandleStream(ctx, conn, s.serviceMetrics.AddOpenTCPConnection(conn))
+						})
 					case listenerTypeUDP:
 						pc, err := lnSet.ListenPacket(lnConfig.Address)
 						if err != nil {
@@ -287,7 +294,9 @@ func (s *OutlineServer) runConfig(config Config) (func() error, error) {
 							}
 							return serviceConfig.Dialer.Fwmark
 						}())
-						go ssService.HandlePacket(pc)
+						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)

cmd/outline-ss-server/metrics.go

@@ -17,6 +17,7 @@ package main
 import (
 	"time"
 
+	"github.com/Jigsaw-Code/outline-ss-server/service"
 	"github.com/prometheus/client_golang/prometheus"
 )
 
@@ -25,12 +26,15 @@ var now = time.Now
 
 type serverMetrics struct {
 	// NOTE: New metrics need to be added to `newPrometheusServerMetrics()`, `Describe()` and `Collect()`.
-	buildInfo  *prometheus.GaugeVec
-	accessKeys prometheus.Gauge
-	ports      prometheus.Gauge
+	buildInfo         *prometheus.GaugeVec
+	accessKeys        prometheus.Gauge
+	ports             prometheus.Gauge
+	addedNatEntries   prometheus.Counter
+	removedNatEntries prometheus.Counter
 }
 
 var _ prometheus.Collector = (*serverMetrics)(nil)
+var _ service.NATMetrics = (*serverMetrics)(nil)
 
 // newPrometheusServerMetrics constructs a Prometheus metrics collector for server
 // related metrics.
@@ -48,19 +52,33 @@ func newPrometheusServerMetrics() *serverMetrics {
 			Name: "ports",
 			Help: "Count of open ports",
 		}),
+		addedNatEntries: prometheus.NewCounter(prometheus.CounterOpts{
+			Subsystem: "udp",
+			Name:      "nat_entries_added",
+			Help:      "Entries added to the UDP NAT table",
+		}),
+		removedNatEntries: prometheus.NewCounter(prometheus.CounterOpts{
+			Subsystem: "udp",
+			Name:      "nat_entries_removed",
+			Help:      "Entries removed from the UDP NAT table",
+		}),
 	}
 }
 
 func (m *serverMetrics) Describe(ch chan<- *prometheus.Desc) {
 	m.buildInfo.Describe(ch)
 	m.accessKeys.Describe(ch)
 	m.ports.Describe(ch)
+	m.addedNatEntries.Describe(ch)
+	m.removedNatEntries.Describe(ch)
 }
 
 func (m *serverMetrics) Collect(ch chan<- prometheus.Metric) {
 	m.buildInfo.Collect(ch)
 	m.accessKeys.Collect(ch)
 	m.ports.Collect(ch)
+	m.addedNatEntries.Collect(ch)
+	m.removedNatEntries.Collect(ch)
 }
 
 func (m *serverMetrics) SetVersion(version string) {
@@ -71,3 +89,11 @@ func (m *serverMetrics) SetNumAccessKeys(numKeys int, ports int) {
 	m.accessKeys.Set(float64(numKeys))
 	m.ports.Set(float64(ports))
 }
+
+func (m *serverMetrics) AddNATEntry() {
+	m.addedNatEntries.Inc()
+}
+
+func (m *serverMetrics) RemoveNATEntry() {
+	m.removedNatEntries.Inc()
+}