Split and move ssh logic to experimental/ssh (#3544)

## Changes

Previous PRs in the stack:
- #3471
- #3475

Move from flat structure to cmd and internal
client/server/proxy/keys/workspace sub packages.

Moved a bunch of hardcoded values to constants (some hardcoded stuff
remains in the `keys` and `workspace` sub packages, will be improved in
a follow up)

Added unit tests for servers connection manager logic and updated proxy
tests.
The coverage isn't great yet, the goal is to avoid unit tests with
excessive mocking and focus on integration tests from now on.

## Why
`/experimental` already has the databricks-bundles python project, seems
like a good place to put ssh stuff too

## Tests
Unit and manual

Author: ilia-db

Makefile

@@ -1,6 +1,6 @@
 default: checks fmt lint
 
-PACKAGES=./acceptance/... ./libs/... ./internal/... ./cmd/... ./bundle/... .
+PACKAGES=./acceptance/... ./libs/... ./internal/... ./cmd/... ./bundle/... ./experimental/ssh/... .
 
 GOTESTSUM_FORMAT ?= pkgname-and-test-fails
 GOTESTSUM_CMD ?= go tool gotestsum --format ${GOTESTSUM_FORMAT} --no-summary=skipped --jsonfile test-output.json
@@ -136,4 +136,4 @@ generate:
 	$(GENKIT_BINARY) update-sdk
 
 
-.PHONY: lint lintfull tidy lintcheck fmt fmtfull test cover showcover build snapshot schema integration integration-short acc-cover acc-showcover docs ws links checks test-update test-update-aws test-update-all generate-validation
+.PHONY: lint lintfull tidy lintcheck fmt fmtfull test cover showcover build snapshot snapshot-release schema integration integration-short acc-cover acc-showcover docs ws links checks test-update test-update-aws test-update-all generate-validation

cmd/cmd.go

@@ -5,7 +5,7 @@ import (
 	"strings"
 
 	"github.com/databricks/cli/cmd/psql"
-	"github.com/databricks/cli/cmd/ssh"
+	ssh "github.com/databricks/cli/experimental/ssh/cmd"
 
 	"github.com/databricks/cli/cmd/account"
 	"github.com/databricks/cli/cmd/api"

cmd/ssh/README.md experimental/ssh/README.mdcmd/ssh/README.md renamed to experimental/ssh/README.md

@@ -4,18 +4,12 @@ import (
 	"time"
 
 	"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/databricks/cli/libs/ssh"
 	"github.com/spf13/cobra"
 )
 
-const (
-	defaultClientPublicKeyName = "client-public-key"
-	defaultShutdownDelay       = 10 * time.Minute
-	defaultHandoverTimeout     = 30 * time.Minute
-	defaultMaxClients          = 10
-)
-
 func newConnectCommand() *cobra.Command {
 	cmd := &cobra.Command{
 		Use:   "connect",
@@ -54,8 +48,8 @@ the SSH server and handling the connection proxy.
 	cmd.PreRunE = root.MustWorkspaceClient
 	cmd.RunE = func(cmd *cobra.Command, args []string) error {
 		ctx := cmd.Context()
-		client := cmdctx.WorkspaceClient(ctx)
-		opts := ssh.ClientOptions{
+		wsClient := cmdctx.WorkspaceClient(ctx)
+		opts := client.ClientOptions{
 			ClusterID:           clusterID,
 			ProxyMode:           proxyMode,
 			ServerMetadata:      serverMetadata,
@@ -65,8 +59,13 @@ the SSH server and handling the connection proxy.
 			ReleasesDir:         releasesDir,
 			AdditionalArgs:      args,
 			ClientPublicKeyName: defaultClientPublicKeyName,
+			ServerTimeout:       serverTimeout,
+		}
+		err := client.RunClient(ctx, wsClient, opts)
+		if err != nil && proxy.IsNormalClosure(err) {
+			return nil
 		}
-		return ssh.RunClient(ctx, client, opts)
+		return err
 	}
 
 	return cmd

cmd/ssh/connect.go experimental/ssh/cmd/connect.gocmd/ssh/connect.go renamed to experimental/ssh/cmd/connect.go

@@ -0,0 +1,17 @@
+package ssh
+
+import "time"
+
+const (
+	defaultServerPort      = 7772
+	defaultMaxClients      = 10
+	defaultShutdownDelay   = 10 * time.Minute
+	defaultHandoverTimeout = 30 * time.Minute
+
+	serverTimeout              = 24 * time.Hour
+	serverPortRange            = 100
+	serverConfigDir            = ".ssh-tunnel"
+	serverPrivateKeyName       = "server-private-key"
+	serverPublicKeyName        = "server-public-key"
+	defaultClientPublicKeyName = "client-public-key"
+)

experimental/ssh/cmd/constants.go

@@ -4,8 +4,8 @@ import (
 	"time"
 
 	"github.com/databricks/cli/cmd/root"
+	"github.com/databricks/cli/experimental/ssh/internal/server"
 	"github.com/databricks/cli/libs/cmdctx"
-	"github.com/databricks/cli/libs/ssh"
 	"github.com/spf13/cobra"
 )
 
@@ -27,33 +27,38 @@ and proxies them to local SSH daemon processes.
 	var shutdownDelay time.Duration
 	var clusterID string
 	var version string
+	var secretsScope string
+	var publicKeySecretName string
 
 	cmd.Flags().StringVar(&clusterID, "cluster", "", "Databricks cluster ID")
 	cmd.MarkFlagRequired("cluster")
-	cmd.Flags().IntVar(&maxClients, "max-clients", 10, "Maximum number of SSH clients")
-	cmd.Flags().DurationVar(&shutdownDelay, "shutdown-delay", 10*time.Minute, "Delay before shutting down after no pings from clients")
+	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().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().StringVar(&version, "version", "", "Client version of the Databricks CLI")
 
 	cmd.PreRunE = root.MustWorkspaceClient
 	cmd.RunE = func(cmd *cobra.Command, args []string) error {
 		ctx := cmd.Context()
 		client := cmdctx.WorkspaceClient(ctx)
-		opts := ssh.ServerOptions{
+		opts := server.ServerOptions{
 			ClusterID:            clusterID,
 			MaxClients:           maxClients,
 			ShutdownDelay:        shutdownDelay,
 			Version:              version,
-			ConfigDir:            ".ssh-tunnel",
-			ServerPrivateKeyName: "server-private-key",
-			ServerPublicKeyName:  "server-public-key",
-			DefaultPort:          7772,
-			PortRange:            100,
-		}
-		err := ssh.RunServer(ctx, client, opts)
-		if err != nil && ssh.IsNormalClosure(err) {
-			return nil
+			ConfigDir:            serverConfigDir,
+			SecretsScope:         secretsScope,
+			ClientPublicKeyName:  publicKeySecretName,
+			ServerPrivateKeyName: serverPrivateKeyName,
+			ServerPublicKeyName:  serverPublicKeyName,
+			DefaultPort:          defaultServerPort,
+			PortRange:            serverPortRange,
 		}
-		return err
+		return server.RunServer(ctx, client, opts)
 	}
 
 	return cmd

cmd/ssh/server.go experimental/ssh/cmd/server.gocmd/ssh/server.go renamed to experimental/ssh/cmd/server.go

@@ -4,8 +4,8 @@ import (
 	"time"
 
 	"github.com/databricks/cli/cmd/root"
+	"github.com/databricks/cli/experimental/ssh/internal/setup"
 	"github.com/databricks/cli/libs/cmdctx"
-	"github.com/databricks/cli/libs/ssh"
 	"github.com/spf13/cobra"
 )
 
@@ -31,20 +31,20 @@ an SSH host configuration to your SSH config file.
 	cmd.Flags().StringVar(&clusterID, "cluster", "", "Databricks cluster ID")
 	cmd.MarkFlagRequired("cluster")
 	cmd.Flags().StringVar(&sshConfigPath, "ssh-config", "", "Path to SSH config file (default ~/.ssh/config)")
-	cmd.Flags().DurationVar(&shutdownDelay, "shutdown-delay", 10*time.Minute, "SSH server will terminate after this delay if there are no active connections")
+	cmd.Flags().DurationVar(&shutdownDelay, "shutdown-delay", defaultShutdownDelay, "SSH server will terminate after this delay if there are no active connections")
 
 	cmd.PreRunE = root.MustWorkspaceClient
 	cmd.RunE = func(cmd *cobra.Command, args []string) error {
 		ctx := cmd.Context()
 		client := cmdctx.WorkspaceClient(ctx)
-		opts := ssh.SetupOptions{
+		opts := setup.SetupOptions{
 			HostName:      hostName,
 			ClusterID:     clusterID,
 			SSHConfigPath: sshConfigPath,
 			ShutdownDelay: shutdownDelay,
 			Profile:       client.Config.Profile,
 		}
-		return ssh.Setup(ctx, client, opts)
+		return setup.Setup(ctx, client, opts)
 	}
 
 	return cmd

cmd/ssh/setup.go experimental/ssh/cmd/setup.gocmd/ssh/setup.go renamed to experimental/ssh/cmd/setup.go

@@ -1,10 +1,9 @@
-package ssh
+package client
 
 import (
 	"context"
 	_ "embed"
 	"encoding/base64"
-	"encoding/json"
 	"errors"
 	"fmt"
 	"io"
@@ -18,6 +17,9 @@ import (
 	"syscall"
 	"time"
 
+	"github.com/databricks/cli/experimental/ssh/internal/keys"
+	"github.com/databricks/cli/experimental/ssh/internal/proxy"
+	sshWorkspace "github.com/databricks/cli/experimental/ssh/internal/workspace"
 	"github.com/databricks/cli/internal/build"
 	"github.com/databricks/cli/libs/cmdio"
 	"github.com/databricks/databricks-sdk-go"
@@ -27,17 +29,11 @@ import (
 	"golang.org/x/sync/errgroup"
 )
 
-type PortMetadata struct {
-	Port int `json:"port"`
-}
-
 //go:embed ssh-server-bootstrap.py
 var sshServerBootstrapScript string
 
 var errServerMetadata = errors.New("server metadata error")
 
-const serverJobTimeoutSeconds = 24 * 60 * 60
-
 type ClientOptions struct {
 	// Id of the cluster to connect to
 	ClusterID string
@@ -54,6 +50,8 @@ type ClientOptions struct {
 	ServerMetadata string
 	// How often the CLI should reconnect to the server with new auth.
 	HandoverTimeout time.Duration
+	// Max amount of time the server process is allowed to live
+	ServerTimeout time.Duration
 	// Directory for local SSH tunnel development releases.
 	// If not present, the CLI will use github releases with the current version.
 	ReleasesDir string
@@ -77,16 +75,17 @@ func RunClient(ctx context.Context, client *databricks.WorkspaceClient, opts Cli
 		cancel()
 	}()
 
-	keyPath, err := getLocalSSHKeyPath(opts.ClusterID, opts.SSHKeysDir)
+	keyPath, err := keys.GetLocalSSHKeyPath(opts.ClusterID, opts.SSHKeysDir)
 	if err != nil {
 		return fmt.Errorf("failed to get local keys folder: %w", err)
 	}
-	privateKeyPath, publicKey, err := checkAndGenerateSSHKeyPair(ctx, keyPath)
+	privateKeyPath, publicKey, err := keys.CheckAndGenerateSSHKeyPair(ctx, keyPath)
 	if err != nil {
 		return fmt.Errorf("failed to check or generate SSH key pair: %w", err)
 	}
+	cmdio.LogString(ctx, "Using SSH key: "+privateKeyPath)
 
-	secretsScopeName, err := putSecretInScope(ctx, client, opts.ClusterID, opts.ClientPublicKeyName, publicKey)
+	secretsScopeName, 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)
 	}
@@ -99,10 +98,10 @@ func RunClient(ctx context.Context, client *databricks.WorkspaceClient, opts Cli
 
 	if opts.ServerMetadata == "" {
 		cmdio.LogString(ctx, "Checking for ssh-tunnel binaries to upload...")
-		if err := uploadTunnelBinaries(ctx, client, version, opts.ReleasesDir); err != nil {
+		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)
+		userName, serverPort, err = ensureSSHServerIsRunning(ctx, client, opts.ClusterID, secretsScopeName, 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)
 		}
@@ -132,36 +131,8 @@ func RunClient(ctx context.Context, client *databricks.WorkspaceClient, opts Cli
 	}
 }
 
-func getWorkspaceMetadata(ctx context.Context, client *databricks.WorkspaceClient, version, clusterID string) (int, error) {
-	contentDir, err := getWorkspaceContentDir(ctx, client, version, clusterID)
-	if err != nil {
-		return 0, fmt.Errorf("failed to get workspace content directory: %w", err)
-	}
-
-	metadataPath := filepath.ToSlash(filepath.Join(contentDir, "metadata.json"))
-
-	content, err := client.Workspace.Download(ctx, metadataPath)
-	if err != nil {
-		return 0, fmt.Errorf("failed to download metadata file: %w", err)
-	}
-	defer content.Close()
-
-	metadataBytes, err := io.ReadAll(content)
-	if err != nil {
-		return 0, fmt.Errorf("failed to read metadata content: %w", err)
-	}
-
-	var metadata PortMetadata
-	err = json.Unmarshal(metadataBytes, &metadata)
-	if err != nil {
-		return 0, fmt.Errorf("failed to parse metadata JSON: %w", err)
-	}
-
-	return metadata.Port, nil
-}
-
 func getServerMetadata(ctx context.Context, client *databricks.WorkspaceClient, clusterID, version string) (int, string, error) {
-	serverPort, err := getWorkspaceMetadata(ctx, client, version, clusterID)
+	serverPort, err := sshWorkspace.GetWorkspaceMetadata(ctx, client, version, clusterID)
 	if err != nil {
 		return 0, "", errors.Join(errServerMetadata, err)
 	}
@@ -194,8 +165,8 @@ 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) (int64, error) {
-	contentDir, err := getWorkspaceContentDir(ctx, client, version, clusterID)
+func submitSSHTunnelJob(ctx context.Context, client *databricks.WorkspaceClient, clusterID, secretsScope, 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)
 	}
@@ -223,7 +194,7 @@ func submitSSHTunnelJob(ctx context.Context, client *databricks.WorkspaceClient,
 
 	submitRun := jobs.SubmitRun{
 		RunName:        sshTunnelJobName,
-		TimeoutSeconds: serverJobTimeoutSeconds,
+		TimeoutSeconds: int(serverTimeout.Seconds()),
 		Tasks: []jobs.SubmitTask{
 			{
 				TaskKey: "start_ssh_server",
@@ -237,7 +208,7 @@ func submitSSHTunnelJob(ctx context.Context, client *databricks.WorkspaceClient,
 						"maxClients":          strconv.Itoa(maxClients),
 					},
 				},
-				TimeoutSeconds:    serverJobTimeoutSeconds,
+				TimeoutSeconds:    int(serverTimeout.Seconds()),
 				ExistingClusterId: clusterID,
 			},
 		},
@@ -286,13 +257,13 @@ func startSSHProxy(ctx context.Context, client *databricks.WorkspaceClient, clus
 	g, gCtx := errgroup.WithContext(ctx)
 
 	cmdio.LogString(ctx, "Establishing SSH proxy connection...")
-	proxy := newProxyConnection(func(ctx context.Context, connID string) (*websocket.Conn, error) {
+	conn := proxy.NewProxyConnection(func(ctx context.Context, connID string) (*websocket.Conn, error) {
 		return createWebsocketConnection(ctx, client, connID, clusterID, serverPort)
 	})
-	if err := proxy.Connect(gCtx); err != nil {
+	if err := conn.Connect(gCtx); err != nil {
 		return fmt.Errorf("failed to connect to proxy: %w", err)
 	}
-	defer proxy.Close()
+	defer conn.Close()
 	cmdio.LogString(ctx, "SSH proxy connection established")
 
 	cmdio.LogString(ctx, fmt.Sprintf("Connection handover timeout: %v", handoverTimeout))
@@ -305,7 +276,7 @@ func startSSHProxy(ctx context.Context, client *databricks.WorkspaceClient, clus
 			case <-gCtx.Done():
 				return gCtx.Err()
 			case <-handoverTicker.C:
-				err := proxy.InitiateHandover(gCtx)
+				err := conn.InitiateHandover(gCtx)
 				if err != nil {
 					return err
 				}
@@ -314,13 +285,13 @@ func startSSHProxy(ctx context.Context, client *databricks.WorkspaceClient, clus
 	})
 
 	g.Go(func() error {
-		return proxy.Start(gCtx, os.Stdin, os.Stdout)
+		return conn.Start(gCtx, os.Stdin, os.Stdout)
 	})
 
 	return g.Wait()
 }
 
-func ensureSSHServerIsRunning(ctx context.Context, client *databricks.WorkspaceClient, clusterID, secretsScope, publicKeySecretName, version string, shutdownDelay time.Duration, maxClients int) (string, int, error) {
+func ensureSSHServerIsRunning(ctx context.Context, client *databricks.WorkspaceClient, clusterID, secretsScope, 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 {
@@ -331,7 +302,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)
+		runID, err := submitSSHTunnelJob(ctx, client, clusterID, secretsScope, publicKeySecretName, version, shutdownDelay, maxClients, serverTimeout)
 		if err != nil {
 			return "", 0, fmt.Errorf("failed to submit ssh server job: %w", err)
 		}