Update description, add comments

Author: ilia-db

cmd/ssh/connect.go

@@ -12,11 +12,11 @@ import (
 func newConnectCommand() *cobra.Command {
 	cmd := &cobra.Command{
 		Use:   "connect",
-		Short: "Connect to a Databricks cluster via SSH",
-		Long: `Connect to a Databricks cluster via SSH.
+		Short: "Connect to Databricks compute via SSH",
+		Long: `Connect to Databricks compute via SSH.
 
-This command establishes an SSH connection to a Databricks cluster, setting up
-the necessary SSH server on the cluster and handling the connection proxy.
+This command establishes an SSH connection to Databricks compute, setting up
+the SSH server and handling the connection proxy.
 
 ` + disclaimer,
 	}

libs/ssh/client.go

@@ -37,16 +37,30 @@ var sshServerBootstrapScript string
 var errServerMetadata = errors.New("server metadata error")
 
 type ClientOptions struct {
-	ClusterID           string
-	ProxyMode           bool
-	ServerMetadata      string
-	ShutdownDelay       time.Duration
-	MaxClients          int
-	HandoverTimeout     time.Duration
-	ReleasesDir         string
-	SSHKeysDir          string
-	AdditionalArgs      []string
+	// Id of the cluster to connect to
+	ClusterID string
+	// Delay before shutting down the server after the last client disconnects
+	ShutdownDelay time.Duration
+	// Maximum number of SSH clients
+	MaxClients int
+	// Indicates that the CLI runs as a ProxyCommand - it should establish ws connection
+	// to the cluster and proxy all traffic through stdin/stdout.
+	// In the non proxy mode the CLI spawns an ssh client with the ProxyCommand config.
+	ProxyMode bool
+	// Expected format: "<user_name>,<port>".
+	// If present, the CLI won't attempt to start the server.
+	ServerMetadata string
+	// How often the CLI should reconnect to the server with new auth.
+	HandoverTimeout time.Duration
+	// Directory for local SSH tunnel development releases.
+	// If not present, the CLI will use github releases with the current version.
+	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.
 	ClientPublicKeyName string
+	// Additional arguments to pass to the SSH client in the non proxy mode.
+	AdditionalArgs []string
 }
 
 func RunClient(ctx context.Context, client *databricks.WorkspaceClient, opts ClientOptions) error {

libs/ssh/proxy.go

@@ -15,8 +15,10 @@ import (
 	"golang.org/x/sync/errgroup"
 )
 
-var (
-	errProxyEOF                      = errors.New("proxy EOF error")
+var errProxyEOF = errors.New("proxy EOF error")
+
+const (
+	proxyBufferSize                  = 4 * 1024 // Same as gorilla/websocket default read/write buffer sizes. Bigger payloads will be split into multiple ws frames.
 	proxyHandoverInitTimeout         = 30 * time.Second
 	proxyHandoverAcceptTimeout       = 25 * time.Second
 	proxyHandoverAckCloseConnTimeout = 15 * time.Second
@@ -85,7 +87,7 @@ func (pc *proxyConnection) runSendingLoop(ctx context.Context, src io.Reader) er
 		if ctx.Err() != nil {
 			return ctx.Err()
 		}
-		b := make([]byte, 32*1024)
+		b := make([]byte, proxyBufferSize)
 		n, readErr := src.Read(b)
 		if n > 0 {
 			// This will block during handover - we stop sending anything except the close message.