[SSH] Fail fast when the driver proxy refuses the tunnel

.nextchanges/cli/ssh-port-forwarding-disabled.md

@@ -0,0 +1 @@
+* `databricks ssh connect` now fails right away with the workspace's reason when port forwarding to serverless compute is disabled, instead of waiting out the server start timeout. ([#7057](https://github.com/databricks/cli/pull/7057))

acceptance/ssh/connect-serverless-feature-disabled/out.connect.txt

@@ -0,0 +1,5 @@
+Connecting to feature-disabled...
+Uploading binaries...
+Starting SSH server...
+Job submitted successfully with run ID: [RUN_ID]
+Error: failed to ensure that ssh server is running: the workspace refused the SSH tunnel connection: FEATURE_DISABLED: An administrator has disabled port forwarding to serverless compute for this workspace, which includes the SSH tunnel used by `databricks ssh connect`

acceptance/ssh/connect-serverless-feature-disabled/out.test.toml

@@ -0,0 +1,5 @@
+Cloud = false
+GOOS.darwin = false
+GOOS.linux = true
+GOOS.windows = false
+EnvMatrix.DATABRICKS_BUNDLE_ENGINE = ["direct"]

acceptance/ssh/connect-serverless-feature-disabled/output.txt

@@ -0,0 +1,13 @@
+
+>>> [CLI] workspace-settings-v2 patch-public-workspace-setting port_fwd_serverless_enabled --json {"name": "port_fwd_serverless_enabled", "boolean_val": {"value": false}}
+
+=== CLI ssh connect
+
+>>> print_requests.py //api/2.2/jobs/runs/cancel
+{
+  "method": "POST",
+  "path": "/api/2.2/jobs/runs/cancel",
+  "body": {
+    "run_id": [RUN_ID]
+  }
+}

acceptance/ssh/connect-serverless-feature-disabled/script

@@ -0,0 +1,17 @@
+# An administrator disabled port forwarding to serverless compute, so the driver proxy refuses
+# every request to the tunnel server. The CLI must report the proxy's reason as soon as it gets
+# the 403 instead of polling until the server start timeout, and cancel the run it submitted
+# since nothing can ever connect to that server.
+# The body repeats the name: with only {"value": false} in it, the API sees an empty message.
+trace $CLI workspace-settings-v2 patch-public-workspace-setting port_fwd_serverless_enabled --json '{"name": "port_fwd_serverless_enabled", "boolean_val": {"value": false}}' >/dev/null
+
+# No release artifacts locally and they're never executed; stand in empty archives.
+mkdir -p releases
+: >releases/databricks_cli_linux_amd64.zip
+: >releases/databricks_cli_linux_arm64.zip
+
+title "CLI ssh connect\n"
+musterr $CLI ssh connect --name feature-disabled --releases-dir="$PWD/releases" -- "echo 'Connection successful'" >out.connect.txt 2>&1
+add_repl "$(grep -o 'run ID: [0-9]*' out.connect.txt | grep -o '[0-9]*$')" RUN_ID
+
+trace print_requests.py //api/2.2/jobs/runs/cancel

acceptance/ssh/connect-serverless-feature-disabled/test.toml

@@ -0,0 +1,16 @@
+Cloud = false
+
+# Record requests so we can assert that the CLI cancels the run it submitted.
+RecordRequests = true
+
+# Stand-in release archives staged by the script; not golden output.
+Ignore = [
+  "releases",
+]
+
+# The rest of acceptance/ssh only runs on Linux; nothing here needs another OS.
+GOOS.linux = true
+GOOS.darwin = false
+GOOS.windows = false
+
+EnvMatrix.DATABRICKS_BUNDLE_ENGINE = ["direct"]

experimental/ssh/internal/client/client.go

@@ -47,6 +47,11 @@ var sshServerBootstrapScript string
 
 var errServerMetadata = errors.New("server metadata error")
 
