refactor: allow hot swapping the config

cmd/outline-ss-server/main.go

@@ -23,6 +23,8 @@ import (
 	"net/http"
 	"os"
 	"os/signal"
+	"strconv"
+	"sync"
 	"syscall"
 	"time"
 
@@ -55,129 +57,202 @@ func init() {
 	)
 }
 
-type ssPort struct {
-	tcpListener *net.TCPListener
-	packetConn  net.PacketConn
-	cipherList  service.CipherList
-}
-
 type SSServer struct {
+	stopConfig  func() error
+	lnManager   service.ListenerManager
 	natTimeout  time.Duration
 	m           *outlineMetrics
 	replayCache service.ReplayCache
-	ports       map[int]*ssPort
 }
 
-func (s *SSServer) startPort(portNum int) error {
-	listener, err := net.ListenTCP("tcp", &net.TCPAddr{Port: portNum})
+func (s *SSServer) loadConfig(filename string) error {
+	config, err := readConfig(filename)
 	if err != nil {
-		//lint:ignore ST1005 Shadowsocks is capitalized.
-		return fmt.Errorf("Shadowsocks TCP service failed to start on port %v: %w", portNum, err)
+		return fmt.Errorf("failed to load config (%v): %w", filename, err)
 	}
-	slog.Info("Shadowsocks TCP service started.", "address", listener.Addr().String())
-	packetConn, err := net.ListenUDP("udp", &net.UDPAddr{Port: portNum})
+	// We hot swap the config by having the old and new listeners both live at
+	// the same time. This means we create listeners for the new config first,
+	// and then close the old ones after.
+	stopConfig, err := s.runConfig(*config)
 	if err != nil {
-		//lint:ignore ST1005 Shadowsocks is capitalized.
-		return fmt.Errorf("Shadowsocks UDP service failed to start on port %v: %w", portNum, err)
+		return err
 	}
-	slog.Info("Shadowsocks UDP service started.", "address", packetConn.LocalAddr().String())
-	port := &ssPort{tcpListener: listener, packetConn: packetConn, cipherList: service.NewCipherList()}
-	authFunc := service.NewShadowsocksStreamAuthenticator(port.cipherList, &s.replayCache, s.m)
-	// TODO: Register initial data metrics at zero.
-	tcpHandler := service.NewTCPHandler(authFunc, s.m, tcpReadTimeout)
-	packetHandler := service.NewPacketHandler(s.natTimeout, port.cipherList, s.m)
-	s.ports[portNum] = port
-	go service.StreamServe(service.WrapStreamListener(listener.AcceptTCP), tcpHandler.Handle)
-	go packetHandler.Handle(port.packetConn)
+	if err := s.Stop(); err != nil {
+		slog.Warn("Failed to stop old config.", "err", err)
+	}
+	s.stopConfig = stopConfig
 	return nil
 }
 
-func (s *SSServer) removePort(portNum int) error {
-	port, ok := s.ports[portNum]
-	if !ok {
-		return fmt.Errorf("port %v doesn't exist", portNum)
-	}
-	tcpErr := port.tcpListener.Close()
-	udpErr := port.packetConn.Close()
-	delete(s.ports, portNum)
-	if tcpErr != nil {
-		//lint:ignore ST1005 Shadowsocks is capitalized.
-		return fmt.Errorf("Shadowsocks TCP service on port %v failed to stop: %w", portNum, tcpErr)
+func (s *SSServer) NewShadowsocksStreamHandler(ciphers service.CipherList) service.StreamHandler {
+	authFunc := service.NewShadowsocksStreamAuthenticator(ciphers, &s.replayCache, s.m)
+	// TODO: Register initial data metrics at zero.
+	return service.NewStreamHandler(authFunc, s.m, tcpReadTimeout)
+}
+
+func (s *SSServer) NewShadowsocksPacketHandler(ciphers service.CipherList) service.PacketHandler {
+	return service.NewPacketHandler(s.natTimeout, ciphers, s.m)
+}
+
+type listenerSet struct {
+	manager            service.ListenerManager
+	listenerCloseFuncs map[string]func() error
+	listenersMu        sync.Mutex
+}
+
+// ListenStream announces on a given network address. Trying to listen for stream connections
+// on the same address twice will result in an error.
+func (ls *listenerSet) ListenStream(addr string) (service.StreamListener, error) {
+	ls.listenersMu.Lock()
+	defer ls.listenersMu.Unlock()
+
+	lnKey := "stream/" + addr
+	if _, exists := ls.listenerCloseFuncs[lnKey]; exists {
+		return nil, fmt.Errorf("stream listener for %s already exists", addr)
 	}
-	slog.Info("Shadowsocks TCP service stopped.", "port", portNum)
-	if udpErr != nil {
-		//lint:ignore ST1005 Shadowsocks is capitalized.
-		return fmt.Errorf("Shadowsocks UDP service on port %v failed to stop: %w", portNum, udpErr)
+	ln, err := ls.manager.ListenStream(addr)
+	if err != nil {
+		return nil, err
 	}
-	slog.Info("Shadowsocks UDP service stopped.", "port", portNum)
-	return nil
+	ls.listenerCloseFuncs[lnKey] = ln.Close
+	return ln, nil
 }
 
-func (s *SSServer) loadConfig(filename string) error {
-	config, err := readConfig(filename)
+// ListenPacket announces on a given network address. Trying to listen for packet connections
+// on the same address twice will result in an error.
+func (ls *listenerSet) ListenPacket(addr string) (net.PacketConn, error) {
+	ls.listenersMu.Lock()
+	defer ls.listenersMu.Unlock()
+
+	lnKey := "packet/" + addr
+	if _, exists := ls.listenerCloseFuncs[lnKey]; exists {
+		return nil, fmt.Errorf("packet listener for %s already exists", addr)
+	}
+	ln, err := ls.manager.ListenPacket(addr)
 	if err != nil {
-		return fmt.Errorf("failed to load config (%v): %w", filename, err)
+		return nil, err
 	}
+	ls.listenerCloseFuncs[lnKey] = ln.Close
+	return ln, nil
+}
 
-	portChanges := make(map[int]int)
-	portCiphers := make(map[int]*list.List) // Values are *List of *CipherEntry.
-	for _, keyConfig := range config.Keys {
-		portChanges[keyConfig.Port] = 1
-		cipherList, ok := portCiphers[keyConfig.Port]
-		if !ok {
-			cipherList = list.New()
-			portCiphers[keyConfig.Port] = cipherList
-		}
-		cryptoKey, err := shadowsocks.NewEncryptionKey(keyConfig.Cipher, keyConfig.Secret)
-		if err != nil {
-			return fmt.Errorf("failed to create encyption key for key %v: %w", keyConfig.ID, err)
+// Close closes all the listeners in the set, after which the set can't be used again.
+func (ls *listenerSet) Close() error {
+	ls.listenersMu.Lock()
+	defer ls.listenersMu.Unlock()
+
+	for addr, listenerCloseFunc := range ls.listenerCloseFuncs {
+		if err := listenerCloseFunc(); err != nil {
+			return fmt.Errorf("listener on address %s failed to stop: %w", addr, err)
 		}
-		entry := service.MakeCipherEntry(keyConfig.ID, cryptoKey, keyConfig.Secret)
-		cipherList.PushBack(&entry)
-	}
-	for port := range s.ports {
-		portChanges[port] = portChanges[port] - 1
 	}
-	for portNum, count := range portChanges {
-		if count == -1 {
-			if err := s.removePort(portNum); err != nil {
-				return fmt.Errorf("failed to remove port %v: %w", portNum, err)
+	ls.listenerCloseFuncs = nil
+	return nil
+}
+
+// Len returns the number of listeners in the set.
+func (ls *listenerSet) Len() int {
+	return len(ls.listenerCloseFuncs)
+}
+
+func (s *SSServer) runConfig(config Config) (func() error, error) {
+	startErrCh := make(chan error)
+	stopErrCh := make(chan error)
+	stopCh := make(chan struct{})
+
+	go func() {
+		lnSet := &listenerSet{
+			manager:            s.lnManager,
+			listenerCloseFuncs: make(map[string]func() error),
+		}
+		defer func() {
+			stopErrCh <- lnSet.Close()
+		}()
+
+		startErrCh <- func() error {
+			portCiphers := make(map[int]*list.List) // Values are *List of *CipherEntry.
+			for _, keyConfig := range config.Keys {
+				cipherList, ok := portCiphers[keyConfig.Port]
+				if !ok {
+					cipherList = list.New()
+					portCiphers[keyConfig.Port] = cipherList
+				}
+				cryptoKey, err := shadowsocks.NewEncryptionKey(keyConfig.Cipher, keyConfig.Secret)
+				if err != nil {
+					return fmt.Errorf("failed to create encyption key for key %v: %w", keyConfig.ID, err)
+				}
+				entry := service.MakeCipherEntry(keyConfig.ID, cryptoKey, keyConfig.Secret)
+				cipherList.PushBack(&entry)
 			}
-		} else if count == +1 {
-			if err := s.startPort(portNum); err != nil {
-				return err
+			for portNum, cipherList := range portCiphers {
+				addr := net.JoinHostPort("::", strconv.Itoa(portNum))
+
+				ciphers := service.NewCipherList()
+				ciphers.Update(cipherList)
+
+				sh := s.NewShadowsocksStreamHandler(ciphers)
+				ln, err := lnSet.ListenStream(addr)
+				if err != nil {
+					return err
+				}
+				slog.Info("Shadowsocks TCP service started.", "address", ln.Addr().String())
+				go service.StreamServe(ln.AcceptStream, sh.Handle)
+
+				pc, err := lnSet.ListenPacket(addr)
+				if err != nil {
+					return err
+				}
+				slog.Info("Shadowsocks UDP service started.", "address", pc.LocalAddr().String())
+				ph := s.NewShadowsocksPacketHandler(ciphers)
+				go ph.Handle(pc)
 			}
-		}
-	}
-	for portNum, cipherList := range portCiphers {
-		s.ports[portNum].cipherList.Update(cipherList)
+			slog.Info("Loaded config.", "access keys", len(config.Keys), "listeners", lnSet.Len())
+			s.m.SetNumAccessKeys(len(config.Keys), lnSet.Len())
+			return nil
+		}()
+
+		<-stopCh
+	}()
+
+	err := <-startErrCh
+	if err != nil {
+		return nil, err
 	}
-	slog.Info("Loaded config.", "access keys", len(config.Keys), "ports", len(s.ports))
-	s.m.SetNumAccessKeys(len(config.Keys), len(portCiphers))
-	return nil
+	return func() error {
+		slog.Info("Stopping running config.")
+		// TODO(sbruens): Actually wait for all handlers to be stopped, e.g. by
+		// using a https://pkg.go.dev/sync#WaitGroup.
+		stopCh <- struct{}{}
+		stopErr := <-stopErrCh
+		return stopErr
+	}, nil
 }
 
-// Stop serving on all ports.
+// Stop stops serving the current config.
 func (s *SSServer) Stop() error {
-	for portNum := range s.ports {
-		if err := s.removePort(portNum); err != nil {
-			return err
-		}
+	stopFunc := s.stopConfig
+	if stopFunc == nil {
+		return nil
+	}
+	if err := stopFunc(); err != nil {
+		slog.Error("Error stopping config.", "err", err)
+		return err
 	}
+	slog.Info("Stopped all listeners for running config.")
 	return nil
 }
 
 // RunSSServer starts a shadowsocks server running, and returns the server or an error.
 func RunSSServer(filename string, natTimeout time.Duration, sm *outlineMetrics, replayHistory int) (*SSServer, error) {
 	server := &SSServer{
+		lnManager:   service.NewListenerManager(),
 		natTimeout:  natTimeout,
 		m:           sm,
 		replayCache: service.NewReplayCache(replayHistory),
-		ports:       make(map[int]*ssPort),
 	}
 	err := server.loadConfig(filename)
 	if err != nil {
-		return nil, fmt.Errorf("failed configure server: %w", err)
+		return nil, fmt.Errorf("failed to configure server: %w", err)
 	}
 	sigHup := make(chan os.Signal, 1)
 	signal.Notify(sigHup, syscall.SIGHUP)

internal/integration_test/integration_test.go

@@ -133,7 +133,7 @@ func TestTCPEcho(t *testing.T) {
 	const testTimeout = 200 * time.Millisecond
 	testMetrics := &service.NoOpTCPMetrics{}
 	authFunc := service.NewShadowsocksStreamAuthenticator(cipherList, &replayCache, testMetrics)
-	handler := service.NewTCPHandler(authFunc, testMetrics, testTimeout)
+	handler := service.NewStreamHandler(authFunc, testMetrics, testTimeout)
 	handler.SetTargetDialer(&transport.TCPDialer{})
 	done := make(chan struct{})
 	go func() {
@@ -202,10 +202,10 @@ func TestRestrictedAddresses(t *testing.T) {
 	const testTimeout = 200 * time.Millisecond
 	testMetrics := &statusMetrics{}
 	authFunc := service.NewShadowsocksStreamAuthenticator(cipherList, nil, testMetrics)
-	handler := service.NewTCPHandler(authFunc, testMetrics, testTimeout)
+	handler := service.NewStreamHandler(authFunc, testMetrics, testTimeout)
 	done := make(chan struct{})
 	go func() {
-		service.StreamServe(service.WrapStreamListener(proxyListener.AcceptTCP), handler.Handle)
+		service.StreamServe(service.WrapStreamAcceptFunc(proxyListener.AcceptTCP), handler.Handle)
 		done <- struct{}{}
 	}()
 
@@ -384,11 +384,11 @@ func BenchmarkTCPThroughput(b *testing.B) {
 	const testTimeout = 200 * time.Millisecond
 	testMetrics := &service.NoOpTCPMetrics{}
 	authFunc := service.NewShadowsocksStreamAuthenticator(cipherList, nil, testMetrics)
-	handler := service.NewTCPHandler(authFunc, testMetrics, testTimeout)
+	handler := service.NewStreamHandler(authFunc, testMetrics, testTimeout)
 	handler.SetTargetDialer(&transport.TCPDialer{})
 	done := make(chan struct{})
 	go func() {
-		service.StreamServe(service.WrapStreamListener(proxyListener.AcceptTCP), handler.Handle)
+		service.StreamServe(service.WrapStreamAcceptFunc(proxyListener.AcceptTCP), handler.Handle)
 		done <- struct{}{}
 	}()
 
@@ -448,11 +448,11 @@ func BenchmarkTCPMultiplexing(b *testing.B) {
 	const testTimeout = 200 * time.Millisecond
 	testMetrics := &service.NoOpTCPMetrics{}
 	authFunc := service.NewShadowsocksStreamAuthenticator(cipherList, &replayCache, testMetrics)
-	handler := service.NewTCPHandler(authFunc, testMetrics, testTimeout)
+	handler := service.NewStreamHandler(authFunc, testMetrics, testTimeout)
 	handler.SetTargetDialer(&transport.TCPDialer{})
 	done := make(chan struct{})
 	go func() {
-		service.StreamServe(service.WrapStreamListener(proxyListener.AcceptTCP), handler.Handle)
+		service.StreamServe(service.WrapStreamAcceptFunc(proxyListener.AcceptTCP), handler.Handle)
 		done <- struct{}{}
 	}()