Skip to content

Commit 95ca5f1

Browse files
jaecho-dbIsaac
andcommitted
sandbox port-forward: address review feedback
- Share sandbox resolution with `sandbox ssh` (resolveSandbox), and reject IDs that are not a single path segment. A 400 from the lookup now fails instead of warning. - Let a loopback Origin take precedence over Sec-Fetch-Site, so a frontend on http://localhost:N can call the printed http://127.0.0.1:N URL. - Send the package's X-Databricks-Org-Id routing header, as the sandbox API does. - Add an acceptance test that forwards a request end to end. Co-authored-by: Isaac <no-reply@databricks.com>
1 parent 1e44928 commit 95ca5f1

9 files changed

Lines changed: 255 additions & 202 deletions

File tree

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1 @@
1+
Cloud = false
Lines changed: 28 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,28 @@
1+
2+
=== The command reports the forward once it is listening
3+
4+
>>> retry --until Press Ctrl-C to stop. cat LOG.forward
5+
Forwarding http://127.0.0.1:[LOCAL_PORT] -> happy-panda-1234:10080
6+
Press Ctrl-C to stop.
7+
8+
=== A local request reaches the sandbox app through the port-forward route
9+
10+
>>> curl -sS -X POST -H X-Databricks-Org-Id: 999 -d ping http://127.0.0.1:[LOCAL_PORT]/api/items?page=2
11+
hello from the sandbox app
12+
13+
=== The workspace receives the CLI's routing header in place of the caller's
14+
15+
>>> print_requests.py //sandbox-proxy
16+
{
17+
"headers": {
18+
"X-Databricks-Org-Id": [
19+
"[WORKSPACE_ID]"
20+
]
21+
},
22+
"method": "POST",
23+
"path": "/sandbox-proxy/sandboxes/happy-panda-1234/ports/10080/api/items",
24+
"q": {
25+
"page": "2"
26+
},
27+
"raw_body": "ping"
28+
}
Lines changed: 20 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,20 @@
1+
# Ports are allocated per run so that tests in parallel worktrees do not collide.
2+
read -r LOCAL_PORT <<< "$(free_port.py 1)"
3+
4+
# The forward runs until interrupted. Start it outside this shell's job table so
5+
# stopping it doesn't print a "Killed" job notice, and stop it by port, which
6+
# also works on Windows.
7+
( $CLI sandbox port-forward happy-panda-1234 -L "$LOCAL_PORT:10080" &> LOG.forward & )
8+
trap 'kill_port.py "$LOCAL_PORT"' EXIT
9+
10+
export RETRY_MAX_ATTEMPTS=30
11+
12+
title "The command reports the forward once it is listening\n"
13+
trace retry --until "Press Ctrl-C to stop." cat LOG.forward
14+
15+
title "A local request reaches the sandbox app through the port-forward route\n"
16+
trace curl -sS -X POST -H "X-Databricks-Org-Id: 999" -d 'ping' "http://127.0.0.1:$LOCAL_PORT/api/items?page=2"
17+
echo
18+
19+
title "The workspace receives the CLI's routing header in place of the caller's\n"
20+
trace print_requests.py //sandbox-proxy
Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,15 @@
1+
RecordRequests = true
2+
IncludeRequestHeaders = ["X-Databricks-Org-Id"]
3+
4+
[[Server]]
5+
Pattern = "GET /api/2.0/lakebox/sandboxes/happy-panda-1234"
6+
Response.Body = '{"sandboxId": "happy-panda-1234", "status": "Running"}'
7+
8+
[[Server]]
9+
Pattern = "POST /sandbox-proxy/sandboxes/happy-panda-1234/ports/10080/{path...}"
10+
Response.Body = 'hello from the sandbox app'
11+
12+
# The local port is allocated per run.
13+
[[Repls]]
14+
Old = '127\.0\.0\.1:\d+'
15+
New = '127.0.0.1:[LOCAL_PORT]'

