From fb572b97fbebe3537055fc281b38668ac52a01a8 Mon Sep 17 00:00:00 2001 From: Anton Nekipelov <226657+anton-107@users.noreply.github.com> Date: Wed, 14 Oct 2026 15:09:10 +0000 Subject: [PATCH 1/2] [SSH] Check port_fwd_serverless_enabled before starting serverless compute A serverless ssh connect now reads the port_fwd_serverless_enabled workspace setting first, in proxy mode too. When an administrator has turned it off, connect stops before it creates the secret scope, uploads binaries or submits the server run, and says which setting is off and who can change it. The driver proxy's gate also depends on an enforcement flag and a port allow-list the CLI cannot see, and the setting is not served through the public API in every environment (production today). A failed read therefore lets the connection go ahead, and the 403 handling from the previous change stays the backstop. The acceptance test for that path now stubs the settings read the way production answers it. DECO-28889 Co-authored-by: Isaac --- .../script | 5 ++- .../test.toml | 13 ++++++ .../out.test.toml | 5 +++ .../output.txt | 10 +++++ .../script | 13 ++++++ .../test.toml | 11 +++++ experimental/ssh/internal/client/client.go | 35 ++++++++++++++++ .../internal/client/client_internal_test.go | 40 +++++++++++++++++++ libs/testserver/handlers.go | 4 ++ libs/testserver/settings.go | 18 +++++++++ libs/testserver/settings_test.go | 34 ++++++++++++++++ 11 files changed, 186 insertions(+), 2 deletions(-) create mode 100644 acceptance/ssh/connect-serverless-port-forwarding-disabled/out.test.toml create mode 100644 acceptance/ssh/connect-serverless-port-forwarding-disabled/output.txt create mode 100644 acceptance/ssh/connect-serverless-port-forwarding-disabled/script create mode 100644 acceptance/ssh/connect-serverless-port-forwarding-disabled/test.toml 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/client.go b/experimental/ssh/internal/client/client.go index 723f6d1e119..e94bdb7e5b6 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,12 @@ func Run(ctx context.Context, client *databricks.WorkspaceClient, opts ClientOpt cmdio.LogString(ctx, fmt.Sprintf("Connecting to %s...", sessionID)) } + if opts.IsServerlessMode() { + 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 +705,26 @@ 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. +func checkPortForwardingToServerless(ctx context.Context, client *databricks.WorkspaceClient) error { + setting, err := client.WorkspaceSettingsV2.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..20dbe2a9bfa 100644 --- a/experimental/ssh/internal/client/client_internal_test.go +++ b/experimental/ssh/internal/client/client_internal_test.go @@ -20,6 +20,7 @@ import ( "github.com/databricks/databricks-sdk-go/service/compute" "github.com/databricks/databricks-sdk-go/service/environments" "github.com/databricks/databricks-sdk-go/service/jobs" + "github.com/databricks/databricks-sdk-go/service/settingsv2" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/mock" "github.com/stretchr/testify/require" @@ -858,3 +859,42 @@ func TestMetadataResponseError(t *testing.T) { }) } } + +func TestCheckPortForwardingToServerless(t *testing.T) { + tests := []struct { + name string + setting *settingsv2.Setting + readErr error + wantErr error + }{ + { + name: "disabled", + setting: &settingsv2.Setting{EffectiveBooleanVal: &settingsv2.BooleanMessage{Value: false}}, + wantErr: errPortForwardingDisabled, + }, + { + name: "enabled", + setting: &settingsv2.Setting{EffectiveBooleanVal: &settingsv2.BooleanMessage{Value: true}}, + }, + { + name: "no effective value", + setting: &settingsv2.Setting{}, + }, + { + // Production does not serve the setting through the public API yet. + name: "read fails", + readErr: errors.New("Setting port_fwd_serverless_enabled is not enabled for public API"), + }, + } + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + ctx := cmdio.MockDiscard(t.Context()) + m := mocks.NewMockWorkspaceClient(t) + m.GetMockWorkspaceSettingsV2API().EXPECT(). + GetPublicWorkspaceSetting(mock.Anything, settingsv2.GetPublicWorkspaceSettingRequest{Name: portFwdServerlessSetting}). + Return(tc.setting, tc.readErr) + + assert.Equal(t, tc.wantErr, checkPortForwardingToServerless(ctx, m.WorkspaceClient)) + }) + } +} 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) +} From a84b403622ff9cabb796b7450b11bb0b0d360afe Mon Sep 17 00:00:00 2001 From: Anton Nekipelov <226657+anton-107@users.noreply.github.com> Date: Thu, 15 Oct 2026 16:01:32 +0000 Subject: [PATCH 2/2] [SSH] Read the serverless port forwarding setting once, without retries The setting check used the default SDK client, which retries 429 and 503 until its retry budget runs out, so a slow settings service could hold up a working connection for minutes. It now reads through the no-retry probe client the AI Gateway check already uses. The check is also skipped when reconnecting with --metadata. The ProxyCommand child runs Run again with the server already running, so the read prevented nothing and was repeated on every IDE reconnect. Co-authored-by: Isaac --- experimental/ssh/internal/client/agentshim.go | 2 +- experimental/ssh/internal/client/client.go | 15 ++++-- .../internal/client/client_internal_test.go | 52 +++++++++++++------ 3 files changed, 49 insertions(+), 20 deletions(-) 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 e94bdb7e5b6..28197d143c5 100644 --- a/experimental/ssh/internal/client/client.go +++ b/experimental/ssh/internal/client/client.go @@ -355,7 +355,9 @@ func Run(ctx context.Context, client *databricks.WorkspaceClient, opts ClientOpt cmdio.LogString(ctx, fmt.Sprintf("Connecting to %s...", sessionID)) } - if opts.IsServerlessMode() { + // 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 } @@ -710,9 +712,16 @@ func metadataResponseError(statusCode int, body []byte) error { // (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. +// (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 { - setting, err := client.WorkspaceSettingsV2.GetPublicWorkspaceSetting(ctx, settingsv2.GetPublicWorkspaceSettingRequest{ + 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 { diff --git a/experimental/ssh/internal/client/client_internal_test.go b/experimental/ssh/internal/client/client_internal_test.go index 20dbe2a9bfa..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" @@ -20,7 +23,6 @@ import ( "github.com/databricks/databricks-sdk-go/service/compute" "github.com/databricks/databricks-sdk-go/service/environments" "github.com/databricks/databricks-sdk-go/service/jobs" - "github.com/databricks/databricks-sdk-go/service/settingsv2" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/mock" "github.com/stretchr/testify/require" @@ -863,38 +865,56 @@ func TestMetadataResponseError(t *testing.T) { func TestCheckPortForwardingToServerless(t *testing.T) { tests := []struct { name string - setting *settingsv2.Setting - readErr error + status int + body string wantErr error }{ { name: "disabled", - setting: &settingsv2.Setting{EffectiveBooleanVal: &settingsv2.BooleanMessage{Value: false}}, + status: http.StatusOK, + body: `{"effective_boolean_val": {"value": false}}`, wantErr: errPortForwardingDisabled, }, { - name: "enabled", - setting: &settingsv2.Setting{EffectiveBooleanVal: &settingsv2.BooleanMessage{Value: true}}, + name: "enabled", + status: http.StatusOK, + body: `{"effective_boolean_val": {"value": true}}`, }, { - name: "no effective value", - setting: &settingsv2.Setting{}, + name: "no effective value", + status: http.StatusOK, + body: `{}`, }, { // Production does not serve the setting through the public API yet. - name: "read fails", - readErr: errors.New("Setting port_fwd_serverless_enabled is not enabled for public API"), + 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) { - ctx := cmdio.MockDiscard(t.Context()) - m := mocks.NewMockWorkspaceClient(t) - m.GetMockWorkspaceSettingsV2API().EXPECT(). - GetPublicWorkspaceSetting(mock.Anything, settingsv2.GetPublicWorkspaceSettingRequest{Name: portFwdServerlessSetting}). - Return(tc.setting, tc.readErr) + 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() - assert.Equal(t, tc.wantErr, checkPortForwardingToServerless(ctx, m.WorkspaceClient)) + ctx := cmdio.MockDiscard(t.Context()) + assert.Equal(t, tc.wantErr, checkPortForwardingToServerless(ctx, newProbeClient(t, srv.URL))) + assert.Equal(t, int32(1), reads.Load()) }) } }