Split and move ssh logic to experimental/ssh

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

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 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 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 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 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)
 		}