Fix race condition in the ssh tunnel handover logic (#3888)

ilia-db · web-flow · commit c2d0e9dcb25b · 2025-11-10T09:22:44.000Z
## Changes
Proxy client can't initiate multiple handovers at the same time from its
point of view, but proxy server can still be in the situation where 2
handovers racing for handoverMutex simultaneously.

The first one can still wait for the close-message ack from the client
(on old ws connection), while the client sends a new handover http
request. Servers receiving loop has blockUntilHandoverComplete, which
tries to grab a handoverMutex, which races for it with the new
acceptHandover call. If the later executed first, we have a deadlock.

The fix is better waiting logic on the receiver loop side - it needs to
wait for a new connection to be activated, and now it does exactly this
using connSwapped channel.

Refactored the overall coordination logic a bit: we now use
handoverState atomic to store handover context (with a timeout), and two
channels necessary to coordinate between receiving loop and the handover
initiator. `isHandover` flag is not gone, as `handoverState` itself is
an indicator of when the handover is in progress.

### Testing
Also ran tests with `-race` and `-count 100` flags
\+ spawning manual ssh sessions with short handover timeout
diff --git a/experimental/ssh/internal/proxy/client_server_test.go b/experimental/ssh/internal/proxy/client_server_test.go
@@ -158,7 +158,7 @@ func TestHandover(t *testing.T) {
 	wg := sync.WaitGroup{}
 	wg.Go(func() {
 		for i := range TOTAL_MESSAGE_COUNT {
-			if i > 0 && i%MESSAGES_PER_CHUNK == 0 {
+			if i > 0 && i%MESSAGES_PER_CHUNK == 0 && i < TOTAL_MESSAGE_COUNT-1 {
 				handoverChan <- time.Now()
 			}
 			message := fmt.Appendf(nil, "message %d\n", i)
@@ -178,3 +178,41 @@ func TestHandover(t *testing.T) {
 	// clientOutput is created by appending incoming messages as they arrive, so we are also test correct order here
 	assert.Equal(t, expectedOutput, clientOutput.String())
 }
+
+// Tests handovers in quick succession with few messages in between.
+// Not a real world scenario, but it can help uncover potential race conditions or deadlocks.
+func TestQuickHandover(t *testing.T) {
+	server := createTestServer(t, 2, time.Hour)
+	defer server.Close()
+
+	handoverChan := make(chan time.Time)
+	requestHandoverTick := func() <-chan time.Time {
+		return handoverChan
+	}
+	clientInputWriter, clientOutput := createTestClient(t, server.URL, requestHandoverTick, nil)
+	defer clientInputWriter.Close()
+
+	expectedOutput := ""
+
+	wg := sync.WaitGroup{}
+	wg.Go(func() {
+		for i := range 16 {
+			if i == 4 || i == 8 || i == 12 {
+				handoverChan <- time.Now()
+			}
+			message := fmt.Appendf(nil, "message %d\n", i)
+			_, err := clientInputWriter.Write(message)
+			if err != nil {
+				t.Errorf("failed to write message %d: %v", i, err)
+			}
+			expectedOutput += string(message)
+		}
+	})
+
+	err := clientOutput.WaitForWrite(fmt.Appendf(nil, "message %d\n", 15))
+	require.NoError(t, err, "failed to receive the last message (%d)", 15)
+
+	wg.Wait()
+
+	assert.Equal(t, expectedOutput, clientOutput.String())
+}
diff --git a/experimental/ssh/internal/proxy/proxy.go b/experimental/ssh/internal/proxy/proxy.go
@@ -18,28 +18,89 @@ import (
 var errProxyEOF = errors.New("proxy EOF error")
 
 const (
-	proxyBufferSize                  = 4 * 1024 // Same as gorilla/websocket default read/write buffer sizes. Bigger payloads will be split into multiple ws frames.
-	proxyHandoverInitTimeout         = 30 * time.Second
-	proxyHandoverAcceptTimeout       = 25 * time.Second
-	proxyHandoverAckCloseConnTimeout = 15 * time.Second
+	// Same as gorilla/websocket default read/write buffer sizes. Bigger payloads will be split into multiple ws frames.
+	proxyBufferSize = 4 * 1024
+	// Timeout for the full handover process, when initiated by the client.
+	proxyHandoverInitTimeout = 30 * time.Second
+	// Timeout for the handover process, when accepted by the server.
+	proxyHandoverAcceptTimeout = 25 * time.Second
 )
 
+// handoverCoordination holds the context and channels used to coordinate a single handover operation
+// between the receiving loop and the handover initiator (initiateHandover or acceptHandover).
+type handoverCoordination struct {
+	// Context with timeout for the entire handover operation.
+	// Shared between the handover initiator and the receiving loop.
+	ctx context.Context
+	// Used by the receiving loop to signal about the closure of the current connection to the handover initiator.
+	// After signalling, the receiving loop will block until connSwapped channel is signaled.
+	connClosed chan error
+	// Used by the handover initiator to signal the receiving loop that it's safe to start reading from the new connection.
+	connSwapped chan struct{}
+}
+
+func (c *handoverCoordination) signalConnectionClosed(err error) error {
+	select {
+	case c.connClosed <- err:
+		return err
+	case <-c.ctx.Done():
+		return c.ctx.Err()
+	}
+}
+
+func (c *handoverCoordination) waitForConnectionToClose() error {
+	select {
+	case err := <-c.connClosed:
+		return err
+	case <-c.ctx.Done():
+		return c.ctx.Err()
+	}
+}
+
+func (c *handoverCoordination) signalConnectionSwapped() error {
+	select {
+	case c.connSwapped <- struct{}{}:
+		return nil
+	case <-c.ctx.Done():
+		return c.ctx.Err()
+	}
+}
+
+func (c *handoverCoordination) waitForConnectionToSwap() error {
+	select {
+	case <-c.connSwapped:
+		return nil
+	case <-c.ctx.Done():
+		return c.ctx.Err()
+	}
+}
+
+// proxyConnection is the main struct that manages the websocket connection and the handover process.
+// It works both on the client and the server side (see internal/client and internal/server packages).
+// It has 3 goroutines:
+// - Sending loop: reads from src and sends to the current connection.
+// - Receiving loop: reads from the current connection and writes to dst.
+// - Main: starts the other two (start method) and initiates or accepts handover (initiateHandover or acceptHandover).
 type proxyConnection struct {
-	connID                    string
-	conn                      atomic.Value // *websocket.Conn
+	// Each connection has a unique ID.
+	connID string
+	// Function to create a new websocket connection. Tests can override this to use a test websocket connection.
 	createWebsocketConnection createWebsocketConnectionFunc
-
-	handoverMutex           sync.Mutex
-	isHandover              atomic.Bool
-	currentConnectionClosed chan error
+	// Atomic that keeps the currently active connection.
+	// Can be swapped during handover.
+	conn atomic.Pointer[websocket.Conn]
+	// Prevents multiple handover processes from running concurrently.
+	// Blocks proxying any outgoing messages during the entire handover in the sending loop.
+	handoverMutex sync.Mutex
+	// Atomic that holds the current handover coordination channels, or nil if no handover is in progress.
+	handoverState atomic.Pointer[handoverCoordination]
 }
 
 type createWebsocketConnectionFunc func(ctx context.Context, connID string) (*websocket.Conn, error)
 
 func newProxyConnection(createConn createWebsocketConnectionFunc) *proxyConnection {
 	return &proxyConnection{
 		connID:                    uuid.NewString(),
-		currentConnectionClosed:   make(chan error),
 		createWebsocketConnection: createConn,
 	}
 }
@@ -114,7 +175,7 @@ func (pc *proxyConnection) runSendingLoop(ctx context.Context, src io.Reader) er
 func (pc *proxyConnection) sendMessage(mt int, data []byte) error {
 	pc.handoverMutex.Lock()
 	defer pc.handoverMutex.Unlock()
-	conn := pc.conn.Load().(*websocket.Conn)
+	conn := pc.conn.Load()
 	return conn.WriteMessage(mt, data)
 }
 
@@ -123,22 +184,23 @@ func (pc *proxyConnection) runReceivingLoop(ctx context.Context, dst io.Writer)
 		if ctx.Err() != nil {
 			return ctx.Err()
 		}
-		conn := pc.conn.Load().(*websocket.Conn)
+		conn := pc.conn.Load()
 		mt, data, err := conn.ReadMessage()
 		if err != nil {
 			// During handover a normal closure is expected, but any other error must stop the read loop (and eventually terminate the ssh session).
-			if pc.isHandover.Load() {
+			if handover := pc.handoverState.Load(); handover != nil {
 				var closeConnSignal error
 				if !websocket.IsCloseError(err, websocket.CloseNormalClosure) {
 					closeConnSignal = fmt.Errorf("failed to read from websocket during handover: %w", err)
 				}
-				if err := pc.signalClosedConnection(closeConnSignal); err != nil {
+				// Signal the current connection is closed to the handover initiator (initiateHandover or acceptHandover).
+				if err := handover.signalConnectionClosed(closeConnSignal); err != nil {
 					return err
 				}
-				// Next time we read, we want to read from the new connection.
+				// Wait for the handover initiator to swap the connection.
 				// While we wait for the handover to complete, the new connection might be getting incoming messages.
 				// They will be buffered by the TCP stack and will be read by us after the handover is complete.
-				if err := pc.blockUntilHandoverComplete(conn); err != nil {
+				if err := handover.waitForConnectionToSwap(); err != nil {
 					return err
 				}
 				// Continue with the receiving loop, pc.conn is now the new connection.
@@ -161,25 +223,6 @@ func (pc *proxyConnection) runReceivingLoop(ctx context.Context, dst io.Writer)
 	}
 }
 
-func (pc *proxyConnection) signalClosedConnection(err error) error {
-	select {
-	case pc.currentConnectionClosed <- err:
-		return err
-	case <-time.After(proxyHandoverAckCloseConnTimeout):
-		return fmt.Errorf("timeout waiting for acknowledgement of old connection closed message: %w", err)
-	}
-}
-
-func (pc *proxyConnection) blockUntilHandoverComplete(conn *websocket.Conn) error {
-	pc.handoverMutex.Lock()
-	defer pc.handoverMutex.Unlock()
-	// Sanity check to ensure we are not in the middle of a handover - this should not happen.
-	if pc.conn.Load() == conn {
-		return errors.New("handover mutex acquired while the old connection is still active")
-	}
-	return nil
-}
-
 func (pc *proxyConnection) close() error {
 	// Keep in mind that pc.sendMessage blocks during handover
 	err := pc.sendMessage(websocket.CloseMessage, websocket.FormatCloseMessage(websocket.CloseNormalClosure, ""))
@@ -198,35 +241,40 @@ func (pc *proxyConnection) initiateHandover(ctx context.Context) error {
 	pc.handoverMutex.Lock()
 	defer pc.handoverMutex.Unlock()
 
-	// When handover flag is set, the receiving loop handles a close message from the current connection
-	// as a signal to finish the handover and switch to the new connection.
-	pc.isHandover.Store(true)
-	defer pc.isHandover.Store(false)
-
-	ctx, cancel := context.WithTimeout(ctx, proxyHandoverInitTimeout)
+	handoverCtx, cancel := context.WithTimeout(ctx, proxyHandoverInitTimeout)
 	defer cancel()
+	handoverState := &handoverCoordination{
+		ctx:         handoverCtx,
+		connClosed:  make(chan error),
+		connSwapped: make(chan struct{}),
+	}
+	// Existence of the handoverState indicates to the receiving loop that we are in the middle of a handover process,
+	// and should treat close messages as a signal to finish the handover instead of erroring out.
+	pc.handoverState.Store(handoverState)
+	defer pc.handoverState.Store(nil)
 
 	// Create a new websocket connection by sending an /ssh?id=<connID> request to the server.
 	// When server realises it's an ID of an existing connection, it will start AcceptHandover process.
-	newConn, err := pc.createWebsocketConnection(ctx, pc.connID)
+	newConn, err := pc.createWebsocketConnection(handoverCtx, pc.connID)
 	if err != nil {
 		return fmt.Errorf("failed to create new websocket connection: %w", err)
 	}
 
 	// Wait for the server to close the old connection
-	select {
-	case err := <-pc.currentConnectionClosed:
-		if err != nil {
-			newConn.Close()
-			return fmt.Errorf("connection handover failed: %w", err)
-		}
-	case <-ctx.Done():
+	// (it does so when it receives an /ssh request with known connection ID and starts AcceptHandover process).
+	// Receiving loop will signal about closed connection to the coord.connClosed channel.
+	if err := handoverState.waitForConnectionToClose(); err != nil {
 		newConn.Close()
-		return ctx.Err()
+		return err
 	}
 
 	pc.conn.Store(newConn)
 
+	// Let the receiving loop know that the current connection is swapped and it's safe to start reading from it.
+	if err := handoverState.signalConnectionSwapped(); err != nil {
+		newConn.Close()
+		return err
+	}
 	return nil
 }
 
@@ -235,13 +283,17 @@ func (pc *proxyConnection) acceptHandover(ctx context.Context, w http.ResponseWr
 	pc.handoverMutex.Lock()
 	defer pc.handoverMutex.Unlock()
 
-	// When handover flag is set, the receiving loop handles a close message from the current connection
-	// as a signal to finish the handover and switch to the new connection.
-	pc.isHandover.Store(true)
-	defer pc.isHandover.Store(false)
-
-	ctx, cancel := context.WithTimeout(ctx, proxyHandoverAcceptTimeout)
+	handoverCtx, cancel := context.WithTimeout(ctx, proxyHandoverAcceptTimeout)
 	defer cancel()
+	handoverState := &handoverCoordination{
+		ctx:         handoverCtx,
+		connClosed:  make(chan error),
+		connSwapped: make(chan struct{}),
+	}
+	// Existence of the handoverState indicates to the receiving loop that we are in the middle of a handover process,
+	// and should treat close messages as a signal to finish the handover instead of erroring out.
+	pc.handoverState.Store(handoverState)
+	defer pc.handoverState.Store(nil)
 
 	newConn, err := pc.acceptWebsocketConnection(w, r)
 	if err != nil {
@@ -250,7 +302,7 @@ func (pc *proxyConnection) acceptHandover(ctx context.Context, w http.ResponseWr
 
 	// Signal the client to complete handover by closing the old connection.
 	// Not using pc.sendMessage here, because it's blocked by the handover mutex.
-	currentConn := pc.conn.Load().(*websocket.Conn)
+	currentConn := pc.conn.Load()
 	err = currentConn.WriteMessage(websocket.CloseMessage, websocket.FormatCloseMessage(websocket.CloseNormalClosure, "handover"))
 	if err != nil {
 		newConn.Close()
@@ -259,20 +311,20 @@ func (pc *proxyConnection) acceptHandover(ctx context.Context, w http.ResponseWr
 
 	// Wait for the client to acknowledge the closure of the old connection.
 	// On the client its done automatically by the websocket library with the default close handler.
-	// On the server we then receive a close error in the RunReceivingLoop and signal about it to the handoverOldConnClosed channel.
-	select {
-	case err := <-pc.currentConnectionClosed:
-		if err != nil {
-			newConn.Close()
-			return fmt.Errorf("connection handover failed: %w", err)
-		}
-	case <-ctx.Done():
+	// On the server we then receive a close error in the RunReceivingLoop and signal about it to the coord.connClosed channel.
+	if err := handoverState.waitForConnectionToClose(); err != nil {
 		newConn.Close()
-		return ctx.Err()
+		return err
 	}
 
 	pc.conn.Store(newConn)
 
+	// Let the receiving loop know that the current connection is swapped and it's safe to start reading from it.
+	if err := handoverState.signalConnectionSwapped(); err != nil {
+		newConn.Close()
+		return err
+	}
+
 	return nil
 }
 
diff --git a/experimental/ssh/internal/proxy/proxy_test.go b/experimental/ssh/internal/proxy/proxy_test.go
@@ -12,6 +12,7 @@ import (
 	"testing"
 	"time"
 
+	"github.com/databricks/cli/libs/log"
 	"github.com/gorilla/websocket"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
@@ -98,6 +99,7 @@ func setupTestServer(ctx context.Context, t *testing.T) *TestProxy {
 	serverOutput := newTestBuffer(t)
 	var serverProxy *proxyConnection
 	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		ctx := log.NewContext(ctx, log.GetLogger(ctx).With("Server", true))
 		if serverProxy != nil {
 			err := serverProxy.acceptHandover(ctx, w, r)
 			if err != nil {
@@ -138,6 +140,7 @@ func createTestWebsocketConnection(url string) (*websocket.Conn, error) {
 }
 
 func setupTestClient(ctx context.Context, t *testing.T, serverURL string) *TestProxy {
+	ctx = log.NewContext(ctx, log.GetLogger(ctx).With("Client", true))
 	clientInput, clientInputWriter := io.Pipe()
 	clientOutput := newTestBuffer(t)
 	wsURL := "ws" + serverURL[4:]
@@ -221,7 +224,7 @@ func TestConnectionHandover(t *testing.T) {
 		for i := range TOTAL_MESSAGE_COUNT {
 			client.Input.Write(createTestMessage("client", i)) // nolint:errcheck
 			server.Input.Write(createTestMessage("server", i)) // nolint:errcheck
-			if i > 0 && i%MESSAGES_PER_CHUNK == 0 {
+			if i > 0 && i%MESSAGES_PER_CHUNK == 0 && i < TOTAL_MESSAGE_COUNT-1 {
 				handoverChan <- struct{}{}
 			}
 		}
