fix(transfer): fix passive port multiplexing races

passive_multiplexer.go

@@ -85,38 +85,20 @@ func (m *passiveListenersManager) close() error {
 
 func (m *passiveListenersManager) getOrCreate(port int) (*sharedPassiveListener, error) {
 	m.mu.Lock()
-	if m.closed {
-		m.mu.Unlock()
+	defer m.mu.Unlock()
 
+	if m.closed {
 		return nil, net.ErrClosed
 	}
 	if listener, ok := m.listeners[port]; ok {
-		m.mu.Unlock()
-
 		return listener, nil
 	}
-	m.mu.Unlock()
 
 	listener, err := newSharedPassiveListener(port, m.logger.With("passivePort", port))
 	if err != nil {
 		return nil, err
 	}
 
-	m.mu.Lock()
-	defer m.mu.Unlock()
-
-	if m.closed {
-		_ = listener.close()
-
-		return nil, net.ErrClosed
-	}
-
-	if existing, ok := m.listeners[port]; ok {
-		_ = listener.close()
-
-		return existing, nil
-	}
-
 	m.listeners[port] = listener
 
 	return listener, nil
@@ -225,13 +207,12 @@ func (l *sharedPassiveListener) dispatch(conn net.Conn) {
 	}
 }
 
-func (l *sharedPassiveListener) release(remoteIP string) {
+func (l *sharedPassiveListener) release(reservation *passiveReservationListener) {
 	l.mu.Lock()
 	defer l.mu.Unlock()
 
-	if reservation, ok := l.reservations[remoteIP]; ok {
-		delete(l.reservations, remoteIP)
-		reservation.markReleased()
+	if l.reservations[reservation.remoteIP] == reservation {
+		delete(l.reservations, reservation.remoteIP)
 	}
 }
 
@@ -271,7 +252,7 @@ type passiveReservationListener struct {
 	closeOnce  sync.Once
 	stateMu    sync.Mutex
 	deadline   time.Time
-	released   bool
+	closing    bool
 	failureErr error
 }
 
@@ -307,8 +288,8 @@ func (l *passiveReservationListener) Accept() (net.Conn, error) {
 
 func (l *passiveReservationListener) Close() error {
 	l.closeOnce.Do(func() {
-		l.parent.release(l.remoteIP)
-		l.markReleased()
+		l.parent.release(l)
+		l.markClosing()
 
 		select {
 		case conn := <-l.connCh:
@@ -337,21 +318,26 @@ func (l *passiveReservationListener) SetDeadline(deadline time.Time) error {
 	return nil
 }
 
+// deliver hands conn over unless Close started: connCh has room for the single connection a reservation
+// gets, so the send cannot block, and Close drains connCh only after setting closing.
 func (l *passiveReservationListener) deliver(conn net.Conn) bool {
-	select {
-	case <-l.closedCh:
+	l.stateMu.Lock()
+	defer l.stateMu.Unlock()
+
+	if l.closing {
 		return false
-	default:
 	}
 
-	select {
-	case l.connCh <- conn:
-		l.markReleased()
+	l.connCh <- conn
 
-		return true
-	case <-l.closedCh:
-		return false
-	}
+	return true
+}
+
+func (l *passiveReservationListener) markClosing() {
+	l.stateMu.Lock()
+	defer l.stateMu.Unlock()
+
+	l.closing = true
 }
 
 func (l *passiveReservationListener) fail(err error) {
@@ -364,13 +350,6 @@ func (l *passiveReservationListener) fail(err error) {
 	_ = l.Close()
 }
 
-func (l *passiveReservationListener) markReleased() {
-	l.stateMu.Lock()
-	defer l.stateMu.Unlock()
-
-	l.released = true
-}
-
 func (l *passiveReservationListener) getDeadline() time.Time {
 	l.stateMu.Lock()
 	defer l.stateMu.Unlock()

passive_multiplexer_test.go

@@ -1,11 +1,14 @@
 package ftpserver
 
 import (
+	"context"
 	"errors"
 	"fmt"
 	"io"
 	"log/slog"
 	"net"
+	"sync"
+	"sync/atomic"
 	"testing"
 	"time"
 
@@ -18,6 +21,11 @@ type trackingTestConn struct {
 	closed bool
 }
 
+type closeCountingConn struct {
+	testNetConn
+	closed atomic.Bool
+}
+
 const testServerListenAddr = "127.0.0.1:0"
 
 var (
@@ -32,6 +40,12 @@ func (c *trackingTestConn) Close() error {
 	return nil
 }
 
+func (c *closeCountingConn) Close() error {
+	c.closed.Store(true)
+
+	return nil
+}
+
 func getFreePassivePort(t *testing.T) int {
 	t.Helper()
 
@@ -107,6 +121,109 @@ func TestPassiveListenersManagerRejectsSameIPForSamePort(t *testing.T) {
 	req.Nil(deadlineSetter)
 }
 
+func TestPassiveReservationCloseKeepsNewerReservation(t *testing.T) {
+	port := getFreePassivePort(t)
+	manager := newPassiveListenersManager(slog.New(slog.DiscardHandler))
+	defer func() {
+		require.NoError(t, manager.close())
+	}()
+
+	portRange := PortRange{Start: port, End: port}
+	clientIP := net.ParseIP("127.0.0.1")
+
+	_, first, _, err := manager.reserve(clientIP, portRange)
+	require.NoError(t, err)
+
+	dialer := net.Dialer{Timeout: time.Second}
+	firstClient, err := dialer.DialContext(context.Background(), "tcp", fmt.Sprintf("127.0.0.1:%d", port))
+	require.NoError(t, err)
+	defer func() { _ = firstClient.Close() }()
+
+	firstConn, err := first.Accept()
+	require.NoError(t, err)
+	defer func() { _ = firstConn.Close() }()
+
+	// The first data connection consumed its reservation, so the port is free for the same IP
+	_, second, secondDeadline, err := manager.reserve(clientIP, portRange)
+	require.NoError(t, err)
+	defer func() { _ = second.Close() }()
+
+	require.NoError(t, first.Close())
+
+	secondClient, err := dialer.DialContext(context.Background(), "tcp", fmt.Sprintf("127.0.0.1:%d", port))
+	require.NoError(t, err)
+	defer func() { _ = secondClient.Close() }()
+
+	require.NoError(t, secondDeadline.SetDeadline(time.Now().Add(5*time.Second)))
+	secondConn, err := second.Accept()
+	require.NoError(t, err)
+	require.NoError(t, secondConn.Close())
+}
+
+func TestPassiveReservationDeliverCloseRace(t *testing.T) {
+	const rounds = 50000
+
+	parent := &sharedPassiveListener{
+		logger:       slog.New(slog.DiscardHandler),
+		reservations: make(map[string]*passiveReservationListener),
+	}
+	clientIP := net.ParseIP("127.0.0.1")
+
+	for range rounds {
+		reservation, err := parent.reserve(clientIP)
+		require.NoError(t, err)
+		delete(parent.reservations, clientIP.String()) // as dispatch does before calling deliver
+
+		conn := &closeCountingConn{}
+		start := make(chan struct{})
+		var delivered bool
+		var waitGroup sync.WaitGroup
+
+		waitGroup.Go(func() {
+			<-start
+			delivered = reservation.deliver(conn)
+		})
+		waitGroup.Go(func() {
+			<-start
+			_ = reservation.Close()
+		})
+		close(start)
+		waitGroup.Wait()
+
+		if delivered {
+			require.True(t, conn.closed.Load(), "the connection was left with nobody to read or close it")
+		}
+		require.Empty(t, reservation.connCh)
+	}
+}
+
+func TestPassiveListenersManagerConcurrentFirstUse(t *testing.T) {
+	const (
+		rounds  = 50
+		clients = 8
+	)
+
+	for range rounds {
+		port := getFreePassivePort(t)
+		manager := newPassiveListenersManager(slog.New(slog.DiscardHandler))
+		portRange := PortRange{Start: port, End: port}
+		errs := make([]error, clients)
+
+		var waitGroup sync.WaitGroup
+		for client := range clients {
+			waitGroup.Go(func() {
+				_, _, _, errs[client] = manager.reserve(net.ParseIP(fmt.Sprintf("127.0.0.%d", client+1)), portRange)
+			})
+		}
+		waitGroup.Wait()
+
+		for _, err := range errs {
+			require.NoError(t, err)
+		}
+		require.NoError(t, manager.close())
+	}
+}
+
 func TestPassiveListenersManagerCloseReleasesReservation(t *testing.T) {
 	req := require.New(t)
 	port := getFreePassivePort(t)