From 0b11b82fd0fbf3610e52686c0c9ef89fb67084c8 Mon Sep 17 00:00:00 2001 From: Hector Castejon Diaz Date: Mon, 3 Mar 2025 12:42:00 +0100 Subject: [PATCH 01/30] WIP --- config/auth_azure_github_oidc.go | 29 +-------------- config/auth_databricks_oidc.go | 64 ++++++++++++++++++++++++++++++++ config/auth_default.go | 2 + config/oidc_github.go | 38 +++++++++++++++++++ 4 files changed, 106 insertions(+), 27 deletions(-) create mode 100644 config/auth_databricks_oidc.go create mode 100644 config/oidc_github.go diff --git a/config/auth_azure_github_oidc.go b/config/auth_azure_github_oidc.go index 2f82214f2..e1e9ee5e8 100644 --- a/config/auth_azure_github_oidc.go +++ b/config/auth_azure_github_oidc.go @@ -8,7 +8,6 @@ import ( "github.com/databricks/databricks-sdk-go/config/credentials" "github.com/databricks/databricks-sdk-go/httpclient" - "github.com/databricks/databricks-sdk-go/logger" "golang.org/x/oauth2" ) @@ -27,8 +26,9 @@ func (c AzureGithubOIDCCredentials) Configure(ctx context.Context, cfg *Config) if !cfg.IsAzure() || cfg.AzureClientID == "" || cfg.Host == "" || cfg.AzureTenantID == "" { return nil, nil } + supplier := GithubOIDCTokenSupplier{cfg: cfg} - idToken, err := requestIDToken(ctx, cfg) + idToken, err := supplier.GetOIDCToken(ctx, "api://AzureADTokenExchange") if err != nil { return nil, err } @@ -47,31 +47,6 @@ func (c AzureGithubOIDCCredentials) Configure(ctx context.Context, cfg *Config) return credentials.NewOAuthCredentialsProvider(refreshableVisitor(ts), ts.Token), nil } -// requestIDToken requests an ID token from the Github Action. -func requestIDToken(ctx context.Context, cfg *Config) (string, error) { - if cfg.ActionsIDTokenRequestURL == "" { - logger.Debugf(ctx, "Missing cfg.ActionsIDTokenRequestURL, likely not calling from a Github action") - return "", nil - } - if cfg.ActionsIDTokenRequestToken == "" { - logger.Debugf(ctx, "Missing cfg.ActionsIDTokenRequestToken, likely not calling from a Github action") - return "", nil - } - - resp := struct { // anonymous struct to parse the response - Value string `json:"value"` - }{} - err := cfg.refreshClient.Do(ctx, "GET", fmt.Sprintf("%s&audience=api://AzureADTokenExchange", cfg.ActionsIDTokenRequestURL), - httpclient.WithRequestHeader("Authorization", fmt.Sprintf("Bearer %s", cfg.ActionsIDTokenRequestToken)), - httpclient.WithResponseUnmarshal(&resp), - ) - if err != nil { - return "", fmt.Errorf("failed to request ID token from %s: %w", cfg.ActionsIDTokenRequestURL, err) - } - - return resp.Value, nil -} - // azureOIDCTokenSource implements [oauth2.TokenSource] to obtain Azure auth // tokens from an ID token. type azureOIDCTokenSource struct { diff --git a/config/auth_databricks_oidc.go b/config/auth_databricks_oidc.go new file mode 100644 index 000000000..3e789e6f8 --- /dev/null +++ b/config/auth_databricks_oidc.go @@ -0,0 +1,64 @@ +package config + +import ( + "context" + "net/url" + "strings" + + "github.com/databricks/databricks-sdk-go/config/credentials" + "github.com/databricks/databricks-sdk-go/httpclient" + "github.com/databricks/databricks-sdk-go/logger" + "golang.org/x/oauth2" + "golang.org/x/oauth2/clientcredentials" +) + +type DatabricksOIDCCredentials struct{} + +// Configure implements CredentialsStrategy. +func (d DatabricksOIDCCredentials) Configure(ctx context.Context, cfg *Config) (credentials.CredentialsProvider, error) { + if cfg.Host == "" || cfg.ClientID == "" { + return nil, nil + } + + // Get the OIDC token from the environment. + audience := strings.TrimPrefix(cfg.CanonicalHostName(), "https://") + if cfg.IsAccountClient() { + audience = cfg.AccountID + } + supplier := GithubOIDCTokenSupplier{ + cfg: cfg, + } + idToken, err := supplier.GetOIDCToken(ctx, audience) + if err != nil { + return nil, err + } + if idToken == "" { + logger.Debugf(ctx, "No OIDC token found") + return nil, nil + } + + endpoints, err := oidcEndpoints(ctx, cfg) + if err != nil { + return nil, err + } + + tsConfig := clientcredentials.Config{ + ClientID: cfg.ClientID, + ClientSecret: "", + AuthStyle: oauth2.AuthStyleInParams, + TokenURL: endpoints.TokenEndpoint, + Scopes: []string{"all-apis"}, + EndpointParams: url.Values{ + "grant_type": {httpclient.JWTGrantType}, + "assertion": {idToken}, + }, + } + ts := tsConfig.TokenSource(ctx) + visitor := refreshableVisitor(ts) + return credentials.NewOAuthCredentialsProvider(visitor, ts.Token), nil +} + +// Name implements CredentialsStrategy. +func (d DatabricksOIDCCredentials) Name() string { + return "inhouse-oidc" +} diff --git a/config/auth_default.go b/config/auth_default.go index 5fc3080f7..6b97653a1 100644 --- a/config/auth_default.go +++ b/config/auth_default.go @@ -25,6 +25,8 @@ var authProviders = []CredentialsStrategy{ // Attempt to configure auth from most specific to most generic (Google Application Default Credentials). GoogleCredentials{}, GoogleDefaultCredentials{}, + + DatabricksOIDCCredentials{}, } type DefaultCredentials struct { diff --git a/config/oidc_github.go b/config/oidc_github.go new file mode 100644 index 000000000..826269b7f --- /dev/null +++ b/config/oidc_github.go @@ -0,0 +1,38 @@ +package config + +import ( + "context" + "fmt" + + "github.com/databricks/databricks-sdk-go/httpclient" + "github.com/databricks/databricks-sdk-go/logger" +) + +type GithubOIDCTokenSupplier struct { + cfg *Config +} + +// requestIDToken requests an ID token from the Github Action. +func (g *GithubOIDCTokenSupplier) GetOIDCToken(ctx context.Context, audience string) (string, error) { + if g.cfg.ActionsIDTokenRequestURL == "" { + logger.Debugf(ctx, "Missing cfg.ActionsIDTokenRequestURL, likely not calling from a Github action") + return "", nil + } + if g.cfg.ActionsIDTokenRequestToken == "" { + logger.Debugf(ctx, "Missing cfg.ActionsIDTokenRequestToken, likely not calling from a Github action") + return "", nil + } + + resp := struct { // anonymous struct to parse the response + Value string `json:"value"` + }{} + err := g.cfg.refreshClient.Do(ctx, "GET", fmt.Sprintf("%s&audience=%s", g.cfg.ActionsIDTokenRequestURL, audience), + httpclient.WithRequestHeader("Authorization", fmt.Sprintf("Bearer %s", g.cfg.ActionsIDTokenRequestToken)), + httpclient.WithResponseUnmarshal(&resp), + ) + if err != nil { + return "", fmt.Errorf("failed to request ID token from %s: %w", g.cfg.ActionsIDTokenRequestURL, err) + } + + return resp.Value, nil +} From 200f16e8e0f26f51e54565d7292a6def5edbcacf Mon Sep 17 00:00:00 2001 From: Hector Castejon Diaz Date: Tue, 18 Mar 2025 12:55:49 +0100 Subject: [PATCH 02/30] Test --- .github/workflows/integration-tests.yml | 4 +- config/auth_databricks_oidc.go | 92 ++++++++++++++++++++++--- config/config.go | 2 + internal/auth_test.go | 87 +++++++++++++++++++++++ 4 files changed, 174 insertions(+), 11 deletions(-) create mode 100644 internal/auth_test.go diff --git a/.github/workflows/integration-tests.yml b/.github/workflows/integration-tests.yml index d0e8707c1..1ab802e51 100644 --- a/.github/workflows/integration-tests.yml +++ b/.github/workflows/integration-tests.yml @@ -2,8 +2,8 @@ name: Integration Tests on: - pull_request: - types: [opened, synchronize] + # pull_request: + # types: [opened, synchronize] merge_group: diff --git a/config/auth_databricks_oidc.go b/config/auth_databricks_oidc.go index 3e789e6f8..c18562d4b 100644 --- a/config/auth_databricks_oidc.go +++ b/config/auth_databricks_oidc.go @@ -2,11 +2,15 @@ package config import ( "context" + "encoding/base64" + "encoding/json" + "fmt" + "net/http" + "net/http/httputil" "net/url" "strings" "github.com/databricks/databricks-sdk-go/config/credentials" - "github.com/databricks/databricks-sdk-go/httpclient" "github.com/databricks/databricks-sdk-go/logger" "golang.org/x/oauth2" "golang.org/x/oauth2/clientcredentials" @@ -14,6 +18,59 @@ import ( type DatabricksOIDCCredentials struct{} +type loggingTransport struct { + Transport http.RoundTripper +} + +func (t *loggingTransport) RoundTrip(req *http.Request) (*http.Response, error) { + // Log the request + dump, err := httputil.DumpRequestOut(req, true) + if err == nil { + fmt.Println("HTTP Request:") + fmt.Println(string(dump)) + } + + // Send the request + resp, err := t.Transport.RoundTrip(req) + if err != nil { + return nil, err + } + + // Log the response + dumpResp, err := httputil.DumpResponse(resp, true) + if err == nil { + fmt.Println("HTTP Response:") + fmt.Println(string(dumpResp)) + } + + return resp, nil +} + +func print(ctx context.Context, token string) { + + parts := strings.Split(token, ".") + if len(parts) != 3 { + fmt.Println("Invalid JWT format") + return + } + // Decode the payload (second part of the JWT) + payload, err := base64.RawURLEncoding.DecodeString(parts[1]) + if err != nil { + logger.Debugf(ctx, "Error decoding JWT payload: %v", err) + return + } + + // Pretty-print JSON + var prettyPayload map[string]interface{} + if err := json.Unmarshal(payload, &prettyPayload); err != nil { + logger.Debugf(ctx, "Error parsing JSON: %v", err) + return + } + + prettyJSON, _ := json.MarshalIndent(prettyPayload, "", " ") + logger.Debugf(ctx, string(prettyJSON)) +} + // Configure implements CredentialsStrategy. func (d DatabricksOIDCCredentials) Configure(ctx context.Context, cfg *Config) (credentials.CredentialsProvider, error) { if cfg.Host == "" || cfg.ClientID == "" { @@ -21,26 +78,32 @@ func (d DatabricksOIDCCredentials) Configure(ctx context.Context, cfg *Config) ( } // Get the OIDC token from the environment. - audience := strings.TrimPrefix(cfg.CanonicalHostName(), "https://") + audience := cfg.CanonicalHostName() if cfg.IsAccountClient() { audience = cfg.AccountID } + + audience = "https://github.com/databricks-eng" + supplier := GithubOIDCTokenSupplier{ cfg: cfg, } - idToken, err := supplier.GetOIDCToken(ctx, audience) + token, err := supplier.GetOIDCToken(ctx, audience) if err != nil { return nil, err } - if idToken == "" { - logger.Debugf(ctx, "No OIDC token found") - return nil, nil + if token == "" { + // logger.Debugf(ctx, "No OIDC token found") + // return nil, nil + token = "MyToken" } + print(ctx, token) endpoints, err := oidcEndpoints(ctx, cfg) if err != nil { return nil, err } + logger.Debugf(ctx, "Getting tokken for client %s", cfg.ClientID) tsConfig := clientcredentials.Config{ ClientID: cfg.ClientID, @@ -49,10 +112,21 @@ func (d DatabricksOIDCCredentials) Configure(ctx context.Context, cfg *Config) ( TokenURL: endpoints.TokenEndpoint, Scopes: []string{"all-apis"}, EndpointParams: url.Values{ - "grant_type": {httpclient.JWTGrantType}, - "assertion": {idToken}, + //"grant_type": {httpclient.JWTGrantType}, + //"assertion": {idToken}, + "subject_token_type": {"urn:ietf:params:oauth:token-type:jwt"}, + "subject_token": {token}, + "grant_type": {"urn:ietf:params:oauth:grant-type:token-exchange"}, }, } + client := &http.Client{Transport: &loggingTransport{Transport: http.DefaultTransport}} + ctx = context.WithValue(ctx, oauth2.HTTPClient, client) + + // Request the token + _, err = tsConfig.Token(ctx) + if err != nil { + return nil, err + } ts := tsConfig.TokenSource(ctx) visitor := refreshableVisitor(ts) return credentials.NewOAuthCredentialsProvider(visitor, ts.Token), nil @@ -60,5 +134,5 @@ func (d DatabricksOIDCCredentials) Configure(ctx context.Context, cfg *Config) ( // Name implements CredentialsStrategy. func (d DatabricksOIDCCredentials) Name() string { - return "inhouse-oidc" + return "databricks-wif" } diff --git a/config/config.go b/config/config.go index b1c9cd75c..41a2c3abd 100644 --- a/config/config.go +++ b/config/config.go @@ -100,6 +100,8 @@ type Config struct { ClientID string `name:"client_id" env:"DATABRICKS_CLIENT_ID" auth:"oauth" auth_types:"oauth-m2m"` ClientSecret string `name:"client_secret" env:"DATABRICKS_CLIENT_SECRET" auth:"oauth,sensitive" auth_types:"oauth-m2m"` + ApplicationID string `name:"client_id" env:"DATABRICKS_APPLICATION_ID" auth:"oauth" auth_types:"oauth-m2m"` + // Path to the Databricks CLI (version >= 0.100.0). DatabricksCliPath string `name:"databricks_cli_path" env:"DATABRICKS_CLI_PATH" auth_types:"databricks-cli"` diff --git a/internal/auth_test.go b/internal/auth_test.go new file mode 100644 index 000000000..832d10bb0 --- /dev/null +++ b/internal/auth_test.go @@ -0,0 +1,87 @@ +package internal + +import ( + "strconv" + "testing" + + "github.com/databricks/databricks-sdk-go" + "github.com/databricks/databricks-sdk-go/apierr" + "github.com/databricks/databricks-sdk-go/service/iam" + "github.com/databricks/databricks-sdk-go/service/oauth2" + "github.com/stretchr/testify/require" +) + +func TestAccWifAuth(t *testing.T) { + ctx, a := ucacctTest(t) + //testWorkspaceId := int64(470576644108500) + //testWorkspaceUrl := "https://dbc-1232e87d-9384.cloud.databricks.com" + + // Create SP + sp, err := a.ServicePrincipals.Create(ctx, iam.ServicePrincipal{ + Active: true, + DisplayName: RandomName("go-sdk-sp-"), + }) + require.NoError(t, err) + t.Cleanup(func() { + err := a.ServicePrincipals.Delete(ctx, iam.DeleteAccountServicePrincipalRequest{Id: sp.Id}) + require.True(t, err == nil || apierr.IsMissing(err)) + }) + + applicationId, err := strconv.ParseInt(sp.Id, 10, 64) + require.NoError(t, err) + + p, err := a.ServicePrincipalFederationPolicy.Create(ctx, oauth2.CreateServicePrincipalFederationPolicyRequest{ + Policy: &oauth2.FederationPolicy{ + OidcPolicy: &oauth2.OidcFederationPolicy{ + Issuer: "https://token.actions.githubusercontent.com", + Audiences: []string{ + "https://github.com/databricks-eng", + }, + Subject: "repo:databricks-eng/eng-dev-ecosystem:environment:integration-tests", + }, + }, + ServicePrincipalId: applicationId, + }) + + require.NoError(t, err) + t.Cleanup(func() { + err := a.ServicePrincipalFederationPolicy.Delete(ctx, oauth2.DeleteServicePrincipalFederationPolicyRequest{ + ServicePrincipalId: applicationId, + PolicyId: p.Uid, + }) + require.True(t, err == nil || apierr.IsMissing(err)) + }) + + // _, err = a.WorkspaceAssignment.Update(ctx, iam.UpdateWorkspaceAssignments{ + // WorkspaceId: testWorkspaceId, + // PrincipalId: applicationId, + // Permissions: []iam.WorkspacePermission{iam.WorkspacePermissionAdmin}, + // }) + + // require.NoError(t, err) + // t.Cleanup(func() { + // err := a.WorkspaceAssignment.Delete(ctx, iam.DeleteWorkspaceAssignmentRequest{ + // PrincipalId: applicationId, + // WorkspaceId: testWorkspaceId, + // }) + // require.True(t, err == nil || apierr.IsMissing(err)) + // }) + + cfg := &databricks.Config{ + //Host: testWorkspaceUrl, + Host: a.Config.Host, + ClientID: sp.Id, + AuthType: "databricks-wif", + //Host: testWorkspaceUrl, + } + + ws, err := databricks.NewAccountClient(cfg) + + require.NoError(t, err) + users := ws.Users.List(ctx, iam.ListAccountUsersRequest{}) + _, err = users.Next(ctx) + require.NoError(t, err) + // testWorkspaceId := GetEnvOrSkipTest(t, "TEST_WORKSPACE_ID") + // testWorkspaceUrl := GetEnvOrSkipTest(t, "TEST_WORKSPACE_HOST") + +} From 452e56fc4df91b0455e34bb3b99b73392a0fed9c Mon Sep 17 00:00:00 2001 From: Hector Castejon Diaz Date: Wed, 19 Mar 2025 09:48:42 +0100 Subject: [PATCH 03/30] CLI test --- config/auth_databricks_oidc.go | 21 +++----- config/config.go | 3 ++ internal/auth_test.go | 87 +++++++++++++++++++++------------- 3 files changed, 64 insertions(+), 47 deletions(-) diff --git a/config/auth_databricks_oidc.go b/config/auth_databricks_oidc.go index c18562d4b..3690d885a 100644 --- a/config/auth_databricks_oidc.go +++ b/config/auth_databricks_oidc.go @@ -73,30 +73,23 @@ func print(ctx context.Context, token string) { // Configure implements CredentialsStrategy. func (d DatabricksOIDCCredentials) Configure(ctx context.Context, cfg *Config) (credentials.CredentialsProvider, error) { - if cfg.Host == "" || cfg.ClientID == "" { + if cfg.Host == "" || cfg.ClientID == "" || cfg.TokenAudience == "" { return nil, nil } // Get the OIDC token from the environment. - audience := cfg.CanonicalHostName() - if cfg.IsAccountClient() { - audience = cfg.AccountID - } - - audience = "https://github.com/databricks-eng" - supplier := GithubOIDCTokenSupplier{ cfg: cfg, } - token, err := supplier.GetOIDCToken(ctx, audience) + token, err := supplier.GetOIDCToken(ctx, cfg.TokenAudience) if err != nil { return nil, err } if token == "" { - // logger.Debugf(ctx, "No OIDC token found") - // return nil, nil - token = "MyToken" + logger.Debugf(ctx, "No OIDC token found") + return nil, nil } + print(ctx, token) endpoints, err := oidcEndpoints(ctx, cfg) @@ -112,8 +105,6 @@ func (d DatabricksOIDCCredentials) Configure(ctx context.Context, cfg *Config) ( TokenURL: endpoints.TokenEndpoint, Scopes: []string{"all-apis"}, EndpointParams: url.Values{ - //"grant_type": {httpclient.JWTGrantType}, - //"assertion": {idToken}, "subject_token_type": {"urn:ietf:params:oauth:token-type:jwt"}, "subject_token": {token}, "grant_type": {"urn:ietf:params:oauth:grant-type:token-exchange"}, @@ -134,5 +125,5 @@ func (d DatabricksOIDCCredentials) Configure(ctx context.Context, cfg *Config) ( // Name implements CredentialsStrategy. func (d DatabricksOIDCCredentials) Name() string { - return "databricks-wif" + return "databricks-oidc" } diff --git a/config/config.go b/config/config.go index 41a2c3abd..f6a284e93 100644 --- a/config/config.go +++ b/config/config.go @@ -135,6 +135,9 @@ type Config struct { // Environment override to return when resolving the current environment. DatabricksEnvironment *environment.DatabricksEnvironment + // OIDC and WIF audience for the token + TokenAudience string `name:"audience" auth:"-"` + Loaders []Loader // marker for configuration resolving diff --git a/internal/auth_test.go b/internal/auth_test.go index 832d10bb0..4162bb489 100644 --- a/internal/auth_test.go +++ b/internal/auth_test.go @@ -13,13 +13,20 @@ import ( func TestAccWifAuth(t *testing.T) { ctx, a := ucacctTest(t) - //testWorkspaceId := int64(470576644108500) - //testWorkspaceUrl := "https://dbc-1232e87d-9384.cloud.databricks.com" - // Create SP + workspaceIdEnvVar := GetEnvOrSkipTest(t, "TEST_WORKSPACE_ID") + workspaceId, err := strconv.ParseInt(workspaceIdEnvVar, 10, 64) + require.NoError(t, err) + + workspaceUrl := GetEnvOrSkipTest(t, "TEST_WORKSPACE_URL") + + // Create SP with access to the workspace sp, err := a.ServicePrincipals.Create(ctx, iam.ServicePrincipal{ Active: true, DisplayName: RandomName("go-sdk-sp-"), + Roles: []iam.ComplexValue{ + {Value: "account_admin"}, // Assigning account-level admin role + }, }) require.NoError(t, err) t.Cleanup(func() { @@ -27,9 +34,25 @@ func TestAccWifAuth(t *testing.T) { require.True(t, err == nil || apierr.IsMissing(err)) }) - applicationId, err := strconv.ParseInt(sp.Id, 10, 64) + spId, err := strconv.ParseInt(sp.Id, 10, 64) + require.NoError(t, err) + + _, err = a.WorkspaceAssignment.Update(ctx, iam.UpdateWorkspaceAssignments{ + WorkspaceId: workspaceId, + PrincipalId: spId, + Permissions: []iam.WorkspacePermission{iam.WorkspacePermissionAdmin}, + }) + require.NoError(t, err) + t.Cleanup(func() { + err := a.WorkspaceAssignment.Delete(ctx, iam.DeleteWorkspaceAssignmentRequest{ + PrincipalId: spId, + WorkspaceId: workspaceId, + }) + require.True(t, err == nil || apierr.IsMissing(err)) + }) + // Setup Federation Policy p, err := a.ServicePrincipalFederationPolicy.Create(ctx, oauth2.CreateServicePrincipalFederationPolicyRequest{ Policy: &oauth2.FederationPolicy{ OidcPolicy: &oauth2.OidcFederationPolicy{ @@ -40,48 +63,48 @@ func TestAccWifAuth(t *testing.T) { Subject: "repo:databricks-eng/eng-dev-ecosystem:environment:integration-tests", }, }, - ServicePrincipalId: applicationId, + ServicePrincipalId: spId, }) require.NoError(t, err) t.Cleanup(func() { err := a.ServicePrincipalFederationPolicy.Delete(ctx, oauth2.DeleteServicePrincipalFederationPolicyRequest{ - ServicePrincipalId: applicationId, + ServicePrincipalId: spId, PolicyId: p.Uid, }) require.True(t, err == nil || apierr.IsMissing(err)) }) - // _, err = a.WorkspaceAssignment.Update(ctx, iam.UpdateWorkspaceAssignments{ - // WorkspaceId: testWorkspaceId, - // PrincipalId: applicationId, - // Permissions: []iam.WorkspacePermission{iam.WorkspacePermissionAdmin}, - // }) - - // require.NoError(t, err) - // t.Cleanup(func() { - // err := a.WorkspaceAssignment.Delete(ctx, iam.DeleteWorkspaceAssignmentRequest{ - // PrincipalId: applicationId, - // WorkspaceId: testWorkspaceId, - // }) - // require.True(t, err == nil || apierr.IsMissing(err)) - // }) - - cfg := &databricks.Config{ - //Host: testWorkspaceUrl, - Host: a.Config.Host, - ClientID: sp.Id, - AuthType: "databricks-wif", - //Host: testWorkspaceUrl, + // Test Workspace Identity Federation at Workspace Level + + wsCfg := &databricks.Config{ + Host: workspaceUrl, + ClientID: sp.ApplicationId, + AuthType: "databricks-oidc", + TokenAudience: "https://github.com/databricks-eng", + } + + wifWsClient, err := databricks.NewWorkspaceClient(wsCfg) + + require.NoError(t, err) + _, err = wifWsClient.CurrentUser.Me(ctx) + require.NoError(t, err) + + // Test Workspace Identity Federation at Account Level + + accCfg := &databricks.Config{ + Host: a.Config.Host, + AccountID: a.Config.AccountID, + ClientID: sp.ApplicationId, + AuthType: "databricks-oidc", + TokenAudience: "https://github.com/databricks-eng", } - ws, err := databricks.NewAccountClient(cfg) + wifAccClient, err := databricks.NewAccountClient(accCfg) require.NoError(t, err) - users := ws.Users.List(ctx, iam.ListAccountUsersRequest{}) - _, err = users.Next(ctx) + it := wifAccClient.Groups.List(ctx, iam.ListAccountGroupsRequest{}) + _, err = it.Next(ctx) require.NoError(t, err) - // testWorkspaceId := GetEnvOrSkipTest(t, "TEST_WORKSPACE_ID") - // testWorkspaceUrl := GetEnvOrSkipTest(t, "TEST_WORKSPACE_HOST") } From 780ddc78076d7de37a6e7926617c5b6bcbe95581 Mon Sep 17 00:00:00 2001 From: Hector Castejon Diaz Date: Thu, 20 Mar 2025 09:52:11 +0100 Subject: [PATCH 04/30] Other stuff --- config/auth_databricks_oidc.go | 22 +++++++++++++++------- 1 file changed, 15 insertions(+), 7 deletions(-) diff --git a/config/auth_databricks_oidc.go b/config/auth_databricks_oidc.go index 3690d885a..e82904eae 100644 --- a/config/auth_databricks_oidc.go +++ b/config/auth_databricks_oidc.go @@ -71,6 +71,17 @@ func print(ctx context.Context, token string) { logger.Debugf(ctx, string(prettyJSON)) } +type databricksOIDCTokenSource struct { + ctx context.Context + jwtTokenSupplicer *GithubOIDCTokenSupplier + workspaceResourceId string + azureTenantId string +} + +func (d *databricksOIDCTokenSource) Token() *oauth2.Token { + return nil +} + // Configure implements CredentialsStrategy. func (d DatabricksOIDCCredentials) Configure(ctx context.Context, cfg *Config) (credentials.CredentialsProvider, error) { if cfg.Host == "" || cfg.ClientID == "" || cfg.TokenAudience == "" { @@ -90,8 +101,6 @@ func (d DatabricksOIDCCredentials) Configure(ctx context.Context, cfg *Config) ( return nil, nil } - print(ctx, token) - endpoints, err := oidcEndpoints(ctx, cfg) if err != nil { return nil, err @@ -99,11 +108,10 @@ func (d DatabricksOIDCCredentials) Configure(ctx context.Context, cfg *Config) ( logger.Debugf(ctx, "Getting tokken for client %s", cfg.ClientID) tsConfig := clientcredentials.Config{ - ClientID: cfg.ClientID, - ClientSecret: "", - AuthStyle: oauth2.AuthStyleInParams, - TokenURL: endpoints.TokenEndpoint, - Scopes: []string{"all-apis"}, + ClientID: cfg.ClientID, + AuthStyle: oauth2.AuthStyleInParams, + TokenURL: endpoints.TokenEndpoint, + Scopes: []string{"all-apis"}, EndpointParams: url.Values{ "subject_token_type": {"urn:ietf:params:oauth:token-type:jwt"}, "subject_token": {token}, From 9494d4096fbaed10d7a0a69c19d957a18a7b5e88 Mon Sep 17 00:00:00 2001 From: Hector Castejon Diaz Date: Thu, 20 Mar 2025 11:22:59 +0100 Subject: [PATCH 05/30] refresh support --- config/auth_databricks_oidc.go | 116 +++++++++------------------------ config/config.go | 2 - config/oidc_github.go | 5 ++ internal/auth_test.go | 85 +++++++++++++++++------- 4 files changed, 97 insertions(+), 111 deletions(-) diff --git a/config/auth_databricks_oidc.go b/config/auth_databricks_oidc.go index e82904eae..0a0d61e22 100644 --- a/config/auth_databricks_oidc.go +++ b/config/auth_databricks_oidc.go @@ -2,13 +2,7 @@ package config import ( "context" - "encoding/base64" - "encoding/json" - "fmt" - "net/http" - "net/http/httputil" "net/url" - "strings" "github.com/databricks/databricks-sdk-go/config/credentials" "github.com/databricks/databricks-sdk-go/logger" @@ -18,97 +12,60 @@ import ( type DatabricksOIDCCredentials struct{} -type loggingTransport struct { - Transport http.RoundTripper -} - -func (t *loggingTransport) RoundTrip(req *http.Request) (*http.Response, error) { - // Log the request - dump, err := httputil.DumpRequestOut(req, true) - if err == nil { - fmt.Println("HTTP Request:") - fmt.Println(string(dump)) +// Configure implements CredentialsStrategy. +func (d DatabricksOIDCCredentials) Configure(ctx context.Context, cfg *Config) (credentials.CredentialsProvider, error) { + if cfg.Host == "" || cfg.ClientID == "" || cfg.TokenAudience == "" { + return nil, nil } - // Send the request - resp, err := t.Transport.RoundTrip(req) - if err != nil { - return nil, err + // Get the OIDC token from the environment. + supplier := GithubOIDCTokenSupplier{ + cfg: cfg, } - // Log the response - dumpResp, err := httputil.DumpResponse(resp, true) - if err == nil { - fmt.Println("HTTP Response:") - fmt.Println(string(dumpResp)) + ts := &databricksOIDCTokenSource{ + ctx: ctx, + jwtTokenSupplicer: &supplier, + audience: cfg.TokenAudience, + clientId: cfg.ClientID, + cfg: cfg, } - return resp, nil + visitor := refreshableVisitor(ts) + return credentials.NewOAuthCredentialsProvider(visitor, ts.Token), nil } -func print(ctx context.Context, token string) { - - parts := strings.Split(token, ".") - if len(parts) != 3 { - fmt.Println("Invalid JWT format") - return - } - // Decode the payload (second part of the JWT) - payload, err := base64.RawURLEncoding.DecodeString(parts[1]) - if err != nil { - logger.Debugf(ctx, "Error decoding JWT payload: %v", err) - return - } - - // Pretty-print JSON - var prettyPayload map[string]interface{} - if err := json.Unmarshal(payload, &prettyPayload); err != nil { - logger.Debugf(ctx, "Error parsing JSON: %v", err) - return - } - - prettyJSON, _ := json.MarshalIndent(prettyPayload, "", " ") - logger.Debugf(ctx, string(prettyJSON)) +// Name implements CredentialsStrategy. +func (d DatabricksOIDCCredentials) Name() string { + return "databricks-oidc" } type databricksOIDCTokenSource struct { - ctx context.Context - jwtTokenSupplicer *GithubOIDCTokenSupplier - workspaceResourceId string - azureTenantId string + ctx context.Context + jwtTokenSupplicer *GithubOIDCTokenSupplier + audience string + clientId string + cfg *Config } -func (d *databricksOIDCTokenSource) Token() *oauth2.Token { - return nil -} - -// Configure implements CredentialsStrategy. -func (d DatabricksOIDCCredentials) Configure(ctx context.Context, cfg *Config) (credentials.CredentialsProvider, error) { - if cfg.Host == "" || cfg.ClientID == "" || cfg.TokenAudience == "" { - return nil, nil - } - - // Get the OIDC token from the environment. - supplier := GithubOIDCTokenSupplier{ - cfg: cfg, - } - token, err := supplier.GetOIDCToken(ctx, cfg.TokenAudience) +func (d *databricksOIDCTokenSource) Token() (*oauth2.Token, error) { + token, err := d.jwtTokenSupplicer.GetOIDCToken(d.ctx, d.audience) if err != nil { return nil, err } if token == "" { - logger.Debugf(ctx, "No OIDC token found") + logger.Debugf(d.ctx, "No OIDC token found") return nil, nil } - endpoints, err := oidcEndpoints(ctx, cfg) + endpoints, err := oidcEndpoints(d.ctx, d.cfg) if err != nil { return nil, err } - logger.Debugf(ctx, "Getting tokken for client %s", cfg.ClientID) + logger.Debugf(d.ctx, "Getting tokken for client %s", d.clientId) tsConfig := clientcredentials.Config{ - ClientID: cfg.ClientID, + ClientID: d.clientId, AuthStyle: oauth2.AuthStyleInParams, TokenURL: endpoints.TokenEndpoint, Scopes: []string{"all-apis"}, @@ -118,20 +75,7 @@ func (d DatabricksOIDCCredentials) Configure(ctx context.Context, cfg *Config) ( "grant_type": {"urn:ietf:params:oauth:grant-type:token-exchange"}, }, } - client := &http.Client{Transport: &loggingTransport{Transport: http.DefaultTransport}} - ctx = context.WithValue(ctx, oauth2.HTTPClient, client) // Request the token - _, err = tsConfig.Token(ctx) - if err != nil { - return nil, err - } - ts := tsConfig.TokenSource(ctx) - visitor := refreshableVisitor(ts) - return credentials.NewOAuthCredentialsProvider(visitor, ts.Token), nil -} - -// Name implements CredentialsStrategy. -func (d DatabricksOIDCCredentials) Name() string { - return "databricks-oidc" + return tsConfig.Token(d.ctx) } diff --git a/config/config.go b/config/config.go index f6a284e93..3839e0792 100644 --- a/config/config.go +++ b/config/config.go @@ -100,8 +100,6 @@ type Config struct { ClientID string `name:"client_id" env:"DATABRICKS_CLIENT_ID" auth:"oauth" auth_types:"oauth-m2m"` ClientSecret string `name:"client_secret" env:"DATABRICKS_CLIENT_SECRET" auth:"oauth,sensitive" auth_types:"oauth-m2m"` - ApplicationID string `name:"client_id" env:"DATABRICKS_APPLICATION_ID" auth:"oauth" auth_types:"oauth-m2m"` - // Path to the Databricks CLI (version >= 0.100.0). DatabricksCliPath string `name:"databricks_cli_path" env:"DATABRICKS_CLI_PATH" auth_types:"databricks-cli"` diff --git a/config/oidc_github.go b/config/oidc_github.go index 826269b7f..a17068574 100644 --- a/config/oidc_github.go +++ b/config/oidc_github.go @@ -3,6 +3,7 @@ package config import ( "context" "fmt" + "time" "github.com/databricks/databricks-sdk-go/httpclient" "github.com/databricks/databricks-sdk-go/logger" @@ -34,5 +35,9 @@ func (g *GithubOIDCTokenSupplier) GetOIDCToken(ctx context.Context, audience str return "", fmt.Errorf("failed to request ID token from %s: %w", g.cfg.ActionsIDTokenRequestURL, err) } + // GitHub issued time is not allways in sync, and can give tokens which are not yet valid. + // TODO: Remove this after Databricks API is updated to handle such cases. + time.Sleep(2 * time.Second) + return resp.Value, nil } diff --git a/internal/auth_test.go b/internal/auth_test.go index 4162bb489..5daa0caa6 100644 --- a/internal/auth_test.go +++ b/internal/auth_test.go @@ -14,6 +14,68 @@ import ( func TestAccWifAuth(t *testing.T) { ctx, a := ucacctTest(t) + // Create SP with access to the workspace + sp, err := a.ServicePrincipals.Create(ctx, iam.ServicePrincipal{ + Active: true, + DisplayName: RandomName("go-sdk-sp-"), + Roles: []iam.ComplexValue{ + {Value: "account_admin"}, // Assigning account-level admin role + }, + }) + require.NoError(t, err) + t.Cleanup(func() { + err := a.ServicePrincipals.Delete(ctx, iam.DeleteAccountServicePrincipalRequest{Id: sp.Id}) + require.True(t, err == nil || apierr.IsMissing(err)) + }) + + spId, err := strconv.ParseInt(sp.Id, 10, 64) + require.NoError(t, err) + + // Setup Federation Policy + p, err := a.ServicePrincipalFederationPolicy.Create(ctx, oauth2.CreateServicePrincipalFederationPolicyRequest{ + Policy: &oauth2.FederationPolicy{ + OidcPolicy: &oauth2.OidcFederationPolicy{ + Issuer: "https://token.actions.githubusercontent.com", + Audiences: []string{ + "https://github.com/databricks-eng", + }, + Subject: "repo:databricks-eng/eng-dev-ecosystem:environment:integration-tests", + }, + }, + ServicePrincipalId: spId, + }) + + require.NoError(t, err) + t.Cleanup(func() { + err := a.ServicePrincipalFederationPolicy.Delete(ctx, oauth2.DeleteServicePrincipalFederationPolicyRequest{ + ServicePrincipalId: spId, + PolicyId: p.Uid, + }) + require.True(t, err == nil || apierr.IsMissing(err)) + }) + + // Test Workspace Identity Federation at Account Level + + accCfg := &databricks.Config{ + Host: a.Config.Host, + AccountID: a.Config.AccountID, + ClientID: sp.ApplicationId, + AuthType: "databricks-oidc", + TokenAudience: "https://github.com/databricks-eng", + } + + wifAccClient, err := databricks.NewAccountClient(accCfg) + + require.NoError(t, err) + it := wifAccClient.Groups.List(ctx, iam.ListAccountGroupsRequest{}) + _, err = it.Next(ctx) + require.NoError(t, err) + +} + +func TestAccWifAuthWorkspace(t *testing.T) { + ctx, a := ucacctTest(t) + workspaceIdEnvVar := GetEnvOrSkipTest(t, "TEST_WORKSPACE_ID") workspaceId, err := strconv.ParseInt(workspaceIdEnvVar, 10, 64) require.NoError(t, err) @@ -24,9 +86,6 @@ func TestAccWifAuth(t *testing.T) { sp, err := a.ServicePrincipals.Create(ctx, iam.ServicePrincipal{ Active: true, DisplayName: RandomName("go-sdk-sp-"), - Roles: []iam.ComplexValue{ - {Value: "account_admin"}, // Assigning account-level admin role - }, }) require.NoError(t, err) t.Cleanup(func() { @@ -75,8 +134,6 @@ func TestAccWifAuth(t *testing.T) { require.True(t, err == nil || apierr.IsMissing(err)) }) - // Test Workspace Identity Federation at Workspace Level - wsCfg := &databricks.Config{ Host: workspaceUrl, ClientID: sp.ApplicationId, @@ -89,22 +146,4 @@ func TestAccWifAuth(t *testing.T) { require.NoError(t, err) _, err = wifWsClient.CurrentUser.Me(ctx) require.NoError(t, err) - - // Test Workspace Identity Federation at Account Level - - accCfg := &databricks.Config{ - Host: a.Config.Host, - AccountID: a.Config.AccountID, - ClientID: sp.ApplicationId, - AuthType: "databricks-oidc", - TokenAudience: "https://github.com/databricks-eng", - } - - wifAccClient, err := databricks.NewAccountClient(accCfg) - - require.NoError(t, err) - it := wifAccClient.Groups.List(ctx, iam.ListAccountGroupsRequest{}) - _, err = it.Next(ctx) - require.NoError(t, err) - } From e94811a17970868989f895ecd5ace0690fc7aff1 Mon Sep 17 00:00:00 2001 From: Hector Castejon Diaz Date: Thu, 20 Mar 2025 12:30:21 +0100 Subject: [PATCH 06/30] Test --- .github/workflows/integration-tests.yml | 4 +- NEXT_CHANGELOG.md | 2 + README.md | 44 ++- ...abricks_oidc.go => auth_databricks_wif.go} | 46 ++- config/auth_databricks_wif_test.go | 297 ++++++++++++++++++ config/auth_default.go | 3 +- internal/auth_test.go | 8 +- 7 files changed, 366 insertions(+), 38 deletions(-) rename config/{auth_databricks_oidc.go => auth_databricks_wif.go} (64%) create mode 100644 config/auth_databricks_wif_test.go diff --git a/.github/workflows/integration-tests.yml b/.github/workflows/integration-tests.yml index 1ab802e51..d0e8707c1 100644 --- a/.github/workflows/integration-tests.yml +++ b/.github/workflows/integration-tests.yml @@ -2,8 +2,8 @@ name: Integration Tests on: - # pull_request: - # types: [opened, synchronize] + pull_request: + types: [opened, synchronize] merge_group: diff --git a/NEXT_CHANGELOG.md b/NEXT_CHANGELOG.md index 50a58aac7..f5b5e615c 100644 --- a/NEXT_CHANGELOG.md +++ b/NEXT_CHANGELOG.md @@ -3,6 +3,8 @@ ## Release v0.61.0 ### New Features and Improvements +* Introduce support for Databricks Workload Identity Federation in GitHub workflows ([1177](https://github.com/databricks/databricks-sdk-go/pull/1177)) + See README.md for instructions. ### Bug Fixes diff --git a/README.md b/README.md index a299507ee..5914b3d89 100644 --- a/README.md +++ b/README.md @@ -14,19 +14,35 @@ The Databricks SDK for Go includes functionality to accelerate development with ## Contents -- [Getting started](#getting-started) -- [Authentication](#authentication) -- [Code examples](#code-examples) -- [Long running operations](#long-running-operations) -- [Paginated responses](#paginated-responses) -- [GetByName utility methods](#getbyname-utility-methods) -- [Node type and Databricks Runtime selectors](#node-type-and-databricks-runtime-selectors) -- [Integration with `io` interfaces for DBFS](#integration-with-io-interfaces-for-dbfs) -- [User Agent Request Attribution](#user-agent-request-attribution) -- [Error Handling](#error-handling) -- [Logging](#logging) +- [Databricks SDK for Go](#databricks-sdk-for-go) + - [Contents](#contents) + - [Getting started](#getting-started) + - [Authentication](#authentication) + - [In this section](#in-this-section) + - [Default authentication flow](#default-authentication-flow) + - [Databricks native authentication](#databricks-native-authentication) + - [Azure native authentication](#azure-native-authentication) + - [Google Cloud Platform native authentication](#google-cloud-platform-native-authentication) + - [Overriding `.databrickscfg`](#overriding-databrickscfg) + - [Additional authentication configuration options](#additional-authentication-configuration-options) + - [Custom credentials provider](#custom-credentials-provider) + - [Code examples](#code-examples) + - [Long-running operations](#long-running-operations) + - [In this section](#in-this-section-1) + - [Command execution on clusters](#command-execution-on-clusters) + - [Cluster library management](#cluster-library-management) + - [Advanced usage](#advanced-usage) + - [Paginated responses](#paginated-responses) + - [`GetByName` utility methods](#getbyname-utility-methods) + - [Node type and Databricks Runtime selectors](#node-type-and-databricks-runtime-selectors) + - [Integration with `io` interfaces for DBFS](#integration-with-io-interfaces-for-dbfs) + - [Reading into and writing from buffers](#reading-into-and-writing-from-buffers) + - [`pflag.Value` for enums](#pflagvalue-for-enums) + - [User Agent Request Attribution](#user-agent-request-attribution) + - [Error handling](#error-handling) + - [Logging](#logging) - [Testing](#testing) -- [Interface stability](#interface-stability) + - [Interface stability](#interface-stability) ## Getting started @@ -158,10 +174,11 @@ Depending on the Databricks authentication method, the SDK uses the following in ### Databricks native authentication -By default, the Databricks SDK for Go initially tries Databricks token authentication (`AuthType: "pat"` in `*databricks.Config`). If the SDK is unsuccessful, it then tries Databricks basic (username/password) authentication (`AuthType: "basic"` in `*databricks.Config`). +By default, the Databricks SDK for Go initially tries Databricks token authentication (`AuthType: "pat"` in `*databricks.Config`). If the SDK is unsuccessful, it then tries Databricks basic (username/password) authentication (`AuthType: "basic"` in `*databricks.Config`). If unsucesful, it then tries Workload Identity Federation (WIF) based authentication(`AuthType: "databricks-wif"` in `*databricks.Config`). Currently, only GitHub provided JWT Tokens is supported. - For Databricks token authentication, you must provide `Host` and `Token`; or their environment variable or `.databrickscfg` file field equivalents. - For Databricks basic authentication, you must provide `Host`, `Username`, and `Password` _(for AWS workspace-level operations)_; or `Host`, `AccountID`, `Username`, and `Password` _(for AWS, Azure, or GCP account-level operations)_; or their environment variable or `.databrickscfg` file field equivalents. +- For Databricks wif authentication, you must provide `Host`, `ClientID` and `TokenAudience`; or their environment variable or `.databrickscfg` file field equivalents. | `*databricks.Config` argument | Description | Environment variable / `.databrickscfg` file field | | ----------------------------- | ------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------ | -------------------------------------------------- | @@ -170,6 +187,7 @@ By default, the Databricks SDK for Go initially tries Databricks token authentic | `Token` | _(String)_ The Databricks personal access token (PAT) _(AWS, Azure, and GCP)_ or Azure Active Directory (Azure AD) token _(Azure)_. | `DATABRICKS_TOKEN` / `token` | | `Username` | _(String)_ The Databricks username part of basic authentication. Only possible when `Host` is `*.cloud.databricks.com` _(AWS)_. | `DATABRICKS_USERNAME` / `username` | | `Password` | _(String)_ The Databricks password part of basic authentication. Only possible when `Host` is `*.cloud.databricks.com` _(AWS)_. | `DATABRICKS_PASSWORD` / `password` | +| `TokenAudience` | _(String)_ The Audience for the JWT Token. | `TOKEN_AUDIENCE` / `token_audience` | For example, to use Databricks token authentication: diff --git a/config/auth_databricks_oidc.go b/config/auth_databricks_wif.go similarity index 64% rename from config/auth_databricks_oidc.go rename to config/auth_databricks_wif.go index 0a0d61e22..cefb0addc 100644 --- a/config/auth_databricks_oidc.go +++ b/config/auth_databricks_wif.go @@ -2,6 +2,7 @@ package config import ( "context" + "errors" "net/url" "github.com/databricks/databricks-sdk-go/config/credentials" @@ -10,25 +11,39 @@ import ( "golang.org/x/oauth2/clientcredentials" ) -type DatabricksOIDCCredentials struct{} +type DatabricksWIFCredentials struct{} // Configure implements CredentialsStrategy. -func (d DatabricksOIDCCredentials) Configure(ctx context.Context, cfg *Config) (credentials.CredentialsProvider, error) { +func (d DatabricksWIFCredentials) Configure(ctx context.Context, cfg *Config) (credentials.CredentialsProvider, error) { if cfg.Host == "" || cfg.ClientID == "" || cfg.TokenAudience == "" { return nil, nil } - // Get the OIDC token from the environment. supplier := GithubOIDCTokenSupplier{ cfg: cfg, } - ts := &databricksOIDCTokenSource{ + // If no supplier can get an IdToken, skip this CredentialsStrategy + idToken, err := supplier.GetOIDCToken(ctx, cfg.TokenAudience) + if err != nil { + return nil, err + } + if idToken == "" { + return nil, nil + } + + endpoints, err := oidcEndpoints(ctx, cfg) + if err != nil { + return nil, err + } + + ts := &databricksWIFTokenSource{ ctx: ctx, jwtTokenSupplicer: &supplier, audience: cfg.TokenAudience, clientId: cfg.ClientID, cfg: cfg, + tokenEndpoint: endpoints.TokenEndpoint, } visitor := refreshableVisitor(ts) @@ -36,38 +51,34 @@ func (d DatabricksOIDCCredentials) Configure(ctx context.Context, cfg *Config) ( } // Name implements CredentialsStrategy. -func (d DatabricksOIDCCredentials) Name() string { - return "databricks-oidc" +func (d DatabricksWIFCredentials) Name() string { + return "databricks-wif" } -type databricksOIDCTokenSource struct { +type databricksWIFTokenSource struct { ctx context.Context jwtTokenSupplicer *GithubOIDCTokenSupplier + tokenEndpoint string audience string clientId string cfg *Config } -func (d *databricksOIDCTokenSource) Token() (*oauth2.Token, error) { +func (d *databricksWIFTokenSource) Token() (*oauth2.Token, error) { token, err := d.jwtTokenSupplicer.GetOIDCToken(d.ctx, d.audience) if err != nil { return nil, err } if token == "" { + // It should not happen, since we check before constructing the token source. logger.Debugf(d.ctx, "No OIDC token found") - return nil, nil + return nil, errors.New("cannot get OIDC token") } - endpoints, err := oidcEndpoints(d.ctx, d.cfg) - if err != nil { - return nil, err - } - logger.Debugf(d.ctx, "Getting tokken for client %s", d.clientId) - tsConfig := clientcredentials.Config{ ClientID: d.clientId, AuthStyle: oauth2.AuthStyleInParams, - TokenURL: endpoints.TokenEndpoint, + TokenURL: d.tokenEndpoint, Scopes: []string{"all-apis"}, EndpointParams: url.Values{ "subject_token_type": {"urn:ietf:params:oauth:token-type:jwt"}, @@ -76,6 +87,7 @@ func (d *databricksOIDCTokenSource) Token() (*oauth2.Token, error) { }, } - // Request the token + logger.Debugf(d.ctx, "Getting tokken for client %s", d.clientId) + return tsConfig.Token(d.ctx) } diff --git a/config/auth_databricks_wif_test.go b/config/auth_databricks_wif_test.go new file mode 100644 index 000000000..a142ed35a --- /dev/null +++ b/config/auth_databricks_wif_test.go @@ -0,0 +1,297 @@ +package config + +import ( + "net/http" + "testing" + + "github.com/databricks/databricks-sdk-go/common/environment" + "github.com/databricks/databricks-sdk-go/httpclient/fixtures" + "github.com/google/go-cmp/cmp" +) + +func TestDatabricksGithubWIFCredentials(t *testing.T) { + testCases := []struct { + desc string + cfg *Config + wantHeaders map[string]string + wantErrPrefix *string + }{ + { + desc: "not an databricks config", + cfg: &Config{ + DatabricksEnvironment: &environment.DatabricksEnvironment{ + Cloud: "foo-bar-cloud", + }, + }, + wantErrPrefix: errPrefix("databricks-wif auth: not configured"), + }, + { + desc: "missing token audience", + cfg: &Config{ + Host: "http://host.com/test", + ClientID: "client-id", + ActionsIDTokenRequestURL: "http://endpoint.com/test?version=1", + ActionsIDTokenRequestToken: "token-1337", + }, + wantErrPrefix: errPrefix("databricks-wif auth: not configured"), + }, + { + desc: "missing host", + cfg: &Config{ + ClientID: "client-id", + TokenAudience: "token-audience", + ActionsIDTokenRequestURL: "http://endpoint.com/test?version=1", + ActionsIDTokenRequestToken: "token-1337", + }, + wantErrPrefix: errPrefix("databricks-wif auth: not configured"), + }, + { + desc: "missing client ID", + cfg: &Config{ + Host: "http://host.com/test", + TokenAudience: "token-audience", + ActionsIDTokenRequestURL: "http://endpoint.com/test?version=1", + ActionsIDTokenRequestToken: "token-1337", + }, + wantErrPrefix: errPrefix("databricks-wif auth: not configured"), + }, + { + desc: "missing env ACTIONS_ID_TOKEN_REQUEST_TOKEN", + cfg: &Config{ + ClientID: "client-id", + Host: "http://host.com/test", + TokenAudience: "token-audience", + ActionsIDTokenRequestURL: "http://endpoint.com/test?version=1", + HTTPTransport: fixtures.MappingTransport{ + "GET /oidc/.well-known/oauth-authorization-server": { + Response: oauthAuthorizationServer{ + AuthorizationEndpoint: "https://host.com/auth", + TokenEndpoint: "https://host.com/oidc/v1/token", + }, + }, + }, + }, + wantErrPrefix: errPrefix("databricks-wif auth: not configured"), + }, + { + desc: "missing env ACTIONS_ID_TOKEN_REQUEST_URL", + cfg: &Config{ + ClientID: "client-id", + Host: "http://host.com/test", + TokenAudience: "token-audience", + ActionsIDTokenRequestToken: "token-1337", + HTTPTransport: fixtures.MappingTransport{ + "GET /oidc/.well-known/oauth-authorization-server": { + Response: oauthAuthorizationServer{ + AuthorizationEndpoint: "https://host.com/auth", + TokenEndpoint: "https://host.com/oidc/v1/token", + }, + }, + }, + }, + wantErrPrefix: errPrefix("databricks-wif auth: not configured"), + }, + { + desc: "databricks token exchange server error", + cfg: &Config{ + ClientID: "client-id", + Host: "http://host.com/test", + TokenAudience: "token-audience", + ActionsIDTokenRequestURL: "http://endpoint.com/test?version=1", + ActionsIDTokenRequestToken: "token-1337", + HTTPTransport: fixtures.MappingTransport{ + "GET /test?version=1&audience=token-audience": { + Status: http.StatusInternalServerError, + ExpectedHeaders: map[string]string{ + "Authorization": "Bearer token-1337", + "Accept": "application/json", + }, + }, + }, + }, + wantErrPrefix: errPrefix("databricks-wif"), + }, + { + desc: "databricks workspace server error", + cfg: &Config{ + ClientID: "client-id", + Host: "http://host.com/test", + TokenAudience: "token-audience", + ActionsIDTokenRequestURL: "http://endpoint.com/test?version=1", + ActionsIDTokenRequestToken: "token-1337", + HTTPTransport: fixtures.MappingTransport{ + "GET /oidc/.well-known/oauth-authorization-server": { + Response: oauthAuthorizationServer{ + AuthorizationEndpoint: "https://host.com/auth", + TokenEndpoint: "https://host.com/oidc/v1/token", + }, + }, + "GET /test?version=1&audience=token-audience": { + Status: http.StatusOK, + ExpectedHeaders: map[string]string{ + "Authorization": "Bearer token-1337", + "Accept": "application/json", + }, + Response: `{"value": "id-token-42"}`, + }, + "POST /oidc/v1/token": { + Status: http.StatusInternalServerError, + ExpectedHeaders: map[string]string{ + "Content-Type": "application/x-www-form-urlencoded", + }, + }, + }, + }, + wantErrPrefix: errPrefix("inner token: Post \"https://host.com/oidc/v1/token\""), + }, + { + desc: "invalid auth token", + cfg: &Config{ + ClientID: "client-id", + Host: "http://host.com/test", + TokenAudience: "token-audience", + ActionsIDTokenRequestURL: "http://endpoint.com/test?version=1", + ActionsIDTokenRequestToken: "token-1337", + HTTPTransport: fixtures.MappingTransport{ + "GET /oidc/.well-known/oauth-authorization-server": { + Response: oauthAuthorizationServer{ + AuthorizationEndpoint: "https://host.com/auth", + TokenEndpoint: "https://host.com/oidc/v1/token", + }, + }, + "GET /test?version=1&audience=token-audience": { + Status: http.StatusOK, + ExpectedHeaders: map[string]string{ + "Authorization": "Bearer token-1337", + "Accept": "application/json", + }, + Response: `{"value": "id-token-42"}`, + }, + "POST /oidc/v1/token": { + Status: http.StatusOK, + ExpectedHeaders: map[string]string{ + "Content-Type": "application/x-www-form-urlencoded", + }, + Response: map[string]string{ + "foo": "bar", + }, + }, + }, + }, + wantErrPrefix: errPrefix("inner token: oauth2: server response missing access_token"), + }, + { + desc: "success workspace", + cfg: &Config{ + ClientID: "client-id", + Host: "http://host.com/test", + TokenAudience: "token-audience", + ActionsIDTokenRequestURL: "http://endpoint.com/test?version=1", + ActionsIDTokenRequestToken: "token-1337", + HTTPTransport: fixtures.MappingTransport{ + "GET /oidc/.well-known/oauth-authorization-server": { + Response: oauthAuthorizationServer{ + AuthorizationEndpoint: "https://host.com/auth", + TokenEndpoint: "https://host.com/oidc/v1/token", + }, + }, + "GET /test?version=1&audience=token-audience": { + Status: http.StatusOK, + ExpectedHeaders: map[string]string{ + "Authorization": "Bearer token-1337", + "Accept": "application/json", + }, + Response: `{"value": "id-token-42"}`, + }, + "POST /oidc/v1/token": { + Status: http.StatusOK, + ExpectedHeaders: map[string]string{ + "Content-Type": "application/x-www-form-urlencoded", + }, + Response: map[string]string{ + "token_type": "access-token", + "access_token": "test-auth-token", + "refresh_token": "refresh", + "expires_on": "0", + }, + }, + }, + }, + wantHeaders: map[string]string{ + "Authorization": "access-token test-auth-token", + }, + }, + { + desc: "success account", + cfg: &Config{ + ClientID: "client-id", + AccountID: "ac123", + Host: "https://accounts.databricks.com", + TokenAudience: "token-audience", + ActionsIDTokenRequestURL: "http://endpoint.com/test?version=1", + ActionsIDTokenRequestToken: "token-1337", + HTTPTransport: fixtures.MappingTransport{ + "GET /test?version=1&audience=token-audience": { + Status: http.StatusOK, + ExpectedHeaders: map[string]string{ + "Authorization": "Bearer token-1337", + "Accept": "application/json", + }, + Response: `{"value": "id-token-42"}`, + }, + "POST /oidc/accounts/ac123/v1/token": { + Status: http.StatusOK, + ExpectedHeaders: map[string]string{ + "Content-Type": "application/x-www-form-urlencoded", + }, + Response: map[string]string{ + "token_type": "access-token", + "access_token": "test-auth-token", + "refresh_token": "refresh", + "expires_on": "0", + }, + }, + }, + }, + wantHeaders: map[string]string{ + "Authorization": "access-token test-auth-token", + }, + }, + } + + for _, tc := range testCases { + t.Run(tc.desc, func(t *testing.T) { + tc.cfg.Credentials = &DatabricksWIFCredentials{} // only test this credential strategy + tc.cfg.DebugHeaders = true + if tc.wantHeaders == nil { + tc.wantHeaders = map[string]string{} + } + + req, _ := http.NewRequest("GET", "http://localhost", nil) + gotErr := tc.cfg.Authenticate(req) + gotHeaders := map[string]string{} + for h := range req.Header { + gotHeaders[h] = req.Header.Get(h) + } + + if tc.wantErrPrefix == nil && gotErr != nil { + t.Errorf("Authenticate(): got error %q, want none", gotErr) + } + if tc.wantErrPrefix != nil && !hasPrefix(gotErr, *tc.wantErrPrefix) { + t.Errorf("Authenticate(): got error %q, want error with prefix %q", gotErr, *tc.wantErrPrefix) + } + if diff := cmp.Diff(tc.wantHeaders, gotHeaders); diff != "" { + t.Errorf("Authenticate(): mismatch (-want +got):\n%s", diff) + } + }) + } +} + +func TestDatabricksWIFCredentials_Name(t *testing.T) { + c := DatabricksWIFCredentials{} + want := "databricks-wif" + + if got := c.Name(); got != want { + t.Errorf("Name(): got %s, want %s", got, want) + } +} diff --git a/config/auth_default.go b/config/auth_default.go index 6b97653a1..aba8fdba6 100644 --- a/config/auth_default.go +++ b/config/auth_default.go @@ -15,6 +15,7 @@ var authProviders = []CredentialsStrategy{ M2mCredentials{}, DatabricksCliCredentials{}, MetadataServiceCredentials{}, + DatabricksWIFCredentials{}, // Attempt to configure auth from most specific to most generic (the Azure CLI). AzureGithubOIDCCredentials{}, @@ -25,8 +26,6 @@ var authProviders = []CredentialsStrategy{ // Attempt to configure auth from most specific to most generic (Google Application Default Credentials). GoogleCredentials{}, GoogleDefaultCredentials{}, - - DatabricksOIDCCredentials{}, } type DefaultCredentials struct { diff --git a/internal/auth_test.go b/internal/auth_test.go index 5daa0caa6..e4e368b72 100644 --- a/internal/auth_test.go +++ b/internal/auth_test.go @@ -11,7 +11,7 @@ import ( "github.com/stretchr/testify/require" ) -func TestAccWifAuth(t *testing.T) { +func TestUcAccWifAuth(t *testing.T) { ctx, a := ucacctTest(t) // Create SP with access to the workspace @@ -60,7 +60,7 @@ func TestAccWifAuth(t *testing.T) { Host: a.Config.Host, AccountID: a.Config.AccountID, ClientID: sp.ApplicationId, - AuthType: "databricks-oidc", + AuthType: "databricks-wif", TokenAudience: "https://github.com/databricks-eng", } @@ -73,7 +73,7 @@ func TestAccWifAuth(t *testing.T) { } -func TestAccWifAuthWorkspace(t *testing.T) { +func TestUcAccWifAuthWorkspace(t *testing.T) { ctx, a := ucacctTest(t) workspaceIdEnvVar := GetEnvOrSkipTest(t, "TEST_WORKSPACE_ID") @@ -137,7 +137,7 @@ func TestAccWifAuthWorkspace(t *testing.T) { wsCfg := &databricks.Config{ Host: workspaceUrl, ClientID: sp.ApplicationId, - AuthType: "databricks-oidc", + AuthType: "databricks-wif", TokenAudience: "https://github.com/databricks-eng", } From bac4c399a72b0d60a7a31f16055e00188d53645d Mon Sep 17 00:00:00 2001 From: Hector Castejon Diaz Date: Mon, 24 Mar 2025 10:15:41 +0100 Subject: [PATCH 07/30] PR Comments --- NEXT_CHANGELOG.md | 6 +++- README.md | 9 ++---- config/auth_databricks_wif.go | 35 +++++++++++++---------- config/auth_databricks_wif_test.go | 45 +++++++++++++++++++++++------- config/config.go | 2 +- config/oidc_github.go | 11 ++++---- internal/auth_test.go | 4 +++ 7 files changed, 73 insertions(+), 39 deletions(-) diff --git a/NEXT_CHANGELOG.md b/NEXT_CHANGELOG.md index f5b5e615c..658271cfc 100644 --- a/NEXT_CHANGELOG.md +++ b/NEXT_CHANGELOG.md @@ -3,8 +3,12 @@ ## Release v0.61.0 ### New Features and Improvements -* Introduce support for Databricks Workload Identity Federation in GitHub workflows ([1177](https://github.com/databricks/databricks-sdk-go/pull/1177)) +* Introduce support for Databricks Workload Identity Federation in GitHub workflows ([1177](https://github.com/databricks/databricks-sdk-go/pull/1177)). See README.md for instructions. + [Breaking] Users running their worklows in GitHub Actions, which use Cloud native authentication and also have a `DATABRICKS_CLIENT_ID` and `DATABRICKS_HOST` + environment variables set may see their authentication start failing due to the order in which the SDK tries different authentication methods. + In such case, the `DATABRICKS_AUTH_TYPE` environment variable must be set to match the previously used authentication method. + ### Bug Fixes diff --git a/README.md b/README.md index 5914b3d89..042e820f7 100644 --- a/README.md +++ b/README.md @@ -174,20 +174,17 @@ Depending on the Databricks authentication method, the SDK uses the following in ### Databricks native authentication -By default, the Databricks SDK for Go initially tries Databricks token authentication (`AuthType: "pat"` in `*databricks.Config`). If the SDK is unsuccessful, it then tries Databricks basic (username/password) authentication (`AuthType: "basic"` in `*databricks.Config`). If unsucesful, it then tries Workload Identity Federation (WIF) based authentication(`AuthType: "databricks-wif"` in `*databricks.Config`). Currently, only GitHub provided JWT Tokens is supported. +By default, the Databricks SDK for Go initially tries Databricks token authentication (`AuthType: "pat"` in `*databricks.Config`). If the SDK is unsuccessful, it then tries Workload Identity Federation (WIF) based authentication(`AuthType: "databricks-wif"` in `*databricks.Config`). Currently, only GitHub provided JWT Tokens is supported. - For Databricks token authentication, you must provide `Host` and `Token`; or their environment variable or `.databrickscfg` file field equivalents. -- For Databricks basic authentication, you must provide `Host`, `Username`, and `Password` _(for AWS workspace-level operations)_; or `Host`, `AccountID`, `Username`, and `Password` _(for AWS, Azure, or GCP account-level operations)_; or their environment variable or `.databrickscfg` file field equivalents. -- For Databricks wif authentication, you must provide `Host`, `ClientID` and `TokenAudience`; or their environment variable or `.databrickscfg` file field equivalents. +- For Databricks wif authentication, you must provide `Host`, `ClientID` and `TokenAudience` _(optional)_; or their environment variable or `.databrickscfg` file field equivalents. | `*databricks.Config` argument | Description | Environment variable / `.databrickscfg` file field | | ----------------------------- | ------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------ | -------------------------------------------------- | | `Host` | _(String)_ The Databricks host URL for either the Databricks workspace endpoint or the Databricks accounts endpoint. | `DATABRICKS_HOST` / `host` | | `AccountID` | _(String)_ The Databricks account ID for the Databricks accounts endpoint. Only has effect when `Host` is either `https://accounts.cloud.databricks.com/` _(AWS)_, `https://accounts.azuredatabricks.net/` _(Azure)_, or `https://accounts.gcp.databricks.com/` _(GCP)_. | `DATABRICKS_ACCOUNT_ID` / `account_id` | | `Token` | _(String)_ The Databricks personal access token (PAT) _(AWS, Azure, and GCP)_ or Azure Active Directory (Azure AD) token _(Azure)_. | `DATABRICKS_TOKEN` / `token` | -| `Username` | _(String)_ The Databricks username part of basic authentication. Only possible when `Host` is `*.cloud.databricks.com` _(AWS)_. | `DATABRICKS_USERNAME` / `username` | -| `Password` | _(String)_ The Databricks password part of basic authentication. Only possible when `Host` is `*.cloud.databricks.com` _(AWS)_. | `DATABRICKS_PASSWORD` / `password` | -| `TokenAudience` | _(String)_ The Audience for the JWT Token. | `TOKEN_AUDIENCE` / `token_audience` | +| `TokenAudience` | _(String)_ When using Workload Identity Federation, the audience to specify when fetching an ID token from the ID token supplier. | `TOKEN_AUDIENCE` / `token_audience` | For example, to use Databricks token authentication: diff --git a/config/auth_databricks_wif.go b/config/auth_databricks_wif.go index cefb0addc..3cd67131a 100644 --- a/config/auth_databricks_wif.go +++ b/config/auth_databricks_wif.go @@ -11,11 +11,16 @@ import ( "golang.org/x/oauth2/clientcredentials" ) +// DatabricksWIFCredentials uses a Token Supplier to get a JWT Token and exchanges +// it for a Databricks Token. +// +// Supported suppliers: +// - GitHub OIDC type DatabricksWIFCredentials struct{} // Configure implements CredentialsStrategy. func (d DatabricksWIFCredentials) Configure(ctx context.Context, cfg *Config) (credentials.CredentialsProvider, error) { - if cfg.Host == "" || cfg.ClientID == "" || cfg.TokenAudience == "" { + if cfg.Host == "" || cfg.ClientID == "" { return nil, nil } @@ -38,12 +43,12 @@ func (d DatabricksWIFCredentials) Configure(ctx context.Context, cfg *Config) (c } ts := &databricksWIFTokenSource{ - ctx: ctx, - jwtTokenSupplicer: &supplier, - audience: cfg.TokenAudience, - clientId: cfg.ClientID, - cfg: cfg, - tokenEndpoint: endpoints.TokenEndpoint, + ctx: ctx, + idTokenSupplier: &supplier, + audience: cfg.TokenAudience, + clientId: cfg.ClientID, + cfg: cfg, + tokenEndpoint: endpoints.TokenEndpoint, } visitor := refreshableVisitor(ts) @@ -56,16 +61,16 @@ func (d DatabricksWIFCredentials) Name() string { } type databricksWIFTokenSource struct { - ctx context.Context - jwtTokenSupplicer *GithubOIDCTokenSupplier - tokenEndpoint string - audience string - clientId string - cfg *Config + ctx context.Context + idTokenSupplier *GithubOIDCTokenSupplier + tokenEndpoint string + audience string + clientId string + cfg *Config } func (d *databricksWIFTokenSource) Token() (*oauth2.Token, error) { - token, err := d.jwtTokenSupplicer.GetOIDCToken(d.ctx, d.audience) + token, err := d.idTokenSupplier.GetOIDCToken(d.ctx, d.audience) if err != nil { return nil, err } @@ -87,7 +92,7 @@ func (d *databricksWIFTokenSource) Token() (*oauth2.Token, error) { }, } - logger.Debugf(d.ctx, "Getting tokken for client %s", d.clientId) + logger.Debugf(d.ctx, "Getting token for client %s", d.clientId) return tsConfig.Token(d.ctx) } diff --git a/config/auth_databricks_wif_test.go b/config/auth_databricks_wif_test.go index a142ed35a..2c27e7389 100644 --- a/config/auth_databricks_wif_test.go +++ b/config/auth_databricks_wif_test.go @@ -25,16 +25,6 @@ func TestDatabricksGithubWIFCredentials(t *testing.T) { }, wantErrPrefix: errPrefix("databricks-wif auth: not configured"), }, - { - desc: "missing token audience", - cfg: &Config{ - Host: "http://host.com/test", - ClientID: "client-id", - ActionsIDTokenRequestURL: "http://endpoint.com/test?version=1", - ActionsIDTokenRequestToken: "token-1337", - }, - wantErrPrefix: errPrefix("databricks-wif auth: not configured"), - }, { desc: "missing host", cfg: &Config{ @@ -257,6 +247,41 @@ func TestDatabricksGithubWIFCredentials(t *testing.T) { "Authorization": "access-token test-auth-token", }, }, + { + desc: "default token audience", + cfg: &Config{ + ClientID: "client-id", + AccountID: "ac123", + Host: "https://accounts.databricks.com", + ActionsIDTokenRequestURL: "http://endpoint.com/test?version=1", + ActionsIDTokenRequestToken: "token-1337", + HTTPTransport: fixtures.MappingTransport{ + "GET /test?version=1": { + Status: http.StatusOK, + ExpectedHeaders: map[string]string{ + "Authorization": "Bearer token-1337", + "Accept": "application/json", + }, + Response: `{"value": "id-token-42"}`, + }, + "POST /oidc/accounts/ac123/v1/token": { + Status: http.StatusOK, + ExpectedHeaders: map[string]string{ + "Content-Type": "application/x-www-form-urlencoded", + }, + Response: map[string]string{ + "token_type": "access-token", + "access_token": "test-auth-token", + "refresh_token": "refresh", + "expires_on": "0", + }, + }, + }, + }, + wantHeaders: map[string]string{ + "Authorization": "access-token test-auth-token", + }, + }, } for _, tc := range testCases { diff --git a/config/config.go b/config/config.go index 3839e0792..7ac908a08 100644 --- a/config/config.go +++ b/config/config.go @@ -133,7 +133,7 @@ type Config struct { // Environment override to return when resolving the current environment. DatabricksEnvironment *environment.DatabricksEnvironment - // OIDC and WIF audience for the token + // When using Workload Identity Federation, the audience to specify when fetching an ID token from the ID token supplier. TokenAudience string `name:"audience" auth:"-"` Loaders []Loader diff --git a/config/oidc_github.go b/config/oidc_github.go index a17068574..f7086766a 100644 --- a/config/oidc_github.go +++ b/config/oidc_github.go @@ -3,7 +3,6 @@ package config import ( "context" "fmt" - "time" "github.com/databricks/databricks-sdk-go/httpclient" "github.com/databricks/databricks-sdk-go/logger" @@ -27,7 +26,11 @@ func (g *GithubOIDCTokenSupplier) GetOIDCToken(ctx context.Context, audience str resp := struct { // anonymous struct to parse the response Value string `json:"value"` }{} - err := g.cfg.refreshClient.Do(ctx, "GET", fmt.Sprintf("%s&audience=%s", g.cfg.ActionsIDTokenRequestURL, audience), + requestUrl := g.cfg.ActionsIDTokenRequestURL + if audience != "" { + requestUrl = fmt.Sprintf("%s&audience=%s", requestUrl, audience) + } + err := g.cfg.refreshClient.Do(ctx, "GET", requestUrl, httpclient.WithRequestHeader("Authorization", fmt.Sprintf("Bearer %s", g.cfg.ActionsIDTokenRequestToken)), httpclient.WithResponseUnmarshal(&resp), ) @@ -35,9 +38,5 @@ func (g *GithubOIDCTokenSupplier) GetOIDCToken(ctx context.Context, audience str return "", fmt.Errorf("failed to request ID token from %s: %w", g.cfg.ActionsIDTokenRequestURL, err) } - // GitHub issued time is not allways in sync, and can give tokens which are not yet valid. - // TODO: Remove this after Databricks API is updated to handle such cases. - time.Sleep(2 * time.Second) - return resp.Value, nil } diff --git a/internal/auth_test.go b/internal/auth_test.go index e4e368b72..0c26e4dbf 100644 --- a/internal/auth_test.go +++ b/internal/auth_test.go @@ -12,6 +12,8 @@ import ( ) func TestUcAccWifAuth(t *testing.T) { + // This test cannot be run locally. It can only be run from GitHub Workflows. + _ = GetEnvOrSkipTest(t, "ACTIONS_ID_TOKEN_REQUEST_URL") ctx, a := ucacctTest(t) // Create SP with access to the workspace @@ -74,6 +76,8 @@ func TestUcAccWifAuth(t *testing.T) { } func TestUcAccWifAuthWorkspace(t *testing.T) { + // This test cannot be run locally. It can only be run from GitHub Workflows. + _ = GetEnvOrSkipTest(t, "ACTIONS_ID_TOKEN_REQUEST_URL") ctx, a := ucacctTest(t) workspaceIdEnvVar := GetEnvOrSkipTest(t, "TEST_WORKSPACE_ID") From 83619c2da4abb8a549e823ff30e0bf4ee7076221 Mon Sep 17 00:00:00 2001 From: Hector Castejon Diaz Date: Tue, 25 Mar 2025 08:32:19 +0100 Subject: [PATCH 08/30] Default values --- NEXT_CHANGELOG.md | 8 ++-- config/auth_databricks_wif.go | 27 +++++++++---- config/auth_databricks_wif_test.go | 61 ++++++++++++++++++++++++++---- 3 files changed, 77 insertions(+), 19 deletions(-) diff --git a/NEXT_CHANGELOG.md b/NEXT_CHANGELOG.md index 0095b8f5e..39cea4ec8 100644 --- a/NEXT_CHANGELOG.md +++ b/NEXT_CHANGELOG.md @@ -5,13 +5,11 @@ ### New Features and Improvements * Introduce support for Databricks Workload Identity Federation in GitHub workflows ([1177](https://github.com/databricks/databricks-sdk-go/pull/1177)). See README.md for instructions. - [Breaking] Users running their worklows in GitHub Actions, which use Cloud native authentication and also have a `DATABRICKS_CLIENT_ID` and `DATABRICKS_HOST` +* [Breaking] Users running their worklows in GitHub Actions, which use Cloud native authentication and also have a `DATABRICKS_CLIENT_ID` and `DATABRICKS_HOST` environment variables set may see their authentication start failing due to the order in which the SDK tries different authentication methods. - In such case, the `DATABRICKS_AUTH_TYPE` environment variable must be set to match the previously used authentication method. - - + In such case, the `DATABRICKS_AUTH_TYPE` environment variable must be set to match the previously used authentication method. * Support user-to-machine authentication in the SDK ([#1108](https://github.com/databricks/databricks-sdk-go/pull/1108)). -- Instances of `ApiClient` now share the same connection pool by default ([PR #1190](https://github.com/databricks/databricks-sdk-go/pull/1190)). +* Instances of `ApiClient` now share the same connection pool by default ([PR #1190](https://github.com/databricks/databricks-sdk-go/pull/1190)). ### Bug Fixes diff --git a/config/auth_databricks_wif.go b/config/auth_databricks_wif.go index 3cd67131a..439c913cc 100644 --- a/config/auth_databricks_wif.go +++ b/config/auth_databricks_wif.go @@ -6,6 +6,7 @@ import ( "net/url" "github.com/databricks/databricks-sdk-go/config/credentials" + "github.com/databricks/databricks-sdk-go/credentials/u2m" "github.com/databricks/databricks-sdk-go/logger" "golang.org/x/oauth2" "golang.org/x/oauth2/clientcredentials" @@ -28,24 +29,26 @@ func (d DatabricksWIFCredentials) Configure(ctx context.Context, cfg *Config) (c cfg: cfg, } - // If no supplier can get an IdToken, skip this CredentialsStrategy - idToken, err := supplier.GetOIDCToken(ctx, cfg.TokenAudience) + endpoints, err := cfg.getOidcEndpoints(ctx) if err != nil { return nil, err } - if idToken == "" { - return nil, nil - } - endpoints, err := oidcEndpoints(ctx, cfg) + audience := d.getAudience(cfg, endpoints) + + // If no supplier can get an IdToken, skip this CredentialsStrategy + idToken, err := supplier.GetOIDCToken(ctx, audience) if err != nil { return nil, err } + if idToken == "" { + return nil, nil + } ts := &databricksWIFTokenSource{ ctx: ctx, idTokenSupplier: &supplier, - audience: cfg.TokenAudience, + audience: audience, clientId: cfg.ClientID, cfg: cfg, tokenEndpoint: endpoints.TokenEndpoint, @@ -55,6 +58,16 @@ func (d DatabricksWIFCredentials) Configure(ctx context.Context, cfg *Config) (c return credentials.NewOAuthCredentialsProvider(visitor, ts.Token), nil } +func (d DatabricksWIFCredentials) getAudience(cfg *Config, endpoints *u2m.OAuthAuthorizationServer) string { + if cfg.TokenAudience != "" { + return cfg.TokenAudience + } + if cfg.IsAccountClient() { + return cfg.AccountID + } + return endpoints.TokenEndpoint +} + // Name implements CredentialsStrategy. func (d DatabricksWIFCredentials) Name() string { return "databricks-wif" diff --git a/config/auth_databricks_wif_test.go b/config/auth_databricks_wif_test.go index 2c27e7389..0dcdfa754 100644 --- a/config/auth_databricks_wif_test.go +++ b/config/auth_databricks_wif_test.go @@ -5,6 +5,7 @@ import ( "testing" "github.com/databricks/databricks-sdk-go/common/environment" + "github.com/databricks/databricks-sdk-go/credentials/u2m" "github.com/databricks/databricks-sdk-go/httpclient/fixtures" "github.com/google/go-cmp/cmp" ) @@ -54,7 +55,7 @@ func TestDatabricksGithubWIFCredentials(t *testing.T) { ActionsIDTokenRequestURL: "http://endpoint.com/test?version=1", HTTPTransport: fixtures.MappingTransport{ "GET /oidc/.well-known/oauth-authorization-server": { - Response: oauthAuthorizationServer{ + Response: u2m.OAuthAuthorizationServer{ AuthorizationEndpoint: "https://host.com/auth", TokenEndpoint: "https://host.com/oidc/v1/token", }, @@ -72,7 +73,7 @@ func TestDatabricksGithubWIFCredentials(t *testing.T) { ActionsIDTokenRequestToken: "token-1337", HTTPTransport: fixtures.MappingTransport{ "GET /oidc/.well-known/oauth-authorization-server": { - Response: oauthAuthorizationServer{ + Response: u2m.OAuthAuthorizationServer{ AuthorizationEndpoint: "https://host.com/auth", TokenEndpoint: "https://host.com/oidc/v1/token", }, @@ -90,6 +91,12 @@ func TestDatabricksGithubWIFCredentials(t *testing.T) { ActionsIDTokenRequestURL: "http://endpoint.com/test?version=1", ActionsIDTokenRequestToken: "token-1337", HTTPTransport: fixtures.MappingTransport{ + "GET /oidc/.well-known/oauth-authorization-server": { + Response: u2m.OAuthAuthorizationServer{ + AuthorizationEndpoint: "https://host.com/auth", + TokenEndpoint: "https://host.com/oidc/v1/token", + }, + }, "GET /test?version=1&audience=token-audience": { Status: http.StatusInternalServerError, ExpectedHeaders: map[string]string{ @@ -111,7 +118,7 @@ func TestDatabricksGithubWIFCredentials(t *testing.T) { ActionsIDTokenRequestToken: "token-1337", HTTPTransport: fixtures.MappingTransport{ "GET /oidc/.well-known/oauth-authorization-server": { - Response: oauthAuthorizationServer{ + Response: u2m.OAuthAuthorizationServer{ AuthorizationEndpoint: "https://host.com/auth", TokenEndpoint: "https://host.com/oidc/v1/token", }, @@ -144,7 +151,7 @@ func TestDatabricksGithubWIFCredentials(t *testing.T) { ActionsIDTokenRequestToken: "token-1337", HTTPTransport: fixtures.MappingTransport{ "GET /oidc/.well-known/oauth-authorization-server": { - Response: oauthAuthorizationServer{ + Response: u2m.OAuthAuthorizationServer{ AuthorizationEndpoint: "https://host.com/auth", TokenEndpoint: "https://host.com/oidc/v1/token", }, @@ -180,7 +187,7 @@ func TestDatabricksGithubWIFCredentials(t *testing.T) { ActionsIDTokenRequestToken: "token-1337", HTTPTransport: fixtures.MappingTransport{ "GET /oidc/.well-known/oauth-authorization-server": { - Response: oauthAuthorizationServer{ + Response: u2m.OAuthAuthorizationServer{ AuthorizationEndpoint: "https://host.com/auth", TokenEndpoint: "https://host.com/oidc/v1/token", }, @@ -248,7 +255,7 @@ func TestDatabricksGithubWIFCredentials(t *testing.T) { }, }, { - desc: "default token audience", + desc: "default token audience account", cfg: &Config{ ClientID: "client-id", AccountID: "ac123", @@ -256,7 +263,7 @@ func TestDatabricksGithubWIFCredentials(t *testing.T) { ActionsIDTokenRequestURL: "http://endpoint.com/test?version=1", ActionsIDTokenRequestToken: "token-1337", HTTPTransport: fixtures.MappingTransport{ - "GET /test?version=1": { + "GET /test?version=1&audience=ac123": { Status: http.StatusOK, ExpectedHeaders: map[string]string{ "Authorization": "Bearer token-1337", @@ -282,6 +289,46 @@ func TestDatabricksGithubWIFCredentials(t *testing.T) { "Authorization": "access-token test-auth-token", }, }, + { + desc: "default token audience workspace", + cfg: &Config{ + ClientID: "client-id", + Host: "https://host.com", + ActionsIDTokenRequestURL: "http://endpoint.com/test?version=1", + ActionsIDTokenRequestToken: "token-1337", + HTTPTransport: fixtures.MappingTransport{ + "GET /oidc/.well-known/oauth-authorization-server": { + Response: u2m.OAuthAuthorizationServer{ + AuthorizationEndpoint: "https://host.com/auth", + TokenEndpoint: "https://host.com/oidc/v1/token", + }, + }, + "GET /test?version=1&audience=https://host.com/oidc/v1/token": { + Status: http.StatusOK, + ExpectedHeaders: map[string]string{ + "Authorization": "Bearer token-1337", + "Accept": "application/json", + }, + Response: `{"value": "id-token-42"}`, + }, + "POST /oidc/v1/token": { + Status: http.StatusOK, + ExpectedHeaders: map[string]string{ + "Content-Type": "application/x-www-form-urlencoded", + }, + Response: map[string]string{ + "token_type": "access-token", + "access_token": "test-auth-token", + "refresh_token": "refresh", + "expires_on": "0", + }, + }, + }, + }, + wantHeaders: map[string]string{ + "Authorization": "access-token test-auth-token", + }, + }, } for _, tc := range testCases { From 9b7f6367dad1b19f39acbc2603bda4a97f4d21cc Mon Sep 17 00:00:00 2001 From: Hector Castejon Diaz Date: Fri, 4 Apr 2025 13:00:26 +0200 Subject: [PATCH 09/30] WIP --- config/.azure/az.json | 1 + config/.azure/az.sess | 1 + config/.azure/azureProfile.json | 1 + config/.azure/commandIndex.json | 1 + config/.azure/config | 3 + config/.azure/versionCheck.json | 1 + config/auth_databricks_wif.go | 139 ++--- config/auth_databricks_wif_test.go | 565 +++++++++--------- config/auth_default.go | 40 +- config/testdata/azure/.azure/az.json | 1 + config/testdata/azure/.azure/az.sess | 1 + .../testdata/azure/.azure/azureProfile.json | 1 + .../testdata/azure/.azure/commandIndex.json | 1 + config/testdata/azure/.azure/config | 3 + .../testdata/azure/.azure/versionCheck.json | 1 + config/token_provider_github_oidc.go | 40 ++ config/token_source_strategy.go | 55 ++ internal/auth_test.go | 4 +- 18 files changed, 484 insertions(+), 375 deletions(-) create mode 100644 config/.azure/az.json create mode 100644 config/.azure/az.sess create mode 100644 config/.azure/azureProfile.json create mode 100644 config/.azure/commandIndex.json create mode 100644 config/.azure/config create mode 100644 config/.azure/versionCheck.json create mode 100644 config/testdata/azure/.azure/az.json create mode 100644 config/testdata/azure/.azure/az.sess create mode 100644 config/testdata/azure/.azure/azureProfile.json create mode 100644 config/testdata/azure/.azure/commandIndex.json create mode 100644 config/testdata/azure/.azure/config create mode 100644 config/testdata/azure/.azure/versionCheck.json create mode 100644 config/token_provider_github_oidc.go create mode 100644 config/token_source_strategy.go diff --git a/config/.azure/az.json b/config/.azure/az.json new file mode 100644 index 000000000..22fdca1b2 --- /dev/null +++ b/config/.azure/az.json @@ -0,0 +1 @@ +{} \ No newline at end of file diff --git a/config/.azure/az.sess b/config/.azure/az.sess new file mode 100644 index 000000000..22fdca1b2 --- /dev/null +++ b/config/.azure/az.sess @@ -0,0 +1 @@ +{} \ No newline at end of file diff --git a/config/.azure/azureProfile.json b/config/.azure/azureProfile.json new file mode 100644 index 000000000..c7306e3ef --- /dev/null +++ b/config/.azure/azureProfile.json @@ -0,0 +1 @@ +{"installationId": "52778700-1129-11f0-95ee-c643c1692dc6"} \ No newline at end of file diff --git a/config/.azure/commandIndex.json b/config/.azure/commandIndex.json new file mode 100644 index 000000000..0870156ec --- /dev/null +++ b/config/.azure/commandIndex.json @@ -0,0 +1 @@ +{"version": "2.69.0", "cloudProfile": "latest", "commandIndex": {"acr": ["azure.cli.command_modules.acr"], "aks": ["azure.cli.command_modules.acs", "azure.cli.command_modules.serviceconnector"], "advisor": ["azure.cli.command_modules.advisor"], "ams": ["azure.cli.command_modules.ams"], "apim": ["azure.cli.command_modules.apim"], "appconfig": ["azure.cli.command_modules.appconfig"], "webapp": ["azure.cli.command_modules.appservice", "azure.cli.command_modules.serviceconnector"], "functionapp": ["azure.cli.command_modules.appservice", "azure.cli.command_modules.serviceconnector"], "appservice": ["azure.cli.command_modules.appservice"], "staticwebapp": ["azure.cli.command_modules.appservice"], "logicapp": ["azure.cli.command_modules.appservice"], "aro": ["azure.cli.command_modules.aro"], "backup": ["azure.cli.command_modules.backup"], "batch": ["azure.cli.command_modules.batch"], "batchai": ["azure.cli.command_modules.batchai"], "billing": ["azure.cli.command_modules.billing"], "bot": ["azure.cli.command_modules.botservice"], "afd": ["azure.cli.command_modules.cdn"], "cdn": ["azure.cli.command_modules.cdn"], "cloud": ["azure.cli.command_modules.cloud"], "cognitiveservices": ["azure.cli.command_modules.cognitiveservices"], "compute-recommender": ["azure.cli.command_modules.compute_recommender"], "compute-fleet": ["azure.cli.command_modules.computefleet"], "config": ["azure.cli.command_modules.config"], "configure": ["azure.cli.command_modules.configure"], "cache": ["azure.cli.command_modules.configure"], "consumption": ["azure.cli.command_modules.consumption"], "container": ["azure.cli.command_modules.container"], "containerapp": ["azure.cli.command_modules.containerapp", "azure.cli.command_modules.serviceconnector"], "cosmosdb": ["azure.cli.command_modules.cosmosdb"], "managed-cassandra": ["azure.cli.command_modules.cosmosdb"], "databoxedge": ["azure.cli.command_modules.databoxedge"], "dls": ["azure.cli.command_modules.dls"], "dms": ["azure.cli.command_modules.dms"], "eventgrid": ["azure.cli.command_modules.eventgrid"], "eventhubs": ["azure.cli.command_modules.eventhubs"], "extension": ["azure.cli.command_modules.extension"], "feedback": ["azure.cli.command_modules.feedback"], "survey": ["azure.cli.command_modules.feedback"], "find": ["azure.cli.command_modules.find"], "hdinsight": ["azure.cli.command_modules.hdinsight"], "identity": ["azure.cli.command_modules.identity"], "interactive": ["azure.cli.command_modules.interactive"], "iot": ["azure.cli.command_modules.iot"], "keyvault": ["azure.cli.command_modules.keyvault"], "lab": ["azure.cli.command_modules.lab"], "managedservices": ["azure.cli.command_modules.managedservices"], "maps": ["azure.cli.command_modules.maps"], "term": ["azure.cli.command_modules.marketplaceordering"], "monitor": ["azure.cli.command_modules.monitor"], "mysql": ["azure.cli.command_modules.mysql", "azure.cli.command_modules.rdbms"], "netappfiles": ["azure.cli.command_modules.netappfiles"], "network": ["azure.cli.command_modules.network", "azure.cli.command_modules.privatedns"], "policy": ["azure.cli.command_modules.policyinsights", "azure.cli.command_modules.resource"], "login": ["azure.cli.command_modules.profile"], "logout": ["azure.cli.command_modules.profile"], "self-test": ["azure.cli.command_modules.profile"], "account": ["azure.cli.command_modules.profile", "azure.cli.command_modules.resource"], "mariadb": ["azure.cli.command_modules.rdbms"], "postgres": ["azure.cli.command_modules.rdbms"], "redis": ["azure.cli.command_modules.redis"], "relay": ["azure.cli.command_modules.relay"], "data-boundary": ["azure.cli.command_modules.resource"], "group": ["azure.cli.command_modules.resource"], "resource": ["azure.cli.command_modules.resource"], "provider": ["azure.cli.command_modules.resource"], "feature": ["azure.cli.command_modules.resource"], "tag": ["azure.cli.command_modules.resource"], "deployment": ["azure.cli.command_modules.resource"], "deployment-scripts": ["azure.cli.command_modules.resource"], "ts": ["azure.cli.command_modules.resource"], "stack": ["azure.cli.command_modules.resource"], "lock": ["azure.cli.command_modules.resource"], "managedapp": ["azure.cli.command_modules.resource"], "bicep": ["azure.cli.command_modules.resource"], "resourcemanagement": ["azure.cli.command_modules.resource"], "private-link": ["azure.cli.command_modules.resource"], "role": ["azure.cli.command_modules.role"], "ad": ["azure.cli.command_modules.role"], "search": ["azure.cli.command_modules.search"], "security": ["azure.cli.command_modules.security"], "servicebus": ["azure.cli.command_modules.servicebus"], "connection": ["azure.cli.command_modules.serviceconnector"], "sf": ["azure.cli.command_modules.servicefabric"], "signalr": ["azure.cli.command_modules.signalr"], "sql": ["azure.cli.command_modules.sql", "azure.cli.command_modules.sqlvm"], "storage": ["azure.cli.command_modules.storage"], "synapse": ["azure.cli.command_modules.synapse"], "rest": ["azure.cli.command_modules.util"], "version": ["azure.cli.command_modules.util"], "upgrade": ["azure.cli.command_modules.util"], "demo": ["azure.cli.command_modules.util"], "snapshot": ["azure.cli.command_modules.vm"], "disk-access": ["azure.cli.command_modules.vm"], "sig": ["azure.cli.command_modules.vm"], "vmss": ["azure.cli.command_modules.vm"], "restore-point": ["azure.cli.command_modules.vm"], "image": ["azure.cli.command_modules.vm"], "capacity": ["azure.cli.command_modules.vm"], "vm": ["azure.cli.command_modules.vm"], "disk": ["azure.cli.command_modules.vm"], "ppg": ["azure.cli.command_modules.vm"], "disk-encryption-set": ["azure.cli.command_modules.vm"], "sshkey": ["azure.cli.command_modules.vm"]}} \ No newline at end of file diff --git a/config/.azure/config b/config/.azure/config new file mode 100644 index 000000000..0ed7f34d6 --- /dev/null +++ b/config/.azure/config @@ -0,0 +1,3 @@ +[cloud] +name = AzureCloud + diff --git a/config/.azure/versionCheck.json b/config/.azure/versionCheck.json new file mode 100644 index 000000000..9be3cb1f5 --- /dev/null +++ b/config/.azure/versionCheck.json @@ -0,0 +1 @@ +{"versions": {"azure-cli": {"local": "2.69.0", "pypi": "2.71.0"}, "core": {"local": "2.69.0", "pypi": "2.71.0"}, "telemetry": {"local": "1.1.0", "pypi": "1.1.0"}}, "update_time": "2025-04-04 09:49:18.938802"} \ No newline at end of file diff --git a/config/auth_databricks_wif.go b/config/auth_databricks_wif.go index 439c913cc..14f1b75e8 100644 --- a/config/auth_databricks_wif.go +++ b/config/auth_databricks_wif.go @@ -5,107 +5,92 @@ import ( "errors" "net/url" - "github.com/databricks/databricks-sdk-go/config/credentials" - "github.com/databricks/databricks-sdk-go/credentials/u2m" "github.com/databricks/databricks-sdk-go/logger" "golang.org/x/oauth2" "golang.org/x/oauth2/clientcredentials" ) -// DatabricksWIFCredentials uses a Token Supplier to get a JWT Token and exchanges -// it for a Databricks Token. -// -// Supported suppliers: -// - GitHub OIDC -type DatabricksWIFCredentials struct{} - -// Configure implements CredentialsStrategy. -func (d DatabricksWIFCredentials) Configure(ctx context.Context, cfg *Config) (credentials.CredentialsProvider, error) { - if cfg.Host == "" || cfg.ClientID == "" { - return nil, nil - } - - supplier := GithubOIDCTokenSupplier{ - cfg: cfg, - } - - endpoints, err := cfg.getOidcEndpoints(ctx) - if err != nil { - return nil, err - } - - audience := d.getAudience(cfg, endpoints) - - // If no supplier can get an IdToken, skip this CredentialsStrategy - idToken, err := supplier.GetOIDCToken(ctx, audience) - if err != nil { - return nil, err - } - if idToken == "" { - return nil, nil +// Constructs all Workload Identity Federation Credentials Strategies +func WifTokenCredentialStrategies(cfg *Config) []CredentialsStrategy { + providers := map[string]TokenProvider{ + "github-oidc-databricks-wif": &GithubProvider{ + cfg: cfg, + }, + // Add new providers at the end of the list } - - ts := &databricksWIFTokenSource{ - ctx: ctx, - idTokenSupplier: &supplier, - audience: audience, - clientId: cfg.ClientID, - cfg: cfg, - tokenEndpoint: endpoints.TokenEndpoint, + strategies := []CredentialsStrategy{} + for name, provider := range providers { + strategies = append(strategies, newWifTokenStrategy(name, cfg, provider)) } - - visitor := refreshableVisitor(ts) - return credentials.NewOAuthCredentialsProvider(visitor, ts.Token), nil + return strategies } -func (d DatabricksWIFCredentials) getAudience(cfg *Config, endpoints *u2m.OAuthAuthorizationServer) string { - if cfg.TokenAudience != "" { - return cfg.TokenAudience - } - if cfg.IsAccountClient() { - return cfg.AccountID +func newWifTokenStrategy( + name string, + cfg *Config, + tokenProvider TokenProvider, +) CredentialsStrategy { + wifTokenExchange := &wifTokenExchange{ + cfg: cfg, + tokenProvider: tokenProvider, } - return endpoints.TokenEndpoint + return NewTokenSourceStrategy(name, wifTokenExchange) } -// Name implements CredentialsStrategy. -func (d DatabricksWIFCredentials) Name() string { - return "databricks-wif" +// wifTokenExchange is a auth.TokenSource which exchanges a token using +// Workload Identity Federation. +type wifTokenExchange struct { + cfg *Config + tokenProvider TokenProvider } -type databricksWIFTokenSource struct { - ctx context.Context - idTokenSupplier *GithubOIDCTokenSupplier - tokenEndpoint string - audience string - clientId string - cfg *Config -} - -func (d *databricksWIFTokenSource) Token() (*oauth2.Token, error) { - token, err := d.idTokenSupplier.GetOIDCToken(d.ctx, d.audience) +func (w *wifTokenExchange) Token(ctx context.Context) (*oauth2.Token, error) { + if w.cfg.ClientID == "" { + logger.Debugf(ctx, "Missing cfg.ClientID") + return nil, errors.New("missing cfg.ClientID") + } + if w.cfg.Host == "" { + logger.Debugf(ctx, "Missing cfg.Host") + return nil, errors.New("missing cfg.Host") + } + audience, err := w.getAudience(ctx) if err != nil { return nil, err } - if token == "" { - // It should not happen, since we check before constructing the token source. - logger.Debugf(d.ctx, "No OIDC token found") - return nil, errors.New("cannot get OIDC token") + idToken, err := w.tokenProvider.IDToken(ctx, audience) + if err != nil { + return nil, err } - - tsConfig := clientcredentials.Config{ - ClientID: d.clientId, + endpoints, err := w.cfg.getOidcEndpoints(ctx) + if err != nil { + return nil, err + } + c := &clientcredentials.Config{ + ClientID: w.cfg.ClientID, AuthStyle: oauth2.AuthStyleInParams, - TokenURL: d.tokenEndpoint, + TokenURL: endpoints.TokenEndpoint, Scopes: []string{"all-apis"}, EndpointParams: url.Values{ "subject_token_type": {"urn:ietf:params:oauth:token-type:jwt"}, - "subject_token": {token}, + "subject_token": {idToken.Value}, "grant_type": {"urn:ietf:params:oauth:grant-type:token-exchange"}, }, } + return c.Token(ctx) +} - logger.Debugf(d.ctx, "Getting token for client %s", d.clientId) - - return tsConfig.Token(d.ctx) +func (w *wifTokenExchange) getAudience(ctx context.Context) (string, error) { + if w.cfg.TokenAudience != "" { + return w.cfg.TokenAudience, nil + } + // For Databricks Accounts, the account id is the default audience. + if w.cfg.IsAccountClient() { + return w.cfg.AccountID, nil + } + // For Databricks Workspaces, the auth endpoint is the default audience. + endpoints, err := w.cfg.getOidcEndpoints(ctx) + if err != nil { + return "", err + } + return endpoints.TokenEndpoint, nil } diff --git a/config/auth_databricks_wif_test.go b/config/auth_databricks_wif_test.go index 0dcdfa754..e96c4c9e9 100644 --- a/config/auth_databricks_wif_test.go +++ b/config/auth_databricks_wif_test.go @@ -4,7 +4,6 @@ import ( "net/http" "testing" - "github.com/databricks/databricks-sdk-go/common/environment" "github.com/databricks/databricks-sdk-go/credentials/u2m" "github.com/databricks/databricks-sdk-go/httpclient/fixtures" "github.com/google/go-cmp/cmp" @@ -17,166 +16,166 @@ func TestDatabricksGithubWIFCredentials(t *testing.T) { wantHeaders map[string]string wantErrPrefix *string }{ - { - desc: "not an databricks config", - cfg: &Config{ - DatabricksEnvironment: &environment.DatabricksEnvironment{ - Cloud: "foo-bar-cloud", - }, - }, - wantErrPrefix: errPrefix("databricks-wif auth: not configured"), - }, - { - desc: "missing host", - cfg: &Config{ - ClientID: "client-id", - TokenAudience: "token-audience", - ActionsIDTokenRequestURL: "http://endpoint.com/test?version=1", - ActionsIDTokenRequestToken: "token-1337", - }, - wantErrPrefix: errPrefix("databricks-wif auth: not configured"), - }, - { - desc: "missing client ID", - cfg: &Config{ - Host: "http://host.com/test", - TokenAudience: "token-audience", - ActionsIDTokenRequestURL: "http://endpoint.com/test?version=1", - ActionsIDTokenRequestToken: "token-1337", - }, - wantErrPrefix: errPrefix("databricks-wif auth: not configured"), - }, - { - desc: "missing env ACTIONS_ID_TOKEN_REQUEST_TOKEN", - cfg: &Config{ - ClientID: "client-id", - Host: "http://host.com/test", - TokenAudience: "token-audience", - ActionsIDTokenRequestURL: "http://endpoint.com/test?version=1", - HTTPTransport: fixtures.MappingTransport{ - "GET /oidc/.well-known/oauth-authorization-server": { - Response: u2m.OAuthAuthorizationServer{ - AuthorizationEndpoint: "https://host.com/auth", - TokenEndpoint: "https://host.com/oidc/v1/token", - }, - }, - }, - }, - wantErrPrefix: errPrefix("databricks-wif auth: not configured"), - }, - { - desc: "missing env ACTIONS_ID_TOKEN_REQUEST_URL", - cfg: &Config{ - ClientID: "client-id", - Host: "http://host.com/test", - TokenAudience: "token-audience", - ActionsIDTokenRequestToken: "token-1337", - HTTPTransport: fixtures.MappingTransport{ - "GET /oidc/.well-known/oauth-authorization-server": { - Response: u2m.OAuthAuthorizationServer{ - AuthorizationEndpoint: "https://host.com/auth", - TokenEndpoint: "https://host.com/oidc/v1/token", - }, - }, - }, - }, - wantErrPrefix: errPrefix("databricks-wif auth: not configured"), - }, - { - desc: "databricks token exchange server error", - cfg: &Config{ - ClientID: "client-id", - Host: "http://host.com/test", - TokenAudience: "token-audience", - ActionsIDTokenRequestURL: "http://endpoint.com/test?version=1", - ActionsIDTokenRequestToken: "token-1337", - HTTPTransport: fixtures.MappingTransport{ - "GET /oidc/.well-known/oauth-authorization-server": { - Response: u2m.OAuthAuthorizationServer{ - AuthorizationEndpoint: "https://host.com/auth", - TokenEndpoint: "https://host.com/oidc/v1/token", - }, - }, - "GET /test?version=1&audience=token-audience": { - Status: http.StatusInternalServerError, - ExpectedHeaders: map[string]string{ - "Authorization": "Bearer token-1337", - "Accept": "application/json", - }, - }, - }, - }, - wantErrPrefix: errPrefix("databricks-wif"), - }, - { - desc: "databricks workspace server error", - cfg: &Config{ - ClientID: "client-id", - Host: "http://host.com/test", - TokenAudience: "token-audience", - ActionsIDTokenRequestURL: "http://endpoint.com/test?version=1", - ActionsIDTokenRequestToken: "token-1337", - HTTPTransport: fixtures.MappingTransport{ - "GET /oidc/.well-known/oauth-authorization-server": { - Response: u2m.OAuthAuthorizationServer{ - AuthorizationEndpoint: "https://host.com/auth", - TokenEndpoint: "https://host.com/oidc/v1/token", - }, - }, - "GET /test?version=1&audience=token-audience": { - Status: http.StatusOK, - ExpectedHeaders: map[string]string{ - "Authorization": "Bearer token-1337", - "Accept": "application/json", - }, - Response: `{"value": "id-token-42"}`, - }, - "POST /oidc/v1/token": { - Status: http.StatusInternalServerError, - ExpectedHeaders: map[string]string{ - "Content-Type": "application/x-www-form-urlencoded", - }, - }, - }, - }, - wantErrPrefix: errPrefix("inner token: Post \"https://host.com/oidc/v1/token\""), - }, - { - desc: "invalid auth token", - cfg: &Config{ - ClientID: "client-id", - Host: "http://host.com/test", - TokenAudience: "token-audience", - ActionsIDTokenRequestURL: "http://endpoint.com/test?version=1", - ActionsIDTokenRequestToken: "token-1337", - HTTPTransport: fixtures.MappingTransport{ - "GET /oidc/.well-known/oauth-authorization-server": { - Response: u2m.OAuthAuthorizationServer{ - AuthorizationEndpoint: "https://host.com/auth", - TokenEndpoint: "https://host.com/oidc/v1/token", - }, - }, - "GET /test?version=1&audience=token-audience": { - Status: http.StatusOK, - ExpectedHeaders: map[string]string{ - "Authorization": "Bearer token-1337", - "Accept": "application/json", - }, - Response: `{"value": "id-token-42"}`, - }, - "POST /oidc/v1/token": { - Status: http.StatusOK, - ExpectedHeaders: map[string]string{ - "Content-Type": "application/x-www-form-urlencoded", - }, - Response: map[string]string{ - "foo": "bar", - }, - }, - }, - }, - wantErrPrefix: errPrefix("inner token: oauth2: server response missing access_token"), - }, + // { + // desc: "not an databricks config", + // cfg: &Config{ + // DatabricksEnvironment: &environment.DatabricksEnvironment{ + // Cloud: "foo-bar-cloud", + // }, + // }, + // wantErrPrefix: errPrefix("federated-oidc-github auth: not configured"), + // }, + // { + // desc: "missing host", + // cfg: &Config{ + // ClientID: "client-id", + // TokenAudience: "token-audience", + // ActionsIDTokenRequestURL: "http://endpoint.com/test?version=1", + // ActionsIDTokenRequestToken: "token-1337", + // }, + // wantErrPrefix: errPrefix("federated-oidc-github auth: not configured"), + // }, + // { + // desc: "missing client ID", + // cfg: &Config{ + // Host: "http://host.com/test", + // TokenAudience: "token-audience", + // ActionsIDTokenRequestURL: "http://endpoint.com/test?version=1", + // ActionsIDTokenRequestToken: "token-1337", + // }, + // wantErrPrefix: errPrefix("federated-oidc-github auth: not configured"), + // }, + // { + // desc: "missing env ACTIONS_ID_TOKEN_REQUEST_TOKEN", + // cfg: &Config{ + // ClientID: "client-id", + // Host: "http://host.com/test", + // TokenAudience: "token-audience", + // ActionsIDTokenRequestURL: "http://endpoint.com/test?version=1", + // HTTPTransport: fixtures.MappingTransport{ + // "GET /oidc/.well-known/oauth-authorization-server": { + // Response: u2m.OAuthAuthorizationServer{ + // AuthorizationEndpoint: "https://host.com/auth", + // TokenEndpoint: "https://host.com/oidc/v1/token", + // }, + // }, + // }, + // }, + // wantErrPrefix: errPrefix("federated-oidc-github auth: not configured"), + // }, + // { + // desc: "missing env ACTIONS_ID_TOKEN_REQUEST_URL", + // cfg: &Config{ + // ClientID: "client-id", + // Host: "http://host.com/test", + // TokenAudience: "token-audience", + // ActionsIDTokenRequestToken: "token-1337", + // HTTPTransport: fixtures.MappingTransport{ + // "GET /oidc/.well-known/oauth-authorization-server": { + // Response: u2m.OAuthAuthorizationServer{ + // AuthorizationEndpoint: "https://host.com/auth", + // TokenEndpoint: "https://host.com/oidc/v1/token", + // }, + // }, + // }, + // }, + // wantErrPrefix: errPrefix("federated-oidc-github auth: not configured"), + // }, + // { + // desc: "databricks token exchange server error", + // cfg: &Config{ + // ClientID: "client-id", + // Host: "http://host.com/test", + // TokenAudience: "token-audience", + // ActionsIDTokenRequestURL: "http://endpoint.com/test?version=1", + // ActionsIDTokenRequestToken: "token-1337", + // HTTPTransport: fixtures.MappingTransport{ + // "GET /oidc/.well-known/oauth-authorization-server": { + // Response: u2m.OAuthAuthorizationServer{ + // AuthorizationEndpoint: "https://host.com/auth", + // TokenEndpoint: "https://host.com/oidc/v1/token", + // }, + // }, + // "GET /test?version=1&audience=token-audience": { + // Status: http.StatusInternalServerError, + // ExpectedHeaders: map[string]string{ + // "Authorization": "Bearer token-1337", + // "Accept": "application/json", + // }, + // }, + // }, + // }, + // wantErrPrefix: errPrefix("federated-oidc-github"), + // }, + // { + // desc: "databricks workspace server error", + // cfg: &Config{ + // ClientID: "client-id", + // Host: "http://host.com/test", + // TokenAudience: "token-audience", + // ActionsIDTokenRequestURL: "http://endpoint.com/test?version=1", + // ActionsIDTokenRequestToken: "token-1337", + // HTTPTransport: fixtures.MappingTransport{ + // "GET /oidc/.well-known/oauth-authorization-server": { + // Response: u2m.OAuthAuthorizationServer{ + // AuthorizationEndpoint: "https://host.com/auth", + // TokenEndpoint: "https://host.com/oidc/v1/token", + // }, + // }, + // "GET /test?version=1&audience=token-audience": { + // Status: http.StatusOK, + // ExpectedHeaders: map[string]string{ + // "Authorization": "Bearer token-1337", + // "Accept": "application/json", + // }, + // Response: `{"value": "id-token-42"}`, + // }, + // "POST /oidc/v1/token": { + // Status: http.StatusInternalServerError, + // ExpectedHeaders: map[string]string{ + // "Content-Type": "application/x-www-form-urlencoded", + // }, + // }, + // }, + // }, + // wantErrPrefix: errPrefix("federated-oidc-github auth: not configured"), + // }, + // { + // desc: "invalid auth token", + // cfg: &Config{ + // ClientID: "client-id", + // Host: "http://host.com/test", + // TokenAudience: "token-audience", + // ActionsIDTokenRequestURL: "http://endpoint.com/test?version=1", + // ActionsIDTokenRequestToken: "token-1337", + // HTTPTransport: fixtures.MappingTransport{ + // "GET /oidc/.well-known/oauth-authorization-server": { + // Response: u2m.OAuthAuthorizationServer{ + // AuthorizationEndpoint: "https://host.com/auth", + // TokenEndpoint: "https://host.com/oidc/v1/token", + // }, + // }, + // "GET /test?version=1&audience=token-audience": { + // Status: http.StatusOK, + // ExpectedHeaders: map[string]string{ + // "Authorization": "Bearer token-1337", + // "Accept": "application/json", + // }, + // Response: `{"value": "id-token-42"}`, + // }, + // "POST /oidc/v1/token": { + // Status: http.StatusOK, + // ExpectedHeaders: map[string]string{ + // "Content-Type": "application/x-www-form-urlencoded", + // }, + // Response: map[string]string{ + // "foo": "bar", + // }, + // }, + // }, + // }, + // wantErrPrefix: errPrefix("federated-oidc-github auth: not configured"), + // }, { desc: "success workspace", cfg: &Config{ @@ -201,6 +200,7 @@ func TestDatabricksGithubWIFCredentials(t *testing.T) { Response: `{"value": "id-token-42"}`, }, "POST /oidc/v1/token": { + Status: http.StatusOK, ExpectedHeaders: map[string]string{ "Content-Type": "application/x-www-form-urlencoded", @@ -218,122 +218,128 @@ func TestDatabricksGithubWIFCredentials(t *testing.T) { "Authorization": "access-token test-auth-token", }, }, - { - desc: "success account", - cfg: &Config{ - ClientID: "client-id", - AccountID: "ac123", - Host: "https://accounts.databricks.com", - TokenAudience: "token-audience", - ActionsIDTokenRequestURL: "http://endpoint.com/test?version=1", - ActionsIDTokenRequestToken: "token-1337", - HTTPTransport: fixtures.MappingTransport{ - "GET /test?version=1&audience=token-audience": { - Status: http.StatusOK, - ExpectedHeaders: map[string]string{ - "Authorization": "Bearer token-1337", - "Accept": "application/json", - }, - Response: `{"value": "id-token-42"}`, - }, - "POST /oidc/accounts/ac123/v1/token": { - Status: http.StatusOK, - ExpectedHeaders: map[string]string{ - "Content-Type": "application/x-www-form-urlencoded", - }, - Response: map[string]string{ - "token_type": "access-token", - "access_token": "test-auth-token", - "refresh_token": "refresh", - "expires_on": "0", - }, - }, - }, - }, - wantHeaders: map[string]string{ - "Authorization": "access-token test-auth-token", - }, - }, - { - desc: "default token audience account", - cfg: &Config{ - ClientID: "client-id", - AccountID: "ac123", - Host: "https://accounts.databricks.com", - ActionsIDTokenRequestURL: "http://endpoint.com/test?version=1", - ActionsIDTokenRequestToken: "token-1337", - HTTPTransport: fixtures.MappingTransport{ - "GET /test?version=1&audience=ac123": { - Status: http.StatusOK, - ExpectedHeaders: map[string]string{ - "Authorization": "Bearer token-1337", - "Accept": "application/json", - }, - Response: `{"value": "id-token-42"}`, - }, - "POST /oidc/accounts/ac123/v1/token": { - Status: http.StatusOK, - ExpectedHeaders: map[string]string{ - "Content-Type": "application/x-www-form-urlencoded", - }, - Response: map[string]string{ - "token_type": "access-token", - "access_token": "test-auth-token", - "refresh_token": "refresh", - "expires_on": "0", - }, - }, - }, - }, - wantHeaders: map[string]string{ - "Authorization": "access-token test-auth-token", - }, - }, - { - desc: "default token audience workspace", - cfg: &Config{ - ClientID: "client-id", - Host: "https://host.com", - ActionsIDTokenRequestURL: "http://endpoint.com/test?version=1", - ActionsIDTokenRequestToken: "token-1337", - HTTPTransport: fixtures.MappingTransport{ - "GET /oidc/.well-known/oauth-authorization-server": { - Response: u2m.OAuthAuthorizationServer{ - AuthorizationEndpoint: "https://host.com/auth", - TokenEndpoint: "https://host.com/oidc/v1/token", - }, - }, - "GET /test?version=1&audience=https://host.com/oidc/v1/token": { - Status: http.StatusOK, - ExpectedHeaders: map[string]string{ - "Authorization": "Bearer token-1337", - "Accept": "application/json", - }, - Response: `{"value": "id-token-42"}`, - }, - "POST /oidc/v1/token": { - Status: http.StatusOK, - ExpectedHeaders: map[string]string{ - "Content-Type": "application/x-www-form-urlencoded", - }, - Response: map[string]string{ - "token_type": "access-token", - "access_token": "test-auth-token", - "refresh_token": "refresh", - "expires_on": "0", - }, - }, - }, - }, - wantHeaders: map[string]string{ - "Authorization": "access-token test-auth-token", - }, - }, + // { + // desc: "success account", + // cfg: &Config{ + // ClientID: "client-id", + // AccountID: "ac123", + // Host: "https://accounts.databricks.com", + // TokenAudience: "token-audience", + // ActionsIDTokenRequestURL: "http://endpoint.com/test?version=1", + // ActionsIDTokenRequestToken: "token-1337", + // HTTPTransport: fixtures.MappingTransport{ + // "GET /test?version=1&audience=token-audience": { + // Status: http.StatusOK, + // ExpectedHeaders: map[string]string{ + // "Authorization": "Bearer token-1337", + // "Accept": "application/json", + // }, + // Response: `{"value": "id-token-42"}`, + // }, + // "POST /oidc/accounts/ac123/v1/token": { + // Status: http.StatusOK, + // ExpectedHeaders: map[string]string{ + // "Content-Type": "application/x-www-form-urlencoded", + // }, + // Response: map[string]string{ + // "token_type": "access-token", + // "access_token": "test-auth-token", + // "refresh_token": "refresh", + // "expires_on": "0", + // }, + // }, + // }, + // }, + // wantHeaders: map[string]string{ + // "Authorization": "access-token test-auth-token", + // }, + // }, + // { + // desc: "default token audience account", + // cfg: &Config{ + // ClientID: "client-id", + // AccountID: "ac123", + // Host: "https://accounts.databricks.com", + // ActionsIDTokenRequestURL: "http://endpoint.com/test?version=1", + // ActionsIDTokenRequestToken: "token-1337", + // HTTPTransport: fixtures.MappingTransport{ + // "GET /test?version=1&audience=ac123": { + // Status: http.StatusOK, + // ExpectedHeaders: map[string]string{ + // "Authorization": "Bearer token-1337", + // "Accept": "application/json", + // }, + // Response: `{"value": "id-token-42"}`, + // }, + // "POST /oidc/accounts/ac123/v1/token": { + // Status: http.StatusOK, + // ExpectedHeaders: map[string]string{ + // "Content-Type": "application/x-www-form-urlencoded", + // }, + // Response: map[string]string{ + // "token_type": "access-token", + // "access_token": "test-auth-token", + // "refresh_token": "refresh", + // "expires_on": "0", + // }, + // }, + // }, + // }, + // wantHeaders: map[string]string{ + // "Authorization": "access-token test-auth-token", + // }, + // }, + // { + // desc: "default token audience workspace", + // cfg: &Config{ + // ClientID: "client-id", + // Host: "https://host.com", + // ActionsIDTokenRequestURL: "http://endpoint.com/test?version=1", + // ActionsIDTokenRequestToken: "token-1337", + // HTTPTransport: fixtures.MappingTransport{ + // "GET /oidc/.well-known/oauth-authorization-server": { + // Response: u2m.OAuthAuthorizationServer{ + // AuthorizationEndpoint: "https://host.com/auth", + // TokenEndpoint: "https://host.com/oidc/v1/token", + // }, + // }, + // "GET /test?version=1&audience=https://host.com/oidc/v1/token": { + // Status: http.StatusOK, + // ExpectedHeaders: map[string]string{ + // "Authorization": "Bearer token-1337", + // "Accept": "application/json", + // }, + // Response: `{"value": "id-token-42"}`, + // }, + // "POST /oidc/v1/token": { + // Status: http.StatusOK, + // ExpectedHeaders: map[string]string{ + // "Content-Type": "application/x-www-form-urlencoded", + // }, + // Response: map[string]string{ + // "token_type": "access-token", + // "access_token": "test-auth-token", + // "refresh_token": "refresh", + // "expires_on": "0", + // }, + // }, + // }, + // }, + // wantHeaders: map[string]string{ + // "Authorization": "access-token test-auth-token", + // }, + // }, } for _, tc := range testCases { t.Run(tc.desc, func(t *testing.T) { - tc.cfg.Credentials = &DatabricksWIFCredentials{} // only test this credential strategy + // if tc.desc != "" { + // t.Skip() + // } + p := &GithubProvider{ + cfg: tc.cfg, + } + tc.cfg.Credentials = newWifTokenStrategy("federated-oidc-github", tc.cfg, p) // only test this credential strategy tc.cfg.DebugHeaders = true if tc.wantHeaders == nil { tc.wantHeaders = map[string]string{} @@ -360,10 +366,13 @@ func TestDatabricksGithubWIFCredentials(t *testing.T) { } func TestDatabricksWIFCredentials_Name(t *testing.T) { - c := DatabricksWIFCredentials{} - want := "databricks-wif" - - if got := c.Name(); got != want { - t.Errorf("Name(): got %s, want %s", got, want) + strategies := WifTokenCredentialStrategies(&Config{}) + expected := []string{"github-oidc-federated-oidc-github"} + found := []string{} + for _, strategy := range strategies { + found = append(found, strategy.Name()) + } + if diff := cmp.Diff(expected, found); diff != "" { + t.Errorf("Strategies mismatch (-want +got):\n%s\n(order must be the same))", diff) } } diff --git a/config/auth_default.go b/config/auth_default.go index 2344a3260..3571b8ba4 100644 --- a/config/auth_default.go +++ b/config/auth_default.go @@ -9,23 +9,27 @@ import ( "github.com/databricks/databricks-sdk-go/logger" ) -var authProviders = []CredentialsStrategy{ - PatCredentials{}, - BasicCredentials{}, - M2mCredentials{}, - DatabricksCliCredentials, - MetadataServiceCredentials{}, - DatabricksWIFCredentials{}, - - // Attempt to configure auth from most specific to most generic (the Azure CLI). - AzureGithubOIDCCredentials{}, - AzureMsiCredentials{}, - AzureClientSecretCredentials{}, - AzureCliCredentials{}, - - // Attempt to configure auth from most specific to most generic (Google Application Default Credentials). - GoogleCredentials{}, - GoogleDefaultCredentials{}, +func buildDefaultStrategies(cfg *Config) []CredentialsStrategy { + return append( + []CredentialsStrategy{ + PatCredentials{}, + BasicCredentials{}, + M2mCredentials{}, + DatabricksCliCredentials, + MetadataServiceCredentials{}, + }, + // append( + // WifTokenCredentialStrategies(cfg), + // Attempt to configure auth from most specific to most generic (the Azure CLI). + AzureGithubOIDCCredentials{}, + AzureMsiCredentials{}, + AzureClientSecretCredentials{}, + AzureCliCredentials{}, + // Attempt to configure auth from most specific to most generic (Google Application Default Credentials). + GoogleCredentials{}, + GoogleDefaultCredentials{}, + //)..., + ) } type DefaultCredentials struct { @@ -46,7 +50,7 @@ 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) { - for _, p := range authProviders { + 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) diff --git a/config/testdata/azure/.azure/az.json b/config/testdata/azure/.azure/az.json new file mode 100644 index 000000000..22fdca1b2 --- /dev/null +++ b/config/testdata/azure/.azure/az.json @@ -0,0 +1 @@ +{} \ No newline at end of file diff --git a/config/testdata/azure/.azure/az.sess b/config/testdata/azure/.azure/az.sess new file mode 100644 index 000000000..22fdca1b2 --- /dev/null +++ b/config/testdata/azure/.azure/az.sess @@ -0,0 +1 @@ +{} \ No newline at end of file diff --git a/config/testdata/azure/.azure/azureProfile.json b/config/testdata/azure/.azure/azureProfile.json new file mode 100644 index 000000000..6dd605248 --- /dev/null +++ b/config/testdata/azure/.azure/azureProfile.json @@ -0,0 +1 @@ +{"installationId": "5277cba2-1129-11f0-86b1-c643c1692dc6"} \ No newline at end of file diff --git a/config/testdata/azure/.azure/commandIndex.json b/config/testdata/azure/.azure/commandIndex.json new file mode 100644 index 000000000..0870156ec --- /dev/null +++ b/config/testdata/azure/.azure/commandIndex.json @@ -0,0 +1 @@ +{"version": "2.69.0", "cloudProfile": "latest", "commandIndex": {"acr": ["azure.cli.command_modules.acr"], "aks": ["azure.cli.command_modules.acs", "azure.cli.command_modules.serviceconnector"], "advisor": ["azure.cli.command_modules.advisor"], "ams": ["azure.cli.command_modules.ams"], "apim": ["azure.cli.command_modules.apim"], "appconfig": ["azure.cli.command_modules.appconfig"], "webapp": ["azure.cli.command_modules.appservice", "azure.cli.command_modules.serviceconnector"], "functionapp": ["azure.cli.command_modules.appservice", "azure.cli.command_modules.serviceconnector"], "appservice": ["azure.cli.command_modules.appservice"], "staticwebapp": ["azure.cli.command_modules.appservice"], "logicapp": ["azure.cli.command_modules.appservice"], "aro": ["azure.cli.command_modules.aro"], "backup": ["azure.cli.command_modules.backup"], "batch": ["azure.cli.command_modules.batch"], "batchai": ["azure.cli.command_modules.batchai"], "billing": ["azure.cli.command_modules.billing"], "bot": ["azure.cli.command_modules.botservice"], "afd": ["azure.cli.command_modules.cdn"], "cdn": ["azure.cli.command_modules.cdn"], "cloud": ["azure.cli.command_modules.cloud"], "cognitiveservices": ["azure.cli.command_modules.cognitiveservices"], "compute-recommender": ["azure.cli.command_modules.compute_recommender"], "compute-fleet": ["azure.cli.command_modules.computefleet"], "config": ["azure.cli.command_modules.config"], "configure": ["azure.cli.command_modules.configure"], "cache": ["azure.cli.command_modules.configure"], "consumption": ["azure.cli.command_modules.consumption"], "container": ["azure.cli.command_modules.container"], "containerapp": ["azure.cli.command_modules.containerapp", "azure.cli.command_modules.serviceconnector"], "cosmosdb": ["azure.cli.command_modules.cosmosdb"], "managed-cassandra": ["azure.cli.command_modules.cosmosdb"], "databoxedge": ["azure.cli.command_modules.databoxedge"], "dls": ["azure.cli.command_modules.dls"], "dms": ["azure.cli.command_modules.dms"], "eventgrid": ["azure.cli.command_modules.eventgrid"], "eventhubs": ["azure.cli.command_modules.eventhubs"], "extension": ["azure.cli.command_modules.extension"], "feedback": ["azure.cli.command_modules.feedback"], "survey": ["azure.cli.command_modules.feedback"], "find": ["azure.cli.command_modules.find"], "hdinsight": ["azure.cli.command_modules.hdinsight"], "identity": ["azure.cli.command_modules.identity"], "interactive": ["azure.cli.command_modules.interactive"], "iot": ["azure.cli.command_modules.iot"], "keyvault": ["azure.cli.command_modules.keyvault"], "lab": ["azure.cli.command_modules.lab"], "managedservices": ["azure.cli.command_modules.managedservices"], "maps": ["azure.cli.command_modules.maps"], "term": ["azure.cli.command_modules.marketplaceordering"], "monitor": ["azure.cli.command_modules.monitor"], "mysql": ["azure.cli.command_modules.mysql", "azure.cli.command_modules.rdbms"], "netappfiles": ["azure.cli.command_modules.netappfiles"], "network": ["azure.cli.command_modules.network", "azure.cli.command_modules.privatedns"], "policy": ["azure.cli.command_modules.policyinsights", "azure.cli.command_modules.resource"], "login": ["azure.cli.command_modules.profile"], "logout": ["azure.cli.command_modules.profile"], "self-test": ["azure.cli.command_modules.profile"], "account": ["azure.cli.command_modules.profile", "azure.cli.command_modules.resource"], "mariadb": ["azure.cli.command_modules.rdbms"], "postgres": ["azure.cli.command_modules.rdbms"], "redis": ["azure.cli.command_modules.redis"], "relay": ["azure.cli.command_modules.relay"], "data-boundary": ["azure.cli.command_modules.resource"], "group": ["azure.cli.command_modules.resource"], "resource": ["azure.cli.command_modules.resource"], "provider": ["azure.cli.command_modules.resource"], "feature": ["azure.cli.command_modules.resource"], "tag": ["azure.cli.command_modules.resource"], "deployment": ["azure.cli.command_modules.resource"], "deployment-scripts": ["azure.cli.command_modules.resource"], "ts": ["azure.cli.command_modules.resource"], "stack": ["azure.cli.command_modules.resource"], "lock": ["azure.cli.command_modules.resource"], "managedapp": ["azure.cli.command_modules.resource"], "bicep": ["azure.cli.command_modules.resource"], "resourcemanagement": ["azure.cli.command_modules.resource"], "private-link": ["azure.cli.command_modules.resource"], "role": ["azure.cli.command_modules.role"], "ad": ["azure.cli.command_modules.role"], "search": ["azure.cli.command_modules.search"], "security": ["azure.cli.command_modules.security"], "servicebus": ["azure.cli.command_modules.servicebus"], "connection": ["azure.cli.command_modules.serviceconnector"], "sf": ["azure.cli.command_modules.servicefabric"], "signalr": ["azure.cli.command_modules.signalr"], "sql": ["azure.cli.command_modules.sql", "azure.cli.command_modules.sqlvm"], "storage": ["azure.cli.command_modules.storage"], "synapse": ["azure.cli.command_modules.synapse"], "rest": ["azure.cli.command_modules.util"], "version": ["azure.cli.command_modules.util"], "upgrade": ["azure.cli.command_modules.util"], "demo": ["azure.cli.command_modules.util"], "snapshot": ["azure.cli.command_modules.vm"], "disk-access": ["azure.cli.command_modules.vm"], "sig": ["azure.cli.command_modules.vm"], "vmss": ["azure.cli.command_modules.vm"], "restore-point": ["azure.cli.command_modules.vm"], "image": ["azure.cli.command_modules.vm"], "capacity": ["azure.cli.command_modules.vm"], "vm": ["azure.cli.command_modules.vm"], "disk": ["azure.cli.command_modules.vm"], "ppg": ["azure.cli.command_modules.vm"], "disk-encryption-set": ["azure.cli.command_modules.vm"], "sshkey": ["azure.cli.command_modules.vm"]}} \ No newline at end of file diff --git a/config/testdata/azure/.azure/config b/config/testdata/azure/.azure/config new file mode 100644 index 000000000..0ed7f34d6 --- /dev/null +++ b/config/testdata/azure/.azure/config @@ -0,0 +1,3 @@ +[cloud] +name = AzureCloud + diff --git a/config/testdata/azure/.azure/versionCheck.json b/config/testdata/azure/.azure/versionCheck.json new file mode 100644 index 000000000..9ad051c80 --- /dev/null +++ b/config/testdata/azure/.azure/versionCheck.json @@ -0,0 +1 @@ +{"versions": {"azure-cli": {"local": "2.69.0", "pypi": "2.71.0"}, "core": {"local": "2.69.0", "pypi": "2.71.0"}, "telemetry": {"local": "1.1.0", "pypi": "1.1.0"}}, "update_time": "2025-04-04 09:49:19.096374"} \ No newline at end of file diff --git a/config/token_provider_github_oidc.go b/config/token_provider_github_oidc.go new file mode 100644 index 000000000..021203e0a --- /dev/null +++ b/config/token_provider_github_oidc.go @@ -0,0 +1,40 @@ +package config + +import ( + "context" + "errors" + "fmt" + + "github.com/databricks/databricks-sdk-go/httpclient" + "github.com/databricks/databricks-sdk-go/logger" +) + +type GithubProvider struct { + cfg *Config +} + +func (g *GithubProvider) IDToken(ctx context.Context, audience string) (*IDToken, error) { + if g.cfg.ActionsIDTokenRequestURL == "" { + logger.Debugf(ctx, "Missing cfg.ActionsIDTokenRequestURL, likely not calling from a Github action") + return nil, errors.New("missing cfg.ActionsIDTokenRequestURL") + } + if g.cfg.ActionsIDTokenRequestToken == "" { + logger.Debugf(ctx, "Missing cfg.ActionsIDTokenRequestToken, likely not calling from a Github action") + return nil, errors.New("missing cfg.ActionsIDTokenRequestToken") + } + + resp := &IDToken{} + requestUrl := g.cfg.ActionsIDTokenRequestURL + if audience != "" { + requestUrl = fmt.Sprintf("%s&audience=%s", requestUrl, audience) + } + err := g.cfg.refreshClient.Do(ctx, "GET", requestUrl, + httpclient.WithRequestHeader("Authorization", fmt.Sprintf("Bearer %s", g.cfg.ActionsIDTokenRequestToken)), + httpclient.WithResponseUnmarshal(resp), + ) + if err != nil { + return nil, fmt.Errorf("failed to request ID token from %s: %w", g.cfg.ActionsIDTokenRequestURL, err) + } + + return resp, nil +} diff --git a/config/token_source_strategy.go b/config/token_source_strategy.go new file mode 100644 index 000000000..4aa07f30d --- /dev/null +++ b/config/token_source_strategy.go @@ -0,0 +1,55 @@ +package config + +import ( + "context" + "fmt" + + "github.com/databricks/databricks-sdk-go/config/credentials" + "github.com/databricks/databricks-sdk-go/config/experimental/auth" + "github.com/databricks/databricks-sdk-go/config/experimental/auth/authconv" + "github.com/databricks/databricks-sdk-go/logger" +) + +type IDToken struct { + Value string +} + +type TokenProvider interface { + // Function to get the token + IDToken(ctx context.Context, audience string) (*IDToken, error) +} + +type TokenSourceStrategy struct { + tokenSource auth.TokenSource + name string +} + +func (t *TokenSourceStrategy) Configure(ctx context.Context, cfg *Config) (credentials.CredentialsProvider, error) { + + // If we cannot get a token, skip this CredentialsStrategy. + // We don't want to fail here because it's possible that the supplier is enabled + // without the user action. For instance, jobs running in GitHub will have + // OIDC environment variables added automatically + if _, err := t.tokenSource.Token(ctx); err != nil { + logger.Debugf(ctx, fmt.Sprintf("Skipping %s due to error: %v", t.name, err)) + return nil, nil + } + + visitor := refreshableVisitor(authconv.OAuth2TokenSource(t.tokenSource)) + return credentials.NewOAuthCredentialsProvider(visitor, authconv.OAuth2TokenSource(t.tokenSource).Token), nil +} + +func (t *TokenSourceStrategy) Name() string { + return t.name +} + +// Creates a CredentialsStrategy from a TokenSource. +func NewTokenSourceStrategy( + name string, + tokenSource auth.TokenSource, +) CredentialsStrategy { + return &TokenSourceStrategy{ + name: name, + tokenSource: tokenSource, + } +} diff --git a/internal/auth_test.go b/internal/auth_test.go index 0c26e4dbf..ef6750e53 100644 --- a/internal/auth_test.go +++ b/internal/auth_test.go @@ -62,7 +62,7 @@ func TestUcAccWifAuth(t *testing.T) { Host: a.Config.Host, AccountID: a.Config.AccountID, ClientID: sp.ApplicationId, - AuthType: "databricks-wif", + AuthType: "github-oidc-databricks-wif", TokenAudience: "https://github.com/databricks-eng", } @@ -141,7 +141,7 @@ func TestUcAccWifAuthWorkspace(t *testing.T) { wsCfg := &databricks.Config{ Host: workspaceUrl, ClientID: sp.ApplicationId, - AuthType: "databricks-wif", + AuthType: "github-oidc-databricks-wif", TokenAudience: "https://github.com/databricks-eng", } From 878c349b12ac3dd5e96b2e440411f6abb71a9def Mon Sep 17 00:00:00 2001 From: Hector Castejon Diaz Date: Fri, 4 Apr 2025 14:31:47 +0200 Subject: [PATCH 10/30] WIP --- config/auth_databricks_wif_test.go | 543 +++++++++++++++-------------- config/oauth_visitors.go | 9 +- config/token_source_strategy.go | 7 +- 3 files changed, 283 insertions(+), 276 deletions(-) diff --git a/config/auth_databricks_wif_test.go b/config/auth_databricks_wif_test.go index e96c4c9e9..363ea97af 100644 --- a/config/auth_databricks_wif_test.go +++ b/config/auth_databricks_wif_test.go @@ -4,6 +4,7 @@ import ( "net/http" "testing" + "github.com/databricks/databricks-sdk-go/common/environment" "github.com/databricks/databricks-sdk-go/credentials/u2m" "github.com/databricks/databricks-sdk-go/httpclient/fixtures" "github.com/google/go-cmp/cmp" @@ -16,166 +17,166 @@ func TestDatabricksGithubWIFCredentials(t *testing.T) { wantHeaders map[string]string wantErrPrefix *string }{ - // { - // desc: "not an databricks config", - // cfg: &Config{ - // DatabricksEnvironment: &environment.DatabricksEnvironment{ - // Cloud: "foo-bar-cloud", - // }, - // }, - // wantErrPrefix: errPrefix("federated-oidc-github auth: not configured"), - // }, - // { - // desc: "missing host", - // cfg: &Config{ - // ClientID: "client-id", - // TokenAudience: "token-audience", - // ActionsIDTokenRequestURL: "http://endpoint.com/test?version=1", - // ActionsIDTokenRequestToken: "token-1337", - // }, - // wantErrPrefix: errPrefix("federated-oidc-github auth: not configured"), - // }, - // { - // desc: "missing client ID", - // cfg: &Config{ - // Host: "http://host.com/test", - // TokenAudience: "token-audience", - // ActionsIDTokenRequestURL: "http://endpoint.com/test?version=1", - // ActionsIDTokenRequestToken: "token-1337", - // }, - // wantErrPrefix: errPrefix("federated-oidc-github auth: not configured"), - // }, - // { - // desc: "missing env ACTIONS_ID_TOKEN_REQUEST_TOKEN", - // cfg: &Config{ - // ClientID: "client-id", - // Host: "http://host.com/test", - // TokenAudience: "token-audience", - // ActionsIDTokenRequestURL: "http://endpoint.com/test?version=1", - // HTTPTransport: fixtures.MappingTransport{ - // "GET /oidc/.well-known/oauth-authorization-server": { - // Response: u2m.OAuthAuthorizationServer{ - // AuthorizationEndpoint: "https://host.com/auth", - // TokenEndpoint: "https://host.com/oidc/v1/token", - // }, - // }, - // }, - // }, - // wantErrPrefix: errPrefix("federated-oidc-github auth: not configured"), - // }, - // { - // desc: "missing env ACTIONS_ID_TOKEN_REQUEST_URL", - // cfg: &Config{ - // ClientID: "client-id", - // Host: "http://host.com/test", - // TokenAudience: "token-audience", - // ActionsIDTokenRequestToken: "token-1337", - // HTTPTransport: fixtures.MappingTransport{ - // "GET /oidc/.well-known/oauth-authorization-server": { - // Response: u2m.OAuthAuthorizationServer{ - // AuthorizationEndpoint: "https://host.com/auth", - // TokenEndpoint: "https://host.com/oidc/v1/token", - // }, - // }, - // }, - // }, - // wantErrPrefix: errPrefix("federated-oidc-github auth: not configured"), - // }, - // { - // desc: "databricks token exchange server error", - // cfg: &Config{ - // ClientID: "client-id", - // Host: "http://host.com/test", - // TokenAudience: "token-audience", - // ActionsIDTokenRequestURL: "http://endpoint.com/test?version=1", - // ActionsIDTokenRequestToken: "token-1337", - // HTTPTransport: fixtures.MappingTransport{ - // "GET /oidc/.well-known/oauth-authorization-server": { - // Response: u2m.OAuthAuthorizationServer{ - // AuthorizationEndpoint: "https://host.com/auth", - // TokenEndpoint: "https://host.com/oidc/v1/token", - // }, - // }, - // "GET /test?version=1&audience=token-audience": { - // Status: http.StatusInternalServerError, - // ExpectedHeaders: map[string]string{ - // "Authorization": "Bearer token-1337", - // "Accept": "application/json", - // }, - // }, - // }, - // }, - // wantErrPrefix: errPrefix("federated-oidc-github"), - // }, - // { - // desc: "databricks workspace server error", - // cfg: &Config{ - // ClientID: "client-id", - // Host: "http://host.com/test", - // TokenAudience: "token-audience", - // ActionsIDTokenRequestURL: "http://endpoint.com/test?version=1", - // ActionsIDTokenRequestToken: "token-1337", - // HTTPTransport: fixtures.MappingTransport{ - // "GET /oidc/.well-known/oauth-authorization-server": { - // Response: u2m.OAuthAuthorizationServer{ - // AuthorizationEndpoint: "https://host.com/auth", - // TokenEndpoint: "https://host.com/oidc/v1/token", - // }, - // }, - // "GET /test?version=1&audience=token-audience": { - // Status: http.StatusOK, - // ExpectedHeaders: map[string]string{ - // "Authorization": "Bearer token-1337", - // "Accept": "application/json", - // }, - // Response: `{"value": "id-token-42"}`, - // }, - // "POST /oidc/v1/token": { - // Status: http.StatusInternalServerError, - // ExpectedHeaders: map[string]string{ - // "Content-Type": "application/x-www-form-urlencoded", - // }, - // }, - // }, - // }, - // wantErrPrefix: errPrefix("federated-oidc-github auth: not configured"), - // }, - // { - // desc: "invalid auth token", - // cfg: &Config{ - // ClientID: "client-id", - // Host: "http://host.com/test", - // TokenAudience: "token-audience", - // ActionsIDTokenRequestURL: "http://endpoint.com/test?version=1", - // ActionsIDTokenRequestToken: "token-1337", - // HTTPTransport: fixtures.MappingTransport{ - // "GET /oidc/.well-known/oauth-authorization-server": { - // Response: u2m.OAuthAuthorizationServer{ - // AuthorizationEndpoint: "https://host.com/auth", - // TokenEndpoint: "https://host.com/oidc/v1/token", - // }, - // }, - // "GET /test?version=1&audience=token-audience": { - // Status: http.StatusOK, - // ExpectedHeaders: map[string]string{ - // "Authorization": "Bearer token-1337", - // "Accept": "application/json", - // }, - // Response: `{"value": "id-token-42"}`, - // }, - // "POST /oidc/v1/token": { - // Status: http.StatusOK, - // ExpectedHeaders: map[string]string{ - // "Content-Type": "application/x-www-form-urlencoded", - // }, - // Response: map[string]string{ - // "foo": "bar", - // }, - // }, - // }, - // }, - // wantErrPrefix: errPrefix("federated-oidc-github auth: not configured"), - // }, + { + desc: "not an databricks config", + cfg: &Config{ + DatabricksEnvironment: &environment.DatabricksEnvironment{ + Cloud: "foo-bar-cloud", + }, + }, + wantErrPrefix: errPrefix("federated-oidc-github auth: not configured"), + }, + { + desc: "missing host", + cfg: &Config{ + ClientID: "client-id", + TokenAudience: "token-audience", + ActionsIDTokenRequestURL: "http://endpoint.com/test?version=1", + ActionsIDTokenRequestToken: "token-1337", + }, + wantErrPrefix: errPrefix("federated-oidc-github auth: not configured"), + }, + { + desc: "missing client ID", + cfg: &Config{ + Host: "http://host.com/test", + TokenAudience: "token-audience", + ActionsIDTokenRequestURL: "http://endpoint.com/test?version=1", + ActionsIDTokenRequestToken: "token-1337", + }, + wantErrPrefix: errPrefix("federated-oidc-github auth: not configured"), + }, + { + desc: "missing env ACTIONS_ID_TOKEN_REQUEST_TOKEN", + cfg: &Config{ + ClientID: "client-id", + Host: "http://host.com/test", + TokenAudience: "token-audience", + ActionsIDTokenRequestURL: "http://endpoint.com/test?version=1", + HTTPTransport: fixtures.MappingTransport{ + "GET /oidc/.well-known/oauth-authorization-server": { + Response: u2m.OAuthAuthorizationServer{ + AuthorizationEndpoint: "https://host.com/auth", + TokenEndpoint: "https://host.com/oidc/v1/token", + }, + }, + }, + }, + wantErrPrefix: errPrefix("federated-oidc-github auth: not configured"), + }, + { + desc: "missing env ACTIONS_ID_TOKEN_REQUEST_URL", + cfg: &Config{ + ClientID: "client-id", + Host: "http://host.com/test", + TokenAudience: "token-audience", + ActionsIDTokenRequestToken: "token-1337", + HTTPTransport: fixtures.MappingTransport{ + "GET /oidc/.well-known/oauth-authorization-server": { + Response: u2m.OAuthAuthorizationServer{ + AuthorizationEndpoint: "https://host.com/auth", + TokenEndpoint: "https://host.com/oidc/v1/token", + }, + }, + }, + }, + wantErrPrefix: errPrefix("federated-oidc-github auth: not configured"), + }, + { + desc: "databricks token exchange server error", + cfg: &Config{ + ClientID: "client-id", + Host: "http://host.com/test", + TokenAudience: "token-audience", + ActionsIDTokenRequestURL: "http://endpoint.com/test?version=1", + ActionsIDTokenRequestToken: "token-1337", + HTTPTransport: fixtures.MappingTransport{ + "GET /oidc/.well-known/oauth-authorization-server": { + Response: u2m.OAuthAuthorizationServer{ + AuthorizationEndpoint: "https://host.com/auth", + TokenEndpoint: "https://host.com/oidc/v1/token", + }, + }, + "GET /test?version=1&audience=token-audience": { + Status: http.StatusInternalServerError, + ExpectedHeaders: map[string]string{ + "Authorization": "Bearer token-1337", + "Accept": "application/json", + }, + }, + }, + }, + wantErrPrefix: errPrefix("federated-oidc-github"), + }, + { + desc: "databricks workspace server error", + cfg: &Config{ + ClientID: "client-id", + Host: "http://host.com/test", + TokenAudience: "token-audience", + ActionsIDTokenRequestURL: "http://endpoint.com/test?version=1", + ActionsIDTokenRequestToken: "token-1337", + HTTPTransport: fixtures.MappingTransport{ + "GET /oidc/.well-known/oauth-authorization-server": { + Response: u2m.OAuthAuthorizationServer{ + AuthorizationEndpoint: "https://host.com/auth", + TokenEndpoint: "https://host.com/oidc/v1/token", + }, + }, + "GET /test?version=1&audience=token-audience": { + Status: http.StatusOK, + ExpectedHeaders: map[string]string{ + "Authorization": "Bearer token-1337", + "Accept": "application/json", + }, + Response: `{"value": "id-token-42"}`, + }, + "POST /oidc/v1/token": { + Status: http.StatusInternalServerError, + ExpectedHeaders: map[string]string{ + "Content-Type": "application/x-www-form-urlencoded", + }, + }, + }, + }, + wantErrPrefix: errPrefix("federated-oidc-github auth: not configured"), + }, + { + desc: "invalid auth token", + cfg: &Config{ + ClientID: "client-id", + Host: "http://host.com/test", + TokenAudience: "token-audience", + ActionsIDTokenRequestURL: "http://endpoint.com/test?version=1", + ActionsIDTokenRequestToken: "token-1337", + HTTPTransport: fixtures.MappingTransport{ + "GET /oidc/.well-known/oauth-authorization-server": { + Response: u2m.OAuthAuthorizationServer{ + AuthorizationEndpoint: "https://host.com/auth", + TokenEndpoint: "https://host.com/oidc/v1/token", + }, + }, + "GET /test?version=1&audience=token-audience": { + Status: http.StatusOK, + ExpectedHeaders: map[string]string{ + "Authorization": "Bearer token-1337", + "Accept": "application/json", + }, + Response: `{"value": "id-token-42"}`, + }, + "POST /oidc/v1/token": { + Status: http.StatusOK, + ExpectedHeaders: map[string]string{ + "Content-Type": "application/x-www-form-urlencoded", + }, + Response: map[string]string{ + "foo": "bar", + }, + }, + }, + }, + wantErrPrefix: errPrefix("federated-oidc-github auth: not configured"), + }, { desc: "success workspace", cfg: &Config{ @@ -218,117 +219,117 @@ func TestDatabricksGithubWIFCredentials(t *testing.T) { "Authorization": "access-token test-auth-token", }, }, - // { - // desc: "success account", - // cfg: &Config{ - // ClientID: "client-id", - // AccountID: "ac123", - // Host: "https://accounts.databricks.com", - // TokenAudience: "token-audience", - // ActionsIDTokenRequestURL: "http://endpoint.com/test?version=1", - // ActionsIDTokenRequestToken: "token-1337", - // HTTPTransport: fixtures.MappingTransport{ - // "GET /test?version=1&audience=token-audience": { - // Status: http.StatusOK, - // ExpectedHeaders: map[string]string{ - // "Authorization": "Bearer token-1337", - // "Accept": "application/json", - // }, - // Response: `{"value": "id-token-42"}`, - // }, - // "POST /oidc/accounts/ac123/v1/token": { - // Status: http.StatusOK, - // ExpectedHeaders: map[string]string{ - // "Content-Type": "application/x-www-form-urlencoded", - // }, - // Response: map[string]string{ - // "token_type": "access-token", - // "access_token": "test-auth-token", - // "refresh_token": "refresh", - // "expires_on": "0", - // }, - // }, - // }, - // }, - // wantHeaders: map[string]string{ - // "Authorization": "access-token test-auth-token", - // }, - // }, - // { - // desc: "default token audience account", - // cfg: &Config{ - // ClientID: "client-id", - // AccountID: "ac123", - // Host: "https://accounts.databricks.com", - // ActionsIDTokenRequestURL: "http://endpoint.com/test?version=1", - // ActionsIDTokenRequestToken: "token-1337", - // HTTPTransport: fixtures.MappingTransport{ - // "GET /test?version=1&audience=ac123": { - // Status: http.StatusOK, - // ExpectedHeaders: map[string]string{ - // "Authorization": "Bearer token-1337", - // "Accept": "application/json", - // }, - // Response: `{"value": "id-token-42"}`, - // }, - // "POST /oidc/accounts/ac123/v1/token": { - // Status: http.StatusOK, - // ExpectedHeaders: map[string]string{ - // "Content-Type": "application/x-www-form-urlencoded", - // }, - // Response: map[string]string{ - // "token_type": "access-token", - // "access_token": "test-auth-token", - // "refresh_token": "refresh", - // "expires_on": "0", - // }, - // }, - // }, - // }, - // wantHeaders: map[string]string{ - // "Authorization": "access-token test-auth-token", - // }, - // }, - // { - // desc: "default token audience workspace", - // cfg: &Config{ - // ClientID: "client-id", - // Host: "https://host.com", - // ActionsIDTokenRequestURL: "http://endpoint.com/test?version=1", - // ActionsIDTokenRequestToken: "token-1337", - // HTTPTransport: fixtures.MappingTransport{ - // "GET /oidc/.well-known/oauth-authorization-server": { - // Response: u2m.OAuthAuthorizationServer{ - // AuthorizationEndpoint: "https://host.com/auth", - // TokenEndpoint: "https://host.com/oidc/v1/token", - // }, - // }, - // "GET /test?version=1&audience=https://host.com/oidc/v1/token": { - // Status: http.StatusOK, - // ExpectedHeaders: map[string]string{ - // "Authorization": "Bearer token-1337", - // "Accept": "application/json", - // }, - // Response: `{"value": "id-token-42"}`, - // }, - // "POST /oidc/v1/token": { - // Status: http.StatusOK, - // ExpectedHeaders: map[string]string{ - // "Content-Type": "application/x-www-form-urlencoded", - // }, - // Response: map[string]string{ - // "token_type": "access-token", - // "access_token": "test-auth-token", - // "refresh_token": "refresh", - // "expires_on": "0", - // }, - // }, - // }, - // }, - // wantHeaders: map[string]string{ - // "Authorization": "access-token test-auth-token", - // }, - // }, + { + desc: "success account", + cfg: &Config{ + ClientID: "client-id", + AccountID: "ac123", + Host: "https://accounts.databricks.com", + TokenAudience: "token-audience", + ActionsIDTokenRequestURL: "http://endpoint.com/test?version=1", + ActionsIDTokenRequestToken: "token-1337", + HTTPTransport: fixtures.MappingTransport{ + "GET /test?version=1&audience=token-audience": { + Status: http.StatusOK, + ExpectedHeaders: map[string]string{ + "Authorization": "Bearer token-1337", + "Accept": "application/json", + }, + Response: `{"value": "id-token-42"}`, + }, + "POST /oidc/accounts/ac123/v1/token": { + Status: http.StatusOK, + ExpectedHeaders: map[string]string{ + "Content-Type": "application/x-www-form-urlencoded", + }, + Response: map[string]string{ + "token_type": "access-token", + "access_token": "test-auth-token", + "refresh_token": "refresh", + "expires_on": "0", + }, + }, + }, + }, + wantHeaders: map[string]string{ + "Authorization": "access-token test-auth-token", + }, + }, + { + desc: "default token audience account", + cfg: &Config{ + ClientID: "client-id", + AccountID: "ac123", + Host: "https://accounts.databricks.com", + ActionsIDTokenRequestURL: "http://endpoint.com/test?version=1", + ActionsIDTokenRequestToken: "token-1337", + HTTPTransport: fixtures.MappingTransport{ + "GET /test?version=1&audience=ac123": { + Status: http.StatusOK, + ExpectedHeaders: map[string]string{ + "Authorization": "Bearer token-1337", + "Accept": "application/json", + }, + Response: `{"value": "id-token-42"}`, + }, + "POST /oidc/accounts/ac123/v1/token": { + Status: http.StatusOK, + ExpectedHeaders: map[string]string{ + "Content-Type": "application/x-www-form-urlencoded", + }, + Response: map[string]string{ + "token_type": "access-token", + "access_token": "test-auth-token", + "refresh_token": "refresh", + "expires_on": "0", + }, + }, + }, + }, + wantHeaders: map[string]string{ + "Authorization": "access-token test-auth-token", + }, + }, + { + desc: "default token audience workspace", + cfg: &Config{ + ClientID: "client-id", + Host: "https://host.com", + ActionsIDTokenRequestURL: "http://endpoint.com/test?version=1", + ActionsIDTokenRequestToken: "token-1337", + HTTPTransport: fixtures.MappingTransport{ + "GET /oidc/.well-known/oauth-authorization-server": { + Response: u2m.OAuthAuthorizationServer{ + AuthorizationEndpoint: "https://host.com/auth", + TokenEndpoint: "https://host.com/oidc/v1/token", + }, + }, + "GET /test?version=1&audience=https://host.com/oidc/v1/token": { + Status: http.StatusOK, + ExpectedHeaders: map[string]string{ + "Authorization": "Bearer token-1337", + "Accept": "application/json", + }, + Response: `{"value": "id-token-42"}`, + }, + "POST /oidc/v1/token": { + Status: http.StatusOK, + ExpectedHeaders: map[string]string{ + "Content-Type": "application/x-www-form-urlencoded", + }, + Response: map[string]string{ + "token_type": "access-token", + "access_token": "test-auth-token", + "refresh_token": "refresh", + "expires_on": "0", + }, + }, + }, + }, + wantHeaders: map[string]string{ + "Authorization": "access-token test-auth-token", + }, + }, } for _, tc := range testCases { diff --git a/config/oauth_visitors.go b/config/oauth_visitors.go index fc7a3d153..10556c19d 100644 --- a/config/oauth_visitors.go +++ b/config/oauth_visitors.go @@ -35,9 +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 { - cts := auth.NewCachedTokenSource(authconv.AuthTokenSource(inner)) + return refreshableAuthVisitor(authconv.AuthTokenSource(inner), context.Background()) +} + +// The same as serviceToServiceVisitor, but without a secondary token source. +func refreshableAuthVisitor(inner auth.TokenSource, ctx context.Context) func(r *http.Request) error { + cts := auth.NewCachedTokenSource(inner) return func(r *http.Request) error { - inner, err := cts.Token(context.Background()) + inner, err := cts.Token(ctx) if err != nil { return fmt.Errorf("inner token: %w", err) } diff --git a/config/token_source_strategy.go b/config/token_source_strategy.go index 4aa07f30d..099e2ae5e 100644 --- a/config/token_source_strategy.go +++ b/config/token_source_strategy.go @@ -30,13 +30,14 @@ func (t *TokenSourceStrategy) Configure(ctx context.Context, cfg *Config) (crede // We don't want to fail here because it's possible that the supplier is enabled // without the user action. For instance, jobs running in GitHub will have // OIDC environment variables added automatically - if _, err := t.tokenSource.Token(ctx); err != nil { + cached := auth.NewCachedTokenSource(t.tokenSource) + if _, err := cached.Token(ctx); err != nil { logger.Debugf(ctx, fmt.Sprintf("Skipping %s due to error: %v", t.name, err)) return nil, nil } - visitor := refreshableVisitor(authconv.OAuth2TokenSource(t.tokenSource)) - return credentials.NewOAuthCredentialsProvider(visitor, authconv.OAuth2TokenSource(t.tokenSource).Token), nil + visitor := refreshableAuthVisitor(cached, ctx) + return credentials.NewOAuthCredentialsProvider(visitor, authconv.OAuth2TokenSource(cached).Token), nil } func (t *TokenSourceStrategy) Name() string { From 8738b46bb38a6e78b3fd955da9190f812dd2651d Mon Sep 17 00:00:00 2001 From: Hector Castejon Diaz Date: Mon, 7 Apr 2025 11:58:01 +0200 Subject: [PATCH 11/30] Remove conf + test --- config/auth_databricks_wif.go | 55 ++- config/auth_databricks_wif_test.go | 645 ++++++++++++++------------- config/token_provider_github_oidc.go | 24 +- 3 files changed, 375 insertions(+), 349 deletions(-) diff --git a/config/auth_databricks_wif.go b/config/auth_databricks_wif.go index 14f1b75e8..e370430fc 100644 --- a/config/auth_databricks_wif.go +++ b/config/auth_databricks_wif.go @@ -5,6 +5,7 @@ import ( "errors" "net/url" + "github.com/databricks/databricks-sdk-go/credentials/u2m" "github.com/databricks/databricks-sdk-go/logger" "golang.org/x/oauth2" "golang.org/x/oauth2/clientcredentials" @@ -14,7 +15,9 @@ import ( func WifTokenCredentialStrategies(cfg *Config) []CredentialsStrategy { providers := map[string]TokenProvider{ "github-oidc-databricks-wif": &GithubProvider{ - cfg: cfg, + actionsIDTokenRequestURL: cfg.ActionsIDTokenRequestURL, + actionsIDTokenRequestToken: cfg.ActionsIDTokenRequestToken, + refreshClient: cfg.refreshClient, }, // Add new providers at the end of the list } @@ -31,8 +34,13 @@ func newWifTokenStrategy( tokenProvider TokenProvider, ) CredentialsStrategy { wifTokenExchange := &wifTokenExchange{ - cfg: cfg, - tokenProvider: tokenProvider, + clientID: cfg.ClientID, + account: cfg.IsAccountClient(), + accountID: cfg.AccountID, + host: cfg.Host, + tokenEndpointProvider: cfg.getOidcEndpoints, + audience: cfg.TokenAudience, + tokenProvider: tokenProvider, } return NewTokenSourceStrategy(name, wifTokenExchange) } @@ -40,20 +48,25 @@ func newWifTokenStrategy( // wifTokenExchange is a auth.TokenSource which exchanges a token using // Workload Identity Federation. type wifTokenExchange struct { - cfg *Config - tokenProvider TokenProvider + clientID string + account bool + accountID string + host string + tokenEndpointProvider func(ctx context.Context) (*u2m.OAuthAuthorizationServer, error) + audience string + tokenProvider TokenProvider } func (w *wifTokenExchange) Token(ctx context.Context) (*oauth2.Token, error) { - if w.cfg.ClientID == "" { - logger.Debugf(ctx, "Missing cfg.ClientID") - return nil, errors.New("missing cfg.ClientID") + if w.clientID == "" { + logger.Debugf(ctx, "Missing ClientID") + return nil, errors.New("missing ClientID") } - if w.cfg.Host == "" { - logger.Debugf(ctx, "Missing cfg.Host") - return nil, errors.New("missing cfg.Host") + if w.host == "" { + logger.Debugf(ctx, "Missing Host") + return nil, errors.New("missing Host") } - audience, err := w.getAudience(ctx) + audience, err := w.determineAudience(ctx) if err != nil { return nil, err } @@ -61,12 +74,12 @@ func (w *wifTokenExchange) Token(ctx context.Context) (*oauth2.Token, error) { if err != nil { return nil, err } - endpoints, err := w.cfg.getOidcEndpoints(ctx) + endpoints, err := w.tokenEndpointProvider(ctx) if err != nil { return nil, err } c := &clientcredentials.Config{ - ClientID: w.cfg.ClientID, + ClientID: w.clientID, AuthStyle: oauth2.AuthStyleInParams, TokenURL: endpoints.TokenEndpoint, Scopes: []string{"all-apis"}, @@ -79,18 +92,18 @@ func (w *wifTokenExchange) Token(ctx context.Context) (*oauth2.Token, error) { return c.Token(ctx) } -func (w *wifTokenExchange) getAudience(ctx context.Context) (string, error) { - if w.cfg.TokenAudience != "" { - return w.cfg.TokenAudience, nil +func (w *wifTokenExchange) determineAudience(ctx context.Context) (string, error) { + if w.audience != "" { + return w.audience, nil } // For Databricks Accounts, the account id is the default audience. - if w.cfg.IsAccountClient() { - return w.cfg.AccountID, nil + if w.account { + return w.accountID, nil } - // For Databricks Workspaces, the auth endpoint is the default audience. - endpoints, err := w.cfg.getOidcEndpoints(ctx) + endpoints, err := w.tokenEndpointProvider(ctx) if err != nil { return "", err } + // For Databricks Workspaces, the auth endpoint is the default audience. return endpoints.TokenEndpoint, nil } diff --git a/config/auth_databricks_wif_test.go b/config/auth_databricks_wif_test.go index 363ea97af..e343c87b3 100644 --- a/config/auth_databricks_wif_test.go +++ b/config/auth_databricks_wif_test.go @@ -1,51 +1,53 @@ package config import ( - "net/http" + "context" "testing" - "github.com/databricks/databricks-sdk-go/common/environment" "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" ) +type mockIdTokenProvider struct { + idToken *IDToken + err error +} + +func (m *mockIdTokenProvider) IDToken(ctx context.Context, audience string) (*IDToken, error) { + return m.idToken, m.err +} + func TestDatabricksGithubWIFCredentials(t *testing.T) { testCases := []struct { - desc string - cfg *Config - wantHeaders map[string]string - wantErrPrefix *string + desc string + cfg *Config + idToken string + tokenProviderError error + wantToken string + wantErrPrefix *string }{ - { - desc: "not an databricks config", - cfg: &Config{ - DatabricksEnvironment: &environment.DatabricksEnvironment{ - Cloud: "foo-bar-cloud", - }, - }, - wantErrPrefix: errPrefix("federated-oidc-github auth: not configured"), - }, - { - desc: "missing host", - cfg: &Config{ - ClientID: "client-id", - TokenAudience: "token-audience", - ActionsIDTokenRequestURL: "http://endpoint.com/test?version=1", - ActionsIDTokenRequestToken: "token-1337", - }, - wantErrPrefix: errPrefix("federated-oidc-github auth: not configured"), - }, - { - desc: "missing client ID", - cfg: &Config{ - Host: "http://host.com/test", - TokenAudience: "token-audience", - ActionsIDTokenRequestURL: "http://endpoint.com/test?version=1", - ActionsIDTokenRequestToken: "token-1337", - }, - wantErrPrefix: errPrefix("federated-oidc-github auth: not configured"), - }, + // { + // desc: "missing host", + // cfg: &Config{ + // ClientID: "client-id", + // TokenAudience: "token-audience", + // ActionsIDTokenRequestURL: "http://endpoint.com/test?version=1", + // ActionsIDTokenRequestToken: "token-1337", + // }, + // wantErrPrefix: errPrefix("missing Host"), + // }, + // { + // desc: "missing client ID", + // cfg: &Config{ + // Host: "http://host.com/test", + // TokenAudience: "token-audience", + // ActionsIDTokenRequestURL: "http://endpoint.com/test?version=1", + // ActionsIDTokenRequestToken: "token-1337", + // }, + // wantErrPrefix: errPrefix("missing ClientID"), + // }, { desc: "missing env ACTIONS_ID_TOKEN_REQUEST_TOKEN", cfg: &Config{ @@ -62,274 +64,274 @@ func TestDatabricksGithubWIFCredentials(t *testing.T) { }, }, }, - wantErrPrefix: errPrefix("federated-oidc-github auth: not configured"), - }, - { - desc: "missing env ACTIONS_ID_TOKEN_REQUEST_URL", - cfg: &Config{ - ClientID: "client-id", - Host: "http://host.com/test", - TokenAudience: "token-audience", - ActionsIDTokenRequestToken: "token-1337", - HTTPTransport: fixtures.MappingTransport{ - "GET /oidc/.well-known/oauth-authorization-server": { - Response: u2m.OAuthAuthorizationServer{ - AuthorizationEndpoint: "https://host.com/auth", - TokenEndpoint: "https://host.com/oidc/v1/token", - }, - }, - }, - }, - wantErrPrefix: errPrefix("federated-oidc-github auth: not configured"), - }, - { - desc: "databricks token exchange server error", - cfg: &Config{ - ClientID: "client-id", - Host: "http://host.com/test", - TokenAudience: "token-audience", - ActionsIDTokenRequestURL: "http://endpoint.com/test?version=1", - ActionsIDTokenRequestToken: "token-1337", - HTTPTransport: fixtures.MappingTransport{ - "GET /oidc/.well-known/oauth-authorization-server": { - Response: u2m.OAuthAuthorizationServer{ - AuthorizationEndpoint: "https://host.com/auth", - TokenEndpoint: "https://host.com/oidc/v1/token", - }, - }, - "GET /test?version=1&audience=token-audience": { - Status: http.StatusInternalServerError, - ExpectedHeaders: map[string]string{ - "Authorization": "Bearer token-1337", - "Accept": "application/json", - }, - }, - }, - }, - wantErrPrefix: errPrefix("federated-oidc-github"), - }, - { - desc: "databricks workspace server error", - cfg: &Config{ - ClientID: "client-id", - Host: "http://host.com/test", - TokenAudience: "token-audience", - ActionsIDTokenRequestURL: "http://endpoint.com/test?version=1", - ActionsIDTokenRequestToken: "token-1337", - HTTPTransport: fixtures.MappingTransport{ - "GET /oidc/.well-known/oauth-authorization-server": { - Response: u2m.OAuthAuthorizationServer{ - AuthorizationEndpoint: "https://host.com/auth", - TokenEndpoint: "https://host.com/oidc/v1/token", - }, - }, - "GET /test?version=1&audience=token-audience": { - Status: http.StatusOK, - ExpectedHeaders: map[string]string{ - "Authorization": "Bearer token-1337", - "Accept": "application/json", - }, - Response: `{"value": "id-token-42"}`, - }, - "POST /oidc/v1/token": { - Status: http.StatusInternalServerError, - ExpectedHeaders: map[string]string{ - "Content-Type": "application/x-www-form-urlencoded", - }, - }, - }, - }, - wantErrPrefix: errPrefix("federated-oidc-github auth: not configured"), - }, - { - desc: "invalid auth token", - cfg: &Config{ - ClientID: "client-id", - Host: "http://host.com/test", - TokenAudience: "token-audience", - ActionsIDTokenRequestURL: "http://endpoint.com/test?version=1", - ActionsIDTokenRequestToken: "token-1337", - HTTPTransport: fixtures.MappingTransport{ - "GET /oidc/.well-known/oauth-authorization-server": { - Response: u2m.OAuthAuthorizationServer{ - AuthorizationEndpoint: "https://host.com/auth", - TokenEndpoint: "https://host.com/oidc/v1/token", - }, - }, - "GET /test?version=1&audience=token-audience": { - Status: http.StatusOK, - ExpectedHeaders: map[string]string{ - "Authorization": "Bearer token-1337", - "Accept": "application/json", - }, - Response: `{"value": "id-token-42"}`, - }, - "POST /oidc/v1/token": { - Status: http.StatusOK, - ExpectedHeaders: map[string]string{ - "Content-Type": "application/x-www-form-urlencoded", - }, - Response: map[string]string{ - "foo": "bar", - }, - }, - }, - }, - wantErrPrefix: errPrefix("federated-oidc-github auth: not configured"), + wantErrPrefix: errPrefix("missing ActionsIdTokenRequestToken"), }, - { - desc: "success workspace", - cfg: &Config{ - ClientID: "client-id", - Host: "http://host.com/test", - TokenAudience: "token-audience", - ActionsIDTokenRequestURL: "http://endpoint.com/test?version=1", - ActionsIDTokenRequestToken: "token-1337", - HTTPTransport: fixtures.MappingTransport{ - "GET /oidc/.well-known/oauth-authorization-server": { - Response: u2m.OAuthAuthorizationServer{ - AuthorizationEndpoint: "https://host.com/auth", - TokenEndpoint: "https://host.com/oidc/v1/token", - }, - }, - "GET /test?version=1&audience=token-audience": { - Status: http.StatusOK, - ExpectedHeaders: map[string]string{ - "Authorization": "Bearer token-1337", - "Accept": "application/json", - }, - Response: `{"value": "id-token-42"}`, - }, - "POST /oidc/v1/token": { + // { + // desc: "missing env ACTIONS_ID_TOKEN_REQUEST_URL", + // cfg: &Config{ + // ClientID: "client-id", + // Host: "http://host.com/test", + // TokenAudience: "token-audience", + // ActionsIDTokenRequestToken: "token-1337", + // HTTPTransport: fixtures.MappingTransport{ + // "GET /oidc/.well-known/oauth-authorization-server": { + // Response: u2m.OAuthAuthorizationServer{ + // AuthorizationEndpoint: "https://host.com/auth", + // TokenEndpoint: "https://host.com/oidc/v1/token", + // }, + // }, + // }, + // }, + // wantErrPrefix: errPrefix("federated-oidc-github auth: not configured"), + // }, + // { + // desc: "databricks token exchange server error", + // cfg: &Config{ + // ClientID: "client-id", + // Host: "http://host.com/test", + // TokenAudience: "token-audience", + // ActionsIDTokenRequestURL: "http://endpoint.com/test?version=1", + // ActionsIDTokenRequestToken: "token-1337", + // HTTPTransport: fixtures.MappingTransport{ + // "GET /oidc/.well-known/oauth-authorization-server": { + // Response: u2m.OAuthAuthorizationServer{ + // AuthorizationEndpoint: "https://host.com/auth", + // TokenEndpoint: "https://host.com/oidc/v1/token", + // }, + // }, + // "GET /test?version=1&audience=token-audience": { + // Status: http.StatusInternalServerError, + // ExpectedHeaders: map[string]string{ + // "Authorization": "Bearer token-1337", + // "Accept": "application/json", + // }, + // }, + // }, + // }, + // wantErrPrefix: errPrefix("federated-oidc-github"), + // }, + // { + // desc: "databricks workspace server error", + // cfg: &Config{ + // ClientID: "client-id", + // Host: "http://host.com/test", + // TokenAudience: "token-audience", + // ActionsIDTokenRequestURL: "http://endpoint.com/test?version=1", + // ActionsIDTokenRequestToken: "token-1337", + // HTTPTransport: fixtures.MappingTransport{ + // "GET /oidc/.well-known/oauth-authorization-server": { + // Response: u2m.OAuthAuthorizationServer{ + // AuthorizationEndpoint: "https://host.com/auth", + // TokenEndpoint: "https://host.com/oidc/v1/token", + // }, + // }, + // "GET /test?version=1&audience=token-audience": { + // Status: http.StatusOK, + // ExpectedHeaders: map[string]string{ + // "Authorization": "Bearer token-1337", + // "Accept": "application/json", + // }, + // Response: `{"value": "id-token-42"}`, + // }, + // "POST /oidc/v1/token": { + // Status: http.StatusInternalServerError, + // ExpectedHeaders: map[string]string{ + // "Content-Type": "application/x-www-form-urlencoded", + // }, + // }, + // }, + // }, + // wantErrPrefix: errPrefix("federated-oidc-github auth: not configured"), + // }, + // { + // desc: "invalid auth token", + // cfg: &Config{ + // ClientID: "client-id", + // Host: "http://host.com/test", + // TokenAudience: "token-audience", + // ActionsIDTokenRequestURL: "http://endpoint.com/test?version=1", + // ActionsIDTokenRequestToken: "token-1337", + // HTTPTransport: fixtures.MappingTransport{ + // "GET /oidc/.well-known/oauth-authorization-server": { + // Response: u2m.OAuthAuthorizationServer{ + // AuthorizationEndpoint: "https://host.com/auth", + // TokenEndpoint: "https://host.com/oidc/v1/token", + // }, + // }, + // "GET /test?version=1&audience=token-audience": { + // Status: http.StatusOK, + // ExpectedHeaders: map[string]string{ + // "Authorization": "Bearer token-1337", + // "Accept": "application/json", + // }, + // Response: `{"value": "id-token-42"}`, + // }, + // "POST /oidc/v1/token": { + // Status: http.StatusOK, + // ExpectedHeaders: map[string]string{ + // "Content-Type": "application/x-www-form-urlencoded", + // }, + // Response: map[string]string{ + // "foo": "bar", + // }, + // }, + // }, + // }, + // wantErrPrefix: errPrefix("federated-oidc-github auth: not configured"), + // }, + // { + // desc: "success workspace", + // cfg: &Config{ + // ClientID: "client-id", + // Host: "http://host.com/test", + // TokenAudience: "token-audience", + // ActionsIDTokenRequestURL: "http://endpoint.com/test?version=1", + // ActionsIDTokenRequestToken: "token-1337", + // HTTPTransport: fixtures.MappingTransport{ + // "GET /oidc/.well-known/oauth-authorization-server": { + // Response: u2m.OAuthAuthorizationServer{ + // AuthorizationEndpoint: "https://host.com/auth", + // TokenEndpoint: "https://host.com/oidc/v1/token", + // }, + // }, + // "GET /test?version=1&audience=token-audience": { + // Status: http.StatusOK, + // ExpectedHeaders: map[string]string{ + // "Authorization": "Bearer token-1337", + // "Accept": "application/json", + // }, + // Response: `{"value": "id-token-42"}`, + // }, + // "POST /oidc/v1/token": { - Status: http.StatusOK, - ExpectedHeaders: map[string]string{ - "Content-Type": "application/x-www-form-urlencoded", - }, - Response: map[string]string{ - "token_type": "access-token", - "access_token": "test-auth-token", - "refresh_token": "refresh", - "expires_on": "0", - }, - }, - }, - }, - wantHeaders: map[string]string{ - "Authorization": "access-token test-auth-token", - }, - }, - { - desc: "success account", - cfg: &Config{ - ClientID: "client-id", - AccountID: "ac123", - Host: "https://accounts.databricks.com", - TokenAudience: "token-audience", - ActionsIDTokenRequestURL: "http://endpoint.com/test?version=1", - ActionsIDTokenRequestToken: "token-1337", - HTTPTransport: fixtures.MappingTransport{ - "GET /test?version=1&audience=token-audience": { - Status: http.StatusOK, - ExpectedHeaders: map[string]string{ - "Authorization": "Bearer token-1337", - "Accept": "application/json", - }, - Response: `{"value": "id-token-42"}`, - }, - "POST /oidc/accounts/ac123/v1/token": { - Status: http.StatusOK, - ExpectedHeaders: map[string]string{ - "Content-Type": "application/x-www-form-urlencoded", - }, - Response: map[string]string{ - "token_type": "access-token", - "access_token": "test-auth-token", - "refresh_token": "refresh", - "expires_on": "0", - }, - }, - }, - }, - wantHeaders: map[string]string{ - "Authorization": "access-token test-auth-token", - }, - }, - { - desc: "default token audience account", - cfg: &Config{ - ClientID: "client-id", - AccountID: "ac123", - Host: "https://accounts.databricks.com", - ActionsIDTokenRequestURL: "http://endpoint.com/test?version=1", - ActionsIDTokenRequestToken: "token-1337", - HTTPTransport: fixtures.MappingTransport{ - "GET /test?version=1&audience=ac123": { - Status: http.StatusOK, - ExpectedHeaders: map[string]string{ - "Authorization": "Bearer token-1337", - "Accept": "application/json", - }, - Response: `{"value": "id-token-42"}`, - }, - "POST /oidc/accounts/ac123/v1/token": { - Status: http.StatusOK, - ExpectedHeaders: map[string]string{ - "Content-Type": "application/x-www-form-urlencoded", - }, - Response: map[string]string{ - "token_type": "access-token", - "access_token": "test-auth-token", - "refresh_token": "refresh", - "expires_on": "0", - }, - }, - }, - }, - wantHeaders: map[string]string{ - "Authorization": "access-token test-auth-token", - }, - }, - { - desc: "default token audience workspace", - cfg: &Config{ - ClientID: "client-id", - Host: "https://host.com", - ActionsIDTokenRequestURL: "http://endpoint.com/test?version=1", - ActionsIDTokenRequestToken: "token-1337", - HTTPTransport: fixtures.MappingTransport{ - "GET /oidc/.well-known/oauth-authorization-server": { - Response: u2m.OAuthAuthorizationServer{ - AuthorizationEndpoint: "https://host.com/auth", - TokenEndpoint: "https://host.com/oidc/v1/token", - }, - }, - "GET /test?version=1&audience=https://host.com/oidc/v1/token": { - Status: http.StatusOK, - ExpectedHeaders: map[string]string{ - "Authorization": "Bearer token-1337", - "Accept": "application/json", - }, - Response: `{"value": "id-token-42"}`, - }, - "POST /oidc/v1/token": { - Status: http.StatusOK, - ExpectedHeaders: map[string]string{ - "Content-Type": "application/x-www-form-urlencoded", - }, - Response: map[string]string{ - "token_type": "access-token", - "access_token": "test-auth-token", - "refresh_token": "refresh", - "expires_on": "0", - }, - }, - }, - }, - wantHeaders: map[string]string{ - "Authorization": "access-token test-auth-token", - }, - }, + // Status: http.StatusOK, + // ExpectedHeaders: map[string]string{ + // "Content-Type": "application/x-www-form-urlencoded", + // }, + // Response: map[string]string{ + // "token_type": "access-token", + // "access_token": "test-auth-token", + // "refresh_token": "refresh", + // "expires_on": "0", + // }, + // }, + // }, + // }, + // wantHeaders: map[string]string{ + // "Authorization": "access-token test-auth-token", + // }, + // }, + // { + // desc: "success account", + // cfg: &Config{ + // ClientID: "client-id", + // AccountID: "ac123", + // Host: "https://accounts.databricks.com", + // TokenAudience: "token-audience", + // ActionsIDTokenRequestURL: "http://endpoint.com/test?version=1", + // ActionsIDTokenRequestToken: "token-1337", + // HTTPTransport: fixtures.MappingTransport{ + // "GET /test?version=1&audience=token-audience": { + // Status: http.StatusOK, + // ExpectedHeaders: map[string]string{ + // "Authorization": "Bearer token-1337", + // "Accept": "application/json", + // }, + // Response: `{"value": "id-token-42"}`, + // }, + // "POST /oidc/accounts/ac123/v1/token": { + // Status: http.StatusOK, + // ExpectedHeaders: map[string]string{ + // "Content-Type": "application/x-www-form-urlencoded", + // }, + // Response: map[string]string{ + // "token_type": "access-token", + // "access_token": "test-auth-token", + // "refresh_token": "refresh", + // "expires_on": "0", + // }, + // }, + // }, + // }, + // wantHeaders: map[string]string{ + // "Authorization": "access-token test-auth-token", + // }, + // }, + // { + // desc: "default token audience account", + // cfg: &Config{ + // ClientID: "client-id", + // AccountID: "ac123", + // Host: "https://accounts.databricks.com", + // ActionsIDTokenRequestURL: "http://endpoint.com/test?version=1", + // ActionsIDTokenRequestToken: "token-1337", + // HTTPTransport: fixtures.MappingTransport{ + // "GET /test?version=1&audience=ac123": { + // Status: http.StatusOK, + // ExpectedHeaders: map[string]string{ + // "Authorization": "Bearer token-1337", + // "Accept": "application/json", + // }, + // Response: `{"value": "id-token-42"}`, + // }, + // "POST /oidc/accounts/ac123/v1/token": { + // Status: http.StatusOK, + // ExpectedHeaders: map[string]string{ + // "Content-Type": "application/x-www-form-urlencoded", + // }, + // Response: map[string]string{ + // "token_type": "access-token", + // "access_token": "test-auth-token", + // "refresh_token": "refresh", + // "expires_on": "0", + // }, + // }, + // }, + // }, + // wantHeaders: map[string]string{ + // "Authorization": "access-token test-auth-token", + // }, + // }, + // { + // desc: "default token audience workspace", + // cfg: &Config{ + // ClientID: "client-id", + // Host: "https://host.com", + // ActionsIDTokenRequestURL: "http://endpoint.com/test?version=1", + // ActionsIDTokenRequestToken: "token-1337", + // HTTPTransport: fixtures.MappingTransport{ + // "GET /oidc/.well-known/oauth-authorization-server": { + // Response: u2m.OAuthAuthorizationServer{ + // AuthorizationEndpoint: "https://host.com/auth", + // TokenEndpoint: "https://host.com/oidc/v1/token", + // }, + // }, + // "GET /test?version=1&audience=https://host.com/oidc/v1/token": { + // Status: http.StatusOK, + // ExpectedHeaders: map[string]string{ + // "Authorization": "Bearer token-1337", + // "Accept": "application/json", + // }, + // Response: `{"value": "id-token-42"}`, + // }, + // "POST /oidc/v1/token": { + // Status: http.StatusOK, + // ExpectedHeaders: map[string]string{ + // "Content-Type": "application/x-www-form-urlencoded", + // }, + // Response: map[string]string{ + // "token_type": "access-token", + // "access_token": "test-auth-token", + // "refresh_token": "refresh", + // "expires_on": "0", + // }, + // }, + // }, + // }, + // wantHeaders: map[string]string{ + // "Authorization": "access-token test-auth-token", + // }, + //}, } for _, tc := range testCases { @@ -337,29 +339,38 @@ func TestDatabricksGithubWIFCredentials(t *testing.T) { // if tc.desc != "" { // t.Skip() // } - p := &GithubProvider{ - cfg: tc.cfg, + p := &mockIdTokenProvider{ + idToken: &IDToken{ + Value: tc.idToken, + }, + err: tc.tokenProviderError, } - tc.cfg.Credentials = newWifTokenStrategy("federated-oidc-github", tc.cfg, p) // only test this credential strategy - tc.cfg.DebugHeaders = true - if tc.wantHeaders == nil { - tc.wantHeaders = map[string]string{} + resolveErr := tc.cfg.EnsureResolved() + if resolveErr != nil { + t.Errorf("EnsureResolved(): got error %q, want none", resolveErr) } - - req, _ := http.NewRequest("GET", "http://localhost", nil) - gotErr := tc.cfg.Authenticate(req) - gotHeaders := map[string]string{} - for h := range req.Header { - gotHeaders[h] = req.Header.Get(h) + ex := &wifTokenExchange{ + clientID: tc.cfg.ClientID, + account: tc.cfg.IsAccountClient(), + accountID: tc.cfg.AccountID, + host: tc.cfg.Host, + tokenEndpointProvider: tc.cfg.getOidcEndpoints, + audience: tc.cfg.TokenAudience, + tokenProvider: p, } - + ctx := context.WithValue(context.Background(), oauth2.HTTPClient, tc.cfg.HTTPTransport) + token, gotErr := ex.Token(ctx) if tc.wantErrPrefix == nil && gotErr != nil { - t.Errorf("Authenticate(): got error %q, want none", gotErr) + t.Errorf("Token(ctx): got error %q, want none", gotErr) } if tc.wantErrPrefix != nil && !hasPrefix(gotErr, *tc.wantErrPrefix) { - t.Errorf("Authenticate(): got error %q, want error with prefix %q", gotErr, *tc.wantErrPrefix) + t.Errorf("Token(ctx): got error %q, want error with prefix %q", gotErr, *tc.wantErrPrefix) + } + tokenValue := "" + if token != nil { + tokenValue = token.AccessToken } - if diff := cmp.Diff(tc.wantHeaders, gotHeaders); diff != "" { + if diff := cmp.Diff(tc.wantToken, tokenValue); diff != "" { t.Errorf("Authenticate(): mismatch (-want +got):\n%s", diff) } }) diff --git a/config/token_provider_github_oidc.go b/config/token_provider_github_oidc.go index 021203e0a..9cec3a39b 100644 --- a/config/token_provider_github_oidc.go +++ b/config/token_provider_github_oidc.go @@ -10,30 +10,32 @@ import ( ) type GithubProvider struct { - cfg *Config + actionsIDTokenRequestURL string + actionsIDTokenRequestToken string + refreshClient *httpclient.ApiClient } func (g *GithubProvider) IDToken(ctx context.Context, audience string) (*IDToken, error) { - if g.cfg.ActionsIDTokenRequestURL == "" { - logger.Debugf(ctx, "Missing cfg.ActionsIDTokenRequestURL, likely not calling from a Github action") - return nil, errors.New("missing cfg.ActionsIDTokenRequestURL") + if g.actionsIDTokenRequestURL == "" { + logger.Debugf(ctx, "Missing ActionsIDTokenRequestURL, likely not calling from a Github action") + return nil, errors.New("missing ActionsIDTokenRequestURL") } - if g.cfg.ActionsIDTokenRequestToken == "" { - logger.Debugf(ctx, "Missing cfg.ActionsIDTokenRequestToken, likely not calling from a Github action") - return nil, errors.New("missing cfg.ActionsIDTokenRequestToken") + if g.actionsIDTokenRequestToken == "" { + logger.Debugf(ctx, "Missing ActionsIDTokenRequestToken, likely not calling from a Github action") + return nil, errors.New("missing ActionsIDTokenRequestToken") } resp := &IDToken{} - requestUrl := g.cfg.ActionsIDTokenRequestURL + requestUrl := g.actionsIDTokenRequestURL if audience != "" { requestUrl = fmt.Sprintf("%s&audience=%s", requestUrl, audience) } - err := g.cfg.refreshClient.Do(ctx, "GET", requestUrl, - httpclient.WithRequestHeader("Authorization", fmt.Sprintf("Bearer %s", g.cfg.ActionsIDTokenRequestToken)), + err := g.refreshClient.Do(ctx, "GET", requestUrl, + httpclient.WithRequestHeader("Authorization", fmt.Sprintf("Bearer %s", g.actionsIDTokenRequestToken)), httpclient.WithResponseUnmarshal(resp), ) if err != nil { - return nil, fmt.Errorf("failed to request ID token from %s: %w", g.cfg.ActionsIDTokenRequestURL, err) + return nil, fmt.Errorf("failed to request ID token from %s: %w", g.actionsIDTokenRequestURL, err) } return resp, nil From 506a146b97ad8a0f5f79b0f2772e48a7ec9cca6c Mon Sep 17 00:00:00 2001 From: Hector Castejon Diaz Date: Mon, 7 Apr 2025 14:26:03 +0200 Subject: [PATCH 12/30] More tests --- config/auth_databricks_wif_test.go | 541 ++++++++++------------ config/token_provider_github_oidc_test.go | 91 ++++ 2 files changed, 329 insertions(+), 303 deletions(-) create mode 100644 config/token_provider_github_oidc_test.go diff --git a/config/auth_databricks_wif_test.go b/config/auth_databricks_wif_test.go index e343c87b3..52154ebcb 100644 --- a/config/auth_databricks_wif_test.go +++ b/config/auth_databricks_wif_test.go @@ -2,6 +2,9 @@ package config import ( "context" + "errors" + "net/http" + "net/url" "testing" "github.com/databricks/databricks-sdk-go/credentials/u2m" @@ -11,12 +14,18 @@ import ( ) type mockIdTokenProvider struct { - idToken *IDToken + // input + audience string + // output + idToken string err error } func (m *mockIdTokenProvider) IDToken(ctx context.Context, audience string) (*IDToken, error) { - return m.idToken, m.err + m.audience = audience + return &IDToken{ + m.idToken, + }, m.err } func TestDatabricksGithubWIFCredentials(t *testing.T) { @@ -24,37 +33,46 @@ func TestDatabricksGithubWIFCredentials(t *testing.T) { desc string cfg *Config idToken string + expectedAudience string tokenProviderError error wantToken string wantErrPrefix *string }{ - // { - // desc: "missing host", - // cfg: &Config{ - // ClientID: "client-id", - // TokenAudience: "token-audience", - // ActionsIDTokenRequestURL: "http://endpoint.com/test?version=1", - // ActionsIDTokenRequestToken: "token-1337", - // }, - // wantErrPrefix: errPrefix("missing Host"), - // }, - // { - // desc: "missing client ID", - // cfg: &Config{ - // Host: "http://host.com/test", - // TokenAudience: "token-audience", - // ActionsIDTokenRequestURL: "http://endpoint.com/test?version=1", - // ActionsIDTokenRequestToken: "token-1337", - // }, - // wantErrPrefix: errPrefix("missing ClientID"), - // }, { - desc: "missing env ACTIONS_ID_TOKEN_REQUEST_TOKEN", + desc: "missing host", cfg: &Config{ - ClientID: "client-id", - Host: "http://host.com/test", - TokenAudience: "token-audience", - ActionsIDTokenRequestURL: "http://endpoint.com/test?version=1", + ClientID: "client-id", + TokenAudience: "token-audience", + }, + wantErrPrefix: errPrefix("missing Host"), + }, + { + desc: "missing client ID", + cfg: &Config{ + Host: "http://host.com/test", + TokenAudience: "token-audience", + }, + wantErrPrefix: errPrefix("missing ClientID"), + }, + { + desc: "auth server error", + cfg: &Config{ + ClientID: "client-id", + Host: "http://host.com/test", + HTTPTransport: fixtures.MappingTransport{ + "GET /oidc/.well-known/oauth-authorization-server": { + Status: http.StatusNotFound, + }, + }, + }, + wantErrPrefix: errPrefix("databricks OAuth is not supported for this host"), + }, + { + desc: "token provider error", + cfg: &Config{ + ClientID: "client-id", + Host: "http://host.com/test", + TokenAudience: "token-audience", HTTPTransport: fixtures.MappingTransport{ "GET /oidc/.well-known/oauth-authorization-server": { Response: u2m.OAuthAuthorizationServer{ @@ -64,274 +82,190 @@ func TestDatabricksGithubWIFCredentials(t *testing.T) { }, }, }, - wantErrPrefix: errPrefix("missing ActionsIdTokenRequestToken"), + expectedAudience: "token-audience", + tokenProviderError: errors.New("error getting id token"), + wantErrPrefix: errPrefix("error getting id token"), }, - // { - // desc: "missing env ACTIONS_ID_TOKEN_REQUEST_URL", - // cfg: &Config{ - // ClientID: "client-id", - // Host: "http://host.com/test", - // TokenAudience: "token-audience", - // ActionsIDTokenRequestToken: "token-1337", - // HTTPTransport: fixtures.MappingTransport{ - // "GET /oidc/.well-known/oauth-authorization-server": { - // Response: u2m.OAuthAuthorizationServer{ - // AuthorizationEndpoint: "https://host.com/auth", - // TokenEndpoint: "https://host.com/oidc/v1/token", - // }, - // }, - // }, - // }, - // wantErrPrefix: errPrefix("federated-oidc-github auth: not configured"), - // }, - // { - // desc: "databricks token exchange server error", - // cfg: &Config{ - // ClientID: "client-id", - // Host: "http://host.com/test", - // TokenAudience: "token-audience", - // ActionsIDTokenRequestURL: "http://endpoint.com/test?version=1", - // ActionsIDTokenRequestToken: "token-1337", - // HTTPTransport: fixtures.MappingTransport{ - // "GET /oidc/.well-known/oauth-authorization-server": { - // Response: u2m.OAuthAuthorizationServer{ - // AuthorizationEndpoint: "https://host.com/auth", - // TokenEndpoint: "https://host.com/oidc/v1/token", - // }, - // }, - // "GET /test?version=1&audience=token-audience": { - // Status: http.StatusInternalServerError, - // ExpectedHeaders: map[string]string{ - // "Authorization": "Bearer token-1337", - // "Accept": "application/json", - // }, - // }, - // }, - // }, - // wantErrPrefix: errPrefix("federated-oidc-github"), - // }, - // { - // desc: "databricks workspace server error", - // cfg: &Config{ - // ClientID: "client-id", - // Host: "http://host.com/test", - // TokenAudience: "token-audience", - // ActionsIDTokenRequestURL: "http://endpoint.com/test?version=1", - // ActionsIDTokenRequestToken: "token-1337", - // HTTPTransport: fixtures.MappingTransport{ - // "GET /oidc/.well-known/oauth-authorization-server": { - // Response: u2m.OAuthAuthorizationServer{ - // AuthorizationEndpoint: "https://host.com/auth", - // TokenEndpoint: "https://host.com/oidc/v1/token", - // }, - // }, - // "GET /test?version=1&audience=token-audience": { - // Status: http.StatusOK, - // ExpectedHeaders: map[string]string{ - // "Authorization": "Bearer token-1337", - // "Accept": "application/json", - // }, - // Response: `{"value": "id-token-42"}`, - // }, - // "POST /oidc/v1/token": { - // Status: http.StatusInternalServerError, - // ExpectedHeaders: map[string]string{ - // "Content-Type": "application/x-www-form-urlencoded", - // }, - // }, - // }, - // }, - // wantErrPrefix: errPrefix("federated-oidc-github auth: not configured"), - // }, - // { - // desc: "invalid auth token", - // cfg: &Config{ - // ClientID: "client-id", - // Host: "http://host.com/test", - // TokenAudience: "token-audience", - // ActionsIDTokenRequestURL: "http://endpoint.com/test?version=1", - // ActionsIDTokenRequestToken: "token-1337", - // HTTPTransport: fixtures.MappingTransport{ - // "GET /oidc/.well-known/oauth-authorization-server": { - // Response: u2m.OAuthAuthorizationServer{ - // AuthorizationEndpoint: "https://host.com/auth", - // TokenEndpoint: "https://host.com/oidc/v1/token", - // }, - // }, - // "GET /test?version=1&audience=token-audience": { - // Status: http.StatusOK, - // ExpectedHeaders: map[string]string{ - // "Authorization": "Bearer token-1337", - // "Accept": "application/json", - // }, - // Response: `{"value": "id-token-42"}`, - // }, - // "POST /oidc/v1/token": { - // Status: http.StatusOK, - // ExpectedHeaders: map[string]string{ - // "Content-Type": "application/x-www-form-urlencoded", - // }, - // Response: map[string]string{ - // "foo": "bar", - // }, - // }, - // }, - // }, - // wantErrPrefix: errPrefix("federated-oidc-github auth: not configured"), - // }, - // { - // desc: "success workspace", - // cfg: &Config{ - // ClientID: "client-id", - // Host: "http://host.com/test", - // TokenAudience: "token-audience", - // ActionsIDTokenRequestURL: "http://endpoint.com/test?version=1", - // ActionsIDTokenRequestToken: "token-1337", - // HTTPTransport: fixtures.MappingTransport{ - // "GET /oidc/.well-known/oauth-authorization-server": { - // Response: u2m.OAuthAuthorizationServer{ - // AuthorizationEndpoint: "https://host.com/auth", - // TokenEndpoint: "https://host.com/oidc/v1/token", - // }, - // }, - // "GET /test?version=1&audience=token-audience": { - // Status: http.StatusOK, - // ExpectedHeaders: map[string]string{ - // "Authorization": "Bearer token-1337", - // "Accept": "application/json", - // }, - // Response: `{"value": "id-token-42"}`, - // }, - // "POST /oidc/v1/token": { + { + desc: "databricks workspace server error", + cfg: &Config{ + ClientID: "client-id", + Host: "http://host.com/test", + TokenAudience: "token-audience", + HTTPTransport: fixtures.MappingTransport{ + "GET /oidc/.well-known/oauth-authorization-server": { + Response: u2m.OAuthAuthorizationServer{ + AuthorizationEndpoint: "https://host.com/auth", + TokenEndpoint: "https://host.com/oidc/v1/token", + }, + }, + "POST /oidc/v1/token": { + Status: http.StatusInternalServerError, + ExpectedHeaders: map[string]string{ + "Content-Type": "application/x-www-form-urlencoded", + }, + }, + }, + }, + expectedAudience: "token-audience", + idToken: "id-token-42", + wantErrPrefix: errPrefix("oauth2: cannot fetch token: Internal Server Error"), + }, + { + desc: "invalid auth token", + cfg: &Config{ + ClientID: "client-id", + Host: "http://host.com/test", + TokenAudience: "token-audience", + HTTPTransport: fixtures.MappingTransport{ + "GET /oidc/.well-known/oauth-authorization-server": { + Response: u2m.OAuthAuthorizationServer{ + AuthorizationEndpoint: "https://host.com/auth", + TokenEndpoint: "https://host.com/oidc/v1/token", + }, + }, + "POST /oidc/v1/token": { + Status: http.StatusOK, + ExpectedHeaders: map[string]string{ + "Content-Type": "application/x-www-form-urlencoded", + }, + Response: map[string]string{ + "foo": "bar", + }, + }, + }, + }, + expectedAudience: "token-audience", + idToken: "id-token-42", + wantErrPrefix: errPrefix("oauth2: server response missing access_token"), + }, + { + desc: "success workspace", + cfg: &Config{ + ClientID: "client-id", + Host: "http://host.com/test", + TokenAudience: "token-audience", + HTTPTransport: fixtures.MappingTransport{ + "GET /oidc/.well-known/oauth-authorization-server": { + Response: u2m.OAuthAuthorizationServer{ + AuthorizationEndpoint: "https://host.com/auth", + TokenEndpoint: "https://host.com/oidc/v1/token", + }, + }, + "POST /oidc/v1/token": { - // Status: http.StatusOK, - // ExpectedHeaders: map[string]string{ - // "Content-Type": "application/x-www-form-urlencoded", - // }, - // Response: map[string]string{ - // "token_type": "access-token", - // "access_token": "test-auth-token", - // "refresh_token": "refresh", - // "expires_on": "0", - // }, - // }, - // }, - // }, - // wantHeaders: map[string]string{ - // "Authorization": "access-token test-auth-token", - // }, - // }, - // { - // desc: "success account", - // cfg: &Config{ - // ClientID: "client-id", - // AccountID: "ac123", - // Host: "https://accounts.databricks.com", - // TokenAudience: "token-audience", - // ActionsIDTokenRequestURL: "http://endpoint.com/test?version=1", - // ActionsIDTokenRequestToken: "token-1337", - // HTTPTransport: fixtures.MappingTransport{ - // "GET /test?version=1&audience=token-audience": { - // Status: http.StatusOK, - // ExpectedHeaders: map[string]string{ - // "Authorization": "Bearer token-1337", - // "Accept": "application/json", - // }, - // Response: `{"value": "id-token-42"}`, - // }, - // "POST /oidc/accounts/ac123/v1/token": { - // Status: http.StatusOK, - // ExpectedHeaders: map[string]string{ - // "Content-Type": "application/x-www-form-urlencoded", - // }, - // Response: map[string]string{ - // "token_type": "access-token", - // "access_token": "test-auth-token", - // "refresh_token": "refresh", - // "expires_on": "0", - // }, - // }, - // }, - // }, - // wantHeaders: map[string]string{ - // "Authorization": "access-token test-auth-token", - // }, - // }, - // { - // desc: "default token audience account", - // cfg: &Config{ - // ClientID: "client-id", - // AccountID: "ac123", - // Host: "https://accounts.databricks.com", - // ActionsIDTokenRequestURL: "http://endpoint.com/test?version=1", - // ActionsIDTokenRequestToken: "token-1337", - // HTTPTransport: fixtures.MappingTransport{ - // "GET /test?version=1&audience=ac123": { - // Status: http.StatusOK, - // ExpectedHeaders: map[string]string{ - // "Authorization": "Bearer token-1337", - // "Accept": "application/json", - // }, - // Response: `{"value": "id-token-42"}`, - // }, - // "POST /oidc/accounts/ac123/v1/token": { - // Status: http.StatusOK, - // ExpectedHeaders: map[string]string{ - // "Content-Type": "application/x-www-form-urlencoded", - // }, - // Response: map[string]string{ - // "token_type": "access-token", - // "access_token": "test-auth-token", - // "refresh_token": "refresh", - // "expires_on": "0", - // }, - // }, - // }, - // }, - // wantHeaders: map[string]string{ - // "Authorization": "access-token test-auth-token", - // }, - // }, - // { - // desc: "default token audience workspace", - // cfg: &Config{ - // ClientID: "client-id", - // Host: "https://host.com", - // ActionsIDTokenRequestURL: "http://endpoint.com/test?version=1", - // ActionsIDTokenRequestToken: "token-1337", - // HTTPTransport: fixtures.MappingTransport{ - // "GET /oidc/.well-known/oauth-authorization-server": { - // Response: u2m.OAuthAuthorizationServer{ - // AuthorizationEndpoint: "https://host.com/auth", - // TokenEndpoint: "https://host.com/oidc/v1/token", - // }, - // }, - // "GET /test?version=1&audience=https://host.com/oidc/v1/token": { - // Status: http.StatusOK, - // ExpectedHeaders: map[string]string{ - // "Authorization": "Bearer token-1337", - // "Accept": "application/json", - // }, - // Response: `{"value": "id-token-42"}`, - // }, - // "POST /oidc/v1/token": { - // Status: http.StatusOK, - // ExpectedHeaders: map[string]string{ - // "Content-Type": "application/x-www-form-urlencoded", - // }, - // Response: map[string]string{ - // "token_type": "access-token", - // "access_token": "test-auth-token", - // "refresh_token": "refresh", - // "expires_on": "0", - // }, - // }, - // }, - // }, - // wantHeaders: map[string]string{ - // "Authorization": "access-token test-auth-token", - // }, - //}, + Status: http.StatusOK, + ExpectedHeaders: map[string]string{ + "Content-Type": "application/x-www-form-urlencoded", + }, + ExpectedRequest: url.Values{ + "client_id": {"client-id"}, + "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", + }, + }, + }, + }, + expectedAudience: "token-audience", + idToken: "id-token-42", + wantToken: "test-auth-token", + }, + { + desc: "success account", + cfg: &Config{ + ClientID: "client-id", + AccountID: "ac123", + Host: "https://accounts.databricks.com", + TokenAudience: "token-audience", + HTTPTransport: fixtures.MappingTransport{ + "POST /oidc/accounts/ac123/v1/token": { + Status: http.StatusOK, + ExpectedHeaders: map[string]string{ + "Content-Type": "application/x-www-form-urlencoded", + }, + ExpectedRequest: url.Values{ + "client_id": {"client-id"}, + "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", + }, + }, + }, + }, + expectedAudience: "token-audience", + idToken: "id-token-42", + wantToken: "test-auth-token", + }, + { + desc: "default token audience account", + cfg: &Config{ + ClientID: "client-id", + AccountID: "ac123", + Host: "https://accounts.databricks.com", + HTTPTransport: fixtures.MappingTransport{ + "POST /oidc/accounts/ac123/v1/token": { + Status: http.StatusOK, + ExpectedHeaders: map[string]string{ + "Content-Type": "application/x-www-form-urlencoded", + }, + Response: map[string]string{ + "token_type": "access-token", + "access_token": "test-auth-token", + "refresh_token": "refresh", + "expires_on": "0", + }, + }, + }, + }, + expectedAudience: "ac123", + idToken: "id-token-42", + wantToken: "test-auth-token", + }, + { + desc: "default token audience workspace", + cfg: &Config{ + ClientID: "client-id", + Host: "https://host.com", + HTTPTransport: fixtures.MappingTransport{ + "GET /oidc/.well-known/oauth-authorization-server": { + Response: u2m.OAuthAuthorizationServer{ + AuthorizationEndpoint: "https://host.com/auth", + TokenEndpoint: "https://host.com/oidc/v1/token", + }, + }, + "POST /oidc/v1/token": { + Status: http.StatusOK, + ExpectedHeaders: map[string]string{ + "Content-Type": "application/x-www-form-urlencoded", + }, + Response: map[string]string{ + "token_type": "access-token", + "access_token": "test-auth-token", + "refresh_token": "refresh", + "expires_on": "0", + }, + }, + }, + }, + expectedAudience: "https://host.com/oidc/v1/token", + idToken: "id-token-42", + wantToken: "test-auth-token", + }, } for _, tc := range testCases { @@ -340,25 +274,23 @@ func TestDatabricksGithubWIFCredentials(t *testing.T) { // t.Skip() // } p := &mockIdTokenProvider{ - idToken: &IDToken{ - Value: tc.idToken, - }, - err: tc.tokenProviderError, - } - resolveErr := tc.cfg.EnsureResolved() - if resolveErr != nil { - t.Errorf("EnsureResolved(): got error %q, want none", resolveErr) + idToken: tc.idToken, + err: tc.tokenProviderError, } + tc.cfg.EnsureResolved() + c := tc.cfg.CanonicalHostName() ex := &wifTokenExchange{ clientID: tc.cfg.ClientID, account: tc.cfg.IsAccountClient(), accountID: tc.cfg.AccountID, - host: tc.cfg.Host, + host: c, tokenEndpointProvider: tc.cfg.getOidcEndpoints, audience: tc.cfg.TokenAudience, tokenProvider: p, } - ctx := context.WithValue(context.Background(), oauth2.HTTPClient, tc.cfg.HTTPTransport) + ctx := context.WithValue(context.Background(), oauth2.HTTPClient, &http.Client{ + Transport: tc.cfg.HTTPTransport, + }) token, gotErr := ex.Token(ctx) if tc.wantErrPrefix == nil && gotErr != nil { t.Errorf("Token(ctx): got error %q, want none", gotErr) @@ -366,6 +298,9 @@ func TestDatabricksGithubWIFCredentials(t *testing.T) { if tc.wantErrPrefix != nil && !hasPrefix(gotErr, *tc.wantErrPrefix) { t.Errorf("Token(ctx): got error %q, want error with prefix %q", gotErr, *tc.wantErrPrefix) } + if tc.expectedAudience != p.audience { + t.Errorf("mockTokenProvider: got audience %s, want %s", p.audience, tc.expectedAudience) + } tokenValue := "" if token != nil { tokenValue = token.AccessToken diff --git a/config/token_provider_github_oidc_test.go b/config/token_provider_github_oidc_test.go new file mode 100644 index 000000000..dba174349 --- /dev/null +++ b/config/token_provider_github_oidc_test.go @@ -0,0 +1,91 @@ +package config + +import ( + "context" + "net/http" + "testing" + + "github.com/databricks/databricks-sdk-go/httpclient" + "github.com/databricks/databricks-sdk-go/httpclient/fixtures" + "github.com/google/go-cmp/cmp" +) + +func TestGithubOIDCProvider(t *testing.T) { + testCases := []struct { + desc string + tokenRequestUrl string + tokenRequestToken string + audience string + httpTransport http.RoundTripper + wantToken *IDToken + wantErrPrefix *string + }{ + { + desc: "missing request token url", + tokenRequestToken: "token-1337", + wantErrPrefix: errPrefix("missing ActionsIDTokenRequestURL"), + }, + { + desc: "missing request token token", + tokenRequestUrl: "http://endpoint.com/test?version=1", + wantErrPrefix: errPrefix("missing ActionsIDTokenRequestToken"), + }, + { + desc: "error getting token", + tokenRequestToken: "token-1337", + tokenRequestUrl: "http://endpoint.com/test?version=1", + httpTransport: fixtures.MappingTransport{ + "GET /test?version=1": { + Status: http.StatusInternalServerError, + ExpectedHeaders: map[string]string{ + "Authorization": "Bearer token-1337", + "Accept": "application/json", + }, + }, + }, + wantErrPrefix: errPrefix("failed to request ID token from"), + }, + { + desc: "success", + tokenRequestToken: "token-1337", + tokenRequestUrl: "http://endpoint.com/test?version=1", + httpTransport: fixtures.MappingTransport{ + "GET /test?version=1": { + Status: http.StatusOK, + ExpectedHeaders: map[string]string{ + "Authorization": "Bearer token-1337", + "Accept": "application/json", + }, + Response: `{"value": "id-token-42"}`, + }, + }, + wantToken: &IDToken{ + Value: "id-token-42", + }, + }, + } + + for _, tc := range testCases { + t.Run(tc.desc, func(t *testing.T) { + cli := httpclient.NewApiClient(httpclient.ClientConfig{ + Transport: tc.httpTransport, + }) + p := &GithubProvider{ + actionsIDTokenRequestURL: tc.tokenRequestUrl, + actionsIDTokenRequestToken: tc.tokenRequestToken, + refreshClient: cli, + } + token, gotErr := p.IDToken(context.Background(), tc.audience) + + if tc.wantErrPrefix == nil && gotErr != nil { + t.Errorf("Authenticate(): got error %q, want none", gotErr) + } + if tc.wantErrPrefix != nil && !hasPrefix(gotErr, *tc.wantErrPrefix) { + t.Errorf("Authenticate(): got error %q, want error with prefix %q", gotErr, *tc.wantErrPrefix) + } + if diff := cmp.Diff(tc.wantToken, token); diff != "" { + t.Errorf("Authenticate(): mismatch (-want +got):\n%s", diff) + } + }) + } +} From b8d28dd47fbb8bb9615d512a607abe80dec46389 Mon Sep 17 00:00:00 2001 From: Hector Castejon Diaz Date: Mon, 7 Apr 2025 14:57:06 +0200 Subject: [PATCH 13/30] almost --- ...abricks_wif.go => auth_databricks_oidc.go} | 20 ++--- ...f_test.go => auth_databricks_oidc_test.go} | 9 +-- config/token_source_strategy_test.go | 79 +++++++++++++++++++ internal/auth_test.go | 4 +- 4 files changed, 94 insertions(+), 18 deletions(-) rename config/{auth_databricks_wif.go => auth_databricks_oidc.go} (81%) rename config/{auth_databricks_wif_test.go => auth_databricks_oidc_test.go} (98%) create mode 100644 config/token_source_strategy_test.go diff --git a/config/auth_databricks_wif.go b/config/auth_databricks_oidc.go similarity index 81% rename from config/auth_databricks_wif.go rename to config/auth_databricks_oidc.go index e370430fc..e574d90b9 100644 --- a/config/auth_databricks_wif.go +++ b/config/auth_databricks_oidc.go @@ -12,9 +12,9 @@ import ( ) // Constructs all Workload Identity Federation Credentials Strategies -func WifTokenCredentialStrategies(cfg *Config) []CredentialsStrategy { +func OidcTokenCredentialStrategies(cfg *Config) []CredentialsStrategy { providers := map[string]TokenProvider{ - "github-oidc-databricks-wif": &GithubProvider{ + "github-oidc": &GithubProvider{ actionsIDTokenRequestURL: cfg.ActionsIDTokenRequestURL, actionsIDTokenRequestToken: cfg.ActionsIDTokenRequestToken, refreshClient: cfg.refreshClient, @@ -23,17 +23,17 @@ func WifTokenCredentialStrategies(cfg *Config) []CredentialsStrategy { } strategies := []CredentialsStrategy{} for name, provider := range providers { - strategies = append(strategies, newWifTokenStrategy(name, cfg, provider)) + strategies = append(strategies, newOidcTokenStrategy(name, cfg, provider)) } return strategies } -func newWifTokenStrategy( +func newOidcTokenStrategy( name string, cfg *Config, tokenProvider TokenProvider, ) CredentialsStrategy { - wifTokenExchange := &wifTokenExchange{ + oidcTokenExchange := &oidcTokenExchange{ clientID: cfg.ClientID, account: cfg.IsAccountClient(), accountID: cfg.AccountID, @@ -42,12 +42,12 @@ func newWifTokenStrategy( audience: cfg.TokenAudience, tokenProvider: tokenProvider, } - return NewTokenSourceStrategy(name, wifTokenExchange) + return NewTokenSourceStrategy(name, oidcTokenExchange) } -// wifTokenExchange is a auth.TokenSource which exchanges a token using +// oidcTokenExchange is a auth.TokenSource which exchanges a token using // Workload Identity Federation. -type wifTokenExchange struct { +type oidcTokenExchange struct { clientID string account bool accountID string @@ -57,7 +57,7 @@ type wifTokenExchange struct { tokenProvider TokenProvider } -func (w *wifTokenExchange) Token(ctx context.Context) (*oauth2.Token, error) { +func (w *oidcTokenExchange) Token(ctx context.Context) (*oauth2.Token, error) { if w.clientID == "" { logger.Debugf(ctx, "Missing ClientID") return nil, errors.New("missing ClientID") @@ -92,7 +92,7 @@ func (w *wifTokenExchange) Token(ctx context.Context) (*oauth2.Token, error) { return c.Token(ctx) } -func (w *wifTokenExchange) determineAudience(ctx context.Context) (string, error) { +func (w *oidcTokenExchange) determineAudience(ctx context.Context) (string, error) { if w.audience != "" { return w.audience, nil } diff --git a/config/auth_databricks_wif_test.go b/config/auth_databricks_oidc_test.go similarity index 98% rename from config/auth_databricks_wif_test.go rename to config/auth_databricks_oidc_test.go index 52154ebcb..28329cdc2 100644 --- a/config/auth_databricks_wif_test.go +++ b/config/auth_databricks_oidc_test.go @@ -28,7 +28,7 @@ func (m *mockIdTokenProvider) IDToken(ctx context.Context, audience string) (*ID }, m.err } -func TestDatabricksGithubWIFCredentials(t *testing.T) { +func TestDatabricksOidcTokenSource(t *testing.T) { testCases := []struct { desc string cfg *Config @@ -270,16 +270,13 @@ func TestDatabricksGithubWIFCredentials(t *testing.T) { for _, tc := range testCases { t.Run(tc.desc, func(t *testing.T) { - // if tc.desc != "" { - // t.Skip() - // } p := &mockIdTokenProvider{ idToken: tc.idToken, err: tc.tokenProviderError, } tc.cfg.EnsureResolved() c := tc.cfg.CanonicalHostName() - ex := &wifTokenExchange{ + ex := &oidcTokenExchange{ clientID: tc.cfg.ClientID, account: tc.cfg.IsAccountClient(), accountID: tc.cfg.AccountID, @@ -313,7 +310,7 @@ func TestDatabricksGithubWIFCredentials(t *testing.T) { } func TestDatabricksWIFCredentials_Name(t *testing.T) { - strategies := WifTokenCredentialStrategies(&Config{}) + strategies := OidcTokenCredentialStrategies(&Config{}) expected := []string{"github-oidc-federated-oidc-github"} found := []string{} for _, strategy := range strategies { diff --git a/config/token_source_strategy_test.go b/config/token_source_strategy_test.go new file mode 100644 index 000000000..2618cdb5b --- /dev/null +++ b/config/token_source_strategy_test.go @@ -0,0 +1,79 @@ +package config + +import ( + "context" + "errors" + "net/http" + "testing" + + "github.com/google/go-cmp/cmp" + "golang.org/x/oauth2" +) + +type staticTokenSource struct { + token *oauth2.Token + err error +} + +func (s *staticTokenSource) Token(ctx context.Context) (*oauth2.Token, error) { + return s.token, s.err +} + +func TestDatabricksTokenSourceStrategy(t *testing.T) { + testCases := []struct { + desc string + token *oauth2.Token + tokenSourceError error + wantHeaders http.Header + }{ + { + desc: "token source error skips", + tokenSourceError: errors.New("random error"), + }, + { + desc: "token source error skips", + token: &oauth2.Token{ + AccessToken: "token-123", + }, + wantHeaders: http.Header{"Authorization": {"Bearer token-123"}}, + }, + } + + for _, tc := range testCases { + t.Run(tc.desc, func(t *testing.T) { + ts := &staticTokenSource{ + token: tc.token, + err: tc.tokenSourceError, + } + strat := &TokenSourceStrategy{ + name: "github-oidc", + tokenSource: ts, + } + provider, err := strat.Configure(context.Background(), &Config{}) + if tc.tokenSourceError == nil && provider == nil { + t.Error("Provider expected to not be nil, but it is") + } + if tc.tokenSourceError != nil && provider != nil { + t.Error("A failure in the TokenSource should cause the provider to be nil, but it's not") + } + if err != nil { + t.Errorf("Configure() got error %q, want none", err) + } + + if provider != nil { + req, _ := http.NewRequest("GET", "http://localhost", nil) + + gotErr := provider.SetHeaders(req) + + if gotErr != nil { + t.Errorf("SetHeaders(): got error %q, want none", gotErr) + } + if diff := cmp.Diff(tc.wantHeaders, req.Header); diff != "" { + t.Errorf("Authenticate(): mismatch (-want +got):\n%s", diff) + } + + } + + }) + } +} diff --git a/internal/auth_test.go b/internal/auth_test.go index ef6750e53..ef12d40a5 100644 --- a/internal/auth_test.go +++ b/internal/auth_test.go @@ -62,7 +62,7 @@ func TestUcAccWifAuth(t *testing.T) { Host: a.Config.Host, AccountID: a.Config.AccountID, ClientID: sp.ApplicationId, - AuthType: "github-oidc-databricks-wif", + AuthType: "github-oidc", TokenAudience: "https://github.com/databricks-eng", } @@ -141,7 +141,7 @@ func TestUcAccWifAuthWorkspace(t *testing.T) { wsCfg := &databricks.Config{ Host: workspaceUrl, ClientID: sp.ApplicationId, - AuthType: "github-oidc-databricks-wif", + AuthType: "github-oidc", TokenAudience: "https://github.com/databricks-eng", } From b2c2ce4aaefedf80ffb0a9821985c493ffd87017 Mon Sep 17 00:00:00 2001 From: Hector Castejon Diaz Date: Mon, 7 Apr 2025 14:57:43 +0200 Subject: [PATCH 14/30] last --- NEXT_CHANGELOG.md | 2 -- README.md | 4 ++-- 2 files changed, 2 insertions(+), 4 deletions(-) diff --git a/NEXT_CHANGELOG.md b/NEXT_CHANGELOG.md index db8a3de55..248628d3c 100644 --- a/NEXT_CHANGELOG.md +++ b/NEXT_CHANGELOG.md @@ -8,8 +8,6 @@ * [Breaking] Users running their worklows in GitHub Actions, which use Cloud native authentication and also have a `DATABRICKS_CLIENT_ID` and `DATABRICKS_HOST` environment variables set may see their authentication start failing due to the order in which the SDK tries different authentication methods. In such case, the `DATABRICKS_AUTH_TYPE` environment variable must be set to match the previously used authentication method. -* Support user-to-machine authentication in the SDK ([#1108](https://github.com/databricks/databricks-sdk-go/pull/1108)). -* Instances of `ApiClient` now share the same connection pool by default ([PR #1190](https://github.com/databricks/databricks-sdk-go/pull/1190)). ### Bug Fixes diff --git a/README.md b/README.md index 042e820f7..32af915ce 100644 --- a/README.md +++ b/README.md @@ -174,10 +174,10 @@ Depending on the Databricks authentication method, the SDK uses the following in ### Databricks native authentication -By default, the Databricks SDK for Go initially tries Databricks token authentication (`AuthType: "pat"` in `*databricks.Config`). If the SDK is unsuccessful, it then tries Workload Identity Federation (WIF) based authentication(`AuthType: "databricks-wif"` in `*databricks.Config`). Currently, only GitHub provided JWT Tokens is supported. +By default, the Databricks SDK for Go initially tries Databricks token authentication (`AuthType: "pat"` in `*databricks.Config`). If the SDK is unsuccessful, it then tries Workload Identity Federation (WIF) based authentication(`AuthType: "github-oidc"` in `*databricks.Config`). Currently, only GitHub provided JWT Tokens is supported. - For Databricks token authentication, you must provide `Host` and `Token`; or their environment variable or `.databrickscfg` file field equivalents. -- For Databricks wif authentication, you must provide `Host`, `ClientID` and `TokenAudience` _(optional)_; or their environment variable or `.databrickscfg` file field equivalents. +- For Databricks OIDC authentication, you must provide the `Host`, `ClientId` and `TokenAudience` _(optional)_ either directly, through the corresponding environment variables, or in your `.databrickscfg` configuration file. | `*databricks.Config` argument | Description | Environment variable / `.databrickscfg` file field | | ----------------------------- | ------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------ | -------------------------------------------------- | From 0bf2d2fcb1a9f7ac9b5dca60ece4fb039b03d9e7 Mon Sep 17 00:00:00 2001 From: Hector Castejon Diaz Date: Mon, 7 Apr 2025 14:59:31 +0200 Subject: [PATCH 15/30] fixes --- config/auth_databricks_oidc_test.go | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/config/auth_databricks_oidc_test.go b/config/auth_databricks_oidc_test.go index 28329cdc2..bea04056b 100644 --- a/config/auth_databricks_oidc_test.go +++ b/config/auth_databricks_oidc_test.go @@ -309,9 +309,9 @@ func TestDatabricksOidcTokenSource(t *testing.T) { } } -func TestDatabricksWIFCredentials_Name(t *testing.T) { +func TestDatabricksOidcCredentials_Name(t *testing.T) { strategies := OidcTokenCredentialStrategies(&Config{}) - expected := []string{"github-oidc-federated-oidc-github"} + expected := []string{"github-oidc"} found := []string{} for _, strategy := range strategies { found = append(found, strategy.Name()) From 3d39805e6f611853006be9ab321b86723893aeca Mon Sep 17 00:00:00 2001 From: Hector Castejon Diaz Date: Mon, 7 Apr 2025 15:09:39 +0200 Subject: [PATCH 16/30] Rename files --- config/.azure/az.json | 1 - config/.azure/az.sess | 1 - config/.azure/azureProfile.json | 1 - config/.azure/commandIndex.json | 1 - config/.azure/config | 3 --- config/.azure/versionCheck.json | 1 - config/testdata/azure/.azure/az.json | 1 - config/testdata/azure/.azure/az.sess | 1 - config/testdata/azure/.azure/azureProfile.json | 1 - config/testdata/azure/.azure/commandIndex.json | 1 - config/testdata/azure/.azure/config | 3 --- config/testdata/azure/.azure/versionCheck.json | 1 - 12 files changed, 16 deletions(-) delete mode 100644 config/.azure/az.json delete mode 100644 config/.azure/az.sess delete mode 100644 config/.azure/azureProfile.json delete mode 100644 config/.azure/commandIndex.json delete mode 100644 config/.azure/config delete mode 100644 config/.azure/versionCheck.json delete mode 100644 config/testdata/azure/.azure/az.json delete mode 100644 config/testdata/azure/.azure/az.sess delete mode 100644 config/testdata/azure/.azure/azureProfile.json delete mode 100644 config/testdata/azure/.azure/commandIndex.json delete mode 100644 config/testdata/azure/.azure/config delete mode 100644 config/testdata/azure/.azure/versionCheck.json diff --git a/config/.azure/az.json b/config/.azure/az.json deleted file mode 100644 index 22fdca1b2..000000000 --- a/config/.azure/az.json +++ /dev/null @@ -1 +0,0 @@ -{} \ No newline at end of file diff --git a/config/.azure/az.sess b/config/.azure/az.sess deleted file mode 100644 index 22fdca1b2..000000000 --- a/config/.azure/az.sess +++ /dev/null @@ -1 +0,0 @@ -{} \ No newline at end of file diff --git a/config/.azure/azureProfile.json b/config/.azure/azureProfile.json deleted file mode 100644 index c7306e3ef..000000000 --- a/config/.azure/azureProfile.json +++ /dev/null @@ -1 +0,0 @@ -{"installationId": "52778700-1129-11f0-95ee-c643c1692dc6"} \ No newline at end of file diff --git a/config/.azure/commandIndex.json b/config/.azure/commandIndex.json deleted file mode 100644 index 0870156ec..000000000 --- a/config/.azure/commandIndex.json +++ /dev/null @@ -1 +0,0 @@ -{"version": "2.69.0", "cloudProfile": "latest", "commandIndex": {"acr": ["azure.cli.command_modules.acr"], "aks": ["azure.cli.command_modules.acs", "azure.cli.command_modules.serviceconnector"], "advisor": ["azure.cli.command_modules.advisor"], "ams": ["azure.cli.command_modules.ams"], "apim": ["azure.cli.command_modules.apim"], "appconfig": ["azure.cli.command_modules.appconfig"], "webapp": ["azure.cli.command_modules.appservice", "azure.cli.command_modules.serviceconnector"], "functionapp": ["azure.cli.command_modules.appservice", "azure.cli.command_modules.serviceconnector"], "appservice": ["azure.cli.command_modules.appservice"], "staticwebapp": ["azure.cli.command_modules.appservice"], "logicapp": ["azure.cli.command_modules.appservice"], "aro": ["azure.cli.command_modules.aro"], "backup": ["azure.cli.command_modules.backup"], "batch": ["azure.cli.command_modules.batch"], "batchai": ["azure.cli.command_modules.batchai"], "billing": ["azure.cli.command_modules.billing"], "bot": ["azure.cli.command_modules.botservice"], "afd": ["azure.cli.command_modules.cdn"], "cdn": ["azure.cli.command_modules.cdn"], "cloud": ["azure.cli.command_modules.cloud"], "cognitiveservices": ["azure.cli.command_modules.cognitiveservices"], "compute-recommender": ["azure.cli.command_modules.compute_recommender"], "compute-fleet": ["azure.cli.command_modules.computefleet"], "config": ["azure.cli.command_modules.config"], "configure": ["azure.cli.command_modules.configure"], "cache": ["azure.cli.command_modules.configure"], "consumption": ["azure.cli.command_modules.consumption"], "container": ["azure.cli.command_modules.container"], "containerapp": ["azure.cli.command_modules.containerapp", "azure.cli.command_modules.serviceconnector"], "cosmosdb": ["azure.cli.command_modules.cosmosdb"], "managed-cassandra": ["azure.cli.command_modules.cosmosdb"], "databoxedge": ["azure.cli.command_modules.databoxedge"], "dls": ["azure.cli.command_modules.dls"], "dms": ["azure.cli.command_modules.dms"], "eventgrid": ["azure.cli.command_modules.eventgrid"], "eventhubs": ["azure.cli.command_modules.eventhubs"], "extension": ["azure.cli.command_modules.extension"], "feedback": ["azure.cli.command_modules.feedback"], "survey": ["azure.cli.command_modules.feedback"], "find": ["azure.cli.command_modules.find"], "hdinsight": ["azure.cli.command_modules.hdinsight"], "identity": ["azure.cli.command_modules.identity"], "interactive": ["azure.cli.command_modules.interactive"], "iot": ["azure.cli.command_modules.iot"], "keyvault": ["azure.cli.command_modules.keyvault"], "lab": ["azure.cli.command_modules.lab"], "managedservices": ["azure.cli.command_modules.managedservices"], "maps": ["azure.cli.command_modules.maps"], "term": ["azure.cli.command_modules.marketplaceordering"], "monitor": ["azure.cli.command_modules.monitor"], "mysql": ["azure.cli.command_modules.mysql", "azure.cli.command_modules.rdbms"], "netappfiles": ["azure.cli.command_modules.netappfiles"], "network": ["azure.cli.command_modules.network", "azure.cli.command_modules.privatedns"], "policy": ["azure.cli.command_modules.policyinsights", "azure.cli.command_modules.resource"], "login": ["azure.cli.command_modules.profile"], "logout": ["azure.cli.command_modules.profile"], "self-test": ["azure.cli.command_modules.profile"], "account": ["azure.cli.command_modules.profile", "azure.cli.command_modules.resource"], "mariadb": ["azure.cli.command_modules.rdbms"], "postgres": ["azure.cli.command_modules.rdbms"], "redis": ["azure.cli.command_modules.redis"], "relay": ["azure.cli.command_modules.relay"], "data-boundary": ["azure.cli.command_modules.resource"], "group": ["azure.cli.command_modules.resource"], "resource": ["azure.cli.command_modules.resource"], "provider": ["azure.cli.command_modules.resource"], "feature": ["azure.cli.command_modules.resource"], "tag": ["azure.cli.command_modules.resource"], "deployment": ["azure.cli.command_modules.resource"], "deployment-scripts": ["azure.cli.command_modules.resource"], "ts": ["azure.cli.command_modules.resource"], "stack": ["azure.cli.command_modules.resource"], "lock": ["azure.cli.command_modules.resource"], "managedapp": ["azure.cli.command_modules.resource"], "bicep": ["azure.cli.command_modules.resource"], "resourcemanagement": ["azure.cli.command_modules.resource"], "private-link": ["azure.cli.command_modules.resource"], "role": ["azure.cli.command_modules.role"], "ad": ["azure.cli.command_modules.role"], "search": ["azure.cli.command_modules.search"], "security": ["azure.cli.command_modules.security"], "servicebus": ["azure.cli.command_modules.servicebus"], "connection": ["azure.cli.command_modules.serviceconnector"], "sf": ["azure.cli.command_modules.servicefabric"], "signalr": ["azure.cli.command_modules.signalr"], "sql": ["azure.cli.command_modules.sql", "azure.cli.command_modules.sqlvm"], "storage": ["azure.cli.command_modules.storage"], "synapse": ["azure.cli.command_modules.synapse"], "rest": ["azure.cli.command_modules.util"], "version": ["azure.cli.command_modules.util"], "upgrade": ["azure.cli.command_modules.util"], "demo": ["azure.cli.command_modules.util"], "snapshot": ["azure.cli.command_modules.vm"], "disk-access": ["azure.cli.command_modules.vm"], "sig": ["azure.cli.command_modules.vm"], "vmss": ["azure.cli.command_modules.vm"], "restore-point": ["azure.cli.command_modules.vm"], "image": ["azure.cli.command_modules.vm"], "capacity": ["azure.cli.command_modules.vm"], "vm": ["azure.cli.command_modules.vm"], "disk": ["azure.cli.command_modules.vm"], "ppg": ["azure.cli.command_modules.vm"], "disk-encryption-set": ["azure.cli.command_modules.vm"], "sshkey": ["azure.cli.command_modules.vm"]}} \ No newline at end of file diff --git a/config/.azure/config b/config/.azure/config deleted file mode 100644 index 0ed7f34d6..000000000 --- a/config/.azure/config +++ /dev/null @@ -1,3 +0,0 @@ -[cloud] -name = AzureCloud - diff --git a/config/.azure/versionCheck.json b/config/.azure/versionCheck.json deleted file mode 100644 index 9be3cb1f5..000000000 --- a/config/.azure/versionCheck.json +++ /dev/null @@ -1 +0,0 @@ -{"versions": {"azure-cli": {"local": "2.69.0", "pypi": "2.71.0"}, "core": {"local": "2.69.0", "pypi": "2.71.0"}, "telemetry": {"local": "1.1.0", "pypi": "1.1.0"}}, "update_time": "2025-04-04 09:49:18.938802"} \ No newline at end of file diff --git a/config/testdata/azure/.azure/az.json b/config/testdata/azure/.azure/az.json deleted file mode 100644 index 22fdca1b2..000000000 --- a/config/testdata/azure/.azure/az.json +++ /dev/null @@ -1 +0,0 @@ -{} \ No newline at end of file diff --git a/config/testdata/azure/.azure/az.sess b/config/testdata/azure/.azure/az.sess deleted file mode 100644 index 22fdca1b2..000000000 --- a/config/testdata/azure/.azure/az.sess +++ /dev/null @@ -1 +0,0 @@ -{} \ No newline at end of file diff --git a/config/testdata/azure/.azure/azureProfile.json b/config/testdata/azure/.azure/azureProfile.json deleted file mode 100644 index 6dd605248..000000000 --- a/config/testdata/azure/.azure/azureProfile.json +++ /dev/null @@ -1 +0,0 @@ -{"installationId": "5277cba2-1129-11f0-86b1-c643c1692dc6"} \ No newline at end of file diff --git a/config/testdata/azure/.azure/commandIndex.json b/config/testdata/azure/.azure/commandIndex.json deleted file mode 100644 index 0870156ec..000000000 --- a/config/testdata/azure/.azure/commandIndex.json +++ /dev/null @@ -1 +0,0 @@ -{"version": "2.69.0", "cloudProfile": "latest", "commandIndex": {"acr": ["azure.cli.command_modules.acr"], "aks": ["azure.cli.command_modules.acs", "azure.cli.command_modules.serviceconnector"], "advisor": ["azure.cli.command_modules.advisor"], "ams": ["azure.cli.command_modules.ams"], "apim": ["azure.cli.command_modules.apim"], "appconfig": ["azure.cli.command_modules.appconfig"], "webapp": ["azure.cli.command_modules.appservice", "azure.cli.command_modules.serviceconnector"], "functionapp": ["azure.cli.command_modules.appservice", "azure.cli.command_modules.serviceconnector"], "appservice": ["azure.cli.command_modules.appservice"], "staticwebapp": ["azure.cli.command_modules.appservice"], "logicapp": ["azure.cli.command_modules.appservice"], "aro": ["azure.cli.command_modules.aro"], "backup": ["azure.cli.command_modules.backup"], "batch": ["azure.cli.command_modules.batch"], "batchai": ["azure.cli.command_modules.batchai"], "billing": ["azure.cli.command_modules.billing"], "bot": ["azure.cli.command_modules.botservice"], "afd": ["azure.cli.command_modules.cdn"], "cdn": ["azure.cli.command_modules.cdn"], "cloud": ["azure.cli.command_modules.cloud"], "cognitiveservices": ["azure.cli.command_modules.cognitiveservices"], "compute-recommender": ["azure.cli.command_modules.compute_recommender"], "compute-fleet": ["azure.cli.command_modules.computefleet"], "config": ["azure.cli.command_modules.config"], "configure": ["azure.cli.command_modules.configure"], "cache": ["azure.cli.command_modules.configure"], "consumption": ["azure.cli.command_modules.consumption"], "container": ["azure.cli.command_modules.container"], "containerapp": ["azure.cli.command_modules.containerapp", "azure.cli.command_modules.serviceconnector"], "cosmosdb": ["azure.cli.command_modules.cosmosdb"], "managed-cassandra": ["azure.cli.command_modules.cosmosdb"], "databoxedge": ["azure.cli.command_modules.databoxedge"], "dls": ["azure.cli.command_modules.dls"], "dms": ["azure.cli.command_modules.dms"], "eventgrid": ["azure.cli.command_modules.eventgrid"], "eventhubs": ["azure.cli.command_modules.eventhubs"], "extension": ["azure.cli.command_modules.extension"], "feedback": ["azure.cli.command_modules.feedback"], "survey": ["azure.cli.command_modules.feedback"], "find": ["azure.cli.command_modules.find"], "hdinsight": ["azure.cli.command_modules.hdinsight"], "identity": ["azure.cli.command_modules.identity"], "interactive": ["azure.cli.command_modules.interactive"], "iot": ["azure.cli.command_modules.iot"], "keyvault": ["azure.cli.command_modules.keyvault"], "lab": ["azure.cli.command_modules.lab"], "managedservices": ["azure.cli.command_modules.managedservices"], "maps": ["azure.cli.command_modules.maps"], "term": ["azure.cli.command_modules.marketplaceordering"], "monitor": ["azure.cli.command_modules.monitor"], "mysql": ["azure.cli.command_modules.mysql", "azure.cli.command_modules.rdbms"], "netappfiles": ["azure.cli.command_modules.netappfiles"], "network": ["azure.cli.command_modules.network", "azure.cli.command_modules.privatedns"], "policy": ["azure.cli.command_modules.policyinsights", "azure.cli.command_modules.resource"], "login": ["azure.cli.command_modules.profile"], "logout": ["azure.cli.command_modules.profile"], "self-test": ["azure.cli.command_modules.profile"], "account": ["azure.cli.command_modules.profile", "azure.cli.command_modules.resource"], "mariadb": ["azure.cli.command_modules.rdbms"], "postgres": ["azure.cli.command_modules.rdbms"], "redis": ["azure.cli.command_modules.redis"], "relay": ["azure.cli.command_modules.relay"], "data-boundary": ["azure.cli.command_modules.resource"], "group": ["azure.cli.command_modules.resource"], "resource": ["azure.cli.command_modules.resource"], "provider": ["azure.cli.command_modules.resource"], "feature": ["azure.cli.command_modules.resource"], "tag": ["azure.cli.command_modules.resource"], "deployment": ["azure.cli.command_modules.resource"], "deployment-scripts": ["azure.cli.command_modules.resource"], "ts": ["azure.cli.command_modules.resource"], "stack": ["azure.cli.command_modules.resource"], "lock": ["azure.cli.command_modules.resource"], "managedapp": ["azure.cli.command_modules.resource"], "bicep": ["azure.cli.command_modules.resource"], "resourcemanagement": ["azure.cli.command_modules.resource"], "private-link": ["azure.cli.command_modules.resource"], "role": ["azure.cli.command_modules.role"], "ad": ["azure.cli.command_modules.role"], "search": ["azure.cli.command_modules.search"], "security": ["azure.cli.command_modules.security"], "servicebus": ["azure.cli.command_modules.servicebus"], "connection": ["azure.cli.command_modules.serviceconnector"], "sf": ["azure.cli.command_modules.servicefabric"], "signalr": ["azure.cli.command_modules.signalr"], "sql": ["azure.cli.command_modules.sql", "azure.cli.command_modules.sqlvm"], "storage": ["azure.cli.command_modules.storage"], "synapse": ["azure.cli.command_modules.synapse"], "rest": ["azure.cli.command_modules.util"], "version": ["azure.cli.command_modules.util"], "upgrade": ["azure.cli.command_modules.util"], "demo": ["azure.cli.command_modules.util"], "snapshot": ["azure.cli.command_modules.vm"], "disk-access": ["azure.cli.command_modules.vm"], "sig": ["azure.cli.command_modules.vm"], "vmss": ["azure.cli.command_modules.vm"], "restore-point": ["azure.cli.command_modules.vm"], "image": ["azure.cli.command_modules.vm"], "capacity": ["azure.cli.command_modules.vm"], "vm": ["azure.cli.command_modules.vm"], "disk": ["azure.cli.command_modules.vm"], "ppg": ["azure.cli.command_modules.vm"], "disk-encryption-set": ["azure.cli.command_modules.vm"], "sshkey": ["azure.cli.command_modules.vm"]}} \ No newline at end of file diff --git a/config/testdata/azure/.azure/config b/config/testdata/azure/.azure/config deleted file mode 100644 index 0ed7f34d6..000000000 --- a/config/testdata/azure/.azure/config +++ /dev/null @@ -1,3 +0,0 @@ -[cloud] -name = AzureCloud - diff --git a/config/testdata/azure/.azure/versionCheck.json b/config/testdata/azure/.azure/versionCheck.json deleted file mode 100644 index 9ad051c80..000000000 --- a/config/testdata/azure/.azure/versionCheck.json +++ /dev/null @@ -1 +0,0 @@ -{"versions": {"azure-cli": {"local": "2.69.0", "pypi": "2.71.0"}, "core": {"local": "2.69.0", "pypi": "2.71.0"}, "telemetry": {"local": "1.1.0", "pypi": "1.1.0"}}, "update_time": "2025-04-04 09:49:19.096374"} \ No newline at end of file From 35df156d1d73ee22d6d078f4dfbaf528b727bc37 Mon Sep 17 00:00:00 2001 From: Hector Castejon Diaz Date: Mon, 7 Apr 2025 15:12:36 +0200 Subject: [PATCH 17/30] Enabe --- config/auth_default.go | 24 +++++++++++------------- 1 file changed, 11 insertions(+), 13 deletions(-) diff --git a/config/auth_default.go b/config/auth_default.go index 3571b8ba4..3709fcc0b 100644 --- a/config/auth_default.go +++ b/config/auth_default.go @@ -10,16 +10,15 @@ import ( ) func buildDefaultStrategies(cfg *Config) []CredentialsStrategy { - return append( - []CredentialsStrategy{ - PatCredentials{}, - BasicCredentials{}, - M2mCredentials{}, - DatabricksCliCredentials, - MetadataServiceCredentials{}, - }, - // append( - // WifTokenCredentialStrategies(cfg), + strategies := []CredentialsStrategy{} + strategies = append(strategies, + PatCredentials{}, + BasicCredentials{}, + M2mCredentials{}, + DatabricksCliCredentials, + MetadataServiceCredentials{}) + strategies = append(strategies, OidcTokenCredentialStrategies(cfg)...) + strategies = append(strategies, // Attempt to configure auth from most specific to most generic (the Azure CLI). AzureGithubOIDCCredentials{}, AzureMsiCredentials{}, @@ -27,9 +26,8 @@ func buildDefaultStrategies(cfg *Config) []CredentialsStrategy { AzureCliCredentials{}, // Attempt to configure auth from most specific to most generic (Google Application Default Credentials). GoogleCredentials{}, - GoogleDefaultCredentials{}, - //)..., - ) + GoogleDefaultCredentials{}) + return strategies } type DefaultCredentials struct { From c5652f3c732a0db1a65d111a19aafabf7b0178ae Mon Sep 17 00:00:00 2001 From: Hector Castejon Diaz Date: Tue, 8 Apr 2025 09:05:10 +0200 Subject: [PATCH 18/30] Remove config from tests --- config/auth_databricks_oidc.go | 11 +- config/auth_databricks_oidc_test.go | 344 ++++++++++++++-------------- 2 files changed, 175 insertions(+), 180 deletions(-) diff --git a/config/auth_databricks_oidc.go b/config/auth_databricks_oidc.go index e574d90b9..0d5ae769f 100644 --- a/config/auth_databricks_oidc.go +++ b/config/auth_databricks_oidc.go @@ -33,10 +33,14 @@ func newOidcTokenStrategy( cfg *Config, tokenProvider TokenProvider, ) CredentialsStrategy { + accountId := "" + // NOTE: cfg.AccountID can be present even if the client is not an account client. + if cfg.IsAccountClient() { + accountId = cfg.AccountID + } oidcTokenExchange := &oidcTokenExchange{ clientID: cfg.ClientID, - account: cfg.IsAccountClient(), - accountID: cfg.AccountID, + accountID: accountId, host: cfg.Host, tokenEndpointProvider: cfg.getOidcEndpoints, audience: cfg.TokenAudience, @@ -49,7 +53,6 @@ func newOidcTokenStrategy( // Workload Identity Federation. type oidcTokenExchange struct { clientID string - account bool accountID string host string tokenEndpointProvider func(ctx context.Context) (*u2m.OAuthAuthorizationServer, error) @@ -97,7 +100,7 @@ func (w *oidcTokenExchange) determineAudience(ctx context.Context) (string, erro return w.audience, nil } // For Databricks Accounts, the account id is the default audience. - if w.account { + if w.accountID != "" { return w.accountID, nil } endpoints, err := w.tokenEndpointProvider(ctx) diff --git a/config/auth_databricks_oidc_test.go b/config/auth_databricks_oidc_test.go index bea04056b..47ffc5e12 100644 --- a/config/auth_databricks_oidc_test.go +++ b/config/auth_databricks_oidc_test.go @@ -30,80 +30,73 @@ func (m *mockIdTokenProvider) IDToken(ctx context.Context, audience string) (*ID func TestDatabricksOidcTokenSource(t *testing.T) { testCases := []struct { - desc string - cfg *Config - idToken string - expectedAudience string - tokenProviderError error - wantToken string - wantErrPrefix *string + desc string + clientID string + accountID string + host string + tokenAudience string + httpTransport http.RoundTripper + oidcEndpointSupplier func(ctx context.Context) (*u2m.OAuthAuthorizationServer, error) + idToken string + expectedAudience string + tokenProviderError error + wantToken string + wantErrPrefix *string }{ { - desc: "missing host", - cfg: &Config{ - ClientID: "client-id", - TokenAudience: "token-audience", - }, + desc: "missing host", + clientID: "client-id", + tokenAudience: "token-audience", wantErrPrefix: errPrefix("missing Host"), }, { - desc: "missing client ID", - cfg: &Config{ - Host: "http://host.com/test", - TokenAudience: "token-audience", - }, + desc: "missing client ID", + host: "http://host.com", + tokenAudience: "token-audience", wantErrPrefix: errPrefix("missing ClientID"), }, { - desc: "auth server error", - cfg: &Config{ - ClientID: "client-id", - Host: "http://host.com/test", - HTTPTransport: fixtures.MappingTransport{ - "GET /oidc/.well-known/oauth-authorization-server": { - Status: http.StatusNotFound, - }, - }, + desc: "auth server error", + clientID: "client-id", + host: "http://host.com", + oidcEndpointSupplier: func(ctx context.Context) (*u2m.OAuthAuthorizationServer, error) { + return nil, errors.New("databricks OAuth is not supported for this host") }, wantErrPrefix: errPrefix("databricks OAuth is not supported for this host"), }, { desc: "token provider error", - cfg: &Config{ - ClientID: "client-id", - Host: "http://host.com/test", - TokenAudience: "token-audience", - HTTPTransport: fixtures.MappingTransport{ - "GET /oidc/.well-known/oauth-authorization-server": { - Response: u2m.OAuthAuthorizationServer{ - AuthorizationEndpoint: "https://host.com/auth", - TokenEndpoint: "https://host.com/oidc/v1/token", - }, - }, - }, + + clientID: "client-id", + host: "http://host.com", + tokenAudience: "token-audience", + oidcEndpointSupplier: func(ctx context.Context) (*u2m.OAuthAuthorizationServer, error) { + return &u2m.OAuthAuthorizationServer{ + AuthorizationEndpoint: "https://host.com/auth", + TokenEndpoint: "https://host.com/oidc/v1/token", + }, nil }, + expectedAudience: "token-audience", tokenProviderError: errors.New("error getting id token"), wantErrPrefix: errPrefix("error getting id token"), }, { - desc: "databricks workspace server error", - cfg: &Config{ - ClientID: "client-id", - Host: "http://host.com/test", - TokenAudience: "token-audience", - HTTPTransport: fixtures.MappingTransport{ - "GET /oidc/.well-known/oauth-authorization-server": { - Response: u2m.OAuthAuthorizationServer{ - AuthorizationEndpoint: "https://host.com/auth", - TokenEndpoint: "https://host.com/oidc/v1/token", - }, - }, - "POST /oidc/v1/token": { - Status: http.StatusInternalServerError, - ExpectedHeaders: map[string]string{ - "Content-Type": "application/x-www-form-urlencoded", - }, + desc: "databricks workspace server error", + clientID: "client-id", + host: "http://host.com", + tokenAudience: "token-audience", + oidcEndpointSupplier: func(ctx context.Context) (*u2m.OAuthAuthorizationServer, error) { + return &u2m.OAuthAuthorizationServer{ + AuthorizationEndpoint: "https://host.com/auth", + TokenEndpoint: "https://host.com/oidc/v1/token", + }, nil + }, + httpTransport: fixtures.MappingTransport{ + "POST /oidc/v1/token": { + Status: http.StatusInternalServerError, + ExpectedHeaders: map[string]string{ + "Content-Type": "application/x-www-form-urlencoded", }, }, }, @@ -112,26 +105,24 @@ func TestDatabricksOidcTokenSource(t *testing.T) { wantErrPrefix: errPrefix("oauth2: cannot fetch token: Internal Server Error"), }, { - desc: "invalid auth token", - cfg: &Config{ - ClientID: "client-id", - Host: "http://host.com/test", - TokenAudience: "token-audience", - HTTPTransport: fixtures.MappingTransport{ - "GET /oidc/.well-known/oauth-authorization-server": { - Response: u2m.OAuthAuthorizationServer{ - AuthorizationEndpoint: "https://host.com/auth", - TokenEndpoint: "https://host.com/oidc/v1/token", - }, + desc: "invalid auth token", + clientID: "client-id", + host: "http://host.com", + tokenAudience: "token-audience", + oidcEndpointSupplier: func(ctx context.Context) (*u2m.OAuthAuthorizationServer, error) { + return &u2m.OAuthAuthorizationServer{ + AuthorizationEndpoint: "https://host.com/auth", + 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", }, - "POST /oidc/v1/token": { - Status: http.StatusOK, - ExpectedHeaders: map[string]string{ - "Content-Type": "application/x-www-form-urlencoded", - }, - Response: map[string]string{ - "foo": "bar", - }, + Response: map[string]string{ + "foo": "bar", }, }, }, @@ -140,37 +131,35 @@ func TestDatabricksOidcTokenSource(t *testing.T) { wantErrPrefix: errPrefix("oauth2: server response missing access_token"), }, { - desc: "success workspace", - cfg: &Config{ - ClientID: "client-id", - Host: "http://host.com/test", - TokenAudience: "token-audience", - HTTPTransport: fixtures.MappingTransport{ - "GET /oidc/.well-known/oauth-authorization-server": { - Response: u2m.OAuthAuthorizationServer{ - AuthorizationEndpoint: "https://host.com/auth", - TokenEndpoint: "https://host.com/oidc/v1/token", - }, - }, - "POST /oidc/v1/token": { + desc: "success workspace", + clientID: "client-id", + host: "http://host.com", + tokenAudience: "token-audience", + oidcEndpointSupplier: func(ctx context.Context) (*u2m.OAuthAuthorizationServer, error) { + return &u2m.OAuthAuthorizationServer{ + AuthorizationEndpoint: "https://host.com/auth", + 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{ - "client_id": {"client-id"}, - "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", - }, + Status: http.StatusOK, + ExpectedHeaders: map[string]string{ + "Content-Type": "application/x-www-form-urlencoded", + }, + ExpectedRequest: url.Values{ + "client_id": {"client-id"}, + "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", }, }, }, @@ -179,31 +168,35 @@ func TestDatabricksOidcTokenSource(t *testing.T) { wantToken: "test-auth-token", }, { - desc: "success account", - cfg: &Config{ - ClientID: "client-id", - AccountID: "ac123", - Host: "https://accounts.databricks.com", - TokenAudience: "token-audience", - HTTPTransport: fixtures.MappingTransport{ - "POST /oidc/accounts/ac123/v1/token": { - Status: http.StatusOK, - ExpectedHeaders: map[string]string{ - "Content-Type": "application/x-www-form-urlencoded", - }, - ExpectedRequest: url.Values{ - "client_id": {"client-id"}, - "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", - }, + desc: "success account", + clientID: "client-id", + accountID: "ac123", + host: "https://accounts.databricks.com", + tokenAudience: "token-audience", + oidcEndpointSupplier: func(ctx context.Context) (*u2m.OAuthAuthorizationServer, error) { + return &u2m.OAuthAuthorizationServer{ + AuthorizationEndpoint: "https://host.com/auth", + TokenEndpoint: "https://host.com/oidc/accounts/ac123/v1/token", + }, nil + }, + httpTransport: fixtures.MappingTransport{ + "POST /oidc/accounts/ac123/v1/token": { + Status: http.StatusOK, + ExpectedHeaders: map[string]string{ + "Content-Type": "application/x-www-form-urlencoded", + }, + ExpectedRequest: url.Values{ + "client_id": {"client-id"}, + "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", }, }, }, @@ -212,23 +205,27 @@ func TestDatabricksOidcTokenSource(t *testing.T) { wantToken: "test-auth-token", }, { - desc: "default token audience account", - cfg: &Config{ - ClientID: "client-id", - AccountID: "ac123", - Host: "https://accounts.databricks.com", - HTTPTransport: fixtures.MappingTransport{ - "POST /oidc/accounts/ac123/v1/token": { - Status: http.StatusOK, - ExpectedHeaders: map[string]string{ - "Content-Type": "application/x-www-form-urlencoded", - }, - Response: map[string]string{ - "token_type": "access-token", - "access_token": "test-auth-token", - "refresh_token": "refresh", - "expires_on": "0", - }, + desc: "default token audience account", + clientID: "client-id", + accountID: "ac123", + host: "https://accounts.databricks.com", + oidcEndpointSupplier: func(ctx context.Context) (*u2m.OAuthAuthorizationServer, error) { + return &u2m.OAuthAuthorizationServer{ + AuthorizationEndpoint: "https://host.com/auth", + TokenEndpoint: "https://host.com/oidc/accounts/ac123/v1/token", + }, nil + }, + httpTransport: fixtures.MappingTransport{ + "POST /oidc/accounts/ac123/v1/token": { + Status: http.StatusOK, + ExpectedHeaders: map[string]string{ + "Content-Type": "application/x-www-form-urlencoded", + }, + Response: map[string]string{ + "token_type": "access-token", + "access_token": "test-auth-token", + "refresh_token": "refresh", + "expires_on": "0", }, }, }, @@ -237,28 +234,26 @@ func TestDatabricksOidcTokenSource(t *testing.T) { wantToken: "test-auth-token", }, { - desc: "default token audience workspace", - cfg: &Config{ - ClientID: "client-id", - Host: "https://host.com", - HTTPTransport: fixtures.MappingTransport{ - "GET /oidc/.well-known/oauth-authorization-server": { - Response: u2m.OAuthAuthorizationServer{ - AuthorizationEndpoint: "https://host.com/auth", - TokenEndpoint: "https://host.com/oidc/v1/token", - }, + desc: "default token audience workspace", + clientID: "client-id", + host: "https://host.com", + oidcEndpointSupplier: func(ctx context.Context) (*u2m.OAuthAuthorizationServer, error) { + return &u2m.OAuthAuthorizationServer{ + AuthorizationEndpoint: "https://host.com/auth", + 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", }, - "POST /oidc/v1/token": { - Status: http.StatusOK, - ExpectedHeaders: map[string]string{ - "Content-Type": "application/x-www-form-urlencoded", - }, - Response: map[string]string{ - "token_type": "access-token", - "access_token": "test-auth-token", - "refresh_token": "refresh", - "expires_on": "0", - }, + Response: map[string]string{ + "token_type": "access-token", + "access_token": "test-auth-token", + "refresh_token": "refresh", + "expires_on": "0", }, }, }, @@ -274,19 +269,16 @@ func TestDatabricksOidcTokenSource(t *testing.T) { idToken: tc.idToken, err: tc.tokenProviderError, } - tc.cfg.EnsureResolved() - c := tc.cfg.CanonicalHostName() ex := &oidcTokenExchange{ - clientID: tc.cfg.ClientID, - account: tc.cfg.IsAccountClient(), - accountID: tc.cfg.AccountID, - host: c, - tokenEndpointProvider: tc.cfg.getOidcEndpoints, - audience: tc.cfg.TokenAudience, + clientID: tc.clientID, + accountID: tc.accountID, + host: tc.host, + tokenEndpointProvider: tc.oidcEndpointSupplier, + audience: tc.tokenAudience, tokenProvider: p, } ctx := context.WithValue(context.Background(), oauth2.HTTPClient, &http.Client{ - Transport: tc.cfg.HTTPTransport, + Transport: tc.httpTransport, }) token, gotErr := ex.Token(ctx) if tc.wantErrPrefix == nil && gotErr != nil { From 635752b380aceb2fdba2e4c7af38c7c146dded89 Mon Sep 17 00:00:00 2001 From: Hector Castejon Diaz Date: Tue, 8 Apr 2025 09:09:58 +0200 Subject: [PATCH 19/30] comments --- config/auth_databricks_oidc.go | 1 + config/token_provider_github_oidc.go | 5 +++++ config/token_source_strategy.go | 4 ++++ 3 files changed, 10 insertions(+) diff --git a/config/auth_databricks_oidc.go b/config/auth_databricks_oidc.go index 0d5ae769f..b012964a7 100644 --- a/config/auth_databricks_oidc.go +++ b/config/auth_databricks_oidc.go @@ -60,6 +60,7 @@ type oidcTokenExchange struct { tokenProvider TokenProvider } +// Token implements [TokenSource.Token] func (w *oidcTokenExchange) Token(ctx context.Context) (*oauth2.Token, error) { if w.clientID == "" { logger.Debugf(ctx, "Missing ClientID") diff --git a/config/token_provider_github_oidc.go b/config/token_provider_github_oidc.go index 9cec3a39b..95b89d977 100644 --- a/config/token_provider_github_oidc.go +++ b/config/token_provider_github_oidc.go @@ -9,12 +9,17 @@ import ( "github.com/databricks/databricks-sdk-go/logger" ) +/** + * GithubProvider retrieves JWT Tokens from Github Actions. + */ type GithubProvider struct { actionsIDTokenRequestURL string actionsIDTokenRequestToken string refreshClient *httpclient.ApiClient } +// IDToken returns a JWT Token for the specified audience. It will return +// an error if not running in GitHub Actions. func (g *GithubProvider) IDToken(ctx context.Context, audience string) (*IDToken, error) { if g.actionsIDTokenRequestURL == "" { logger.Debugf(ctx, "Missing ActionsIDTokenRequestURL, likely not calling from a Github action") diff --git a/config/token_source_strategy.go b/config/token_source_strategy.go index 099e2ae5e..888076cc5 100644 --- a/config/token_source_strategy.go +++ b/config/token_source_strategy.go @@ -14,16 +14,19 @@ type IDToken struct { Value string } +// TokenProvider is an interface for a Token Provider with an audience. type TokenProvider interface { // Function to get the token IDToken(ctx context.Context, audience string) (*IDToken, error) } +// TokenSourceStrategy is wrapper on a auth.TokenSource which converts it into a CredentialsStrategy type TokenSourceStrategy struct { tokenSource auth.TokenSource name string } +// Configure implements [CredentialsStrategy.Configure]. func (t *TokenSourceStrategy) Configure(ctx context.Context, cfg *Config) (credentials.CredentialsProvider, error) { // If we cannot get a token, skip this CredentialsStrategy. @@ -40,6 +43,7 @@ func (t *TokenSourceStrategy) Configure(ctx context.Context, cfg *Config) (crede return credentials.NewOAuthCredentialsProvider(visitor, authconv.OAuth2TokenSource(cached).Token), nil } +// Name implements [CredentialsStrategy.Name]. func (t *TokenSourceStrategy) Name() string { return t.name } From 78db7658746b70270f1da64a0a00e21d53732676 Mon Sep 17 00:00:00 2001 From: Hector Castejon Diaz Date: Tue, 8 Apr 2025 13:54:42 +0200 Subject: [PATCH 20/30] PR comments --- client/client_test.go | 6 ++ config/auth_databricks_oidc.go | 88 +++++++----------- config/auth_databricks_oidc_test.go | 132 ++++++++------------------- config/auth_default.go | 43 ++++++++- config/token_source_strategy.go | 34 +++---- config/token_source_strategy_test.go | 2 +- 6 files changed, 132 insertions(+), 173 deletions(-) diff --git a/client/client_test.go b/client/client_test.go index 267635946..0263a0e8f 100644 --- a/client/client_test.go +++ b/client/client_test.go @@ -33,6 +33,10 @@ func TestSimpleRequestFailsURLError(t *testing.T) { ConfigFile: "/dev/null", RetryTimeoutSeconds: 1, HTTPTransport: hc(func(r *http.Request) (*http.Response, error) { + // Unrelated to this test. This URL is called during Client setup. + if r.URL.Path == "/oidc/.well-known/oauth-authorization-server" { + return nil, nil + } require.Equal(t, "GET", r.Method) require.Equal(t, "/a/b", r.URL.Path) require.Equal(t, "c=d", r.URL.RawQuery) @@ -312,6 +316,8 @@ func TestHttpTransport(t *testing.T) { func TestDoRemovesDoubleSlashesFromFilesAPI(t *testing.T) { i := 0 expectedPaths := []string{ + // Unrelated to this test. This URL is called during Client setup. + "/oidc/.well-known/oauth-authorization-server", "/api/2.0/fs/files/Volumes/abc/def/ghi", "/api/2.0/anotherservice//test", } diff --git a/config/auth_databricks_oidc.go b/config/auth_databricks_oidc.go index b012964a7..85188658a 100644 --- a/config/auth_databricks_oidc.go +++ b/config/auth_databricks_oidc.go @@ -5,59 +5,44 @@ import ( "errors" "net/url" - "github.com/databricks/databricks-sdk-go/credentials/u2m" + "github.com/databricks/databricks-sdk-go/config/experimental/auth" "github.com/databricks/databricks-sdk-go/logger" "golang.org/x/oauth2" "golang.org/x/oauth2/clientcredentials" ) -// Constructs all Workload Identity Federation Credentials Strategies -func OidcTokenCredentialStrategies(cfg *Config) []CredentialsStrategy { - providers := map[string]TokenProvider{ - "github-oidc": &GithubProvider{ - actionsIDTokenRequestURL: cfg.ActionsIDTokenRequestURL, - actionsIDTokenRequestToken: cfg.ActionsIDTokenRequestToken, - refreshClient: cfg.refreshClient, - }, - // Add new providers at the end of the list - } - strategies := []CredentialsStrategy{} - for name, provider := range providers { - strategies = append(strategies, newOidcTokenStrategy(name, cfg, provider)) +func NewOIDCTokenExchange( + cfg *OIDCTokenExchangeConfig, + idTokenProvider IDTokenSource, +) auth.TokenSource { + return &oidcTokenExchange{ + clientID: cfg.ClientID, + accountID: cfg.AccountID, + host: cfg.Host, + tokenEndpoint: cfg.TokenEndpoint, + audience: cfg.Audience, + tokenProvider: idTokenProvider, } - return strategies } -func newOidcTokenStrategy( - name string, - cfg *Config, - tokenProvider TokenProvider, -) CredentialsStrategy { - accountId := "" - // NOTE: cfg.AccountID can be present even if the client is not an account client. - if cfg.IsAccountClient() { - accountId = cfg.AccountID - } - oidcTokenExchange := &oidcTokenExchange{ - clientID: cfg.ClientID, - accountID: accountId, - host: cfg.Host, - tokenEndpointProvider: cfg.getOidcEndpoints, - audience: cfg.TokenAudience, - tokenProvider: tokenProvider, - } - return NewTokenSourceStrategy(name, oidcTokenExchange) +type OIDCTokenExchangeConfig struct { + ClientID string + AccountID string + Host string + TokenEndpoint string + Audience string + IdTokenProvider IDTokenSource } // oidcTokenExchange is a auth.TokenSource which exchanges a token using // Workload Identity Federation. type oidcTokenExchange struct { - clientID string - accountID string - host string - tokenEndpointProvider func(ctx context.Context) (*u2m.OAuthAuthorizationServer, error) - audience string - tokenProvider TokenProvider + clientID string + accountID string + host string + tokenEndpoint string + audience string + tokenProvider IDTokenSource } // Token implements [TokenSource.Token] @@ -70,22 +55,15 @@ func (w *oidcTokenExchange) Token(ctx context.Context) (*oauth2.Token, error) { logger.Debugf(ctx, "Missing Host") return nil, errors.New("missing Host") } - audience, err := w.determineAudience(ctx) - if err != nil { - return nil, err - } + audience := w.determineAudience() idToken, err := w.tokenProvider.IDToken(ctx, audience) if err != nil { return nil, err } - endpoints, err := w.tokenEndpointProvider(ctx) - if err != nil { - return nil, err - } c := &clientcredentials.Config{ ClientID: w.clientID, AuthStyle: oauth2.AuthStyleInParams, - TokenURL: endpoints.TokenEndpoint, + TokenURL: w.tokenEndpoint, Scopes: []string{"all-apis"}, EndpointParams: url.Values{ "subject_token_type": {"urn:ietf:params:oauth:token-type:jwt"}, @@ -96,18 +74,14 @@ func (w *oidcTokenExchange) Token(ctx context.Context) (*oauth2.Token, error) { return c.Token(ctx) } -func (w *oidcTokenExchange) determineAudience(ctx context.Context) (string, error) { +func (w *oidcTokenExchange) determineAudience() string { if w.audience != "" { - return w.audience, nil + return w.audience } // For Databricks Accounts, the account id is the default audience. if w.accountID != "" { - return w.accountID, nil - } - endpoints, err := w.tokenEndpointProvider(ctx) - if err != nil { - return "", err + return w.accountID } // For Databricks Workspaces, the auth endpoint is the default audience. - return endpoints.TokenEndpoint, nil + return w.tokenEndpoint } diff --git a/config/auth_databricks_oidc_test.go b/config/auth_databricks_oidc_test.go index 47ffc5e12..3ca0ba07b 100644 --- a/config/auth_databricks_oidc_test.go +++ b/config/auth_databricks_oidc_test.go @@ -7,7 +7,6 @@ import ( "net/url" "testing" - "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" @@ -30,18 +29,18 @@ func (m *mockIdTokenProvider) IDToken(ctx context.Context, audience string) (*ID func TestDatabricksOidcTokenSource(t *testing.T) { testCases := []struct { - desc string - clientID string - accountID string - host string - tokenAudience string - httpTransport http.RoundTripper - oidcEndpointSupplier func(ctx context.Context) (*u2m.OAuthAuthorizationServer, error) - idToken string - expectedAudience string - tokenProviderError error - wantToken string - wantErrPrefix *string + desc string + clientID string + accountID string + host string + tokenAudience string + httpTransport http.RoundTripper + oidcEndpoint string + idToken string + expectedAudience string + tokenProviderError error + wantToken string + wantErrPrefix *string }{ { desc: "missing host", @@ -55,28 +54,13 @@ func TestDatabricksOidcTokenSource(t *testing.T) { tokenAudience: "token-audience", wantErrPrefix: errPrefix("missing ClientID"), }, - { - desc: "auth server error", - clientID: "client-id", - host: "http://host.com", - oidcEndpointSupplier: func(ctx context.Context) (*u2m.OAuthAuthorizationServer, error) { - return nil, errors.New("databricks OAuth is not supported for this host") - }, - wantErrPrefix: errPrefix("databricks OAuth is not supported for this host"), - }, { desc: "token provider error", - clientID: "client-id", - host: "http://host.com", - tokenAudience: "token-audience", - oidcEndpointSupplier: func(ctx context.Context) (*u2m.OAuthAuthorizationServer, error) { - return &u2m.OAuthAuthorizationServer{ - AuthorizationEndpoint: "https://host.com/auth", - TokenEndpoint: "https://host.com/oidc/v1/token", - }, nil - }, - + clientID: "client-id", + host: "http://host.com", + tokenAudience: "token-audience", + oidcEndpoint: "https://host.com/oidc/v1/token", expectedAudience: "token-audience", tokenProviderError: errors.New("error getting id token"), wantErrPrefix: errPrefix("error getting id token"), @@ -86,12 +70,7 @@ func TestDatabricksOidcTokenSource(t *testing.T) { clientID: "client-id", host: "http://host.com", tokenAudience: "token-audience", - oidcEndpointSupplier: func(ctx context.Context) (*u2m.OAuthAuthorizationServer, error) { - return &u2m.OAuthAuthorizationServer{ - AuthorizationEndpoint: "https://host.com/auth", - TokenEndpoint: "https://host.com/oidc/v1/token", - }, nil - }, + oidcEndpoint: "https://host.com/oidc/v1/token", httpTransport: fixtures.MappingTransport{ "POST /oidc/v1/token": { Status: http.StatusInternalServerError, @@ -109,12 +88,7 @@ func TestDatabricksOidcTokenSource(t *testing.T) { clientID: "client-id", host: "http://host.com", tokenAudience: "token-audience", - oidcEndpointSupplier: func(ctx context.Context) (*u2m.OAuthAuthorizationServer, error) { - return &u2m.OAuthAuthorizationServer{ - AuthorizationEndpoint: "https://host.com/auth", - TokenEndpoint: "https://host.com/oidc/v1/token", - }, nil - }, + oidcEndpoint: "https://host.com/oidc/v1/token", httpTransport: fixtures.MappingTransport{ "POST /oidc/v1/token": { Status: http.StatusOK, @@ -135,12 +109,7 @@ func TestDatabricksOidcTokenSource(t *testing.T) { clientID: "client-id", host: "http://host.com", tokenAudience: "token-audience", - oidcEndpointSupplier: func(ctx context.Context) (*u2m.OAuthAuthorizationServer, error) { - return &u2m.OAuthAuthorizationServer{ - AuthorizationEndpoint: "https://host.com/auth", - TokenEndpoint: "https://host.com/oidc/v1/token", - }, nil - }, + oidcEndpoint: "https://host.com/oidc/v1/token", httpTransport: fixtures.MappingTransport{ "POST /oidc/v1/token": { @@ -173,14 +142,9 @@ func TestDatabricksOidcTokenSource(t *testing.T) { accountID: "ac123", host: "https://accounts.databricks.com", tokenAudience: "token-audience", - oidcEndpointSupplier: func(ctx context.Context) (*u2m.OAuthAuthorizationServer, error) { - return &u2m.OAuthAuthorizationServer{ - AuthorizationEndpoint: "https://host.com/auth", - TokenEndpoint: "https://host.com/oidc/accounts/ac123/v1/token", - }, nil - }, + oidcEndpoint: "https://host.com/oidc/v1/token", httpTransport: fixtures.MappingTransport{ - "POST /oidc/accounts/ac123/v1/token": { + "POST /oidc/v1/token": { Status: http.StatusOK, ExpectedHeaders: map[string]string{ "Content-Type": "application/x-www-form-urlencoded", @@ -205,18 +169,13 @@ func TestDatabricksOidcTokenSource(t *testing.T) { wantToken: "test-auth-token", }, { - desc: "default token audience account", - clientID: "client-id", - accountID: "ac123", - host: "https://accounts.databricks.com", - oidcEndpointSupplier: func(ctx context.Context) (*u2m.OAuthAuthorizationServer, error) { - return &u2m.OAuthAuthorizationServer{ - AuthorizationEndpoint: "https://host.com/auth", - TokenEndpoint: "https://host.com/oidc/accounts/ac123/v1/token", - }, nil - }, + desc: "default token audience account", + clientID: "client-id", + accountID: "ac123", + host: "https://accounts.databricks.com", + oidcEndpoint: "https://host.com/oidc/v1/token", httpTransport: fixtures.MappingTransport{ - "POST /oidc/accounts/ac123/v1/token": { + "POST /oidc/v1/token": { Status: http.StatusOK, ExpectedHeaders: map[string]string{ "Content-Type": "application/x-www-form-urlencoded", @@ -234,15 +193,10 @@ func TestDatabricksOidcTokenSource(t *testing.T) { wantToken: "test-auth-token", }, { - desc: "default token audience workspace", - clientID: "client-id", - host: "https://host.com", - oidcEndpointSupplier: func(ctx context.Context) (*u2m.OAuthAuthorizationServer, error) { - return &u2m.OAuthAuthorizationServer{ - AuthorizationEndpoint: "https://host.com/auth", - TokenEndpoint: "https://host.com/oidc/v1/token", - }, nil - }, + desc: "default token audience workspace", + clientID: "client-id", + host: "https://host.com", + oidcEndpoint: "https://host.com/oidc/v1/token", httpTransport: fixtures.MappingTransport{ "POST /oidc/v1/token": { Status: http.StatusOK, @@ -270,12 +224,12 @@ func TestDatabricksOidcTokenSource(t *testing.T) { err: tc.tokenProviderError, } ex := &oidcTokenExchange{ - clientID: tc.clientID, - accountID: tc.accountID, - host: tc.host, - tokenEndpointProvider: tc.oidcEndpointSupplier, - audience: tc.tokenAudience, - tokenProvider: p, + clientID: tc.clientID, + accountID: tc.accountID, + host: tc.host, + tokenEndpoint: tc.oidcEndpoint, + audience: tc.tokenAudience, + tokenProvider: p, } ctx := context.WithValue(context.Background(), oauth2.HTTPClient, &http.Client{ Transport: tc.httpTransport, @@ -300,15 +254,3 @@ func TestDatabricksOidcTokenSource(t *testing.T) { }) } } - -func TestDatabricksOidcCredentials_Name(t *testing.T) { - strategies := OidcTokenCredentialStrategies(&Config{}) - expected := []string{"github-oidc"} - found := []string{} - for _, strategy := range strategies { - found = append(found, strategy.Name()) - } - if diff := cmp.Diff(expected, found); diff != "" { - t.Errorf("Strategies mismatch (-want +got):\n%s\n(order must be the same))", diff) - } -} diff --git a/config/auth_default.go b/config/auth_default.go index 3709fcc0b..e2104d009 100644 --- a/config/auth_default.go +++ b/config/auth_default.go @@ -9,7 +9,40 @@ import ( "github.com/databricks/databricks-sdk-go/logger" ) -func buildDefaultStrategies(cfg *Config) []CredentialsStrategy { +// Constructs all Databricks OIDC Credentials Strategies +func buildOidcTokenCredentialStrategies(ctx context.Context, cfg *Config) ([]CredentialsStrategy, error) { + oidcEndpoints, err := cfg.getOidcEndpoints(ctx) + if err != nil { + return nil, err + } + providers := map[string]IDTokenSource{ + "github-oidc": &GithubProvider{ + actionsIDTokenRequestURL: cfg.ActionsIDTokenRequestURL, + actionsIDTokenRequestToken: cfg.ActionsIDTokenRequestToken, + refreshClient: cfg.refreshClient, + }, + // Add new providers at the end of the list + } + + strategies := []CredentialsStrategy{} + for name, provider := range providers { + oidcConfig := &OIDCTokenExchangeConfig{ + ClientID: cfg.ClientID, + Host: cfg.Host, + TokenEndpoint: oidcEndpoints.TokenEndpoint, + Audience: cfg.TokenAudience, + IdTokenProvider: provider, + } + if cfg.IsAccountClient() { + oidcConfig.AccountID = cfg.AccountID + } + tokenSource := NewOIDCTokenExchange(oidcConfig, provider) + strategies = append(strategies, NewTokenSourceStrategy(name, tokenSource)) + } + return strategies, nil +} + +func buildDefaultStrategies(ctx context.Context, cfg *Config) []CredentialsStrategy { strategies := []CredentialsStrategy{} strategies = append(strategies, PatCredentials{}, @@ -17,7 +50,11 @@ func buildDefaultStrategies(cfg *Config) []CredentialsStrategy { M2mCredentials{}, DatabricksCliCredentials, MetadataServiceCredentials{}) - strategies = append(strategies, OidcTokenCredentialStrategies(cfg)...) + oidcCredentialStrategies, err := buildOidcTokenCredentialStrategies(ctx, cfg) + if err != nil { + logger.Debugf(ctx, "Cannot set up OIDC Credentials Strategy: %v. Skipping", err) + } + strategies = append(strategies, oidcCredentialStrategies...) strategies = append(strategies, // Attempt to configure auth from most specific to most generic (the Azure CLI). AzureGithubOIDCCredentials{}, @@ -48,7 +85,7 @@ 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) { - for _, p := range buildDefaultStrategies(cfg) { + for _, p := range buildDefaultStrategies(ctx, 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) diff --git a/config/token_source_strategy.go b/config/token_source_strategy.go index 888076cc5..b69483199 100644 --- a/config/token_source_strategy.go +++ b/config/token_source_strategy.go @@ -10,24 +10,35 @@ import ( "github.com/databricks/databricks-sdk-go/logger" ) +// Creates a CredentialsStrategy from a TokenSource. +func NewTokenSourceStrategy( + name string, + tokenSource auth.TokenSource, +) CredentialsStrategy { + return &tokenSourceStrategy{ + name: name, + tokenSource: tokenSource, + } +} + type IDToken struct { Value string } -// TokenProvider is an interface for a Token Provider with an audience. -type TokenProvider interface { +// IDTokenSource is an interface for an IDToken Source with an audience. +type IDTokenSource interface { // Function to get the token IDToken(ctx context.Context, audience string) (*IDToken, error) } -// TokenSourceStrategy is wrapper on a auth.TokenSource which converts it into a CredentialsStrategy -type TokenSourceStrategy struct { +// tokenSourceStrategy is wrapper on a auth.TokenSource which converts it into a CredentialsStrategy +type tokenSourceStrategy struct { tokenSource auth.TokenSource name string } // Configure implements [CredentialsStrategy.Configure]. -func (t *TokenSourceStrategy) Configure(ctx context.Context, cfg *Config) (credentials.CredentialsProvider, error) { +func (t *tokenSourceStrategy) Configure(ctx context.Context, cfg *Config) (credentials.CredentialsProvider, error) { // If we cannot get a token, skip this CredentialsStrategy. // We don't want to fail here because it's possible that the supplier is enabled @@ -44,17 +55,6 @@ func (t *TokenSourceStrategy) Configure(ctx context.Context, cfg *Config) (crede } // Name implements [CredentialsStrategy.Name]. -func (t *TokenSourceStrategy) Name() string { +func (t *tokenSourceStrategy) Name() string { return t.name } - -// Creates a CredentialsStrategy from a TokenSource. -func NewTokenSourceStrategy( - name string, - tokenSource auth.TokenSource, -) CredentialsStrategy { - return &TokenSourceStrategy{ - name: name, - tokenSource: tokenSource, - } -} diff --git a/config/token_source_strategy_test.go b/config/token_source_strategy_test.go index 2618cdb5b..7741c9cad 100644 --- a/config/token_source_strategy_test.go +++ b/config/token_source_strategy_test.go @@ -45,7 +45,7 @@ func TestDatabricksTokenSourceStrategy(t *testing.T) { token: tc.token, err: tc.tokenSourceError, } - strat := &TokenSourceStrategy{ + strat := &tokenSourceStrategy{ name: "github-oidc", tokenSource: ts, } From 47f7a2d4d9fd201e952a92a664b0434aeb20cf85 Mon Sep 17 00:00:00 2001 From: Hector Castejon Diaz Date: Tue, 8 Apr 2025 14:11:39 +0200 Subject: [PATCH 21/30] Undo endpoint supply --- client/client_test.go | 6 -- config/auth_databricks_oidc.go | 52 ++++++++------- config/auth_databricks_oidc_test.go | 99 +++++++++++++++++++---------- config/auth_default.go | 28 +++----- 4 files changed, 103 insertions(+), 82 deletions(-) diff --git a/client/client_test.go b/client/client_test.go index 0263a0e8f..267635946 100644 --- a/client/client_test.go +++ b/client/client_test.go @@ -33,10 +33,6 @@ func TestSimpleRequestFailsURLError(t *testing.T) { ConfigFile: "/dev/null", RetryTimeoutSeconds: 1, HTTPTransport: hc(func(r *http.Request) (*http.Response, error) { - // Unrelated to this test. This URL is called during Client setup. - if r.URL.Path == "/oidc/.well-known/oauth-authorization-server" { - return nil, nil - } require.Equal(t, "GET", r.Method) require.Equal(t, "/a/b", r.URL.Path) require.Equal(t, "c=d", r.URL.RawQuery) @@ -316,8 +312,6 @@ func TestHttpTransport(t *testing.T) { func TestDoRemovesDoubleSlashesFromFilesAPI(t *testing.T) { i := 0 expectedPaths := []string{ - // Unrelated to this test. This URL is called during Client setup. - "/oidc/.well-known/oauth-authorization-server", "/api/2.0/fs/files/Volumes/abc/def/ghi", "/api/2.0/anotherservice//test", } diff --git a/config/auth_databricks_oidc.go b/config/auth_databricks_oidc.go index 85188658a..a229eb4a6 100644 --- a/config/auth_databricks_oidc.go +++ b/config/auth_databricks_oidc.go @@ -6,6 +6,7 @@ import ( "net/url" "github.com/databricks/databricks-sdk-go/config/experimental/auth" + "github.com/databricks/databricks-sdk-go/credentials/u2m" "github.com/databricks/databricks-sdk-go/logger" "golang.org/x/oauth2" "golang.org/x/oauth2/clientcredentials" @@ -16,33 +17,33 @@ func NewOIDCTokenExchange( idTokenProvider IDTokenSource, ) auth.TokenSource { return &oidcTokenExchange{ - clientID: cfg.ClientID, - accountID: cfg.AccountID, - host: cfg.Host, - tokenEndpoint: cfg.TokenEndpoint, - audience: cfg.Audience, - tokenProvider: idTokenProvider, + clientID: cfg.ClientID, + accountID: cfg.AccountID, + host: cfg.Host, + tokenEndpointProvider: cfg.TokenEndpointProvider, + audience: cfg.Audience, + idTokenProvider: idTokenProvider, } } type OIDCTokenExchangeConfig struct { - ClientID string - AccountID string - Host string - TokenEndpoint string - Audience string - IdTokenProvider IDTokenSource + ClientID string + AccountID string + Host string + TokenEndpointProvider func(ctx context.Context) (*u2m.OAuthAuthorizationServer, error) + Audience string + IdTokenProvider IDTokenSource } // oidcTokenExchange is a auth.TokenSource which exchanges a token using // Workload Identity Federation. type oidcTokenExchange struct { - clientID string - accountID string - host string - tokenEndpoint string - audience string - tokenProvider IDTokenSource + clientID string + accountID string + host string + tokenEndpointProvider func(ctx context.Context) (*u2m.OAuthAuthorizationServer, error) + audience string + idTokenProvider IDTokenSource } // Token implements [TokenSource.Token] @@ -55,15 +56,20 @@ func (w *oidcTokenExchange) Token(ctx context.Context) (*oauth2.Token, error) { logger.Debugf(ctx, "Missing Host") return nil, errors.New("missing Host") } - audience := w.determineAudience() - idToken, err := w.tokenProvider.IDToken(ctx, audience) + endpoints, err := w.tokenEndpointProvider(ctx) if err != nil { return nil, err } + audience := w.determineAudience(endpoints) + idToken, err := w.idTokenProvider.IDToken(ctx, audience) + if err != nil { + return nil, err + } + c := &clientcredentials.Config{ ClientID: w.clientID, AuthStyle: oauth2.AuthStyleInParams, - TokenURL: w.tokenEndpoint, + TokenURL: endpoints.TokenEndpoint, Scopes: []string{"all-apis"}, EndpointParams: url.Values{ "subject_token_type": {"urn:ietf:params:oauth:token-type:jwt"}, @@ -74,7 +80,7 @@ func (w *oidcTokenExchange) Token(ctx context.Context) (*oauth2.Token, error) { return c.Token(ctx) } -func (w *oidcTokenExchange) determineAudience() string { +func (w *oidcTokenExchange) determineAudience(endpoints *u2m.OAuthAuthorizationServer) string { if w.audience != "" { return w.audience } @@ -83,5 +89,5 @@ func (w *oidcTokenExchange) determineAudience() string { return w.accountID } // For Databricks Workspaces, the auth endpoint is the default audience. - return w.tokenEndpoint + return endpoints.TokenEndpoint } diff --git a/config/auth_databricks_oidc_test.go b/config/auth_databricks_oidc_test.go index 3ca0ba07b..fb86d6c5c 100644 --- a/config/auth_databricks_oidc_test.go +++ b/config/auth_databricks_oidc_test.go @@ -7,6 +7,7 @@ import ( "net/url" "testing" + "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" @@ -29,18 +30,18 @@ func (m *mockIdTokenProvider) IDToken(ctx context.Context, audience string) (*ID func TestDatabricksOidcTokenSource(t *testing.T) { testCases := []struct { - desc string - clientID string - accountID string - host string - tokenAudience string - httpTransport http.RoundTripper - oidcEndpoint string - idToken string - expectedAudience string - tokenProviderError error - wantToken string - wantErrPrefix *string + desc string + clientID string + accountID string + host string + tokenAudience string + httpTransport http.RoundTripper + oidcEndpointProvider func(context.Context) (*u2m.OAuthAuthorizationServer, error) + idToken string + expectedAudience string + tokenProviderError error + wantToken string + wantErrPrefix *string }{ { desc: "missing host", @@ -57,10 +58,14 @@ func TestDatabricksOidcTokenSource(t *testing.T) { { desc: "token provider error", - clientID: "client-id", - host: "http://host.com", - tokenAudience: "token-audience", - oidcEndpoint: "https://host.com/oidc/v1/token", + clientID: "client-id", + 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 + }, expectedAudience: "token-audience", tokenProviderError: errors.New("error getting id token"), wantErrPrefix: errPrefix("error getting id token"), @@ -70,7 +75,11 @@ func TestDatabricksOidcTokenSource(t *testing.T) { clientID: "client-id", host: "http://host.com", tokenAudience: "token-audience", - oidcEndpoint: "https://host.com/oidc/v1/token", + 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.StatusInternalServerError, @@ -88,7 +97,11 @@ func TestDatabricksOidcTokenSource(t *testing.T) { clientID: "client-id", host: "http://host.com", tokenAudience: "token-audience", - oidcEndpoint: "https://host.com/oidc/v1/token", + 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, @@ -109,7 +122,11 @@ func TestDatabricksOidcTokenSource(t *testing.T) { clientID: "client-id", host: "http://host.com", tokenAudience: "token-audience", - oidcEndpoint: "https://host.com/oidc/v1/token", + 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": { @@ -142,7 +159,11 @@ func TestDatabricksOidcTokenSource(t *testing.T) { accountID: "ac123", host: "https://accounts.databricks.com", tokenAudience: "token-audience", - oidcEndpoint: "https://host.com/oidc/v1/token", + 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, @@ -169,11 +190,15 @@ func TestDatabricksOidcTokenSource(t *testing.T) { wantToken: "test-auth-token", }, { - desc: "default token audience account", - clientID: "client-id", - accountID: "ac123", - host: "https://accounts.databricks.com", - oidcEndpoint: "https://host.com/oidc/v1/token", + desc: "default token audience account", + clientID: "client-id", + accountID: "ac123", + host: "https://accounts.databricks.com", + 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, @@ -193,10 +218,14 @@ func TestDatabricksOidcTokenSource(t *testing.T) { wantToken: "test-auth-token", }, { - desc: "default token audience workspace", - clientID: "client-id", - host: "https://host.com", - oidcEndpoint: "https://host.com/oidc/v1/token", + desc: "default token audience workspace", + clientID: "client-id", + host: "https://host.com", + 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, @@ -224,12 +253,12 @@ func TestDatabricksOidcTokenSource(t *testing.T) { err: tc.tokenProviderError, } ex := &oidcTokenExchange{ - clientID: tc.clientID, - accountID: tc.accountID, - host: tc.host, - tokenEndpoint: tc.oidcEndpoint, - audience: tc.tokenAudience, - tokenProvider: p, + clientID: tc.clientID, + accountID: tc.accountID, + host: tc.host, + tokenEndpointProvider: tc.oidcEndpointProvider, + audience: tc.tokenAudience, + idTokenProvider: p, } ctx := context.WithValue(context.Background(), oauth2.HTTPClient, &http.Client{ Transport: tc.httpTransport, diff --git a/config/auth_default.go b/config/auth_default.go index e2104d009..79bd96959 100644 --- a/config/auth_default.go +++ b/config/auth_default.go @@ -10,11 +10,7 @@ import ( ) // Constructs all Databricks OIDC Credentials Strategies -func buildOidcTokenCredentialStrategies(ctx context.Context, cfg *Config) ([]CredentialsStrategy, error) { - oidcEndpoints, err := cfg.getOidcEndpoints(ctx) - if err != nil { - return nil, err - } +func buildOidcTokenCredentialStrategies(cfg *Config) []CredentialsStrategy { providers := map[string]IDTokenSource{ "github-oidc": &GithubProvider{ actionsIDTokenRequestURL: cfg.ActionsIDTokenRequestURL, @@ -27,11 +23,11 @@ func buildOidcTokenCredentialStrategies(ctx context.Context, cfg *Config) ([]Cre strategies := []CredentialsStrategy{} for name, provider := range providers { oidcConfig := &OIDCTokenExchangeConfig{ - ClientID: cfg.ClientID, - Host: cfg.Host, - TokenEndpoint: oidcEndpoints.TokenEndpoint, - Audience: cfg.TokenAudience, - IdTokenProvider: provider, + ClientID: cfg.ClientID, + Host: cfg.Host, + TokenEndpointProvider: cfg.getOidcEndpoints, + Audience: cfg.TokenAudience, + IdTokenProvider: provider, } if cfg.IsAccountClient() { oidcConfig.AccountID = cfg.AccountID @@ -39,10 +35,10 @@ func buildOidcTokenCredentialStrategies(ctx context.Context, cfg *Config) ([]Cre tokenSource := NewOIDCTokenExchange(oidcConfig, provider) strategies = append(strategies, NewTokenSourceStrategy(name, tokenSource)) } - return strategies, nil + return strategies } -func buildDefaultStrategies(ctx context.Context, cfg *Config) []CredentialsStrategy { +func buildDefaultStrategies(cfg *Config) []CredentialsStrategy { strategies := []CredentialsStrategy{} strategies = append(strategies, PatCredentials{}, @@ -50,11 +46,7 @@ func buildDefaultStrategies(ctx context.Context, cfg *Config) []CredentialsStrat M2mCredentials{}, DatabricksCliCredentials, MetadataServiceCredentials{}) - oidcCredentialStrategies, err := buildOidcTokenCredentialStrategies(ctx, cfg) - if err != nil { - logger.Debugf(ctx, "Cannot set up OIDC Credentials Strategy: %v. Skipping", err) - } - strategies = append(strategies, oidcCredentialStrategies...) + strategies = append(strategies, buildOidcTokenCredentialStrategies(cfg)...) strategies = append(strategies, // Attempt to configure auth from most specific to most generic (the Azure CLI). AzureGithubOIDCCredentials{}, @@ -85,7 +77,7 @@ 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) { - for _, p := range buildDefaultStrategies(ctx, cfg) { + 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) From 10e5615147d4ba879a32c6eb6f655414381dad7d Mon Sep 17 00:00:00 2001 From: Hector Castejon Diaz Date: Tue, 8 Apr 2025 14:20:22 +0200 Subject: [PATCH 22/30] Cleanup --- config/auth_databricks_oidc.go | 23 ++++++++++++----------- config/auth_databricks_oidc_test.go | 2 +- config/auth_default.go | 6 +++--- 3 files changed, 16 insertions(+), 15 deletions(-) diff --git a/config/auth_databricks_oidc.go b/config/auth_databricks_oidc.go index a229eb4a6..fd9114981 100644 --- a/config/auth_databricks_oidc.go +++ b/config/auth_databricks_oidc.go @@ -12,32 +12,33 @@ import ( "golang.org/x/oauth2/clientcredentials" ) -func NewOIDCTokenExchange( - cfg *OIDCTokenExchangeConfig, - idTokenProvider IDTokenSource, +// Creates a new Databricks OIDC TokenSource. +func NewDatabricksOIDCTokenSource( + cfg *DatabricksOIDCTokenSourceConfig, ) auth.TokenSource { - return &oidcTokenExchange{ + return &databricksOIDCTokenSource{ clientID: cfg.ClientID, accountID: cfg.AccountID, host: cfg.Host, tokenEndpointProvider: cfg.TokenEndpointProvider, audience: cfg.Audience, - idTokenProvider: idTokenProvider, + idTokenProvider: cfg.IdTokenSource, } } -type OIDCTokenExchangeConfig struct { +// Config for Databricks OIDC TokenSource. +type DatabricksOIDCTokenSourceConfig struct { ClientID string AccountID string Host string TokenEndpointProvider func(ctx context.Context) (*u2m.OAuthAuthorizationServer, error) Audience string - IdTokenProvider IDTokenSource + IdTokenSource IDTokenSource } -// oidcTokenExchange is a auth.TokenSource which exchanges a token using +// databricksOIDCTokenSource is a auth.TokenSource which exchanges a token using // Workload Identity Federation. -type oidcTokenExchange struct { +type databricksOIDCTokenSource struct { clientID string accountID string host string @@ -47,7 +48,7 @@ type oidcTokenExchange struct { } // Token implements [TokenSource.Token] -func (w *oidcTokenExchange) Token(ctx context.Context) (*oauth2.Token, error) { +func (w *databricksOIDCTokenSource) Token(ctx context.Context) (*oauth2.Token, error) { if w.clientID == "" { logger.Debugf(ctx, "Missing ClientID") return nil, errors.New("missing ClientID") @@ -80,7 +81,7 @@ func (w *oidcTokenExchange) Token(ctx context.Context) (*oauth2.Token, error) { return c.Token(ctx) } -func (w *oidcTokenExchange) determineAudience(endpoints *u2m.OAuthAuthorizationServer) string { +func (w *databricksOIDCTokenSource) determineAudience(endpoints *u2m.OAuthAuthorizationServer) string { if w.audience != "" { return w.audience } diff --git a/config/auth_databricks_oidc_test.go b/config/auth_databricks_oidc_test.go index fb86d6c5c..23f4a2b0b 100644 --- a/config/auth_databricks_oidc_test.go +++ b/config/auth_databricks_oidc_test.go @@ -252,7 +252,7 @@ func TestDatabricksOidcTokenSource(t *testing.T) { idToken: tc.idToken, err: tc.tokenProviderError, } - ex := &oidcTokenExchange{ + ex := &databricksOIDCTokenSource{ clientID: tc.clientID, accountID: tc.accountID, host: tc.host, diff --git a/config/auth_default.go b/config/auth_default.go index 79bd96959..680b6a272 100644 --- a/config/auth_default.go +++ b/config/auth_default.go @@ -22,17 +22,17 @@ func buildOidcTokenCredentialStrategies(cfg *Config) []CredentialsStrategy { strategies := []CredentialsStrategy{} for name, provider := range providers { - oidcConfig := &OIDCTokenExchangeConfig{ + oidcConfig := &DatabricksOIDCTokenSourceConfig{ ClientID: cfg.ClientID, Host: cfg.Host, TokenEndpointProvider: cfg.getOidcEndpoints, Audience: cfg.TokenAudience, - IdTokenProvider: provider, + IdTokenSource: provider, } if cfg.IsAccountClient() { oidcConfig.AccountID = cfg.AccountID } - tokenSource := NewOIDCTokenExchange(oidcConfig, provider) + tokenSource := NewDatabricksOIDCTokenSource(oidcConfig) strategies = append(strategies, NewTokenSourceStrategy(name, tokenSource)) } return strategies From 2431d67d4102467e33be8b10f6c6055a031fc99e Mon Sep 17 00:00:00 2001 From: Hector Castejon Diaz Date: Wed, 9 Apr 2025 08:37:00 +0200 Subject: [PATCH 23/30] PR comments --- config/auth_databricks_oidc.go | 55 +++++++++++++--------------- config/auth_databricks_oidc_test.go | 43 ++++++++++++---------- config/token_source_strategy.go | 22 ++++++----- config/token_source_strategy_test.go | 20 +++------- 4 files changed, 67 insertions(+), 73 deletions(-) diff --git a/config/auth_databricks_oidc.go b/config/auth_databricks_oidc.go index fd9114981..2156604c5 100644 --- a/config/auth_databricks_oidc.go +++ b/config/auth_databricks_oidc.go @@ -13,62 +13,59 @@ import ( ) // Creates a new Databricks OIDC TokenSource. -func NewDatabricksOIDCTokenSource( - cfg *DatabricksOIDCTokenSourceConfig, -) auth.TokenSource { +func NewDatabricksOIDCTokenSource(cfg *DatabricksOIDCTokenSourceConfig) auth.TokenSource { return &databricksOIDCTokenSource{ - clientID: cfg.ClientID, - accountID: cfg.AccountID, - host: cfg.Host, - tokenEndpointProvider: cfg.TokenEndpointProvider, - audience: cfg.Audience, - idTokenProvider: cfg.IdTokenSource, + cfg: cfg, } } // Config for Databricks OIDC TokenSource. type DatabricksOIDCTokenSourceConfig struct { - ClientID string - AccountID string - Host string + // ClientID is the client ID of the Databricks OIDC application. For + // Databricks Service Principal, this is the Application ID of the Service Principal. + ClientID string + // [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 string + // TokenEndpointProvider is a function that returns the token endpoint for the Databricks OIDC application. TokenEndpointProvider func(ctx context.Context) (*u2m.OAuthAuthorizationServer, error) - Audience string - IdTokenSource IDTokenSource + // 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 IDTokenSource } // databricksOIDCTokenSource is a auth.TokenSource which exchanges a token using // Workload Identity Federation. type databricksOIDCTokenSource struct { - clientID string - accountID string - host string - tokenEndpointProvider func(ctx context.Context) (*u2m.OAuthAuthorizationServer, error) - audience string - idTokenProvider IDTokenSource + cfg *DatabricksOIDCTokenSourceConfig } // Token implements [TokenSource.Token] func (w *databricksOIDCTokenSource) Token(ctx context.Context) (*oauth2.Token, error) { - if w.clientID == "" { + if w.cfg.ClientID == "" { logger.Debugf(ctx, "Missing ClientID") return nil, errors.New("missing ClientID") } - if w.host == "" { + if w.cfg.Host == "" { logger.Debugf(ctx, "Missing Host") return nil, errors.New("missing Host") } - endpoints, err := w.tokenEndpointProvider(ctx) + endpoints, err := w.cfg.TokenEndpointProvider(ctx) if err != nil { return nil, err } audience := w.determineAudience(endpoints) - idToken, err := w.idTokenProvider.IDToken(ctx, audience) + idToken, err := w.cfg.IdTokenSource.IDToken(ctx, audience) if err != nil { return nil, err } c := &clientcredentials.Config{ - ClientID: w.clientID, + ClientID: w.cfg.ClientID, AuthStyle: oauth2.AuthStyleInParams, TokenURL: endpoints.TokenEndpoint, Scopes: []string{"all-apis"}, @@ -82,12 +79,12 @@ func (w *databricksOIDCTokenSource) Token(ctx context.Context) (*oauth2.Token, e } func (w *databricksOIDCTokenSource) determineAudience(endpoints *u2m.OAuthAuthorizationServer) string { - if w.audience != "" { - return w.audience + if w.cfg.Audience != "" { + return w.cfg.Audience } // For Databricks Accounts, the account id is the default audience. - if w.accountID != "" { - return w.accountID + if w.cfg.AccountID != "" { + return w.cfg.AccountID } // For Databricks Workspaces, the auth endpoint is the default audience. return endpoints.TokenEndpoint diff --git a/config/auth_databricks_oidc_test.go b/config/auth_databricks_oidc_test.go index 23f4a2b0b..40f7da164 100644 --- a/config/auth_databricks_oidc_test.go +++ b/config/auth_databricks_oidc_test.go @@ -10,7 +10,6 @@ import ( "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" ) type mockIdTokenProvider struct { @@ -23,9 +22,7 @@ type mockIdTokenProvider struct { func (m *mockIdTokenProvider) IDToken(ctx context.Context, audience string) (*IDToken, error) { m.audience = audience - return &IDToken{ - m.idToken, - }, m.err + return &IDToken{Value: m.idToken}, m.err } func TestDatabricksOidcTokenSource(t *testing.T) { @@ -252,23 +249,31 @@ func TestDatabricksOidcTokenSource(t *testing.T) { idToken: tc.idToken, err: tc.tokenProviderError, } - ex := &databricksOIDCTokenSource{ - clientID: tc.clientID, - accountID: tc.accountID, - host: tc.host, - tokenEndpointProvider: tc.oidcEndpointProvider, - audience: tc.tokenAudience, - idTokenProvider: p, + + cfg := &DatabricksOIDCTokenSourceConfig{ + ClientID: tc.clientID, + AccountID: tc.accountID, + Host: tc.host, + TokenEndpointProvider: tc.oidcEndpointProvider, + Audience: tc.tokenAudience, + IdTokenSource: p, + } + + ts := NewDatabricksOIDCTokenSource(cfg) + if tc.httpTransport != nil { + ts.(*databricksOIDCTokenSource).cfg.TokenEndpointProvider = func(ctx context.Context) (*u2m.OAuthAuthorizationServer, error) { + return &u2m.OAuthAuthorizationServer{ + TokenEndpoint: "https://host.com/oidc/v1/token", + }, nil + } } - ctx := context.WithValue(context.Background(), oauth2.HTTPClient, &http.Client{ - Transport: tc.httpTransport, - }) - token, gotErr := ex.Token(ctx) - if tc.wantErrPrefix == nil && gotErr != nil { - t.Errorf("Token(ctx): got error %q, want none", gotErr) + + token, err := ts.Token(context.Background()) + if tc.wantErrPrefix == nil && err != nil { + t.Errorf("Token(ctx): got error %q, want none", err) } - if tc.wantErrPrefix != nil && !hasPrefix(gotErr, *tc.wantErrPrefix) { - t.Errorf("Token(ctx): got error %q, want error with prefix %q", gotErr, *tc.wantErrPrefix) + if tc.wantErrPrefix != nil && !hasPrefix(err, *tc.wantErrPrefix) { + t.Errorf("Token(ctx): got error %q, want error with prefix %q", err, *tc.wantErrPrefix) } if tc.expectedAudience != p.audience { t.Errorf("mockTokenProvider: got audience %s, want %s", p.audience, tc.expectedAudience) diff --git a/config/token_source_strategy.go b/config/token_source_strategy.go index b69483199..7e464c73b 100644 --- a/config/token_source_strategy.go +++ b/config/token_source_strategy.go @@ -10,6 +10,18 @@ import ( "github.com/databricks/databricks-sdk-go/logger" ) +// IDToken is a token that can be exchanged for a an access token. +// Value is the token string. +type IDToken struct { + Value string +} + +// IDTokenSource is anything that returns an IDToken given an audience. +type IDTokenSource interface { + // Function to get the token + IDToken(ctx context.Context, audience string) (*IDToken, error) +} + // Creates a CredentialsStrategy from a TokenSource. func NewTokenSourceStrategy( name string, @@ -21,16 +33,6 @@ func NewTokenSourceStrategy( } } -type IDToken struct { - Value string -} - -// IDTokenSource is an interface for an IDToken Source with an audience. -type IDTokenSource interface { - // Function to get the token - IDToken(ctx context.Context, audience string) (*IDToken, error) -} - // tokenSourceStrategy is wrapper on a auth.TokenSource which converts it into a CredentialsStrategy type tokenSourceStrategy struct { tokenSource auth.TokenSource diff --git a/config/token_source_strategy_test.go b/config/token_source_strategy_test.go index 7741c9cad..48ecc38c7 100644 --- a/config/token_source_strategy_test.go +++ b/config/token_source_strategy_test.go @@ -6,19 +6,11 @@ import ( "net/http" "testing" + "github.com/databricks/databricks-sdk-go/config/experimental/auth" "github.com/google/go-cmp/cmp" "golang.org/x/oauth2" ) -type staticTokenSource struct { - token *oauth2.Token - err error -} - -func (s *staticTokenSource) Token(ctx context.Context) (*oauth2.Token, error) { - return s.token, s.err -} - func TestDatabricksTokenSourceStrategy(t *testing.T) { testCases := []struct { desc string @@ -41,13 +33,11 @@ func TestDatabricksTokenSourceStrategy(t *testing.T) { for _, tc := range testCases { t.Run(tc.desc, func(t *testing.T) { - ts := &staticTokenSource{ - token: tc.token, - err: tc.tokenSourceError, - } strat := &tokenSourceStrategy{ - name: "github-oidc", - tokenSource: ts, + name: "github-oidc", + tokenSource: auth.TokenSourceFn(func(_ context.Context) (*oauth2.Token, error) { + return tc.token, tc.tokenSourceError + }), } provider, err := strat.Configure(context.Background(), &Config{}) if tc.tokenSourceError == nil && provider == nil { From a6e1769bf7c1bcd429c54ad6fb4ddd18c48ea050 Mon Sep 17 00:00:00 2001 From: Hector Castejon Diaz Date: Wed, 9 Apr 2025 08:39:39 +0200 Subject: [PATCH 24/30] Name --- config/auth_default.go | 2 +- config/token_provider_github_oidc.go | 6 +++--- config/token_provider_github_oidc_test.go | 2 +- 3 files changed, 5 insertions(+), 5 deletions(-) diff --git a/config/auth_default.go b/config/auth_default.go index 680b6a272..d832ca137 100644 --- a/config/auth_default.go +++ b/config/auth_default.go @@ -12,7 +12,7 @@ import ( // Constructs all Databricks OIDC Credentials Strategies func buildOidcTokenCredentialStrategies(cfg *Config) []CredentialsStrategy { providers := map[string]IDTokenSource{ - "github-oidc": &GithubProvider{ + "github-oidc": &githubIDTokenSource{ actionsIDTokenRequestURL: cfg.ActionsIDTokenRequestURL, actionsIDTokenRequestToken: cfg.ActionsIDTokenRequestToken, refreshClient: cfg.refreshClient, diff --git a/config/token_provider_github_oidc.go b/config/token_provider_github_oidc.go index 95b89d977..219726cc0 100644 --- a/config/token_provider_github_oidc.go +++ b/config/token_provider_github_oidc.go @@ -10,9 +10,9 @@ import ( ) /** - * GithubProvider retrieves JWT Tokens from Github Actions. + * githubIDTokenSource retrieves JWT Tokens from Github Actions. */ -type GithubProvider struct { +type githubIDTokenSource struct { actionsIDTokenRequestURL string actionsIDTokenRequestToken string refreshClient *httpclient.ApiClient @@ -20,7 +20,7 @@ type GithubProvider struct { // IDToken returns a JWT Token for the specified audience. It will return // an error if not running in GitHub Actions. -func (g *GithubProvider) IDToken(ctx context.Context, audience string) (*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") diff --git a/config/token_provider_github_oidc_test.go b/config/token_provider_github_oidc_test.go index dba174349..65e9486f9 100644 --- a/config/token_provider_github_oidc_test.go +++ b/config/token_provider_github_oidc_test.go @@ -70,7 +70,7 @@ func TestGithubOIDCProvider(t *testing.T) { cli := httpclient.NewApiClient(httpclient.ClientConfig{ Transport: tc.httpTransport, }) - p := &GithubProvider{ + p := &githubIDTokenSource{ actionsIDTokenRequestURL: tc.tokenRequestUrl, actionsIDTokenRequestToken: tc.tokenRequestToken, refreshClient: cli, From 72d715eae54c446f4a22c9dc4cb469ed9d85a027 Mon Sep 17 00:00:00 2001 From: Hector Castejon Diaz Date: Wed, 9 Apr 2025 08:42:13 +0200 Subject: [PATCH 25/30] Revert mistake --- config/auth_databricks_oidc_test.go | 10 +++++++++- 1 file changed, 9 insertions(+), 1 deletion(-) diff --git a/config/auth_databricks_oidc_test.go b/config/auth_databricks_oidc_test.go index 40f7da164..eb664f87a 100644 --- a/config/auth_databricks_oidc_test.go +++ b/config/auth_databricks_oidc_test.go @@ -10,6 +10,7 @@ import ( "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" ) type mockIdTokenProvider struct { @@ -268,7 +269,14 @@ func TestDatabricksOidcTokenSource(t *testing.T) { } } - token, err := ts.Token(context.Background()) + ctx := context.Background() + if tc.httpTransport != nil { + ctx = context.WithValue(ctx, oauth2.HTTPClient, &http.Client{ + Transport: tc.httpTransport, + }) + } + + token, err := ts.Token(ctx) if tc.wantErrPrefix == nil && err != nil { t.Errorf("Token(ctx): got error %q, want none", err) } From 7f32a6787ee46f58684e4db225e739ac8da25235 Mon Sep 17 00:00:00 2001 From: Hector Castejon Diaz Date: Wed, 9 Apr 2025 09:49:43 +0200 Subject: [PATCH 26/30] Ensure credentials order --- config/auth_default.go | 33 +++++++++++++++++++++++++-------- 1 file changed, 25 insertions(+), 8 deletions(-) diff --git a/config/auth_default.go b/config/auth_default.go index d832ca137..a5b6c9f5e 100644 --- a/config/auth_default.go +++ b/config/auth_default.go @@ -10,8 +10,13 @@ import ( ) // Constructs all Databricks OIDC Credentials Strategies -func buildOidcTokenCredentialStrategies(cfg *Config) []CredentialsStrategy { - providers := map[string]IDTokenSource{ +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 + } + idTokenSources := map[string]IDTokenSource{ "github-oidc": &githubIDTokenSource{ actionsIDTokenRequestURL: cfg.ActionsIDTokenRequestURL, actionsIDTokenRequestToken: cfg.ActionsIDTokenRequestToken, @@ -21,7 +26,11 @@ func buildOidcTokenCredentialStrategies(cfg *Config) []CredentialsStrategy { } strategies := []CredentialsStrategy{} - for name, provider := range providers { + for _, name := range idTokenSourceOrder { + provider, ok := idTokenSources[name] + if !ok { + return nil, fmt.Errorf("no provider found for %s", name) + } oidcConfig := &DatabricksOIDCTokenSourceConfig{ ClientID: cfg.ClientID, Host: cfg.Host, @@ -35,10 +44,10 @@ func buildOidcTokenCredentialStrategies(cfg *Config) []CredentialsStrategy { tokenSource := NewDatabricksOIDCTokenSource(oidcConfig) strategies = append(strategies, NewTokenSourceStrategy(name, tokenSource)) } - return strategies + return strategies, nil } -func buildDefaultStrategies(cfg *Config) []CredentialsStrategy { +func buildDefaultStrategies(cfg *Config) ([]CredentialsStrategy, error) { strategies := []CredentialsStrategy{} strategies = append(strategies, PatCredentials{}, @@ -46,7 +55,11 @@ func buildDefaultStrategies(cfg *Config) []CredentialsStrategy { M2mCredentials{}, DatabricksCliCredentials, MetadataServiceCredentials{}) - strategies = append(strategies, buildOidcTokenCredentialStrategies(cfg)...) + oidcStrategies, err := buildOidcTokenCredentialStrategies(cfg) + if err != nil { + return nil, err + } + strategies = append(strategies, oidcStrategies...) strategies = append(strategies, // Attempt to configure auth from most specific to most generic (the Azure CLI). AzureGithubOIDCCredentials{}, @@ -56,7 +69,7 @@ func buildDefaultStrategies(cfg *Config) []CredentialsStrategy { // Attempt to configure auth from most specific to most generic (Google Application Default Credentials). GoogleCredentials{}, GoogleDefaultCredentials{}) - return strategies + return strategies, nil } type DefaultCredentials struct { @@ -77,7 +90,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) { - for _, p := range buildDefaultStrategies(cfg) { + strategies, err := buildDefaultStrategies(cfg) + if err != nil { + return nil, err + } + for _, p := range strategies { 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) From b09f4e855443c7978e443555705fcc58f76e899b Mon Sep 17 00:00:00 2001 From: Hector Castejon Diaz Date: Wed, 9 Apr 2025 12:44:13 +0200 Subject: [PATCH 27/30] More comments --- README.md | 2 +- config/auth_databricks_oidc.go | 10 +++--- config/auth_databricks_oidc_test.go | 2 +- config/auth_default.go | 53 +++++++++++++---------------- config/oauth_visitors.go | 6 ++-- config/token_source_strategy.go | 2 +- 6 files changed, 34 insertions(+), 41 deletions(-) diff --git a/README.md b/README.md index 32af915ce..488ef11bf 100644 --- a/README.md +++ b/README.md @@ -177,7 +177,7 @@ Depending on the Databricks authentication method, the SDK uses the following in By default, the Databricks SDK for Go initially tries Databricks token authentication (`AuthType: "pat"` in `*databricks.Config`). If the SDK is unsuccessful, it then tries Workload Identity Federation (WIF) based authentication(`AuthType: "github-oidc"` in `*databricks.Config`). Currently, only GitHub provided JWT Tokens is supported. - For Databricks token authentication, you must provide `Host` and `Token`; or their environment variable or `.databrickscfg` file field equivalents. -- For Databricks OIDC authentication, you must provide the `Host`, `ClientId` and `TokenAudience` _(optional)_ either directly, through the corresponding environment variables, or in your `.databrickscfg` configuration file. +- For Databricks OIDC authentication, you must provide the `Host`, `ClientId` and `TokenAudience` _(optional)_ either directly, through the corresponding environment variables, or in your `.databrickscfg` configuration file. More information can be found in [Databricks Documentation](https://docs.databricks.com/aws/en/dev-tools/auth/oauth-federation#workload-identity-federation) | `*databricks.Config` argument | Description | Environment variable / `.databrickscfg` file field | | ----------------------------- | ------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------ | -------------------------------------------------- | diff --git a/config/auth_databricks_oidc.go b/config/auth_databricks_oidc.go index 2156604c5..434eb49b1 100644 --- a/config/auth_databricks_oidc.go +++ b/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] diff --git a/config/auth_databricks_oidc_test.go b/config/auth_databricks_oidc_test.go index eb664f87a..388766e14 100644 --- a/config/auth_databricks_oidc_test.go +++ b/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, diff --git a/config/auth_default.go b/config/auth_default.go index a5b6c9f5e..a49757d02 100644 --- a/config/auth_default.go +++ b/config/auth_default.go @@ -10,44 +10,41 @@ 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{}, @@ -55,11 +52,7 @@ func buildDefaultStrategies(cfg *Config) ([]CredentialsStrategy, error) { 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) diff --git a/config/oauth_visitors.go b/config/oauth_visitors.go index 10556c19d..69fadc03f 100644 --- a/config/oauth_visitors.go +++ b/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) } diff --git a/config/token_source_strategy.go b/config/token_source_strategy.go index 7e464c73b..fd5d995ce 100644 --- a/config/token_source_strategy.go +++ b/config/token_source_strategy.go @@ -52,7 +52,7 @@ func (t *tokenSourceStrategy) Configure(ctx context.Context, cfg *Config) (crede return nil, nil } - visitor := refreshableAuthVisitor(cached, ctx) + visitor := refreshableAuthVisitor(cached) return credentials.NewOAuthCredentialsProvider(visitor, authconv.OAuth2TokenSource(cached).Token), nil } From 3e040c9875a74cea9ae56720398f76962b1a8d4d Mon Sep 17 00:00:00 2001 From: Hector Castejon Diaz Date: Wed, 9 Apr 2025 15:02:25 +0200 Subject: [PATCH 28/30] Rename --- README.md | 2 +- config/config.go | 2 +- ...n_provider_github_oidc.go => id_token_source_github_oidc.go} | 0 ..._github_oidc_test.go => id_token_source_github_oidc_test.go} | 2 +- 4 files changed, 3 insertions(+), 3 deletions(-) rename config/{token_provider_github_oidc.go => id_token_source_github_oidc.go} (100%) rename config/{token_provider_github_oidc_test.go => id_token_source_github_oidc_test.go} (98%) diff --git a/README.md b/README.md index 488ef11bf..ae8413273 100644 --- a/README.md +++ b/README.md @@ -184,7 +184,7 @@ By default, the Databricks SDK for Go initially tries Databricks token authentic | `Host` | _(String)_ The Databricks host URL for either the Databricks workspace endpoint or the Databricks accounts endpoint. | `DATABRICKS_HOST` / `host` | | `AccountID` | _(String)_ The Databricks account ID for the Databricks accounts endpoint. Only has effect when `Host` is either `https://accounts.cloud.databricks.com/` _(AWS)_, `https://accounts.azuredatabricks.net/` _(Azure)_, or `https://accounts.gcp.databricks.com/` _(GCP)_. | `DATABRICKS_ACCOUNT_ID` / `account_id` | | `Token` | _(String)_ The Databricks personal access token (PAT) _(AWS, Azure, and GCP)_ or Azure Active Directory (Azure AD) token _(Azure)_. | `DATABRICKS_TOKEN` / `token` | -| `TokenAudience` | _(String)_ When using Workload Identity Federation, the audience to specify when fetching an ID token from the ID token supplier. | `TOKEN_AUDIENCE` / `token_audience` | +| `TokenAudience` | _(String)_ When using Workload Identity Federation, the audience to specify when fetching an ID token from the ID token supplier. | `DATABRICKS_TOKEN_AUDIENCE` / `token_audience` | For example, to use Databricks token authentication: diff --git a/config/config.go b/config/config.go index b7f7f95ab..6346d6cf8 100644 --- a/config/config.go +++ b/config/config.go @@ -135,7 +135,7 @@ type Config struct { DatabricksEnvironment *environment.DatabricksEnvironment // When using Workload Identity Federation, the audience to specify when fetching an ID token from the ID token supplier. - TokenAudience string `name:"audience" auth:"-"` + TokenAudience string `name:"audience" env:"DATABRICKS_TOKEN_AUDIENCE" auth:"-"` Loaders []Loader diff --git a/config/token_provider_github_oidc.go b/config/id_token_source_github_oidc.go similarity index 100% rename from config/token_provider_github_oidc.go rename to config/id_token_source_github_oidc.go diff --git a/config/token_provider_github_oidc_test.go b/config/id_token_source_github_oidc_test.go similarity index 98% rename from config/token_provider_github_oidc_test.go rename to config/id_token_source_github_oidc_test.go index 65e9486f9..58a1bbc2b 100644 --- a/config/token_provider_github_oidc_test.go +++ b/config/id_token_source_github_oidc_test.go @@ -10,7 +10,7 @@ import ( "github.com/google/go-cmp/cmp" ) -func TestGithubOIDCProvider(t *testing.T) { +func TestGithubIDTokenSource(t *testing.T) { testCases := []struct { desc string tokenRequestUrl string From 4f59d7d09aca3a166ee0f7901aa3e4216ea4eb2a Mon Sep 17 00:00:00 2001 From: Hector Castejon Diaz Date: Thu, 10 Apr 2025 09:31:59 +0200 Subject: [PATCH 29/30] Remove duplicate code --- config/auth_azure_github_oidc.go | 13 ++++++---- config/oidc_github.go | 42 -------------------------------- 2 files changed, 8 insertions(+), 47 deletions(-) delete mode 100644 config/oidc_github.go diff --git a/config/auth_azure_github_oidc.go b/config/auth_azure_github_oidc.go index e1e9ee5e8..7be69563f 100644 --- a/config/auth_azure_github_oidc.go +++ b/config/auth_azure_github_oidc.go @@ -23,16 +23,19 @@ func (c AzureGithubOIDCCredentials) Name() string { // Configure implements [CredentialsStrategy.Configure]. func (c AzureGithubOIDCCredentials) Configure(ctx context.Context, cfg *Config) (credentials.CredentialsProvider, error) { // Sanity check that the config is configured for Azure Databricks. - if !cfg.IsAzure() || cfg.AzureClientID == "" || cfg.Host == "" || cfg.AzureTenantID == "" { + if !cfg.IsAzure() || cfg.AzureClientID == "" || cfg.Host == "" || cfg.AzureTenantID == "" || cfg.ActionsIDTokenRequestURL == "" || cfg.ActionsIDTokenRequestToken == "" { return nil, nil } - supplier := GithubOIDCTokenSupplier{cfg: cfg} + supplier := githubIDTokenSource{actionsIDTokenRequestURL: cfg.ActionsIDTokenRequestURL, + actionsIDTokenRequestToken: cfg.ActionsIDTokenRequestToken, + refreshClient: cfg.refreshClient, + } - idToken, err := supplier.GetOIDCToken(ctx, "api://AzureADTokenExchange") + idToken, err := supplier.IDToken(ctx, "api://AzureADTokenExchange") if err != nil { return nil, err } - if idToken == "" { + if idToken.Value == "" { return nil, nil } @@ -40,7 +43,7 @@ func (c AzureGithubOIDCCredentials) Configure(ctx context.Context, cfg *Config) aadEndpoint: fmt.Sprintf("%s%s/oauth2/token", cfg.Environment().AzureActiveDirectoryEndpoint(), cfg.AzureTenantID), clientID: cfg.AzureClientID, applicationID: cfg.Environment().AzureApplicationID, - idToken: idToken, + idToken: idToken.Value, httpClient: cfg.refreshClient, } diff --git a/config/oidc_github.go b/config/oidc_github.go deleted file mode 100644 index f7086766a..000000000 --- a/config/oidc_github.go +++ /dev/null @@ -1,42 +0,0 @@ -package config - -import ( - "context" - "fmt" - - "github.com/databricks/databricks-sdk-go/httpclient" - "github.com/databricks/databricks-sdk-go/logger" -) - -type GithubOIDCTokenSupplier struct { - cfg *Config -} - -// requestIDToken requests an ID token from the Github Action. -func (g *GithubOIDCTokenSupplier) GetOIDCToken(ctx context.Context, audience string) (string, error) { - if g.cfg.ActionsIDTokenRequestURL == "" { - logger.Debugf(ctx, "Missing cfg.ActionsIDTokenRequestURL, likely not calling from a Github action") - return "", nil - } - if g.cfg.ActionsIDTokenRequestToken == "" { - logger.Debugf(ctx, "Missing cfg.ActionsIDTokenRequestToken, likely not calling from a Github action") - return "", nil - } - - resp := struct { // anonymous struct to parse the response - Value string `json:"value"` - }{} - requestUrl := g.cfg.ActionsIDTokenRequestURL - if audience != "" { - requestUrl = fmt.Sprintf("%s&audience=%s", requestUrl, audience) - } - err := g.cfg.refreshClient.Do(ctx, "GET", requestUrl, - httpclient.WithRequestHeader("Authorization", fmt.Sprintf("Bearer %s", g.cfg.ActionsIDTokenRequestToken)), - httpclient.WithResponseUnmarshal(&resp), - ) - if err != nil { - return "", fmt.Errorf("failed to request ID token from %s: %w", g.cfg.ActionsIDTokenRequestURL, err) - } - - return resp.Value, nil -} From fa9fd253ddaa38b477eeb3e3789b6856b67e50fc Mon Sep 17 00:00:00 2001 From: Hector Castejon Diaz Date: Thu, 10 Apr 2025 11:44:17 +0200 Subject: [PATCH 30/30] Fix comment --- config/id_token_source_github_oidc.go | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/config/id_token_source_github_oidc.go b/config/id_token_source_github_oidc.go index 219726cc0..6f4048226 100644 --- a/config/id_token_source_github_oidc.go +++ b/config/id_token_source_github_oidc.go @@ -9,9 +9,7 @@ import ( "github.com/databricks/databricks-sdk-go/logger" ) -/** - * githubIDTokenSource retrieves JWT Tokens from Github Actions. - */ +// githubIDTokenSource retrieves JWT Tokens from Github Actions. type githubIDTokenSource struct { actionsIDTokenRequestURL string actionsIDTokenRequestToken string