Move pure ssh client and server logic to proxy package and test it (#3569)

## Changes
Based on #3544

Move pure client and server logic to the proxy package and test it. 

Now `proxyConnection` is completely private, and `proxy` package exposes
only 2 public functions:
- `RunClientProxy` function
- `NewProxyServer` constructor

Meanwhile the internal logic doesn't depend on any ssh stuff and can be
unit tested.

`internal/client` and `internal/server` packages are now much smaller
and contain mainly ssh/jobs/cluster glue.

Also adding `bench.sh` until we have proper e2e tests.

## Why
Tests

## Tests
Unit and manual tests

<!-- If your PR needs to be included in the release notes for next
release,
add a separate entry in NEXT_CHANGELOG.md as part of your PR. -->

Author: ilia-db

experimental/ssh/bench.sh

@@ -0,0 +1,79 @@
+#!/bin/bash
+
+# SSH Tunnel Performance Test Script
+# Usage:
+# 1. Setup ssh config: ./cli ssh setup --cluster --name
+# 2. Run: ./experimental/ssh/bench.sh <cluster-id> <ssh-config-hostname> [ssh-tunnel-binary-path] [profile]
+
+set -e
+
+CLUSTER_ID="$1"
+HOSTNAME="$2"
+CLI=${3:-./cli}
+PROFILE="${4:-DEFAULT}"
+
+TEST_SIZES=(300 600)  # MB
+
+if [ -z "$CLUSTER_ID" ] || [ -z "$HOSTNAME" ]; then
+    echo "Usage: $0 <cluster-id> <hostname> [ssh-tunnel-binary-path] [profile]"
+    exit 1
+fi
+
+echo "=== SSH Tunnel Performance Test ==="
+echo "Cluster ID: $CLUSTER_ID"
+echo "Hostname: $HOSTNAME"
+echo "Profile: $PROFILE"
+echo "Start time: $(date)"
+echo "SSH Tunnel: $CLI"
+echo
+
+# Test basic connectivity
+echo "🔍 Testing basic connectivity..."
+
+start_time=$(date +%s.%N)
+if ! $CLI ssh connect --cluster="$CLUSTER_ID" --profile="$PROFILE" --releases-dir=./dist -- "echo 'Connection successful'"; then
+    echo "❌ Failed to connect to the cluster"
+    exit 1
+fi
+end_time=$(date +%s.%N)
+duration=$(echo "$end_time - $start_time" | bc)
+duration_ms=$(echo "$duration * 1000" | bc)
+echo "✅ Basic connectivity OK ($duration_ms ms)"
+echo
+
+# Throughput Tests
+echo "⚡ Testing Throughput..."
+
+# Create test files
+echo "Creating test files..."
+for size in "${TEST_SIZES[@]}"; do
+    if [ ! -f "test_${size}mb.dat" ]; then
+        dd if=/dev/zero of="test_${size}mb.dat" bs=1M count=$size 2>/dev/null
+        echo "  Created test_${size}mb.dat"
+    fi
+done
+echo
+
+# Upload tests
+echo "📤 Upload Speed Tests:"
+for size in "${TEST_SIZES[@]}"; do
+    echo -n "  ${size}MB file: "
+    scp "test_${size}mb.dat" "$HOSTNAME:/tmp/test_upload_${size}mb.dat"
+done
+echo
+
+# Download tests
+echo "📥 Download Speed Tests:"
+for size in "${TEST_SIZES[@]}"; do
+    echo -n "  ${size}MB file: "
+    scp "$HOSTNAME:/tmp/test_upload_${size}mb.dat" "./test_download_${size}mb.dat"
+done
+echo
+
+# Cleanup
+echo "🧹 Cleaning up..."
+for size in "${TEST_SIZES[@]}"; do
+    rm -f "test_${size}mb.dat" "test_download_${size}mb.dat"
+    $CLI ssh connect --cluster="$CLUSTER_ID" --profile="$PROFILE" -- "rm -f /tmp/test_upload_${size}mb.dat" 2>/dev/null || true
+done
+echo

experimental/ssh/cmd/connect.go

@@ -5,7 +5,6 @@ import (
 
 	"github.com/databricks/cli/cmd/root"
 	"github.com/databricks/cli/experimental/ssh/internal/client"
-	"github.com/databricks/cli/experimental/ssh/internal/proxy"
 	"github.com/databricks/cli/libs/cmdctx"
 	"github.com/spf13/cobra"
 )
@@ -61,11 +60,7 @@ the SSH server and handling the connection proxy.
 			ClientPublicKeyName: defaultClientPublicKeyName,
 			ServerTimeout:       serverTimeout,
 		}
