From e385a448a29b36d78f091cd3881409c73d88c668 Mon Sep 17 00:00:00 2001 From: "jae.cho" Date: Wed, 4 Nov 2026 23:17:31 +0000 Subject: [PATCH 1/2] Add `databricks sandbox port-forward` Forward local ports to apps in a sandbox: each -L LOCAL:REMOTE serves http://127.0.0.1:LOCAL and relays HTTP and WebSocket traffic through the workspace's /sandbox-proxy/sandboxes//ports// route with the caller's credentials. Forwards accept only localhost-addressed requests and refuse dot-segment paths, so neither a DNS-rebinding page nor a crafted path can use those credentials outside the port route. Co-authored-by: Isaac --- .../invalid-forward/out.test.toml | 1 + .../port-forward/invalid-forward/output.txt | 3 + .../port-forward/invalid-forward/script | 6 + .../port-forward/no-default/out.test.toml | 1 + .../port-forward/no-default/output.txt | 1 + .../sandbox/port-forward/no-default/script | 2 + .../port-forward/not-found/out.test.toml | 1 + .../sandbox/port-forward/not-found/output.txt | 1 + .../cmd/sandbox/port-forward/not-found/script | 1 + .../sandbox/port-forward/not-found/test.toml | 4 + cmd/sandbox/portforward.go | 332 +++++++++++++++++ cmd/sandbox/portforward_test.go | 352 ++++++++++++++++++ cmd/sandbox/sandbox.go | 2 + 13 files changed, 707 insertions(+) create mode 100644 acceptance/cmd/sandbox/port-forward/invalid-forward/out.test.toml create mode 100644 acceptance/cmd/sandbox/port-forward/invalid-forward/output.txt create mode 100644 acceptance/cmd/sandbox/port-forward/invalid-forward/script create mode 100644 acceptance/cmd/sandbox/port-forward/no-default/out.test.toml create mode 100644 acceptance/cmd/sandbox/port-forward/no-default/output.txt create mode 100644 acceptance/cmd/sandbox/port-forward/no-default/script create mode 100644 acceptance/cmd/sandbox/port-forward/not-found/out.test.toml create mode 100644 acceptance/cmd/sandbox/port-forward/not-found/output.txt create mode 100644 acceptance/cmd/sandbox/port-forward/not-found/script create mode 100644 acceptance/cmd/sandbox/port-forward/not-found/test.toml create mode 100644 cmd/sandbox/portforward.go create mode 100644 cmd/sandbox/portforward_test.go diff --git a/acceptance/cmd/sandbox/port-forward/invalid-forward/out.test.toml b/acceptance/cmd/sandbox/port-forward/invalid-forward/out.test.toml new file mode 100644 index 00000000000..18b1a88417e --- /dev/null +++ b/acceptance/cmd/sandbox/port-forward/invalid-forward/out.test.toml @@ -0,0 +1 @@ +Cloud = false diff --git a/acceptance/cmd/sandbox/port-forward/invalid-forward/output.txt b/acceptance/cmd/sandbox/port-forward/invalid-forward/output.txt new file mode 100644 index 00000000000..accf4593bef --- /dev/null +++ b/acceptance/cmd/sandbox/port-forward/invalid-forward/output.txt @@ -0,0 +1,3 @@ +Error: invalid -L "8080": expected LOCAL:REMOTE ports, e.g. 8080:10080 +Error: invalid -L "8080:0": expected LOCAL:REMOTE ports, e.g. 8080:10080 +Error: required flag(s) "forward" not set diff --git a/acceptance/cmd/sandbox/port-forward/invalid-forward/script b/acceptance/cmd/sandbox/port-forward/invalid-forward/script new file mode 100644 index 00000000000..681604533c9 --- /dev/null +++ b/acceptance/cmd/sandbox/port-forward/invalid-forward/script @@ -0,0 +1,6 @@ +# Malformed -L specs are rejected before any API call, so no [[Server]] stub is needed. +musterr $CLI sandbox port-forward happy-panda-1234 -L 8080 +musterr $CLI sandbox port-forward happy-panda-1234 -L 8080:0 + +# At least one -L is required. +musterr $CLI sandbox port-forward happy-panda-1234 diff --git a/acceptance/cmd/sandbox/port-forward/no-default/out.test.toml b/acceptance/cmd/sandbox/port-forward/no-default/out.test.toml new file mode 100644 index 00000000000..18b1a88417e --- /dev/null +++ b/acceptance/cmd/sandbox/port-forward/no-default/out.test.toml @@ -0,0 +1 @@ +Cloud = false diff --git a/acceptance/cmd/sandbox/port-forward/no-default/output.txt b/acceptance/cmd/sandbox/port-forward/no-default/output.txt new file mode 100644 index 00000000000..5905d0f9bdd --- /dev/null +++ b/acceptance/cmd/sandbox/port-forward/no-default/output.txt @@ -0,0 +1 @@ +Error: no default sandbox configured — run `databricks sandbox create` to provision one, or `databricks sandbox default ` to point at an existing sandbox diff --git a/acceptance/cmd/sandbox/port-forward/no-default/script b/acceptance/cmd/sandbox/port-forward/no-default/script new file mode 100644 index 00000000000..b256b754414 --- /dev/null +++ b/acceptance/cmd/sandbox/port-forward/no-default/script @@ -0,0 +1,2 @@ +# Without a sandbox ID the default sandbox is used; the fresh home has none. +musterr $CLI sandbox port-forward -L 0:10080 diff --git a/acceptance/cmd/sandbox/port-forward/not-found/out.test.toml b/acceptance/cmd/sandbox/port-forward/not-found/out.test.toml new file mode 100644 index 00000000000..18b1a88417e --- /dev/null +++ b/acceptance/cmd/sandbox/port-forward/not-found/out.test.toml @@ -0,0 +1 @@ +Cloud = false diff --git a/acceptance/cmd/sandbox/port-forward/not-found/output.txt b/acceptance/cmd/sandbox/port-forward/not-found/output.txt new file mode 100644 index 00000000000..235a49bc669 --- /dev/null +++ b/acceptance/cmd/sandbox/port-forward/not-found/output.txt @@ -0,0 +1 @@ +Error: no sandbox named "no-such-box" — `databricks sandbox list` shows available IDs diff --git a/acceptance/cmd/sandbox/port-forward/not-found/script b/acceptance/cmd/sandbox/port-forward/not-found/script new file mode 100644 index 00000000000..5988fa053bd --- /dev/null +++ b/acceptance/cmd/sandbox/port-forward/not-found/script @@ -0,0 +1 @@ +musterr $CLI sandbox port-forward no-such-box -L 0:10080 diff --git a/acceptance/cmd/sandbox/port-forward/not-found/test.toml b/acceptance/cmd/sandbox/port-forward/not-found/test.toml new file mode 100644 index 00000000000..3d2b831b415 --- /dev/null +++ b/acceptance/cmd/sandbox/port-forward/not-found/test.toml @@ -0,0 +1,4 @@ +[[Server]] +Pattern = "GET /api/2.0/lakebox/sandboxes/no-such-box" +Response.StatusCode = 404 +Response.Body = '{"error_code": "NOT_FOUND", "message": "Sandbox not found"}' diff --git a/cmd/sandbox/portforward.go b/cmd/sandbox/portforward.go new file mode 100644 index 00000000000..942f3df1c0c --- /dev/null +++ b/cmd/sandbox/portforward.go @@ -0,0 +1,332 @@ +package sandbox + +import ( + "context" + "crypto/tls" + "errors" + "fmt" + "net" + "net/http" + "net/http/httputil" + "net/url" + "os" + "os/signal" + "strconv" + "strings" + "syscall" + "time" + + "github.com/databricks/cli/cmd/root" + "github.com/databricks/cli/libs/auth" + "github.com/databricks/cli/libs/cmdctx" + "github.com/databricks/cli/libs/cmdio" + "github.com/databricks/cli/libs/log" + "github.com/databricks/databricks-sdk-go/apierr" + "github.com/databricks/databricks-sdk-go/config" + "github.com/spf13/cobra" + "golang.org/x/sync/errgroup" +) + +// portForwardRoute is the workspace route that relays a request to an app +// listening on a sandbox port: //ports//. +const portForwardRoute = "/sandbox-proxy/sandboxes/" + +// portForwardBindHost keeps forwards reachable only from this machine, since +// every request through them carries the caller's credentials. +const portForwardBindHost = "127.0.0.1" + +// portForwardReadHeaderTimeout bounds how long a local client may take to send +// request headers. Bodies and upgraded connections are not bounded. +const portForwardReadHeaderTimeout = time.Minute + +// databricksHeaderPrefix marks the platform's auth and routing headers. Only +// the CLI's own values may reach the workspace. +const databricksHeaderPrefix = "X-Databricks-" + +// portForward is one parsed -L LOCAL:REMOTE spec. +type portForward struct { + localPort int + remotePort int +} + +// portForwardTarget is where and how forwarded requests are sent. +type portForwardTarget struct { + workspace *url.URL + sandboxID string + authenticate func(*http.Request) error + headers map[string]string + transport http.RoundTripper +} + +func newPortForwardCommand() *cobra.Command { + var specs []string + + cmd := &cobra.Command{ + Use: "port-forward [sandbox-id] -L :...", + Short: "Forward local ports to apps in a Sandbox environment", + Long: `Forward local ports to apps running in a Sandbox environment. + +Each -L LOCAL:REMOTE serves http://127.0.0.1:LOCAL and relays HTTP and +WebSocket traffic to an app listening on 0.0.0.0:REMOTE in the sandbox, +authenticated as you. Sandboxes accept remote ports from 10000 to 65535 and +refuse others. Pass 0 as the local port to pick a free one. Runs until +interrupted. + +Examples: + databricks sandbox port-forward -L 8080:10080 # default sandbox + databricks sandbox port-forward happy-panda-1234 -L 8080:10080 -L 8081:10081`, + Args: cobra.MaximumNArgs(1), + PreRunE: root.MustWorkspaceClient, + ValidArgsFunction: completeSandboxIDs, + RunE: func(cmd *cobra.Command, args []string) error { + forwards, err := parseForwards(specs) + if err != nil { + return err + } + + ctx := cmd.Context() + w := cmdctx.WorkspaceClient(ctx) + + profile := w.Config.Profile + if profile == "" { + profile = w.Config.Host + } + + api, err := newSandboxAPI(w) + if err != nil { + return err + } + sandboxID, err := resolvePortForwardSandbox(ctx, api, profile, args) + if err != nil { + return err + } + + workspace, err := url.Parse(w.Config.Host) + if err != nil { + return fmt.Errorf("failed to parse workspace host %q: %w", w.Config.Host, err) + } + transport, err := newPortForwardTransport(w.Config) + if err != nil { + return err + } + target := portForwardTarget{ + workspace: workspace, + sandboxID: sandboxID, + authenticate: w.Config.Authenticate, + headers: auth.WorkspaceIDHeaders(w.Config), + transport: transport, + } + + ctx, stop := signal.NotifyContext(ctx, os.Interrupt, syscall.SIGTERM) + defer stop() + return servePortForwards(ctx, target, forwards) + }, + } + + cmd.Flags().StringArrayVarP(&specs, "forward", "L", nil, "Forward LOCAL:REMOTE ports, e.g. 8080:10080 (repeatable)") + _ = cmd.MarkFlagRequired("forward") + + return cmd +} + +// parseForwards parses -L LOCAL:REMOTE specs. LOCAL may be 0 to pick a free +// port. The sandbox enforces which remote ports it accepts. +func parseForwards(specs []string) ([]portForward, error) { + forwards := make([]portForward, 0, len(specs)) + for _, spec := range specs { + local, remote, found := strings.Cut(spec, ":") + localPort, localErr := strconv.Atoi(local) + remotePort, remoteErr := strconv.Atoi(remote) + if !found || localErr != nil || remoteErr != nil || localPort < 0 || localPort > 65535 || remotePort < 1 || remotePort > 65535 { + return nil, fmt.Errorf("invalid -L %q: expected LOCAL:REMOTE ports, e.g. 8080:10080", spec) + } + forwards = append(forwards, portForward{localPort: localPort, remotePort: remotePort}) + } + return forwards, nil +} + +// resolvePortForwardSandbox returns the sandbox named in args, or the default +// sandbox when args is empty. Like `sandbox ssh`, it clears a default that no +// longer exists and only warns on other lookup errors, since the forward may +// still work. +func resolvePortForwardSandbox(ctx context.Context, api *sandboxAPI, profile string, args []string) (string, error) { + var sandboxID string + if len(args) == 0 { + if sandboxID = getDefault(ctx, profile); sandboxID == "" { + return "", errors.New("no default sandbox configured — run `databricks sandbox create` to provision one, or `databricks sandbox default ` to point at an existing sandbox") + } + } else { + resolved, err := resolveLocalID(ctx, profile, args[0]) + if err != nil { + return "", err + } + sandboxID = resolved + } + + // A "." or ".." ID would move every forwarded request off the port-forward route. + if hasDotSegment("/" + url.PathEscape(sandboxID)) { + return "", fmt.Errorf("invalid sandbox ID %q", sandboxID) + } + + if _, err := api.get(ctx, sandboxID); err != nil { + switch { + case !errors.Is(err, apierr.ErrNotFound): + warn(ctx, fmt.Sprintf("Could not validate sandbox %s: %v", sandboxID, err)) + case len(args) == 0: + _ = clearDefault(ctx, profile) + return "", fmt.Errorf("saved default %q no longer exists (cleared) — run `databricks sandbox create` to provision a new one, or `databricks sandbox default ` to point at an existing sandbox", sandboxID) + default: + return "", fmt.Errorf("no sandbox named %q — `databricks sandbox list` shows available IDs", sandboxID) + } + } + return sandboxID, nil +} + +// newPortForwardTransport returns the workspace config's transport, or the +// default one with the config's TLS verification setting. +func newPortForwardTransport(cfg *config.Config) (http.RoundTripper, error) { + clientCfg, err := config.HTTPClientConfigFromConfig(cfg) + if err != nil { + return nil, fmt.Errorf("failed to configure port forward transport: %w", err) + } + if clientCfg.Transport != nil { + return clientCfg.Transport, nil + } + transport := http.DefaultTransport.(*http.Transport).Clone() + if clientCfg.InsecureSkipVerify { + transport.TLSClientConfig = &tls.Config{InsecureSkipVerify: true} + } + return transport, nil +} + +// servePortForwards listens on every forward's local port before serving any, +// so a port that is already in use fails the command up front. +func servePortForwards(ctx context.Context, target portForwardTarget, forwards []portForward) error { + var lc net.ListenConfig + listeners := make([]net.Listener, 0, len(forwards)) + for _, f := range forwards { + ln, err := lc.Listen(ctx, "tcp", net.JoinHostPort(portForwardBindHost, strconv.Itoa(f.localPort))) + if err != nil { + for _, open := range listeners { + open.Close() + } + return fmt.Errorf("failed to listen on local port %d: %w", f.localPort, err) + } + listeners = append(listeners, ln) + } + + g, ctx := errgroup.WithContext(ctx) + for i, f := range forwards { + srv := &http.Server{ + Handler: newPortForwardHandler(target, f.remotePort), + ReadHeaderTimeout: portForwardReadHeaderTimeout, + } + cmdio.LogString(ctx, fmt.Sprintf("Forwarding http://%s -> %s:%d", listeners[i].Addr(), target.sandboxID, f.remotePort)) + g.Go(func() error { + if err := srv.Serve(listeners[i]); !errors.Is(err, http.ErrServerClosed) { + return fmt.Errorf("port forward on %s failed: %w", listeners[i].Addr(), err) + } + return nil + }) + // An interactive interrupt stops at once: WebSockets and streams never + // go idle, so a graceful Shutdown would only delay the exit. + g.Go(func() error { + <-ctx.Done() + return srv.Close() + }) + } + cmdio.LogString(ctx, "Press Ctrl-C to stop.") + return g.Wait() +} + +// newPortForwardHandler relays each request to remotePort in the sandbox +// through the workspace's port-forward route. +func newPortForwardHandler(target portForwardTarget, remotePort int) http.Handler { + base := strings.TrimSuffix(target.workspace.Path, "/") + portForwardRoute + port := "/ports/" + strconv.Itoa(remotePort) + proxy := &httputil.ReverseProxy{ + Rewrite: func(pr *httputil.ProxyRequest) { + pr.Out.URL.Scheme = target.workspace.Scheme + pr.Out.URL.Host = target.workspace.Host + pr.Out.URL.Path = base + target.sandboxID + port + pr.In.URL.Path + pr.Out.URL.RawPath = base + url.PathEscape(target.sandboxID) + port + pr.In.URL.EscapedPath() + pr.Out.Host = "" + pr.Out.Header.Del("Authorization") + for name := range pr.Out.Header { + if strings.HasPrefix(name, databricksHeaderPrefix) { + pr.Out.Header.Del(name) + } + } + for name, value := range target.headers { + pr.Out.Header.Set(name, value) + } + }, + Transport: &authenticatingTransport{authenticate: target.authenticate, base: target.transport}, + ErrorHandler: func(w http.ResponseWriter, r *http.Request, err error) { + log.Debugf(r.Context(), "port forward to %s:%d failed: %v", target.sandboxID, remotePort, err) + http.Error(w, "port forward failed: "+err.Error(), http.StatusBadGateway) + }, + } + return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + switch { + case !isLoopbackHost(r.Host): + http.Error(w, "port forward only accepts requests addressed to localhost", http.StatusForbidden) + case hasDotSegment(r.URL.EscapedPath()): + http.Error(w, "request path must not contain . or .. segments", http.StatusBadRequest) + default: + proxy.ServeHTTP(w, r) + } + }) +} + +// isLoopbackHost reports whether a Host header names this machine. A +// DNS-rebinding page in a browser sends its own hostname, so refusing other +// hosts keeps it from acting with the caller's credentials. +func isLoopbackHost(hostport string) bool { + host, _, err := net.SplitHostPort(hostport) + if err != nil { + host = hostport + } + if strings.EqualFold(host, "localhost") { + return true + } + ip := net.ParseIP(host) + return ip != nil && ip.IsLoopback() +} + +// hasDotSegment reports whether an escaped URL path contains a "." or ".." +// segment, including percent-encoded and backslash-separated forms. The +// workspace would resolve such a path out of the port-forward route and onto +// another API, with the caller's credentials attached. +func hasDotSegment(escapedPath string) bool { + decoded, err := url.PathUnescape(escapedPath) + if err != nil { + return true + } + for _, segment := range strings.FieldsFunc(decoded, isPathSeparator) { + if segment == "." || segment == ".." { + return true + } + } + return false +} + +func isPathSeparator(r rune) bool { + return r == '/' || r == '\\' +} + +// authenticatingTransport adds the CLI's credentials to each outgoing request, +// so OAuth tokens refresh during a long-running forward. +type authenticatingTransport struct { + authenticate func(*http.Request) error + base http.RoundTripper +} + +func (t *authenticatingTransport) RoundTrip(req *http.Request) (*http.Response, error) { + // A RoundTripper must not modify the request it is given. + req = req.Clone(req.Context()) + if err := t.authenticate(req); err != nil { + return nil, fmt.Errorf("failed to authenticate with the workspace: %w", err) + } + return t.base.RoundTrip(req) +} diff --git a/cmd/sandbox/portforward_test.go b/cmd/sandbox/portforward_test.go new file mode 100644 index 00000000000..be9b491d2ef --- /dev/null +++ b/cmd/sandbox/portforward_test.go @@ -0,0 +1,352 @@ +package sandbox + +import ( + "bufio" + "context" + "encoding/json" + "errors" + "io" + "net" + "net/http" + "net/http/httptest" + "net/url" + "strconv" + "strings" + "sync/atomic" + "testing" + + "github.com/databricks/cli/libs/cmdio" + "github.com/databricks/databricks-sdk-go" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +const testSandboxID = "happy-panda-1234" + +func TestParseForwards(t *testing.T) { + forwards, err := parseForwards([]string{"8080:10080", "0:10081"}) + require.NoError(t, err) + assert.Equal(t, []portForward{{localPort: 8080, remotePort: 10080}, {localPort: 0, remotePort: 10081}}, forwards) +} + +func TestParseForwardsRejectsInvalidSpecs(t *testing.T) { + for _, spec := range []string{"8080", "8080:", ":10080", "a:10080", "8080:b", "-1:10080", "65536:10080", "8080:0", "8080:65536", "8080:10080:1"} { + t.Run(spec, func(t *testing.T) { + _, err := parseForwards([]string{spec}) + assert.EqualError(t, err, `invalid -L "`+spec+`": expected LOCAL:REMOTE ports, e.g. 8080:10080`) + }) + } +} + +func TestHasDotSegment(t *testing.T) { + tests := []struct { + path string + want bool + }{ + {"/", false}, + {"/hello/a%2Fb", false}, + {"/a.b/.c/d..", false}, + {"/...", false}, + {"/../api/2.0/clusters/list", true}, + {"/a/./b", true}, + {"/a/..", true}, + {"/%2e%2e/api", true}, + {"/%2E/api", true}, + {"/a/..%5c..%5capi", true}, + {"/a%2F..%2Fapi", true}, + } + for _, tt := range tests { + t.Run(tt.path, func(t *testing.T) { + assert.Equal(t, tt.want, hasDotSegment(tt.path)) + }) + } +} + +// fakeAuth stands in for the SDK's config.Authenticate. +func fakeAuth(r *http.Request) error { + r.Header.Set("Authorization", "Bearer cli-token") + return nil +} + +// startForward serves a port-forward handler for remote port 10080 in front of upstream. +func startForward(t *testing.T, upstream *httptest.Server, authenticate func(*http.Request) error) *httptest.Server { + workspace, err := url.Parse(upstream.URL) + require.NoError(t, err) + target := portForwardTarget{ + workspace: workspace, + sandboxID: testSandboxID, + authenticate: authenticate, + headers: map[string]string{"X-Databricks-Workspace-Id": "123"}, + transport: http.DefaultTransport, + } + forward := httptest.NewServer(newPortForwardHandler(target, 10080)) + t.Cleanup(forward.Close) + return forward +} + +// rawRequest sends request verbatim, so paths reach the handler without +// client-side cleaning, and returns the response status. The reader and conn +// stay usable after a 101, whose response has no body. +func rawRequest(t *testing.T, addr, request string) (int, *bufio.Reader, net.Conn) { + conn, err := net.Dial("tcp", addr) + require.NoError(t, err) + t.Cleanup(func() { conn.Close() }) + _, err = io.WriteString(conn, request) + require.NoError(t, err) + reader := bufio.NewReader(conn) + resp, err := http.ReadResponse(reader, nil) + require.NoError(t, err) + require.NoError(t, resp.Body.Close()) + return resp.StatusCode, reader, conn +} + +func TestPortForwardHandlerRewritesRequest(t *testing.T) { + var got *http.Request + var gotBody string + upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + body, _ := io.ReadAll(r.Body) + got, gotBody = r, string(body) + w.Header().Set("Set-Cookie", "app=1") + w.WriteHeader(http.StatusCreated) + _, _ = io.WriteString(w, "from app") + })) + defer upstream.Close() + forward := startForward(t, upstream, fakeAuth) + + req, err := http.NewRequestWithContext(t.Context(), http.MethodPost, forward.URL+"/hello/a%2Fb?x=1", strings.NewReader("payload")) + require.NoError(t, err) + req.Header.Set("Authorization", "Bearer local-secret") + req.Header.Set("Cookie", "sid=abc") + req.Header.Set("X-Databricks-Org-Id", "999") + req.Header.Set("X-Databricks-Workspace-Id", "999") + resp, err := http.DefaultClient.Do(req) + require.NoError(t, err) + defer resp.Body.Close() + body, err := io.ReadAll(resp.Body) + require.NoError(t, err) + + assert.Equal(t, http.StatusCreated, resp.StatusCode) + assert.Equal(t, "app=1", resp.Header.Get("Set-Cookie")) + assert.Equal(t, "from app", string(body)) + + require.NotNil(t, got) + assert.Equal(t, http.MethodPost, got.Method) + assert.Equal(t, "/sandbox-proxy/sandboxes/"+testSandboxID+"/ports/10080/hello/a%2Fb", got.URL.EscapedPath()) + assert.Equal(t, "x=1", got.URL.RawQuery) + assert.Equal(t, strings.TrimPrefix(upstream.URL, "http://"), got.Host) + assert.Equal(t, "Bearer cli-token", got.Header.Get("Authorization")) + assert.Equal(t, "123", got.Header.Get("X-Databricks-Workspace-Id")) + assert.Empty(t, got.Header.Get("X-Databricks-Org-Id")) + assert.Equal(t, "sid=abc", got.Header.Get("Cookie")) + assert.Empty(t, got.Header.Get("X-Forwarded-For")) + assert.Equal(t, "payload", gotBody) +} + +func TestPortForwardHandlerRejectsDotSegments(t *testing.T) { + var calls atomic.Int32 + upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + calls.Add(1) + })) + defer upstream.Close() + forward := startForward(t, upstream, fakeAuth) + addr := strings.TrimPrefix(forward.URL, "http://") + + for _, path := range []string{"/../../api/2.0/clusters/list", "/%2e%2e/api/2.0/clusters/list", "/a/..%5c..%5capi"} { + t.Run(path, func(t *testing.T) { + status, _, _ := rawRequest(t, addr, "GET "+path+" HTTP/1.1\r\nHost: "+addr+"\r\nConnection: close\r\n\r\n") + assert.Equal(t, http.StatusBadRequest, status) + }) + } + assert.Zero(t, calls.Load()) +} + +func TestPortForwardHandlerRejectsForeignHost(t *testing.T) { + var calls atomic.Int32 + upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + calls.Add(1) + })) + defer upstream.Close() + forward := startForward(t, upstream, fakeAuth) + addr := strings.TrimPrefix(forward.URL, "http://") + + status, _, _ := rawRequest(t, addr, "GET / HTTP/1.1\r\nHost: rebound.example:8080\r\nConnection: close\r\n\r\n") + assert.Equal(t, http.StatusForbidden, status) + assert.Zero(t, calls.Load()) +} + +func TestIsLoopbackHost(t *testing.T) { + tests := []struct { + host string + want bool + }{ + {"127.0.0.1:8080", true}, + {"localhost:8080", true}, + {"LOCALHOST:8080", true}, + {"[::1]:8080", true}, + {"localhost", true}, + {"rebound.example:8080", false}, + {"10.0.0.1:8080", false}, + {"", false}, + } + for _, tt := range tests { + t.Run(tt.host, func(t *testing.T) { + assert.Equal(t, tt.want, isLoopbackHost(tt.host)) + }) + } +} + +func TestPortForwardHandlerRelaysUpgrade(t *testing.T) { + upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.Header.Get("Upgrade") != "websocket" { + http.Error(w, "expected an upgrade", http.StatusBadRequest) + return + } + conn, buf, err := http.NewResponseController(w).Hijack() + if err != nil { + return + } + defer conn.Close() + _, _ = buf.WriteString("HTTP/1.1 101 Switching Protocols\r\nUpgrade: websocket\r\nConnection: Upgrade\r\n\r\n") + _ = buf.Flush() + _, _ = io.Copy(conn, buf) + })) + defer upstream.Close() + forward := startForward(t, upstream, fakeAuth) + addr := strings.TrimPrefix(forward.URL, "http://") + + status, reader, conn := rawRequest(t, addr, "GET /ws HTTP/1.1\r\nHost: "+addr+"\r\nUpgrade: websocket\r\nConnection: Upgrade\r\n\r\n") + require.Equal(t, http.StatusSwitchingProtocols, status) + + _, err := io.WriteString(conn, "ping") + require.NoError(t, err) + echo := make([]byte, 4) + _, err = io.ReadFull(reader, echo) + require.NoError(t, err) + assert.Equal(t, "ping", string(echo)) +} + +func TestPortForwardHandlerReportsAuthFailure(t *testing.T) { + var calls atomic.Int32 + upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + calls.Add(1) + })) + defer upstream.Close() + forward := startForward(t, upstream, func(*http.Request) error { return errors.New("token expired") }) + + resp, err := http.Get(forward.URL + "/") + require.NoError(t, err) + defer resp.Body.Close() + body, err := io.ReadAll(resp.Body) + require.NoError(t, err) + + assert.Equal(t, http.StatusBadGateway, resp.StatusCode) + assert.Contains(t, string(body), "failed to authenticate with the workspace: token expired") + assert.Zero(t, calls.Load()) +} + +func TestServePortForwardsFailsWhenLocalPortInUse(t *testing.T) { + taken, err := net.Listen("tcp", net.JoinHostPort(portForwardBindHost, "0")) + require.NoError(t, err) + defer taken.Close() + port := taken.Addr().(*net.TCPAddr).Port + + err = servePortForwards(cmdio.MockDiscard(t.Context()), testTarget(), []portForward{{localPort: 0, remotePort: 10080}, {localPort: port, remotePort: 10081}}) + assert.ErrorContains(t, err, "failed to listen on local port "+strconv.Itoa(port)) +} + +func TestServePortForwardsStopsWhenContextIsDone(t *testing.T) { + ctx, cancel := context.WithCancel(cmdio.MockDiscard(t.Context())) + cancel() + + err := servePortForwards(ctx, testTarget(), []portForward{{localPort: 0, remotePort: 10080}}) + assert.NoError(t, err) +} + +func testTarget() portForwardTarget { + return portForwardTarget{ + workspace: &url.URL{Scheme: "https", Host: "workspace.test"}, + sandboxID: testSandboxID, + authenticate: fakeAuth, + transport: http.DefaultTransport, + } +} + +// newLookupAPI returns a sandbox API whose GET for testSandboxID answers with status. +func newLookupAPI(t *testing.T, status int) *sandboxAPI { + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.URL.Path == "/.well-known/databricks-config" { + assert.NoError(t, json.NewEncoder(w).Encode(map[string]string{ + "oidc_endpoint": "https://workspace.example.test/oidc", + })) + return + } + assert.Equal(t, sandboxPath(testSandboxID), r.URL.Path) + w.WriteHeader(status) + if status == http.StatusOK { + assert.NoError(t, json.NewEncoder(w).Encode(sandboxEntry{SandboxID: testSandboxID, Status: "Running"})) + return + } + _, _ = io.WriteString(w, `{"error_code": "ERROR", "message": "lookup failed"}`) + })) + t.Cleanup(server.Close) + + workspaceClient, err := databricks.NewWorkspaceClient(&databricks.Config{Host: server.URL, Token: "test-token"}) + require.NoError(t, err) + api, err := newSandboxAPI(workspaceClient) + require.NoError(t, err) + return api +} + +func TestResolvePortForwardSandboxUsesDefault(t *testing.T) { + ctx, _ := stateCtx(t) + ctx = cmdio.MockDiscard(ctx) + require.NoError(t, setDefault(ctx, "profile", testSandboxID)) + + got, err := resolvePortForwardSandbox(ctx, newLookupAPI(t, http.StatusOK), "profile", nil) + require.NoError(t, err) + assert.Equal(t, testSandboxID, got) +} + +func TestResolvePortForwardSandboxClearsMissingDefault(t *testing.T) { + ctx, _ := stateCtx(t) + ctx = cmdio.MockDiscard(ctx) + require.NoError(t, setDefault(ctx, "profile", testSandboxID)) + + _, err := resolvePortForwardSandbox(ctx, newLookupAPI(t, http.StatusNotFound), "profile", nil) + assert.ErrorContains(t, err, `saved default "`+testSandboxID+`" no longer exists (cleared)`) + assert.Empty(t, getDefault(ctx, "profile")) +} + +func TestResolvePortForwardSandboxRejectsUnknownID(t *testing.T) { + ctx, _ := stateCtx(t) + ctx = cmdio.MockDiscard(ctx) + + _, err := resolvePortForwardSandbox(ctx, newLookupAPI(t, http.StatusNotFound), "profile", []string{testSandboxID}) + assert.ErrorContains(t, err, `no sandbox named "`+testSandboxID+`"`) +} + +func TestResolvePortForwardSandboxContinuesOnOtherLookupErrors(t *testing.T) { + ctx, _ := stateCtx(t) + ctx = cmdio.MockDiscard(ctx) + + got, err := resolvePortForwardSandbox(ctx, newLookupAPI(t, http.StatusForbidden), "profile", []string{testSandboxID}) + require.NoError(t, err) + assert.Equal(t, testSandboxID, got) +} + +func TestResolvePortForwardSandboxRejectsDotSegmentIDs(t *testing.T) { + ctx, _ := stateCtx(t) + for _, id := range []string{".", ".."} { + t.Run(id, func(t *testing.T) { + _, err := resolvePortForwardSandbox(ctx, nil, "profile", []string{id}) + assert.EqualError(t, err, `invalid sandbox ID "`+id+`"`) + }) + } +} + +func TestResolvePortForwardSandboxRequiresDefault(t *testing.T) { + ctx, _ := stateCtx(t) + + _, err := resolvePortForwardSandbox(ctx, nil, "profile", nil) + assert.ErrorContains(t, err, "no default sandbox configured") +} diff --git a/cmd/sandbox/sandbox.go b/cmd/sandbox/sandbox.go index 39ff208129b..b03ebf49732 100644 --- a/cmd/sandbox/sandbox.go +++ b/cmd/sandbox/sandbox.go @@ -30,6 +30,7 @@ Common workflows: databricks sandbox create # create a new sandbox databricks sandbox delete my-project # delete a sandbox databricks sandbox status my-project # show sandbox status + databricks sandbox port-forward -L 8080:10080 # open sandbox port 10080 at localhost:8080 `, } @@ -37,6 +38,7 @@ Common workflows: cmd.AddCommand(newSSHKeyCommand()) cmd.AddCommand(newSetDefaultCommand()) cmd.AddCommand(newSSHCommand()) + cmd.AddCommand(newPortForwardCommand()) cmd.AddCommand(newListCommand()) cmd.AddCommand(newCreateCommand()) cmd.AddCommand(newDeleteCommand()) From 024f57d965a14b3a87de50db580c88f0340ef955 Mon Sep 17 00:00:00 2001 From: "jae.cho" Date: Thu, 5 Nov 2026 00:21:24 +0000 Subject: [PATCH 2/2] sandbox port-forward: refuse cross-site requests, start stopped sandboxes Refuse browser requests from other sites (Origin, Sec-Fetch-Site) so a web page can't use a forward, including over WebSockets. Reject remote ports outside 10000-65535 at parse time, start a stopped sandbox before forwarding like `sandbox ssh`, also serve the IPv6 loopback, and keep net/http's error reports in the debug log. Co-authored-by: Isaac --- .../port-forward/invalid-forward/output.txt | 2 +- .../port-forward/invalid-forward/script | 4 +- cmd/sandbox/portforward.go | 127 +++++++++++++----- cmd/sandbox/portforward_test.go | 69 ++++++++-- 4 files changed, 156 insertions(+), 46 deletions(-) diff --git a/acceptance/cmd/sandbox/port-forward/invalid-forward/output.txt b/acceptance/cmd/sandbox/port-forward/invalid-forward/output.txt index accf4593bef..067a76be641 100644 --- a/acceptance/cmd/sandbox/port-forward/invalid-forward/output.txt +++ b/acceptance/cmd/sandbox/port-forward/invalid-forward/output.txt @@ -1,3 +1,3 @@ Error: invalid -L "8080": expected LOCAL:REMOTE ports, e.g. 8080:10080 -Error: invalid -L "8080:0": expected LOCAL:REMOTE ports, e.g. 8080:10080 +Error: invalid -L "8080:8080": remote port must be between 10000 and 65535 Error: required flag(s) "forward" not set diff --git a/acceptance/cmd/sandbox/port-forward/invalid-forward/script b/acceptance/cmd/sandbox/port-forward/invalid-forward/script index 681604533c9..1f72d2890e2 100644 --- a/acceptance/cmd/sandbox/port-forward/invalid-forward/script +++ b/acceptance/cmd/sandbox/port-forward/invalid-forward/script @@ -1,6 +1,8 @@ # Malformed -L specs are rejected before any API call, so no [[Server]] stub is needed. musterr $CLI sandbox port-forward happy-panda-1234 -L 8080 -musterr $CLI sandbox port-forward happy-panda-1234 -L 8080:0 + +# Remote ports outside 10000-65535 are refused up front. +musterr $CLI sandbox port-forward happy-panda-1234 -L 8080:8080 # At least one -L is required. musterr $CLI sandbox port-forward happy-panda-1234 diff --git a/cmd/sandbox/portforward.go b/cmd/sandbox/portforward.go index 942f3df1c0c..984a316e7b4 100644 --- a/cmd/sandbox/portforward.go +++ b/cmd/sandbox/portforward.go @@ -5,6 +5,7 @@ import ( "crypto/tls" "errors" "fmt" + "log/slog" "net" "net/http" "net/http/httputil" @@ -31,9 +32,19 @@ import ( // listening on a sandbox port: //ports//. const portForwardRoute = "/sandbox-proxy/sandboxes/" -// portForwardBindHost keeps forwards reachable only from this machine, since -// every request through them carries the caller's credentials. -const portForwardBindHost = "127.0.0.1" +// Forwards listen only on loopback, since every request through them carries +// the caller's credentials. Some clients (e.g. Node 17+) resolve localhost to +// ::1 first, so the IPv6 loopback is served too when it is available. +const ( + portForwardBindHost = "127.0.0.1" + portForwardBindHostIPv6 = "::1" +) + +// The sandbox only forwards to remote ports in this range. +const ( + minRemotePort = 10000 + maxRemotePort = 65535 +) // portForwardReadHeaderTimeout bounds how long a local client may take to send // request headers. Bodies and upgraded connections are not bounded. @@ -68,9 +79,9 @@ func newPortForwardCommand() *cobra.Command { Each -L LOCAL:REMOTE serves http://127.0.0.1:LOCAL and relays HTTP and WebSocket traffic to an app listening on 0.0.0.0:REMOTE in the sandbox, -authenticated as you. Sandboxes accept remote ports from 10000 to 65535 and -refuse others. Pass 0 as the local port to pick a free one. Runs until -interrupted. +authenticated as you. Remote ports must be between 10000 and 65535. Pass 0 +as the local port to pick a free one. A stopped sandbox is started first. +Runs until interrupted. Examples: databricks sandbox port-forward -L 8080:10080 # default sandbox @@ -96,10 +107,15 @@ Examples: if err != nil { return err } - sandboxID, err := resolvePortForwardSandbox(ctx, api, profile, args) + sandboxID, status, err := resolvePortForwardSandbox(ctx, api, profile, args) if err != nil { return err } + if status != "" { + if _, err := ensureRunning(ctx, api, sandboxID, status); err != nil { + return err + } + } workspace, err := url.Parse(w.Config.Host) if err != nil { @@ -130,56 +146,62 @@ Examples: } // parseForwards parses -L LOCAL:REMOTE specs. LOCAL may be 0 to pick a free -// port. The sandbox enforces which remote ports it accepts. +// port. A few platform ports inside the remote range are refused by the sandbox. func parseForwards(specs []string) ([]portForward, error) { forwards := make([]portForward, 0, len(specs)) for _, spec := range specs { local, remote, found := strings.Cut(spec, ":") localPort, localErr := strconv.Atoi(local) remotePort, remoteErr := strconv.Atoi(remote) - if !found || localErr != nil || remoteErr != nil || localPort < 0 || localPort > 65535 || remotePort < 1 || remotePort > 65535 { + if !found || localErr != nil || remoteErr != nil || localPort < 0 || localPort > 65535 { return nil, fmt.Errorf("invalid -L %q: expected LOCAL:REMOTE ports, e.g. 8080:10080", spec) } + if remotePort < minRemotePort || remotePort > maxRemotePort { + return nil, fmt.Errorf("invalid -L %q: remote port must be between %d and %d", spec, minRemotePort, maxRemotePort) + } forwards = append(forwards, portForward{localPort: localPort, remotePort: remotePort}) } return forwards, nil } // resolvePortForwardSandbox returns the sandbox named in args, or the default -// sandbox when args is empty. Like `sandbox ssh`, it clears a default that no -// longer exists and only warns on other lookup errors, since the forward may -// still work. -func resolvePortForwardSandbox(ctx context.Context, api *sandboxAPI, profile string, args []string) (string, error) { +// sandbox when args is empty, and its status. Like `sandbox ssh`, it clears a +// default that no longer exists. Other lookup errors only warn and return an +// empty status, since the forward may still work. +func resolvePortForwardSandbox(ctx context.Context, api *sandboxAPI, profile string, args []string) (string, string, error) { var sandboxID string if len(args) == 0 { if sandboxID = getDefault(ctx, profile); sandboxID == "" { - return "", errors.New("no default sandbox configured — run `databricks sandbox create` to provision one, or `databricks sandbox default ` to point at an existing sandbox") + return "", "", errors.New("no default sandbox configured — run `databricks sandbox create` to provision one, or `databricks sandbox default ` to point at an existing sandbox") } } else { resolved, err := resolveLocalID(ctx, profile, args[0]) if err != nil { - return "", err + return "", "", err } sandboxID = resolved } - // A "." or ".." ID would move every forwarded request off the port-forward route. - if hasDotSegment("/" + url.PathEscape(sandboxID)) { - return "", fmt.Errorf("invalid sandbox ID %q", sandboxID) + // The ID is a segment of every forwarded path, so "." or ".." would move + // requests off the port-forward route. + if sandboxID == "." || sandboxID == ".." { + return "", "", fmt.Errorf("invalid sandbox ID %q", sandboxID) } - if _, err := api.get(ctx, sandboxID); err != nil { + entry, err := api.get(ctx, sandboxID) + if err != nil { switch { case !errors.Is(err, apierr.ErrNotFound): warn(ctx, fmt.Sprintf("Could not validate sandbox %s: %v", sandboxID, err)) + return sandboxID, "", nil case len(args) == 0: _ = clearDefault(ctx, profile) - return "", fmt.Errorf("saved default %q no longer exists (cleared) — run `databricks sandbox create` to provision a new one, or `databricks sandbox default ` to point at an existing sandbox", sandboxID) + return "", "", fmt.Errorf("saved default %q no longer exists (cleared) — run `databricks sandbox create` to provision a new one, or `databricks sandbox default ` to point at an existing sandbox", sandboxID) default: - return "", fmt.Errorf("no sandbox named %q — `databricks sandbox list` shows available IDs", sandboxID) + return "", "", fmt.Errorf("no sandbox named %q — `databricks sandbox list` shows available IDs", sandboxID) } } - return sandboxID, nil + return sandboxID, entry.Status, nil } // newPortForwardTransport returns the workspace config's transport, or the @@ -203,31 +225,46 @@ func newPortForwardTransport(cfg *config.Config) (http.RoundTripper, error) { // so a port that is already in use fails the command up front. func servePortForwards(ctx context.Context, target portForwardTarget, forwards []portForward) error { var lc net.ListenConfig - listeners := make([]net.Listener, 0, len(forwards)) + listeners := make([][]net.Listener, 0, len(forwards)) for _, f := range forwards { ln, err := lc.Listen(ctx, "tcp", net.JoinHostPort(portForwardBindHost, strconv.Itoa(f.localPort))) if err != nil { - for _, open := range listeners { - open.Close() + for _, group := range listeners { + for _, open := range group { + open.Close() + } } return fmt.Errorf("failed to listen on local port %d: %w", f.localPort, err) } - listeners = append(listeners, ln) + group := []net.Listener{ln} + port := strconv.Itoa(ln.Addr().(*net.TCPAddr).Port) + if ln6, err := lc.Listen(ctx, "tcp", net.JoinHostPort(portForwardBindHostIPv6, port)); err == nil { + group = append(group, ln6) + } else { + log.Debugf(ctx, "not forwarding [%s]:%s: %v", portForwardBindHostIPv6, port, err) + } + listeners = append(listeners, group) } + // net/http reports routine client disconnects through ErrorLog; keep them at + // debug level in the CLI log instead of on stderr. + logger := log.GetLogger(ctx) g, ctx := errgroup.WithContext(ctx) for i, f := range forwards { srv := &http.Server{ - Handler: newPortForwardHandler(target, f.remotePort), + Handler: newPortForwardHandler(target, f.remotePort, logger), ReadHeaderTimeout: portForwardReadHeaderTimeout, + ErrorLog: slog.NewLogLogger(logger.Handler(), slog.LevelDebug), + } + cmdio.LogString(ctx, fmt.Sprintf("Forwarding http://%s -> %s:%d", listeners[i][0].Addr(), target.sandboxID, f.remotePort)) + for _, ln := range listeners[i] { + g.Go(func() error { + if err := srv.Serve(ln); !errors.Is(err, http.ErrServerClosed) { + return fmt.Errorf("port forward on %s failed: %w", ln.Addr(), err) + } + return nil + }) } - cmdio.LogString(ctx, fmt.Sprintf("Forwarding http://%s -> %s:%d", listeners[i].Addr(), target.sandboxID, f.remotePort)) - g.Go(func() error { - if err := srv.Serve(listeners[i]); !errors.Is(err, http.ErrServerClosed) { - return fmt.Errorf("port forward on %s failed: %w", listeners[i].Addr(), err) - } - return nil - }) // An interactive interrupt stops at once: WebSockets and streams never // go idle, so a graceful Shutdown would only delay the exit. g.Go(func() error { @@ -241,7 +278,7 @@ func servePortForwards(ctx context.Context, target portForwardTarget, forwards [ // newPortForwardHandler relays each request to remotePort in the sandbox // through the workspace's port-forward route. -func newPortForwardHandler(target portForwardTarget, remotePort int) http.Handler { +func newPortForwardHandler(target portForwardTarget, remotePort int, logger *slog.Logger) http.Handler { base := strings.TrimSuffix(target.workspace.Path, "/") + portForwardRoute port := "/ports/" + strconv.Itoa(remotePort) proxy := &httputil.ReverseProxy{ @@ -262,6 +299,7 @@ func newPortForwardHandler(target portForwardTarget, remotePort int) http.Handle } }, Transport: &authenticatingTransport{authenticate: target.authenticate, base: target.transport}, + ErrorLog: slog.NewLogLogger(logger.Handler(), slog.LevelDebug), ErrorHandler: func(w http.ResponseWriter, r *http.Request, err error) { log.Debugf(r.Context(), "port forward to %s:%d failed: %v", target.sandboxID, remotePort, err) http.Error(w, "port forward failed: "+err.Error(), http.StatusBadGateway) @@ -271,6 +309,8 @@ func newPortForwardHandler(target portForwardTarget, remotePort int) http.Handle switch { case !isLoopbackHost(r.Host): http.Error(w, "port forward only accepts requests addressed to localhost", http.StatusForbidden) + case isCrossSiteRequest(r): + http.Error(w, "port forward does not accept cross-site browser requests", http.StatusForbidden) case hasDotSegment(r.URL.EscapedPath()): http.Error(w, "request path must not contain . or .. segments", http.StatusBadRequest) default: @@ -294,6 +334,23 @@ func isLoopbackHost(hostport string) bool { return ip != nil && ip.IsLoopback() } +// isCrossSiteRequest reports whether a browser marked the request as coming +// from a page outside this machine. Browsers send Sec-Fetch-Site with every +// request and Origin with cross-origin ones; WebSocket handshakes bypass CORS, +// so this is their only guard. Pages from any localhost port, such as a local +// frontend, are allowed. +func isCrossSiteRequest(r *http.Request) bool { + if r.Header.Get("Sec-Fetch-Site") == "cross-site" { + return true + } + origin := r.Header.Get("Origin") + if origin == "" { + return false + } + u, err := url.Parse(origin) + return err != nil || !isLoopbackHost(u.Host) +} + // hasDotSegment reports whether an escaped URL path contains a "." or ".." // segment, including percent-encoded and backslash-separated forms. The // workspace would resolve such a path out of the port-forward route and onto diff --git a/cmd/sandbox/portforward_test.go b/cmd/sandbox/portforward_test.go index be9b491d2ef..d174704103c 100644 --- a/cmd/sandbox/portforward_test.go +++ b/cmd/sandbox/portforward_test.go @@ -6,6 +6,7 @@ import ( "encoding/json" "errors" "io" + "log/slog" "net" "net/http" "net/http/httptest" @@ -29,8 +30,8 @@ func TestParseForwards(t *testing.T) { assert.Equal(t, []portForward{{localPort: 8080, remotePort: 10080}, {localPort: 0, remotePort: 10081}}, forwards) } -func TestParseForwardsRejectsInvalidSpecs(t *testing.T) { - for _, spec := range []string{"8080", "8080:", ":10080", "a:10080", "8080:b", "-1:10080", "65536:10080", "8080:0", "8080:65536", "8080:10080:1"} { +func TestParseForwardsRejectsMalformedSpecs(t *testing.T) { + for _, spec := range []string{"8080", "8080:", ":10080", "a:10080", "8080:b", "-1:10080", "65536:10080", "8080:10080:1"} { t.Run(spec, func(t *testing.T) { _, err := parseForwards([]string{spec}) assert.EqualError(t, err, `invalid -L "`+spec+`": expected LOCAL:REMOTE ports, e.g. 8080:10080`) @@ -38,6 +39,15 @@ func TestParseForwardsRejectsInvalidSpecs(t *testing.T) { } } +func TestParseForwardsRejectsRemotePortsOutsideRange(t *testing.T) { + for _, spec := range []string{"8080:0", "8080:8080", "8080:9999", "8080:65536"} { + t.Run(spec, func(t *testing.T) { + _, err := parseForwards([]string{spec}) + assert.EqualError(t, err, `invalid -L "`+spec+`": remote port must be between 10000 and 65535`) + }) + } +} + func TestHasDotSegment(t *testing.T) { tests := []struct { path string @@ -79,7 +89,7 @@ func startForward(t *testing.T, upstream *httptest.Server, authenticate func(*ht headers: map[string]string{"X-Databricks-Workspace-Id": "123"}, transport: http.DefaultTransport, } - forward := httptest.NewServer(newPortForwardHandler(target, 10080)) + forward := httptest.NewServer(newPortForwardHandler(target, 10080, slog.Default())) t.Cleanup(forward.Close) return forward } @@ -174,6 +184,45 @@ func TestPortForwardHandlerRejectsForeignHost(t *testing.T) { assert.Zero(t, calls.Load()) } +func TestPortForwardHandlerRejectsCrossSiteRequests(t *testing.T) { + var calls atomic.Int32 + upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + calls.Add(1) + })) + defer upstream.Close() + forward := startForward(t, upstream, fakeAuth) + addr := strings.TrimPrefix(forward.URL, "http://") + + for name, headers := range map[string]string{ + "foreign origin": "Origin: https://attacker.example\r\n", + "null origin": "Origin: null\r\n", + "cross-site fetch": "Sec-Fetch-Site: cross-site\r\n", + "cross-site websocket": "Origin: https://attacker.example\r\nUpgrade: websocket\r\nConnection: Upgrade\r\n", + } { + t.Run(name, func(t *testing.T) { + status, _, _ := rawRequest(t, addr, "GET / HTTP/1.1\r\nHost: "+addr+"\r\n"+headers+"Connection: close\r\n\r\n") + assert.Equal(t, http.StatusForbidden, status) + }) + } + assert.Zero(t, calls.Load()) +} + +func TestPortForwardHandlerAllowsLocalOrigins(t *testing.T) { + var calls atomic.Int32 + upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + calls.Add(1) + })) + defer upstream.Close() + forward := startForward(t, upstream, fakeAuth) + addr := strings.TrimPrefix(forward.URL, "http://") + + for _, headers := range []string{"Origin: http://localhost:3000\r\nSec-Fetch-Site: same-site\r\n", "Origin: http://" + addr + "\r\nSec-Fetch-Site: same-origin\r\n", "Sec-Fetch-Site: none\r\n"} { + status, _, _ := rawRequest(t, addr, "POST / HTTP/1.1\r\nHost: "+addr+"\r\n"+headers+"Content-Length: 0\r\nConnection: close\r\n\r\n") + assert.Equal(t, http.StatusOK, status, headers) + } + assert.Equal(t, int32(3), calls.Load()) +} + func TestIsLoopbackHost(t *testing.T) { tests := []struct { host string @@ -302,9 +351,10 @@ func TestResolvePortForwardSandboxUsesDefault(t *testing.T) { ctx = cmdio.MockDiscard(ctx) require.NoError(t, setDefault(ctx, "profile", testSandboxID)) - got, err := resolvePortForwardSandbox(ctx, newLookupAPI(t, http.StatusOK), "profile", nil) + got, status, err := resolvePortForwardSandbox(ctx, newLookupAPI(t, http.StatusOK), "profile", nil) require.NoError(t, err) assert.Equal(t, testSandboxID, got) + assert.Equal(t, "Running", status) } func TestResolvePortForwardSandboxClearsMissingDefault(t *testing.T) { @@ -312,7 +362,7 @@ func TestResolvePortForwardSandboxClearsMissingDefault(t *testing.T) { ctx = cmdio.MockDiscard(ctx) require.NoError(t, setDefault(ctx, "profile", testSandboxID)) - _, err := resolvePortForwardSandbox(ctx, newLookupAPI(t, http.StatusNotFound), "profile", nil) + _, _, err := resolvePortForwardSandbox(ctx, newLookupAPI(t, http.StatusNotFound), "profile", nil) assert.ErrorContains(t, err, `saved default "`+testSandboxID+`" no longer exists (cleared)`) assert.Empty(t, getDefault(ctx, "profile")) } @@ -321,7 +371,7 @@ func TestResolvePortForwardSandboxRejectsUnknownID(t *testing.T) { ctx, _ := stateCtx(t) ctx = cmdio.MockDiscard(ctx) - _, err := resolvePortForwardSandbox(ctx, newLookupAPI(t, http.StatusNotFound), "profile", []string{testSandboxID}) + _, _, err := resolvePortForwardSandbox(ctx, newLookupAPI(t, http.StatusNotFound), "profile", []string{testSandboxID}) assert.ErrorContains(t, err, `no sandbox named "`+testSandboxID+`"`) } @@ -329,16 +379,17 @@ func TestResolvePortForwardSandboxContinuesOnOtherLookupErrors(t *testing.T) { ctx, _ := stateCtx(t) ctx = cmdio.MockDiscard(ctx) - got, err := resolvePortForwardSandbox(ctx, newLookupAPI(t, http.StatusForbidden), "profile", []string{testSandboxID}) + got, status, err := resolvePortForwardSandbox(ctx, newLookupAPI(t, http.StatusForbidden), "profile", []string{testSandboxID}) require.NoError(t, err) assert.Equal(t, testSandboxID, got) + assert.Empty(t, status) } func TestResolvePortForwardSandboxRejectsDotSegmentIDs(t *testing.T) { ctx, _ := stateCtx(t) for _, id := range []string{".", ".."} { t.Run(id, func(t *testing.T) { - _, err := resolvePortForwardSandbox(ctx, nil, "profile", []string{id}) + _, _, err := resolvePortForwardSandbox(ctx, nil, "profile", []string{id}) assert.EqualError(t, err, `invalid sandbox ID "`+id+`"`) }) } @@ -347,6 +398,6 @@ func TestResolvePortForwardSandboxRejectsDotSegmentIDs(t *testing.T) { func TestResolvePortForwardSandboxRequiresDefault(t *testing.T) { ctx, _ := stateCtx(t) - _, err := resolvePortForwardSandbox(ctx, nil, "profile", nil) + _, _, err := resolvePortForwardSandbox(ctx, nil, "profile", nil) assert.ErrorContains(t, err, "no default sandbox configured") }