Pass public key name in every client ssh connection (#3775)

## Changes
Before we were passing a client key name (a secret name that contains
the actual client public key) to the ssh server job. After it's started,
the ssh server was only accepting one client with one specific key. This
is a problem when you try to ssh to the same cluster under the same
databricks user from multiple different client machines (that have
different ssh keys).

Our integration tests can't work consistently because of this, and we've
already got one customer report about this problem.

The main change here is that now we don't always generate key-pairs on
the client, but check if the pair already exists in the secret scope and
use it. And if not, we generate the pair and save both keys to the scope
(and before we were only saving public key there)

Author: ilia-db

experimental/ssh/cmd/connect.go

@@ -61,18 +61,19 @@ the SSH server and handling the connection proxy.
 		ctx := cmd.Context()
 		wsClient := cmdctx.WorkspaceClient(ctx)
 		opts := client.ClientOptions{
-			ClusterID:           clusterID,
-			ProxyMode:           proxyMode,
-			ServerMetadata:      serverMetadata,
-			ShutdownDelay:       shutdownDelay,
-			MaxClients:          maxClients,
-			HandoverTimeout:     handoverTimeout,
-			ReleasesDir:         releasesDir,
-			AdditionalArgs:      args,
-			ClientPublicKeyName: defaultClientPublicKeyName,
-			ServerTimeout:       serverTimeout,
-			AutoStartCluster:    autoStartCluster,
-			Profile:             wsClient.Config.Profile,
+			Profile:              wsClient.Config.Profile,
+			ClusterID:            clusterID,
+			ProxyMode:            proxyMode,
+			ServerMetadata:       serverMetadata,
+			ShutdownDelay:        shutdownDelay,
+			MaxClients:           maxClients,
+			HandoverTimeout:      handoverTimeout,
+			ReleasesDir:          releasesDir,
+			ServerTimeout:        serverTimeout,
+			AutoStartCluster:     autoStartCluster,
+			ClientPublicKeyName:  clientPublicKeyName,
+			ClientPrivateKeyName: clientPrivateKeyName,
+			AdditionalArgs:       args,
 		}
 		return client.Run(ctx, wsClient, opts)
 	}

experimental/ssh/cmd/constants.go

@@ -8,10 +8,11 @@ const (
 	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"
+	serverTimeout        = 24 * time.Hour
+	serverPortRange      = 100
+	serverConfigDir      = ".ssh-tunnel"
+	serverPrivateKeyName = "server-private-key"
+	serverPublicKeyName  = "server-public-key"
+	clientPrivateKeyName = "client-private-key"
+	clientPublicKeyName  = "client-public-key"
 )

experimental/ssh/cmd/server.go

@@ -27,14 +27,14 @@ and proxies them to local SSH daemon processes.
 	var shutdownDelay time.Duration
 	var clusterID string
 	var version string
-	var keysSecretScopeName string
-	var authorizedKeyName string
+	var secretScopeName string
+	var authorizedKeySecretName string
 
 	cmd.Flags().StringVar(&clusterID, "cluster", "", "Databricks cluster ID")
 	cmd.MarkFlagRequired("cluster")
-	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.Flags().StringVar(&secretScopeName, "secret-scope-name", "", "Databricks secret scope name to store SSH keys")
+	cmd.MarkFlagRequired("secret-scope-name")
+	cmd.Flags().StringVar(&authorizedKeySecretName, "authorized-key-secret-name", "", "Name of the secret containing the client public key")
 	cmd.MarkFlagRequired("authorized-key-secret-name")
 
 	cmd.Flags().IntVar(&maxClients, "max-clients", defaultMaxClients, "Maximum number of SSH clients")
@@ -55,17 +55,17 @@ and proxies them to local SSH daemon processes.
 		ctx := cmd.Context()
 		wsc := cmdctx.WorkspaceClient(ctx)
 		opts := server.ServerOptions{
-			ClusterID:            clusterID,
-			MaxClients:           maxClients,
-			ShutdownDelay:        shutdownDelay,
-			Version:              version,
-			ConfigDir:            serverConfigDir,
-			KeysSecretScopeName:  keysSecretScopeName,
-			AuthorizedKeyName:    authorizedKeyName,
-			ServerPrivateKeyName: serverPrivateKeyName,
-			ServerPublicKeyName:  serverPublicKeyName,
-			DefaultPort:          defaultServerPort,
-			PortRange:            serverPortRange,
+			ClusterID:               clusterID,
+			MaxClients:              maxClients,
+			ShutdownDelay:           shutdownDelay,
+			Version:                 version,
+			ConfigDir:               serverConfigDir,
+			SecretScopeName:         secretScopeName,
+			ServerPrivateKeyName:    serverPrivateKeyName,
+			ServerPublicKeyName:     serverPublicKeyName,
+			AuthorizedKeySecretName: authorizedKeySecretName,
+			DefaultPort:             defaultServerPort,
+			PortRange:               serverPortRange,
 		}
 		return server.Run(ctx, wsc, opts)
 	}

