diff --git a/.nextchanges/cli/ssh-tunnel-keepalive.md b/.nextchanges/cli/ssh-tunnel-keepalive.md new file mode 100644 index 00000000000..71802b97adb --- /dev/null +++ b/.nextchanges/cli/ssh-tunnel-keepalive.md @@ -0,0 +1 @@ +Fixed idle `databricks ssh connect` sessions disconnecting after a few minutes. The tunnel now sends a websocket keepalive every 20 seconds, so a session nobody is typing into stays connected without setting `ServerAliveInterval` in the SSH client config. diff --git a/experimental/ssh/cmd/connect.go b/experimental/ssh/cmd/connect.go index 2c50b871902..7502d70f560 100644 --- a/experimental/ssh/cmd/connect.go +++ b/experimental/ssh/cmd/connect.go @@ -119,6 +119,7 @@ Connect to a dedicated cluster: ShutdownDelay: shutdownDelay, MaxClients: maxClients, HandoverTimeout: handoverTimeout, + KeepaliveInterval: defaultKeepaliveInterval, ReleasesDir: releasesDir, ServerTimeout: max(serverTimeout, shutdownDelay), TaskStartupTimeout: startupTimeout, diff --git a/experimental/ssh/cmd/constants.go b/experimental/ssh/cmd/constants.go index 64c99b5bd48..dd5d3b2fdf2 100644 --- a/experimental/ssh/cmd/constants.go +++ b/experimental/ssh/cmd/constants.go @@ -3,10 +3,15 @@ package ssh import "time" const ( - defaultServerPort = 7772 - defaultMaxClients = 10 - defaultShutdownDelay = 10 * time.Minute - defaultHandoverTimeout = 30 * time.Minute + defaultServerPort = 7772 + defaultMaxClients = 10 + defaultShutdownDelay = 10 * time.Minute + defaultHandoverTimeout = 30 * time.Minute + // How often the client pings the tunnel websocket so an idle SSH session keeps the transport + // alive. Matches the keepalive interval of the vite bridge (libs/apps/vite/bridge.go), and sits + // well under both the ~9 minutes after which idle sessions were observed to drop and the + // 30 second SSH-level keepalive that was verified to prevent it. + defaultKeepaliveInterval = 20 * time.Second defaultEnvironmentVersion = 4 serverTimeout = 24 * time.Hour diff --git a/experimental/ssh/cmd/server.go b/experimental/ssh/cmd/server.go index 21c651b2365..3675a4a7fe8 100644 --- a/experimental/ssh/cmd/server.go +++ b/experimental/ssh/cmd/server.go @@ -41,7 +41,7 @@ and proxies them to local SSH daemon processes.`, cmd.MarkFlagRequired("authorized-key-secret-name") cmd.Flags().IntVar(&maxClients, "max-clients", defaultMaxClients, "Maximum number of SSH clients") - cmd.Flags().DurationVar(&shutdownDelay, "shutdown-delay", defaultShutdownDelay, "Delay before shutting down after no pings from clients") + cmd.Flags().DurationVar(&shutdownDelay, "shutdown-delay", defaultShutdownDelay, "Delay before shutting down the server when there are no active connections") cmd.Flags().StringVar(&version, "version", "", "Client version of the Databricks CLI") cmd.Flags().BoolVar(&serverless, "serverless", false, "Enable serverless mode for Jupyter initialization") cmd.Flags().StringVar(&usagePolicyID, "usage-policy-id", "", "Usage policy ID the job was submitted with") diff --git a/experimental/ssh/internal/client/client.go b/experimental/ssh/internal/client/client.go index 90e99f6b590..b2a8c9d0b7f 100644 --- a/experimental/ssh/internal/client/client.go +++ b/experimental/ssh/internal/client/client.go @@ -86,6 +86,8 @@ type ClientOptions struct { ServerMetadata string // How often the CLI should reconnect to the server with new auth. HandoverTimeout time.Duration + // How often the CLI pings the tunnel websocket to keep an idle session alive. + KeepaliveInterval time.Duration // Max amount of time the server process is allowed to live ServerTimeout time.Duration // Max amount of time to wait for the SSH server task to reach RUNNING state @@ -895,7 +897,7 @@ func runSSHProxy(ctx context.Context, client *databricks.WorkspaceClient, server requestHandoverTick := func() <-chan time.Time { return time.After(opts.HandoverTimeout) } - return proxy.RunClientProxy(ctx, os.Stdin, os.Stdout, requestHandoverTick, createConn) + return proxy.RunClientProxy(ctx, os.Stdin, os.Stdout, requestHandoverTick, opts.KeepaliveInterval, createConn) } // accessModeUILabel maps a cluster's access mode to the name shown in the Databricks UI. diff --git a/experimental/ssh/internal/proxy/client.go b/experimental/ssh/internal/proxy/client.go index a1e8389e7ff..9d40cac5169 100644 --- a/experimental/ssh/internal/proxy/client.go +++ b/experimental/ssh/internal/proxy/client.go @@ -9,6 +9,7 @@ import ( "time" "github.com/databricks/cli/libs/log" + "github.com/gorilla/websocket" "golang.org/x/sync/errgroup" ) @@ -37,8 +38,25 @@ func (f *firstByteWriter) Write(p []byte) (int, error) { return f.w.Write(p) } -func RunClientProxy(ctx context.Context, src io.ReadCloser, dst io.Writer, requestHandoverTick func() <-chan time.Time, createConn createWebsocketConnectionFunc) error { - proxy := newProxyConnection(createConn) +// logPongs wraps a connection factory so every connection it creates — the initial one and each +// one a handover creates — logs the pongs coming back for our keepalive pings. Debug visibility only: +// the receiving loop stays the only judge of whether a connection is alive. +func logPongs(ctx context.Context, createConn createWebsocketConnectionFunc) createWebsocketConnectionFunc { + return func(connCtx context.Context, connID string) (*websocket.Conn, error) { + conn, err := createConn(connCtx, connID) + if err != nil { + return nil, err + } + conn.SetPongHandler(func(string) error { + log.Debugf(ctx, "Received websocket keepalive pong") + return nil + }) + return conn, nil + } +} + +func RunClientProxy(ctx context.Context, src io.ReadCloser, dst io.Writer, requestHandoverTick func() <-chan time.Time, keepaliveInterval time.Duration, createConn createWebsocketConnectionFunc) error { + proxy := newProxyConnection(logPongs(ctx, createConn)) log.Infof(ctx, "Establishing SSH proxy connection...") ctx, cancel := context.WithCancel(ctx) defer cancel() @@ -68,9 +86,38 @@ func RunClientProxy(ctx context.Context, src io.ReadCloser, dst io.Writer, reque } } }) + g.Go(func() error { + // Keep the websocket carrying traffic while the SSH session is idle. Both proxy loops + // are data-driven, so an idle session puts no frames on the connection at all and the + // server side reaps the stream it then considers dead (websocket close 4000). + ticker := time.NewTicker(keepaliveInterval) + defer ticker.Stop() + for { + select { + case <-gCtx.Done(): + return gCtx.Err() + case <-ticker.C: + // A ping that ticks during a handover goes to the connection being replaced and + // may simply fail. Harmless: a handover establishes a fresh connection, which + // resets the peer's idle clock anyway, and the next tick uses the new one. + if err := proxy.sendPing(); err != nil { + // Not fatal, but not harmless either: gorilla puts the connection into a + // permanent write-error state after any failed write, so nothing more can + // be sent on it. Reads are unaffected and may still be delivering output + // the user is waiting on, so the session is left to end the way it would + // anyway — the next write fails and the sending loop reports it. + log.Warnf(gCtx, "Failed to send websocket keepalive ping, the connection can no longer send: %v", err) + } else { + // The driver proxy does not return pongs (verified end to end), so this + // line is the only evidence in a customer's log that pings were flowing. + log.Debugf(gCtx, "Sent websocket keepalive ping") + } + } + } + }) g.Go(func() error { // When proxy.start returns (EOF from ssh, or the server closing the connection), - // cancel so the handover goroutine stops too and g.Wait can return. + // cancel so the handover and keepalive goroutines stop too and g.Wait can return. defer cancel() return proxy.start(gCtx, src, wrappedDst) }) diff --git a/experimental/ssh/internal/proxy/client_server_test.go b/experimental/ssh/internal/proxy/client_server_test.go index 1915cc07c88..7aa45d35e81 100644 --- a/experimental/ssh/internal/proxy/client_server_test.go +++ b/experimental/ssh/internal/proxy/client_server_test.go @@ -38,7 +38,7 @@ type testClient struct { Cleanup func() } -func createTestClient(t *testing.T, serverURL string, requestHandoverTick func() <-chan time.Time, errChan chan error) *testClient { +func createTestClient(t *testing.T, serverURL string, requestHandoverTick func() <-chan time.Time, keepaliveInterval time.Duration, errChan chan error) *testClient { ctx := cmdio.MockDiscard(t.Context()) clientInput, clientInputWriter := io.Pipe() clientOutput := newTestBuffer(t) @@ -49,13 +49,11 @@ func createTestClient(t *testing.T, serverURL string, requestHandoverTick func() return conn, err } if requestHandoverTick == nil { - requestHandoverTick = func() <-chan time.Time { - return time.After(time.Hour) - } + requestHandoverTick = neverTick } wg := sync.WaitGroup{} wg.Go(func() { - err := RunClientProxy(ctx, clientInput, clientOutput, requestHandoverTick, createConn) + err := RunClientProxy(ctx, clientInput, clientOutput, requestHandoverTick, keepaliveInterval, createConn) if err != nil && !errors.Is(err, context.Canceled) && !errors.Is(err, io.ErrClosedPipe) { if errChan != nil { errChan <- err @@ -78,7 +76,7 @@ func createTestClient(t *testing.T, serverURL string, requestHandoverTick func() func TestClientServerEcho(t *testing.T) { server := createTestServer(t, 2, time.Hour) defer server.Close() - client := createTestClient(t, server.URL, nil, nil) + client := createTestClient(t, server.URL, nil, time.Hour, nil) defer client.Cleanup() testMsg1 := []byte("test message 1\n") @@ -100,9 +98,9 @@ func TestClientServerEcho(t *testing.T) { func TestMultipleClients(t *testing.T) { server := createTestServer(t, 2, time.Hour) defer server.Close() - client1 := createTestClient(t, server.URL, nil, nil) + client1 := createTestClient(t, server.URL, nil, time.Hour, nil) defer client1.Cleanup() - client2 := createTestClient(t, server.URL, nil, nil) + client2 := createTestClient(t, server.URL, nil, time.Hour, nil) defer client2.Cleanup() messageCount := 10 @@ -131,9 +129,9 @@ func TestMaxClients(t *testing.T) { maxClients := 2 server := createTestServer(t, maxClients, time.Hour) defer server.Close() - client1 := createTestClient(t, server.URL, nil, nil) + client1 := createTestClient(t, server.URL, nil, time.Hour, nil) defer client1.Cleanup() - client2 := createTestClient(t, server.URL, nil, nil) + client2 := createTestClient(t, server.URL, nil, time.Hour, nil) defer client2.Cleanup() testMsg1 := []byte("test message 1\n") @@ -147,7 +145,7 @@ func TestMaxClients(t *testing.T) { require.NoError(t, err) errChan := make(chan error, 1) - client3 := createTestClient(t, server.URL, nil, errChan) + client3 := createTestClient(t, server.URL, nil, time.Hour, errChan) defer client3.Cleanup() select { case err = <-errChan: @@ -158,6 +156,17 @@ func TestMaxClients(t *testing.T) { } func TestHandover(t *testing.T) { + t.Run("without keepalive", func(t *testing.T) { + runHandoverExchange(t, time.Hour) + }) + // Pings and the data stream share the connection's write lock: they must not corrupt or reorder + // the stream, nor trip gorilla's concurrent-write panic. + t.Run("with keepalive", func(t *testing.T) { + runHandoverExchange(t, time.Millisecond) + }) +} + +func runHandoverExchange(t *testing.T, keepaliveInterval time.Duration) { server := createTestServer(t, 2, time.Hour) defer server.Close() @@ -165,7 +174,7 @@ func TestHandover(t *testing.T) { requestHandoverTick := func() <-chan time.Time { return handoverChan } - client := createTestClient(t, server.URL, requestHandoverTick, nil) + client := createTestClient(t, server.URL, requestHandoverTick, keepaliveInterval, nil) defer client.Cleanup() var expectedOutput []byte @@ -204,7 +213,7 @@ func TestQuickHandover(t *testing.T) { requestHandoverTick := func() <-chan time.Time { return handoverChan } - client := createTestClient(t, server.URL, requestHandoverTick, nil) + client := createTestClient(t, server.URL, requestHandoverTick, time.Hour, nil) defer client.Cleanup() var expectedOutput []byte @@ -256,7 +265,7 @@ func TestClientExitsWhenServerCommandFails(t *testing.T) { done := make(chan error, 1) go func() { - done <- RunClientProxy(ctx, src, io.Discard, requestHandoverTick, createConn) + done <- RunClientProxy(ctx, src, io.Discard, requestHandoverTick, time.Hour, createConn) }() select { @@ -303,7 +312,7 @@ func TestClientTimesOutWhenServerSendsNothing(t *testing.T) { done := make(chan error, 1) go func() { - done <- RunClientProxy(ctx, src, io.Discard, requestHandoverTick, createConn) + done <- RunClientProxy(ctx, src, io.Discard, requestHandoverTick, time.Hour, createConn) }() select { diff --git a/experimental/ssh/internal/proxy/keepalive_test.go b/experimental/ssh/internal/proxy/keepalive_test.go new file mode 100644 index 00000000000..c972624f6f1 --- /dev/null +++ b/experimental/ssh/internal/proxy/keepalive_test.go @@ -0,0 +1,325 @@ +package proxy + +import ( + "context" + "errors" + "fmt" + "io" + "net" + "net/http" + "net/http/httptest" + "os" + "sync" + "sync/atomic" + "testing" + "time" + + "github.com/databricks/cli/libs/cmdio" + "github.com/gorilla/websocket" + "github.com/stretchr/testify/require" +) + +// Write modes for pausableConn. +const ( + connWriteOK = iota + connWriteFail + connWritePark +) + +var errTestWriteFailed = errors.New("test: socket write failed") + +// unboundedParkDuration stands in for "parks until the kernel gives up", which is what a write with +// no deadline does on a stalled socket. It must outlast the assertions of any test that parks a +// write, so that a write path which forgets to bound itself is measured as a stall, not as a failure. +const unboundedParkDuration = 30 * time.Second + +// pausableConn emulates the socket conditions a keepalive meets on a stalled peer: writes can be +// made to fail outright, or to park the way a full send buffer does — until the write's deadline +// expires, or effectively forever if it has none. Wrapping the socket rather than the websocket keeps +// the production write path (gorilla's own locking and deadline handling) in the test. +type pausableConn struct { + net.Conn + mode atomic.Int32 + deadline atomic.Pointer[time.Time] + parked chan struct{} + signalParked func() +} + +func newPausableConn(conn net.Conn) *pausableConn { + parked := make(chan struct{}) + return &pausableConn{ + Conn: conn, + parked: parked, + signalParked: sync.OnceFunc(func() { close(parked) }), + } +} + +func (c *pausableConn) SetWriteDeadline(t time.Time) error { + c.deadline.Store(&t) + return c.Conn.SetWriteDeadline(t) +} + +func (c *pausableConn) Write(p []byte) (int, error) { + switch c.mode.Load() { + case connWriteFail: + return 0, errTestWriteFailed + case connWritePark: + c.signalParked() + if d := c.deadline.Load(); d != nil && !d.IsZero() { + time.Sleep(time.Until(*d)) + } else { + time.Sleep(unboundedParkDuration) + } + return 0, os.ErrDeadlineExceeded + } + return c.Conn.Write(p) +} + +// startKeepaliveTestServer stands up a websocket peer that behaves like the SSH server side of the +// tunnel for an idle session: it sends the first bytes (which the client waits for before it +// considers the session established), then only reads. Pings it receives are reported on the +// returned channel and answered with a pong, the same way gorilla's default ping handler does. +func startKeepaliveTestServer(t *testing.T) (*httptest.Server, <-chan struct{}) { + pings := make(chan struct{}, 1) + upgrader := websocket.Upgrader{} + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + conn, err := upgrader.Upgrade(w, r, nil) + if err != nil { + return + } + defer conn.Close() + conn.SetPingHandler(func(appData string) error { + select { + case pings <- struct{}{}: + default: + } + return conn.WriteControl(websocket.PongMessage, []byte(appData), time.Now().Add(time.Second)) + }) + if err := conn.WriteMessage(websocket.BinaryMessage, []byte("SSH-2.0-test\r\n")); err != nil { + return + } + // Ping handlers only run while a read is in progress, so keep reading until the client goes away. + for { + if _, _, err := conn.ReadMessage(); err != nil { + return + } + } + })) + return server, pings +} + +// keepaliveTestDialer returns a connection factory for RunClientProxy. onNetConn, when set, receives +// each connection's underlying socket so a test can control how its writes behave. +func keepaliveTestDialer(serverURL string, onNetConn func(*pausableConn)) createWebsocketConnectionFunc { + wsURL := "ws" + serverURL[4:] + dialer := websocket.Dialer{ + NetDialContext: func(ctx context.Context, network, addr string) (net.Conn, error) { + conn, err := net.Dial(network, addr) + if err != nil { + return nil, err + } + wrapped := newPausableConn(conn) + if onNetConn != nil { + onNetConn(wrapped) + } + return wrapped, nil + }, + } + return func(ctx context.Context, connID string) (*websocket.Conn, error) { + conn, _, err := dialer.DialContext(ctx, fmt.Sprintf("%s?id=%s", wsURL, connID), nil) // nolint:bodyclose + return conn, err + } +} + +// TestKeepalivePingReachesServer covers the fix itself: an idle session sends no data, so without +// the keepalive nothing at all crosses the websocket and the server side eventually reaps the +// stream it considers dead. +func TestKeepalivePingReachesServer(t *testing.T) { + ctx := cmdio.MockDiscard(t.Context()) + server, pings := startKeepaliveTestServer(t) + defer server.Close() + + // Never written to: the session stays idle for the whole test. + src, _ := io.Pipe() + done := make(chan error, 1) + go func() { + done <- RunClientProxy(ctx, src, io.Discard, neverTick, 20*time.Millisecond, keepaliveTestDialer(server.URL, nil)) + }() + + select { + case <-pings: + case err := <-done: + t.Fatalf("session ended before a keepalive ping arrived: %v", err) + case <-time.After(10 * time.Second): + t.Fatal("no keepalive ping arrived at the server") + } +} + +// TestKeepalivePingFailureDoesNotEndSession asserts a keepalive can never be the thing that ends a +// session: it has no information the data loops lack, and it runs in the errgroup that drives them, +// so a returned error would tear the session down. +func TestKeepalivePingFailureDoesNotEndSession(t *testing.T) { + ctx := cmdio.MockDiscard(t.Context()) + server, _ := startKeepaliveTestServer(t) + defer server.Close() + + var socket atomic.Pointer[pausableConn] + src, _ := io.Pipe() + done := make(chan error, 1) + go func() { + done <- RunClientProxy(ctx, src, io.Discard, neverTick, 20*time.Millisecond, + keepaliveTestDialer(server.URL, func(c *pausableConn) { socket.Store(c) })) + }() + + // Fail every write once the connection is up, leaving reads healthy: the connection is otherwise + // fine and the session must survive the pings that then fail. + require.Eventually(t, func() bool { return socket.Load() != nil }, 10*time.Second, 10*time.Millisecond) + socket.Load().mode.Store(connWriteFail) + + // Long enough for tens of pings to be attempted and fail at the 20ms interval above. + select { + case err := <-done: + t.Fatalf("session ended after a failed keepalive ping: %v", err) + case <-time.After(time.Second): + } +} + +// TestKeepalivePingParkedInWriteDoesNotStallClose covers the hazard of writing from a third +// goroutine: a ping on a stalled or half-open connection parks in the socket write while holding the +// websocket's write lock, which the closing handshake also needs. Its deadline is what keeps that +// from lasting until the kernel abandons its retransmits, minutes later. +// +// A ping write with no deadline parks for unboundedParkDuration here, so this fails if the ping ever +// goes back to a write path that does not bound itself. +func TestKeepalivePingParkedInWriteDoesNotStallClose(t *testing.T) { + ctx := cmdio.MockDiscard(t.Context()) + server, _ := startKeepaliveTestServer(t) + defer server.Close() + + var socket atomic.Pointer[pausableConn] + proxy := newProxyConnection(keepaliveTestDialer(server.URL, func(c *pausableConn) { socket.Store(c) })) + require.NoError(t, proxy.connect(ctx)) + require.NotNil(t, socket.Load()) + + socket.Load().mode.Store(connWritePark) + pingDone := make(chan error, 1) + go func() { + pingDone <- proxy.sendPing() + }() + select { + case <-socket.Load().parked: + case <-time.After(10 * time.Second): + t.Fatal("the keepalive ping did not park in the socket write") + } + + start := time.Now() + closeErr := proxy.close() + require.Less(t, time.Since(start), proxyPingWriteTimeout+5*time.Second, + "the closing handshake waited on the parked keepalive ping for longer than its write deadline allows") + // The write itself fails, which close() reports; the point is that it was not held indefinitely. + require.Error(t, closeErr) + + select { + case err := <-pingDone: + require.Error(t, err, "a parked ping write must end in an error, not succeed") + case <-time.After(10 * time.Second): + t.Fatal("the keepalive ping never returned from its parked write") + } +} + +// TestKeepalivePingFailureDoesNotHangTheSession covers what a failed ping actually costs. gorilla +// puts the connection into a permanent write-error state after any failed write, so one timed-out +// keepalive stops the close message going out too — and a session whose close message never reaches +// the peer used to wait forever for the peer to close the connection, which is the same silent +// black hole this feature exists to remove. The session must end, promptly and with an error. +func TestKeepalivePingFailureDoesNotHangTheSession(t *testing.T) { + ctx := cmdio.MockDiscard(t.Context()) + server, _ := startKeepaliveTestServer(t) + defer server.Close() + + var socket atomic.Pointer[pausableConn] + src, srcWriter := io.Pipe() + done := make(chan error, 1) + go func() { + done <- RunClientProxy(ctx, src, io.Discard, neverTick, 20*time.Millisecond, + keepaliveTestDialer(server.URL, func(c *pausableConn) { socket.Store(c) })) + }() + + require.Eventually(t, func() bool { return socket.Load() != nil }, 10*time.Second, 10*time.Millisecond) + socket.Load().mode.Store(connWriteFail) + // Let a ping fail, which is what poisons the connection. + time.Sleep(100 * time.Millisecond) + + // The session ends the way it would without a keepalive at all: the next write fails and the + // sending loop reports it. Before, the teardown could not close the connection and this hung. + _, err := srcWriter.Write([]byte("keystroke")) + require.NoError(t, err) + + select { + case err := <-done: + require.Error(t, err, "a session that can no longer send must end with an error, not silently") + case <-time.After(30 * time.Second): + t.Fatal("session hung after a failed keepalive ping poisoned the connection") + } +} + +// TestKeepalivePingDuringHandoverDoesNotDisruptIt asserts the two periodic behaviours of the tunnel +// stay independent: a ping sent while a handover is in flight neither waits for the handover nor +// breaks it. The keepalive writes from a third goroutine, so nothing else guarantees this. +func TestKeepalivePingDuringHandoverDoesNotDisruptIt(t *testing.T) { + ctx := t.Context() + + server := setupTestServer(ctx, t) + defer server.Cleanup() + + // Holds the handover open at the point where it has taken the handover mutex but has not yet + // swapped the connection, which is when a ping must neither block nor interfere. + handoverDialing := make(chan struct{}) + releaseHandover := make(chan struct{}) + release := sync.OnceFunc(func() { close(releaseHandover) }) + var dials atomic.Int32 + + client := setupTestClientWithDialHook(ctx, t, server.URL, func() { + if dials.Add(1) > 1 { + close(handoverDialing) + <-releaseHandover + } + }) + defer client.Cleanup() + // Deferred after client.Cleanup so it runs before it: a t.Fatal below would otherwise leave the + // handover parked in the dial hook, and cleanup waits on proxy loops that cannot finish until it + // is released — which wedges the whole package instead of failing one test. + defer release() + + handoverDone := make(chan error, 1) + go func() { + handoverDone <- client.Proxy.initiateHandover(ctx) + }() + <-handoverDialing + + pingDone := make(chan error, 1) + go func() { + pingDone <- client.Proxy.sendPing() + }() + + select { + case err := <-pingDone: + require.NoError(t, err) + case <-time.After(proxyPingWriteTimeout): + t.Fatal("keepalive ping waited on the in-flight handover") + } + + release() + + select { + case err := <-handoverDone: + require.NoError(t, err) + case <-time.After(10 * time.Second): + t.Fatal("handover did not complete after a concurrent keepalive ping") + } + + // The tunnel still carries data on the connection the handover installed. + _, err := client.Input.Write(createTestMessage("client", 1)) + require.NoError(t, err) + require.NoError(t, server.Output.WaitForWrite(createTestMessage("client", 1))) +} diff --git a/experimental/ssh/internal/proxy/proxy.go b/experimental/ssh/internal/proxy/proxy.go index e761cbf54c4..685b02a9b46 100644 --- a/experimental/ssh/internal/proxy/proxy.go +++ b/experimental/ssh/internal/proxy/proxy.go @@ -5,6 +5,7 @@ import ( "errors" "fmt" "io" + "net" "net/http" "os" "sync" @@ -29,6 +30,10 @@ const ( proxyHandoverInitTimeout = 30 * time.Second // Timeout for the handover process, when accepted by the server. proxyHandoverAcceptTimeout = 25 * time.Second + // Bounds how long a keepalive ping may hold the websocket's write lock. A stalled or half-open + // connection parks a write until the kernel gives up retransmitting (~15 minutes with Linux + // defaults), and close() and the sending loop need that same lock, so the ping caps its wait. + proxyPingWriteTimeout = 5 * time.Second ) // handoverCoordination holds the context and channels used to coordinate a single handover operation @@ -134,7 +139,7 @@ func (pc *proxyConnection) start(ctx context.Context, src io.ReadCloser, dst io. // Both loops can still be stuck on conn.ReadMessage or src.Read and won't notice context cancellation, // so we close the connection and the source (sshd stdout pipe or ssh client stdio) to unblock them. <-gCtx.Done() - return errors.Join(pc.close(), pc.closeSource(src)) + return errors.Join(pc.close(), pc.closeConnection(), pc.closeSource(src)) }) err := g.Wait() if err == nil || isNormalClosure(err) { @@ -204,6 +209,15 @@ func (pc *proxyConnection) sendMessage(mt int, data []byte) error { return conn.WriteMessage(mt, data) } +// sendPing writes a keepalive ping on the current connection. Unlike sendMessage it takes neither +// the handover mutex nor an unbounded wait: gorilla permits WriteControl concurrently with the data +// writes, and its deadline bounds how long a stalled socket holds the connection's write lock, which +// close() and the sending loop also need. +func (pc *proxyConnection) sendPing() error { + conn := pc.conn.Load() + return conn.WriteControl(websocket.PingMessage, nil, time.Now().Add(proxyPingWriteTimeout)) +} + func (pc *proxyConnection) runReceivingLoop(ctx context.Context, dst io.Writer) error { for { if ctx.Err() != nil { @@ -261,6 +275,19 @@ func (pc *proxyConnection) close() error { return nil } +// closeConnection closes the underlying websocket. The close message pc.close sends only ends the +// session if the peer is still there to react to it by closing the connection, and it does not even +// go out once a failed write has put the connection into gorilla's permanent write-error state (one +// timed-out keepalive ping is enough). Without this the receiving loop stays blocked in ReadMessage +// and the session hangs instead of exiting. +func (pc *proxyConnection) closeConnection() error { + err := pc.conn.Load().Close() + if errors.Is(err, net.ErrClosed) { + return nil + } + return err +} + func (pc *proxyConnection) closeSource(src io.ReadCloser) error { err := src.Close() if err != nil && (errors.Is(err, os.ErrClosed) || errors.Is(err, io.ErrClosedPipe)) { diff --git a/experimental/ssh/internal/proxy/proxy_test.go b/experimental/ssh/internal/proxy/proxy_test.go index 0e1db9021e2..1e3af34f40d 100644 --- a/experimental/ssh/internal/proxy/proxy_test.go +++ b/experimental/ssh/internal/proxy/proxy_test.go @@ -139,12 +139,26 @@ func createTestWebsocketConnection(url string) (*websocket.Conn, error) { return conn, err } +// neverTick is a tick channel that never fires, for tests that don't exercise a periodic behaviour. +func neverTick() <-chan time.Time { + return time.After(time.Hour) +} + func setupTestClient(ctx context.Context, t *testing.T, serverURL string) *TestProxy { + return setupTestClientWithDialHook(ctx, t, serverURL, nil) +} + +// setupTestClientWithDialHook is setupTestClient with a hook called on every websocket dial: the +// initial connection and each one a handover creates. +func setupTestClientWithDialHook(ctx context.Context, t *testing.T, serverURL string, onDial func()) *TestProxy { ctx = log.NewContext(ctx, log.GetLogger(ctx).With("Client", true)) clientInput, clientInputWriter := io.Pipe() clientOutput := newTestBuffer(t) wsURL := "ws" + serverURL[4:] clientProxy := newProxyConnection(func(ctx context.Context, connID string) (*websocket.Conn, error) { + if onDial != nil { + onDial() + } return createTestWebsocketConnection(wsURL) }) err := clientProxy.connect(ctx)