diff --git a/NEXT_CHANGELOG.md b/NEXT_CHANGELOG.md index b9ae889e7a..27ce418ded 100644 --- a/NEXT_CHANGELOG.md +++ b/NEXT_CHANGELOG.md @@ -10,6 +10,8 @@ * Accept `yes` in addition to `y` for confirmation prompts, and show `[y/N]` to indicate that no is the default. * Deprecated `auth env`. The command is hidden from help listings and prints a deprecation warning to stderr; it will be removed in a future release. * Moved file-based OAuth token cache management from the SDK to the CLI. No user-visible change; part of a three-PR sequence that makes the CLI the sole owner of its token cache. +* Added experimental OS-native secure token storage behind the `--secure-storage` flag on `databricks auth login` and the `DATABRICKS_AUTH_STORAGE=secure` environment variable. Hidden from help during MS1. Legacy file-backed token storage remains the default. +* Added experimental OS-native secure token storage opt-in via `DATABRICKS_AUTH_STORAGE=secure` or `[__settings__].auth_storage = secure` in `.databrickscfg`. Legacy file-backed token storage remains the default. ### Bundles * Remove `experimental-jobs-as-code` template, superseded by `pydabs` ([#4999](https://github.com/databricks/cli/pull/4999)). diff --git a/acceptance/cmd/auth/storage-modes/invalid-env/out.test.toml b/acceptance/cmd/auth/storage-modes/invalid-env/out.test.toml new file mode 100644 index 0000000000..d560f1de04 --- /dev/null +++ b/acceptance/cmd/auth/storage-modes/invalid-env/out.test.toml @@ -0,0 +1,5 @@ +Local = true +Cloud = false + +[EnvMatrix] + DATABRICKS_BUNDLE_ENGINE = ["terraform", "direct"] diff --git a/acceptance/cmd/auth/storage-modes/invalid-env/output.txt b/acceptance/cmd/auth/storage-modes/invalid-env/output.txt new file mode 100644 index 0000000000..5723269bb1 --- /dev/null +++ b/acceptance/cmd/auth/storage-modes/invalid-env/output.txt @@ -0,0 +1,5 @@ + +>>> [CLI] auth token --profile nonexistent +Error: DATABRICKS_AUTH_STORAGE: unknown storage mode "bogus" (want legacy, secure, or plaintext) + +Exit code: 1 diff --git a/acceptance/cmd/auth/storage-modes/invalid-env/script b/acceptance/cmd/auth/storage-modes/invalid-env/script new file mode 100644 index 0000000000..8b00e9d54d --- /dev/null +++ b/acceptance/cmd/auth/storage-modes/invalid-env/script @@ -0,0 +1,6 @@ +export DATABRICKS_AUTH_STORAGE=bogus + +# Any auth command that resolves the storage mode must surface the error. +# auth token is the smallest reproducer because it doesn't perform any +# network I/O before resolving the mode. +trace $CLI auth token --profile nonexistent diff --git a/acceptance/cmd/auth/storage-modes/legacy-env-default/out.test.toml b/acceptance/cmd/auth/storage-modes/legacy-env-default/out.test.toml new file mode 100644 index 0000000000..d560f1de04 --- /dev/null +++ b/acceptance/cmd/auth/storage-modes/legacy-env-default/out.test.toml @@ -0,0 +1,5 @@ +Local = true +Cloud = false + +[EnvMatrix] + DATABRICKS_BUNDLE_ENGINE = ["terraform", "direct"] diff --git a/acceptance/cmd/auth/storage-modes/legacy-env-default/output.txt b/acceptance/cmd/auth/storage-modes/legacy-env-default/output.txt new file mode 100644 index 0000000000..8994c41ebb --- /dev/null +++ b/acceptance/cmd/auth/storage-modes/legacy-env-default/output.txt @@ -0,0 +1,11 @@ + +=== Token cache keys before logout +[ + "dev" +] + +>>> [CLI] auth logout --profile dev --auto-approve +Logged out of profile "dev". Use --delete to also remove it from the config file. + +=== Token cache keys after logout (should be empty) +[] diff --git a/acceptance/cmd/auth/storage-modes/legacy-env-default/script b/acceptance/cmd/auth/storage-modes/legacy-env-default/script new file mode 100644 index 0000000000..37f367fd73 --- /dev/null +++ b/acceptance/cmd/auth/storage-modes/legacy-env-default/script @@ -0,0 +1,29 @@ +export DATABRICKS_AUTH_STORAGE=legacy + +cat > "./home/.databrickscfg" < "./home/.databricks/token-cache.json" < 0 { persistentAuthOpts = append(persistentAuthOpts, u2m.WithScopes(scopesList)) @@ -254,6 +280,7 @@ a new profile is created. if err = persistentAuth.Challenge(); err != nil { return err } + dualWriteLegacyHostKey(ctx, tokenCache, oauthArgument, mode) // At this point, an OAuth token has been successfully minted and stored // in the CLI cache. The rest of the command focuses on: // 1. Workspace selection for SPOG hosts (best-effort); @@ -567,35 +594,48 @@ func validateDiscoveryFlagCompatibility(cmd *cobra.Command) error { return nil } +// discoveryLoginInputs groups the dependencies of discoveryLogin. +// See https://google.github.io/styleguide/go/best-practices#option-structure. +type discoveryLoginInputs struct { + dc discoveryClient + profileName string + timeout time.Duration + scopes string + existingProfile *profile.Profile + browserFunc func(string) error + tokenCache cache.TokenCache + mode storage.StorageMode +} + // discoveryLogin runs the login.databricks.com discovery flow. The user // authenticates in the browser, selects a workspace, and the CLI receives // the workspace host from the OAuth callback's iss parameter. -func discoveryLogin(ctx context.Context, dc discoveryClient, tokenCache cache.TokenCache, profileName string, timeout time.Duration, scopes string, existingProfile *profile.Profile, browserFunc func(string) error) error { - arg, err := dc.NewOAuthArgument(profileName) +func discoveryLogin(ctx context.Context, in discoveryLoginInputs) error { + arg, err := in.dc.NewOAuthArgument(in.profileName) if err != nil { return discoveryErr("setting up login.databricks.com", err) } - scopesList := splitScopes(scopes) - if len(scopesList) == 0 && existingProfile != nil && existingProfile.Scopes != "" { - scopesList = splitScopes(existingProfile.Scopes) + scopesList := splitScopes(in.scopes) + if len(scopesList) == 0 && in.existingProfile != nil && in.existingProfile.Scopes != "" { + scopesList = splitScopes(in.existingProfile.Scopes) } opts := []u2m.PersistentAuthOption{ - u2m.WithTokenCache(storage.NewDualWritingTokenCache(tokenCache, arg)), u2m.WithOAuthArgument(arg), - u2m.WithBrowser(browserFunc), + u2m.WithBrowser(in.browserFunc), u2m.WithDiscoveryLogin(), + u2m.WithTokenCache(in.tokenCache), } if len(scopesList) > 0 { opts = append(opts, u2m.WithScopes(scopesList)) } // Apply timeout before creating PersistentAuth so Challenge() respects it. - ctx, cancel := context.WithTimeout(ctx, timeout) + ctx, cancel := context.WithTimeout(ctx, in.timeout) defer cancel() - persistentAuth, err := dc.NewPersistentAuth(ctx, opts...) + persistentAuth, err := in.dc.NewPersistentAuth(ctx, opts...) if err != nil { return discoveryErr("setting up login.databricks.com", err) } @@ -605,6 +645,7 @@ func discoveryLogin(ctx context.Context, dc discoveryClient, tokenCache cache.To if err := persistentAuth.Challenge(); err != nil { return discoveryErr("login via login.databricks.com failed", err) } + dualWriteLegacyHostKey(ctx, in.tokenCache, arg, in.mode) discoveredHost := arg.GetDiscoveredHost() if discoveredHost == "" { @@ -627,7 +668,7 @@ func discoveryLogin(ctx context.Context, dc discoveryClient, tokenCache cache.To return fmt.Errorf("retrieving token after login: %w", err) } - introspection, err := dc.IntrospectToken(ctx, discoveredHost, tok.AccessToken) + introspection, err := in.dc.IntrospectToken(ctx, discoveredHost, tok.AccessToken) if err != nil { log.Debugf(ctx, "token introspection failed (non-fatal): %v", err) } else { @@ -638,10 +679,10 @@ func discoveryLogin(ctx context.Context, dc discoveryClient, tokenCache cache.To accountID = introspection.AccountID } - if existingProfile != nil && existingProfile.AccountID != "" && introspection.AccountID != "" && - existingProfile.AccountID != introspection.AccountID { + if in.existingProfile != nil && in.existingProfile.AccountID != "" && introspection.AccountID != "" && + in.existingProfile.AccountID != introspection.AccountID { log.Warnf(ctx, "detected account ID %q differs from existing profile account ID %q", - introspection.AccountID, existingProfile.AccountID) + introspection.AccountID, in.existingProfile.AccountID) } } @@ -660,7 +701,7 @@ func discoveryLogin(ctx context.Context, dc discoveryClient, tokenCache cache.To "serverless_compute_id", ) err = databrickscfg.SaveToProfile(ctx, &config.Config{ - Profile: profileName, + Profile: in.profileName, Host: discoveredHost, AuthType: authTypeDatabricksCLI, AccountID: accountID, @@ -670,12 +711,12 @@ func discoveryLogin(ctx context.Context, dc discoveryClient, tokenCache cache.To }, clearKeys...) if err != nil { if configFile != "" { - return fmt.Errorf("saving profile %q to %s: %w", profileName, configFile, err) + return fmt.Errorf("saving profile %q to %s: %w", in.profileName, configFile, err) } - return fmt.Errorf("saving profile %q: %w", profileName, err) + return fmt.Errorf("saving profile %q: %w", in.profileName, err) } - cmdio.LogString(ctx, fmt.Sprintf("Profile %s was successfully saved", profileName)) + cmdio.LogString(ctx, fmt.Sprintf("Profile %s was successfully saved", in.profileName)) return nil } diff --git a/cmd/auth/login_test.go b/cmd/auth/login_test.go index 2b8d473f51..b5d8a39f43 100644 --- a/cmd/auth/login_test.go +++ b/cmd/auth/login_test.go @@ -630,7 +630,14 @@ func TestDiscoveryLogin_IntrospectionFailureStillSavesProfile(t *testing.T) { } ctx, _ := cmdio.NewTestContextWithStdout(t.Context()) - err = discoveryLogin(ctx, dc, newTestTokenCache(), "DISCOVERY", time.Second, "all-apis, ,sql,", nil, func(string) error { return nil }) + err = discoveryLogin(ctx, discoveryLoginInputs{ + dc: dc, + profileName: "DISCOVERY", + timeout: time.Second, + scopes: "all-apis, ,sql,", + browserFunc: func(string) error { return nil }, + tokenCache: newTestTokenCache(), + }) require.NoError(t, err) assert.Equal(t, "https://workspace.example.com", dc.introspectHost) @@ -678,7 +685,14 @@ func TestDiscoveryLogin_AccountIDMismatchWarning(t *testing.T) { AccountID: "old-account-id", } - err = discoveryLogin(ctx, dc, newTestTokenCache(), "DISCOVERY", time.Second, "", existingProfile, func(string) error { return nil }) + err = discoveryLogin(ctx, discoveryLoginInputs{ + dc: dc, + profileName: "DISCOVERY", + timeout: time.Second, + existingProfile: existingProfile, + browserFunc: func(string) error { return nil }, + tokenCache: newTestTokenCache(), + }) require.NoError(t, err) // Verify warning about mismatched account IDs was logged. @@ -726,7 +740,14 @@ func TestDiscoveryLogin_NoWarningWhenAccountIDsMatch(t *testing.T) { AccountID: "same-account-id", } - err = discoveryLogin(ctx, dc, newTestTokenCache(), "DISCOVERY", time.Second, "", existingProfile, func(string) error { return nil }) + err = discoveryLogin(ctx, discoveryLoginInputs{ + dc: dc, + profileName: "DISCOVERY", + timeout: time.Second, + existingProfile: existingProfile, + browserFunc: func(string) error { return nil }, + tokenCache: newTestTokenCache(), + }) require.NoError(t, err) // No warning should be logged when account IDs match. @@ -746,7 +767,13 @@ func TestDiscoveryLogin_EmptyDiscoveredHostReturnsError(t *testing.T) { } ctx, _ := cmdio.NewTestContextWithStdout(t.Context()) - err = discoveryLogin(ctx, dc, newTestTokenCache(), "DISCOVERY", time.Second, "", nil, func(string) error { return nil }) + err = discoveryLogin(ctx, discoveryLoginInputs{ + dc: dc, + profileName: "DISCOVERY", + timeout: time.Second, + browserFunc: func(string) error { return nil }, + tokenCache: newTestTokenCache(), + }) require.Error(t, err) assert.Contains(t, err.Error(), "no workspace host was discovered") } @@ -778,7 +805,14 @@ func TestDiscoveryLogin_ReloginPreservesExistingProfileScopes(t *testing.T) { // No --scopes flag (empty string), should fall back to existing profile scopes. ctx, _ := cmdio.NewTestContextWithStdout(t.Context()) - err = discoveryLogin(ctx, dc, newTestTokenCache(), "DISCOVERY", time.Second, "", existingProfile, func(string) error { return nil }) + err = discoveryLogin(ctx, discoveryLoginInputs{ + dc: dc, + profileName: "DISCOVERY", + timeout: time.Second, + existingProfile: existingProfile, + browserFunc: func(string) error { return nil }, + tokenCache: newTestTokenCache(), + }) require.NoError(t, err) savedProfile, err := loadProfileByName(ctx, "DISCOVERY", profile.DefaultProfiler) @@ -815,7 +849,15 @@ func TestDiscoveryLogin_ExplicitScopesOverrideExistingProfile(t *testing.T) { // Explicit --scopes flag should override existing profile scopes. ctx, _ := cmdio.NewTestContextWithStdout(t.Context()) - err = discoveryLogin(ctx, dc, newTestTokenCache(), "DISCOVERY", time.Second, "all-apis", existingProfile, func(string) error { return nil }) + err = discoveryLogin(ctx, discoveryLoginInputs{ + dc: dc, + profileName: "DISCOVERY", + timeout: time.Second, + scopes: "all-apis", + existingProfile: existingProfile, + browserFunc: func(string) error { return nil }, + tokenCache: newTestTokenCache(), + }) require.NoError(t, err) savedProfile, err := loadProfileByName(ctx, "DISCOVERY", profile.DefaultProfiler) @@ -855,7 +897,13 @@ func TestDiscoveryLogin_SPOGHostPopulatesAccountIDFromDiscovery(t *testing.T) { } ctx, _ := cmdio.NewTestContextWithStdout(t.Context()) - err = discoveryLogin(ctx, dc, newTestTokenCache(), "DISCOVERY", time.Second, "", nil, func(string) error { return nil }) + err = discoveryLogin(ctx, discoveryLoginInputs{ + dc: dc, + profileName: "DISCOVERY", + timeout: time.Second, + browserFunc: func(string) error { return nil }, + tokenCache: newTestTokenCache(), + }) require.NoError(t, err) savedProfile, err := loadProfileByName(ctx, "DISCOVERY", profile.DefaultProfiler) @@ -890,7 +938,13 @@ func TestDiscoveryLogin_IntrospectionFallsBackWhenDiscoveryFails(t *testing.T) { } ctx, _ := cmdio.NewTestContextWithStdout(t.Context()) - err = discoveryLogin(ctx, dc, newTestTokenCache(), "DISCOVERY", time.Second, "", nil, func(string) error { return nil }) + err = discoveryLogin(ctx, discoveryLoginInputs{ + dc: dc, + profileName: "DISCOVERY", + timeout: time.Second, + browserFunc: func(string) error { return nil }, + tokenCache: newTestTokenCache(), + }) require.NoError(t, err) savedProfile, err := loadProfileByName(ctx, "DISCOVERY", profile.DefaultProfiler) @@ -939,7 +993,14 @@ auth_type = databricks-cli } ctx, _ := cmdio.NewTestContextWithStdout(t.Context()) - err = discoveryLogin(ctx, dc, newTestTokenCache(), "DISCOVERY", time.Second, "", existingProfile, func(string) error { return nil }) + err = discoveryLogin(ctx, discoveryLoginInputs{ + dc: dc, + profileName: "DISCOVERY", + timeout: time.Second, + existingProfile: existingProfile, + browserFunc: func(string) error { return nil }, + tokenCache: newTestTokenCache(), + }) require.NoError(t, err) savedProfile, err := loadProfileByName(ctx, "DISCOVERY", profile.DefaultProfiler) @@ -989,7 +1050,14 @@ auth_type = databricks-cli } ctx, _ := cmdio.NewTestContextWithStdout(t.Context()) - err = discoveryLogin(ctx, dc, newTestTokenCache(), "DISCOVERY", time.Second, "", existingProfile, func(string) error { return nil }) + err = discoveryLogin(ctx, discoveryLoginInputs{ + dc: dc, + profileName: "DISCOVERY", + timeout: time.Second, + existingProfile: existingProfile, + browserFunc: func(string) error { return nil }, + tokenCache: newTestTokenCache(), + }) require.NoError(t, err) savedProfile, err := loadProfileByName(ctx, "DISCOVERY", profile.DefaultProfiler) diff --git a/cmd/auth/logout.go b/cmd/auth/logout.go index 67829ec169..bdd0f75430 100644 --- a/cmd/auth/logout.go +++ b/cmd/auth/logout.go @@ -133,9 +133,9 @@ to specify it explicitly. profileName = selected } - tokenCache, err := storage.NewFileTokenCache(ctx) + tokenCache, _, err := storage.ResolveCache(ctx, "") if err != nil { - return fmt.Errorf("failed to open token cache, please check if the file version is up-to-date and that the file is not corrupted: %w", err) + return fmt.Errorf("failed to open token cache: %w", err) } return runLogout(ctx, logoutArgs{ diff --git a/cmd/auth/token.go b/cmd/auth/token.go index d82bcb2c9a..da954b6318 100644 --- a/cmd/auth/token.go +++ b/cmd/auth/token.go @@ -79,9 +79,9 @@ and secret is not supported.`, ctx := cmd.Context() profileName := cmd.Flag("profile").Value.String() - tokenCache, err := storage.NewFileTokenCache(ctx) + tokenCache, mode, err := storage.ResolveCache(ctx, "") if err != nil { - return fmt.Errorf("opening token cache: %w", err) + return err } t, err := loadToken(ctx, loadTokenArgs{ @@ -92,6 +92,7 @@ and secret is not supported.`, forceRefresh: forceRefresh, profiler: profile.DefaultProfiler, tokenCache: tokenCache, + mode: mode, persistentAuthOpts: nil, }) if err != nil { @@ -144,6 +145,11 @@ type loadTokenArgs struct { // responsible for construction so that tests can substitute an in-memory cache. tokenCache cache.TokenCache + // mode is the resolved storage mode. When set to StorageModeLegacy, login + // paths mirror freshly minted tokens under the legacy host-based key so + // older SDKs that still look up by host continue to find them. + mode storage.StorageMode + // persistentAuthOpts are the options to pass to the persistent auth client. persistentAuthOpts []u2m.PersistentAuthOption } @@ -195,7 +201,7 @@ func loadToken(ctx context.Context, args loadTokenArgs) (*oauth2.Token, error) { // resolve the target through environment variables or interactive profile selection. if args.profileName == "" && args.authArguments.Host == "" && len(args.args) == 0 { var resolvedProfile string - resolvedProfile, existingProfile, err = resolveNoArgsToken(ctx, args.profiler, args.authArguments, args.tokenCache) + resolvedProfile, existingProfile, err = resolveNoArgsToken(ctx, args.profiler, args.authArguments, args.tokenCache, args.mode) if err != nil { return nil, err } @@ -284,8 +290,7 @@ func loadToken(ctx context.Context, args loadTokenArgs) (*oauth2.Token, error) { if err != nil { return nil, err } - wrappedCache := storage.NewDualWritingTokenCache(args.tokenCache, oauthArgument) - allArgs := append([]u2m.PersistentAuthOption{u2m.WithTokenCache(wrappedCache)}, args.persistentAuthOpts...) + allArgs := append([]u2m.PersistentAuthOption{u2m.WithTokenCache(args.tokenCache)}, args.persistentAuthOpts...) allArgs = append(allArgs, u2m.WithOAuthArgument(oauthArgument)) persistentAuth, err := u2m.NewPersistentAuth(ctx, allArgs...) if err != nil { @@ -327,7 +332,7 @@ func loadToken(ctx context.Context, args loadTokenArgs) (*oauth2.Token, error) { // // Returns the resolved profile name and profile (if any). The host and related // fields on authArgs are updated in place when resolved via environment variables. -func resolveNoArgsToken(ctx context.Context, profiler profile.Profiler, authArgs *auth.AuthArguments, tokenCache cache.TokenCache) (string, *profile.Profile, error) { +func resolveNoArgsToken(ctx context.Context, profiler profile.Profiler, authArgs *auth.AuthArguments, tokenCache cache.TokenCache, mode storage.StorageMode) (string, *profile.Profile, error) { // Step 1: Try DATABRICKS_HOST env var (highest priority). if envHost := env.Get(ctx, "DATABRICKS_HOST"); envHost != "" { authArgs.Host = envHost @@ -376,7 +381,7 @@ func resolveNoArgsToken(ctx context.Context, profiler profile.Profiler, authArgs // Fall through — setHostAndAccountId will prompt for the host. return "", nil, nil case createNewSelected: - return runInlineLogin(ctx, profiler, tokenCache) + return runInlineLogin(ctx, profiler, tokenCache, mode) default: p, err := loadProfileByName(ctx, selectedName, profiler) if err != nil { @@ -440,7 +445,7 @@ func promptForProfileSelection(ctx context.Context, profiles profile.Profiles) ( // runInlineLogin runs a minimal interactive login flow: prompts for a profile // name and host, performs the OAuth challenge, saves the profile to // .databrickscfg, and returns the new profile name and profile. -func runInlineLogin(ctx context.Context, profiler profile.Profiler, tokenCache cache.TokenCache) (string, *profile.Profile, error) { +func runInlineLogin(ctx context.Context, profiler profile.Profiler, tokenCache cache.TokenCache, mode storage.StorageMode) (string, *profile.Profile, error) { profileName, err := promptForProfile(ctx, "DEFAULT") if err != nil { return "", nil, err @@ -473,9 +478,9 @@ func runInlineLogin(ctx context.Context, profiler profile.Profiler, tokenCache c return "", nil, err } persistentAuthOpts := []u2m.PersistentAuthOption{ - u2m.WithTokenCache(storage.NewDualWritingTokenCache(tokenCache, oauthArgument)), u2m.WithOAuthArgument(oauthArgument), u2m.WithBrowser(func(url string) error { return browser.Open(ctx, url) }), + u2m.WithTokenCache(tokenCache), } if len(scopesList) > 0 { persistentAuthOpts = append(persistentAuthOpts, u2m.WithScopes(scopesList)) @@ -492,6 +497,7 @@ func runInlineLogin(ctx context.Context, profiler profile.Profiler, tokenCache c if err = persistentAuth.Challenge(); err != nil { return "", nil, err } + dualWriteLegacyHostKey(ctx, tokenCache, oauthArgument, mode) clearKeys := oauthLoginClearKeys() if !loginArgs.IsUnifiedHost { diff --git a/libs/auth/credentials.go b/libs/auth/credentials.go index a406955dd3..6b95177353 100644 --- a/libs/auth/credentials.go +++ b/libs/auth/credentials.go @@ -3,7 +3,6 @@ package auth import ( "context" "errors" - "fmt" "github.com/databricks/cli/libs/auth/storage" "github.com/databricks/databricks-sdk-go/config" @@ -99,7 +98,19 @@ func (c CLICredentials) Configure(ctx context.Context, cfg *config.Config) (cred if err != nil { return nil, err } - ts, err := c.persistentAuth(ctx, oauthArg) + // Without WithTokenCache, u2m.NewPersistentAuth falls back to the SDK's + // default file cache. For secure-storage users that would split tokens + // across two backends: login writes to the keyring, but every workspace + // client built through this strategy would read an empty file cache and + // fail with "cache: token not found". + tokenCache, _, err := storage.ResolveCache(ctx, "") + if err != nil { + return nil, err + } + ts, err := c.persistentAuth(ctx, + u2m.WithOAuthArgument(oauthArg), + u2m.WithTokenCache(tokenCache), + ) if err != nil { return nil, err } @@ -109,22 +120,17 @@ func (c CLICredentials) Configure(ctx context.Context, cfg *config.Config) (cred return cp, nil } -// persistentAuth returns a token source. It wraps the file-backed token -// cache with a dual-writing cache so every token write (Challenge, refresh, -// discovery) mirrors to the legacy host key for cross-SDK compatibility. -// The persistentAuthFn override is used in tests. -func (c CLICredentials) persistentAuth(ctx context.Context, arg u2m.OAuthArgument) (auth.TokenSource, error) { +// persistentAuth returns a token source. It is a convenience function that +// overrides the default implementation of the persistent auth client if +// an alternative implementation is provided for testing. The caller is +// responsible for supplying the token cache via u2m.WithTokenCache; Configure +// does this via storage.ResolveCache so login, refresh, and all workspace +// clients share the same backend. +func (c CLICredentials) persistentAuth(ctx context.Context, opts ...u2m.PersistentAuthOption) (auth.TokenSource, error) { if c.persistentAuthFn != nil { - return c.persistentAuthFn(ctx, u2m.WithOAuthArgument(arg)) + return c.persistentAuthFn(ctx, opts...) } - tc, err := storage.NewFileTokenCache(ctx) - if err != nil { - return nil, fmt.Errorf("opening token cache: %w", err) - } - ts, err := u2m.NewPersistentAuth(ctx, - u2m.WithTokenCache(storage.NewDualWritingTokenCache(tc, arg)), - u2m.WithOAuthArgument(arg), - ) + ts, err := u2m.NewPersistentAuth(ctx, opts...) if err != nil { return nil, err } diff --git a/libs/auth/credentials_test.go b/libs/auth/credentials_test.go index 1bc70b63ab..9412d16418 100644 --- a/libs/auth/credentials_test.go +++ b/libs/auth/credentials_test.go @@ -4,15 +4,28 @@ import ( "context" "errors" "net/http" + "os" + "path/filepath" "slices" "testing" + "github.com/databricks/cli/libs/auth/storage" "github.com/databricks/databricks-sdk-go/config" "github.com/databricks/databricks-sdk-go/config/experimental/auth" "github.com/databricks/databricks-sdk-go/credentials/u2m" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" "golang.org/x/oauth2" ) +// hermeticAuthStorage isolates the test from the caller's real env vars and +// .databrickscfg so storage.ResolveCache sees a clean default. +func hermeticAuthStorage(t *testing.T) { + t.Helper() + t.Setenv(storage.EnvVar, "") + t.Setenv("DATABRICKS_CONFIG_FILE", filepath.Join(t.TempDir(), "databrickscfg")) +} + // TestCredentialChainOrder purely exists as an extra measure to catch // accidental change in the ordering. func TestCredentialChainOrder(t *testing.T) { @@ -163,6 +176,7 @@ func TestCLICredentialsConfigure(t *testing.T) { for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { + hermeticAuthStorage(t) ctx := t.Context() c := CLICredentials{persistentAuthFn: tt.persistentAuthFn} @@ -190,3 +204,86 @@ func TestCLICredentialsConfigure(t *testing.T) { }) } } + +// TestCLICredentialsConfigure_ThreadsResolvedTokenCache guards against a +// regression where Configure forgot to pass u2m.WithTokenCache. Without it, +// the SDK's NewPersistentAuth silently defaulted to the file cache, so users +// who opted into secure storage saw "cache: token not found" on every command +// other than auth login/token/logout. +func TestCLICredentialsConfigure_ThreadsResolvedTokenCache(t *testing.T) { + hermeticAuthStorage(t) + + var receivedOpts []u2m.PersistentAuthOption + c := CLICredentials{ + persistentAuthFn: func(_ context.Context, opts ...u2m.PersistentAuthOption) (auth.TokenSource, error) { + receivedOpts = opts + return auth.TokenSourceFn(func(_ context.Context) (*oauth2.Token, error) { + return &oauth2.Token{AccessToken: "tok"}, nil + }), nil + }, + } + + _, err := c.Configure(t.Context(), &config.Config{Host: "https://x.cloud.databricks.com"}) + require.NoError(t, err) + + // Two opts expected: WithOAuthArgument and WithTokenCache. The length + // check is the most resilient way to assert both were passed without + // poking at u2m's unexported state. + assert.Len(t, receivedOpts, 2) +} + +// TestCLICredentialsConfigure_PropagatesStorageResolutionError confirms +// Configure surfaces invalid DATABRICKS_AUTH_STORAGE values instead of +// silently falling back to the file cache. If Configure ever stops calling +// storage.ResolveCache, this test will catch it. +func TestCLICredentialsConfigure_PropagatesStorageResolutionError(t *testing.T) { + hermeticAuthStorage(t) + t.Setenv(storage.EnvVar, "bogus") + + c := CLICredentials{ + persistentAuthFn: func(_ context.Context, _ ...u2m.PersistentAuthOption) (auth.TokenSource, error) { + t.Fatal("persistentAuthFn must not be called when cache resolution fails") + return nil, nil + }, + } + + _, err := c.Configure(t.Context(), &config.Config{Host: "https://x.cloud.databricks.com"}) + require.Error(t, err) + assert.Contains(t, err.Error(), "DATABRICKS_AUTH_STORAGE") +} + +// Writing a throwaway config file is verbose enough that future tests may +// want it too. Keeping the helper scoped here so it stays close to use. +func writeAuthStorageConfig(t *testing.T, mode string) { + t.Helper() + dir := t.TempDir() + configPath := filepath.Join(dir, "databrickscfg") + body := "[__settings__]\nauth_storage = " + mode + "\n" + require.NoError(t, os.WriteFile(configPath, []byte(body), 0o600)) + t.Setenv("DATABRICKS_CONFIG_FILE", configPath) + t.Setenv(storage.EnvVar, "") +} + +// TestCLICredentialsConfigure_HonorsConfigFileSecureMode proves that +// Configure picks up auth_storage = secure from .databrickscfg, not just +// from DATABRICKS_AUTH_STORAGE. Both sources flow through the same resolver, +// but the PR's user-facing docs promise both work and nothing was asserting +// that for this call site. +func TestCLICredentialsConfigure_HonorsConfigFileSecureMode(t *testing.T) { + writeAuthStorageConfig(t, "secure") + + c := CLICredentials{ + persistentAuthFn: func(_ context.Context, opts ...u2m.PersistentAuthOption) (auth.TokenSource, error) { + // The presence of the second opt is verified by the sibling + // test; here we just need Configure to succeed end-to-end when + // the config file selects secure storage. + assert.Len(t, opts, 2) + return auth.TokenSourceFn(func(_ context.Context) (*oauth2.Token, error) { + return &oauth2.Token{AccessToken: "tok"}, nil + }), nil + }, + } + + _, err := c.Configure(t.Context(), &config.Config{Host: "https://x.cloud.databricks.com"}) + require.NoError(t, err) +} diff --git a/libs/auth/storage/cache.go b/libs/auth/storage/cache.go new file mode 100644 index 0000000000..7a8bb775ab --- /dev/null +++ b/libs/auth/storage/cache.go @@ -0,0 +1,61 @@ +package storage + +import ( + "context" + "fmt" + + "github.com/databricks/databricks-sdk-go/credentials/u2m/cache" +) + +// cacheFactories bundles the constructors ResolveCache depends on. Extracted +// so unit tests can inject stubs without hitting the real OS keyring or +// filesystem. Production code uses defaultCacheFactories(). +type cacheFactories struct { + newFile func(context.Context) (cache.TokenCache, error) + newKeyring func() cache.TokenCache +} + +// defaultCacheFactories returns the production factory set. +func defaultCacheFactories() cacheFactories { + return cacheFactories{ + newFile: func(ctx context.Context) (cache.TokenCache, error) { return NewFileTokenCache(ctx) }, + newKeyring: NewKeyringCache, + } +} + +// ResolveCache resolves the storage mode for this invocation and returns +// the corresponding token cache plus the resolved mode (so callers can log +// or surface it). +// +// override is usually the command-level flag value. Pass "" when the command +// has no flag; precedence then falls through to env -> config -> default. +// +// Every CLI code path that calls u2m.NewPersistentAuth must route the result +// through u2m.WithTokenCache, otherwise the SDK defaults to the file cache +// and splits the user's tokens across two backends. +func ResolveCache(ctx context.Context, override StorageMode) (cache.TokenCache, StorageMode, error) { + return resolveCacheWith(ctx, override, defaultCacheFactories()) +} + +// resolveCacheWith is the pure form of ResolveCache. It takes the factory +// set as a parameter so tests can inject stubs. +func resolveCacheWith(ctx context.Context, override StorageMode, f cacheFactories) (cache.TokenCache, StorageMode, error) { + mode, err := ResolveStorageMode(ctx, override) + if err != nil { + return nil, "", err + } + switch mode { + case StorageModeSecure: + return f.newKeyring(), mode, nil + case StorageModeLegacy, StorageModePlaintext: + // Plaintext currently maps to the file cache; a dedicated + // plaintext backend (no host-keyed dual-writes) is a follow-up. + c, err := f.newFile(ctx) + if err != nil { + return nil, "", fmt.Errorf("open file token cache: %w", err) + } + return c, mode, nil + default: + return nil, "", fmt.Errorf("unsupported storage mode %q", string(mode)) + } +} diff --git a/libs/auth/storage/cache_test.go b/libs/auth/storage/cache_test.go new file mode 100644 index 0000000000..6a967891b9 --- /dev/null +++ b/libs/auth/storage/cache_test.go @@ -0,0 +1,116 @@ +package storage + +import ( + "context" + "errors" + "path/filepath" + "testing" + + "github.com/databricks/cli/libs/env" + "github.com/databricks/databricks-sdk-go/credentials/u2m/cache" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "golang.org/x/oauth2" +) + +// stubCache is a test double for cache.TokenCache that records the source +// it was constructed from. It lets the tests confirm which factory ran. +type stubCache struct{ source string } + +func (stubCache) Store(string, *oauth2.Token) error { return nil } +func (stubCache) Lookup(string) (*oauth2.Token, error) { return nil, cache.ErrNotFound } + +func fakeFactories(t *testing.T) cacheFactories { + t.Helper() + return cacheFactories{ + newFile: func(context.Context) (cache.TokenCache, error) { return stubCache{source: "file"}, nil }, + newKeyring: func() cache.TokenCache { return stubCache{source: "keyring"} }, + } +} + +// hermetic isolates the test from the caller's real env vars and +// .databrickscfg so ResolveStorageMode starts from a clean default. +func hermetic(t *testing.T) { + t.Helper() + t.Setenv(EnvVar, "") + t.Setenv("DATABRICKS_CONFIG_FILE", filepath.Join(t.TempDir(), "databrickscfg")) +} + +func TestResolveCache_DefaultsToLegacyFile(t *testing.T) { + hermetic(t) + ctx := t.Context() + + got, mode, err := resolveCacheWith(ctx, "", fakeFactories(t)) + + require.NoError(t, err) + assert.Equal(t, StorageModeLegacy, mode) + assert.Equal(t, "file", got.(stubCache).source) +} + +func TestResolveCache_OverrideSecureUsesKeyring(t *testing.T) { + hermetic(t) + ctx := t.Context() + + got, mode, err := resolveCacheWith(ctx, StorageModeSecure, fakeFactories(t)) + + require.NoError(t, err) + assert.Equal(t, StorageModeSecure, mode) + assert.Equal(t, "keyring", got.(stubCache).source) +} + +func TestResolveCache_EnvVarSelectsSecure(t *testing.T) { + hermetic(t) + ctx := env.Set(t.Context(), EnvVar, "secure") + + got, mode, err := resolveCacheWith(ctx, "", fakeFactories(t)) + + require.NoError(t, err) + assert.Equal(t, StorageModeSecure, mode) + assert.Equal(t, "keyring", got.(stubCache).source) +} + +func TestResolveCache_PlaintextFallsBackToFile(t *testing.T) { + hermetic(t) + ctx := t.Context() + + got, mode, err := resolveCacheWith(ctx, StorageModePlaintext, fakeFactories(t)) + + require.NoError(t, err) + assert.Equal(t, StorageModePlaintext, mode) + assert.Equal(t, "file", got.(stubCache).source) +} + +func TestResolveCache_InvalidOverrideReturnsError(t *testing.T) { + hermetic(t) + ctx := t.Context() + + _, _, err := resolveCacheWith(ctx, StorageMode("bogus"), fakeFactories(t)) + + require.Error(t, err) + assert.Contains(t, err.Error(), `unknown storage mode "bogus"`) +} + +func TestResolveCache_InvalidEnvReturnsError(t *testing.T) { + hermetic(t) + ctx := env.Set(t.Context(), EnvVar, "bogus") + + _, _, err := resolveCacheWith(ctx, "", fakeFactories(t)) + + require.Error(t, err) + assert.Contains(t, err.Error(), "DATABRICKS_AUTH_STORAGE") +} + +func TestResolveCache_FileFactoryErrorPropagates(t *testing.T) { + hermetic(t) + ctx := t.Context() + boom := errors.New("disk full") + factories := cacheFactories{ + newFile: func(context.Context) (cache.TokenCache, error) { return nil, boom }, + newKeyring: func() cache.TokenCache { return stubCache{source: "keyring"} }, + } + + _, _, err := resolveCacheWith(ctx, StorageModeLegacy, factories) + + require.Error(t, err) + assert.ErrorIs(t, err, boom) +}