From bc78c11bb9ec6a96fb423b69e501ee698403ee72 Mon Sep 17 00:00:00 2001 From: Ubuntu Date: Fri, 9 May 2025 09:17:13 +0000 Subject: [PATCH 1/2] Add OIDC accout-wide --- config/auth_azure_github_oidc.go | 10 ++-- config/auth_default.go | 16 ++--- .../auth/oidc/github.go} | 18 ++++-- .../auth/oidc/github_test.go} | 7 +-- .../auth/oidc/tokensource.go} | 56 ++++++++++------- .../auth/oidc/tokensource_test.go} | 60 +++++++++++++++---- 6 files changed, 113 insertions(+), 54 deletions(-) rename config/{id_token_source_github_oidc.go => experimental/auth/oidc/github.go} (70%) rename config/{id_token_source_github_oidc_test.go => experimental/auth/oidc/github_test.go} (94%) rename config/{auth_databricks_oidc.go => experimental/auth/oidc/tokensource.go} (67%) rename config/{auth_databricks_oidc_test.go => experimental/auth/oidc/tokensource_test.go} (85%) diff --git a/config/auth_azure_github_oidc.go b/config/auth_azure_github_oidc.go index 7be69563f..8120f3b7b 100644 --- a/config/auth_azure_github_oidc.go +++ b/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 { diff --git a/config/auth_default.go b/config/auth_default.go index a53665bc4..e2be578fe 100644 --- a/config/auth_default.go +++ b/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 diff --git a/config/id_token_source_github_oidc.go b/config/experimental/auth/oidc/github.go similarity index 70% rename from config/id_token_source_github_oidc.go rename to config/experimental/auth/oidc/github.go index 93b9bb11c..5e7cfc961 100644 --- a/config/id_token_source_github_oidc.go +++ b/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) diff --git a/config/id_token_source_github_oidc_test.go b/config/experimental/auth/oidc/github_test.go similarity index 94% rename from config/id_token_source_github_oidc_test.go rename to config/experimental/auth/oidc/github_test.go index af71858d3..33a1fb213 100644 --- a/config/id_token_source_github_oidc_test.go +++ b/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", }, }, diff --git a/config/auth_databricks_oidc.go b/config/experimental/auth/oidc/tokensource.go similarity index 67% rename from config/auth_databricks_oidc.go rename to config/experimental/auth/oidc/tokensource.go index 7a5f35ef2..25eb1b6f4 100644 --- a/config/auth_databricks_oidc.go +++ b/config/experimental/auth/oidc/tokensource.go @@ -1,4 +1,4 @@ -package config +package oidc import ( "context" @@ -6,37 +6,44 @@ import ( "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 } diff --git a/config/auth_databricks_oidc_test.go b/config/experimental/auth/oidc/tokensource_test.go similarity index 85% rename from config/auth_databricks_oidc_test.go rename to config/experimental/auth/oidc/tokensource_test.go index 64077bfab..3410978c6 100644 --- a/config/auth_databricks_oidc_test.go +++ b/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 }), } From ab703f1fbf1f34e15b86f321e014915e99301403 Mon Sep 17 00:00:00 2001 From: Ubuntu Date: Fri, 9 May 2025 09:18:57 +0000 Subject: [PATCH 2/2] Update changelogs --- NEXT_CHANGELOG.md | 3 +++ 1 file changed, 3 insertions(+) diff --git a/NEXT_CHANGELOG.md b/NEXT_CHANGELOG.md index 9c0031211..fa2120b59 100644 --- a/NEXT_CHANGELOG.md +++ b/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