Add support for OIDC ID token authentication using an environment variable. #1215
There are no files selected for viewing
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -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 | ||
| }) | ||
| } |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -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) | ||
| } | ||
| }) | ||
| } | ||
| } |
Oops, something went wrong.
Oops, something went wrong.
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Uh oh!
There was an error while loading. Please reload this page.