diff --git a/NEXT_CHANGELOG.md b/NEXT_CHANGELOG.md index b4585a98d..339656483 100644 --- a/NEXT_CHANGELOG.md +++ b/NEXT_CHANGELOG.md @@ -4,6 +4,9 @@ ### New Features and Improvements +- Add support for OIDC ID token authentication using an environment variable + ([PR #1215](https://github.com/databricks/databricks-sdk-go/pull/1215)). + ### Bug Fixes ### Documentation diff --git a/config/auth_databricks_oidc.go b/config/auth_databricks_oidc.go index 434eb49b1..7a5f35ef2 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/config/experimental/auth/oidc" "github.com/databricks/databricks-sdk-go/credentials/u2m" "github.com/databricks/databricks-sdk-go/logger" "golang.org/x/oauth2" @@ -35,7 +36,7 @@ type DatabricksOIDCTokenSourceConfig struct { // This is only used for Workspace level tokens. Audience string // IdTokenSource returns the IDToken to be used for the token exchange. - IdTokenSource IDTokenSource + IdTokenSource oidc.IDTokenSource } // databricksOIDCTokenSource is a auth.TokenSource which exchanges a token using diff --git a/config/auth_databricks_oidc_test.go b/config/auth_databricks_oidc_test.go index 388766e14..64077bfab 100644 --- a/config/auth_databricks_oidc_test.go +++ b/config/auth_databricks_oidc_test.go @@ -7,25 +7,13 @@ import ( "net/url" "testing" + "github.com/databricks/databricks-sdk-go/config/experimental/auth/oidc" "github.com/databricks/databricks-sdk-go/credentials/u2m" "github.com/databricks/databricks-sdk-go/httpclient/fixtures" "github.com/google/go-cmp/cmp" "golang.org/x/oauth2" ) -type mockIdTokenProvider struct { - // input - audience string - // output - idToken string - err error -} - -func (m *mockIdTokenProvider) IDToken(ctx context.Context, audience string) (*IDToken, error) { - m.audience = audience - return &IDToken{Value: m.idToken}, m.err -} - func TestDatabricksOidcTokenSource(t *testing.T) { testCases := []struct { desc string @@ -36,7 +24,7 @@ func TestDatabricksOidcTokenSource(t *testing.T) { httpTransport http.RoundTripper oidcEndpointProvider func(context.Context) (*u2m.OAuthAuthorizationServer, error) idToken string - expectedAudience string + wantAudience string tokenProviderError error wantToken string wantErrPrefix *string @@ -64,7 +52,7 @@ func TestDatabricksOidcTokenSource(t *testing.T) { TokenEndpoint: "https://host.com/oidc/v1/token", }, nil }, - expectedAudience: "token-audience", + wantAudience: "token-audience", tokenProviderError: errors.New("error getting id token"), wantErrPrefix: errPrefix("error getting id token"), }, @@ -86,9 +74,9 @@ func TestDatabricksOidcTokenSource(t *testing.T) { }, }, }, - expectedAudience: "token-audience", - idToken: "id-token-42", - wantErrPrefix: errPrefix("oauth2: cannot fetch token: Internal Server Error"), + wantAudience: "token-audience", + idToken: "id-token-42", + wantErrPrefix: errPrefix("oauth2: cannot fetch token: Internal Server Error"), }, { desc: "invalid auth token", @@ -111,9 +99,9 @@ func TestDatabricksOidcTokenSource(t *testing.T) { }, }, }, - expectedAudience: "token-audience", - idToken: "id-token-42", - wantErrPrefix: errPrefix("oauth2: server response missing access_token"), + wantAudience: "token-audience", + idToken: "id-token-42", + wantErrPrefix: errPrefix("oauth2: server response missing access_token"), }, { desc: "success workspace", @@ -147,9 +135,9 @@ func TestDatabricksOidcTokenSource(t *testing.T) { }, }, }, - expectedAudience: "token-audience", - idToken: "id-token-42", - wantToken: "test-auth-token", + wantAudience: "token-audience", + idToken: "id-token-42", + wantToken: "test-auth-token", }, { desc: "success account", @@ -183,9 +171,9 @@ func TestDatabricksOidcTokenSource(t *testing.T) { }, }, }, - expectedAudience: "token-audience", - idToken: "id-token-42", - wantToken: "test-auth-token", + wantAudience: "token-audience", + idToken: "id-token-42", + wantToken: "test-auth-token", }, { desc: "default token audience account", @@ -211,9 +199,9 @@ func TestDatabricksOidcTokenSource(t *testing.T) { }, }, }, - expectedAudience: "ac123", - idToken: "id-token-42", - wantToken: "test-auth-token", + wantAudience: "ac123", + idToken: "id-token-42", + wantToken: "test-auth-token", }, { desc: "default token audience workspace", @@ -238,26 +226,25 @@ func TestDatabricksOidcTokenSource(t *testing.T) { }, }, }, - expectedAudience: "https://host.com/oidc/v1/token", - idToken: "id-token-42", - wantToken: "test-auth-token", + wantAudience: "https://host.com/oidc/v1/token", + idToken: "id-token-42", + wantToken: "test-auth-token", }, } for _, tc := range testCases { t.Run(tc.desc, func(t *testing.T) { - p := &mockIdTokenProvider{ - idToken: tc.idToken, - err: tc.tokenProviderError, - } - + var gotAudience string // set when IDTokenSource is called cfg := DatabricksOIDCTokenSourceConfig{ ClientID: tc.clientID, AccountID: tc.accountID, Host: tc.host, TokenEndpointProvider: tc.oidcEndpointProvider, Audience: tc.tokenAudience, - IdTokenSource: p, + IdTokenSource: oidc.IDTokenSourceFn(func(ctx context.Context, aud string) (*oidc.IDToken, error) { + gotAudience = aud + return &oidc.IDToken{Value: tc.idToken}, tc.tokenProviderError + }), } ts := NewDatabricksOIDCTokenSource(cfg) @@ -283,8 +270,8 @@ func TestDatabricksOidcTokenSource(t *testing.T) { 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) + if tc.wantAudience != gotAudience { + t.Errorf("mockTokenProvider: got audience %s, want %s", gotAudience, tc.wantAudience) } tokenValue := "" if token != nil { diff --git a/config/auth_default.go b/config/auth_default.go index a49757d02..15083ec02 100644 --- a/config/auth_default.go +++ b/config/auth_default.go @@ -6,6 +6,7 @@ import ( "fmt" "github.com/databricks/databricks-sdk-go/config/credentials" + "github.com/databricks/databricks-sdk-go/config/experimental/auth/oidc" "github.com/databricks/databricks-sdk-go/logger" ) @@ -13,9 +14,21 @@ import ( func buildOidcTokenCredentialStrategies(cfg *Config) []CredentialsStrategy { type namedIdTokenSource struct { name string - tokenSource IDTokenSource + tokenSource oidc.IDTokenSource } idTokenSources := []namedIdTokenSource{ + { + name: "env-oidc", + // If the OIDCTokenEnv is not set, use DATABRICKS_OIDC_TOKEN as + // default value. + tokenSource: func() oidc.IDTokenSource { + v := cfg.OIDCTokenEnv + if v == "" { + v = "DATABRICKS_OIDC_TOKEN" + } + return oidc.NewEnvIDTokenSource(v) + }(), + }, { name: "github-oidc", tokenSource: &githubIDTokenSource{ @@ -26,6 +39,7 @@ func buildOidcTokenCredentialStrategies(cfg *Config) []CredentialsStrategy { }, // Add new providers at the end of the list } + strategies := []CredentialsStrategy{} for _, idTokenSource := range idTokenSources { oidcConfig := DatabricksOIDCTokenSourceConfig{ diff --git a/config/config.go b/config/config.go index c81f03610..912fac5c0 100644 --- a/config/config.go +++ b/config/config.go @@ -108,6 +108,9 @@ type Config struct { // specified by this argument. This argument also holds currently selected auth. AuthType string `name:"auth_type" env:"DATABRICKS_AUTH_TYPE" auth:"-"` + // Environment variable name that contains an OIDC ID token. + OIDCTokenEnv string `name:"oidc_token_env" env:"DATABRICKS_OIDC_TOKEN_ENV" auth:"-"` + // Skip SSL certificate verification for HTTP calls. // Use at your own risk or for unit testing purposes. InsecureSkipVerify bool `name:"skip_verify" auth:"-"` diff --git a/config/experimental/auth/oidc/oidc.go b/config/experimental/auth/oidc/oidc.go new file mode 100644 index 000000000..ff671b8dc --- /dev/null +++ b/config/experimental/auth/oidc/oidc.go @@ -0,0 +1,50 @@ +// Package oidc provides utilities for working with OIDC ID tokens. +// +// This package is experimental and subject to change. +package oidc + +import ( + "context" + "fmt" + "os" +) + +// IDToken represents an OIDC ID token that can be exchanged for a Databricks +// access token. +type IDToken struct { + Value string +} + +// IDTokenSource is anything that returns an IDToken given an audience. +type IDTokenSource interface { + IDToken(ctx context.Context, audience string) (*IDToken, error) +} + +// IDTokenSourceFn is an adapter to allow the use of ordinary functions as +// IDTokenSource. +// +// Example: +// +// ts := IDTokenSourceFn(func(ctx context.Context, aud string) (*IDToken, error) { +// return &IDToken{}, nil +// }) +type IDTokenSourceFn func(ctx context.Context, audience string) (*IDToken, error) + +func (fn IDTokenSourceFn) IDToken(ctx context.Context, audience string) (*IDToken, error) { + return fn(ctx, audience) +} + +// NewEnvIDTokenSource returns an IDTokenSource that reads the token from +// environment variable env. +// +// Note that the IDTokenSource does not cache the token and will read the token +// from environment variable env each time. +func NewEnvIDTokenSource(env string) IDTokenSource { + return IDTokenSourceFn(func(ctx context.Context, _ string) (*IDToken, error) { + t := os.Getenv(env) + if t == "" { + return nil, fmt.Errorf("missing env var %q", env) + } + return &IDToken{Value: t}, nil + }) +} diff --git a/config/experimental/auth/oidc/oidc_test.go b/config/experimental/auth/oidc/oidc_test.go new file mode 100644 index 000000000..9c34174cc --- /dev/null +++ b/config/experimental/auth/oidc/oidc_test.go @@ -0,0 +1,106 @@ +package oidc + +import ( + "context" + "fmt" + "testing" + + "github.com/google/go-cmp/cmp" +) + +func TestIDTokenSourceFn(t *testing.T) { + wantToken := &IDToken{Value: "from-func"} + wantErr := fmt.Errorf("test error") + wantAud := "func-audience" + wantCtx := context.Background() + + ts := IDTokenSourceFn(func(gotCtx context.Context, gotAud string) (*IDToken, error) { + if gotCtx != wantCtx { + t.Errorf("unexpected context: got %v, want %v", gotCtx, wantCtx) + } + if gotAud != wantAud { + t.Errorf("unexpected audience: got %q, want %q", gotAud, wantAud) + } + return wantToken, wantErr + }) + + gotToken, gotErr := ts.IDToken(wantCtx, wantAud) + + if gotErr != wantErr { + t.Errorf("IDToken() want error: %v, got error: %v", wantErr, gotErr) + } + if !cmp.Equal(gotToken, wantToken) { + t.Errorf("IDToken() token = %v, want %v", gotToken, wantToken) + } +} + +func TestNewEnvIDTokenSource(t *testing.T) { + testCases := []struct { + desc string + envName string + envValue string + audience string + want *IDToken + wantErr bool + }{ + { + desc: "Success - variable set", + envName: "OIDC_TEST_TOKEN_SUCCESS", + envValue: "test-token-123", + audience: "test-audience-1", + want: &IDToken{Value: "test-token-123"}, + wantErr: false, + }, + { + desc: "Failure - variable not set", + envName: "OIDC_TEST_TOKEN_MISSING", + envValue: "", + audience: "test-audience-2", + want: nil, + wantErr: true, + }, + { + desc: "Failure - variable set to empty string", + envName: "OIDC_TEST_TOKEN_EMPTY", + envValue: "", + audience: "test-audience-3", + want: nil, + wantErr: true, + }, + { + desc: "Success - different variable name", + envName: "ANOTHER_OIDC_TOKEN", + envValue: "another-token-456", + audience: "test-audience-4", + want: &IDToken{Value: "another-token-456"}, + wantErr: false, + }, + { + desc: "Success - empty audience string", + envName: "OIDC_TEST_TOKEN_NO_AUDIENCE", + envValue: "token-no-audience", + audience: "", + want: &IDToken{Value: "token-no-audience"}, + wantErr: false, + }, + } + + for _, tc := range testCases { + t.Run(tc.desc, func(t *testing.T) { + t.Setenv(tc.envName, tc.envValue) + + ts := NewEnvIDTokenSource(tc.envName) + got, gotErr := ts.IDToken(context.Background(), tc.audience) + + if tc.wantErr && gotErr == nil { + t.Fatalf("IDToken() want error, got none") + } + if !tc.wantErr && gotErr != nil { + t.Fatalf("IDToken() want no error, got error: %v", gotErr) + } + if !cmp.Equal(got, tc.want) { + t.Errorf("IDToken() token = %v, want %v", got, tc.want) + } + }) + } +} diff --git a/config/id_token_source_github_oidc.go b/config/id_token_source_github_oidc.go index 6f4048226..93b9bb11c 100644 --- a/config/id_token_source_github_oidc.go +++ b/config/id_token_source_github_oidc.go @@ -5,6 +5,7 @@ import ( "errors" "fmt" + "github.com/databricks/databricks-sdk-go/config/experimental/auth/oidc" "github.com/databricks/databricks-sdk-go/httpclient" "github.com/databricks/databricks-sdk-go/logger" ) @@ -18,7 +19,7 @@ type githubIDTokenSource struct { // IDToken returns a JWT Token for the specified audience. It will return // an error if not running in GitHub Actions. -func (g *githubIDTokenSource) IDToken(ctx context.Context, audience string) (*IDToken, error) { +func (g *githubIDTokenSource) IDToken(ctx context.Context, audience string) (*oidc.IDToken, error) { if g.actionsIDTokenRequestURL == "" { logger.Debugf(ctx, "Missing ActionsIDTokenRequestURL, likely not calling from a Github action") return nil, errors.New("missing ActionsIDTokenRequestURL") @@ -28,7 +29,7 @@ func (g *githubIDTokenSource) IDToken(ctx context.Context, audience string) (*ID return nil, errors.New("missing ActionsIDTokenRequestToken") } - resp := &IDToken{} + resp := &oidc.IDToken{} requestUrl := g.actionsIDTokenRequestURL if audience != "" { requestUrl = fmt.Sprintf("%s&audience=%s", requestUrl, audience) diff --git a/config/id_token_source_github_oidc_test.go b/config/id_token_source_github_oidc_test.go index 58a1bbc2b..af71858d3 100644 --- a/config/id_token_source_github_oidc_test.go +++ b/config/id_token_source_github_oidc_test.go @@ -5,6 +5,7 @@ import ( "net/http" "testing" + "github.com/databricks/databricks-sdk-go/config/experimental/auth/oidc" "github.com/databricks/databricks-sdk-go/httpclient" "github.com/databricks/databricks-sdk-go/httpclient/fixtures" "github.com/google/go-cmp/cmp" @@ -17,7 +18,7 @@ func TestGithubIDTokenSource(t *testing.T) { tokenRequestToken string audience string httpTransport http.RoundTripper - wantToken *IDToken + wantToken *oidc.IDToken wantErrPrefix *string }{ { @@ -59,7 +60,7 @@ func TestGithubIDTokenSource(t *testing.T) { Response: `{"value": "id-token-42"}`, }, }, - wantToken: &IDToken{ + wantToken: &oidc.IDToken{ Value: "id-token-42", }, }, diff --git a/config/token_source_strategy.go b/config/token_source_strategy.go index fd5d995ce..45393ca31 100644 --- a/config/token_source_strategy.go +++ b/config/token_source_strategy.go @@ -10,18 +10,6 @@ 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,