+// errDriverProxyForbidden is a 403 from the driver proxy. Unlike the other non-200 responses, which a
+// server that is still starting produces too, retrying cannot change it: the workspace forbids the
+// connection, e.g. because an administrator disabled port forwarding to serverless compute.
+var errDriverProxyForbidden = errors.New("the workspace refused the SSH tunnel connection")
+
 var connectionNameRegex = regexp.MustCompile(`^[a-zA-Z0-9][a-zA-Z0-9_-]*$`)
 
 const (
@@ -653,7 +658,7 @@ func getServerMetadata(ctx context.Context, client *databricks.WorkspaceClient,
 	log.Debugf(ctx, "Metadata response status code: %d", resp.StatusCode)
 
 	if resp.StatusCode != http.StatusOK {
-		return serverMetadata{}, errors.Join(errServerMetadata, fmt.Errorf("server is not ok, status code %d", resp.StatusCode))
+		return serverMetadata{}, metadataResponseError(resp.StatusCode, bodyBytes)
 	}
 
 	return serverMetadata{
@@ -665,6 +670,26 @@ func getServerMetadata(ctx context.Context, client *databricks.WorkspaceClient,
 	}, nil
 }
 
+// metadataResponseError classifies a non-200 response from the server's /metadata endpoint. A 403
+// is final and carries the driver proxy's own explanation; anything else is errServerMetadata,
+// which callers treat as a server that is not running yet.
+func metadataResponseError(statusCode int, body []byte) error {
+	if statusCode != http.StatusForbidden {
+		return errors.Join(errServerMetadata, fmt.Errorf("server is not ok, status code %d", statusCode))
+	}
+	// The driver proxy writes its error as a JSON string ("CODE: message") rather than the usual
+	// {"error_code", "message"} object.
+	var message string
+	if json.Unmarshal(body, &message) != nil {
+		message = string(body)
+	}
+	message = strings.TrimSpace(message)
+	if message == "" {
+		message = fmt.Sprintf("status code %d", statusCode)
+	}
+	return fmt.Errorf("%w: %s", errDriverProxyForbidden, message)
+}
+
 // newDriverProxyRequest builds an authenticated GET request to one of the SSH server's
 // HTTP endpoints behind the workspace driver proxy.
 func newDriverProxyRequest(ctx context.Context, client *databricks.WorkspaceClient, clusterID string, port int, endpoint, liteswap string) (*http.Request, error) {
@@ -1292,6 +1317,15 @@ func describeRunFailure(ctx context.Context, client *databricks.WorkspaceClient,
 	return strings.TrimRight(b.String(), "\n")
 }
 
+// cancelSSHServerRun cancels the bootstrap run this connection submitted once no client can reach
+// its server. Best effort: the server also shuts itself down after its no-client delay, so a
+// failure here must not hide the error that made the connection give up.
+func cancelSSHServerRun(ctx context.Context, client *databricks.WorkspaceClient, runID int64) {
+	if _, err := client.Jobs.CancelRun(ctx, jobs.CancelRun{RunId: runID}); err != nil {
+		log.Debugf(ctx, "Failed to cancel the ssh server job run %d: %v", runID, err)
+	}
+}
+
 // runFailureIfTerminated reports whether the bootstrap run has reached a terminal state (so the
 // SSH server will never come up), returning a formatted failure description when it has.
 func runFailureIfTerminated(ctx context.Context, client *databricks.WorkspaceClient, runID int64) (string, bool) {
@@ -1381,6 +1415,7 @@ func ensureSSHServerIsRunning(ctx context.Context, client *databricks.WorkspaceC
 	clusterID := opts.ClusterID
 
 	meta, err := getServerMetadata(ctx, client, sessionID, clusterID, version, opts.Liteswap)
+	// This includes errDriverProxyForbidden: a new server would be refused the same way.
 	if err != nil && !errors.Is(err, errServerMetadata) {
 		return "", 0, "", err
 	}
@@ -1410,6 +1445,10 @@ func ensureSSHServerIsRunning(ctx context.Context, client *databricks.WorkspaceC
 				return "", 0, "", ctx.Err()
 			}
 			meta, err = getServerMetadata(ctx, client, sessionID, clusterID, version, opts.Liteswap)
+			if errors.Is(err, errDriverProxyForbidden) {
+				cancelSSHServerRun(ctx, client, runID)
+				return "", 0, "", err
+			}
 			// Accept only once metadata reflects the requested usage policy, so we don't latch
 			// onto a server a previous connection started under a different policy before our new
 			// server has overwritten metadata.json.

experimental/ssh/internal/client/client_internal_test.go

@@ -803,3 +803,58 @@ func TestBuildSshTunnelEventReportsFailure(t *testing.T) {
 	require.NoError(t, err)
 	assert.Contains(t, string(b), `"is_success":false`)
 }
+
+func TestMetadataResponseError(t *testing.T) {
+	tests := []struct {
+		name       string
+		statusCode int
+		body       string
+		forbidden  bool
+		message    string
+	}{
+		{
+			name:       "feature disabled",
+			statusCode: 403,
+			body:       `"FEATURE_DISABLED: An administrator has disabled port forwarding to serverless compute for this workspace"`,
+			forbidden:  true,
+			message:    "the workspace refused the SSH tunnel connection: FEATURE_DISABLED: An administrator has disabled port forwarding to serverless compute for this workspace",
+		},
+		{
+			name:       "forbidden with a body that is not a JSON string",
+			statusCode: 403,
+			body:       "Forbidden\n",
+			forbidden:  true,
+			message:    "the workspace refused the SSH tunnel connection: Forbidden",
+		},
+		{
+			name:       "forbidden with an empty body",
+			statusCode: 403,
+			body:       "",
+			forbidden:  true,
+			message:    "the workspace refused the SSH tunnel connection: status code 403",
+		},
+		{
+			// A stale metadata.json points at a terminated cluster; the caller must start a new server.
+			name:       "terminated cluster is retried",
+			statusCode: 400,
+			body:       `"INVALID_STATE: Cluster 1014-135824-abcdefgh-v2n is in Terminated state but Spark driver is not healthy"`,
+			forbidden:  false,
+			message:    "server metadata error\nserver is not ok, status code 400",
+		},
+		{
+			name:       "server still starting is retried",
+			statusCode: 502,
+			body:       "",
+			forbidden:  false,
+			message:    "server metadata error\nserver is not ok, status code 502",
+		},
+	}
+	for _, tc := range tests {
+		t.Run(tc.name, func(t *testing.T) {
+			err := metadataResponseError(tc.statusCode, []byte(tc.body))
+			assert.Equal(t, tc.forbidden, errors.Is(err, errDriverProxyForbidden))
+			assert.Equal(t, !tc.forbidden, errors.Is(err, errServerMetadata))
+			assert.EqualError(t, err, tc.message)
+		})
+	}
+}

libs/testserver/fake_workspace.go

@@ -34,6 +34,7 @@ import (
 	"github.com/databricks/databricks-sdk-go/service/ml"
 	"github.com/databricks/databricks-sdk-go/service/pipelines"
 	"github.com/databricks/databricks-sdk-go/service/serving"
+	"github.com/databricks/databricks-sdk-go/service/settingsv2"
 	"github.com/databricks/databricks-sdk-go/service/sql"
 	"github.com/databricks/databricks-sdk-go/service/vectorsearch"
 	"github.com/databricks/databricks-sdk-go/service/workspace"
@@ -241,6 +242,9 @@ type FakeWorkspace struct {
 	Acls         map[string][]workspace.AclItem
 	UCSecrets    map[string]catalog.Secret // full_name -> secret (Unity Catalog secrets)
 
+	// WorkspaceSettings holds settings written through the settings v2 public API, keyed by name.
+	WorkspaceSettings map[string]settingsv2.Setting
+
 	// Generic permissions storage: key is "{object_type}:{object_id}"
 	Permissions map[string]iam.ObjectPermissions
 
@@ -525,6 +529,7 @@ func NewFakeWorkspace(url, token string) *FakeWorkspace {
 		SecretScopes:              map[string]workspace.SecretScope{},
 		Secrets:                   map[string]map[string]string{},
 		Acls:                      map[string][]workspace.AclItem{},
+		WorkspaceSettings:         map[string]settingsv2.Setting{},
 		Permissions:               map[string]iam.ObjectPermissions{},
 		Groups:                    map[string]iam.Group{},
 		DatabaseInstances:         map[string]database.DatabaseInstance{},

libs/testserver/handlers.go

@@ -1047,6 +1047,10 @@ func AddDefaultHandlers(server *Server) {
 	// user, /logs is a best-effort error tail, /ssh hijacks the connection for a
 	// websocket (so it's raw, not a JSON response).
 	server.Handle("GET", "/driver-proxy-api/o/{workspace_id}/{cluster_id}/{port}/metadata", func(req Request) any {
+		defer req.Workspace.LockUnlock()()
+		if req.Vars["cluster_id"] == sshTunnelClusterID && req.Workspace.portForwardingToServerlessDisabled() {
+			return Response{StatusCode: 403, Body: sshTunnelFeatureDisabledBody}
+		}
 		return Response{Body: sshTunnelRemoteUser()}
 	})
 
@@ -1065,6 +1069,10 @@ func AddDefaultHandlers(server *Server) {
 
 	server.HandleRaw("GET", "/driver-proxy-api/o/{workspace_id}/{cluster_id}/{port}/ssh", server.sshTunnelHandler)
 
+	server.Handle("PATCH", "/api/2.1/settings/{name}", func(req Request) any {
+		return req.Workspace.WorkspaceSettingPatch(req)
+	})
+
 	// Secrets ACLs:
 	server.Handle("GET", "/api/2.0/secrets/acls/get", func(req Request) any {
 		return req.Workspace.SecretsAclsGet(req)

libs/testserver/jobs.go

@@ -693,6 +693,11 @@ const (
 	sshTunnelBootstrapNotebook  = "ssh-server-bootstrap"
 	sshTunnelServerPort         = 7772
 	sshTunnelClusterID          = "1234-567890-serverless"
+
+	// sshTunnelFeatureDisabledBody is the driver proxy's refusal while port forwarding to
+	// serverless compute is disabled: a JSON string, not the usual error object.
+	sshTunnelFeatureDisabledBody = `"FEATURE_DISABLED: An administrator has disabled port forwarding to serverless compute ` +
+		"for this workspace, which includes the SSH tunnel used by `databricks ssh connect`\""
 )
 
 // writeSSHTunnelMetadata publishes the metadata.json a real tunnel server would

libs/testserver/settings.go

@@ -0,0 +1,49 @@
+package testserver
+
+import (
+	"encoding/json"
+	"fmt"
+
+	"github.com/databricks/databricks-sdk-go/service/settingsv2"
+)
+
+// portFwdServerlessSetting lets an administrator disable port forwarding to serverless compute,
+// which the driver proxy enforces for the SSH tunnel.
+const portFwdServerlessSetting = "port_fwd_serverless_enabled"
+
+func (s *FakeWorkspace) WorkspaceSettingPatch(req Request) Response {
+	var setting settingsv2.Setting
+	if err := json.Unmarshal(req.Body, &setting); err != nil {
+		return Response{
+			StatusCode: 400,
+			Body:       fmt.Sprintf("request parsing error: %s", err),
+		}
+	}
+
+	// The backend decodes the body as a protobuf message, in which false is the default value, so a
+	// body that carries only {"boolean_val": {"value": false}} arrives empty and is rejected.
+	if setting.Name == "" && (setting.BooleanVal == nil || !setting.BooleanVal.Value) {
+		return Response{
+			StatusCode: 400,
+			Body: map[string]string{
+				"error_code": "INVALID_PARAMETER_VALUE",
+				"message":    "Field 'setting' is required and must contain at least one subfield with a non-default value!",
+			},
+		}
+	}
+
+	defer s.LockUnlock()()
+
+	setting.Name = req.Vars["name"]
+	// Account-level values are not modelled, so the workspace value is the effective one.
+	setting.EffectiveBooleanVal = setting.BooleanVal
+	s.WorkspaceSettings[setting.Name] = setting
+	return Response{Body: setting}
+}
+
+// portForwardingToServerlessDisabled reports whether port_fwd_serverless_enabled was set to false.
+// Callers must hold the workspace lock.
+func (s *FakeWorkspace) portForwardingToServerlessDisabled() bool {
+	setting, ok := s.WorkspaceSettings[portFwdServerlessSetting]
+	return ok && setting.EffectiveBooleanVal != nil && !setting.EffectiveBooleanVal.Value
+}

libs/testserver/settings_test.go

@@ -0,0 +1,32 @@
+package testserver
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func patchSetting(workspace *FakeWorkspace, body string) Response {
+	return workspace.WorkspaceSettingPatch(Request{
+		Body: []byte(body),
+		Vars: map[string]string{"name": portFwdServerlessSetting},
+	})
+}
+
+func TestWorkspaceSettingPatchRejectsBodyWithOnlyDefaults(t *testing.T) {
+	workspace := NewFakeWorkspace("http://test.test", "token")
+
+	response := patchSetting(workspace, `{"boolean_val": {"value": false}}`)
+
+	assert.Equal(t, 400, response.StatusCode)
+	assert.False(t, workspace.portForwardingToServerlessDisabled())
+}
+
+func TestWorkspaceSettingPatchDisablesPortForwarding(t *testing.T) {
+	workspace := NewFakeWorkspace("http://test.test", "token")
+
+	response := patchSetting(workspace, `{"name": "port_fwd_serverless_enabled", "boolean_val": {"value": false}}`)
+
+	assert.Equal(t, 0, response.StatusCode)
+	assert.True(t, workspace.portForwardingToServerlessDisabled())
+}