refactor: split out association handling in packet handler

caddy/shadowsocks_handler.go

@@ -16,19 +16,26 @@ package caddy
 
 import (
 	"container/list"
+	"context"
 	"fmt"
 	"log/slog"
 	"net"
 
 	"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"
 
+type OutlineService interface {
+	HandleStream(ctx context.Context, conn transport.StreamConn)
+	HandleAssociation(conn net.Conn) error
+}
+
 func init() {
 	caddy.RegisterModule(ModuleRegistration{
 		ID:  ssModuleName,
@@ -45,7 +52,7 @@ type KeyConfig struct {
 type ShadowsocksHandler struct {
 	Keys []KeyConfig `json:"keys,omitempty"`
 
-	service outline.Service
+	service OutlineService
 	logger  *slog.Logger
 }
 
@@ -115,8 +122,10 @@ func (h *ShadowsocksHandler) Handle(cx *layer4.Connection, _ layer4.Handler) err
 	switch conn := cx.Conn.(type) {
 	case transport.StreamConn:
 		h.service.HandleStream(cx.Context, conn)
-	case net.PacketConn:
-		h.service.HandlePacket(conn)
+	case net.Conn:
+		if err := h.service.HandleAssociation(conn); err != nil {
+			return err
+		}
 	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"
@@ -60,6 +62,11 @@ func init() {
 	)
 }
 
+type OutlineService interface {
+	HandleStream(ctx context.Context, conn transport.StreamConn)
+	NewPacketAssociation(conn net.Conn) (service.PacketAssociation, error)
+}
+
 type OutlineServer struct {
 	stopConfig     func() error
 	lnManager      service.ListenerManager
@@ -223,11 +230,12 @@ func (s *OutlineServer) runConfig(config Config) (func() error, error) {
 				ciphers := service.NewCipherList()
 				ciphers.Update(cipherList)
 
+				var ssService OutlineService
 				ssService, err := service.NewShadowsocksService(
 					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)
@@ -242,21 +250,21 @@ func (s *OutlineServer) runConfig(config Config) (func() error, error) {
 					return err
 				}
 				slog.Info("UDP service started.", "address", pc.LocalAddr().String())
-				go ssService.HandlePacket(pc)
+				go service.PacketServe(pc, ssService.NewPacketAssociation, 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(
+				var ssService OutlineService
+				ssService, err = service.NewShadowsocksService(
 					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 {
@@ -287,7 +295,7 @@ func (s *OutlineServer) runConfig(config Config) (func() error, error) {
 							}
 							return serviceConfig.Dialer.Fwmark
 						}())
-						go ssService.HandlePacket(pc)
+						go service.PacketServe(pc, ssService.NewPacketAssociation, 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()
+}