fix(ssh): bound tunnel client handshake timeout

cmd/internal/ssh_server.go

@@ -87,6 +87,7 @@ func (cmd *sshServerCmd) run(ctx context.Context) error {
 	// internal SSH cleanup could run. Liveness is decided via a per-directory
 	// flock the owning process holds for its lifetime; the kernel releases
 	// the flock on any process exit (including SIGKILL).
+	log.Debugf("starting ssh server: stdio=%v", cmd.stdio)
 	sshserver.SweepStaleAgentSockets()
 
 	server, err := sshserver.NewServer(
@@ -95,6 +96,7 @@ func (cmd *sshServerCmd) run(ctx context.Context) error {
 	if err != nil {
 		return fmt.Errorf("create ssh server: %w", err)
 	}
+	log.Debug("ssh server initialized")
 
 	if cmd.stdio {
 		return cmd.serveStdio(ctx, server)
@@ -108,6 +110,7 @@ func (cmd *sshServerCmd) serveStdio(ctx context.Context, server sshserver.Server
 	}
 	go shutdownOnCancel(ctx, server) // #nosec G118 -- see shutdownOnCancel.
 	lis := stdio.NewStdioListener(os.Stdin, os.Stdout)
+	log.Debug("serving ssh on stdio")
 	return ignoreServerClosed(server.Serve(lis))
 }
 

e2e/framework/retry.go

@@ -53,6 +53,7 @@ var retryableSSHPatterns = []string{
 	"connection timed out",
 	"broken pipe",
 	"workspace not found",
+	"ssh handshake made no progress",
 }
 
 // isRetryableSSHError returns true when the error indicates a transient SSH
@@ -91,8 +92,10 @@ func isRetryableDockerError(stderr string) bool {
 	return false
 }
 
-// execWithDockerRetry runs fn and retries if stderr indicates a transient
-// Docker registry error. Returns the last stdout, stderr, and error.
+// execWithDockerRetry runs fn and retries if the failure looks transient:
+// either a Docker registry error or a retryable SSH setup error in stderr
+// (up commands build the tunnel before the workspace starts). Returns the
+// last stdout, stderr, and error.
 func execWithDockerRetry(
 	ctx context.Context,
 	fn func(ctx context.Context) (stdout, stderr string, err error),
@@ -109,19 +112,20 @@ func execWithDockerRetry(
 		if lastErr == nil {
 			return lastStdout, lastStderr, nil
 		}
-		if !isRetryableDockerError(lastStderr) || attempt == dockerPullBackoff.Steps {
+		retryable := isRetryableDockerError(lastStderr) || isRetryableSSHError(lastErr, lastStderr)
+		if !retryable || attempt == dockerPullBackoff.Steps {
 			break
 		}
 		delay := nextBackoffDelay(dockerPullBackoff, attempt)
 		if !retryFitsBudget(ctx, delay) {
 			return lastStdout, lastStderr, fmt.Errorf(
-				"after %d attempts: retryable Docker error; retry not attempted because "+
+				"after %d attempts: retryable transient error; retry not attempted because "+
 					"remaining deadline budget was insufficient (next retry delay: %s): %w",
 				attempt, delay, lastErr,
 			)
 		}
 		ginkgo.GinkgoWriter.Printf(
-			"[retry] attempt %d failed with transient Docker error, retrying after %s: %s\n",
+			"[retry] attempt %d failed with transient error, retrying after %s: %s\n",
 			attempt, delay, lastErr,
 		)
 		if err := waitForRetry(ctx, delay); err != nil {

e2e/framework/retry_test.go

@@ -256,6 +256,37 @@ func TestExecWithDockerRetry_RetryThenSuccess(t *testing.T) {
 	assert.Equal(t, "", stderr)
 }
 
+func TestExecWithDockerRetry_HandshakeStallRetried(t *testing.T) {
+	withFastBackoffs(t)
+	calls := 0
+	stderrMsg := "start workspace: failed to create SSH client: ssh handshake made no progress for 15s"
+	out, _, err := execWithDockerRetry(context.Background(),
+		func(context.Context) (string, string, error) {
+			calls++
+			if calls == 1 {
+				return "", stderrMsg, transientExitErr(t)
+			}
+			return "ok", "", nil
+		},
+	)
+	require.NoError(t, err)
+	assert.Equal(t, 2, calls)
+	assert.Equal(t, "ok", out)
+}
+
+func TestExecWithDockerRetry_NonRetryableExitNotRetried(t *testing.T) {
+	withFastBackoffs(t)
+	calls := 0
+	_, _, err := execWithDockerRetry(context.Background(),
+		func(context.Context) (string, string, error) {
+			calls++
+			return "", "provider \"docker123\" does not exist", transientExitErr(t)
+		},
+	)
+	require.Error(t, err)
+	assert.Equal(t, 1, calls)
+}
+
 func TestExecWithDockerRetry_RetryExhausted(t *testing.T) {
 	withFastBackoffs(t)
 	calls := 0
@@ -365,6 +396,25 @@ func TestExecWithSSHRetry_RetryThenSuccess(t *testing.T) {
 	assert.Equal(t, "ok", out)
 }
 
+func TestExecWithSSHRetry_RetryHandshakeStallThenSuccess(t *testing.T) {
+	withFastBackoffs(t)
+	calls := 0
+	out, err := execWithSSHRetry(context.Background(), "ws",
+		func(context.Context) (string, string, error) {
+			calls++
+			if calls == 1 {
+				return "", "run in container: ssh client: ssh handshake made no progress for 15s", transientExitErr(
+					t,
+				)
+			}
+			return "ok", "", nil
+		},
+	)
+	require.NoError(t, err)
+	assert.Equal(t, 2, calls)
+	assert.Equal(t, "ok", out)
+}
+
 func TestExecWithSSHRetry_RetryExhausted(t *testing.T) {
 	withFastBackoffs(t)
 	calls := 0

pkg/ssh/helper.go

@@ -6,6 +6,8 @@ import (
 	"fmt"
 	"io"
 	"net"
+	"sync/atomic"
+	"time"
 
 	"github.com/devsy-org/devsy/pkg/stdio"
 	"golang.org/x/crypto/ssh"
@@ -90,8 +92,83 @@ func StdioClientFromKeyBytesWithUser(
 	return ClientFromConn(conn, user, keyBytes)
 }
 
+// HandshakeIdleTimeout fails a handshake that makes no progress for this
+// long. Progress-based instead of absolute so slow-but-alive links (high
+// latency is the norm for tunneled handshakes) always complete while a peer
+// that stops sending entirely surfaces an error to callers. Zero disables.
+const HandshakeIdleTimeout = 15 * time.Second
+
+// handshakeMaxTimeout caps a handshake whose peer never goes silent, e.g.
+// one that dribbles a byte at a time.
+const handshakeMaxTimeout = 2 * time.Minute
+
+// HandshakeTimeoutError reports a handshake that stopped making progress.
+type HandshakeTimeoutError struct{ idle time.Duration }
+
+func (e *HandshakeTimeoutError) Error() string {
+	return fmt.Sprintf("ssh handshake made no progress for %s", e.idle)
+}
+
+func (e *HandshakeTimeoutError) Timeout() bool { return true }
+
+// Temporary reports the error as transient so callers with net.Error retry
+// handling treat it as worth retrying.
+func (e *HandshakeTimeoutError) Temporary() bool { return true }
+
 // ClientFromConn creates an SSH client over an existing network connection.
+// The handshake is bounded by HandshakeIdleTimeout so a stalled peer surfaces
+// an error instead of blocking the tunnel forever. No deadlines are set on
+// the conn, so the established session is unaffected.
 func ClientFromConn(conn net.Conn, user string, keyBytes []byte) (*ssh.Client, error) {
+	return clientFromConn(conn, user, keyBytes, HandshakeIdleTimeout)
+}
+
+// activityConn records the time of the last successful read or write. The
+// stored time keeps its monotonic reading so idle measurement is immune to
+// wall-clock corrections.
+type activityConn struct {
+	net.Conn
+	tracking     atomic.Bool
+	lastActivity atomic.Pointer[time.Time]
+}
+
+func (c *activityConn) Read(b []byte) (int, error) {
+	n, err := c.Conn.Read(b)
+	if n > 0 {
+		c.recordActivity()
+	}
+	return n, err
+}
+
+func (c *activityConn) Write(b []byte) (int, error) {
+	n, err := c.Conn.Write(b)
+	if n > 0 {
+		c.recordActivity()
+	}
+	return n, err
+}
+
+func (c *activityConn) recordActivity() {
+	if !c.tracking.Load() {
+		return
+	}
+	now := time.Now()
+	c.lastActivity.Store(&now)
+}
+
+type handshakeResult struct {
+	c     ssh.Conn
+	chans <-chan ssh.NewChannel
+	reqs  <-chan *ssh.Request
+	err   error
+}
+
+func clientFromConn(
+	conn net.Conn,
+	user string,
+	keyBytes []byte,
+	handshakeIdleTimeout time.Duration,
+) (*ssh.Client, error) {
 	if conn == nil {
 		return nil, fmt.Errorf("connection is required")
 	}
@@ -101,14 +178,58 @@ func ClientFromConn(conn net.Conn, user string, keyBytes []byte) (*ssh.Client, e
 	}
 
 	clientConfig.User = user
+	if handshakeIdleTimeout > 0 {
+		return handshakeWithIdleTimeout(conn, clientConfig, handshakeIdleTimeout)
+	}
 	c, chans, req, err := ssh.NewClientConn(conn, "stdio", clientConfig)
 	if err != nil {
 		return nil, err
 	}
-
 	return ssh.NewClient(c, chans, handleKeepAliveRequests(req)), nil
 }
 
+// handshakeWithIdleTimeout runs the handshake in a goroutine and closes the
+// conn when it stops making progress, so stall detection works on conns that
+// do not support deadlines.
+func handshakeWithIdleTimeout(
+	conn net.Conn,
+	clientConfig *ssh.ClientConfig,
+	idle time.Duration,
+) (*ssh.Client, error) {
+	tracked := &activityConn{Conn: conn}
+	tracked.tracking.Store(true)
+	tracked.recordActivity()
+	result := make(chan handshakeResult, 1)
+	go func() {
+		c, chans, req, err := ssh.NewClientConn(tracked, "stdio", clientConfig)
+		result <- handshakeResult{c: c, chans: chans, reqs: req, err: err}
+	}()
+
+	ticker := time.NewTicker(idle / 4)
+	defer ticker.Stop()
+	deadline := time.After(handshakeMaxTimeout)
+	for {
+		select {
+		case res := <-result:
+			tracked.tracking.Store(false)
+			if res.err != nil {
+				return nil, res.err
+			}
+			return ssh.NewClient(res.c, res.chans, handleKeepAliveRequests(res.reqs)), nil
+		case <-ticker.C:
+			if time.Since(*tracked.lastActivity.Load()) > idle {
+				tracked.tracking.Store(false)
+				_ = conn.Close()
+				return nil, &HandshakeTimeoutError{idle: idle}
+			}
+		case <-deadline:
+			tracked.tracking.Store(false)
+			_ = conn.Close()
+			return nil, &HandshakeTimeoutError{idle: handshakeMaxTimeout}
+		}
+	}
+}
+
 func ConfigFromKeyBytes(keyBytes []byte) (*ssh.ClientConfig, error) {
 	clientConfig := &ssh.ClientConfig{
 		Auth:            []ssh.AuthMethod{},