diff --git a/cmd/auth/login.go b/cmd/auth/login.go index bd3123b7772..eff56db2f76 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 @@ -279,6 +282,9 @@ a new profile is created. if clientID == "" { clientID = u2mClientIDFromProfile(existingProfile) } + if len(resources) == 0 { + resources = u2mResourcesFromProfile(existingProfile) + } // If no host is available from any source, use the discovery flow // via login.databricks.com. @@ -331,6 +337,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 @@ -422,6 +431,11 @@ a new profile is created. if err != nil { return err } + // `resources` is a CLI-only profile key (not an SDK config attribute), so + // it is written separately from SaveToProfile. + if err := databrickscfg.SaveResourcesToProfile(ctx, profileName, env.Get(ctx, "DATABRICKS_CONFIG_FILE"), resources); err != nil { + return err + } } if err := storeLoginToken(ctx, tokenStore, mode, oauthArgument, token); err != nil { @@ -643,6 +657,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 @@ -825,6 +843,15 @@ func u2mClientIDFromProfile(p *profile.Profile) string { return p.ClientID } +// u2mResourcesFromProfile returns the RFC 8707 resource indicators saved on a +// databricks-cli-auth profile, so `login`/`auth token` re-request the same ones. +func u2mResourcesFromProfile(p *profile.Profile) []string { + if p == nil || p.AuthType != authTypeDatabricksCLI { + return nil + } + return splitScopes(p.Resources) +} + // promptForWorkspaceSelection lists workspaces for a SPOG account and lets the // user pick one. Returns the selected workspace ID or empty string if skipped. // This is best-effort: errors are returned to the caller for logging, not shown diff --git a/cmd/auth/login_test.go b/cmd/auth/login_test.go index 688b0a66dbc..1c4428596dd 100644 --- a/cmd/auth/login_test.go +++ b/cmd/auth/login_test.go @@ -516,6 +516,40 @@ func TestU2MClientIDFromProfile(t *testing.T) { } } +func TestU2MResourcesFromProfile(t *testing.T) { + tests := []struct { + name string + profile *profile.Profile + want []string + }{ + {name: "no profile"}, + { + name: "implicit auth type", + profile: &profile.Profile{Resources: "https://workspace.test/ai-gateway/mcp/system.ai.github"}, + }, + { + name: "M2M auth type", + profile: &profile.Profile{AuthType: "oauth-m2m", Resources: "https://workspace.test/ai-gateway/mcp/system.ai.github"}, + }, + { + name: "U2M auth type single resource", + profile: &profile.Profile{AuthType: authTypeDatabricksCLI, Resources: "https://workspace.test/ai-gateway/mcp/system.ai.github"}, + want: []string{"https://workspace.test/ai-gateway/mcp/system.ai.github"}, + }, + { + name: "U2M auth type multiple resources", + profile: &profile.Profile{AuthType: authTypeDatabricksCLI, Resources: "https://workspace.test/ai-gateway/mcp/system.ai.github, https://workspace.test/ai-gateway/mcp/system.ai.slack"}, + want: []string{"https://workspace.test/ai-gateway/mcp/system.ai.github", "https://workspace.test/ai-gateway/mcp/system.ai.slack"}, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + assert.Equal(t, tt.want, u2mResourcesFromProfile(tt.profile)) + }) + } +} + func TestRunHostDiscovery_NoHost(t *testing.T) { ctx := t.Context() args := &auth.AuthArguments{} @@ -796,6 +830,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 +847,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/cmd/auth/token.go b/cmd/auth/token.go index 300ebe8986c..047370b4bed 100644 --- a/cmd/auth/token.go +++ b/cmd/auth/token.go @@ -267,6 +267,9 @@ func loadToken(ctx context.Context, args loadTokenArgs) (*oauth2.Token, error) { if clientID := u2mClientIDFromProfile(existingProfile); clientID != "" { allArgs = append(allArgs, u2m.WithClientID(clientID)) } + if resources := u2mResourcesFromProfile(existingProfile); len(resources) > 0 { + allArgs = append(allArgs, u2m.WithResources(resources)) + } allArgs = append(allArgs, u2m.WithOAuthArgument(oauthArgument)) persistentAuth, err := u2m.NewPersistentAuth(ctx, allArgs...) if err != nil { @@ -423,6 +426,10 @@ func runInlineLogin(ctx context.Context, profiler profile.Profiler, tokenStore s scopesList = splitScopes(existingProfile.Scopes) } + // Preserve RFC 8707 resource indicators from the existing profile so the + // inline login requests the same resources the user previously configured. + resourcesList := u2mResourcesFromProfile(existingProfile) + oauthArgument, err := loginArgs.ToOAuthArgument() if err != nil { return "", nil, err @@ -437,6 +444,9 @@ func runInlineLogin(ctx context.Context, profiler profile.Profiler, tokenStore s if len(scopesList) > 0 { persistentAuthOpts = append(persistentAuthOpts, u2m.WithScopes(scopesList)) } + if len(resourcesList) > 0 { + persistentAuthOpts = append(persistentAuthOpts, u2m.WithResources(resourcesList)) + } persistentAuth, err := u2m.NewPersistentAuth(ctx, persistentAuthOpts...) if err != nil { return "", nil, err @@ -467,6 +477,11 @@ func runInlineLogin(ctx context.Context, profiler profile.Profiler, tokenStore s if err != nil { return "", nil, err } + // `resources` is a CLI-only profile key (not an SDK config attribute), so + // it is written separately from SaveToProfile. + if err := databrickscfg.SaveResourcesToProfile(ctx, profileName, env.Get(ctx, "DATABRICKS_CONFIG_FILE"), resourcesList); err != nil { + return "", nil, err + } if err := storeLoginToken(ctx, tokenStore, mode, oauthArgument, token); err != nil { return "", nil, err } 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..d4cecd0f200 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,60 @@ 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" diff --git a/libs/databrickscfg/ops.go b/libs/databrickscfg/ops.go index 2eb58c2cc78..9ff963f0165 100644 --- a/libs/databrickscfg/ops.go +++ b/libs/databrickscfg/ops.go @@ -258,6 +258,30 @@ func SetConfiguredAuthStorage(ctx context.Context, value, configFilePath string) return writeConfigFile(ctx, configFile) } +// SaveResourcesToProfile writes (or clears) the `resources` key on a profile +// section. `resources` is a comma-separated list of RFC 8707 resource indicators +// requested during U2M login. It is a CLI-only profile key (not a +// databricks-sdk-go config attribute), so it is written here rather than through +// SaveToProfile: `databricks auth token` re-reads it to request the same +// resources, and a change in the requested resources is reflected in the profile. +// An empty list removes the key. The profile section must already exist. +func SaveResourcesToProfile(ctx context.Context, profileName, configFilePath string, resources []string) error { + configFile, err := loadOrCreateConfigFile(ctx, configFilePath) + if err != nil { + return err + } + section, err := configFile.GetSection(profileName) + if err != nil { + return fmt.Errorf("profile %q not found: %w", profileName, err) + } + if len(resources) == 0 { + section.DeleteKey("resources") + } else { + section.Key("resources").SetValue(strings.Join(resources, ",")) + } + return writeConfigFile(ctx, configFile) +} + // ClearDefaultProfile removes the default_profile key from the [__settings__] // section if the current default matches the given profile name. func ClearDefaultProfile(ctx context.Context, profileName, configFilePath string) error { diff --git a/libs/databrickscfg/ops_test.go b/libs/databrickscfg/ops_test.go index 480054bfed6..8fe1956a81f 100644 --- a/libs/databrickscfg/ops_test.go +++ b/libs/databrickscfg/ops_test.go @@ -185,6 +185,40 @@ default_profile = DEFAULT `, string(contents)) } +func TestSaveResourcesToProfile(t *testing.T) { + ctx := t.Context() + path := filepath.Join(t.TempDir(), "databrickscfg") + + require.NoError(t, SaveToProfile(ctx, &config.Config{ + ConfigFile: path, + Profile: "u2m", + Host: "https://foo", + })) + + resources := []string{ + "https://foo/ai-gateway/mcp/system.ai.github", + "https://foo/ai-gateway/mcp/system.ai.slack", + } + require.NoError(t, SaveResourcesToProfile(ctx, "u2m", path, resources)) + + contents, err := os.ReadFile(path) + require.NoError(t, err) + assert.Contains(t, string(contents), "resources = https://foo/ai-gateway/mcp/system.ai.github,https://foo/ai-gateway/mcp/system.ai.slack") + + // An empty list clears the key. + require.NoError(t, SaveResourcesToProfile(ctx, "u2m", path, nil)) + contents, err = os.ReadFile(path) + require.NoError(t, err) + assert.NotContains(t, string(contents), "resources") +} + +func TestSaveResourcesToProfile_MissingProfile(t *testing.T) { + ctx := t.Context() + path := filepath.Join(t.TempDir(), "databrickscfg") + err := SaveResourcesToProfile(ctx, "does-not-exist", path, []string{"https://foo/ai-gateway/mcp/system.ai.github"}) + assert.ErrorContains(t, err, `profile "does-not-exist" not found`) +} + func TestGetDefaultProfile(t *testing.T) { testCases := []struct { name string diff --git a/libs/databrickscfg/profile/file.go b/libs/databrickscfg/profile/file.go index 529b2a0c34b..0844794df3a 100644 --- a/libs/databrickscfg/profile/file.go +++ b/libs/databrickscfg/profile/file.go @@ -88,6 +88,7 @@ func (f FileProfilerImpl) LoadProfiles(ctx context.Context, fn ProfileMatchFunct HasClientCredentials: all["client_id"] != "" && all["client_secret"] != "", ClientID: all["client_id"], Scopes: all["scopes"], + Resources: all["resources"], AuthType: all["auth_type"], } if fn(profile) { diff --git a/libs/databrickscfg/profile/file_test.go b/libs/databrickscfg/profile/file_test.go index a59b2202693..82243842a94 100644 --- a/libs/databrickscfg/profile/file_test.go +++ b/libs/databrickscfg/profile/file_test.go @@ -82,6 +82,22 @@ client_id = custom-client-id assert.False(t, profiles[0].HasClientCredentials) } +func TestLoadProfilesResources(t *testing.T) { + configPath := filepath.Join(t.TempDir(), ".databrickscfg") + err := os.WriteFile(configPath, []byte(`[u2m] +host = https://workspace.test +auth_type = databricks-cli +resources = https://workspace.test/ai-gateway/mcp/system.ai.github,https://workspace.test/ai-gateway/mcp/system.ai.slack +`), 0o600) + require.NoError(t, err) + + ctx := env.Set(t.Context(), "DATABRICKS_CONFIG_FILE", configPath) + profiles, err := (FileProfilerImpl{}).LoadProfiles(ctx, MatchAllProfiles) + require.NoError(t, err) + require.Len(t, profiles, 1) + assert.Equal(t, "https://workspace.test/ai-gateway/mcp/system.ai.github,https://workspace.test/ai-gateway/mcp/system.ai.slack", profiles[0].Resources) +} + func TestLoadProfilesMatchWorkspace(t *testing.T) { ctx := t.Context() ctx = env.Set(ctx, "DATABRICKS_CONFIG_FILE", "./testdata/databrickscfg") diff --git a/libs/databrickscfg/profile/profile.go b/libs/databrickscfg/profile/profile.go index 6a2ff7b5e1e..3d32989e13d 100644 --- a/libs/databrickscfg/profile/profile.go +++ b/libs/databrickscfg/profile/profile.go @@ -19,6 +19,7 @@ type Profile struct { HasClientCredentials bool ClientID string Scopes string + Resources string AuthType string }