Inject OBO access token in apps run-local

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
 

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)

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) {

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,14 +48,40 @@ 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)
 	}
 	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)

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")
+}

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",