concurrent Serve calls

Author: linkdata

README.md

@@ -19,6 +19,8 @@ supports multiple concurrent `Accept()` calls, allowing you to reverse-proxy a s
 
 ## Server
 
+The server can listen on multiple listeners concurrently.
+
 The server provides two abstractions to customize it's behavior.
 
 The `Authenticator` interface allows custom authentication methods, and comes with implementations for

server/listener.go

@@ -16,7 +16,7 @@ type listener struct {
 
 func (l *listener) Close() (err error) {
 	if refs := l.refs.Add(-1); refs < 1 {
-		died := int64(time.Since(l.srv.Started))
+		died := int64(time.Since(l.srv.started))
 		l.died.Store(died)
 		_ = l.srv.Debug && l.srv.LogDebug("listener deref", "key", l.key, "refs", refs, "died", died)
 	}

server/server.go

@@ -5,16 +5,13 @@ import (
 	"io"
 	"net"
 	"sync"
-	"sync/atomic"
 	"time"
 
 	"github.com/linkdata/socks5"
 )
 
 // Server is a SOCKS5 proxy server.
 type Server struct {
-	Started time.Time // time when Server.Serve() was called
-
 	// List of authentication providers. If nil, uses NoAuthAuthenticator.
 	// Order matters; they are tried in the given order.
 	Authenticators []Authenticator
@@ -26,9 +23,10 @@ type Server struct {
 	Logger socks5.Logger // If not nil, use this Logger (compatible with log/slog)
 	Debug  bool          // If true, output debug logging using Logger.Info
 
-	closed    atomic.Bool
 	mu        sync.Mutex // protects following
+	serving   int
 	listeners map[string]*listener
+	started   time.Time // time when Server.Serve() was called
 }
 
 func listenKey(client net.Conn, address string) (key string) {
@@ -47,7 +45,7 @@ func listenKey(client net.Conn, address string) (key string) {
 
 func (s *Server) getListener(ctx context.Context, client net.Conn, bindaddress string) (nl net.Listener, err error) {
 	err = net.ErrClosed
-	if !s.closed.Load() {
+	if s.Serving() > 0 {
 		err = nil
 		key := listenKey(client, bindaddress)
 		var lc net.ListenConfig
@@ -113,9 +111,10 @@ func (s *Server) maybeLogError(err error, msg string, keyvaluepairs ...any) {
 }
 
 func (s *Server) close() {
-	if !s.closed.Swap(true) {
-		s.mu.Lock()
-		defer s.mu.Unlock()
+	s.mu.Lock()
+	defer s.mu.Unlock()
+	s.serving--
+	if s.serving < 1 {
 		for _, l := range s.listeners {
 			_ = s.Debug && s.LogDebug("Server.close(): listener stop", "address", l.key)
 			l.refs.Store(0)
@@ -125,14 +124,26 @@ func (s *Server) close() {
 	}
 }
 
+// Serving returns the number of active calls to Serve()
+func (s *Server) Serving() (n int) {
+	s.mu.Lock()
+	n = s.serving
+	s.mu.Unlock()
+	return
+}
+
 // Serve accepts and handles incoming connections on the given listener.
 func (s *Server) Serve(ctx context.Context, l net.Listener) (err error) {
-	defer l.Close()
 	defer s.close()
-	s.Started = time.Now()
+	s.mu.Lock()
+	s.serving++
+	if s.listeners == nil {
+		s.listeners = make(map[string]*listener)
+		s.started = time.Now()
+	}
+	s.mu.Unlock()
 	errchan := make(chan error, 1)
 	s.LogInfo("listening", "addr", l.Addr())
-	s.listeners = make(map[string]*listener)
 	go s.listenerMaintenance(ctx)
 	go s.listen(ctx, errchan, l)
 	select {
@@ -145,7 +156,7 @@ func (s *Server) Serve(ctx context.Context, l net.Listener) (err error) {
 func (s *Server) listenerCleanup() {
 	s.mu.Lock()
 	defer s.mu.Unlock()
-	deadline := int64(time.Since(s.Started) - socks5.ListenerTimeout)
+	deadline := int64(time.Since(s.started) - socks5.ListenerTimeout)
 	for k, l := range s.listeners {
 		if refs := l.refs.Load(); refs < 1 {
 			if died := l.died.Load(); died < deadline {

server/server_test.go

@@ -8,6 +8,7 @@ import (
 	"log/slog"
 	"net"
 	"testing"
+	"time"
 
 	"github.com/linkdata/socks5"
 	"github.com/linkdata/socks5/client"
@@ -54,12 +55,72 @@ func TestServer_Logging(t *testing.T) {
 		t.Fatal(err)
 	}
 	defer listen.Close()
-	proxy := server.Server{Logger: slog.Default(), Debug: true}
-	proxy.LogDebug("debug")
-	proxy.LogInfo("info")
-	proxy.LogError("error")
+	srv := server.Server{Logger: slog.Default(), Debug: true}
+	srv.LogDebug("debug")
+	srv.LogInfo("info")
+	srv.LogError("error")
 }
 
 func TestServer_DialerSelector(t *testing.T) {
 	socks5test.InvalidCommand(t, srvfn, clifn)
 }
+
+func TestServer_Serve_CancelContext(t *testing.T) {
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+
+	listen, err := net.Listen("tcp", ":0")
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer listen.Close()
+	srv := server.Server{Logger: slog.Default(), Debug: true}
+	doneCh := make(chan struct{})
+	go func() {
+		defer close(doneCh)
+		srv.Serve(ctx, listen)
+	}()
+	cancel()
+	select {
+	case <-time.NewTimer(time.Second).C:
+		t.Error("timeout")
+	case <-doneCh:
+	}
+}
+
+func TestServer_Serve_TwoListeners(t *testing.T) {
+	ctx, cancel := context.WithTimeout(context.Background(), time.Second)
+	defer cancel()
+
+	listen1, err := net.Listen("tcp", ":0")
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer listen1.Close()
+
+	listen2, err := net.Listen("tcp", ":0")
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer listen2.Close()
+
+	srv := server.Server{Logger: slog.Default(), Debug: true}
+	go srv.Serve(ctx, listen1)
+	go srv.Serve(ctx, listen2)
+
+	for ctx.Err() == nil && srv.Serving() != 2 {
+		time.Sleep(time.Millisecond)
+	}
+
+	listen1.Close()
+
+	for ctx.Err() == nil && srv.Serving() != 1 {
+		time.Sleep(time.Millisecond)
+	}
+
+	listen2.Close()
+
+	for ctx.Err() == nil && srv.Serving() != 0 {
+		time.Sleep(time.Millisecond)
+	}
+}