Allow OAuth U2M logins to send RFC 8707 resource indicators

.nextchanges/cli/u2m-resource-indicator.md

@@ -0,0 +1 @@
+* Allow OAuth U2M logins to send RFC 8707 `resource` indicators on the authorization request with the repeatable `--resource` flag (requires `--host`). ([#6621](https://github.com/databricks/cli/pull/6621))

cmd/auth/login.go

@@ -133,6 +133,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,
@@ -145,6 +146,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
 
@@ -315,6 +318,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
@@ -624,6 +630,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

cmd/auth/login_test.go

@@ -767,6 +767,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",
 		},
@@ -778,6 +784,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))

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) {
@@ -626,10 +643,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,
 		},
@@ -638,6 +659,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 {

libs/auth/u2m/persistent_auth_test.go

@@ -8,6 +8,7 @@ import (
 	"net/http"
 	"net/http/httptest"
 	"net/url"
+	"reflect"
 	"strings"
 	"testing"
 	"time"
@@ -165,6 +166,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"