experimental/ssh/internal/client/client.go

@@ -58,8 +58,10 @@ type ClientOptions struct {
 	ReleasesDir string
 	// Directory for local SSH keys. Defaults to ~/.databricks/ssh-tunnel-keys
 	SSHKeysDir string
-	// Name of the client public key file to be used in the ssh-tunnel secrets scope.
+	// Client public key name located in the ssh-tunnel secrets scope.
 	ClientPublicKeyName string
+	// Client private key name located in the ssh-tunnel secrets scope.
+	ClientPrivateKeyName string
 	// If true, the CLI will attempt to start the cluster if it is not running.
 	AutoStartCluster bool
 	// Optional auth profile name. If present, will be added as --profile flag to the ProxyCommand while spawning ssh client.
@@ -85,21 +87,27 @@ func Run(ctx context.Context, client *databricks.WorkspaceClient, opts ClientOpt
 		return err
 	}
 
-	keyPath, err := keys.GetLocalSSHKeyPath(opts.ClusterID, opts.SSHKeysDir)
+	secretScopeName, err := keys.CreateKeysSecretScope(ctx, client, opts.ClusterID)
 	if err != nil {
-		return fmt.Errorf("failed to get local keys folder: %w", err)
+		return fmt.Errorf("failed to create secret scope: %w", err)
+	}
+
+	privateKeyBytes, publicKeyBytes, err := keys.CheckAndGenerateSSHKeyPairFromSecrets(ctx, client, opts.ClusterID, secretScopeName, opts.ClientPrivateKeyName, opts.ClientPublicKeyName)
+	if err != nil {
+		return fmt.Errorf("failed to get or generate SSH key pair from secrets: %w", err)
 	}
-	privateKeyPath, publicKey, err := keys.CheckAndGenerateSSHKeyPair(ctx, keyPath)
+
+	keyPath, err := keys.GetLocalSSHKeyPath(opts.ClusterID, opts.SSHKeysDir)
 	if err != nil {
-		return fmt.Errorf("failed to check or generate SSH key pair: %w", err)
+		return fmt.Errorf("failed to get local keys folder: %w", err)
 	}
-	cmdio.LogString(ctx, "Using SSH key: "+privateKeyPath)
 
-	keysSecretScopeName, err := keys.PutSecretInScope(ctx, client, opts.ClusterID, opts.ClientPublicKeyName, publicKey)
+	err = keys.SaveSSHKeyPair(keyPath, privateKeyBytes, publicKeyBytes)
 	if err != nil {
-		return fmt.Errorf("failed to store public key in secret scope: %w", err)
+		return fmt.Errorf("failed to save SSH key pair locally: %w", err)
 	}
-	cmdio.LogString(ctx, fmt.Sprintf("Secrets scope: %s, key name: %s", keysSecretScopeName, opts.ClientPublicKeyName))
+	cmdio.LogString(ctx, "Using SSH key: "+keyPath)
+	cmdio.LogString(ctx, fmt.Sprintf("Secrets scope: %s, key name: %s", secretScopeName, opts.ClientPublicKeyName))
 
 	var userName string
 	var serverPort int
@@ -111,7 +119,7 @@ func Run(ctx context.Context, client *databricks.WorkspaceClient, opts ClientOpt
 		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, version, keysSecretScopeName, opts)
+		userName, serverPort, err = ensureSSHServerIsRunning(ctx, client, version, secretScopeName, opts)
 		if err != nil {
 			return fmt.Errorf("failed to ensure that ssh server is running: %w", err)
 		}
@@ -137,7 +145,7 @@ func Run(ctx context.Context, client *databricks.WorkspaceClient, opts ClientOpt
 		return runSSHProxy(ctx, client, serverPort, opts)
 	} else {
 		cmdio.LogString(ctx, fmt.Sprintf("Additional SSH arguments: %v", opts.AdditionalArgs))
-		return spawnSSHClient(ctx, userName, privateKeyPath, serverPort, opts)
+		return spawnSSHClient(ctx, userName, keyPath, serverPort, opts)
 	}
 }
 
@@ -175,7 +183,7 @@ func getServerMetadata(ctx context.Context, client *databricks.WorkspaceClient,
 	return serverPort, string(bodyBytes), nil
 }
 