-		err := client.RunClient(ctx, wsClient, opts)
-		if err != nil && proxy.IsNormalClosure(err) {
-			return nil
-		}
-		return err
+		return client.Run(ctx, wsClient, opts)
 	}
 
 	return cmd

experimental/ssh/cmd/server.go

@@ -27,15 +27,15 @@ and proxies them to local SSH daemon processes.
 	var shutdownDelay time.Duration
 	var clusterID string
 	var version string
-	var secretsScope string
-	var publicKeySecretName string
+	var keysSecretScopeName string
+	var authorizedKeyName string
 
 	cmd.Flags().StringVar(&clusterID, "cluster", "", "Databricks cluster ID")
 	cmd.MarkFlagRequired("cluster")
-	cmd.Flags().StringVar(&secretsScope, "secrets-scope-name", "", "Databricks secrets scope name")
-	cmd.MarkFlagRequired("secrets-scope-name")
-	cmd.Flags().StringVar(&publicKeySecretName, "client-key-name", "", "Databricks client key name")
-	cmd.MarkFlagRequired("client-key-name")
+	cmd.Flags().StringVar(&keysSecretScopeName, "keys-secret-scope-name", "", "Databricks secret scope name to store SSH keys")
+	cmd.MarkFlagRequired("keys-secret-scope-name")
+	cmd.Flags().StringVar(&authorizedKeyName, "authorized-key-secret-name", "", "Authorized key secret name in the secret scope")
+	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")
@@ -44,21 +44,21 @@ and proxies them to local SSH daemon processes.
 	cmd.PreRunE = root.MustWorkspaceClient
 	cmd.RunE = func(cmd *cobra.Command, args []string) error {
 		ctx := cmd.Context()
-		client := cmdctx.WorkspaceClient(ctx)
+		wsc := cmdctx.WorkspaceClient(ctx)
 		opts := server.ServerOptions{
 			ClusterID:            clusterID,
 			MaxClients:           maxClients,
 			ShutdownDelay:        shutdownDelay,
 			Version:              version,
 			ConfigDir:            serverConfigDir,
-			SecretsScope:         secretsScope,
-			ClientPublicKeyName:  publicKeySecretName,
+			KeysSecretScopeName:  keysSecretScopeName,
+			AuthorizedKeyName:    authorizedKeyName,
 			ServerPrivateKeyName: serverPrivateKeyName,
 			ServerPublicKeyName:  serverPublicKeyName,
 			DefaultPort:          defaultServerPort,
 			PortRange:            serverPortRange,
 		}
-		return server.RunServer(ctx, client, opts)
+		return server.Run(ctx, wsc, opts)
 	}
 
 	return cmd

experimental/ssh/internal/client/client.go

@@ -26,7 +26,6 @@ import (
 	"github.com/databricks/databricks-sdk-go/service/jobs"
 	"github.com/databricks/databricks-sdk-go/service/workspace"
 	"github.com/gorilla/websocket"
-	"golang.org/x/sync/errgroup"
 )
 
 //go:embed ssh-server-bootstrap.py
@@ -63,7 +62,7 @@ type ClientOptions struct {
 	AdditionalArgs []string
 }
 
