diff --git a/NEXT_CHANGELOG.md b/NEXT_CHANGELOG.md index 715aad31819..a3ce99a68aa 100644 --- a/NEXT_CHANGELOG.md +++ b/NEXT_CHANGELOG.md @@ -7,6 +7,7 @@ ### CLI * An explicitly selected profile (`--profile` or a bundle's `workspace.profile`) now takes precedence over auth environment variables (`DATABRICKS_HOST`, `DATABRICKS_TOKEN`, etc.) instead of being silently shadowed by them; env vars still fill auth fields the profile leaves empty ([#5096](https://github.com/databricks/cli/issues/5096)). +* `databricks apps run-local` now injects an `X-Forwarded-Access-Token` header minted from your CLI credentials, matching the deployed Apps OAuth2 proxy, so on-behalf-of (OBO) code paths can be exercised locally. When the current credentials are not OAuth-based (e.g. a PAT), the header is omitted with a warning and the proxy keeps working as before ([#5795](https://github.com/databricks/cli/pull/5795)). ### Bundles diff --git a/cmd/apps/run_local.go b/cmd/apps/run_local.go index d9671276f0f..dd8bd2205ae 100644 --- a/cmd/apps/run_local.go +++ b/cmd/apps/run_local.go @@ -17,6 +17,7 @@ import ( "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" "github.com/databricks/databricks-sdk-go/service/iam" "github.com/spf13/cobra" @@ -126,6 +127,27 @@ func setupProxy(ctx context.Context, cmd *cobra.Command, config *runlocal.Config proxy.InjectHeader(key, value) } + // The deployed Apps OAuth2 proxy injects the caller's OBO token as + // X-Forwarded-Access-Token; locally that header is absent, so OBO paths + // can't be run. Inject a token from the CLI's own credentials, resolved + // per request so a browser session outliving the token still refreshes. + // + // GetTokenSource only yields OAuth tokens; for PAT/basic auth it returns a + // source that errors. Probe it once here so those users still get a working + // proxy (without the OBO header, as before) rather than a 502 per request. + tokenSource := w.Config.GetTokenSource() + if _, err := tokenSource.Token(ctx); err != nil { + log.Warnf(ctx, "Not injecting %s: could not mint an OAuth token from the current credentials (%v). OBO code paths will not work locally; log in with OAuth (databricks auth login) to exercise them.", runlocal.HeaderForwardedAccessToken, err) + } else { + proxy.InjectHeaderFunc(runlocal.HeaderForwardedAccessToken, func(ctx context.Context) (string, error) { + token, err := tokenSource.Token(ctx) + if err != nil { + return "", err + } + return token.AccessToken, nil + }) + } + proxyAddr := fmt.Sprintf("localhost:%d", port) // Bind synchronously so a taken port fails the command instead of only printing an error from the goroutine. ln, err := proxy.Listen(proxyAddr) diff --git a/cmd/apps/run_local_test.go b/cmd/apps/run_local_test.go index 94a0b56ec70..b4e4c4d0dd6 100644 --- a/cmd/apps/run_local_test.go +++ b/cmd/apps/run_local_test.go @@ -1,13 +1,19 @@ package apps import ( + "io" "net" + "net/http" + "net/http/httptest" "os" "os/exec" + "strconv" "testing" "time" "github.com/databricks/cli/libs/apps/runlocal" + "github.com/databricks/cli/libs/cmdio" + "github.com/databricks/databricks-sdk-go/config" "github.com/databricks/databricks-sdk-go/experimental/mocks" "github.com/databricks/databricks-sdk-go/service/iam" "github.com/spf13/cobra" @@ -23,12 +29,48 @@ func TestSetupProxyPortInUse(t *testing.T) { m := mocks.NewMockWorkspaceClient(t) m.GetMockCurrentUserAPI().EXPECT().Me(mock.Anything, mock.Anything).Return(&iam.User{UserName: "test-user"}, nil) + // setupProxy reads a token source off the config; the real command + // always has a resolved config here. + m.WorkspaceClient.Config = &config.Config{Host: "https://workspace.databricks.test", Token: "token"} - config := runlocal.NewConfig("https://workspace.databricks.test", "123", t.TempDir(), runlocal.DEFAULT_HOST, runlocal.DEFAULT_PORT) - err = setupProxy(t.Context(), &cobra.Command{}, config, m.WorkspaceClient, port, false) + cfg := runlocal.NewConfig("https://workspace.databricks.test", "123", t.TempDir(), runlocal.DEFAULT_HOST, runlocal.DEFAULT_PORT) + err = setupProxy(t.Context(), &cobra.Command{}, cfg, m.WorkspaceClient, port, false) require.ErrorContains(t, err, "failed to start app proxy") } +func TestSetupProxyPATOmitsTokenHeader(t *testing.T) { + // A PAT config has no OAuth token source, so setupProxy must forward + // requests without the OBO token header rather than 502 on every one. + backend := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(http.StatusOK) + _, _ = io.WriteString(w, r.Header.Get(runlocal.HeaderForwardedAccessToken)) + })) + defer backend.Close() + + ln, err := net.Listen("tcp", "localhost:0") + require.NoError(t, err) + port := ln.Addr().(*net.TCPAddr).Port + require.NoError(t, ln.Close()) + + m := mocks.NewMockWorkspaceClient(t) + m.GetMockCurrentUserAPI().EXPECT().Me(mock.Anything, mock.Anything).Return(&iam.User{UserName: "test-user"}, nil) + m.WorkspaceClient.Config = &config.Config{Host: "https://workspace.databricks.test", Token: "dapi-token"} + + cfg := runlocal.NewConfig("https://workspace.databricks.test", "123", t.TempDir(), runlocal.DEFAULT_HOST, runlocal.DEFAULT_PORT) + cfg.AppURL = backend.URL + + require.NoError(t, setupProxy(cmdio.MockDiscard(t.Context()), &cobra.Command{}, cfg, m.WorkspaceClient, port, false)) + + resp, err := http.Get("http://" + net.JoinHostPort("localhost", strconv.Itoa(port))) + require.NoError(t, err) + defer resp.Body.Close() + body, err := io.ReadAll(resp.Body) + require.NoError(t, err) + + require.Equal(t, http.StatusOK, resp.StatusCode) + require.Empty(t, string(body)) +} + // TestAppHelperProcess is not a real test: TestKillAppProcess re-invokes the // test binary with -test.run targeting it to get a long-running child process. func TestAppHelperProcess(t *testing.T) { diff --git a/libs/appproxy/appproxy.go b/libs/appproxy/appproxy.go index bfe5150ff96..4cb29316537 100644 --- a/libs/appproxy/appproxy.go +++ b/libs/appproxy/appproxy.go @@ -17,6 +17,7 @@ type Proxy struct { client *http.Client server *http.Server headerToInject map[string]string + dynamicHeaders map[string]func(context.Context) (string, error) } // Creates a new proxy instance that will forward all requests to the targetURL @@ -47,7 +48,7 @@ func (p *Proxy) Stop() error { return p.server.Shutdown(p.ctx) } -// InjectHeader injects a header that will be added to all requests forwarded by the proxy +// InjectHeader injects a static header that will be added to all requests forwarded by the proxy. func (p *Proxy) InjectHeader(key, value string) { if p.headerToInject == nil { p.headerToInject = make(map[string]string) @@ -55,6 +56,32 @@ func (p *Proxy) InjectHeader(key, value string) { p.headerToInject[key] = value } +// InjectHeaderFunc injects a header whose value is resolved per request by fn, +// for values that expire such as an OAuth token. If fn errors the request +// fails with 502 rather than forwarding without the header. +func (p *Proxy) InjectHeaderFunc(key string, fn func(context.Context) (string, error)) { + if p.dynamicHeaders == nil { + p.dynamicHeaders = make(map[string]func(context.Context) (string, error)) + } + p.dynamicHeaders[key] = fn +} + +// applyInjectedHeaders adds the static and dynamic injected headers to r, +// resolving dynamic headers against r's context. +func (p *Proxy) applyInjectedHeaders(r *http.Request) error { + for k, v := range p.headerToInject { + r.Header.Add(k, v) + } + for k, fn := range p.dynamicHeaders { + v, err := fn(r.Context()) + if err != nil { + return err + } + r.Header.Set(k, v) + } + return nil +} + func (p *Proxy) proxyHandler(w http.ResponseWriter, r *http.Request) { if p.isWebSocketRequest(r) { p.handleWebSocket(w, r) @@ -82,6 +109,15 @@ func (p *Proxy) handleWebSocket(w http.ResponseWriter, r *http.Request) { return } + // Inject headers on the upgrade request too, matching the HTTP path. + // This must happen before Hijack: once the connection is hijacked, w is + // detached and http.Error no longer reaches the client, so the error + // branch would otherwise send nothing. + if err := p.applyInjectedHeaders(r); err != nil { + http.Error(w, "Error resolving injected header: "+err.Error(), http.StatusBadGateway) + return + } + // We need to hijack the connection to be able to proxy the request // to the websocket server. We can use the hijacked connection to // read and write messages back and forth from the client to the app. @@ -129,8 +165,9 @@ func (p *Proxy) handleHTTP(w http.ResponseWriter, r *http.Request) { r.Host = p.targetURL.Host // Inject additional headers - for k, v := range p.headerToInject { - r.Header.Add(k, v) + if err := p.applyInjectedHeaders(r); err != nil { + http.Error(w, "Error resolving injected header: "+err.Error(), http.StatusBadGateway) + return } resp, err := p.client.Do(r) diff --git a/libs/appproxy/appproxy_test.go b/libs/appproxy/appproxy_test.go index 2f90ce06d0e..db3684357b7 100644 --- a/libs/appproxy/appproxy_test.go +++ b/libs/appproxy/appproxy_test.go @@ -2,9 +2,12 @@ package appproxy import ( "bytes" + "context" + "errors" "fmt" "io" "net/http" + "net/http/httptest" "strings" "testing" @@ -152,3 +155,88 @@ func TestProxyHandleWebSocket(t *testing.T) { t.Errorf("Expected one of the expected errors, got %s", err) } } + +func TestProxyInjectHeaderFunc(t *testing.T) { + // Echo the injected header back so we can assert its per-request value. + server := httptest.NewUnstartedServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(http.StatusOK) + _, _ = io.WriteString(w, r.Header.Get("X-Forwarded-Access-Token")) + })) + go server.Start() + defer server.Close() + + proxy, addr := startProxy(t, server.Listener.Addr().String()) + defer func() { + require.NoError(t, proxy.Stop()) + }() + + calls := 0 + proxy.InjectHeaderFunc("X-Forwarded-Access-Token", func(context.Context) (string, error) { + calls++ + return fmt.Sprintf("token-%d", calls), nil + }) + + code, body := sendTestRequest(t, "http://"+addr, "/") + require.Equal(t, http.StatusOK, code) + require.Equal(t, "token-1", string(body)) + + // A second request resolves the value again rather than reusing it. + code, body = sendTestRequest(t, "http://"+addr, "/") + require.Equal(t, http.StatusOK, code) + require.Equal(t, "token-2", string(body)) +} + +func TestProxyInjectHeaderFuncError(t *testing.T) { + // The app would 200 if reached; a resolution error must 502 instead. + server := httptest.NewUnstartedServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(http.StatusOK) + })) + go server.Start() + defer server.Close() + + proxy, addr := startProxy(t, server.Listener.Addr().String()) + defer func() { + require.NoError(t, proxy.Stop()) + }() + + proxy.InjectHeaderFunc("X-Forwarded-Access-Token", func(context.Context) (string, error) { + return "", errors.New("token refresh failed") + }) + + code, body := sendTestRequest(t, "http://"+addr, "/") + require.Equal(t, http.StatusBadGateway, code) + require.Contains(t, string(body), "token refresh failed") +} + +func TestProxyInjectHeaderFuncErrorWebSocket(t *testing.T) { + // On the WebSocket upgrade path the 502 must be written before Hijack, + // otherwise the client sees a closed connection instead of the error. + server := httptest.NewUnstartedServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(http.StatusOK) + })) + go server.Start() + defer server.Close() + + proxy, addr := startProxy(t, server.Listener.Addr().String()) + defer func() { + require.NoError(t, proxy.Stop()) + }() + + proxy.InjectHeaderFunc("X-Forwarded-Access-Token", func(context.Context) (string, error) { + return "", errors.New("token refresh failed") + }) + + req, err := http.NewRequest(http.MethodGet, "http://"+addr+"/", nil) + require.NoError(t, err) + req.Header.Set("Upgrade", "websocket") + req.Header.Set("Connection", "Upgrade") + + resp, err := http.DefaultClient.Do(req) + require.NoError(t, err) + defer resp.Body.Close() + body, err := io.ReadAll(resp.Body) + require.NoError(t, err) + + require.Equal(t, http.StatusBadGateway, resp.StatusCode) + require.Contains(t, string(body), "token refresh failed") +} diff --git a/libs/apps/runlocal/headers.go b/libs/apps/runlocal/headers.go index 9655e51ba28..2eb8743f9fa 100644 --- a/libs/apps/runlocal/headers.go +++ b/libs/apps/runlocal/headers.go @@ -5,6 +5,10 @@ import ( "github.com/google/uuid" ) +// HeaderForwardedAccessToken carries the caller's OBO access token, matching +// the deployed Apps OAuth2 proxy. +const HeaderForwardedAccessToken = "X-Forwarded-Access-Token" + func GetXHeaders(user *iam.User) map[string]string { return map[string]string{ "X-Forwarded-Host": "localhost",