More comments

Author: hectorcast-db

config/auth_databricks_oidc.go

@@ -13,7 +13,7 @@ import (
 )
 
 // Creates a new Databricks OIDC TokenSource.
-func NewDatabricksOIDCTokenSource(cfg *DatabricksOIDCTokenSourceConfig) auth.TokenSource {
+func NewDatabricksOIDCTokenSource(cfg DatabricksOIDCTokenSourceConfig) auth.TokenSource {
 	return &databricksOIDCTokenSource{
 		cfg: cfg,
 	}
@@ -27,21 +27,21 @@ type DatabricksOIDCTokenSourceConfig struct {
 	// [Optional] AccountID is the account ID of the Databricks Account.
 	// This is only used for Account level tokens.
 	AccountID string
-	// Host is the host of the Databricks cluster.
+	// Host is the host of the Databricks account or workspace.
 	Host string
-	// TokenEndpointProvider is a function that 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 is a function that returns the ID token to be used for the token exchange.
+	// IdTokenSource returns the IDToken to be used for the token exchange.
 	IdTokenSource IDTokenSource
 }
 
 // databricksOIDCTokenSource is a auth.TokenSource which exchanges a token using
 // Workload Identity Federation.
 type databricksOIDCTokenSource struct {
-	cfg *DatabricksOIDCTokenSourceConfig
+	cfg DatabricksOIDCTokenSourceConfig
 }
 
 // Token implements [TokenSource.Token]

config/auth_databricks_oidc_test.go

@@ -251,7 +251,7 @@ func TestDatabricksOidcTokenSource(t *testing.T) {
 				err:     tc.tokenProviderError,
 			}
 
-			cfg := &DatabricksOIDCTokenSourceConfig{
+			cfg := DatabricksOIDCTokenSourceConfig{
 				ClientID:              tc.clientID,
 				AccountID:             tc.accountID,
 				Host:                  tc.host,

config/auth_default.go

@@ -10,56 +10,49 @@ import (
 )
 
 // Constructs all Databricks OIDC Credentials Strategies
-func buildOidcTokenCredentialStrategies(cfg *Config) ([]CredentialsStrategy, error) {
-	// Maps in Go are unordered, so we need to maintain an order of the strategies.
-	idTokenSourceOrder := []string{
-		"github-oidc",
-		// Add new providers at the end of the list
+func buildOidcTokenCredentialStrategies(cfg *Config) []CredentialsStrategy {
+	type namedIdTokenSource struct {
+		name        string
+		tokenSource IDTokenSource
 	}
-	idTokenSources := map[string]IDTokenSource{
-		"github-oidc": &githubIDTokenSource{
-			actionsIDTokenRequestURL:   cfg.ActionsIDTokenRequestURL,
-			actionsIDTokenRequestToken: cfg.ActionsIDTokenRequestToken,
-			refreshClient:              cfg.refreshClient,
+	idTokenSources := []namedIdTokenSource{
+		{
+			name: "github-oidc",
+			tokenSource: &githubIDTokenSource{
+				actionsIDTokenRequestURL:   cfg.ActionsIDTokenRequestURL,
+				actionsIDTokenRequestToken: cfg.ActionsIDTokenRequestToken,
+				refreshClient:              cfg.refreshClient,
+			},
 		},
 		// Add new providers at the end of the list
 	}
-
 	strategies := []CredentialsStrategy{}
-	for _, name := range idTokenSourceOrder {
-		provider, ok := idTokenSources[name]
-		if !ok {
-			return nil, fmt.Errorf("no provider found for %s", name)
-		}
-		oidcConfig := &DatabricksOIDCTokenSourceConfig{
+	for _, idTokenSource := range idTokenSources {
+		oidcConfig := DatabricksOIDCTokenSourceConfig{
 			ClientID:              cfg.ClientID,
-			Host:                  cfg.Host,
+			Host:                  cfg.CanonicalHostName(),
 			TokenEndpointProvider: cfg.getOidcEndpoints,
 			Audience:              cfg.TokenAudience,
-			IdTokenSource:         provider,
+			IdTokenSource:         idTokenSource.tokenSource,
 		}
 		if cfg.IsAccountClient() {
 			oidcConfig.AccountID = cfg.AccountID
 		}
 		tokenSource := NewDatabricksOIDCTokenSource(oidcConfig)
-		strategies = append(strategies, NewTokenSourceStrategy(name, tokenSource))
+		strategies = append(strategies, NewTokenSourceStrategy(idTokenSource.name, tokenSource))
 	}
-	return strategies, nil
+	return strategies
 }
 
-func buildDefaultStrategies(cfg *Config) ([]CredentialsStrategy, error) {
+func buildDefaultStrategies(cfg *Config) []CredentialsStrategy {
 	strategies := []CredentialsStrategy{}
 	strategies = append(strategies,
 		PatCredentials{},
 		BasicCredentials{},
 		M2mCredentials{},
 		DatabricksCliCredentials,
 		MetadataServiceCredentials{})
-	oidcStrategies, err := buildOidcTokenCredentialStrategies(cfg)
-	if err != nil {
-		return nil, err
-	}
-	strategies = append(strategies, oidcStrategies...)
+	strategies = append(strategies, buildOidcTokenCredentialStrategies(cfg)...)
 	strategies = append(strategies,
 		// Attempt to configure auth from most specific to most generic (the Azure CLI).
 		AzureGithubOIDCCredentials{},
@@ -69,7 +62,7 @@ func buildDefaultStrategies(cfg *Config) ([]CredentialsStrategy, error) {
 		// Attempt to configure auth from most specific to most generic (Google Application Default Credentials).
 		GoogleCredentials{},
 		GoogleDefaultCredentials{})
-	return strategies, nil
+	return strategies
 }
 
 type DefaultCredentials struct {
@@ -90,11 +83,11 @@ var errorMessage = fmt.Sprintf("cannot configure default credentials, please che
 var ErrCannotConfigureAuth = errors.New(errorMessage)
 
 func (c *DefaultCredentials) Configure(ctx context.Context, cfg *Config) (credentials.CredentialsProvider, error) {
-	strategies, err := buildDefaultStrategies(cfg)
+	err := cfg.EnsureResolved()
 	if err != nil {
 		return nil, err
 	}
-	for _, p := range strategies {
+	for _, p := range buildDefaultStrategies(cfg) {
 		if cfg.AuthType != "" && p.Name() != cfg.AuthType {
 			// ignore other auth types if one is explicitly enforced
 			logger.Infof(ctx, "Ignoring %s auth, because %s is preferred", p.Name(), cfg.AuthType)

config/oauth_visitors.go

@@ -35,14 +35,14 @@ func serviceToServiceVisitor(primary, secondary oauth2.TokenSource, secondaryHea
 
 // The same as serviceToServiceVisitor, but without a secondary token source.
 func refreshableVisitor(inner oauth2.TokenSource) func(r *http.Request) error {
-	return refreshableAuthVisitor(authconv.AuthTokenSource(inner), context.Background())
+	return refreshableAuthVisitor(authconv.AuthTokenSource(inner))
 }
 
 // The same as serviceToServiceVisitor, but without a secondary token source.
-func refreshableAuthVisitor(inner auth.TokenSource, ctx context.Context) func(r *http.Request) error {
+func refreshableAuthVisitor(inner auth.TokenSource) func(r *http.Request) error {
 	cts := auth.NewCachedTokenSource(inner)
 	return func(r *http.Request) error {
-		inner, err := cts.Token(ctx)
+		inner, err := cts.Token(context.Background())
 		if err != nil {
 			return fmt.Errorf("inner token: %w", err)
 		}