diff --git a/acceptance/ssh/connect-serverless-feature-disabled/script b/acceptance/ssh/connect-serverless-feature-disabled/script index 2121593dc7f..9fb5bb05b06 100644 --- a/acceptance/ssh/connect-serverless-feature-disabled/script +++ b/acceptance/ssh/connect-serverless-feature-disabled/script @@ -1,6 +1,7 @@ # 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 +# every request to the tunnel server. The setting itself cannot be read here (see test.toml), so +# the CLI goes ahead and starts the server. It must then 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 diff --git a/acceptance/ssh/connect-serverless-feature-disabled/test.toml b/acceptance/ssh/connect-serverless-feature-disabled/test.toml index 52d2432304d..36257965b75 100644 --- a/acceptance/ssh/connect-serverless-feature-disabled/test.toml +++ b/acceptance/ssh/connect-serverless-feature-disabled/test.toml @@ -14,3 +14,16 @@ GOOS.darwin = false GOOS.windows = false EnvMatrix.DATABRICKS_BUNDLE_ENGINE = ["direct"] + +# Production does not serve port_fwd_serverless_enabled through the public API, so the CLI cannot +# check it up front there and only learns about it from the driver proxy's 403. This reproduces +# that environment, which the fake workspace's settings handler does not model. +[[Server]] +Pattern = "GET /api/2.1/settings/port_fwd_serverless_enabled" +Response.StatusCode = 404 +Response.Body = ''' +{ + "error_code": "NOT_FOUND", + "message": "Setting port_fwd_serverless_enabled is not enabled for public API" +} +''' diff --git a/acceptance/ssh/connect-serverless-port-forwarding-disabled/out.test.toml b/acceptance/ssh/connect-serverless-port-forwarding-disabled/out.test.toml new file mode 100644 index 00000000000..c826e1c76b6 --- /dev/null +++ b/acceptance/ssh/connect-serverless-port-forwarding-disabled/out.test.toml @@ -0,0 +1,5 @@ +Cloud = false +GOOS.darwin = false +GOOS.linux = true +GOOS.windows = false +EnvMatrix.DATABRICKS_BUNDLE_ENGINE = ["direct"] diff --git a/acceptance/ssh/connect-serverless-port-forwarding-disabled/output.txt b/acceptance/ssh/connect-serverless-port-forwarding-disabled/output.txt new file mode 100644 index 00000000000..3b3f5fea6f9 --- /dev/null +++ b/acceptance/ssh/connect-serverless-port-forwarding-disabled/output.txt @@ -0,0 +1,10 @@ + +>>> [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 + +>>> musterr [CLI] ssh connect --name port-forwarding-disabled -- echo 'Connection successful' +Connecting to port-forwarding-disabled... +Error: port forwarding to serverless compute is disabled for this workspace (setting port_fwd_serverless_enabled), and the SSH tunnel needs it. Ask a workspace or account admin to enable it, or connect to a dedicated cluster with --cluster + +>>> print_requests.py //api diff --git a/acceptance/ssh/connect-serverless-port-forwarding-disabled/script b/acceptance/ssh/connect-serverless-port-forwarding-disabled/script new file mode 100644 index 00000000000..6c250774c4b --- /dev/null +++ b/acceptance/ssh/connect-serverless-port-forwarding-disabled/script @@ -0,0 +1,13 @@ +# An administrator disabled port forwarding to serverless compute. The CLI reads the setting +# before anything else, so it refuses the connection without creating a secret scope, uploading +# binaries or starting compute that nobody could reach. +# 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 +# Drop the setup request so the check below only sees what ssh connect sent. +print_requests.py //api/2.1/settings >/dev/null + +title "CLI ssh connect\n" +trace musterr $CLI ssh connect --name port-forwarding-disabled -- "echo 'Connection successful'" + +# Expect no requests other than reads: the connection stops at the settings check. +trace print_requests.py //api diff --git a/acceptance/ssh/connect-serverless-port-forwarding-disabled/test.toml b/acceptance/ssh/connect-serverless-port-forwarding-disabled/test.toml new file mode 100644 index 00000000000..10c3ef353f7 --- /dev/null +++ b/acceptance/ssh/connect-serverless-port-forwarding-disabled/test.toml @@ -0,0 +1,11 @@ +Cloud = false + +# Record requests so we can assert that nothing is created before the CLI gives up. +RecordRequests = true + +# 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"] diff --git a/experimental/ssh/internal/client/agentshim.go b/experimental/ssh/internal/client/agentshim.go index 58663eb0d01..de9198e53c8 100644 --- a/experimental/ssh/internal/client/agentshim.go +++ b/experimental/ssh/internal/client/agentshim.go @@ -360,7 +360,7 @@ func probeAIGateway(ctx context.Context, wsclient *databricks.WorkspaceClient) e } // newProbeAPIClient builds a low-level client from cfg with the SDK's retries -// turned off, so the gateway preflight is one-shot. Clearing ErrorRetriable and +// turned off, so a preflight check is one-shot. Clearing ErrorRetriable and // TransientErrors makes a retriable status (429, 503, ...) surface on the first // attempt rather than looping until the retry budget expires; it mirrors // client.New, which builds the retrying client the rest of the CLI uses. diff --git a/experimental/ssh/internal/client/client.go b/experimental/ssh/internal/client/client.go index 723f6d1e119..28197d143c5 100644 --- a/experimental/ssh/internal/client/client.go +++ b/experimental/ssh/internal/client/client.go @@ -38,6 +38,7 @@ import ( "github.com/databricks/databricks-sdk-go/service/environments" "github.com/databricks/databricks-sdk-go/service/iam" "github.com/databricks/databricks-sdk-go/service/jobs" + "github.com/databricks/databricks-sdk-go/service/settingsv2" "github.com/databricks/databricks-sdk-go/service/workspace" "github.com/gorilla/websocket" ) @@ -47,6 +48,14 @@ var sshServerBootstrapScript string var errServerMetadata = errors.New("server metadata error") +// portFwdServerlessSetting is the workspace setting an administrator turns off to disable port +// forwarding to serverless compute, which the tunnel needs. +const portFwdServerlessSetting = "port_fwd_serverless_enabled" + +var errPortForwardingDisabled = errors.New("port forwarding to serverless compute is disabled for this workspace " + + "(setting " + portFwdServerlessSetting + "), and the SSH tunnel needs it. " + + "Ask a workspace or account admin to enable it, or connect to a dedicated cluster with --cluster") + // 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. @@ -346,6 +355,14 @@ func Run(ctx context.Context, client *databricks.WorkspaceClient, opts ClientOpt cmdio.LogString(ctx, fmt.Sprintf("Connecting to %s...", sessionID)) } + // Skipped when reconnecting with --metadata (the ProxyCommand child): the server is already + // running, so there is no compute left to spare, and every IDE reconnect would repeat the read. + if opts.IsServerlessMode() && opts.ServerMetadata == "" { + if err := checkPortForwardingToServerless(ctx, client); err != nil { + return err + } + } + if opts.IDE != "" && !opts.ProxyMode { if err := vscode.CheckIDECommand(opts.IDE); err != nil { outcome.errorCategory = protos.SshTunnelErrorCategoryIDECommandNotOnPath @@ -690,6 +707,33 @@ func metadataResponseError(statusCode int, body []byte) error { return fmt.Errorf("%w: %s", errDriverProxyForbidden, message) } +// checkPortForwardingToServerless refuses a serverless connection before any compute is started when +// the workspace setting that the driver proxy enforces for the tunnel is off. The proxy's gate +// (PortForwardingToServerlessGate in data-plane-router) also depends on an enforcement flag and a port +// allow-list that the CLI cannot see, and the setting is not readable through the public API in every +// environment. A read that fails therefore lets the connection go ahead, and the proxy's own 403 +// (errDriverProxyForbidden) remains the backstop. The read goes through a client with retries turned +// off, so a throttled or unavailable settings service cannot hold up the connection for the SDK's +// retry budget. +func checkPortForwardingToServerless(ctx context.Context, client *databricks.WorkspaceClient) error { + apiClient, err := newProbeAPIClient(client.Config) + if err != nil { + log.Debugf(ctx, "Could not build a client to read the %s setting, continuing: %v", portFwdServerlessSetting, err) + return nil + } + setting, err := settingsv2.NewWorkspaceSettingsV2(apiClient).GetPublicWorkspaceSetting(ctx, settingsv2.GetPublicWorkspaceSettingRequest{ + Name: portFwdServerlessSetting, + }) + if err != nil { + log.Debugf(ctx, "Could not read the %s setting, continuing: %v", portFwdServerlessSetting, err) + return nil + } + if setting.EffectiveBooleanVal != nil && !setting.EffectiveBooleanVal.Value { + return errPortForwardingDisabled + } + return nil +} + // 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) { diff --git a/experimental/ssh/internal/client/client_internal_test.go b/experimental/ssh/internal/client/client_internal_test.go index 64454328693..746f5d58f98 100644 --- a/experimental/ssh/internal/client/client_internal_test.go +++ b/experimental/ssh/internal/client/client_internal_test.go @@ -5,9 +5,12 @@ import ( "encoding/json" "errors" "fmt" + "net/http" + "net/http/httptest" "slices" "strconv" "strings" + "sync/atomic" "testing" "time" @@ -858,3 +861,60 @@ func TestMetadataResponseError(t *testing.T) { }) } } + +func TestCheckPortForwardingToServerless(t *testing.T) { + tests := []struct { + name string + status int + body string + wantErr error + }{ + { + name: "disabled", + status: http.StatusOK, + body: `{"effective_boolean_val": {"value": false}}`, + wantErr: errPortForwardingDisabled, + }, + { + name: "enabled", + status: http.StatusOK, + body: `{"effective_boolean_val": {"value": true}}`, + }, + { + name: "no effective value", + status: http.StatusOK, + body: `{}`, + }, + { + // Production does not serve the setting through the public API yet. + name: "read fails", + status: http.StatusBadRequest, + body: `{"error_code": "INVALID_PARAMETER_VALUE", "message": "Setting port_fwd_serverless_enabled is not enabled for public API"}`, + }, + { + // A retriable status is not retried, so a throttled service cannot hold up the connection. + name: "throttled", + status: http.StatusTooManyRequests, + body: `{"error_code": "TOO_MANY_REQUESTS", "message": "slow down"}`, + }, + } + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + var reads atomic.Int32 + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.URL.Path != "/api/2.1/settings/"+portFwdServerlessSetting { + w.WriteHeader(http.StatusNotFound) + return + } + reads.Add(1) + w.WriteHeader(tc.status) + _, _ = w.Write([]byte(tc.body)) + })) + defer srv.Close() + + ctx := cmdio.MockDiscard(t.Context()) + assert.Equal(t, tc.wantErr, checkPortForwardingToServerless(ctx, newProbeClient(t, srv.URL))) + assert.Equal(t, int32(1), reads.Load()) + }) + } +} diff --git a/libs/testserver/handlers.go b/libs/testserver/handlers.go index 3a9a538313e..f2d9cff5109 100644 --- a/libs/testserver/handlers.go +++ b/libs/testserver/handlers.go @@ -1070,6 +1070,10 @@ func AddDefaultHandlers(server *Server) { server.HandleRaw("GET", "/driver-proxy-api/o/{workspace_id}/{cluster_id}/{port}/ssh", server.sshTunnelHandler) + server.Handle("GET", "/api/2.1/settings/{name}", func(req Request) any { + return req.Workspace.WorkspaceSettingGet(req) + }) + server.Handle("PATCH", "/api/2.1/settings/{name}", func(req Request) any { return req.Workspace.WorkspaceSettingPatch(req) }) diff --git a/libs/testserver/settings.go b/libs/testserver/settings.go index 517d4725c42..958f0626f0b 100644 --- a/libs/testserver/settings.go +++ b/libs/testserver/settings.go @@ -11,6 +11,24 @@ import ( // which the driver proxy enforces for the SSH tunnel. const portFwdServerlessSetting = "port_fwd_serverless_enabled" +func (s *FakeWorkspace) WorkspaceSettingGet(req Request) Response { + defer s.LockUnlock()() + + name := req.Vars["name"] + if setting, ok := s.WorkspaceSettings[name]; ok { + return Response{Body: setting} + } + // The setting's own default is true: port forwarding stays on until an administrator turns it off. + if name == portFwdServerlessSetting { + enabled := &settingsv2.BooleanMessage{Value: true} + return Response{Body: settingsv2.Setting{Name: name, EffectiveBooleanVal: enabled}} + } + return Response{ + StatusCode: 404, + Body: map[string]string{"error_code": "RESOURCE_DOES_NOT_EXIST", "message": fmt.Sprintf("Setting %s does not exist", name)}, + } +} + func (s *FakeWorkspace) WorkspaceSettingPatch(req Request) Response { var setting settingsv2.Setting if err := json.Unmarshal(req.Body, &setting); err != nil { diff --git a/libs/testserver/settings_test.go b/libs/testserver/settings_test.go index e4a5c63d844..781c5326051 100644 --- a/libs/testserver/settings_test.go +++ b/libs/testserver/settings_test.go @@ -3,7 +3,9 @@ package testserver import ( "testing" + "github.com/databricks/databricks-sdk-go/service/settingsv2" "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" ) func patchSetting(workspace *FakeWorkspace, body string) Response { @@ -30,3 +32,35 @@ func TestWorkspaceSettingPatchDisablesPortForwarding(t *testing.T) { assert.Equal(t, 0, response.StatusCode) assert.True(t, workspace.portForwardingToServerlessDisabled()) } + +func TestWorkspaceSettingGetDefaultsPortForwardingToEnabled(t *testing.T) { + workspace := NewFakeWorkspace("http://test.test", "token") + + response := workspace.WorkspaceSettingGet(Request{Vars: map[string]string{"name": portFwdServerlessSetting}}) + + assert.Equal(t, 0, response.StatusCode) + setting, ok := response.Body.(settingsv2.Setting) + require.True(t, ok) + require.NotNil(t, setting.EffectiveBooleanVal) + assert.True(t, setting.EffectiveBooleanVal.Value) +} + +func TestWorkspaceSettingGetReturnsPatchedValue(t *testing.T) { + workspace := NewFakeWorkspace("http://test.test", "token") + patchSetting(workspace, `{"name": "port_fwd_serverless_enabled", "boolean_val": {"value": false}}`) + + response := workspace.WorkspaceSettingGet(Request{Vars: map[string]string{"name": portFwdServerlessSetting}}) + + setting, ok := response.Body.(settingsv2.Setting) + require.True(t, ok) + require.NotNil(t, setting.EffectiveBooleanVal) + assert.False(t, setting.EffectiveBooleanVal.Value) +} + +func TestWorkspaceSettingGetUnknownSetting(t *testing.T) { + workspace := NewFakeWorkspace("http://test.test", "token") + + response := workspace.WorkspaceSettingGet(Request{Vars: map[string]string{"name": "no_such_setting"}}) + + assert.Equal(t, 404, response.StatusCode) +}