Add support to authenticate with Account-wide token federation (#1219)

## What changes are proposed in this pull request?

This PR adds support to authenticate with [Account-wide token
federation](https://docs.databricks.com/aws/en/dev-tools/auth/oauth-federation#account-wide-token-federation)
from the following auth methods: `env-oidc`, `file-oidc`, and
`github-oidc`.

The PR also slightly re-organize the code by moving the OIDC token
source and Github IDTokenSource in the `oidc` package.

## How is this tested?

Unit test + local validation.

Author: renaudhartert-db

NEXT_CHANGELOG.md

@@ -4,6 +4,9 @@
 
 ### New Features and Improvements
 
+- Add support to authenticate with Account-wide token federation from the 
+  following auth methods: `env-oidc`, `file-oidc`, and `github-oidc`.  
+
 ### Bug Fixes
 
 ### Documentation

config/auth_azure_github_oidc.go

@@ -7,6 +7,7 @@ import (
 	"time"
 
 	"github.com/databricks/databricks-sdk-go/config/credentials"
+	"github.com/databricks/databricks-sdk-go/config/experimental/auth/oidc"
 	"github.com/databricks/databricks-sdk-go/httpclient"
 	"golang.org/x/oauth2"
 )
@@ -26,10 +27,11 @@ func (c AzureGithubOIDCCredentials) Configure(ctx context.Context, cfg *Config)
 	if !cfg.IsAzure() || cfg.AzureClientID == "" || cfg.Host == "" || cfg.AzureTenantID == "" || cfg.ActionsIDTokenRequestURL == "" || cfg.ActionsIDTokenRequestToken == "" {
 		return nil, nil
 	}
-	supplier := githubIDTokenSource{actionsIDTokenRequestURL: cfg.ActionsIDTokenRequestURL,
-		actionsIDTokenRequestToken: cfg.ActionsIDTokenRequestToken,
-		refreshClient:              cfg.refreshClient,
-	}
+	supplier := oidc.NewGithubIDTokenSource(
+		cfg.refreshClient,
+		cfg.ActionsIDTokenRequestURL,
+		cfg.ActionsIDTokenRequestToken,
+	)
 
 	idToken, err := supplier.IDToken(ctx, "api://AzureADTokenExchange")
 	if err != nil {

config/auth_default.go

@@ -35,28 +35,28 @@ func buildOidcTokenCredentialStrategies(cfg *Config) []CredentialsStrategy {
 		},
 		{
 			name: "github-oidc",
-			tokenSource: &githubIDTokenSource{
-				actionsIDTokenRequestURL:   cfg.ActionsIDTokenRequestURL,
-				actionsIDTokenRequestToken: cfg.ActionsIDTokenRequestToken,
-				refreshClient:              cfg.refreshClient,
-			},
+			tokenSource: oidc.NewGithubIDTokenSource(
+				cfg.refreshClient,
+				cfg.ActionsIDTokenRequestURL,
+				cfg.ActionsIDTokenRequestToken,
+			),
 		},
 		// Add new providers at the end of the list
 	}
 
 	strategies := []CredentialsStrategy{}
 	for _, idTokenSource := range idTokenSources {
-		oidcConfig := DatabricksOIDCTokenSourceConfig{
+		oidcConfig := oidc.DatabricksOIDCTokenSourceConfig{
 			ClientID:              cfg.ClientID,
 			Host:                  cfg.CanonicalHostName(),
 			TokenEndpointProvider: cfg.getOidcEndpoints,
 			Audience:              cfg.TokenAudience,
-			IdTokenSource:         idTokenSource.tokenSource,
+			IDTokenSource:         idTokenSource.tokenSource,
 		}
 		if cfg.IsAccountClient() {
 			oidcConfig.AccountID = cfg.AccountID
 		}
-		tokenSource := NewDatabricksOIDCTokenSource(oidcConfig)
+		tokenSource := oidc.NewDatabricksOIDCTokenSource(oidcConfig)
 		strategies = append(strategies, NewTokenSourceStrategy(idTokenSource.name, tokenSource))
 	}
 	return strategies

config/id_token_source_github_oidc.go renamed to config/experimental/auth/oidc/github.go

@@ -1,15 +1,25 @@
-package config
+package oidc
 
 import (
 	"context"
 	"errors"
 	"fmt"
 
-	"github.com/databricks/databricks-sdk-go/config/experimental/auth/oidc"
 	"github.com/databricks/databricks-sdk-go/httpclient"
 	"github.com/databricks/databricks-sdk-go/logger"
 )
 
+// NewGithubIDTokenSource returns a new IDTokenSource that retrieves an IDToken
+// from the Github Actions environment. This IDTokenSource is only valid when
+// running in Github Actions with OIDC enabled.
+func NewGithubIDTokenSource(client *httpclient.ApiClient, actionsIDTokenRequestURL, actionsIDTokenRequestToken string) IDTokenSource {
+	return &githubIDTokenSource{
+		actionsIDTokenRequestURL:   actionsIDTokenRequestURL,
+		actionsIDTokenRequestToken: actionsIDTokenRequestToken,
+		refreshClient:              client,
+	}
+}
+
 // githubIDTokenSource retrieves JWT Tokens from Github Actions.
 type githubIDTokenSource struct {
 	actionsIDTokenRequestURL   string
@@ -19,7 +29,7 @@ type githubIDTokenSource struct {
 
 // IDToken returns a JWT Token for the specified audience. It will return
 // an error if not running in GitHub Actions.
-func (g *githubIDTokenSource) IDToken(ctx context.Context, audience string) (*oidc.IDToken, error) {
+func (g *githubIDTokenSource) IDToken(ctx context.Context, audience string) (*IDToken, error) {
 	if g.actionsIDTokenRequestURL == "" {
 		logger.Debugf(ctx, "Missing ActionsIDTokenRequestURL, likely not calling from a Github action")
 		return nil, errors.New("missing ActionsIDTokenRequestURL")
@@ -29,7 +39,7 @@ func (g *githubIDTokenSource) IDToken(ctx context.Context, audience string) (*oi
 		return nil, errors.New("missing ActionsIDTokenRequestToken")
 	}
 
-	resp := &oidc.IDToken{}
+	resp := &IDToken{}
 	requestUrl := g.actionsIDTokenRequestURL
 	if audience != "" {
 		requestUrl = fmt.Sprintf("%s&audience=%s", requestUrl, audience)

config/id_token_source_github_oidc_test.go renamed to config/experimental/auth/oidc/github_test.go

@@ -1,11 +1,10 @@
-package config
+package oidc
 
 import (
 	"context"
 	"net/http"
 	"testing"
 
-	"github.com/databricks/databricks-sdk-go/config/experimental/auth/oidc"
 	"github.com/databricks/databricks-sdk-go/httpclient"
 	"github.com/databricks/databricks-sdk-go/httpclient/fixtures"
 	"github.com/google/go-cmp/cmp"
@@ -18,7 +17,7 @@ func TestGithubIDTokenSource(t *testing.T) {
 		tokenRequestToken string
 		audience          string
 		httpTransport     http.RoundTripper
-		wantToken         *oidc.IDToken
+		wantToken         *IDToken
 		wantErrPrefix     *string
 	}{
 		{
@@ -60,7 +59,7 @@ func TestGithubIDTokenSource(t *testing.T) {
 					Response: `{"value": "id-token-42"}`,
 				},
 			},
-			wantToken: &oidc.IDToken{
+			wantToken: &IDToken{
 				Value: "id-token-42",
 			},
 		},

config/auth_databricks_oidc.go renamed to config/experimental/auth/oidc/tokensource.go

@@ -1,42 +1,49 @@
-package config
+package oidc
 
 import (
 	"context"
 	"errors"
 	"net/url"
 
 	"github.com/databricks/databricks-sdk-go/config/experimental/auth"
-	"github.com/databricks/databricks-sdk-go/config/experimental/auth/oidc"
 	"github.com/databricks/databricks-sdk-go/credentials/u2m"
 	"github.com/databricks/databricks-sdk-go/logger"
 	"golang.org/x/oauth2"
 	"golang.org/x/oauth2/clientcredentials"
 )
 
-// Creates a new Databricks OIDC TokenSource.
-func NewDatabricksOIDCTokenSource(cfg DatabricksOIDCTokenSourceConfig) auth.TokenSource {
-	return &databricksOIDCTokenSource{
-		cfg: cfg,
-	}
-}
-
-// Config for Databricks OIDC TokenSource.
+// DatabricksOIDCTokenSourceConfig is the configuration for a Databricks OIDC
+// TokenSource.
 type DatabricksOIDCTokenSourceConfig struct {
-	// ClientID is the client ID of the Databricks OIDC application. For
-	// Databricks Service Principal, this is the Application ID of the Service Principal.
+	// ClientID of the Databricks OIDC application. It corresponds to the
+	// Application ID of the Databricks Service Principal.
+	//
+	// This field is only required for Workload Identity Federation and should
+	// be empty for Account-wide token federation.
 	ClientID string
-	// [Optional] AccountID is the account ID of the Databricks Account.
-	// This is only used for Account level tokens.
+
+	// AccountID is the account ID of the Databricks Account. This field is
+	// only required for Account-wide token federation.
 	AccountID string
+
 	// Host is the host of the Databricks account or workspace.
 	Host string
-	// TokenEndpointProvider returns the token endpoint for the Databricks OIDC application.
+
+	// TokenEndpointProvider returns the token endpoint for the Databricks OIDC
+	// application.
 	TokenEndpointProvider func(ctx context.Context) (*u2m.OAuthAuthorizationServer, error)
+
 	// Audience is the audience of the Databricks OIDC application.
 	// This is only used for Workspace level tokens.
 	Audience string
-	// IdTokenSource returns the IDToken to be used for the token exchange.
-	IdTokenSource oidc.IDTokenSource
+
+	// IDTokenSource returns the IDToken to be used for the token exchange.
+	IDTokenSource IDTokenSource
+}
+
+// NewDatabricksOIDCTokenSource returns a new Databricks OIDC TokenSource.
+func NewDatabricksOIDCTokenSource(cfg DatabricksOIDCTokenSourceConfig) auth.TokenSource {
+	return &databricksOIDCTokenSource{cfg: cfg}
 }
 
 // databricksOIDCTokenSource is a auth.TokenSource which exchanges a token using
@@ -47,10 +54,6 @@ type databricksOIDCTokenSource struct {
 
 // Token implements [TokenSource.Token]
 func (w *databricksOIDCTokenSource) Token(ctx context.Context) (*oauth2.Token, error) {
-	if w.cfg.ClientID == "" {
-		logger.Debugf(ctx, "Missing ClientID")
-		return nil, errors.New("missing ClientID")
-	}
 	if w.cfg.Host == "" {
 		logger.Debugf(ctx, "Missing Host")
 		return nil, errors.New("missing Host")
@@ -59,8 +62,17 @@ func (w *databricksOIDCTokenSource) Token(ctx context.Context) (*oauth2.Token, e
 	if err != nil {
 		return nil, err
 	}
+
+	if w.cfg.ClientID == "" {
+		logger.Debugf(ctx, "No ClientID provided, authenticating with Account-wide token federation")
+	} else {
+		logger.Debugf(ctx, "Client ID provided, authenticating with Workload Identity Federation")
+	}
+
+	// TODO: The audience is a concept of the IDToken that should likely be
+	// configured when the IDTokenSource is created.
 	audience := w.determineAudience(endpoints)
-	idToken, err := w.cfg.IdTokenSource.IDToken(ctx, audience)
+	idToken, err := w.cfg.IDTokenSource.IDToken(ctx, audience)
 	if err != nil {
 		return nil, err
 	}

config/auth_databricks_oidc_test.go renamed to config/experimental/auth/oidc/tokensource_test.go

@@ -1,19 +1,27 @@
-package config
+package oidc
 
 import (
 	"context"
 	"errors"
 	"net/http"
 	"net/url"
+	"strings"
 	"testing"
 
-	"github.com/databricks/databricks-sdk-go/config/experimental/auth/oidc"
 	"github.com/databricks/databricks-sdk-go/credentials/u2m"
 	"github.com/databricks/databricks-sdk-go/httpclient/fixtures"
 	"github.com/google/go-cmp/cmp"
 	"golang.org/x/oauth2"
 )
 
+func errPrefix(s string) *string {
+	return &s
+}
+
+func hasPrefix(err error, prefix string) bool {
+	return strings.HasPrefix(err.Error(), prefix)
+}
+
 func TestDatabricksOidcTokenSource(t *testing.T) {
 	testCases := []struct {
 		desc                 string
@@ -35,12 +43,6 @@ func TestDatabricksOidcTokenSource(t *testing.T) {
 			tokenAudience: "token-audience",
 			wantErrPrefix: errPrefix("missing Host"),
 		},
-		{
-			desc:          "missing client ID",
-			host:          "http://host.com",
-			tokenAudience: "token-audience",
-			wantErrPrefix: errPrefix("missing ClientID"),
-		},
 		{
 			desc: "token provider error",
 
@@ -104,7 +106,7 @@ func TestDatabricksOidcTokenSource(t *testing.T) {
 			wantErrPrefix: errPrefix("oauth2: server response missing access_token"),
 		},
 		{
-			desc:          "success workspace",
+			desc:          "success WIF workspace",
 			clientID:      "client-id",
 			host:          "http://host.com",
 			tokenAudience: "token-audience",
@@ -140,7 +142,7 @@ func TestDatabricksOidcTokenSource(t *testing.T) {
 			wantToken:    "test-auth-token",
 		},
 		{
-			desc:          "success account",
+			desc:          "success WIF account",
 			clientID:      "client-id",
 			accountID:     "ac123",
 			host:          "https://accounts.databricks.com",
@@ -230,6 +232,40 @@ func TestDatabricksOidcTokenSource(t *testing.T) {
 			idToken:      "id-token-42",
 			wantToken:    "test-auth-token",
 		},
+		{
+			desc:          "success account-wide",
+			host:          "http://host.com",
+			tokenAudience: "token-audience",
+			oidcEndpointProvider: func(ctx context.Context) (*u2m.OAuthAuthorizationServer, error) {
+				return &u2m.OAuthAuthorizationServer{
+					TokenEndpoint: "https://host.com/oidc/v1/token",
+				}, nil
+			},
+			httpTransport: fixtures.MappingTransport{
+				"POST /oidc/v1/token": {
+
+					Status: http.StatusOK,
+					ExpectedHeaders: map[string]string{
+						"Content-Type": "application/x-www-form-urlencoded",
+					},
+					ExpectedRequest: url.Values{
+						"scope":              {"all-apis"},
+						"subject_token_type": {"urn:ietf:params:oauth:token-type:jwt"},
+						"subject_token":      {"id-token-42"},
+						"grant_type":         {"urn:ietf:params:oauth:grant-type:token-exchange"},
+					},
+					Response: map[string]string{
+						"token_type":    "access-token",
+						"access_token":  "test-auth-token",
+						"refresh_token": "refresh",
+						"expires_on":    "0",
+					},
+				},
+			},
+			wantAudience: "token-audience",
+			idToken:      "id-token-42",
+			wantToken:    "test-auth-token",
+		},
 	}
 
 	for _, tc := range testCases {
@@ -241,9 +277,9 @@ func TestDatabricksOidcTokenSource(t *testing.T) {
 				Host:                  tc.host,
 				TokenEndpointProvider: tc.oidcEndpointProvider,
 				Audience:              tc.tokenAudience,
-				IdTokenSource: oidc.IDTokenSourceFn(func(ctx context.Context, aud string) (*oidc.IDToken, error) {
+				IDTokenSource: IDTokenSourceFn(func(ctx context.Context, aud string) (*IDToken, error) {
 					gotAudience = aud
-					return &oidc.IDToken{Value: tc.idToken}, tc.tokenProviderError
+					return &IDToken{Value: tc.idToken}, tc.tokenProviderError
 				}),
 			}