Fix race condition in the ssh tunnel handover logic

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 recieving 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 waitig logic on the reciever loop side - it cares to wait for a new connection to be activated,
and now it does exactly this using connChanged condition variable.

Author: ilia-db

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())
+}

experimental/ssh/internal/proxy/proxy.go

@@ -27,6 +27,7 @@ const (
 type proxyConnection struct {
 	connID                    string
 	conn                      atomic.Value // *websocket.Conn
+	connChanged               sync.Cond
 	createWebsocketConnection createWebsocketConnectionFunc
 
 	handoverMutex           sync.Mutex
@@ -39,6 +40,7 @@ type createWebsocketConnectionFunc func(ctx context.Context, connID string) (*we
 func newProxyConnection(createConn createWebsocketConnectionFunc) *proxyConnection {
 	return &proxyConnection{
 		connID:                    uuid.NewString(),
+		connChanged:               sync.Cond{L: &sync.Mutex{}},
 		currentConnectionClosed:   make(chan error),
 		createWebsocketConnection: createConn,
 	}
@@ -65,6 +67,7 @@ func (pc *proxyConnection) connect(ctx context.Context) error {
 		return err
 	}
 	pc.conn.Store(conn)
+	pc.connChanged.Broadcast()
 	return nil
 }
 
@@ -74,6 +77,7 @@ func (pc *proxyConnection) accept(w http.ResponseWriter, r *http.Request) error
 		return err
 	}
 	pc.conn.Store(conn)
+	pc.connChanged.Broadcast()
 	return nil
 }
 
@@ -138,9 +142,7 @@ func (pc *proxyConnection) runReceivingLoop(ctx context.Context, dst io.Writer)
 				// Next time we read, we want to read from the new 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 {
-					return err
-				}
+				pc.waitForNewConnection(conn)
 				// Continue with the receiving loop, pc.conn is now the new connection.
 				continue
 			} else {
@@ -170,14 +172,12 @@ func (pc *proxyConnection) signalClosedConnection(err error) error {
 	}
 }
 
-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")
+func (pc *proxyConnection) waitForNewConnection(conn *websocket.Conn) {
+	pc.connChanged.L.Lock()
+	defer pc.connChanged.L.Unlock()
+	for pc.conn.Load() == conn {
+		pc.connChanged.Wait()
 	}
-	return nil
 }
 
 func (pc *proxyConnection) close() error {
@@ -226,6 +226,7 @@ func (pc *proxyConnection) initiateHandover(ctx context.Context) error {
 	}
 
 	pc.conn.Store(newConn)
+	pc.connChanged.Broadcast()
 
 	return nil
 }
@@ -272,6 +273,7 @@ func (pc *proxyConnection) acceptHandover(ctx context.Context, w http.ResponseWr
 	}
 
 	pc.conn.Store(newConn)
+	pc.connChanged.Broadcast()
 
 	return nil
 }

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{}{}
 			}
 		}