-func RunClient(ctx context.Context, client *databricks.WorkspaceClient, opts ClientOptions) error {
+func Run(ctx context.Context, client *databricks.WorkspaceClient, opts ClientOptions) error {
 	ctx, cancel := context.WithCancel(ctx)
 	defer cancel()
 
@@ -85,11 +84,11 @@ func RunClient(ctx context.Context, client *databricks.WorkspaceClient, opts Cli
 	}
 	cmdio.LogString(ctx, "Using SSH key: "+privateKeyPath)
 
-	secretsScopeName, err := keys.PutSecretInScope(ctx, client, opts.ClusterID, opts.ClientPublicKeyName, publicKey)
+	keysSecretScopeName, err := keys.PutSecretInScope(ctx, client, opts.ClusterID, opts.ClientPublicKeyName, publicKey)
 	if err != nil {
 		return fmt.Errorf("failed to store public key in secret scope: %w", err)
 	}
-	cmdio.LogString(ctx, fmt.Sprintf("Secrets scope: %s, key name: %s", secretsScopeName, opts.ClientPublicKeyName))
+	cmdio.LogString(ctx, fmt.Sprintf("Secrets scope: %s, key name: %s", keysSecretScopeName, opts.ClientPublicKeyName))
 
 	var userName string
 	var serverPort int
@@ -101,7 +100,7 @@ func RunClient(ctx context.Context, client *databricks.WorkspaceClient, opts Cli
 		if err := UploadTunnelReleases(ctx, client, version, opts.ReleasesDir); err != nil {
 			return fmt.Errorf("failed to upload ssh-tunnel binaries: %w", err)
 		}
-		userName, serverPort, err = ensureSSHServerIsRunning(ctx, client, opts.ClusterID, secretsScopeName, opts.ClientPublicKeyName, version, opts.ShutdownDelay, opts.MaxClients, opts.ServerTimeout)
+		userName, serverPort, err = ensureSSHServerIsRunning(ctx, client, opts.ClusterID, keysSecretScopeName, opts.ClientPublicKeyName, version, opts.ShutdownDelay, opts.MaxClients, opts.ServerTimeout)
 		if err != nil {
 			return fmt.Errorf("failed to ensure that ssh server is running: %w", err)
 		}
@@ -124,7 +123,7 @@ func RunClient(ctx context.Context, client *databricks.WorkspaceClient, opts Cli
 	cmdio.LogString(ctx, fmt.Sprintf("Server port: %d", serverPort))
 
 	if opts.ProxyMode {
-		return startSSHProxy(ctx, client, opts.ClusterID, serverPort, opts.HandoverTimeout)
+		return runSSHProxy(ctx, client, opts.ClusterID, serverPort, opts.HandoverTimeout)
 	} else {
 		cmdio.LogString(ctx, fmt.Sprintf("Additional SSH arguments: %v", opts.AdditionalArgs))
 		return spawnSSHClient(ctx, opts.ClusterID, userName, privateKeyPath, serverPort, opts.HandoverTimeout, opts.AdditionalArgs)
@@ -165,7 +164,7 @@ func getServerMetadata(ctx context.Context, client *databricks.WorkspaceClient,
 	return serverPort, string(bodyBytes), nil
 }
 
-func submitSSHTunnelJob(ctx context.Context, client *databricks.WorkspaceClient, clusterID, secretsScope, publicKeySecretName, version string, shutdownDelay time.Duration, maxClients int, serverTimeout time.Duration) (int64, error) {
+func submitSSHTunnelJob(ctx context.Context, client *databricks.WorkspaceClient, clusterID, keysSecretScopeName, publicKeySecretName, version string, shutdownDelay time.Duration, maxClients int, serverTimeout time.Duration) (int64, error) {
 	contentDir, err := sshWorkspace.GetWorkspaceContentDir(ctx, client, version, clusterID)
 	if err != nil {
 		return 0, fmt.Errorf("failed to get workspace content directory: %w", err)
@@ -201,11 +200,11 @@ func submitSSHTunnelJob(ctx context.Context, client *databricks.WorkspaceClient,
 				NotebookTask: &jobs.NotebookTask{
 					NotebookPath: jobNotebookPath,
 					BaseParameters: map[string]string{
-						"version":             version,
-						"secretsScope":        secretsScope,
-						"publicKeySecretName": publicKeySecretName,
-						"shutdownDelay":       shutdownDelay.String(),
-						"maxClients":          strconv.Itoa(maxClients),
+						"version":                 version,
+						"keysSecretScopeName":     keysSecretScopeName,
+						"authorizedKeySecretName": publicKeySecretName,
+						"shutdownDelay":           shutdownDelay.String(),
+						"maxClients":              strconv.Itoa(maxClients),
 					},
 				},
 				TimeoutSeconds:    int(serverTimeout.Seconds()),
@@ -253,45 +252,14 @@ func spawnSSHClient(ctx context.Context, clusterID, userName, privateKeyPath str
 	return sshCmd.Run()
 }
 
-func startSSHProxy(ctx context.Context, client *databricks.WorkspaceClient, clusterID string, serverPort int, handoverTimeout time.Duration) error {
-	g, gCtx := errgroup.WithContext(ctx)
-
-	cmdio.LogString(ctx, "Establishing SSH proxy connection...")
-	conn := proxy.NewProxyConnection(func(ctx context.Context, connID string) (*websocket.Conn, error) {
+func runSSHProxy(ctx context.Context, client *databricks.WorkspaceClient, clusterID string, serverPort int, handoverTimeout time.Duration) error {
+	createConn := func(ctx context.Context, connID string) (*websocket.Conn, error) {
 		return createWebsocketConnection(ctx, client, connID, clusterID, serverPort)
-	})
-	if err := conn.Connect(gCtx); err != nil {
-		return fmt.Errorf("failed to connect to proxy: %w", err)
 	}
-	defer conn.Close()
-	cmdio.LogString(ctx, "SSH proxy connection established")
-
-	cmdio.LogString(ctx, fmt.Sprintf("Connection handover timeout: %v", handoverTimeout))
-	handoverTicker := time.NewTicker(handoverTimeout)
-	defer handoverTicker.Stop()
-
-	g.Go(func() error {
-		for {
-			select {
-			case <-gCtx.Done():
-				return gCtx.Err()
-			case <-handoverTicker.C:
-				err := conn.InitiateHandover(gCtx)
-				if err != nil {
-					return err
-				}
-			}
-		}
-	})
-
-	g.Go(func() error {
-		return conn.Start(gCtx, os.Stdin, os.Stdout)
-	})
-
-	return g.Wait()
+	return proxy.RunClientProxy(ctx, os.Stdin, os.Stdout, handoverTimeout, createConn)
 }
 
-func ensureSSHServerIsRunning(ctx context.Context, client *databricks.WorkspaceClient, clusterID, secretsScope, publicKeySecretName, version string, shutdownDelay time.Duration, maxClients int, serverTimeout time.Duration) (string, int, error) {
+func ensureSSHServerIsRunning(ctx context.Context, client *databricks.WorkspaceClient, clusterID, keysSecretScopeName, publicKeySecretName, version string, shutdownDelay time.Duration, maxClients int, serverTimeout time.Duration) (string, int, error) {
 	cmdio.LogString(ctx, "Ensuring the cluster is running: "+clusterID)
 	err := client.Clusters.EnsureClusterIsRunning(ctx, clusterID)
 	if err != nil {
@@ -302,7 +270,7 @@ func ensureSSHServerIsRunning(ctx context.Context, client *databricks.WorkspaceC
 	if errors.Is(err, errServerMetadata) {
 		cmdio.LogString(ctx, "SSH server is not running, starting it now...")
 
-		runID, err := submitSSHTunnelJob(ctx, client, clusterID, secretsScope, publicKeySecretName, version, shutdownDelay, maxClients, serverTimeout)
+		runID, err := submitSSHTunnelJob(ctx, client, clusterID, keysSecretScopeName, publicKeySecretName, version, shutdownDelay, maxClients, serverTimeout)
 		if err != nil {
 			return "", 0, fmt.Errorf("failed to submit ssh server job: %w", err)
 		}

experimental/ssh/internal/client/ssh-server-bootstrap.py

@@ -13,8 +13,8 @@
 SSH_TUNNEL_BASENAME = "databricks_cli"
 
 dbutils.widgets.text("version", "")
-dbutils.widgets.text("secretsScope", "")
-dbutils.widgets.text("publicKeySecretName", "")
+dbutils.widgets.text("keysSecretScopeName", "")
+dbutils.widgets.text("authorizedKeySecretName", "")
 dbutils.widgets.text("maxClients", "10")
 dbutils.widgets.text("shutdownDelay", "10m")
 
@@ -86,13 +86,13 @@ def run_ssh_server():
     if os.environ.get("VIRTUAL_ENV") is None:
         os.environ["VIRTUAL_ENV"] = sys.executable
 
-    secrets_scope = dbutils.widgets.get("secretsScope")
+    secrets_scope = dbutils.widgets.get("keysSecretScopeName")
     if not secrets_scope:
-        raise RuntimeError("Secrets scope is required. Please provide it using the 'secretsScope' widget.")
+        raise RuntimeError("Secrets scope is required. Please provide it using the 'keysSecretScopeName' widget.")
 
-    public_key_secret_name = dbutils.widgets.get("publicKeySecretName")
+    public_key_secret_name = dbutils.widgets.get("authorizedKeySecretName")
     if not public_key_secret_name:
-        raise RuntimeError("Public key secret name is required. Please provide it using the 'publicKeySecretName' widget.")
+        raise RuntimeError("Public key secret name is required. Please provide it using the 'authorizedKeySecretName' widget.")
 
     version = dbutils.widgets.get("version")
     if not version:
@@ -123,8 +123,8 @@ def run_ssh_server():
                 "ssh",
                 "server",
                 f"--cluster={ctx.clusterId}",
-                f"--secrets-scope-name={secrets_scope}",
-                f"--client-key-name={public_key_secret_name}",
+                f"--keys-secret-scope-name={secrets_scope}",
+                f"--authorized-key-secret-name={public_key_secret_name}",
                 f"--max-clients={max_clients}",
                 f"--shutdown-delay={shutdown_delay}",
                 f"--version={version}",

experimental/ssh/internal/keys/keys.go

@@ -125,27 +125,27 @@ func CheckAndGenerateSSHKeyPair(ctx context.Context, keyPath string) (string, st
 	return keyPath, strings.TrimSpace(string(publicKeyBytes)), nil
 }
 
-func CheckAndGenerateSSHKeyPairFromSecrets(ctx context.Context, client *databricks.WorkspaceClient, clusterID, secretsScopeName, privateKeyName, publicKeyName string) ([]byte, []byte, error) {
-	privateKeyBytes, err := GetSecret(ctx, client, secretsScopeName, privateKeyName)
+func CheckAndGenerateSSHKeyPairFromSecrets(ctx context.Context, client *databricks.WorkspaceClient, clusterID, secretScopeName, privateKeyName, publicKeyName string) ([]byte, []byte, error) {
+	privateKeyBytes, err := GetSecret(ctx, client, secretScopeName, privateKeyName)
 	if err != nil {
 		privateKeyBytes, publicKeyBytes, err := generateSSHKeyPair()
 		if err != nil {
 			return nil, nil, fmt.Errorf("failed to generate SSH key pair: %w", err)
 		}
 
-		err = putSecret(ctx, client, secretsScopeName, privateKeyName, string(privateKeyBytes))
+		err = putSecret(ctx, client, secretScopeName, privateKeyName, string(privateKeyBytes))
 		if err != nil {
 			return nil, nil, err
 		}
 
-		err = putSecret(ctx, client, secretsScopeName, publicKeyName, string(publicKeyBytes))
+		err = putSecret(ctx, client, secretScopeName, publicKeyName, string(publicKeyBytes))
 		if err != nil {
 			return nil, nil, err
 		}
 
 		return privateKeyBytes, publicKeyBytes, nil
 	} else {
-		publicKeyBytes, err := GetSecret(ctx, client, secretsScopeName, publicKeyName)
+		publicKeyBytes, err := GetSecret(ctx, client, secretScopeName, publicKeyName)
 		if err != nil {
 			return nil, nil, fmt.Errorf("failed to get public key from secrets scope: %w", err)
 		}

experimental/ssh/internal/keys/secrets.go

@@ -10,19 +10,19 @@ import (
 	"github.com/databricks/databricks-sdk-go/service/workspace"
 )
 
-func createSecretsScope(ctx context.Context, client *databricks.WorkspaceClient, clusterID string) (string, error) {
+func createKeysSecretScope(ctx context.Context, client *databricks.WorkspaceClient, clusterID string) (string, error) {
 	me, err := client.CurrentUser.Me(ctx)
 	if err != nil {
 		return "", fmt.Errorf("failed to get current user: %w", err)
 	}
-	secretsScope := fmt.Sprintf("%s-%s-ssh-tunnel-keys", me.UserName, clusterID)
+	secretScopeName := fmt.Sprintf("%s-%s-ssh-tunnel-keys", me.UserName, clusterID)
 	err = client.Secrets.CreateScope(ctx, workspace.CreateScope{
-		Scope: secretsScope,
+		Scope: secretScopeName,
 	})
 	if err != nil && !errors.Is(err, databricks.ErrResourceAlreadyExists) {
 		return "", fmt.Errorf("failed to create secrets scope: %w", err)
 	}
-	return secretsScope, nil
+	return secretScopeName, nil
 }
 
 func GetSecret(ctx context.Context, client *databricks.WorkspaceClient, scope, key string) ([]byte, error) {
@@ -54,7 +54,7 @@ func putSecret(ctx context.Context, client *databricks.WorkspaceClient, scope, k
 }
 
 func PutSecretInScope(ctx context.Context, client *databricks.WorkspaceClient, clusterID, key, value string) (string, error) {
-	scopeName, err := createSecretsScope(ctx, client, clusterID)
+	scopeName, err := createKeysSecretScope(ctx, client, clusterID)
 	if err != nil {
 		return "", err
 	}