diff --git a/cmd/auth/login.go b/cmd/auth/login.go index bd3123b7772..ff86f3ec717 100644 --- a/cmd/auth/login.go +++ b/cmd/auth/login.go @@ -150,6 +150,7 @@ a new profile is created. var skipWorkspace bool var scopes string var clientID string + var resources []string cmd.Flags().DurationVar(&loginTimeout, "timeout", defaultTimeout, "Timeout for completing login challenge in the browser") cmd.Flags().BoolVar(&configureCluster, "configure-cluster", false, @@ -162,6 +163,8 @@ a new profile is created. "Comma-separated list of OAuth scopes to request (defaults to 'all-apis')") cmd.Flags().StringVar(&clientID, "client-id", "", "OAuth client ID to use for U2M authentication") + cmd.Flags().StringArrayVar(&resources, "resource", nil, + "RFC 8707 resource indicator to scope the login to (repeatable). Requires --host.") cmd.PreRunE = profileHostConflictCheck @@ -331,6 +334,9 @@ a new profile is created. if len(scopesList) > 0 { persistentAuthOpts = append(persistentAuthOpts, u2m.WithScopes(scopesList)) } + if len(resources) > 0 { + persistentAuthOpts = append(persistentAuthOpts, u2m.WithResources(resources)) + } persistentAuth, err := u2m.NewPersistentAuth(ctx, persistentAuthOpts...) if err != nil { return err @@ -643,6 +649,10 @@ var discoveryIncompatibleFlags = []string{ "workspace-id", "configure-cluster", "configure-serverless", + // A resource indicator scopes the login to a protected resource on a + // specific workspace's /oidc, which the login.databricks.com discovery + // flow does not target. + "resource", } // validateDiscoveryFlagCompatibility returns an error if any flags that require diff --git a/cmd/auth/login_test.go b/cmd/auth/login_test.go index 688b0a66dbc..6fcbdde2b52 100644 --- a/cmd/auth/login_test.go +++ b/cmd/auth/login_test.go @@ -796,6 +796,12 @@ func TestValidateDiscoveryFlagCompatibility(t *testing.T) { flagVal: "true", wantErr: "--configure-serverless requires --host to be specified", }, + { + name: "resource is incompatible", + setFlag: "resource", + flagVal: "https://workspace.test/ai-gateway/mcp-services/system.ai.github", + wantErr: "--resource requires --host to be specified", + }, { name: "no flags set is ok", }, @@ -807,6 +813,7 @@ func TestValidateDiscoveryFlagCompatibility(t *testing.T) { cmd.Flags().String("workspace-id", "", "") cmd.Flags().Bool("configure-cluster", false, "") cmd.Flags().Bool("configure-serverless", false, "") + cmd.Flags().StringArray("resource", nil, "") if tt.setFlag != "" { require.NoError(t, cmd.Flags().Set(tt.setFlag, tt.flagVal)) diff --git a/libs/auth/u2m/persistent_auth.go b/libs/auth/u2m/persistent_auth.go index 8f5bd723fe7..a87e37783f4 100644 --- a/libs/auth/u2m/persistent_auth.go +++ b/libs/auth/u2m/persistent_auth.go @@ -10,6 +10,7 @@ import ( "fmt" "net" "net/http" + "net/url" "strings" "time" @@ -112,6 +113,13 @@ type PersistentAuth struct { // scopes is the list of OAuth scopes to request. scopes []string + // resources is the list of RFC 8707 resource indicators to send on the + // authorization request. The Databricks /oidc authorize endpoint reads + // these to scope the login to a specific protected resource (e.g. an AI + // Gateway MCP connection), so it can drive that resource's own login + // before issuing the authorization code. Empty means an unrestricted login. + resources []string + // disableOfflineAccess controls whether offline_access scope is requested. // When true, offline_access will NOT be automatically added to scopes, // meaning the token will not include a refresh token. @@ -194,6 +202,15 @@ func WithScopes(scopes []string) PersistentAuthOption { } } +// WithResources sets the RFC 8707 resource indicators for the PersistentAuth. +// Each value is added as a `resource` query parameter on the authorization +// request. +func WithResources(resources []string) PersistentAuthOption { + return func(a *PersistentAuth) { + a.resources = resources + } +} + // WithDisableOfflineAccess controls whether offline_access scope is requested. func WithDisableOfflineAccess(disable bool) PersistentAuthOption { return func(a *PersistentAuth) { @@ -623,10 +640,14 @@ func (a *PersistentAuth) oauth2Config() (*oauth2.Config, error) { if err != nil { return nil, fmt.Errorf("fetching OAuth endpoints: %w", err) } + authURL, err := appendResources(endpoints.AuthorizationEndpoint, a.resources) + if err != nil { + return nil, err + } return &oauth2.Config{ ClientID: a.clientID, Endpoint: oauth2.Endpoint{ - AuthURL: endpoints.AuthorizationEndpoint, + AuthURL: authURL, TokenURL: endpoints.TokenEndpoint, AuthStyle: oauth2.AuthStyleInParams, }, @@ -635,6 +656,26 @@ func (a *PersistentAuth) oauth2Config() (*oauth2.Config, error) { }, nil } +// appendResources adds RFC 8707 `resource` indicators to an authorization +// endpoint URL, preserving any query parameters the endpoint already carries. +// The oauth2 library appends its own parameters (client_id, PKCE, etc.) after +// these when it builds the final authorization URL. +func appendResources(authURL string, resources []string) (string, error) { + if len(resources) == 0 { + return authURL, nil + } + u, err := url.Parse(authURL) + if err != nil { + return "", fmt.Errorf("parsing authorization endpoint: %w", err) + } + q := u.Query() + for _, r := range resources { + q.Add("resource", r) + } + u.RawQuery = q.Encode() + return u.String(), nil +} + func (a *PersistentAuth) stateAndPKCE() (string, *authhandler.PKCEParams, error) { verifier, err := a.randomString(64) if err != nil { diff --git a/libs/auth/u2m/persistent_auth_test.go b/libs/auth/u2m/persistent_auth_test.go index 1ec02cf1750..22ccfd43828 100644 --- a/libs/auth/u2m/persistent_auth_test.go +++ b/libs/auth/u2m/persistent_auth_test.go @@ -8,6 +8,7 @@ import ( "net/http" "net/http/httptest" "net/url" + "reflect" "strings" "testing" "time" @@ -168,6 +169,58 @@ func TestPersistentAuthClientID(t *testing.T) { } } +func TestPersistentAuthResources(t *testing.T) { + tests := []struct { + name string + opts []PersistentAuthOption + want []string + }{ + { + name: "none", + want: nil, + }, + { + name: "single", + opts: []PersistentAuthOption{WithResources([]string{"https://workspace.test/ai-gateway/mcp-services/system.ai.github"})}, + want: []string{"https://workspace.test/ai-gateway/mcp-services/system.ai.github"}, + }, + { + name: "multiple", + opts: []PersistentAuthOption{WithResources([]string{"https://a.test/r1", "https://b.test/r2"})}, + want: []string{"https://a.test/r1", "https://b.test/r2"}, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + arg, err := NewBasicWorkspaceOAuthArgument("https://workspace.test") + if err != nil { + t.Fatalf("NewBasicWorkspaceOAuthArgument(): %v", err) + } + opts := append([]PersistentAuthOption{ + WithOAuthArgument(arg), + WithOAuthEndpointSupplier(MockOAuthEndpointSupplier{}), + }, tt.opts...) + p, err := NewPersistentAuth(t.Context(), opts...) + if err != nil { + t.Fatalf("NewPersistentAuth(): %v", err) + } + cfg, err := p.oauth2Config() + if err != nil { + t.Fatalf("oauth2Config(): %v", err) + } + parsed, err := url.Parse(cfg.Endpoint.AuthURL) + if err != nil { + t.Fatalf("parsing AuthURL %q: %v", cfg.Endpoint.AuthURL, err) + } + got := parsed.Query()["resource"] + if !reflect.DeepEqual(got, tt.want) { + t.Errorf("resource params = %v, want %v (AuthURL=%q)", got, tt.want, cfg.Endpoint.AuthURL) + } + }) + } +} + func TestToken_RefreshesExpiredAccessToken(t *testing.T) { ctx := t.Context() expectedKey := "https://accounts.cloud.databricks.test/oidc/accounts/xyz"