diff --git a/.nextchanges/cli/ssh-tunnel-keepalive.md b/.nextchanges/cli/ssh-tunnel-keepalive.md index 71802b97adb..cf1d33a18e5 100644 --- a/.nextchanges/cli/ssh-tunnel-keepalive.md +++ b/.nextchanges/cli/ssh-tunnel-keepalive.md @@ -1 +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. +Fixed idle `databricks ssh connect` sessions disconnecting after a few minutes, on dedicated clusters and on serverless. The tunnel now keeps itself warm: the SSH client and the SSH server on the compute exchange keepalives every 30 seconds, and the CLI's proxy pings the tunnel's websocket every 20 seconds. A session nobody is typing into stays connected, with no need to set `ServerAliveInterval` by hand. diff --git a/experimental/ssh/internal/client/client.go b/experimental/ssh/internal/client/client.go index b2a8c9d0b7f..031bbd4d7af 100644 --- a/experimental/ssh/internal/client/client.go +++ b/experimental/ssh/internal/client/client.go @@ -820,6 +820,7 @@ func buildSSHArgs(userName, privateKeyPath, proxyCommand, hostName, wsHome strin "-o", "IdentitiesOnly=yes", "-o", "StrictHostKeyChecking=accept-new", "-o", "ConnectTimeout=360", + "-o", "ServerAliveInterval=" + strconv.Itoa(sshconfig.ServerAliveIntervalSeconds), "-o", "ProxyCommand=" + proxyCommand, } if opts.UserKnownHostsFile != "" { diff --git a/experimental/ssh/internal/client/client_internal_test.go b/experimental/ssh/internal/client/client_internal_test.go index 1e3accbaf46..cb47cd71bb8 100644 --- a/experimental/ssh/internal/client/client_internal_test.go +++ b/experimental/ssh/internal/client/client_internal_test.go @@ -5,10 +5,13 @@ import ( "encoding/json" "errors" "fmt" + "slices" + "strconv" "strings" "testing" "time" + "github.com/databricks/cli/experimental/ssh/internal/sshconfig" "github.com/databricks/cli/libs/cmdio" "github.com/databricks/cli/libs/telemetry/protos" "github.com/databricks/databricks-sdk-go/experimental/mocks" @@ -373,6 +376,17 @@ func TestBuildRemoteShellArgs(t *testing.T) { }) } +func TestBuildSSHArgsSetsServerAliveInterval(t *testing.T) { + args := buildSSHArgs("user", "/key", "proxy command", "myhost", "", ClientOptions{}) + + // ssh stops parsing options at the destination, so an option placed after the host would be + // treated as part of the remote command rather than as an ssh option. + optIdx := slices.Index(args, "ServerAliveInterval="+strconv.Itoa(sshconfig.ServerAliveIntervalSeconds)) + require.NotEqual(t, -1, optIdx, "ssh must be asked to send keepalives") + require.Equal(t, "-o", args[optIdx-1]) + assert.Less(t, optIdx, slices.Index(args, "myhost"), "the option must precede the destination host") +} + func TestBuildSSHArgsPTYPlacement(t *testing.T) { indexOf := func(args []string, want string) int { for i, a := range args { diff --git a/experimental/ssh/internal/server/sshd.go b/experimental/ssh/internal/server/sshd.go index bfafbbe5212..01fd0e3cdb0 100644 --- a/experimental/ssh/internal/server/sshd.go +++ b/experimental/ssh/internal/server/sshd.go @@ -9,6 +9,7 @@ import ( "os/exec" "path" "path/filepath" + "strconv" "strings" "github.com/databricks/cli/experimental/ssh/internal/keys" @@ -17,6 +18,18 @@ import ( "github.com/databricks/databricks-sdk-go" ) +// clientAliveIntervalSeconds is how often sshd asks the client to confirm it is still there. It +// drives the keepalive from the server end of the tunnel; sshconfig.ServerAliveIntervalSeconds +// documents why an SSH keepalive is what keeps the leg past the driver proxy from being reaped. +// Configuring it here covers clients the CLI does not configure — a hand-written ProxyCommand host +// block, or an IDE that supplies its own ssh options — where nothing sets ServerAliveInterval. +// The two intervals are deliberately independent: neither package imports the other, and each end +// keeps its own leg warm, so they need not track a single shared value. +// +// It also brings in ClientAliveCountMax (OpenSSH default 3), so sshd reclaims a session whose +// client has gone away after ~90s instead of holding it open until the server's shutdown delay. +const clientAliveIntervalSeconds = 30 + func prepareSSHDConfig(ctx context.Context, client *databricks.WorkspaceClient, opts ServerOptions) (string, error) { clientPublicKey, err := keys.GetSecret(ctx, client, opts.SecretScopeName, opts.AuthorizedKeySecretName) if err != nil { @@ -76,15 +89,7 @@ func prepareSSHDConfig(ctx context.Context, client *databricks.WorkspaceClient, } setEnv := setEnvBuf.String() - sshdConfigContent := "PubkeyAuthentication yes\n" + - "PasswordAuthentication no\n" + - "ChallengeResponseAuthentication no\n" + - "Subsystem sftp internal-sftp\n" + - "HostKey " + keyPath + "\n" + - "AuthorizedKeysFile " + authKeysPath + "\n" + - setEnv + "\n" - - if err := os.WriteFile(sshdConfig, []byte(sshdConfigContent), 0o600); err != nil { + if err := os.WriteFile(sshdConfig, []byte(sshdConfigContent(keyPath, authKeysPath, setEnv)), 0o600); err != nil { return "", err } @@ -97,6 +102,18 @@ func prepareSSHDConfig(ctx context.Context, client *databricks.WorkspaceClient, return sshdConfig, nil } +// sshdConfigContent assembles the configuration the tunnel's sshd runs with. +func sshdConfigContent(hostKeyPath, authorizedKeysPath, setEnv string) string { + return "PubkeyAuthentication yes\n" + + "PasswordAuthentication no\n" + + "ChallengeResponseAuthentication no\n" + + "ClientAliveInterval " + strconv.Itoa(clientAliveIntervalSeconds) + "\n" + + "Subsystem sftp internal-sftp\n" + + "HostKey " + hostKeyPath + "\n" + + "AuthorizedKeysFile " + authorizedKeysPath + "\n" + + setEnv + "\n" +} + func createSSHDProcess(ctx context.Context, configPath string) *exec.Cmd { return exec.CommandContext(ctx, "/usr/sbin/sshd", "-f", configPath, "-i") } diff --git a/experimental/ssh/internal/server/sshd_test.go b/experimental/ssh/internal/server/sshd_test.go index a453d987a00..4887f9895d8 100644 --- a/experimental/ssh/internal/server/sshd_test.go +++ b/experimental/ssh/internal/server/sshd_test.go @@ -1,6 +1,7 @@ package server import ( + "strconv" "testing" "github.com/stretchr/testify/assert" @@ -71,3 +72,18 @@ func TestEscapeEnvValue(t *testing.T) { }) } } + +func TestSSHDConfigSetsClientAliveInterval(t *testing.T) { + config := sshdConfigContent("/keys/server-private-key", "/keys/authorized_keys", `SetEnv FOO="bar"`) + + // This is the half of the keepalive that reaches clients the CLI never configures, so sshd has + // to drive it: without ClientAliveInterval sshd sends nothing on an idle session. + assert.Contains(t, config, "\nClientAliveInterval "+strconv.Itoa(clientAliveIntervalSeconds)+"\n") + + // The interval has to fire well inside the ~8 minute reap window, and ClientAliveCountMax + // (OpenSSH default 3) intervals have to outlast the longest legitimate pause on a healthy + // tunnel — the up to 30s a handover can hold the sending loop. The 30 below is + // proxy.proxyHandoverInitTimeout's current value; it is unexported, so it can't be referenced. + assert.Less(t, clientAliveIntervalSeconds, 8*60) + assert.Greater(t, 3*clientAliveIntervalSeconds, 30) +} diff --git a/experimental/ssh/internal/sshconfig/sshconfig.go b/experimental/ssh/internal/sshconfig/sshconfig.go index ad8ca0ee2a7..fbbf3b42fee 100644 --- a/experimental/ssh/internal/sshconfig/sshconfig.go +++ b/experimental/ssh/internal/sshconfig/sshconfig.go @@ -19,6 +19,18 @@ const ( configDirName = ".databricks/ssh-tunnel-configs" ) +// ServerAliveIntervalSeconds is how often the ssh client asks the SSH server to confirm it is +// still there. The reply is a real SSH packet, so the keepalive puts payload bytes on every hop +// of the tunnel — and payload is what an idle session needs. The driver proxy terminates +// websocket control frames itself, so the proxy's own websocket ping never becomes payload on +// the leg past it, and that leg is reaped after ~8 minutes without any. +// +// It also brings in ServerAliveCountMax (OpenSSH default 3), so a tunnel that stops responding +// is torn down after ~90s with ssh's own "server not responding" message instead of hanging. +// That is well clear of the up to 30s a handover can hold the sending loop +// (proxyHandoverInitTimeout), the longest legitimate pause on a healthy tunnel. +const ServerAliveIntervalSeconds = 30 + func GetConfigDir(ctx context.Context) (string, error) { homeDir, err := env.UserHomeDir(ctx) if err != nil { @@ -206,9 +218,10 @@ func GenerateHostConfig(hostName, userName, identityFile, proxyCommand string) s Host %s User %s ConnectTimeout 360 + ServerAliveInterval %d StrictHostKeyChecking accept-new IdentitiesOnly yes IdentityFile %q ProxyCommand %s -`, hostName, userName, identityFile, proxyCommand) +`, hostName, userName, ServerAliveIntervalSeconds, identityFile, proxyCommand) } diff --git a/experimental/ssh/internal/sshconfig/sshconfig_test.go b/experimental/ssh/internal/sshconfig/sshconfig_test.go index 6c453910cdc..23abcbc34d7 100644 --- a/experimental/ssh/internal/sshconfig/sshconfig_test.go +++ b/experimental/ssh/internal/sshconfig/sshconfig_test.go @@ -1,6 +1,7 @@ package sshconfig import ( + "fmt" "os" "path/filepath" "testing" @@ -10,6 +11,21 @@ import ( "github.com/stretchr/testify/require" ) +func TestGenerateHostConfigSetsServerAliveInterval(t *testing.T) { + config := GenerateHostConfig("myhost", "root", "/keys/myhost", "databricks ssh connect --proxy") + + // `ssh setup` and `--ide` reach ssh through this block and nothing else, so the option has to + // be in it. + assert.Contains(t, config, fmt.Sprintf("\n ServerAliveInterval %d\n", ServerAliveIntervalSeconds)) + + // The interval has to fire well inside the ~8 minute reap window, and ServerAliveCountMax + // (OpenSSH default 3) intervals have to outlast the longest legitimate pause on a healthy + // tunnel — the up to 30s a handover can hold the sending loop. The 30 below is + // proxy.proxyHandoverInitTimeout's current value; it is unexported, so it can't be referenced. + assert.Less(t, ServerAliveIntervalSeconds, 8*60) + assert.Greater(t, 3*ServerAliveIntervalSeconds, 30) +} + func TestGetConfigDir(t *testing.T) { dir, err := GetConfigDir(t.Context()) assert.NoError(t, err)