‎cmd/sandbox/portforward.go‎

Lines changed: 18 additions & 59 deletions
Original file line numberDiff line numberDiff line change
@@ -18,11 +18,9 @@ import (
1818
"time"
1919

2020
"github.com/databricks/cli/cmd/root"
21-
"github.com/databricks/cli/libs/auth"
2221
"github.com/databricks/cli/libs/cmdctx"
2322
"github.com/databricks/cli/libs/cmdio"
2423
"github.com/databricks/cli/libs/log"
25-
"github.com/databricks/databricks-sdk-go/apierr"
2624
"github.com/databricks/databricks-sdk-go/config"
2725
"github.com/spf13/cobra"
2826
"golang.org/x/sync/errgroup"
@@ -107,12 +105,16 @@ Examples:
107105
if err != nil {
108106
return err
109107
}
110-
sandboxID, status, err := resolvePortForwardSandbox(ctx, api, profile, args)
108+
var arg string
109+
if len(args) > 0 {
110+
arg = args[0]
111+
}
112+
sandboxID, sb, err := resolveSandbox(ctx, api, profile, arg)
111113
if err != nil {
112114
return err
113115
}
114-
if status != "" {
115-
if _, err := ensureRunning(ctx, api, sandboxID, status); err != nil {
116+
if sb != nil {
117+
if _, err := ensureRunning(ctx, api, sandboxID, sb.Status); err != nil {
116118
return err
117119
}
118120
}
@@ -129,7 +131,7 @@ Examples:
129131
workspace: workspace,
130132
sandboxID: sandboxID,
131133
authenticate: w.Config.Authenticate,
132-
headers: auth.WorkspaceIDHeaders(w.Config),
134+
headers: api.headers(),
133135
transport: transport,
134136
}
135137

@@ -164,46 +166,6 @@ func parseForwards(specs []string) ([]portForward, error) {
164166
return forwards, nil
165167
}
166168

167-
// resolvePortForwardSandbox returns the sandbox named in args, or the default
168-
// sandbox when args is empty, and its status. Like `sandbox ssh`, it clears a
169-
// default that no longer exists. Other lookup errors only warn and return an
170-
// empty status, since the forward may still work.
171-
func resolvePortForwardSandbox(ctx context.Context, api *sandboxAPI, profile string, args []string) (string, string, error) {
172-
var sandboxID string
173-
if len(args) == 0 {
174-
if sandboxID = getDefault(ctx, profile); sandboxID == "" {
175-
return "", "", errors.New("no default sandbox configured — run `databricks sandbox create` to provision one, or `databricks sandbox default <id>` to point at an existing sandbox")
176-
}
177-
} else {
178-
resolved, err := resolveLocalID(ctx, profile, args[0])
179-
if err != nil {
180-
return "", "", err
181-
}
182-
sandboxID = resolved
183-
}
184-
185-
// The ID is a segment of every forwarded path, so "." or ".." would move
186-
// requests off the port-forward route.
187-
if sandboxID == "." || sandboxID == ".." {
188-
return "", "", fmt.Errorf("invalid sandbox ID %q", sandboxID)
189-
}
190-
191-
entry, err := api.get(ctx, sandboxID)
192-
if err != nil {
193-
switch {
194-
case !errors.Is(err, apierr.ErrNotFound):
195-
warn(ctx, fmt.Sprintf("Could not validate sandbox %s: %v", sandboxID, err))
196-
return sandboxID, "", nil
197-
case len(args) == 0:
198-
_ = clearDefault(ctx, profile)
199-
return "", "", fmt.Errorf("saved default %q no longer exists (cleared) — run `databricks sandbox create` to provision a new one, or `databricks sandbox default <id>` to point at an existing sandbox", sandboxID)
200-
default:
201-
return "", "", fmt.Errorf("no sandbox named %q — `databricks sandbox list` shows available IDs", sandboxID)
202-
}
203-
}
204-
return sandboxID, entry.Status, nil
205-
}
206-
207169
// newPortForwardTransport returns the workspace config's transport, or the
208170
// default one with the config's TLS verification setting.
209171
func newPortForwardTransport(cfg *config.Config) (http.RoundTripper, error) {
@@ -334,21 +296,18 @@ func isLoopbackHost(hostport string) bool {
334296
return ip != nil && ip.IsLoopback()
335297
}
336298

337-
// isCrossSiteRequest reports whether a browser marked the request as coming
338-
// from a page outside this machine. Browsers send Sec-Fetch-Site with every
339-
// request and Origin with cross-origin ones; WebSocket handshakes bypass CORS,
340-
// so this is their only guard. Pages from any localhost port, such as a local
341-
// frontend, are allowed.
299+
// isCrossSiteRequest reports whether a browser sent the request from a page
300+
// outside this machine. WebSocket handshakes bypass CORS, so this is their only
301+
// guard. Pages from any loopback origin, such as a local frontend, are allowed:
302+
// browsers treat localhost and 127.0.0.1 as different sites, so their Origin
303+
// decides rather than Sec-Fetch-Site. Requests without an Origin, such as
304+
// navigations and image loads, fall back to Sec-Fetch-Site.
342305
func isCrossSiteRequest(r *http.Request) bool {
343-
if r.Header.Get("Sec-Fetch-Site") == "cross-site" {
344-
return true
345-
}
346-
origin := r.Header.Get("Origin")
347-
if origin == "" {
348-
return false
306+
if origin := r.Header.Get("Origin"); origin != "" {
307+
u, err := url.Parse(origin)
308+
return err != nil || !isLoopbackHost(u.Host)
349309
}
350-
u, err := url.Parse(origin)
351-
return err != nil || !isLoopbackHost(u.Host)
310+
return r.Header.Get("Sec-Fetch-Site") == "cross-site"
352311
}
353312

354313
// hasDotSegment reports whether an escaped URL path contains a "." or ".."

‎cmd/sandbox/portforward_test.go‎

Lines changed: 8 additions & 86 deletions
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,6 @@ package sandbox
33
import (
44
"bufio"
55
"context"
6-
"encoding/json"
76
"errors"
87
"io"
98
"log/slog"
@@ -17,7 +16,6 @@ import (
1716
"testing"
1817

1918
"github.com/databricks/cli/libs/cmdio"
20-
"github.com/databricks/databricks-sdk-go"
2119
"github.com/stretchr/testify/assert"
2220
"github.com/stretchr/testify/require"
2321
)
@@ -216,11 +214,17 @@ func TestPortForwardHandlerAllowsLocalOrigins(t *testing.T) {
216214
forward := startForward(t, upstream, fakeAuth)
217215
addr := strings.TrimPrefix(forward.URL, "http://")
218216

219-
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"} {
217+
for _, headers := range []string{
218+
"Origin: http://localhost:3000\r\nSec-Fetch-Site: same-site\r\n",
219+
"Origin: http://" + addr + "\r\nSec-Fetch-Site: same-origin\r\n",
220+
// Browsers treat localhost and 127.0.0.1 as different sites.
221+
"Origin: http://localhost:5173\r\nSec-Fetch-Site: cross-site\r\n",
222+
"Sec-Fetch-Site: none\r\n",
223+
} {
220224
status, _, _ := rawRequest(t, addr, "POST / HTTP/1.1\r\nHost: "+addr+"\r\n"+headers+"Content-Length: 0\r\nConnection: close\r\n\r\n")
221225
assert.Equal(t, http.StatusOK, status, headers)
222226
}
223-
assert.Equal(t, int32(3), calls.Load())
227+
assert.Equal(t, int32(4), calls.Load())
224228
}
225229

226230
func TestIsLoopbackHost(t *testing.T) {
@@ -319,85 +323,3 @@ func testTarget() portForwardTarget {
319323
transport: http.DefaultTransport,
320324
}
321325
}
322-
323-
// newLookupAPI returns a sandbox API whose GET for testSandboxID answers with status.
324-
func newLookupAPI(t *testing.T, status int) *sandboxAPI {
325-
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
326-
if r.URL.Path == "/.well-known/databricks-config" {
327-
assert.NoError(t, json.NewEncoder(w).Encode(map[string]string{
328-
"oidc_endpoint": "https://workspace.example.test/oidc",
329-
}))
330-
return
331-
}
332-
assert.Equal(t, sandboxPath(testSandboxID), r.URL.Path)
333-
w.WriteHeader(status)
334-
if status == http.StatusOK {
335-
assert.NoError(t, json.NewEncoder(w).Encode(sandboxEntry{SandboxID: testSandboxID, Status: "Running"}))
336-
return
337-
}
338-
_, _ = io.WriteString(w, `{"error_code": "ERROR", "message": "lookup failed"}`)
339-
}))
340-
t.Cleanup(server.Close)
341-
342-
workspaceClient, err := databricks.NewWorkspaceClient(&databricks.Config{Host: server.URL, Token: "test-token"})
343-
require.NoError(t, err)
344-
api, err := newSandboxAPI(workspaceClient)
345-
require.NoError(t, err)
346-
return api
347-
}
348-
349-
func TestResolvePortForwardSandboxUsesDefault(t *testing.T) {
350-
ctx, _ := stateCtx(t)
351-
ctx = cmdio.MockDiscard(ctx)
352-
require.NoError(t, setDefault(ctx, "profile", testSandboxID))
353-
354-
got, status, err := resolvePortForwardSandbox(ctx, newLookupAPI(t, http.StatusOK), "profile", nil)
355-
require.NoError(t, err)
356-
assert.Equal(t, testSandboxID, got)
357-
assert.Equal(t, "Running", status)
358-
}
359-
360-
func TestResolvePortForwardSandboxClearsMissingDefault(t *testing.T) {
361-
ctx, _ := stateCtx(t)
362-
ctx = cmdio.MockDiscard(ctx)
363-
require.NoError(t, setDefault(ctx, "profile", testSandboxID))
364-
365-
_, _, err := resolvePortForwardSandbox(ctx, newLookupAPI(t, http.StatusNotFound), "profile", nil)
366-
assert.ErrorContains(t, err, `saved default "`+testSandboxID+`" no longer exists (cleared)`)
367-
assert.Empty(t, getDefault(ctx, "profile"))
368-
}
369-
370-
func TestResolvePortForwardSandboxRejectsUnknownID(t *testing.T) {
371-
ctx, _ := stateCtx(t)
372-
ctx = cmdio.MockDiscard(ctx)
373-
374-
_, _, err := resolvePortForwardSandbox(ctx, newLookupAPI(t, http.StatusNotFound), "profile", []string{testSandboxID})
375-
assert.ErrorContains(t, err, `no sandbox named "`+testSandboxID+`"`)
376-
}
377-
378-
func TestResolvePortForwardSandboxContinuesOnOtherLookupErrors(t *testing.T) {
379-
ctx, _ := stateCtx(t)
380-
ctx = cmdio.MockDiscard(ctx)
381-
382-
got, status, err := resolvePortForwardSandbox(ctx, newLookupAPI(t, http.StatusForbidden), "profile", []string{testSandboxID})
383-
require.NoError(t, err)
384-
assert.Equal(t, testSandboxID, got)
385-
assert.Empty(t, status)
386-
}
387-
388-
func TestResolvePortForwardSandboxRejectsDotSegmentIDs(t *testing.T) {
389-
ctx, _ := stateCtx(t)
390-
for _, id := range []string{".", ".."} {
391-
t.Run(id, func(t *testing.T) {
392-
_, _, err := resolvePortForwardSandbox(ctx, nil, "profile", []string{id})
393-
assert.EqualError(t, err, `invalid sandbox ID "`+id+`"`)
394-
})
395-
}
396-
}
397-
398-
func TestResolvePortForwardSandboxRequiresDefault(t *testing.T) {
399-
ctx, _ := stateCtx(t)
400-
401-
_, _, err := resolvePortForwardSandbox(ctx, nil, "profile", nil)
402-
assert.ErrorContains(t, err, "no default sandbox configured")
403-
}

‎cmd/sandbox/resolve.go‎

Lines changed: 55 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,8 @@ import (
55
"errors"
66
"fmt"
77
"strings"
8+
9+
"github.com/databricks/databricks-sdk-go/apierr"
810
)
911

1012
// resolveLocalID maps a user-typed argument to a sandbox ID using only
@@ -61,3 +63,56 @@ func resolveLocalID(ctx context.Context, profile, arg string) (string, error) {
6163
)
6264
}
6365
}
66+
67+
// resolveSandbox returns the ID of the sandbox named by arg, or of the
68+
// profile's default sandbox when arg is empty, and the sandbox as the server
69+
// reports it. The entry is nil when the lookup failed with an error other than
70+
// not-found or a rejected request: those only warn, so a transient API hiccup
71+
// doesn't block a connection the gateway can still route.
72+
func resolveSandbox(ctx context.Context, api *sandboxAPI, profile, arg string) (string, *sandboxEntry, error) {
73+
var sandboxID string
74+
if arg == "" {
75+
if sandboxID = getDefault(ctx, profile); sandboxID == "" {
76+
return "", nil, errors.New("no default sandbox configured — run `databricks sandbox create` to provision one, or `databricks sandbox default <id>` to point at an existing sandbox")
77+
}
78+
} else {
79+
resolved, err := resolveLocalID(ctx, profile, arg)
80+
if err != nil {
81+
return "", nil, err
82+
}
83+
sandboxID = resolved
84+
}
85+
if err := validateSandboxID(sandboxID); err != nil {
86+
return "", nil, err
87+
}
88+
89+
sb, err := api.get(ctx, sandboxID)
90+
switch {
91+
case err == nil:
92+
return sandboxID, sb, nil
93+
case errors.Is(err, apierr.ErrNotFound) && arg == "":
94+
// A stale default is almost always user-correctable (deleted from
95+
// another machine, reaped by an admin); fail fast with an actionable
96+
// message.
97+
_ = clearDefault(ctx, profile)
98+
return "", nil, fmt.Errorf("saved default %q no longer exists (cleared) — run `databricks sandbox create` to provision a new one, or `databricks sandbox default <id>` to point at an existing sandbox", sandboxID)
99+
case errors.Is(err, apierr.ErrNotFound):
100+
return "", nil, fmt.Errorf("no sandbox named %q — `databricks sandbox list` shows available IDs", sandboxID)
101+
case errors.Is(err, apierr.ErrBadRequest):
102+
return "", nil, fmt.Errorf("failed to look up sandbox %s: %w", sandboxID, err)
103+
default:
104+
warn(ctx, fmt.Sprintf("Could not validate sandbox %s: %v", sandboxID, err))
105+
return sandboxID, nil, nil
106+
}
107+
}
108+
109+
// validateSandboxID rejects IDs that are not a single path segment. Commands
110+
// put the ID in workspace URL paths, and the workspace decodes an escaped "/",
111+
// so a separator or dot segment would move a request onto another API with the
112+
// caller's credentials.
113+
func validateSandboxID(id string) error {
114+
if id == "." || id == ".." || strings.ContainsAny(id, `/\`) {
115+
return fmt.Errorf("invalid sandbox ID %q", id)
116+
}
117+
return nil
118+
}

0 commit comments

Comments
 (0)