-func submitSSHTunnelJob(ctx context.Context, client *databricks.WorkspaceClient, version, keysSecretScopeName string, opts ClientOptions) (int64, error) {
+func submitSSHTunnelJob(ctx context.Context, client *databricks.WorkspaceClient, version, secretScopeName string, opts ClientOptions) (int64, error) {
 	contentDir, err := sshWorkspace.GetWorkspaceContentDir(ctx, client, version, opts.ClusterID)
 	if err != nil {
 		return 0, fmt.Errorf("failed to get workspace content directory: %w", err)
@@ -212,7 +220,7 @@ func submitSSHTunnelJob(ctx context.Context, client *databricks.WorkspaceClient,
 					NotebookPath: jobNotebookPath,
 					BaseParameters: map[string]string{
 						"version":                 version,
-						"keysSecretScopeName":     keysSecretScopeName,
+						"secretScopeName":         secretScopeName,
 						"authorizedKeySecretName": opts.ClientPublicKeyName,
 						"shutdownDelay":           opts.ShutdownDelay.String(),
 						"maxClients":              strconv.Itoa(opts.MaxClients),
@@ -290,12 +298,12 @@ func checkClusterState(ctx context.Context, client *databricks.WorkspaceClient,
 	return nil
 }
 
-func ensureSSHServerIsRunning(ctx context.Context, client *databricks.WorkspaceClient, version, keysSecretScopeName string, opts ClientOptions) (string, int, error) {
+func ensureSSHServerIsRunning(ctx context.Context, client *databricks.WorkspaceClient, version, secretScopeName string, opts ClientOptions) (string, int, error) {
 	serverPort, userName, err := getServerMetadata(ctx, client, opts.ClusterID, version)
 	if errors.Is(err, errServerMetadata) {
 		cmdio.LogString(ctx, "SSH server is not running, starting it now...")
 
-		runID, err := submitSSHTunnelJob(ctx, client, version, keysSecretScopeName, opts)
+		runID, err := submitSSHTunnelJob(ctx, client, version, secretScopeName, opts)
 		if err != nil {
 			return "", 0, fmt.Errorf("failed to submit ssh server job: %w", err)
 		}

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

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

experimental/ssh/internal/keys/keys.go

@@ -9,8 +9,6 @@ import (
 	"fmt"
 	"os"
 	"path/filepath"
-	"runtime"
-	"strings"
 
 	"github.com/databricks/databricks-sdk-go"
 	"golang.org/x/crypto/ssh"
@@ -70,61 +68,6 @@ func SaveSSHKeyPair(keyPath string, privateKeyBytes, publicKeyBytes []byte) erro
 	return nil
 }
 
-func checkSSHKeyPairPermissions(keyPath string) error {
-	dir := filepath.Dir(keyPath)
-	dirInfo, err := os.Stat(dir)
-	if err != nil {
-		if os.IsNotExist(err) {
-			return fmt.Errorf("key directory does not exist: %s", dir)
-		} else {
-			return fmt.Errorf("failed to stat key directory: %w", err)
-		}
-	}
-	if runtime.GOOS != "windows" && dirInfo.Mode().Perm() != 0o700 {
-		return fmt.Errorf("key directory permissions are not set to 0700, current permissions: %o", dirInfo.Mode().Perm())
-	}
-
-	info, err := os.Stat(keyPath)
-	if err != nil {
-		return fmt.Errorf("failed to stat key file: %w", err)
-	}
-
-	if runtime.GOOS != "windows" && info.Mode().Perm() != 0o600 {
-		return fmt.Errorf("private key permissions are not set to 0600, current permissions: %o", info.Mode().Perm())
-	}
-
-	pubKeyPath := keyPath + ".pub"
-	pubInfo, err := os.Stat(pubKeyPath)
-	if err != nil {
-		return fmt.Errorf("failed to stat public key file: %w", err)
-	}
-
-	if runtime.GOOS != "windows" && pubInfo.Mode().Perm() != 0o644 {
-		return fmt.Errorf("public key permissions are not set to 0644, current permissions: %o", pubInfo.Mode().Perm())
-	}
-
-	return nil
-}
-
-func CheckAndGenerateSSHKeyPair(ctx context.Context, keyPath string) (string, string, error) {
-	if err := checkSSHKeyPairPermissions(keyPath); err != nil {
-		privateKeyBytes, publicKeyBytes, err := generateSSHKeyPair()
-		if err != nil {
-			return "", "", err
-		}
-		if err := SaveSSHKeyPair(keyPath, privateKeyBytes, publicKeyBytes); err != nil {
-			return "", "", err
-		}
-	}
-
-	publicKeyBytes, err := os.ReadFile(keyPath + ".pub")
-	if err != nil {
-		return "", "", fmt.Errorf("failed to read public key: %w", err)
-	}
-
-	return keyPath, strings.TrimSpace(string(publicKeyBytes)), nil
-}
-
 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 {

experimental/ssh/internal/keys/secrets.go

@@ -10,7 +10,7 @@ import (
 	"github.com/databricks/databricks-sdk-go/service/workspace"
 )
 
-func createKeysSecretScope(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)
@@ -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 := createKeysSecretScope(ctx, client, clusterID)
+	scopeName, err := CreateKeysSecretScope(ctx, client, clusterID)
 	if err != nil {
 		return "", err
 	}

experimental/ssh/internal/server/server.go

@@ -34,13 +34,13 @@ type ServerOptions struct {
 	// The directory to store sshd configuration
 	ConfigDir string
 	// The name of the secrets scope to use for client and server keys
-	KeysSecretScopeName string
-	// The name of a secret containing the client's public key value
-	AuthorizedKeyName string
+	SecretScopeName string
 	// The name of a secret containing the server's private key value
 	ServerPrivateKeyName string
 	// The name of a secret containing the server's public key value
 	ServerPublicKeyName string
+	// The name of a secret containing the client's public key (authorized key)
+	AuthorizedKeySecretName string
 	// The default port to listen on (for /ssh and /metadata requests from the clients)
 	DefaultPort int
 	// If the default port is taken, the server will try to listen on the first free port in the DefaultPort + PortRange range
@@ -71,10 +71,10 @@ func Run(ctx context.Context, client *databricks.WorkspaceClient, opts ServerOpt
 		return fmt.Errorf("failed to save Jupyter init script: %w", err)
 	}
 
-	connections := proxy.NewConnectionsManager(opts.MaxClients, opts.ShutdownDelay)
 	createServerCommand := func(ctx context.Context) *exec.Cmd {
 		return createSSHDProcess(ctx, sshdConfigPath)
 	}
+	connections := proxy.NewConnectionsManager(opts.MaxClients, opts.ShutdownDelay)
 	http.Handle("/ssh", proxy.NewProxyServer(ctx, connections, createServerCommand))
 	http.HandleFunc("/metadata", serveMetadata)
 	go handleTimeout(ctx, connections.TimedOut, opts.ShutdownDelay)

experimental/ssh/internal/server/sshd.go

@@ -15,7 +15,7 @@ import (
 )
 
 func prepareSSHDConfig(ctx context.Context, client *databricks.WorkspaceClient, opts ServerOptions) (string, error) {
-	clientPublicKey, err := keys.GetSecret(ctx, client, opts.KeysSecretScopeName, opts.AuthorizedKeyName)
+	clientPublicKey, err := keys.GetSecret(ctx, client, opts.SecretScopeName, opts.AuthorizedKeySecretName)
 	if err != nil {
 		return "", fmt.Errorf("failed to get client public key: %w", err)
 	}
@@ -36,7 +36,7 @@ func prepareSSHDConfig(ctx context.Context, client *databricks.WorkspaceClient,
 		return "", fmt.Errorf("failed to create SSH directory: %w", err)
 	}
 
-	privateKeyBytes, publicKeyBytes, err := keys.CheckAndGenerateSSHKeyPairFromSecrets(ctx, client, opts.ClusterID, opts.KeysSecretScopeName, opts.ServerPrivateKeyName, opts.ServerPublicKeyName)
+	privateKeyBytes, publicKeyBytes, err := keys.CheckAndGenerateSSHKeyPairFromSecrets(ctx, client, opts.ClusterID, opts.SecretScopeName, opts.ServerPrivateKeyName, opts.ServerPublicKeyName)
 	if err != nil {
 		return "", fmt.Errorf("failed to get SSH key pair from secrets: %w", err)
 	}
@@ -47,8 +47,8 @@ func prepareSSHDConfig(ctx context.Context, client *databricks.WorkspaceClient,
 	}
 
 	sshdConfig := filepath.Join(sshDir, "sshd_config")
-	authKeys := filepath.Join(sshDir, "authorized_keys")
-	if err := os.WriteFile(authKeys, clientPublicKey, 0o600); err != nil {
+	authKeysPath := filepath.Join(sshDir, "authorized_keys")
+	if err := os.WriteFile(authKeysPath, clientPublicKey, 0o600); err != nil {
 		return "", err
 	}
 
@@ -75,7 +75,7 @@ func prepareSSHDConfig(ctx context.Context, client *databricks.WorkspaceClient,
 		"ChallengeResponseAuthentication no\n" +
 		"Subsystem sftp internal-sftp\n" +
 		"HostKey " + keyPath + "\n" +
-		"AuthorizedKeysFile " + authKeys + "\n" +
+		"AuthorizedKeysFile " + authKeysPath + "\n" +
 		setEnv + "\n"
 
 	if err := os.WriteFile(sshdConfig, []byte(sshdConfigContent), 0o600); err != nil {