Move pure ssh client and server logic to proxy package and test it

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

@@ -44,7 +44,7 @@ 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,
@@ -58,7 +58,7 @@ and proxies them to local SSH daemon processes.
 			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()
 
@@ -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)
@@ -253,42 +252,11 @@ 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) {

experimental/ssh/internal/proxy/client.go

@@ -0,0 +1,46 @@
+package proxy
+
+import (
+	"context"
+	"fmt"
+	"io"
+	"time"
+
+	"github.com/databricks/cli/libs/cmdio"
+	"golang.org/x/sync/errgroup"
+)
+
+func RunClientProxy(ctx context.Context, src io.Reader, dst io.Writer, handoverTimeout time.Duration, createConn createWebsocketConnectionFunc) error {
+	proxy := newProxyConnection(createConn)
+	cmdio.LogString(ctx, "Establishing SSH proxy connection...")
+	g, gCtx := errgroup.WithContext(ctx)
+	if err := proxy.connect(gCtx); err != nil {
+		return fmt.Errorf("failed to connect to proxy: %w", err)
+	}
+	defer proxy.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 := proxy.initiateHandover(gCtx)
+				if err != nil {
+					return err
+				}
+			}
+		}
+	})
+
+	g.Go(func() error {
+		return proxy.start(gCtx, src, dst)
+	})
+
+	return g.Wait()
+}