diff --git a/.nextchanges/cli/file-backed-oauth-cache-transactions.md b/.nextchanges/cli/file-backed-oauth-cache-transactions.md new file mode 100644 index 00000000000..ba143e26d69 --- /dev/null +++ b/.nextchanges/cli/file-backed-oauth-cache-transactions.md @@ -0,0 +1 @@ +* Fix concurrent file-backed OAuth cache updates and prevent automatic reuse of uncertain refresh credentials. ([#7396](https://github.com/databricks/cli/pull/7396)) diff --git a/acceptance/bundle/deploy/mlops-stacks/script b/acceptance/bundle/deploy/mlops-stacks/script index a1770334913..72c45924559 100644 --- a/acceptance/bundle/deploy/mlops-stacks/script +++ b/acceptance/bundle/deploy/mlops-stacks/script @@ -36,7 +36,7 @@ if [ -z "$CLOUD_ENV" ]; then else trace $CLI bundle init mlops-stacks --config-file config.json fi -trace cat test_repo_mlops_stacks/README.md | head -n 4 +trace cat test_repo_mlops_stacks/README.md | sed -n '1,4p' cd "test_repo_mlops_stacks/project_name_${UNIQUE_NAME}" || exit 1 diff --git a/acceptance/cmd/auth/docker/invalid-refresh-token/output.txt b/acceptance/cmd/auth/docker/invalid-refresh-token/output.txt index a206e114758..e27c711aa98 100644 --- a/acceptance/cmd/auth/docker/invalid-refresh-token/output.txt +++ b/acceptance/cmd/auth/docker/invalid-refresh-token/output.txt @@ -6,21 +6,21 @@ Error: A new access token could not be retrieved because the refresh token is in $ databricks auth login --profile docker-test >>> musterr [CLI] auth docker configure docker-test -Error: A new access token could not be retrieved because the refresh token is invalid. To reauthenticate, run the following command: +Error: A previous OAuth refresh has an uncertain outcome; automatic refresh is blocked. To reauthenticate, run the following command: $ databricks auth login --profile docker-test === Stored workspace ID >>> musterr [CLI] auth docker host --profile docker-test -Error: A new access token could not be retrieved because the refresh token is invalid. To reauthenticate, run the following command: +Error: A previous OAuth refresh has an uncertain outcome; automatic refresh is blocked. To reauthenticate, run the following command: $ databricks auth login --profile docker-test >>> musterr [CLI] auth docker configure docker-test -Error: A new access token could not be retrieved because the refresh token is invalid. To reauthenticate, run the following command: +Error: A previous OAuth refresh has an uncertain outcome; automatic refresh is blocked. To reauthenticate, run the following command: $ databricks auth login --profile docker-test >>> musterr [CLI] auth docker token -Error: A new access token could not be retrieved because the refresh token is invalid. To reauthenticate, run the following command: +Error: A previous OAuth refresh has an uncertain outcome; automatic refresh is blocked. To reauthenticate, run the following command: $ databricks auth login --profile docker-test === Docker config not created diff --git a/acceptance/cmd/auth/login/device-code/non-interactive/nominal/output.txt b/acceptance/cmd/auth/login/device-code/non-interactive/nominal/output.txt index 234e8e2ab3a..d308f2b1d25 100644 --- a/acceptance/cmd/auth/login/device-code/non-interactive/nominal/output.txt +++ b/acceptance/cmd/auth/login/device-code/non-interactive/nominal/output.txt @@ -16,6 +16,7 @@ provided in your browser. >>> find.py home home/.databricks/token-cache.json +home/.databricks/token-cache.json.lock >>> cache_keys [ diff --git a/acceptance/cmd/auth/token/force-refresh-invalid-refresh-token/output.txt b/acceptance/cmd/auth/token/force-refresh-invalid-refresh-token/output.txt index d975fe88448..16cf71ecf32 100644 --- a/acceptance/cmd/auth/token/force-refresh-invalid-refresh-token/output.txt +++ b/acceptance/cmd/auth/token/force-refresh-invalid-refresh-token/output.txt @@ -2,5 +2,5 @@ Error: A new access token could not be retrieved because the refresh token is in $ databricks auth login --profile test-profile { "error_code": "UNAUTHENTICATED", - "message": "A new access token could not be retrieved because the refresh token is invalid. To reauthenticate, run the following command:\n $ databricks auth login --profile test-profile" + "message": "A previous OAuth refresh has an uncertain outcome; automatic refresh is blocked. To reauthenticate, run the following command:\n $ databricks auth login --profile test-profile" } diff --git a/acceptance/cmd/completion/script b/acceptance/cmd/completion/script index d315d83839a..c22a95939ac 100644 --- a/acceptance/cmd/completion/script +++ b/acceptance/cmd/completion/script @@ -22,10 +22,10 @@ trace $CLI completion uninstall --shell zsh --auto-approve trace $CLI completion status --shell zsh # Test shell subcommands produce output -$CLI completion bash 2>&1 | head -1 -$CLI completion zsh 2>&1 | head -1 -$CLI completion fish 2>&1 | head -1 -$CLI completion powershell 2>&1 | head -1 +$CLI completion bash 2>&1 | sed -n '1p' +$CLI completion zsh 2>&1 | sed -n '1p' +$CLI completion fish 2>&1 | sed -n '1p' +$CLI completion powershell 2>&1 | sed -n '1p' # Test bash --no-descriptions -$CLI completion bash --no-descriptions 2>&1 | head -1 +$CLI completion bash --no-descriptions 2>&1 | sed -n '1p' diff --git a/bundle/bundle_test.go b/bundle/bundle_test.go index 69c8fa5bdcb..b1134429dc2 100644 --- a/bundle/bundle_test.go +++ b/bundle/bundle_test.go @@ -8,6 +8,7 @@ import ( "github.com/databricks/cli/bundle/config" "github.com/databricks/cli/bundle/config/resources" "github.com/databricks/cli/bundle/env" + "github.com/databricks/cli/internal/testutil" "github.com/databricks/cli/libs/diag" "github.com/databricks/cli/libs/dyn" "github.com/databricks/cli/libs/logdiag" @@ -178,6 +179,10 @@ func TestBundleGetResourceConfigJobsPointer(t *testing.T) { } func TestClearWorkspaceClient(t *testing.T) { + configPath := filepath.Join(t.TempDir(), "databrickscfg") + testutil.WriteFile(t, configPath, "") + t.Setenv("DATABRICKS_CONFIG_FILE", configPath) + ctx := t.Context() // First attempt: profile "profile-A" doesn't exist → error mentions "profile-A". diff --git a/bundle/config/workspace_test.go b/bundle/config/workspace_test.go index 3fb86da2b64..e2268c70782 100644 --- a/bundle/config/workspace_test.go +++ b/bundle/config/workspace_test.go @@ -1,9 +1,12 @@ package config import ( + "context" + "errors" "io/fs" "path/filepath" "runtime" + "slices" "testing" "github.com/databricks/cli/bundle/env" @@ -14,6 +17,21 @@ import ( "github.com/stretchr/testify/require" ) +func stubUnavailableWorkspaceMetadata(t *testing.T, hosts ...string) { + t.Helper() + previous := config.DefaultHostMetadataResolverFactory + config.DefaultHostMetadataResolverFactory = func(*config.Config) config.HostMetadataResolver { + return func(_ context.Context, host string) (*config.HostMetadata, error) { + if !slices.Contains(hosts, host) { + t.Error("unexpected fixture metadata host") + return nil, errors.New("unexpected fixture metadata host") + } + return nil, nil + } + } + t.Cleanup(func() { config.DefaultHostMetadataResolverFactory = previous }) +} + func setupWorkspaceTest(t *testing.T) string { testutil.CleanupEnvironment(t) @@ -27,6 +45,7 @@ func setupWorkspaceTest(t *testing.T) string { } func TestWorkspaceResolveProfileFromHost(t *testing.T) { + stubUnavailableWorkspaceMetadata(t, "https://abc.cloud.databricks.com") // If only a workspace host is specified, try to find a profile that uses // the same workspace host (unambiguously). w := Workspace{ @@ -125,6 +144,7 @@ func TestWorkspaceNormalizeHostURL(t *testing.T) { } func TestWorkspaceClientNormalizesHostBeforeProfileResolution(t *testing.T) { + stubUnavailableWorkspaceMetadata(t, "https://spog.databricks.com") // Regression test: Client() must normalize the host URL (strip ?o= and // populate WorkspaceID) before building the SDK config and resolving // profiles. This ensures workspace_id is available for disambiguation. @@ -156,6 +176,7 @@ func TestWorkspaceClientNormalizesHostBeforeProfileResolution(t *testing.T) { } func TestWorkspaceClientProfileOverridesAuthEnv(t *testing.T) { + stubUnavailableWorkspaceMetadata(t, "https://tst.cloud.databricks.test") // An explicit profile must win over auth env vars (#5096). setupWorkspaceTest(t) @@ -179,6 +200,7 @@ func TestWorkspaceClientProfileOverridesAuthEnv(t *testing.T) { } func TestWorkspaceClientProfileFillsAuthFromEnv(t *testing.T) { + stubUnavailableWorkspaceMetadata(t, "https://tst.cloud.databricks.test") // Host-only profile: the profile wins for the host, but env fills the token // it omits (#5096). setupWorkspaceTest(t) @@ -201,6 +223,7 @@ func TestWorkspaceClientProfileFillsAuthFromEnv(t *testing.T) { } func TestWorkspaceClientHostAndProfileOverridesAuthEnv(t *testing.T) { + stubUnavailableWorkspaceMetadata(t, "https://tst.cloud.databricks.test") // Bundle pins both workspace.host and workspace.profile: the profile wins // for auth and the host check passes because they agree (#5096). setupWorkspaceTest(t) @@ -247,6 +270,7 @@ func TestWorkspaceConfigHTTPTimeout(t *testing.T) { } func TestWorkspaceVerifyProfileForHost(t *testing.T) { + stubUnavailableWorkspaceMetadata(t, "https://abc.cloud.databricks.com", "https://def.cloud.databricks.com") // If both a workspace host and a profile are specified, // verify that the host configured in the profile matches // the host configured in the bundle configuration. diff --git a/cmd/auth/auth_test.go b/cmd/auth/auth_test.go index 10cb67409b9..78a73ddda34 100644 --- a/cmd/auth/auth_test.go +++ b/cmd/auth/auth_test.go @@ -10,6 +10,7 @@ import ( "github.com/databricks/cli/libs/cmdctx" "github.com/databricks/cli/libs/cmdio/cmdiotest/termtest" "github.com/databricks/cli/libs/databrickscfg/profile" + "github.com/spf13/cobra" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" ) @@ -113,13 +114,9 @@ func TestProfileHostConflictTokenViaCobra(t *testing.T) { assert.Contains(t, err.Error(), `--profile "profile-1" has host "https://www.host1.test", which conflicts with --host "https://other.host.test"`) } -// TestProfileHostCompatibleViaCobra verifies that matching --profile and --host -// pass the conflict check (the command will fail later for other reasons, but -// NOT with a conflict error). +// Matching flags reach the stubbed run boundary after the real conflict check. func TestProfileHostCompatibleViaCobra(t *testing.T) { - // Copy the fixture into a temp directory so the auth login flow's writes - // (e.g. silent plaintext-fallback persistence on CI runners without a - // usable keyring) cannot dirty the checked-in fixture. + // Keep Cobra's profile fixture in the owned test directory. configPath := filepath.Join(t.TempDir(), ".databrickscfg") fixture, err := os.ReadFile("./testdata/.databrickscfg") require.NoError(t, err) @@ -129,6 +126,10 @@ func TestProfileHostCompatibleViaCobra(t *testing.T) { ctx := cmdctx.GenerateExecId(t.Context()) cli := root.New(ctx) cli.AddCommand(New()) + login, _, err := cli.Find([]string{"auth", "login"}) + require.NoError(t, err) + reachedRun := false + login.RunE = func(*cobra.Command, []string) error { reachedRun = true; return nil } cli.SetArgs([]string{ "auth", "login", @@ -137,8 +138,7 @@ func TestProfileHostCompatibleViaCobra(t *testing.T) { }) _, err = cli.ExecuteContextC(ctx) - // The command may fail for other reasons (no browser, non-interactive, etc.) - // but it should NOT fail with a conflict error. + assert.True(t, reachedRun) if err != nil { assert.NotContains(t, err.Error(), "conflicts with --host") } diff --git a/cmd/auth/login.go b/cmd/auth/login.go index 067d9e6f7be..f8f02ad325b 100644 --- a/cmd/auth/login.go +++ b/cmd/auth/login.go @@ -126,6 +126,23 @@ func (d *defaultDiscoveryClient) IntrospectToken(ctx context.Context, host, acce func storeLoginToken(ctx context.Context, tokenStore storage.Store, mode storage.StorageMode, arg u2m.OAuthArgument, token *oauth2.Token) error { tokenStore = storage.WrapForOAuthArgument(ctx, tokenStore, mode, arg) key := arg.GetCacheKey() + handled, err := storage.WithFileTransaction(ctx, tokenStore, + func(view storage.Store) error { + entry := storage.Entry{Token: token} + if err := storage.RecordFileLoginResult(view, key, entry); err != nil { + return err + } + return view.Put(key, entry) + }, + func(view storage.Store) error { return view.Delete(key) }, + ) + if handled { + if err != nil { + return fmt.Errorf("store token: %w", err) + } + storage.PinSecureMode(ctx, mode, storage.StorageModeUnknown) + return nil + } if err := tokenStore.Put(key, storage.Entry{Token: token}); err != nil { storeErr := fmt.Errorf("store token: %w", err) // The profile is already saved, but the old token may still be present. diff --git a/cmd/auth/login_device_code.go b/cmd/auth/login_device_code.go index 4346bf61404..4963e63e520 100644 --- a/cmd/auth/login_device_code.go +++ b/cmd/auth/login_device_code.go @@ -194,7 +194,7 @@ func (l nonInteractiveDeviceCodeLogin) start(ctx context.Context) error { if err != nil { return err } - err = l.tokenStore.Put(l.key, storage.Entry{PendingLogin: &storage.PendingLogin{ + err = storage.PutWithContext(ctx, l.tokenStore, l.key, storage.Entry{PendingLogin: &storage.PendingLogin{ DeviceCode: da.DeviceCode, UserCode: da.UserCode, Host: l.host, @@ -274,7 +274,7 @@ func (l nonInteractiveDeviceCodeLogin) complete(ctx context.Context, userCode, c // deletePendingLogin deletes the pending login. A pending login that is left // over only fails a rerun of the second step, so an error is only logged. func (l nonInteractiveDeviceCodeLogin) deletePendingLogin(ctx context.Context) { - if err := l.tokenStore.Delete(l.key); err != nil { + if err := storage.DeleteWithContext(ctx, l.tokenStore, l.key); err != nil { log.Debugf(ctx, "Failed to delete the pending device code login: %v", err) } } diff --git a/cmd/auth/login_device_code_test.go b/cmd/auth/login_device_code_test.go index 01cae5fc2d4..9c83036abfa 100644 --- a/cmd/auth/login_device_code_test.go +++ b/cmd/auth/login_device_code_test.go @@ -241,6 +241,9 @@ func TestIsDeviceCodeSupportedForHost(t *testing.T) { } for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { + if tt.args.Host == "https://accounts.cloud.databricks.test" { + stubUnavailableFixtureMetadata(t, []string{tt.args.Host}) + } args := tt.args require.NoError(t, setHost(t.Context(), nil, &args, nil)) assert.Equal(t, tt.want, isDeviceCodeSupportedForHost(args.Host, args.DiscoveryURL)) diff --git a/cmd/auth/login_intent_test.go b/cmd/auth/login_intent_test.go new file mode 100644 index 00000000000..dbd5ccc2b47 --- /dev/null +++ b/cmd/auth/login_intent_test.go @@ -0,0 +1,70 @@ +package auth + +import ( + "crypto/sha256" + "errors" + "path/filepath" + "testing" + "time" + + "github.com/databricks/cli/libs/auth/storage" + "github.com/databricks/cli/libs/auth/u2m" + "golang.org/x/oauth2" +) + +func TestStoreLoginAccessOnlyReplacesJournalBoundGeneration(t *testing.T) { + for _, fail := range []bool{false, true} { + t.Run(map[bool]string{false: "publish", true: "failed-commit-cleanup"}[fail], func(t *testing.T) { + path := filepath.Join(t.TempDir(), "tokens.json") + store, err := storage.NewFileStore(t.Context(), storage.WithFileLocation(path)) + if err != nil { + t.Fatal("store setup failed") + } + const host = "https://workspace.test" + old := storage.Entry{Token: &oauth2.Token{AccessToken: "old", RefreshToken: "synthetic-refresh", Expiry: time.Now().Add(time.Hour)}} + for _, key := range []string{"profile", host, "copy"} { + if err := store.Put(key, old); err != nil { + t.Fatal("seed failed") + } + } + _, err = storage.WithFileTransaction(t.Context(), store, func(view storage.Store) error { + if _, err := storage.BeginFileRefresh(t.Context(), view, "profile", sha256.Sum256([]byte("route"))); err != nil { + return err + } + return errors.New("closed unknown outcome") + }, nil) + if err == nil { + t.Fatal("prepared setup failed") + } + arg, err := u2m.NewProfileWorkspaceOAuthArgument(host, "profile") + if err != nil { + t.Fatal("argument setup failed") + } + subject := store + if fail { + subject = obstructFileCommitStore{Store: store, path: path} + } + err = storeLoginToken(t.Context(), subject, storage.StorageModePlaintext, arg, &oauth2.Token{AccessToken: "access-only"}) + if (err != nil) != fail { + t.Fatal("access-only publication outcome changed") + } + for _, key := range []string{"profile", host} { + entry, lookupErr := store.Lookup(key) + if fail { + if !errors.Is(lookupErr, storage.ErrNotFound) { + t.Fatal("failed publication left stale entry") + } + } else if lookupErr != nil || entry.Token == nil || entry.Token.AccessToken != "access-only" || entry.Token.RefreshToken != "" { + t.Fatal("access-only result changed") + } + } + _, err = storage.WithFileTransaction(t.Context(), store, func(view storage.Store) error { + _, err := storage.BeginFileRefresh(t.Context(), view, "copy", sha256.Sum256([]byte("route"))) + return err + }, nil) + if !errors.Is(err, storage.ErrRefreshRecoveryRequired) { + t.Fatal("replacement reopened copied old generation") + } + }) + } +} diff --git a/cmd/auth/login_test.go b/cmd/auth/login_test.go index 97b8d82cc9f..3830720de65 100644 --- a/cmd/auth/login_test.go +++ b/cmd/auth/login_test.go @@ -73,6 +73,76 @@ func TestStoreLoginTokenDeletesStaleTokenOnFailure(t *testing.T) { assert.ErrorIs(t, err, storage.ErrNotFound) } +type obstructFileCommitStore struct { + storage.Store + path string +} + +func (s obstructFileCommitStore) WithFileTransaction(ctx context.Context, update, cleanup func(storage.Store) error) (bool, error) { + backup := s.path + ".saved" + return storage.WithFileTransaction(ctx, s.Store, func(view storage.Store) error { + if err := update(view); err != nil { + return err + } + if err := os.Rename(s.path, backup); err != nil { + return err + } + return os.Mkdir(s.path, 0o700) + }, func(view storage.Store) error { + if err := os.Remove(s.path); err != nil { + return err + } + if err := os.Rename(backup, s.path); err != nil { + return err + } + return cleanup(view) + }) +} + +func TestStoreLoginTokenFileBackedPublication(t *testing.T) { + for _, failCommit := range []bool{false, true} { + t.Run(map[bool]string{false: "publish", true: "failed-commit-cleanup"}[failCommit], func(t *testing.T) { + path := filepath.Join(t.TempDir(), "tokens.json") + store, err := storage.NewFileStore(t.Context(), storage.WithFileLocation(path)) + require.NoError(t, err) + const host = "https://workspace.test" + for _, key := range []string{"profile", host, "unrelated"} { + require.NoError(t, store.Put(key, storage.Entry{Token: &oauth2.Token{AccessToken: "old"}})) + } + pending := &storage.PendingLogin{DeviceCode: "synthetic-device", UserCode: "synthetic-user", Host: host, Scopes: []string{"scope"}} + require.NoError(t, store.Put("unrelated-pending", storage.Entry{PendingLogin: pending})) + arg, err := u2m.NewProfileWorkspaceOAuthArgument(host, "profile") + require.NoError(t, err) + subject := store + if failCommit { + subject = obstructFileCommitStore{Store: store, path: path} + } + err = storeLoginToken(t.Context(), subject, storage.StorageModePlaintext, arg, &oauth2.Token{AccessToken: "new"}) + if failCommit { + require.Error(t, err) + assert.ErrorContains(t, err, "cache update") + } else { + require.NoError(t, err) + } + for _, key := range []string{"profile", host} { + value, err := store.Lookup(key) + if failCommit { + assert.ErrorIs(t, err, storage.ErrNotFound) + } else { + require.NoError(t, err) + assert.Equal(t, "new", value.Token.AccessToken) + } + } + other, err := store.Lookup("unrelated") + require.NoError(t, err) + assert.Equal(t, "old", other.Token.AccessToken) + other, err = store.Lookup("unrelated-pending") + require.NoError(t, err) + assert.Equal(t, pending, other.PendingLogin) + }) + } +} + // logBuffer is a thread-safe bytes.Buffer for capturing log output in tests. type logBuffer struct { mu sync.Mutex @@ -147,6 +217,7 @@ func (f *fakeDiscoveryClient) IntrospectToken(ctx context.Context, host, accessT } func TestSetHostDoesNotFailWithNoDatabrickscfg(t *testing.T) { + stubUnavailableFixtureMetadata(t, []string{"test"}) ctx := t.Context() ctx = env.Set(ctx, "DATABRICKS_CONFIG_FILE", "./imaginary-file/databrickscfg") @@ -158,6 +229,7 @@ func TestSetHostDoesNotFailWithNoDatabrickscfg(t *testing.T) { } func TestSetHost(t *testing.T) { + stubUnavailableFixtureMetadata(t, []string{"https://www.host-from-flag.test", "https://www.host-from-arg.test", "https://www.host1.test", "https://www.host2.test", "HTTPS://WWW.HOST1.TEST", "www.host1.test:8080"}) var authArguments auth.AuthArguments t.Setenv("DATABRICKS_CONFIG_FILE", "./testdata/.databrickscfg") ctx, _ := cmdio.SetupTest(t.Context(), cmdio.TestOptions{}) @@ -279,6 +351,7 @@ func TestSetHost(t *testing.T) { } func TestSetAccountId(t *testing.T) { + stubUnavailableFixtureMetadata(t, []string{"https://accounts.cloud.databricks.com"}) var authArguments auth.AuthArguments t.Setenv("DATABRICKS_CONFIG_FILE", "./testdata/.databrickscfg") ctx, _ := cmdio.SetupTest(t.Context(), cmdio.TestOptions{}) @@ -307,6 +380,7 @@ func TestSetAccountId(t *testing.T) { } func TestSetWorkspaceIDForUnifiedHost(t *testing.T) { + stubUnavailableFixtureMetadata(t, []string{"https://unified.databricks.com"}) var authArguments auth.AuthArguments t.Setenv("DATABRICKS_CONFIG_FILE", "./testdata/.databrickscfg") ctx, _ := cmdio.SetupTest(t.Context(), cmdio.TestOptions{}) @@ -628,6 +702,7 @@ func TestRunHostDiscovery_NoHost(t *testing.T) { } func TestRunHostDiscovery_ExplicitFieldsNotOverridden(t *testing.T) { + stubUnavailableFixtureMetadata(t, []string{"https://nonexistent.example.com"}) ctx := t.Context() args := &auth.AuthArguments{ Host: "https://nonexistent.example.com", @@ -691,6 +766,7 @@ func TestRunHostDiscovery_ClassicWorkspaceDoesNotSetAccountID(t *testing.T) { } func TestSetHostAndAccountId_WorkspaceIDNoneSentinelInherited(t *testing.T) { + stubUnavailableFixtureMetadata(t, []string{"https://spog.example.com"}) t.Setenv("DATABRICKS_CONFIG_FILE", "./testdata/.databrickscfg") ctx, _ := cmdio.SetupTest(t.Context(), cmdio.TestOptions{}) @@ -802,6 +878,7 @@ func TestShouldPromptWorkspace(t *testing.T) { } func TestSetHostAndAccountId_URLParamsOverrideProfile(t *testing.T) { + stubUnavailableFixtureMetadata(t, []string{"https://unified.databricks.com?o=99999"}) t.Setenv("DATABRICKS_CONFIG_FILE", "./testdata/.databrickscfg") ctx, _ := cmdio.SetupTest(t.Context(), cmdio.TestOptions{}) @@ -936,6 +1013,7 @@ func TestValidateDiscoveryFlagCompatibility(t *testing.T) { } func TestDiscoveryLogin_IntrospectionFailureStillSavesProfile(t *testing.T) { + stubUnavailableFixtureMetadata(t, []string{"https://workspace.example.com"}) tmpDir := t.TempDir() configPath := filepath.Join(tmpDir, ".databrickscfg") err := os.WriteFile(configPath, []byte(""), 0o600) @@ -978,6 +1056,7 @@ func TestDiscoveryLogin_IntrospectionFailureStillSavesProfile(t *testing.T) { } func TestDiscoveryLogin_ProfileSaveFailureDoesNotStoreToken(t *testing.T) { + stubUnavailableFixtureMetadata(t, []string{"https://workspace.example.test"}) tmpDir := t.TempDir() parentFile := filepath.Join(tmpDir, "not-a-directory") require.NoError(t, os.WriteFile(parentFile, nil, 0o600)) @@ -1010,6 +1089,7 @@ func TestDiscoveryLogin_ProfileSaveFailureDoesNotStoreToken(t *testing.T) { } func TestDiscoveryLogin_AccountIDMismatchWarning(t *testing.T) { + stubUnavailableFixtureMetadata(t, []string{"https://workspace.example.com"}) tmpDir := t.TempDir() configPath := filepath.Join(tmpDir, ".databrickscfg") err := os.WriteFile(configPath, []byte(""), 0o600) @@ -1066,6 +1146,7 @@ func TestDiscoveryLogin_AccountIDMismatchWarning(t *testing.T) { } func TestDiscoveryLogin_NoWarningWhenAccountIDsMatch(t *testing.T) { + stubUnavailableFixtureMetadata(t, []string{"https://workspace.example.com"}) tmpDir := t.TempDir() configPath := filepath.Join(tmpDir, ".databrickscfg") err := os.WriteFile(configPath, []byte(""), 0o600) @@ -1136,6 +1217,7 @@ func TestDiscoveryLogin_EmptyDiscoveredHostReturnsError(t *testing.T) { } func TestDiscoveryLogin_ReloginPreservesExistingProfileScopes(t *testing.T) { + stubUnavailableFixtureMetadata(t, []string{"https://workspace.example.com"}) tmpDir := t.TempDir() configPath := filepath.Join(tmpDir, ".databrickscfg") err := os.WriteFile(configPath, []byte(""), 0o600) @@ -1184,6 +1266,7 @@ func TestDiscoveryLogin_ReloginPreservesExistingProfileScopes(t *testing.T) { } func TestDiscoveryLogin_ExplicitClientIDOverridesExistingProfile(t *testing.T) { + stubUnavailableFixtureMetadata(t, []string{"https://workspace.example.com"}) tmpDir := t.TempDir() configPath := filepath.Join(tmpDir, ".databrickscfg") err := os.WriteFile(configPath, []byte(""), 0o600) @@ -1228,6 +1311,7 @@ func TestDiscoveryLogin_ExplicitClientIDOverridesExistingProfile(t *testing.T) { } func TestDiscoveryLogin_ExplicitScopesOverrideExistingProfile(t *testing.T) { + stubUnavailableFixtureMetadata(t, []string{"https://workspace.example.com"}) tmpDir := t.TempDir() configPath := filepath.Join(tmpDir, ".databrickscfg") err := os.WriteFile(configPath, []byte(""), 0o600) @@ -1320,6 +1404,7 @@ func TestDiscoveryLogin_SPOGHostPopulatesAccountIDFromDiscovery(t *testing.T) { } func TestDiscoveryLogin_IntrospectionFallsBackWhenDiscoveryFails(t *testing.T) { + stubUnavailableFixtureMetadata(t, []string{"https://workspace.example.com"}) tmpDir := t.TempDir() configPath := filepath.Join(tmpDir, ".databrickscfg") err := os.WriteFile(configPath, []byte(""), 0o600) @@ -1361,6 +1446,7 @@ func TestDiscoveryLogin_IntrospectionFallsBackWhenDiscoveryFails(t *testing.T) { } func TestDiscoveryLogin_ClearsStaleRoutingFieldsFromUnifiedProfile(t *testing.T) { + stubUnavailableFixtureMetadata(t, []string{"https://new-workspace.example.com"}) tmpDir := t.TempDir() configPath := filepath.Join(tmpDir, ".databrickscfg") @@ -1422,6 +1508,7 @@ auth_type = databricks-cli } func TestDiscoveryLogin_IntrospectionWritesFreshWorkspaceID(t *testing.T) { + stubUnavailableFixtureMetadata(t, []string{"https://new-workspace.example.com"}) tmpDir := t.TempDir() configPath := filepath.Join(tmpDir, ".databrickscfg") @@ -1477,6 +1564,7 @@ auth_type = databricks-cli } func TestDiscoveryLogin_OverridesHostFromEnv(t *testing.T) { + stubUnavailableFixtureMetadata(t, []string{"https://workspace.example.com"}) tmpDir := t.TempDir() configPath := filepath.Join(tmpDir, ".databrickscfg") err := os.WriteFile(configPath, []byte(""), 0o600) diff --git a/cmd/auth/logout.go b/cmd/auth/logout.go index 94813abbf4b..f022660d61d 100644 --- a/cmd/auth/logout.go +++ b/cmd/auth/logout.go @@ -215,7 +215,7 @@ func runLogout(ctx context.Context, args logoutArgs) error { // A device code login started for the profile could still complete // and save a token for it. The profile does not use the token store // otherwise, so an unavailable store does not fail its logout. - if err := args.tokenStore.Delete(pendingLoginKey(matchedProfile.Name)); err != nil { + if err := storage.DeleteWithContext(ctx, args.tokenStore, pendingLoginKey(matchedProfile.Name)); err != nil { log.Debugf(ctx, "Failed to delete the pending device code login: %v", err) } } @@ -278,39 +278,65 @@ func getMatchingProfile(ctx context.Context, profileName string, profiler profil // profiles, the cache key includes the OIDC path // (host/oidc/accounts/). func clearTokenStore(ctx context.Context, p profile.Profile, profiler profile.Profiler, tokenStore storage.Store) error { - if err := tokenStore.Delete(p.Name); err != nil { - return fmt.Errorf("failed to delete profile-keyed token for profile %q: %w", p.Name, err) - } - // A pending device code login could still complete and restore the - // credentials that the user is logging out of. - if err := tokenStore.Delete(pendingLoginKey(p.Name)); err != nil { - return fmt.Errorf("failed to delete pending device code login for profile %q: %w", p.Name, err) - } - - hostCacheKey, matchFn := hostCacheKeyAndMatchFn(p) - if hostCacheKey == "" { - return fmt.Errorf("failed to get host-based cache key for profile %q", p.Name) + deletePrimary := func(view storage.Store) error { + if err := view.Delete(p.Name); err != nil { + return fmt.Errorf("failed to delete profile-keyed token for profile %q: %w", p.Name, err) + } + if err := view.Delete(pendingLoginKey(p.Name)); err != nil { + return fmt.Errorf("failed to delete pending login for profile %q: %w", p.Name, err) + } + return nil } - - // Only preserve the host-keyed token if another U2M profile shares the - // same host. Non-U2M profiles (PAT, M2M, etc.) never use the OAuth - // token cache, so they should not prevent cleanup. - otherProfiles, err := profiler.LoadProfiles(ctx, func(candidate profile.Profile) bool { - return candidate.Name != p.Name && candidate.AuthType == "databricks-cli" && matchFn(candidate) - }) - if err != nil { - return fmt.Errorf("failed to load profiles for host cache key %q: %w", hostCacheKey, err) + var hostKey string + prepareHostCleanup := func() error { + key, matchFn := hostCacheKeyAndMatchFn(p) + if key == "" { + return fmt.Errorf("failed to get host-based cache key for profile %q", p.Name) + } + otherProfiles, err := profiler.LoadProfiles(ctx, func(candidate profile.Profile) bool { + return candidate.Name != p.Name && candidate.AuthType == "databricks-cli" && matchFn(candidate) + }) + if err != nil { + return fmt.Errorf("failed to load profiles for host cache key %q: %w", key, err) + } + if len(otherProfiles) == 0 { + hostKey = key + } + return nil } - - if len(otherProfiles) == 0 { - if err := tokenStore.Delete(hostCacheKey); err != nil { - return fmt.Errorf("failed to delete host-keyed token for %q: %w", hostCacheKey, err) + deleteHost := func(view storage.Store) error { + if hostKey == "" { + return nil + } + if err := view.Delete(hostKey); err != nil { + return fmt.Errorf("failed to delete host-keyed token for %q: %w", hostKey, err) } - if err := tokenStore.Delete(pendingLoginKey(hostCacheKey)); err != nil { - return fmt.Errorf("failed to delete pending device code login for %q: %w", hostCacheKey, err) + if err := view.Delete(pendingLoginKey(hostKey)); err != nil { + return fmt.Errorf("failed to delete pending login for host %q: %w", hostKey, err) } + return nil } - return nil + if storage.SupportsFileTransactions(tokenStore) { + // Resolve routing before holding the file lease; publish all deletes together. + if err := prepareHostCleanup(); err != nil { + return err + } + _, err := storage.WithFileTransaction(ctx, tokenStore, func(view storage.Store) error { + if err := deletePrimary(view); err != nil { + return err + } + return deleteHost(view) + }, nil) + return err + } + // Preserve mutation and error ordering for other backends. + if err := deletePrimary(tokenStore); err != nil { + return err + } + if err := prepareHostCleanup(); err != nil { + return err + } + return deleteHost(tokenStore) } // hostCacheKeyAndMatchFn returns the host-based token cache key and a profile diff --git a/cmd/auth/logout_test.go b/cmd/auth/logout_test.go index 2e2428e6cc9..c28546f58a0 100644 --- a/cmd/auth/logout_test.go +++ b/cmd/auth/logout_test.go @@ -1,6 +1,7 @@ package auth import ( + "context" "encoding/json" "errors" "maps" @@ -10,9 +11,12 @@ import ( "path/filepath" "testing" + "github.com/databricks/cli/libs/auth" + "github.com/databricks/cli/libs/auth/storage" "github.com/databricks/cli/libs/cmdio" "github.com/databricks/cli/libs/databrickscfg" "github.com/databricks/cli/libs/databrickscfg/profile" + "github.com/databricks/databricks-sdk-go/config" "github.com/spf13/cobra" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" @@ -74,6 +78,49 @@ func writeTempConfig(t *testing.T, content string) string { return path } +func stubFixtureHostMetadata(t *testing.T, endpoints map[string]string) { + t.Helper() + previous := config.DefaultHostMetadataResolverFactory + config.DefaultHostMetadataResolverFactory = func(*config.Config) config.HostMetadataResolver { + return func(_ context.Context, host string) (*config.HostMetadata, error) { + endpoint, ok := endpoints[host] + if !ok { + t.Error("unexpected fixture metadata host") + return nil, errors.New("unexpected fixture metadata host") + } + if endpoint == "" { + return nil, nil + } + return &config.HostMetadata{OIDCEndpoint: endpoint}, nil + } + } + t.Cleanup(func() { config.DefaultHostMetadataResolverFactory = previous }) +} + +// These cases test input selection, not a metadata response or HTTP transport. +func stubUnavailableFixtureMetadata(t *testing.T, inputs []string) { + t.Helper() + declared := make(map[string]string, len(inputs)) + for _, input := range inputs { + copy := auth.AuthArguments{Host: input} + require.NoError(t, normalizeLoginHost(©)) + cfg := &config.Config{Host: copy.Host} + declared[cfg.CanonicalHostName()] = "" + } + stubFixtureHostMetadata(t, declared) +} + +func stubLogoutHostMetadata(t *testing.T) { + t.Helper() + stubFixtureHostMetadata(t, map[string]string{ + "https://my-workspace.cloud.databricks.com": "https://my-workspace.cloud.databricks.com/oidc", + "https://my-unique-workspace.cloud.databricks.com": "https://my-unique-workspace.cloud.databricks.com/oidc", + "https://stale-account.cloud.databricks.com": "https://stale-account.cloud.databricks.com/oidc", + "https://accounts.cloud.databricks.com": "https://accounts.cloud.databricks.com/oidc/accounts/{account_id}", + "https://my-m2m.cloud.databricks.com": "https://my-m2m.cloud.databricks.com/oidc", + }) +} + func TestLogout(t *testing.T) { cases := []struct { name string @@ -162,6 +209,7 @@ func TestLogout(t *testing.T) { for _, tc := range cases { t.Run(tc.name, func(t *testing.T) { ctx := cmdio.MockDiscard(t.Context()) + stubLogoutHostMetadata(t) configPath := writeTempConfig(t, logoutTestConfig) t.Setenv("DATABRICKS_CONFIG_FILE", configPath) @@ -211,6 +259,7 @@ func TestLogout(t *testing.T) { func TestLogoutNoTokens(t *testing.T) { ctx := cmdio.MockDiscard(t.Context()) + stubLogoutHostMetadata(t) configPath := writeTempConfig(t, logoutTestConfig) t.Setenv("DATABRICKS_CONFIG_FILE", configPath) @@ -235,6 +284,7 @@ func TestLogoutNoTokens(t *testing.T) { func TestLogoutNoTokensWithDelete(t *testing.T) { ctx := cmdio.MockDiscard(t.Context()) + stubLogoutHostMetadata(t) configPath := writeTempConfig(t, logoutTestConfig) t.Setenv("DATABRICKS_CONFIG_FILE", configPath) @@ -270,6 +320,7 @@ func (failingDeleteStore) Delete(string) error { // an unavailable store does not fail its logout. func TestLogoutNonLoginProfileIgnoresTokenStoreError(t *testing.T) { ctx := cmdio.MockDiscard(t.Context()) + stubLogoutHostMetadata(t) configPath := writeTempConfig(t, logoutTestConfig) t.Setenv("DATABRICKS_CONFIG_FILE", configPath) @@ -299,6 +350,10 @@ func TestLogoutProfileFlagAndPositionalArgConflict(t *testing.T) { } func TestLogoutDeleteClearsDefaultProfile(t *testing.T) { + stubFixtureHostMetadata(t, map[string]string{ + "https://my-workspace.cloud.databricks.com": "https://my-workspace.cloud.databricks.com/oidc", + "https://other-workspace.cloud.databricks.com": "https://other-workspace.cloud.databricks.com/oidc", + }) configWithDefault := `[DEFAULT] [my-workspace] host = https://my-workspace.cloud.databricks.com @@ -412,6 +467,43 @@ auth_type = databricks-cli assert.Nil(t, tokenStore.Tokens[hostKey]) } +func TestClearFileTokenStorePreservesOtherEntries(t *testing.T) { + for _, shared := range []bool{false, true} { + t.Run(map[bool]string{false: "unique-host", true: "shared-host"}[shared], func(t *testing.T) { + server := newWellKnownServer(t, false, "synthetic-account") + owner := profile.Profile{Name: "owner", Host: server.URL, AuthType: "databricks-cli"} + profiles := profile.Profiles{owner} + if shared { + profiles = append(profiles, profile.Profile{Name: "sibling", Host: server.URL, AuthType: "databricks-cli"}) + } + path := filepath.Join(t.TempDir(), "tokens.json") + store, err := storage.NewFileStore(t.Context(), storage.WithFileLocation(path)) + require.NoError(t, err) + for _, key := range []string{owner.Name, server.URL, "unrelated"} { + require.NoError(t, store.Put(key, storage.Entry{Token: &oauth2.Token{AccessToken: "synthetic"}})) + require.NoError(t, store.Put(pendingLoginKey(key), storage.Entry{PendingLogin: &storage.PendingLogin{UserCode: "synthetic"}})) + } + require.NoError(t, clearTokenStore(t.Context(), owner, profile.InMemoryProfiler{Profiles: profiles}, store)) + for _, key := range []string{owner.Name, pendingLoginKey(owner.Name)} { + _, err := store.Lookup(key) + assert.ErrorIs(t, err, storage.ErrNotFound) + } + for _, key := range []string{server.URL, pendingLoginKey(server.URL)} { + _, err := store.Lookup(key) + if shared { + assert.NoError(t, err) + } else { + assert.ErrorIs(t, err, storage.ErrNotFound) + } + } + for _, key := range []string{"unrelated", pendingLoginKey("unrelated")} { + _, err := store.Lookup(key) + assert.NoError(t, err) + } + }) + } +} + func TestHostCacheKeyAndMatchFn(t *testing.T) { wsServer := newWellKnownServer(t, false, "ws-account") spogServer := newWellKnownServer(t, true, "spog-account") @@ -468,6 +560,11 @@ func TestHostCacheKeyAndMatchFn(t *testing.T) { for _, tc := range cases { t.Run(tc.name, func(t *testing.T) { + if tc.profile.Host == "https://accounts.cloud.databricks.com" { + stubFixtureHostMetadata(t, map[string]string{ + tc.profile.Host: tc.profile.Host + "/oidc/accounts/{account_id}", + }) + } key, matchFn := hostCacheKeyAndMatchFn(tc.profile) if tc.wantKeyEmpty { assert.Empty(t, key) diff --git a/cmd/auth/token.go b/cmd/auth/token.go index bbde8dc5eff..ae8a4a44b4a 100644 --- a/cmd/auth/token.go +++ b/cmd/auth/token.go @@ -80,7 +80,7 @@ and secret is not supported.`, }) if err != nil { if cmd.Flag("output").Changed && root.OutputType(cmd) == flags.OutputJSON { - if _, ok := errors.AsType[*u2m.InvalidRefreshTokenError](err); ok { + if _, ok := errors.AsType[*u2m.InvalidRefreshTokenError](err); ok || errors.Is(err, storage.ErrRefreshRecoveryRequired) { if outputErr := writeTokenErrorOutput(cmd.OutOrStdout(), err); outputErr != nil { return outputErr } diff --git a/cmd/auth/token_test.go b/cmd/auth/token_test.go index 001aa527034..db5e64fa0fc 100644 --- a/cmd/auth/token_test.go +++ b/cmd/auth/token_test.go @@ -116,6 +116,18 @@ func (m *MockApiClient) GetEndpointsFromURL(_ context.Context, _ string) (*u2m.O var _ u2m.OAuthEndpointSupplier = (*MockApiClient)(nil) func TestToken_loadToken(t *testing.T) { + stubFixtureHostMetadata(t, map[string]string{ + "https://accounts.cloud.databricks.com": "https://accounts.cloud.databricks.com/oidc/accounts/{account_id}", + "https://dev.cloud.databricks.com": "https://dev.cloud.databricks.com/oidc", + "https://legacy-ws.cloud.databricks.com": "https://legacy-ws.cloud.databricks.com/oidc", + "https://m2m.cloud.databricks.com": "https://m2m.cloud.databricks.com/oidc", + "https://no-profile.cloud.databricks.com": "https://no-profile.cloud.databricks.com/oidc", + "https://nonexistent.cloud.databricks.com": "https://nonexistent.cloud.databricks.com/oidc", + "https://shared.cloud.databricks.com": "https://shared.cloud.databricks.com/oidc", + "https://unique-ws.cloud.databricks.com": "https://unique-ws.cloud.databricks.com/oidc", + "https://valid-token.cloud.databricks.com": "https://valid-token.cloud.databricks.com/oidc", + "https://workspace-a.cloud.databricks.com": "https://workspace-a.cloud.databricks.com/oidc", + }) profiler := profile.InMemoryProfiler{ Profiles: profile.Profiles{ { diff --git a/libs/auth/arguments_test.go b/libs/auth/arguments_test.go index 873ce557dfe..25f95abb3e6 100644 --- a/libs/auth/arguments_test.go +++ b/libs/auth/arguments_test.go @@ -1,17 +1,37 @@ package auth import ( + "context" "encoding/json" + "errors" "net/http" "net/http/httptest" + "slices" "testing" "github.com/databricks/cli/libs/auth/u2m" + "github.com/databricks/databricks-sdk-go/config" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" ) +func stubUnavailableAuthArgumentMetadata(t *testing.T, hosts ...string) { + t.Helper() + previous := config.DefaultHostMetadataResolverFactory + config.DefaultHostMetadataResolverFactory = func(*config.Config) config.HostMetadataResolver { + return func(_ context.Context, host string) (*config.HostMetadata, error) { + if !slices.Contains(hosts, host) { + t.Error("unexpected fixture metadata host") + return nil, errors.New("unexpected fixture metadata host") + } + return nil, nil + } + } + t.Cleanup(func() { config.DefaultHostMetadataResolverFactory = previous }) +} + func TestToOAuthArgument(t *testing.T) { + stubUnavailableAuthArgumentMetadata(t, "https://my-workspace.cloud.databricks.com", "https://accounts.cloud.databricks.com") tests := []struct { name string args AuthArguments diff --git a/libs/auth/credentials_test.go b/libs/auth/credentials_test.go index 2802c9b4279..1525c7f4750 100644 --- a/libs/auth/credentials_test.go +++ b/libs/auth/credentials_test.go @@ -123,6 +123,7 @@ func TestAuthArgumentsFromConfig(t *testing.T) { } func TestCLICredentialsConfigure(t *testing.T) { + stubUnavailableAuthArgumentMetadata(t, "https://myworkspace.cloud.databricks.com", "https://accounts.cloud.databricks.com") testErr := errors.New("test error") tests := []struct { @@ -210,6 +211,7 @@ func TestCLICredentialsConfigure(t *testing.T) { // NewPersistentAuth defaults to an in-memory store, so it cannot see the token // persisted by auth login. func TestCLICredentialsConfigure_ThreadsResolvedTokenStore(t *testing.T) { + stubUnavailableAuthArgumentMetadata(t, "https://x.cloud.databricks.com") hermeticAuthStorage(t) var receivedOpts []u2m.PersistentAuthOption @@ -232,6 +234,7 @@ func TestCLICredentialsConfigure_ThreadsResolvedTokenStore(t *testing.T) { } func TestCLICredentialsConfigure_ClientID(t *testing.T) { + stubUnavailableAuthArgumentMetadata(t, "https://workspace.test") tests := []struct { name string authType string @@ -301,6 +304,7 @@ func TestCLICredentialsConfigure_ClientID(t *testing.T) { // silently falling back to the file cache. If Configure ever stops calling // storage.ResolveStore, this test will catch it. func TestCLICredentialsConfigure_PropagatesStorageResolutionError(t *testing.T) { + stubUnavailableAuthArgumentMetadata(t, "https://x.cloud.databricks.com") hermeticAuthStorage(t) t.Setenv(storage.EnvVar, "bogus") @@ -334,6 +338,7 @@ func writeAuthStorageConfig(t *testing.T, mode string) { // 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) { + stubUnavailableAuthArgumentMetadata(t, "https://x.cloud.databricks.com") writeAuthStorageConfig(t, "secure") c := CLICredentials{ diff --git a/libs/auth/error.go b/libs/auth/error.go index fdb5f256f43..47ea69d58bd 100644 --- a/libs/auth/error.go +++ b/libs/auth/error.go @@ -7,6 +7,7 @@ import ( "net/http" "strings" + "github.com/databricks/cli/libs/auth/storage" "github.com/databricks/cli/libs/auth/u2m" "github.com/databricks/databricks-sdk-go/apierr" "github.com/databricks/databricks-sdk-go/config" @@ -64,7 +65,7 @@ func (e *rewrittenAuthError) Unwrap() error { return e.cause } -// RewriteAuthError rewrites the error message for invalid refresh token error. +// RewriteAuthError adds actionable guidance to closed OAuth refresh failures. // It returns whether the error was rewritten and the rewritten error. func RewriteAuthError(ctx context.Context, host, accountId, profile string, err error) (bool, error) { if _, ok := errors.AsType[*u2m.InvalidRefreshTokenError](err); ok { @@ -76,6 +77,23 @@ func RewriteAuthError(ctx context.Context, host, accountId, profile string, err return false, argErr } msg := `A new access token could not be retrieved because the refresh token is invalid. To reauthenticate, run the following command: + $ ` + BuildLoginCommand(ctx, profile, oauthArgument) + return true, &rewrittenAuthError{message: msg, cause: err} + } + if errors.Is(err, storage.ErrRefreshRecoveryRequired) { + var oauthArgument u2m.OAuthArgument + if profile == "" && host != "" { + if accountId != "" { + if arg, argErr := u2m.NewBasicAccountOAuthArgument(host, accountId); argErr == nil { + oauthArgument = arg + } + } else { + if arg, argErr := u2m.NewBasicWorkspaceOAuthArgument(host); argErr == nil { + oauthArgument = arg + } + } + } + msg := `A previous OAuth refresh has an uncertain outcome; automatic refresh is blocked. To reauthenticate, run the following command: $ ` + BuildLoginCommand(ctx, profile, oauthArgument) return true, &rewrittenAuthError{message: msg, cause: err} } diff --git a/libs/auth/error_intent_test.go b/libs/auth/error_intent_test.go new file mode 100644 index 00000000000..4441cd2854a --- /dev/null +++ b/libs/auth/error_intent_test.go @@ -0,0 +1,58 @@ +package auth + +import ( + "context" + "errors" + "fmt" + "strings" + "testing" + + "github.com/databricks/cli/libs/auth/storage" + "github.com/databricks/databricks-sdk-go/config" + "github.com/stretchr/testify/require" +) + +func TestRewriteRefreshRecoveryPreservesCauseAndGuidance(t *testing.T) { + previous := config.DefaultHostMetadataResolverFactory + metadataCalls := 0 + config.DefaultHostMetadataResolverFactory = func(*config.Config) config.HostMetadataResolver { + metadataCalls++ + return func(context.Context, string) (*config.HostMetadata, error) { + return nil, errors.New("unexpected guidance metadata call") + } + } + t.Cleanup(func() { config.DefaultHostMetadataResolverFactory = previous }) + cases := []struct{ name, host, account, profile, command string }{ + {name: "profile", profile: "test-profile", command: "databricks auth login --profile test-profile"}, + {name: "no-profile", command: "databricks auth login"}, + {name: "unnamed-workspace", host: "https://workspace.example.test", command: "databricks auth login --host https://workspace.example.test"}, + {name: "unnamed-account", host: "https://accounts.example.test", account: "account-fixture", command: "databricks auth login --host https://accounts.example.test --account-id account-fixture"}, + {name: "unnamed-unified-host", host: "https://unified.example.test", account: "account-fixture", command: "databricks auth login --host https://unified.example.test --account-id account-fixture"}, + {name: "profile-precedence", host: "https://workspace.example.test", account: "account-fixture", profile: "test-profile", command: "databricks auth login --profile test-profile"}, + {name: "invalid-workspace-host", host: "invalid-host-sentinel", command: "databricks auth login"}, + {name: "invalid-account-host", host: "https://invalid.example.test/", account: "account-fixture", command: "databricks auth login"}, + {name: "account-without-host", account: "account-fixture", command: "databricks auth login"}, + } + for _, test := range cases { + t.Run(test.name, func(t *testing.T) { + cause := fmt.Errorf("outer: %w", storage.ErrRefreshRecoveryRequired) + rewritten, err := RewriteAuthError(t.Context(), test.host, test.account, test.profile, cause) + if !rewritten || !errors.Is(err, storage.ErrRefreshRecoveryRequired) || !strings.Contains(err.Error(), "uncertain outcome; automatic refresh is blocked") { + t.Fatal("recovery guidance lost its typed cause") + } + _, command, found := strings.Cut(err.Error(), "\n $ ") + require.True(t, found) + require.Equal(t, test.command, command) + if strings.Contains(err.Error(), "refresh token is invalid") || strings.Contains(err.Error(), "report this issue") { + t.Fatal("uncertain refresh was mislabeled") + } + }) + } + require.Zero(t, metadataCalls, "static recovery guidance must not resolve metadata") + cause := errors.New("unrelated failure") + rewritten, err := RewriteAuthError(t.Context(), "", "", "", cause) + if rewritten { + t.Fatal("unrelated error behavior changed") + } + require.Same(t, cause, err) +} diff --git a/libs/auth/error_test.go b/libs/auth/error_test.go index ce8f0df1b63..13fa931571c 100644 --- a/libs/auth/error_test.go +++ b/libs/auth/error_test.go @@ -64,6 +64,7 @@ func TestEnrichAuthError_PreservesOriginalError(t *testing.T) { } func TestEnrichAuthError(t *testing.T) { + stubUnavailableAuthArgumentMetadata(t, "https://accounts.cloud.databricks.com") tests := []struct { name string cfg *config.Config diff --git a/libs/auth/storage/dual_writing_store.go b/libs/auth/storage/dual_writing_store.go index c6e20021313..b476b4528c7 100644 --- a/libs/auth/storage/dual_writing_store.go +++ b/libs/auth/storage/dual_writing_store.go @@ -1,5 +1,7 @@ package storage +import "context" + // TokenKeyProvider supplies the primary key used to store an OAuth token. type TokenKeyProvider interface { GetCacheKey() string @@ -43,6 +45,11 @@ func NewDualWritingStore(inner Store, arg TokenKeyProvider) *DualWritingStore { // whole Put call would break primary login over a non-essential mirror, // so a stale host-key entry is the lesser harm. func (s *DualWritingStore) Put(key string, e Entry) error { + // The contextless API has no request budget; callbacks use staged inner stores. + handled, err := WithFileTransaction(context.Background(), s, func(view Store) error { return view.Put(key, e) }, nil) //nolint:gocritic // Legacy Store interface has no operation context. + if handled { + return err + } if err := s.inner.Put(key, e); err != nil { return err } @@ -66,6 +73,11 @@ func (s *DualWritingStore) Lookup(key string) (Entry, error) { // Delete implements [Store]. Deletes under the primary key are also mirrored // under the host key, using the same best-effort compatibility policy as Put. func (s *DualWritingStore) Delete(key string) error { + // The contextless API has no request budget; callbacks use staged inner stores. + handled, err := WithFileTransaction(context.Background(), s, func(view Store) error { return view.Delete(key) }, nil) //nolint:gocritic // Legacy Store interface has no operation context. + if handled { + return err + } if err := s.inner.Delete(key); err != nil { return err } @@ -81,6 +93,21 @@ func (s *DualWritingStore) Delete(key string) error { return nil } +// WithFileTransaction keeps the primary and legacy mirror in one staged map. +func (s *DualWritingStore) WithFileTransaction(ctx context.Context, update, cleanup func(Store) error) (bool, error) { + wrap := func(mutate func(Store) error) func(Store) error { + if mutate == nil { + return nil + } + return func(view Store) error { + wrapped := *s + wrapped.inner = view + return mutate(&wrapped) + } + } + return WithFileTransaction(ctx, s.inner, wrap(update), wrap(cleanup)) +} + // hostCacheKey mirrors the SDK's former PersistentAuth.hostCacheKey: // discovery arguments expose the host via GetDiscoveredHost (populated by // Challenge); static arguments expose it via HostCacheKeyProvider. diff --git a/libs/auth/storage/file_transaction.go b/libs/auth/storage/file_transaction.go new file mode 100644 index 00000000000..5f94b4dfe9b --- /dev/null +++ b/libs/auth/storage/file_transaction.go @@ -0,0 +1,116 @@ +package storage + +import "context" + +type fileTransactionStore interface { + WithFileTransaction(context.Context, func(Store) error, func(Store) error) (bool, error) +} + +// SupportsFileTransactions reports whether a known wrapper owns a file cache. +func SupportsFileTransactions(store Store) bool { + switch store := store.(type) { + case *fileStore: + return true + case *DualWritingStore: + return SupportsFileTransactions(store.inner) + case *notFoundHintStore: + return SupportsFileTransactions(store.inner) + default: + return false + } +} + +// WithFileTransaction stages file-backed mutations and publishes them atomically. +// A false result leaves callbacks uncalled: other backends retain their behavior. +// Cleanup runs under the same lease only when a successful callback cannot commit. +func WithFileTransaction(ctx context.Context, store Store, update, cleanup func(Store) error) (bool, error) { + owner, ok := store.(fileTransactionStore) + if !ok { + return false, nil + } + return owner.WithFileTransaction(ctx, update, cleanup) +} + +// PutWithContext preserves the caller's context for file-backed mutations. +func PutWithContext(ctx context.Context, store Store, key string, entry Entry) error { + handled, err := WithFileTransaction(ctx, store, func(view Store) error { return view.Put(key, entry) }, nil) + if handled { + return err + } + return store.Put(key, entry) +} + +// DeleteWithContext preserves the caller's context for file-backed mutations. +func DeleteWithContext(ctx context.Context, store Store, key string) error { + handled, err := WithFileTransaction(ctx, store, func(view Store) error { return view.Delete(key) }, nil) + if handled { + return err + } + return store.Delete(key) +} + +type fileTransaction struct { + file *tokenStoreFile + dirty bool + owner *fileStore + ctx context.Context + original *tokenStoreFile + originalRaw []byte + journal *refreshJournal + touched map[string]bool + live bool + attempt *FileRefreshAttempt + beforeBindings map[string]*refreshBinding + beforeGenerations map[string]*refreshGeneration +} + +func (t *fileTransaction) Lookup(key string) (Entry, error) { + if !t.live { + return Entry{}, intentError(errRefreshOperation, nil) + } + entry, ok := t.file.Tokens[key] + if !ok || entry == nil { + return Entry{}, ErrNotFound + } + return cloneEntry(Entry{Token: entry.Token, PendingLogin: entry.PendingLogin}), nil +} + +func (t *fileTransaction) Put(key string, entry Entry) error { + if !t.live { + return intentError(errRefreshOperation, nil) + } + t.touched[key] = true + entry = cloneEntry(entry) + if t.file.Tokens == nil { + t.file.Tokens = make(map[string]*fileEntry) + } + t.file.Tokens[key] = &fileEntry{Token: entry.Token, PendingLogin: entry.PendingLogin} + t.dirty = true + return nil +} + +func (t *fileTransaction) Delete(key string) error { + if !t.live { + return intentError(errRefreshOperation, nil) + } + t.touched[key] = true + if _, ok := t.file.Tokens[key]; ok { + delete(t.file.Tokens, key) + t.dirty = true + } + return nil +} + +func cloneStoreFile(file *tokenStoreFile) *tokenStoreFile { + clone := &tokenStoreFile{Version: file.Version, Tokens: make(map[string]*fileEntry, len(file.Tokens))} + for key, entry := range file.Tokens { + // Preserve JSON null entries accepted by the cache loader. + if entry == nil { + clone.Tokens[key] = nil + continue + } + value := cloneEntry(Entry{Token: entry.Token, PendingLogin: entry.PendingLogin}) + clone.Tokens[key] = &fileEntry{Token: value.Token, PendingLogin: value.PendingLogin} + } + return clone +} diff --git a/libs/auth/storage/file_transaction_test.go b/libs/auth/storage/file_transaction_test.go new file mode 100644 index 00000000000..7e3fb6c6e74 --- /dev/null +++ b/libs/auth/storage/file_transaction_test.go @@ -0,0 +1,387 @@ +package storage_test + +import ( + "bufio" + "context" + "encoding/json" + "errors" + "io" + "os" + "os/exec" + "path/filepath" + "sync" + "testing" + "time" + + "github.com/databricks/cli/libs/auth/storage" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "golang.org/x/oauth2" +) + +func newFileStore(t *testing.T) (storage.Store, string) { + t.Helper() + path := filepath.Join(t.TempDir(), "tokens.json") + store, err := storage.NewFileStore(t.Context(), storage.WithFileLocation(path)) + require.NoError(t, err) + return store, path +} + +func entry(value string) storage.Entry { + return storage.Entry{Token: &oauth2.Token{AccessToken: value}} +} + +func TestFileTransactionPublishesOnlyCommittedState(t *testing.T) { + for _, abort := range []bool{false, true} { + t.Run(map[bool]string{false: "commit", true: "abort"}[abort], func(t *testing.T) { + store, path := newFileStore(t) + require.NoError(t, store.Put("primary", entry("old-primary"))) + require.NoError(t, store.Put("legacy", entry("old-legacy"))) + before, err := os.ReadFile(path) + require.NoError(t, err) + callbackErr := errors.New("abort staged update") + handled, err := storage.WithFileTransaction(t.Context(), store, func(view storage.Store) error { + require.NoError(t, view.Put("primary", entry("new-primary"))) + visible, err := store.Lookup("primary") + require.NoError(t, err) + assert.Equal(t, "old-primary", visible.Token.AccessToken) + require.NoError(t, view.Put("legacy", entry("new-legacy"))) + visible, err = store.Lookup("legacy") + require.NoError(t, err) + assert.Equal(t, "old-legacy", visible.Token.AccessToken) + if abort { + return callbackErr + } + return nil + }, nil) + require.True(t, handled) + if abort { + require.ErrorIs(t, err, callbackErr) + after, err := os.ReadFile(path) + require.NoError(t, err) + assert.Equal(t, before, after) + } else { + require.NoError(t, err) + for _, key := range []string{"primary", "legacy"} { + value, err := store.Lookup(key) + require.NoError(t, err) + assert.Equal(t, "new-"+key, value.Token.AccessToken) + } + } + }) + } +} + +func TestFileTransactionCommitFailureCleanupUsesOriginalMap(t *testing.T) { + store, path := newFileStore(t) + require.NoError(t, store.Put("login", entry("stale"))) + require.NoError(t, store.Put("unrelated", entry("keep"))) + backup := path + ".saved" + cleanupErr := errors.New("cleanup callback failed") + for _, failCleanup := range []bool{false, true} { + t.Run(map[bool]string{false: "cleanup", true: "cleanup-error"}[failCleanup], func(t *testing.T) { + require.NoError(t, store.Put("login", entry("stale"))) + handled, err := storage.WithFileTransaction(t.Context(), store, func(view storage.Store) error { + require.NoError(t, view.Put("login", entry("new"))) + require.NoError(t, os.Rename(path, backup)) + require.NoError(t, os.Mkdir(path, 0o700)) + return nil + }, func(view storage.Store) error { + require.NoError(t, os.Remove(path)) + require.NoError(t, os.Rename(backup, path)) + if failCleanup { + return cleanupErr + } + return view.Delete("login") + }) + require.True(t, handled) + require.Error(t, err) + assert.ErrorContains(t, err, "cache update") + if failCleanup { + assert.ErrorIs(t, err, cleanupErr) + value, err := store.Lookup("login") + require.NoError(t, err) + assert.Equal(t, "stale", value.Token.AccessToken) + } else { + _, err := store.Lookup("login") + assert.ErrorIs(t, err, storage.ErrNotFound) + } + other, err := store.Lookup("unrelated") + require.NoError(t, err) + assert.Equal(t, "keep", other.Token.AccessToken) + }) + } +} + +func TestFileTransactionPreservesPendingLoginAndClonesViews(t *testing.T) { + store, _ := newFileStore(t) + pending := &storage.PendingLogin{DeviceCode: "synthetic-device", UserCode: "synthetic-user", Host: "https://host.test", ClientID: "client", Scopes: []string{"scope"}} + require.NoError(t, store.Put("pending", storage.Entry{PendingLogin: pending})) + handled, err := storage.WithFileTransaction(t.Context(), store, func(view storage.Store) error { + value, err := view.Lookup("pending") + require.NoError(t, err) + value.PendingLogin.Scopes[0] = "unpublished" + value, err = view.Lookup("pending") + require.NoError(t, err) + assert.Equal(t, pending, value.PendingLogin) + return view.Put("unrelated", entry("new")) + }, nil) + require.True(t, handled) + require.NoError(t, err) + value, err := store.Lookup("pending") + require.NoError(t, err) + assert.Equal(t, pending, value.PendingLogin) +} + +func TestFileTransactionUsesOriginalContextAndDifferentStoresRemainIndependent(t *testing.T) { + store, path := newFileStore(t) + other, _ := newFileStore(t) + held, release, done := make(chan struct{}), make(chan struct{}), make(chan error, 1) + releaseLease := sync.OnceFunc(func() { close(release) }) + t.Cleanup(releaseLease) + go func() { + _, err := storage.WithFileTransaction(t.Context(), store, func(view storage.Store) error { + close(held) + <-release + return view.Put("first", entry("first")) + }, nil) + done <- err + }() + <-held + ctx, cancel := context.WithTimeout(t.Context(), 100*time.Millisecond) + defer cancel() + second, err := storage.NewFileStore(ctx, storage.WithFileLocation(path)) + require.NoError(t, err) + err = storage.PutWithContext(ctx, second, "second", entry("second")) + assert.ErrorIs(t, err, context.DeadlineExceeded) + require.NoError(t, other.Put("independent", entry("independent"))) + releaseLease() + require.NoError(t, <-done) + require.NoError(t, store.Put("second", entry("second"))) + for _, key := range []string{"first", "second"} { + _, err := store.Lookup(key) + assert.NoError(t, err) + } +} + +func TestFileStoreLegacyWritesOutliveInitializationContext(t *testing.T) { + ctx, cancel := context.WithCancel(t.Context()) + path := filepath.Join(t.TempDir(), "tokens.json") + store, err := storage.NewFileStore(ctx, storage.WithFileLocation(path)) + require.NoError(t, err) + cancel() + require.NoError(t, store.Put("legacy", entry("value"))) + value, err := store.Lookup("legacy") + require.NoError(t, err) + assert.Equal(t, "value", value.Token.AccessToken) + require.NoError(t, store.Delete("legacy")) + _, err = store.Lookup("legacy") + assert.ErrorIs(t, err, storage.ErrNotFound) +} + +func TestFileStoreCancelledOperationDoesNotMutate(t *testing.T) { + store, path := newFileStore(t) + require.NoError(t, store.Put("keep", entry("old"))) + before, err := os.ReadFile(path) + require.NoError(t, err) + ctx, cancel := context.WithCancel(t.Context()) + cancel() + assert.ErrorIs(t, storage.PutWithContext(ctx, store, "keep", entry("new")), context.Canceled) + assert.ErrorIs(t, storage.DeleteWithContext(ctx, store, "keep"), context.Canceled) + called := false + handled, err := storage.WithFileTransaction(ctx, store, func(storage.Store) error { called = true; return nil }, nil) + assert.True(t, handled) + assert.ErrorIs(t, err, context.Canceled) + assert.False(t, called) + after, err := os.ReadFile(path) + require.NoError(t, err) + assert.Equal(t, before, after) +} + +func TestFileStoreConstructorJoinsTheSameLease(t *testing.T) { + store, path := newFileStore(t) + require.NoError(t, store.Put("original", entry("keep"))) + held, release, done := make(chan struct{}), make(chan struct{}), make(chan error, 1) + releaseLease := sync.OnceFunc(func() { close(release) }) + t.Cleanup(releaseLease) + go func() { + _, err := storage.WithFileTransaction(t.Context(), store, func(view storage.Store) error { + if err := os.Remove(path); err != nil { + return err + } + close(held) + <-release + return view.Put("restored", entry("restored")) + }, nil) + done <- err + }() + <-held + ctx, cancel := context.WithTimeout(t.Context(), 100*time.Millisecond) + defer cancel() + _, err := storage.NewFileStore(ctx, storage.WithFileLocation(path)) + assert.ErrorIs(t, err, context.DeadlineExceeded) + _, err = os.Stat(path) + assert.ErrorIs(t, err, os.ErrNotExist) + releaseLease() + require.NoError(t, <-done) + for _, key := range []string{"original", "restored"} { + _, err := store.Lookup(key) + assert.NoError(t, err) + } +} + +func TestUnsupportedStoreDoesNotInvokeFileCallbacks(t *testing.T) { + store := storage.NewMemoryStore() + require.NoError(t, store.Put("keep", entry("keep"))) + for _, subject := range []storage.Store{store, storage.WrapForOAuthArgument(t.Context(), store, storage.StorageModeSecure, aliasKey{})} { + called := false + handled, err := storage.WithFileTransaction(t.Context(), subject, func(storage.Store) error { called = true; return nil }, nil) + assert.False(t, handled) + assert.False(t, called) + assert.NoError(t, err) + value, err := subject.Lookup("keep") + require.NoError(t, err) + assert.Equal(t, "keep", value.Token.AccessToken) + } +} + +type aliasKey struct{} + +func (aliasKey) GetCacheKey() string { return "primary" } +func (aliasKey) GetHostCacheKey() string { return "https://host.test" } + +func TestFileTransactionWrappersShareOneStagedView(t *testing.T) { + store, _ := newFileStore(t) + wrapped := storage.WrapForOAuthArgument(t.Context(), store, storage.StorageModePlaintext, aliasKey{}) + require.NoError(t, wrapped.Put("primary", entry("old"))) + handled, err := storage.WithFileTransaction(t.Context(), wrapped, func(view storage.Store) error { + require.NoError(t, view.Put("primary", entry("new"))) + for _, key := range []string{"primary", "https://host.test"} { + committed, err := store.Lookup(key) + require.NoError(t, err) + assert.Equal(t, "old", committed.Token.AccessToken) + staged, err := view.Lookup(key) + require.NoError(t, err) + assert.Equal(t, "new", staged.Token.AccessToken) + } + _, err := view.Lookup("missing") + assert.ErrorIs(t, err, storage.ErrNotFound) + return nil + }, nil) + require.True(t, handled) + require.NoError(t, err) + for _, key := range []string{"primary", "https://host.test"} { + value, err := store.Lookup(key) + require.NoError(t, err) + assert.Equal(t, "new", value.Token.AccessToken) + } +} + +const fileLeaseActor = "DATABRICKS_TEST_FILE_LEASE_ACTOR" + +func TestFileTransactionOwnerExitReleasesStableSidecar(t *testing.T) { + if os.Getenv(fileLeaseActor) == "1" { + var path string + require.NoError(t, json.NewDecoder(os.Stdin).Decode(&path)) + store, err := storage.NewFileStore(t.Context(), storage.WithFileLocation(path)) + require.NoError(t, err) + _, err = storage.WithFileTransaction(t.Context(), store, func(storage.Store) error { + _, err := io.WriteString(os.Stdout, "held\n") + if err != nil { + return err + } + _, err = io.Copy(io.Discard, os.Stdin) + return err + }, nil) + require.NoError(t, err) + return + } + store, path := newFileStore(t) + before, err := os.Stat(path + ".lock") + require.NoError(t, err) + command := exec.CommandContext(t.Context(), os.Args[0], "-test.run=^TestFileTransactionOwnerExitReleasesStableSidecar$", "-test.timeout=30s") + home := t.TempDir() + command.Env = []string{fileLeaseActor + "=1", "HOME=" + home, "USERPROFILE=" + home, "PATH=" + os.Getenv("PATH"), "TMPDIR=" + os.TempDir(), "TMP=" + os.TempDir(), "TEMP=" + os.TempDir()} + input, err := command.StdinPipe() + require.NoError(t, err) + output, err := command.StdoutPipe() + require.NoError(t, err) + command.Stderr = io.Discard + require.NoError(t, command.Start()) + t.Cleanup(func() { _ = input.Close(); _ = command.Process.Kill(); _ = command.Wait() }) + require.NoError(t, json.NewEncoder(input).Encode(path)) + ready, err := bufio.NewReader(output).ReadString('\n') + require.NoError(t, err) + require.Equal(t, "held\n", ready) + ctx, cancel := context.WithTimeout(t.Context(), 100*time.Millisecond) + defer cancel() + _, err = storage.WithFileTransaction(ctx, store, func(storage.Store) error { return nil }, nil) + assert.ErrorIs(t, err, context.DeadlineExceeded) + require.NoError(t, command.Process.Kill()) + require.Error(t, command.Wait()) + require.NoError(t, store.Put("after-exit", entry("after-exit"))) + after, err := os.Stat(path + ".lock") + require.NoError(t, err) + assert.True(t, os.SameFile(before, after)) +} + +func TestFileTransactionPreservesNullEntries(t *testing.T) { + for _, test := range []struct { + name string + initial string + afterPut string + }{ + { + name: "healthy", + initial: `{"version":1,"tokens":{"healthy":{"access_token":"synthetic-before","expiry":"0001-01-01T00:00:00Z"}}}`, + afterPut: `{"version":1,"tokens":{"healthy":{"access_token":"synthetic-before","expiry":"0001-01-01T00:00:00Z"},"added":{"access_token":"synthetic-added","expiry":"0001-01-01T00:00:00Z"}}}`, + }, + { + name: "unrelated_null", + initial: `{"version":1,"tokens":{"healthy":{"access_token":"synthetic-before","expiry":"0001-01-01T00:00:00Z"},"unrelated":null}}`, + afterPut: `{"version":1,"tokens":{"healthy":{"access_token":"synthetic-before","expiry":"0001-01-01T00:00:00Z"},"unrelated":null,"added":{"access_token":"synthetic-added","expiry":"0001-01-01T00:00:00Z"}}}`, + }, + } { + t.Run(test.name, func(t *testing.T) { + path := filepath.Join(t.TempDir(), "tokens.json") + require.NoError(t, os.WriteFile(path, []byte(test.initial), 0o600)) + store, err := storage.NewFileStore(t.Context(), storage.WithFileLocation(path)) + require.NoError(t, err) + _, err = store.Lookup("unrelated") + require.ErrorIs(t, err, storage.ErrNotFound) + + handled, err := storage.WithFileTransaction(t.Context(), store, func(view storage.Store) error { + _, err := view.Lookup("unrelated") + require.ErrorIs(t, err, storage.ErrNotFound) + healthy, err := view.Lookup("healthy") + require.NoError(t, err) + require.NotNil(t, healthy.Token) + assert.Equal(t, "synthetic-before", healthy.Token.AccessToken) + return nil + }, nil) + require.True(t, handled) + require.NoError(t, err) + after, err := os.ReadFile(path) + require.NoError(t, err) + assert.Equal(t, []byte(test.initial), after) + + require.NoError(t, store.Put("added", entry("synthetic-added"))) + after, err = os.ReadFile(path) + require.NoError(t, err) + assert.JSONEq(t, test.afterPut, string(after)) + healthy, err := store.Lookup("healthy") + require.NoError(t, err) + require.NotNil(t, healthy.Token) + assert.Equal(t, "synthetic-before", healthy.Token.AccessToken) + + require.NoError(t, store.Delete("missing")) + unchanged, err := os.ReadFile(path) + require.NoError(t, err) + assert.Equal(t, after, unchanged) + require.NoError(t, store.Delete("added")) + after, err = os.ReadFile(path) + require.NoError(t, err) + assert.JSONEq(t, test.initial, string(after)) + }) + } +} diff --git a/libs/auth/storage/filestore.go b/libs/auth/storage/filestore.go index 64e4b02ab2e..0538910b5ec 100644 --- a/libs/auth/storage/filestore.go +++ b/libs/auth/storage/filestore.go @@ -75,6 +75,8 @@ type fileStore struct { // locker protects the token store file from concurrent reads and writes. locker sync.Mutex + + beforeIntentCompletion func() } // NewFileStore creates a new file-backed Store. By default, tokens are stored @@ -99,17 +101,7 @@ func NewFileStore(ctx context.Context, opts ...FileStoreOption) (Store, error) { // Put implements the Store interface. func (c *fileStore) Put(key string, e Entry) error { - c.locker.Lock() - defer c.locker.Unlock() - f, err := c.load() - if err != nil { - return fmt.Errorf("load: %w", err) - } - if f.Tokens == nil { - f.Tokens = map[string]*fileEntry{} - } - f.Tokens[key] = &fileEntry{Token: e.Token, PendingLogin: e.PendingLogin} - return c.write(f) + return PutWithContext(context.Background(), c, key, e) //nolint:gocritic // Legacy Store interface has no operation context. } // Lookup implements the Store interface. @@ -121,7 +113,7 @@ func (c *fileStore) Lookup(key string) (Entry, error) { return Entry{}, fmt.Errorf("load: %w", err) } fe, ok := f.Tokens[key] - if !ok { + if !ok || fe == nil { return Entry{}, ErrNotFound } return Entry{Token: fe.Token, PendingLogin: fe.PendingLogin}, nil @@ -129,25 +121,58 @@ func (c *fileStore) Lookup(key string) (Entry, error) { // Delete implements the Store interface. Removing a missing key is a no-op. func (c *fileStore) Delete(key string) error { - c.locker.Lock() - defer c.locker.Unlock() - f, err := c.load() + return DeleteWithContext(context.Background(), c, key) //nolint:gocritic // Legacy Store interface has no operation context. +} + +// WithFileTransaction publishes one staged map while holding the cache lease. +func (c *fileStore) WithFileTransaction(ctx context.Context, update, cleanup func(Store) error) (bool, error) { + unlock, err := lockTokenStore(ctx, c.fileLocation+".lock") + if err != nil { + return true, fmt.Errorf("token store lock: %w", err) + } + defer unlock() + file, err := c.load() + if err != nil { + return true, fmt.Errorf("load: %w", err) + } + raw, err := os.ReadFile(c.fileLocation) + if err != nil { + return true, fmt.Errorf("load: %w", err) + } + journal, err := c.loadRefreshJournal(ctx) if err != nil { - return fmt.Errorf("load: %w", err) + return true, err + } + if err := c.reconcileRefreshJournal(ctx, journal, raw); err != nil { + return true, err } - if _, ok := f.Tokens[key]; !ok { - return nil + staged := newFileTransaction(ctx, c, file, raw, journal) + defer func() { staged.live = false }() + if err := update(staged); err != nil { + return true, err } - delete(f.Tokens, key) - return c.write(f) + staged.live = false + return true, c.publishFileTransaction(staged, cleanup) } // write marshals f and atomically replaces the store file. -func (c *fileStore) write(f *tokenStoreFile) error { - raw, err := json.MarshalIndent(f, "", " ") +func encodeTokenStoreFile(file *tokenStoreFile) ([]byte, error) { + raw, err := json.MarshalIndent(file, "", " ") if err != nil { - return fmt.Errorf("marshal: %w", err) + return nil, fmt.Errorf("marshal: %w", err) } + return raw, nil +} + +func (c *fileStore) write(file *tokenStoreFile) error { + raw, err := encodeTokenStoreFile(file) + if err != nil { + return err + } + return c.writeTokenStoreBytes(raw) +} + +func (c *fileStore) writeTokenStoreBytes(raw []byte) error { if err := atomicfile.Write(c.fileLocation, raw, ownerReadWrite); err != nil { return fmt.Errorf("error storing token in local cache: %w", err) } @@ -165,14 +190,40 @@ func (c *fileStore) init(ctx context.Context) error { } c.fileLocation = filepath.Join(home, tokenStoreFilePath) } + if err := ctx.Err(); err != nil { + return err + } + location, err := filepath.Abs(c.fileLocation) + if err != nil { + return fmt.Errorf("resolve file location: %w", err) + } + if err := os.MkdirAll(filepath.Dir(location), ownerExecReadWrite); err != nil { + return fmt.Errorf("mkdir: %w", err) + } + parent, err := filepath.EvalSymlinks(filepath.Dir(location)) + if err != nil { + return fmt.Errorf("resolve store directory: %w", err) + } + c.fileLocation = filepath.Join(parent, filepath.Base(location)) + if target, err := filepath.EvalSymlinks(c.fileLocation); err == nil { + c.fileLocation = target + } else if !errors.Is(err, fs.ErrNotExist) { + return fmt.Errorf("resolve store file: %w", err) + } // Create the store file if it does not exist. if _, err := os.Stat(c.fileLocation); err != nil { if !errors.Is(err, fs.ErrNotExist) { return fmt.Errorf("stat file: %w", err) } - // Create the parent directories if needed. - if err := os.MkdirAll(filepath.Dir(c.fileLocation), ownerExecReadWrite); err != nil { - return fmt.Errorf("mkdir: %w", err) + unlock, err := lockTokenStore(ctx, c.fileLocation+".lock") + if err != nil { + return fmt.Errorf("token store lock: %w", err) + } + defer unlock() + if _, err := os.Stat(c.fileLocation); err == nil { + return nil + } else if !errors.Is(err, fs.ErrNotExist) { + return fmt.Errorf("stat file: %w", err) } // Create an empty store file. diff --git a/libs/auth/storage/lock.go b/libs/auth/storage/lock.go new file mode 100644 index 00000000000..c3e148443f8 --- /dev/null +++ b/libs/auth/storage/lock.go @@ -0,0 +1,64 @@ +package storage + +import ( + "context" + "fmt" + "os" + "path/filepath" + "time" +) + +const ( + // lockRetryInterval is how long to wait before retrying a contended lock. + lockRetryInterval = 20 * time.Millisecond +) + +// lockTokenStore acquires the cross-process lock that serializes token store +// mutations. The stable sidecar survives atomic replacement of the cache. +func lockTokenStore(ctx context.Context, path string) (func(), error) { + if err := ctx.Err(); err != nil { + return nil, err + } + if err := os.MkdirAll(filepath.Dir(path), ownerExecReadWrite); err != nil { + return nil, fmt.Errorf("mkdir: %w", err) + } + + f, err := os.OpenFile(path, os.O_CREATE|os.O_RDWR, ownerReadWrite) + if err != nil { + return nil, fmt.Errorf("open lock file: %w", err) + } + + for { + if err := ctx.Err(); err != nil { + f.Close() + return nil, err + } + locked, err := tryLock(f) + if err != nil { + f.Close() + return nil, fmt.Errorf("lock %s: %w", filepath.ToSlash(path), err) + } + if locked { + if err := ctx.Err(); err != nil { + unlock(f) + f.Close() + return nil, err + } + return func() { + // Closing the file releases the lock on both platforms; the + // explicit unlock keeps the two operations from drifting apart. + unlock(f) + f.Close() + }, nil + } + + retry := time.NewTimer(lockRetryInterval) + select { + case <-ctx.Done(): + retry.Stop() + f.Close() + return nil, ctx.Err() + case <-retry.C: + } + } +} diff --git a/libs/auth/storage/lock_unix.go b/libs/auth/storage/lock_unix.go new file mode 100644 index 00000000000..d2d5127d730 --- /dev/null +++ b/libs/auth/storage/lock_unix.go @@ -0,0 +1,29 @@ +//go:build !windows + +package storage + +import ( + "errors" + "os" + + "golang.org/x/sys/unix" +) + +// tryLock takes an exclusive advisory lock on f without blocking. It reports +// whether the lock was taken. +func tryLock(f *os.File) (bool, error) { + err := unix.Flock(int(f.Fd()), unix.LOCK_EX|unix.LOCK_NB) + switch { + case err == nil: + return true, nil + case errors.Is(err, unix.EWOULDBLOCK): + return false, nil + default: + return false, err + } +} + +// unlock releases the advisory lock held on f. +func unlock(f *os.File) { + _ = unix.Flock(int(f.Fd()), unix.LOCK_UN) +} diff --git a/libs/auth/storage/lock_windows.go b/libs/auth/storage/lock_windows.go new file mode 100644 index 00000000000..5a607e8d2c0 --- /dev/null +++ b/libs/auth/storage/lock_windows.go @@ -0,0 +1,40 @@ +package storage + +import ( + "errors" + "os" + + "golang.org/x/sys/windows" +) + +// lockRegionLength is the number of bytes LockFileEx covers. The lock file has +// no contents, so locking a single byte is enough to make holders exclusive. +const lockRegionLength = 1 + +// tryLock takes an exclusive lock on f without blocking. It reports whether the +// lock was taken. +func tryLock(f *os.File) (bool, error) { + var overlapped windows.Overlapped + err := windows.LockFileEx( + windows.Handle(f.Fd()), + windows.LOCKFILE_EXCLUSIVE_LOCK|windows.LOCKFILE_FAIL_IMMEDIATELY, + 0, + lockRegionLength, + 0, + &overlapped, + ) + switch { + case err == nil: + return true, nil + case errors.Is(err, windows.ERROR_LOCK_VIOLATION): + return false, nil + default: + return false, err + } +} + +// unlock releases the lock held on f. +func unlock(f *os.File) { + var overlapped windows.Overlapped + _ = windows.UnlockFileEx(windows.Handle(f.Fd()), 0, lockRegionLength, 0, &overlapped) +} diff --git a/libs/auth/storage/not_found_hint.go b/libs/auth/storage/not_found_hint.go index 76ac854fa56..f2cc15539ea 100644 --- a/libs/auth/storage/not_found_hint.go +++ b/libs/auth/storage/not_found_hint.go @@ -43,6 +43,21 @@ func (s *notFoundHintStore) Lookup(key string) (Entry, error) { func (s *notFoundHintStore) Delete(key string) error { return s.inner.Delete(key) } +// WithFileTransaction preserves lookup hints on the staged file view. +func (s *notFoundHintStore) WithFileTransaction(ctx context.Context, update, cleanup func(Store) error) (bool, error) { + wrap := func(mutate func(Store) error) func(Store) error { + if mutate == nil { + return nil + } + return func(view Store) error { + wrapped := *s + wrapped.inner = view + return mutate(&wrapped) + } + } + return WithFileTransaction(ctx, s.inner, wrap(update), wrap(cleanup)) +} + // notFoundHint replaces ErrNotFound's terse "token not found" string // with an actionable message while still satisfying errors.Is(err, // ErrNotFound). PersistentAuth.loadToken wraps every store error with diff --git a/libs/auth/storage/refresh_intent.go b/libs/auth/storage/refresh_intent.go new file mode 100644 index 00000000000..610202efccd --- /dev/null +++ b/libs/auth/storage/refresh_intent.go @@ -0,0 +1,734 @@ +package storage + +import ( + "context" + "crypto/rand" + "crypto/sha256" + "encoding/binary" + "encoding/hex" + "encoding/json" + "errors" + "fmt" + "slices" +) + +var ( + ErrRefreshRecoveryRequired = errors.New("cached OAuth refresh outcome is uncertain; automatic refresh is blocked") + ErrRefreshIntentInvalid = errors.New("cached OAuth refresh metadata is invalid; automatic refresh is blocked") + ErrRefreshIntentIO = errors.New("cannot persist OAuth refresh state") + errRefreshOperation = errors.New("invalid file refresh operation") +) + +// FileRefreshStateError exposes a closed reason without rendering credential metadata. +type FileRefreshStateError struct{ reason, cause error } + +func (e *FileRefreshStateError) Error() string { + if e == nil || e.reason == nil { + return errRefreshOperation.Error() + } + return e.reason.Error() +} + +func (e *FileRefreshStateError) Unwrap() []error { + if e == nil || e.reason == nil { + return []error{errRefreshOperation} + } + if e.cause == nil { + return []error{e.reason} + } + return []error{e.reason, e.cause} +} + +func intentError(reason, cause error) error { + return &FileRefreshStateError{reason: reason, cause: cause} +} + +func intentHash(parts ...string) string { + h := sha256.New() + var size [8]byte + for _, part := range parts { + binary.BigEndian.PutUint64(size[:], uint64(len(part))) + _, _ = h.Write(size[:]) + _, _ = h.Write([]byte(part)) + } + return hex.EncodeToString(h.Sum(nil)) +} + +func intentNonce() (string, error) { + var nonce [32]byte + if _, err := rand.Read(nonce[:]); err != nil { + return "", intentError(ErrRefreshIntentIO, err) + } + return hex.EncodeToString(nonce[:]), nil +} + +func nullable(value string) *string { + if value == "" { + return nil + } + return &value +} + +func deref(value *string) string { + if value == nil { + return "" + } + return *value +} + +type refreshGeneration struct { + Consumed bool `json:"consumed"` + ActiveEpoch *string `json:"active_epoch"` + Reservation *string `json:"reservation"` +} +type refreshBinding struct { + Generation string `json:"generation"` + Epoch string `json:"epoch"` + Entry string `json:"entry_d"` +} +type refreshWrite struct { + Key string `json:"key_d"` + BeforeSlot string `json:"before_slot_d"` + BeforeGeneration *string `json:"before_generation"` + BeforeEpoch *string `json:"before_epoch"` + BeforeBindingEntry *string `json:"before_binding_entry_d"` + BeforeReservation *string `json:"before_reservation"` + AfterEntry string `json:"after_entry_d"` + AfterGeneration *string `json:"after_generation"` + AfterEpoch *string `json:"after_epoch"` +} +type refreshAttempt struct { + Kind string `json:"kind"` + Phase string `json:"phase"` + InputGeneration *string `json:"input_generation"` + InputEpoch *string `json:"input_epoch"` + InputKey string `json:"input_key_d"` + InputEntry *string `json:"input_entry_d"` + PreCache string `json:"pre_cache_d"` + Request *string `json:"request_d"` + ResponseEntry *string `json:"response_entry_d"` + ResultGeneration *string `json:"result_generation"` + PostCache *string `json:"post_cache_d"` + Writes []refreshWrite `json:"writes"` + Proof *string `json:"response_proof_d"` +} +type refreshJournal struct { + retirements []refreshRetirement `json:"-"` + disk *refreshJournal `json:"-"` + sequence uint64 `json:"-"` + chain [32]byte `json:"-"` + end int64 `json:"-"` + failed error `json:"-"` + Version int `json:"version"` + Scope string `json:"store_scope_d"` + Nonce string `json:"store_nonce"` + Generations map[string]*refreshGeneration `json:"generations"` + Bindings map[string]*refreshBinding `json:"bindings"` + Attempts map[string]*refreshAttempt `json:"attempts"` +} + +func newRefreshJournal(location string) (*refreshJournal, error) { + nonce, err := intentNonce() + if err != nil { + return nil, err + } + return &refreshJournal{ + Version: 1, Scope: intentHash("cli-file-intent/store/v1", location), Nonce: nonce, + Generations: map[string]*refreshGeneration{}, Bindings: map[string]*refreshBinding{}, Attempts: map[string]*refreshAttempt{}, + }, nil +} + +func (j *refreshJournal) key(key string) string { + return intentHash("cli-file-intent/key/v1", j.Nonce, key) +} + +func (j *refreshJournal) generation(entry Entry) string { + if entry.Token == nil || entry.Token.RefreshToken == "" { + return "" + } + return intentHash("cli-file-intent/generation/v1", j.Nonce, entry.Token.RefreshToken) +} + +func entryIntentDigest(entry *fileEntry) (string, error) { + raw, err := json.Marshal(entry) + if err != nil { + return "", intentError(ErrRefreshIntentInvalid, err) + } + return intentHash("cli-file-intent/entry/v1", string(raw)), nil +} + +func slotIntentDigest(file *tokenStoreFile, key string) (string, error) { + entry, ok := file.Tokens[key] + if !ok { + return intentHash("cli-file-intent/slot/v1", "absent"), nil + } + digest, err := entryIntentDigest(entry) + if err != nil { + return "", err + } + return intentHash("cli-file-intent/slot/v1", "present", digest), nil +} + +func cacheIntentDigest(raw []byte) string { return intentHash("cli-file-intent/cache/v1", string(raw)) } + +func attemptProof(nonce string, a *refreshAttempt) string { + copy := *a + copy.Phase = "" + copy.Proof = nil + raw, _ := json.Marshal(copy) + return intentHash("cli-file-intent/response/v1", nonce, string(raw)) +} + +// FileRefreshAttempt is valid only during its owning staged transaction callback. +type FileRefreshAttempt struct { + view *fileTransaction + nonce string + key string + response *Entry +} + +func transactionIntentView(view Store) *fileTransaction { + switch view := view.(type) { + case *fileTransaction: + return view + case *DualWritingStore: + return transactionIntentView(view.inner) + case *notFoundHintStore: + return transactionIntentView(view.inner) + default: + return nil + } +} + +func (t *fileTransaction) ensureJournal() error { + if t.journal != nil { + return nil + } + journal, err := newRefreshJournal(t.owner.fileLocation) + if err != nil { + return err + } + t.journal = journal + return nil +} + +// BeginFileRefresh records consumption before the caller exchanges a cached credential. +// The view must come from the current WithFileTransaction callback. +func BeginFileRefresh(ctx context.Context, view Store, key string, requestDigest [32]byte) (*FileRefreshAttempt, error) { + t := transactionIntentView(view) + if t == nil || !t.live || t.attempt != nil { + return nil, intentError(errRefreshOperation, nil) + } + if err := ctx.Err(); err != nil { + return nil, err + } + if err := t.ctx.Err(); err != nil { + return nil, err + } + current, err := t.Lookup(key) + if err != nil { + return nil, err + } + if current.Token == nil || current.Token.RefreshToken == "" { + return nil, intentError(errRefreshOperation, nil) + } + if err := t.ensureJournal(); err != nil { + return nil, err + } + j := t.journal + g := j.generation(current) + k := j.key(key) + entry, err := entryIntentDigest(t.file.Tokens[key]) + if err != nil { + return nil, err + } + gen := j.Generations[g] + if gen != nil { + binding := j.Bindings[k] + if gen.Reservation != nil || gen.ActiveEpoch == nil || binding == nil || binding.Generation != g || binding.Epoch != deref(gen.ActiveEpoch) || binding.Entry != entry { + return nil, intentError(ErrRefreshRecoveryRequired, nil) + } + } else { + gen = &refreshGeneration{} + j.Generations[g] = gen + } + nonce, err := intentNonce() + if err != nil { + return nil, err + } + a := &refreshAttempt{ + Kind: "refresh", Phase: "prepared", InputGeneration: nullable(g), InputEpoch: gen.ActiveEpoch, + InputKey: k, InputEntry: nullable(entry), PreCache: cacheIntentDigest(t.originalRaw), Request: nullable(hex.EncodeToString(requestDigest[:])), Writes: []refreshWrite{}, + } + j.Attempts[nonce] = a + gen.Consumed = true + gen.ActiveEpoch = nil + gen.Reservation = nullable(nonce) + if err := t.owner.saveRefreshJournal(t.ctx, j); err != nil { + return nil, err + } + attempt := &FileRefreshAttempt{view: t, nonce: nonce, key: key} + t.attempt = attempt + return attempt, nil +} + +// RecordResponse binds an actual successful exchange to this live attempt. +func (a *FileRefreshAttempt) RecordResponse(entry Entry) error { + if a == nil || a.view == nil || !a.view.live || a.view.attempt != a || a.response != nil || entry.Token == nil || entry.Token.AccessToken == "" { + return intentError(errRefreshOperation, nil) + } + if _, err := entryIntentDigest(&fileEntry{Token: entry.Token, PendingLogin: entry.PendingLogin}); err != nil { + return err + } + copy := cloneEntry(entry) + a.response = © + return nil +} + +// RecordFileLoginResult marks an actual fresh-login result without inventing a refresh generation. +func RecordFileLoginResult(view Store, key string, entry Entry) error { + t := transactionIntentView(view) + if t == nil || !t.live || t.attempt != nil || entry.Token == nil || entry.Token.AccessToken == "" { + return intentError(errRefreshOperation, nil) + } + if _, err := entryIntentDigest(&fileEntry{Token: entry.Token, PendingLogin: entry.PendingLogin}); err != nil { + return err + } + // Access-only login needs no companion until existing generation metadata exists. + if t.journal == nil && entry.Token.RefreshToken == "" { + return nil + } + if err := t.ensureJournal(); err != nil { + return err + } + nonce, err := intentNonce() + if err != nil { + return err + } + prior, _ := t.Lookup(key) + k := t.journal.key(key) + var epoch *string + if b := t.journal.Bindings[k]; b != nil { + epoch = nullable(b.Epoch) + } + inputEntry, err := entryIntentDigest(t.file.Tokens[key]) + if err != nil { + return err + } + a := &refreshAttempt{ + Kind: "login", Phase: "prepared", InputGeneration: nullable(t.journal.generation(prior)), InputEpoch: epoch, + InputKey: k, InputEntry: nullable(inputEntry), PreCache: cacheIntentDigest(t.originalRaw), Writes: []refreshWrite{}, + } + if _, present := t.file.Tokens[key]; !present { + a.InputEntry = nil + } + t.journal.Attempts[nonce] = a + copy := cloneEntry(entry) + t.attempt = &FileRefreshAttempt{view: t, nonce: nonce, key: key, response: ©} + return nil +} + +func (t *fileTransaction) responseProof(raw []byte) error { + if t.attempt == nil { + return nil + } + handle := t.attempt + a := t.journal.Attempts[handle.nonce] + if handle.response == nil { + return intentError(errRefreshOperation, nil) + } + response := &fileEntry{Token: handle.response.Token, PendingLogin: handle.response.PendingLogin} + responseEntry, err := entryIntentDigest(response) + if err != nil { + return err + } + primary, err := entryIntentDigest(t.file.Tokens[t.primaryKey()]) + if err != nil { + return err + } + if primary != responseEntry { + return intentError(errRefreshOperation, nil) + } + a.ResponseEntry = nullable(responseEntry) + a.ResultGeneration = nullable(t.journal.generation(*handle.response)) + a.PostCache = nullable(cacheIntentDigest(raw)) + a.Writes = []refreshWrite{} + keys := make([]string, 0, len(t.touched)) + for key := range t.touched { + keys = append(keys, key) + } + slices.Sort(keys) + for _, key := range keys { + k := t.journal.key(key) + old := t.original.Tokens[key] + oldEntry := Entry{} + if old != nil { + oldEntry = Entry{Token: old.Token, PendingLogin: old.PendingLogin} + } + g := t.journal.generation(oldEntry) + var epoch, reservation, bindingEntry *string + if b := t.beforeBindings[k]; b != nil { + epoch = nullable(b.Epoch) + bindingEntry = nullable(b.Entry) + } + if gen := t.beforeGenerations[g]; gen != nil { + reservation = gen.Reservation + } + after := t.file.Tokens[key] + afterEntry := Entry{} + if after != nil { + afterEntry = Entry{Token: after.Token, PendingLogin: after.PendingLogin} + } + beforeSlot, err := slotIntentDigest(t.original, key) + if err != nil { + return err + } + afterDigest, err := entryIntentDigest(after) + if err != nil { + return err + } + w := refreshWrite{ + Key: k, BeforeSlot: beforeSlot, BeforeGeneration: nullable(g), BeforeEpoch: epoch, BeforeBindingEntry: bindingEntry, BeforeReservation: reservation, + AfterEntry: afterDigest, AfterGeneration: nullable(t.journal.generation(afterEntry)), + } + if after != nil && responseEntry == w.AfterEntry && w.AfterGeneration != nil { + w.AfterEpoch = nullable(handle.nonce) + } + a.Writes = append(a.Writes, w) + } + if !t.journal.canComplete(handle.nonce, a) { + return intentError(ErrRefreshRecoveryRequired, nil) + } + a.Phase = "response_ready" + a.Proof = nullable(attemptProof(handle.nonce, a)) + return t.owner.saveRefreshJournal(t.ctx, t.journal) +} + +func (t *fileTransaction) primaryKey() string { return t.attempt.key } + +func (j *refreshJournal) capturedReservation(a *refreshAttempt, g, nonce string) bool { + owner := j.Attempts[nonce] + if owner == nil || owner.Kind != "refresh" || deref(owner.InputGeneration) != g { + return false + } + for _, w := range a.Writes { + if w.Key == owner.InputKey && deref(w.BeforeGeneration) == g && deref(w.BeforeReservation) == nonce && deref(w.BeforeEpoch) == deref(owner.InputEpoch) && owner.InputEntry != nil && w.BeforeSlot == intentHash("cli-file-intent/slot/v1", "present", *owner.InputEntry) { + return true + } + } + return false +} + +func (j *refreshJournal) canComplete(nonce string, a *refreshAttempt) bool { + if a.Kind == "refresh" { + gen := j.Generations[deref(a.InputGeneration)] + if gen == nil || deref(gen.Reservation) != nonce { + return false + } + } + for _, w := range a.Writes { + b := j.Bindings[w.Key] + if b != nil && (b.Generation != deref(w.BeforeGeneration) || b.Epoch != deref(w.BeforeEpoch) || b.Entry != deref(w.BeforeBindingEntry)) { + return false + } + if b == nil && (w.BeforeEpoch != nil || w.BeforeBindingEntry != nil) { + return false + } + g := deref(w.BeforeGeneration) + if g != "" { + current := "" + if gen := j.Generations[g]; gen != nil { + current = deref(gen.Reservation) + } + if (a.Kind != "refresh" || g != deref(a.InputGeneration) || current != nonce) && current != deref(w.BeforeReservation) { + return false + } + if a.Kind == "login" && current != "" { + owner := j.Attempts[current] + if owner != nil && owner.InputKey == w.Key && !j.capturedReservation(a, g, current) { + return false + } + } + } + } + result := deref(a.ResultGeneration) + if result != "" { + if gen := j.Generations[result]; gen != nil && gen.Reservation != nil && deref(gen.Reservation) != nonce { + if a.Kind != "login" || !j.capturedReservation(a, result, *gen.Reservation) { + return false + } + } + } + return true +} + +func (j *refreshJournal) complete(nonce string, a *refreshAttempt) error { + if !j.canComplete(nonce, a) { + return intentError(ErrRefreshRecoveryRequired, nil) + } + if a.Kind == "refresh" { + j.Generations[deref(a.InputGeneration)].Reservation = nil + } + if a.Kind == "login" { + for _, w := range a.Writes { + g := deref(w.BeforeGeneration) + reservation := deref(w.BeforeReservation) + if g != "" && reservation != "" && j.capturedReservation(a, g, reservation) { + gen := j.Generations[g] + if gen != nil && deref(gen.Reservation) == reservation { + j.Attempts[reservation].Phase = "uncertain" + gen.Reservation = nil + } + } + } + } + for _, w := range a.Writes { + delete(j.Bindings, w.Key) + } + if g := deref(a.ResultGeneration); g != "" { + gen := j.Generations[g] + if gen == nil { + gen = &refreshGeneration{} + j.Generations[g] = gen + } + if gen.Reservation != nil && deref(gen.Reservation) != nonce { + return intentError(ErrRefreshRecoveryRequired, nil) + } + gen.Reservation = nil + gen.ActiveEpoch = nullable(nonce) + for _, w := range a.Writes { + if w.AfterEpoch != nil && deref(w.AfterGeneration) == g { + j.Bindings[w.Key] = &refreshBinding{Generation: g, Epoch: nonce, Entry: w.AfterEntry} + } + } + } + a.Phase = "completed" + return nil +} + +func (c *fileStore) reconcileRefreshJournal(ctx context.Context, j *refreshJournal, raw []byte) error { + if j == nil { + return nil + } + keys := make([]string, 0, len(j.Attempts)) + for nonce := range j.Attempts { + keys = append(keys, nonce) + } + slices.Sort(keys) + for _, nonce := range keys { + a := j.Attempts[nonce] + if a.Phase != "response_ready" { + continue + } + if deref(a.PostCache) == cacheIntentDigest(raw) && j.canComplete(nonce, a) { + if err := j.complete(nonce, a); err != nil { + return err + } + } else { + a.Phase = "uncertain" + } + if err := c.saveRefreshJournal(ctx, j); err != nil { + return err + } + } + return nil +} + +func (t *fileTransaction) retireTouchedBindings() error { + if t.journal == nil { + return nil + } + t.journal.retirements = nil + var proofs []refreshRetirement + keys := make([]string, 0, len(t.touched)) + for key := range t.touched { + keys = append(keys, key) + } + slices.Sort(keys) + for _, key := range keys { + oldSlot, err := slotIntentDigest(t.original, key) + if err != nil { + return err + } + newSlot, err := slotIntentDigest(t.file, key) + if err != nil { + return err + } + if oldSlot == newSlot { + continue + } + entry := t.original.Tokens[key] + if entry == nil || entry.Token == nil || entry.RefreshToken == "" { + continue + } + k := t.journal.key(key) + g := t.journal.generation(Entry{Token: entry.Token, PendingLogin: entry.PendingLogin}) + digest, err := entryIntentDigest(entry) + if err != nil { + return err + } + proof := refreshRetirement{Key: k, Generation: g, Entry: digest} + if b := t.beforeBindings[k]; b != nil { + proof.Epoch = nullable(b.Epoch) + proof.BindingEntry = nullable(b.Entry) + } + if gen := t.beforeGenerations[g]; gen != nil { + proof.Reservation = gen.Reservation + } + proofs = append(proofs, proof) + } + used, err := t.journal.retireCapturedSet(proofs) + if err != nil { + return err + } + t.journal.retirements = used + return nil +} + +func (t *fileTransaction) finishJournal() error { + if t.journal == nil { + return nil + } + if t.attempt != nil { + a := t.journal.Attempts[t.attempt.nonce] + if err := t.journal.complete(t.attempt.nonce, a); err != nil { + return err + } + } else { + if err := t.retireTouchedBindings(); err != nil { + return err + } + } + return t.owner.saveRefreshJournal(t.ctx, t.journal) +} + +func (t *fileTransaction) abandonPublication() error { + if t.attempt == nil { + return nil + } + if t.journal.failed != nil { + return intentError(ErrRefreshIntentIO, t.journal.failed) + } + if t.journal.disk == nil { + delete(t.journal.Attempts, t.attempt.nonce) + return nil + } + prior := t.journal.disk.Attempts[t.attempt.nonce] + if prior == nil { + delete(t.journal.Attempts, t.attempt.nonce) + return nil + } + if prior.Phase == "uncertain" { + return nil + } + copy := *prior + copy.Phase = "uncertain" + t.journal.Attempts[t.attempt.nonce] = © + return t.owner.saveRefreshJournal(t.ctx, t.journal) +} + +func copyBindings(j *refreshJournal) map[string]*refreshBinding { + result := map[string]*refreshBinding{} + if j != nil { + for k, b := range j.Bindings { + copy := *b + result[k] = © + } + } + return result +} + +func copyGenerations(j *refreshJournal) map[string]*refreshGeneration { + result := map[string]*refreshGeneration{} + if j != nil { + for k, g := range j.Generations { + copy := *g + result[k] = © + } + } + return result +} + +func newFileTransaction(ctx context.Context, c *fileStore, file *tokenStoreFile, raw []byte, j *refreshJournal) *fileTransaction { + return &fileTransaction{ + owner: c, ctx: ctx, file: cloneStoreFile(file), original: cloneStoreFile(file), originalRaw: raw, journal: j, + touched: map[string]bool{}, live: true, beforeBindings: copyBindings(j), beforeGenerations: copyGenerations(j), + } +} + +// RefreshRequestDigest binds evidence to the actual configured refresh route. +func RefreshRequestDigest(endpoint, clientID string) [32]byte { + text := intentHash("cli-file-intent/request/v1", endpoint, clientID) + bytes, _ := hex.DecodeString(text) + var result [32]byte + copy(result[:], bytes) + return result +} + +func intentStateReason(err error) error { + if err == nil { + return nil + } + return fmt.Errorf("cache update: %w", err) +} + +func (c *fileStore) publishFileTransaction(t *fileTransaction, cleanup func(Store) error) error { + if !t.dirty { + if t.attempt != nil { + return intentError(errRefreshOperation, nil) + } + return nil + } + raw, err := encodeTokenStoreFile(t.file) + if err != nil { + return c.failedFilePublication(t, cleanup, err) + } + if t.attempt != nil { + if err := t.responseProof(raw); err != nil { + return c.failedFilePublication(t, cleanup, err) + } + } else if t.journal != nil { + if err := t.retireTouchedBindings(); err != nil { + return c.failedFilePublication(t, cleanup, err) + } + if err := c.saveRefreshJournal(t.ctx, t.journal); err != nil { + return c.failedFilePublication(t, cleanup, err) + } + } + if err := c.writeTokenStoreBytes(raw); err != nil { + return c.failedFilePublication(t, cleanup, err) + } + if t.attempt != nil { + if c.beforeIntentCompletion != nil { + c.beforeIntentCompletion() + } + return t.finishJournal() + } + return nil +} + +func (c *fileStore) failedFilePublication(t *fileTransaction, cleanup func(Store) error, cause error) error { + publicationErr := intentStateReason(cause) + stateErr := t.abandonPublication() + if cleanup == nil { + return errors.Join(publicationErr, stateErr) + } + rollback := newFileTransaction(t.ctx, c, t.original, t.originalRaw, t.journal) + defer func() { rollback.live = false }() + cleanupErr := cleanup(rollback) + rollback.live = false + if cleanupErr == nil && rollback.dirty { + // An uncertain append is never followed by another journal append. + // Existing failure cleanup may still remove the stale cache entries. + if rollback.journal != nil && rollback.journal.failed == nil && rollback.journal.disk != nil { + cleanupErr = rollback.retireTouchedBindings() + if cleanupErr == nil { + cleanupErr = c.saveRefreshJournal(t.ctx, rollback.journal) + } + } + if cleanupErr == nil { + cleanupErr = c.write(rollback.file) + } + } + return errors.Join(publicationErr, stateErr, cleanupErr) +} diff --git a/libs/auth/storage/refresh_intent_file.go b/libs/auth/storage/refresh_intent_file.go new file mode 100644 index 00000000000..a884bc5ea84 --- /dev/null +++ b/libs/auth/storage/refresh_intent_file.go @@ -0,0 +1,636 @@ +package storage + +import ( + "bytes" + "context" + "crypto/sha256" + "encoding/binary" + "encoding/hex" + "encoding/json" + "errors" + "io" + "math" + "os" + "reflect" +) + +const ( + refreshIntentMagic = "CLIRFI1\n" + refreshIntentSuffix = ".refresh-intents.json" +) + +type refreshHeader struct { + Version int `json:"version"` + Scope string `json:"store_scope_d"` + Nonce string `json:"store_nonce"` +} +type generationChange struct { + Expect *string `json:"expect_d"` + Value *refreshGeneration `json:"value"` +} +type bindingChange struct { + Expect *string `json:"expect_d"` + Value *refreshBinding `json:"value"` +} +type attemptChange struct { + Expect *string `json:"expect_d"` + Value *refreshAttempt `json:"value"` +} +type refreshDelta struct { + Retirements []refreshRetirement `json:"retirements,omitempty"` + Sequence uint64 `json:"seq"` + Transition string `json:"transition"` + Attempt *string `json:"attempt_nonce"` + Generations map[string]generationChange `json:"generations"` + Bindings map[string]bindingChange `json:"bindings"` + Attempts map[string]attemptChange `json:"attempts"` +} + +func digestValid(s string) bool { + if len(s) != 64 { + return false + } + b, err := hex.DecodeString(s) + return err == nil && hex.EncodeToString(b) == s +} +func optionalValid(s *string) bool { return s == nil || digestValid(*s) } +func invalidIntent(cause error) error { return intentError(ErrRefreshIntentInvalid, cause) } +func canonicalDecode(raw []byte, value any) error { + decoder := json.NewDecoder(bytes.NewReader(raw)) + decoder.DisallowUnknownFields() + if err := decoder.Decode(value); err != nil { + return invalidIntent(err) + } + var extra any + if err := decoder.Decode(&extra); err != io.EOF { + return invalidIntent(nil) + } + encoded, err := json.Marshal(value) + if err != nil || !bytes.Equal(raw, encoded) { + return invalidIntent(err) + } + return nil +} + +func cloneRefreshJournal(j *refreshJournal) (*refreshJournal, error) { + if j == nil { + return nil, invalidIntent(nil) + } + clone := &refreshJournal{Version: j.Version, Scope: j.Scope, Nonce: j.Nonce, Generations: map[string]*refreshGeneration{}, Bindings: map[string]*refreshBinding{}, Attempts: map[string]*refreshAttempt{}} + for k, v := range j.Generations { + if v == nil { + return nil, invalidIntent(nil) + } + copy := *v + clone.Generations[k] = © + } + for k, v := range j.Bindings { + if v == nil { + return nil, invalidIntent(nil) + } + copy := *v + clone.Bindings[k] = © + } + for k, v := range j.Attempts { + if v == nil { + return nil, invalidIntent(nil) + } + cloned := *v + if v.Writes != nil { + cloned.Writes = make([]refreshWrite, len(v.Writes)) + copy(cloned.Writes, v.Writes) + } + clone.Attempts[k] = &cloned + } + return clone, nil +} + +func recordDigest(kind, key string, value any) (string, error) { + raw, err := json.Marshal(value) + if err != nil { + return "", invalidIntent(err) + } + h := sha256.New() + _, _ = h.Write([]byte("cli-file-intent/record/v1\x00")) + for _, s := range []string{kind, key} { + var length [4]byte + binary.BigEndian.PutUint32(length[:], uint32(len(s))) + _, _ = h.Write(length[:]) + _, _ = h.Write([]byte(s)) + } + _, _ = h.Write(raw) + return hex.EncodeToString(h.Sum(nil)), nil +} + +func expectedRecord(kind, key string, value any) (*string, error) { + if reflect.ValueOf(value).IsNil() { + return nil, nil + } + digest, err := recordDigest(kind, key, value) + if err != nil { + return nil, err + } + return nullable(digest), nil +} + +func expectationMatches(kind, key string, value any, expect *string) bool { + digest, err := expectedRecord(kind, key, value) + return err == nil && reflect.DeepEqual(digest, expect) +} + +func validAttempt(nonce string, a *refreshAttempt) bool { + if a == nil || !digestValid(nonce) || (a.Kind != "refresh" && a.Kind != "login") || !digestValid(a.InputKey) || !digestValid(a.PreCache) { + return false + } + for _, v := range []*string{a.InputGeneration, a.InputEpoch, a.InputEntry, a.Request, a.ResponseEntry, a.ResultGeneration, a.PostCache, a.Proof} { + if !optionalValid(v) { + return false + } + } + if a.Kind == "refresh" && (a.InputGeneration == nil || a.InputEntry == nil || a.Request == nil) { + return false + } + if a.Kind == "login" && a.Request != nil { + return false + } + switch a.Phase { + case "prepared", "response_ready", "completed", "uncertain": + default: + return false + } + if a.Proof == nil { + return a.Kind == "refresh" && (a.Phase == "prepared" || a.Phase == "uncertain") && a.ResponseEntry == nil && a.ResultGeneration == nil && a.PostCache == nil && len(a.Writes) == 0 + } + if a.Phase == "prepared" || a.ResponseEntry == nil || a.PostCache == nil || len(a.Writes) == 0 || deref(a.Proof) != attemptProof(nonce, a) { + return false + } + seen := map[string]bool{} + primary := false + for _, w := range a.Writes { + if !digestValid(w.Key) || !digestValid(w.BeforeSlot) || !digestValid(w.AfterEntry) || seen[w.Key] { + return false + } + seen[w.Key] = true + for _, v := range []*string{w.BeforeGeneration, w.BeforeEpoch, w.BeforeBindingEntry, w.BeforeReservation, w.AfterGeneration, w.AfterEpoch} { + if !optionalValid(v) { + return false + } + } + if (w.BeforeEpoch == nil) != (w.BeforeBindingEntry == nil) || (w.BeforeEpoch != nil && w.BeforeGeneration == nil) || (w.BeforeReservation != nil && w.BeforeGeneration == nil) { + return false + } + if (w.AfterGeneration == nil) != (w.AfterEpoch == nil) { + return false + } + if w.AfterEpoch != nil && (deref(w.AfterEpoch) != nonce || deref(w.AfterGeneration) != deref(a.ResultGeneration) || w.AfterEntry != deref(a.ResponseEntry)) { + return false + } + if w.Key == a.InputKey && w.AfterEntry == deref(a.ResponseEntry) { + primary = true + } + } + return primary +} + +func validateRefreshJournal(j *refreshJournal) error { + if j == nil || j.Version != 1 || !digestValid(j.Scope) || !digestValid(j.Nonce) || j.Generations == nil || j.Bindings == nil || j.Attempts == nil { + return invalidIntent(nil) + } + for nonce, a := range j.Attempts { + if !validAttempt(nonce, a) { + return invalidIntent(nil) + } + } + for g, v := range j.Generations { + if !digestValid(g) || v == nil || !optionalValid(v.ActiveEpoch) || !optionalValid(v.Reservation) || (v.ActiveEpoch != nil && v.Reservation != nil) { + return invalidIntent(nil) + } + if v.ActiveEpoch != nil { + a := j.Attempts[*v.ActiveEpoch] + if a == nil || a.Phase != "completed" || deref(a.ResultGeneration) != g { + return invalidIntent(nil) + } + } + if v.Reservation != nil { + a := j.Attempts[*v.Reservation] + if !v.Consumed || a == nil || a.Kind != "refresh" || deref(a.InputGeneration) != g || a.Phase == "completed" { + return invalidIntent(nil) + } + } + } + indexes := map[string]map[string]refreshWrite{} + for k, b := range j.Bindings { + if !digestValid(k) || b == nil || !digestValid(b.Generation) || !digestValid(b.Epoch) || !digestValid(b.Entry) || j.Generations[b.Generation] == nil { + return invalidIntent(nil) + } + a := j.Attempts[b.Epoch] + if a == nil || a.Phase != "completed" { + return invalidIntent(nil) + } + index := indexes[b.Epoch] + if index == nil { + index = map[string]refreshWrite{} + for _, w := range a.Writes { + index[w.Key] = w + } + indexes[b.Epoch] = index + } + w, found := index[k] + if !found || deref(w.AfterGeneration) != b.Generation || deref(w.AfterEpoch) != b.Epoch || w.AfterEntry != b.Entry { + return invalidIntent(nil) + } + } + return nil +} + +func sameRefreshState(a, b *refreshJournal) bool { + return a.Version == b.Version && a.Scope == b.Scope && a.Nonce == b.Nonce && reflect.DeepEqual(a.Generations, b.Generations) && reflect.DeepEqual(a.Bindings, b.Bindings) && reflect.DeepEqual(a.Attempts, b.Attempts) +} + +func skeleton(a *refreshAttempt) refreshAttempt { + copy := *a + copy.Phase = "" + copy.ResponseEntry = nil + copy.ResultGeneration = nil + copy.PostCache = nil + copy.Writes = nil + copy.Proof = nil + return copy +} + +func applyRefreshDelta(old *refreshJournal, d *refreshDelta) (*refreshJournal, error) { + if err := validateRefreshDelta(old, d); err != nil { + return nil, err + } + // A failed load is discarded; all checks precede these map updates. + for k, c := range d.Generations { + old.Generations[k] = c.Value + } + for k, c := range d.Bindings { + if c.Value == nil { + delete(old.Bindings, k) + } else { + old.Bindings[k] = c.Value + } + } + for k, c := range d.Attempts { + old.Attempts[k] = c.Value + } + return old, nil +} + +func makeRefreshDelta(old, next *refreshJournal, seq uint64) (*refreshDelta, error) { + d := &refreshDelta{Sequence: seq, Retirements: next.retirements, Generations: map[string]generationChange{}, Bindings: map[string]bindingChange{}, Attempts: map[string]attemptChange{}} + for k := range old.Generations { + if next.Generations[k] == nil { + return nil, invalidIntent(nil) + } + } + for k := range old.Attempts { + if next.Attempts[k] == nil { + return nil, invalidIntent(nil) + } + } + for k, v := range next.Generations { + if !reflect.DeepEqual(old.Generations[k], v) { + e, err := expectedRecord("generation", k, old.Generations[k]) + if err != nil { + return nil, err + } + d.Generations[k] = generationChange{Expect: e, Value: v} + } + } + for k, v := range next.Attempts { + if !reflect.DeepEqual(old.Attempts[k], v) { + e, err := expectedRecord("attempt", k, old.Attempts[k]) + if err != nil { + return nil, err + } + d.Attempts[k] = attemptChange{Expect: e, Value: v} + } + } + keys := map[string]bool{} + for k := range old.Bindings { + keys[k] = true + } + for k := range next.Bindings { + keys[k] = true + } + for k := range keys { + if !reflect.DeepEqual(old.Bindings[k], next.Bindings[k]) { + e, err := expectedRecord("binding", k, old.Bindings[k]) + if err != nil { + return nil, err + } + d.Bindings[k] = bindingChange{Expect: e, Value: next.Bindings[k]} + } + } + if len(d.Generations)+len(d.Bindings)+len(d.Attempts) == 0 { + return nil, nil + } + if len(d.Retirements) != 0 { + d.Transition = "retire_bindings" + } + for k, c := range d.Attempts { + if d.Transition == "retire_bindings" { + break + } + if c.Value.Phase == "completed" { + d.Transition = "complete" + d.Attempt = nullable(k) + } + } + if d.Transition == "" { + if len(d.Attempts) == 0 { + d.Transition = "retire_bindings" + } else if len(d.Attempts) == 1 { + for k, c := range d.Attempts { + d.Attempt = nullable(k) + switch c.Value.Phase { + case "prepared": + d.Transition = "prepare_refresh" + case "response_ready": + d.Transition = "response_ready" + case "uncertain": + d.Transition = "uncertain" + } + } + } + } + if err := validateRefreshDelta(old, d); err != nil { + return nil, err + } + return d, nil +} + +func refreshFrameHash(domain string, prior, length, payload []byte) [32]byte { + h := sha256.New() + _, _ = h.Write([]byte(domain)) + _, _ = h.Write(prior) + _, _ = h.Write(length) + _, _ = h.Write(payload) + var result [32]byte + copy(result[:], h.Sum(nil)) + return result +} + +func encodeRefreshFrame(payload, prior []byte, header bool) ([]byte, [32]byte, error) { + if uint64(len(payload)) > math.MaxUint32 { + return nil, [32]byte{}, invalidIntent(nil) + } + var length [4]byte + binary.BigEndian.PutUint32(length[:], uint32(len(payload))) + domain := "cli-file-intent/delta/v1\x00" + var prefix []byte + if header { + domain = "cli-file-intent/header/v1\x00" + prefix = []byte(refreshIntentMagic) + } + digest := refreshFrameHash(domain, prior, length[:], payload) + result := append(prefix, length[:]...) + result = append(result, payload...) + result = append(result, digest[:]...) + return result, digest, nil +} + +func openRegularRefreshFile(path string, flag int, perm os.FileMode) (*os.File, error) { + before, err := os.Lstat(path) + if err != nil { + return nil, err + } + if !before.Mode().IsRegular() { + return nil, invalidIntent(nil) + } + file, err := os.OpenFile(path, flag, perm) + if err != nil { + return nil, err + } + after, err := file.Stat() + if err != nil || !after.Mode().IsRegular() || !os.SameFile(before, after) { + _ = file.Close() + return nil, invalidIntent(err) + } + return file, nil +} + +func (c *fileStore) loadRefreshJournal(ctx context.Context) (*refreshJournal, error) { + if err := ctx.Err(); err != nil { + return nil, err + } + file, err := openRegularRefreshFile(c.fileLocation+refreshIntentSuffix, os.O_RDWR, 0) + if errors.Is(err, os.ErrNotExist) { + return nil, nil + } + if err != nil { + if errors.Is(err, ErrRefreshIntentInvalid) { + return nil, err + } + return nil, intentError(ErrRefreshIntentIO, err) + } + raw, readErr := io.ReadAll(file) + closeErr := file.Close() + if readErr != nil || closeErr != nil { + return nil, intentError(ErrRefreshIntentIO, errors.Join(readErr, closeErr)) + } + return decodeRefreshJournal(raw, c.fileLocation) +} + +func decodeRefreshJournal(raw []byte, location string) (*refreshJournal, error) { + if len(raw) < len(refreshIntentMagic) || string(raw[:len(refreshIntentMagic)]) != refreshIntentMagic { + return nil, invalidIntent(nil) + } + offset := len(refreshIntentMagic) + readFrame := func(prior []byte, header bool) ([]byte, [32]byte, error) { + if len(raw)-offset < 4 { + return nil, [32]byte{}, invalidIntent(nil) + } + length := binary.BigEndian.Uint32(raw[offset : offset+4]) + start := offset + offset += 4 + if uint64(length)+32 > uint64(len(raw)-offset) { + return nil, [32]byte{}, invalidIntent(nil) + } + payload := raw[offset : offset+int(length)] + offset += int(length) + var got [32]byte + copy(got[:], raw[offset:offset+32]) + offset += 32 + domain := "cli-file-intent/delta/v1\x00" + if header { + domain = "cli-file-intent/header/v1\x00" + } + want := refreshFrameHash(domain, prior, raw[start:start+4], payload) + if got != want { + return nil, [32]byte{}, invalidIntent(nil) + } + return payload, got, nil + } + payload, chain, err := readFrame([]byte(refreshIntentMagic), true) + if err != nil { + return nil, err + } + var header refreshHeader + if err := canonicalDecode(payload, &header); err != nil { + return nil, err + } + if header.Version != 1 || header.Scope != intentHash("cli-file-intent/store/v1", location) || !digestValid(header.Nonce) { + return nil, invalidIntent(nil) + } + j := &refreshJournal{Version: 1, Scope: header.Scope, Nonce: header.Nonce, Generations: map[string]*refreshGeneration{}, Bindings: map[string]*refreshBinding{}, Attempts: map[string]*refreshAttempt{}} + var sequence uint64 + for offset < len(raw) { + payload, nextChain, err := readFrame(chain[:], false) + if err != nil { + return nil, err + } + var delta refreshDelta + if err := canonicalDecode(payload, &delta); err != nil { + return nil, err + } + if sequence == math.MaxUint64 || delta.Sequence != sequence+1 { + return nil, invalidIntent(nil) + } + next, err := applyRefreshDelta(j, &delta) + if err != nil { + return nil, err + } + j = next + sequence = delta.Sequence + chain = nextChain + } + if err := validateRefreshJournal(j); err != nil { + return nil, err + } + disk, err := cloneRefreshJournal(j) + if err != nil { + return nil, err + } + j.disk = disk + j.sequence = sequence + j.chain = chain + j.end = int64(len(raw)) + return j, nil +} + +func checkedRefreshWrite(file *os.File, raw []byte) error { + n, err := file.Write(raw) + if err != nil { + return err + } + if n != len(raw) { + return io.ErrShortWrite + } + return nil +} + +func (c *fileStore) initializeRefreshJournal(ctx context.Context, j *refreshJournal) error { + if err := ctx.Err(); err != nil { + return err + } + header := refreshHeader{Version: 1, Scope: j.Scope, Nonce: j.Nonce} + payload, err := json.Marshal(header) + if err != nil { + return invalidIntent(err) + } + raw, chain, err := encodeRefreshFrame(payload, []byte(refreshIntentMagic), true) + if err != nil { + return err + } + file, err := os.OpenFile(c.fileLocation+refreshIntentSuffix, os.O_RDWR|os.O_CREATE|os.O_EXCL, ownerReadWrite) + if err != nil { + return intentError(ErrRefreshIntentIO, err) + } + err = checkedRefreshWrite(file, raw) + if err == nil { + err = file.Sync() + } + err = errors.Join(err, file.Close()) + if err != nil { + j.failed = err + return intentError(ErrRefreshIntentIO, err) + } + j.chain = chain + j.end = int64(len(raw)) + j.disk = &refreshJournal{Version: 1, Scope: j.Scope, Nonce: j.Nonce, Generations: map[string]*refreshGeneration{}, Bindings: map[string]*refreshBinding{}, Attempts: map[string]*refreshAttempt{}} + return nil +} + +func (c *fileStore) saveRefreshJournal(ctx context.Context, j *refreshJournal) error { + if j.failed != nil { + return intentError(ErrRefreshIntentIO, j.failed) + } + if err := ctx.Err(); err != nil { + return err + } + if j.disk == nil { + if err := c.initializeRefreshJournal(ctx, j); err != nil { + return err + } + } + if j.sequence == math.MaxUint64 { + return invalidIntent(nil) + } + delta, err := makeRefreshDelta(j.disk, j, j.sequence+1) + if err != nil { + return err + } + if delta == nil { + return nil + } + payload, err := json.Marshal(delta) + if err != nil { + return invalidIntent(err) + } + raw, chain, err := encodeRefreshFrame(payload, j.chain[:], false) + if err != nil { + return err + } + file, err := openRegularRefreshFile(c.fileLocation+refreshIntentSuffix, os.O_RDWR, 0) + if err != nil { + return intentError(ErrRefreshIntentIO, err) + } + // Revalidate the complete file before appending; no valid-prefix authorization. + existing, readErr := io.ReadAll(file) + if readErr != nil { + closeErr := file.Close() + return intentError(ErrRefreshIntentIO, errors.Join(readErr, closeErr)) + } + loaded, decodeErr := decodeRefreshJournal(existing, c.fileLocation) + if decodeErr != nil { + _ = file.Close() + return decodeErr + } + if loaded == nil || loaded.end != j.end || loaded.sequence != j.sequence || loaded.chain != j.chain || !sameRefreshState(loaded, j.disk) { + _ = file.Close() + return invalidIntent(nil) + } + end, err := file.Seek(0, io.SeekEnd) + if err != nil || end != j.end { + _ = file.Close() + return intentError(ErrRefreshIntentIO, err) + } + if err := ctx.Err(); err != nil { + _ = file.Close() + return err + } + err = checkedRefreshWrite(file, raw) + if err == nil { + err = file.Sync() + } + err = errors.Join(err, file.Close()) + if err != nil { + j.failed = err + return intentError(ErrRefreshIntentIO, err) + } + disk, err := cloneRefreshJournal(j) + if err != nil { + j.failed = err + return err + } + j.disk = disk + j.retirements = nil + j.sequence++ + j.chain = chain + j.end += int64(len(raw)) + return nil +} diff --git a/libs/auth/storage/refresh_intent_process_test.go b/libs/auth/storage/refresh_intent_process_test.go new file mode 100644 index 00000000000..1d2840d902e --- /dev/null +++ b/libs/auth/storage/refresh_intent_process_test.go @@ -0,0 +1,140 @@ +package storage + +import ( + "bytes" + "encoding/json" + "errors" + "io" + "os" + "os/exec" + "testing" +) + +const refreshIntentActor = "DATABRICKS_TEST_REFRESH_INTENT_ACTOR" + +type refreshIntentSetup struct { + Path string + Nonrotating bool +} +type refreshIntentMessage struct { + Kind string + Generation int +} + +func TestRefreshIntentProcessLoss(t *testing.T) { + if os.Getenv(refreshIntentActor) == "1" { + var setup refreshIntentSetup + if err := json.NewDecoder(os.Stdin).Decode(&setup); err != nil { + t.Fatal("closed actor setup failure") + } + store, err := NewFileStore(t.Context(), WithFileLocation(setup.Path)) + if err != nil { + t.Fatal("closed actor store failure") + } + c := store.(*fileStore) + result := intentToken("1") + generation := 1 + if setup.Nonrotating { + entry, err := c.Lookup("primary") + if err != nil { + t.Fatal("closed actor original entry failure") + } + result = entry + generation = 0 + } + c.beforeIntentCompletion = func() { + if err := json.NewEncoder(os.Stdout).Encode(refreshIntentMessage{Kind: "cache-published", Generation: generation}); err != nil { + os.Exit(3) + } + _, _ = io.Copy(io.Discard, os.Stdin) + } + calls := 0 + if err := intentRefresh(t, c, "primary", result, &calls); err != nil { + t.Fatal("closed actor publication failure") + } + t.Fatal("actor passed the intentional process-loss boundary") + } + for _, nonrotating := range []bool{false, true} { + t.Run(map[bool]string{false: "rotating", true: "nonrotating-identical-entry"}[nonrotating], func(t *testing.T) { + c := intentFixture(t) + original := intentToken("0") + seedIntent(t, c, "primary", original) + seedIntent(t, c, "copy", original) + home := t.TempDir() + command := exec.CommandContext(t.Context(), os.Args[0], "-test.run=^TestRefreshIntentProcessLoss$", "-test.timeout=30s") + command.Env = []string{refreshIntentActor + "=1", "HOME=" + home, "USERPROFILE=" + home, "PATH=" + os.Getenv("PATH"), "TMPDIR=" + os.TempDir(), "TMP=" + os.TempDir(), "TEMP=" + os.TempDir()} + input, err := command.StdinPipe() + if err != nil { + t.Fatal("closed actor input failure") + } + output, err := command.StdoutPipe() + if err != nil { + t.Fatal("closed actor output failure") + } + command.Stderr = &bytes.Buffer{} + if err := command.Start(); err != nil { + t.Fatal("closed actor start failure") + } + waited := false + t.Cleanup(func() { + _ = input.Close() + if !waited { + _ = command.Process.Kill() + _ = command.Wait() + } + }) + if err := json.NewEncoder(input).Encode(refreshIntentSetup{Path: c.fileLocation, Nonrotating: nonrotating}); err != nil { + t.Fatal("closed actor send failure") + } + var message refreshIntentMessage + if err := json.NewDecoder(output).Decode(&message); err != nil || message.Kind != "cache-published" { + t.Fatal("closed actor did not reach exact publication boundary") + } + expected := 1 + if nonrotating { + expected = 0 + } + if message.Generation != expected { + t.Fatal("closed actor generation mismatch") + } + if err := command.Process.Kill(); err != nil { + t.Fatal("closed actor process loss failed") + } + err = command.Wait() + waited = true + if _, ok := errors.AsType[*exec.ExitError](err); !ok { + t.Fatal("actor did not terminate by process loss") + } + store, err := NewFileStore(t.Context(), WithFileLocation(c.fileLocation)) + if err != nil { + t.Fatal("fresh actor store failed") + } + fresh := store.(*fileStore) + _, err = WithFileTransaction(t.Context(), fresh, func(view Store) error { return nil }, nil) + if err != nil { + t.Fatal("actual committed cache proof failed reconciliation") + } + j := readIntent(t, fresh) + completed := 0 + for _, a := range j.Attempts { + if a.Phase == "completed" { + completed++ + } + } + if completed != 1 { + t.Fatal("response proof was not completed exactly once") + } + calls := 0 + if !errors.Is(intentRefresh(t, fresh, "copy", intentToken("2"), &calls), ErrRefreshRecoveryRequired) || calls != 0 { + t.Fatal("stale alias replayed after process loss") + } + entry, err := fresh.Lookup("primary") + if err != nil { + t.Fatal("fresh completed entry missing") + } + if err := intentRefresh(t, fresh, "primary", entry, &calls); err != nil || calls != 1 { + t.Fatal("proven current epoch could not refresh") + } + }) + } +} diff --git a/libs/auth/storage/refresh_intent_test.go b/libs/auth/storage/refresh_intent_test.go new file mode 100644 index 00000000000..8148192bf62 --- /dev/null +++ b/libs/auth/storage/refresh_intent_test.go @@ -0,0 +1,654 @@ +package storage + +import ( + "bytes" + "context" + "crypto/sha256" + "encoding/json" + "errors" + "os" + "path/filepath" + "reflect" + "testing" + "time" + + "golang.org/x/oauth2" +) + +func intentFixture(t *testing.T) *fileStore { + t.Helper() + store, err := NewFileStore(t.Context(), WithFileLocation(filepath.Join(t.TempDir(), "tokens.json"))) + if err != nil { + t.Fatal("file fixture setup failed") + } + return store.(*fileStore) +} + +func intentToken(g string) Entry { + return Entry{Token: &oauth2.Token{AccessToken: "access-" + g, RefreshToken: "refresh-" + g, Expiry: time.Now().Add(time.Hour)}} +} + +func intentRefresh(t *testing.T, c *fileStore, key string, result Entry, calls *int) error { + t.Helper() + _, err := WithFileTransaction(t.Context(), c, func(view Store) error { + a, err := BeginFileRefresh(t.Context(), view, key, sha256.Sum256([]byte("synthetic-route"))) + if err != nil { + return err + } + *calls++ + if err := a.RecordResponse(result); err != nil { + return err + } + return view.Put(key, result) + }, nil) + return err +} + +func seedIntent(t *testing.T, c *fileStore, key string, e Entry) { + t.Helper() + if err := c.Put(key, e); err != nil { + t.Fatal("seed failed") + } +} + +func readIntent(t *testing.T, c *fileStore) *refreshJournal { + t.Helper() + j, err := c.loadRefreshJournal(t.Context()) + if err != nil || j == nil { + t.Fatal("intent load failed") + } + return j +} + +func TestRefreshIntentPrepareFailureHasNoExchange(t *testing.T) { + c := intentFixture(t) + seedIntent(t, c, "primary", intentToken("0")) + before, _ := os.ReadFile(c.fileLocation) + calls := 0 + _, err := WithFileTransaction(t.Context(), c, func(view Store) error { + if err := os.Mkdir(c.fileLocation+refreshIntentSuffix, 0o700); err != nil { + t.Fatal("fault setup failed") + } + _, err := BeginFileRefresh(t.Context(), view, "primary", sha256.Sum256([]byte("synthetic-route"))) + if err != nil { + return err + } + calls++ + return nil + }, nil) + after, _ := os.ReadFile(c.fileLocation) + if !errors.Is(err, ErrRefreshIntentIO) || calls != 0 || !bytes.Equal(before, after) { + t.Fatal("prepare failure permitted exchange or cache mutation") + } +} + +func TestRefreshIntentAliasesAndIndependentGeneration(t *testing.T) { + c := intentFixture(t) + old := intentToken("0") + for _, key := range []string{"profile", "legacy", "copy"} { + seedIntent(t, c, key, old) + } + seedIntent(t, c, "other", intentToken("9")) + _, err := WithFileTransaction(t.Context(), c, func(view Store) error { + _, err := BeginFileRefresh(t.Context(), view, "profile", sha256.Sum256([]byte("route"))) + if err != nil { + return err + } + return errors.New("closed uncertain exchange") + }, nil) + if err == nil { + t.Fatal("uncertain fixture did not fail") + } + j := readIntent(t, c) + g := j.generation(old) + prior := *j.Generations[g] + calls := 0 + for _, key := range []string{"profile", "legacy", "copy"} { + if !errors.Is(intentRefresh(t, c, key, intentToken("1"), &calls), ErrRefreshRecoveryRequired) { + t.Fatal("alias reopened consumed credential") + } + } + if err := intentRefresh(t, c, "other", intentToken("10"), &calls); err != nil { + t.Fatal("independent generation blocked") + } + current := readIntent(t, c) + if calls != 1 || !reflect.DeepEqual(prior, *current.Generations[g]) { + t.Fatal("independent refresh changed consumed intent") + } + // Generic access-only and pending writes cannot open an epoch. + seedIntent(t, c, "profile", Entry{Token: &oauth2.Token{AccessToken: "access-only"}}) + seedIntent(t, c, "pending", Entry{PendingLogin: &PendingLogin{DeviceCode: "synthetic", Host: "synthetic"}}) + seedIntent(t, c, "copy", old) + if !errors.Is(intentRefresh(t, c, "copy", intentToken("1"), &calls), ErrRefreshRecoveryRequired) || calls != 1 { + t.Fatal("generic writer reopened consumed credential") + } +} + +type intentPause struct{} + +func captureLoginReady(t *testing.T) (*fileStore, *refreshJournal, string) { + t.Helper() + c := intentFixture(t) + seedIntent(t, c, "primary", intentToken("0")) + calls := 0 + if err := intentRefresh(t, c, "primary", intentToken("1"), &calls); err != nil { + t.Fatal("initial completed generation failed") + } + _, err := WithFileTransaction(t.Context(), c, func(view Store) error { + _, err := BeginFileRefresh(t.Context(), view, "primary", sha256.Sum256([]byte("route"))) + if err != nil { + return err + } + return errors.New("closed stop") + }, nil) + if err == nil { + t.Fatal("pending fixture failed") + } + c.beforeIntentCompletion = func() { panic(intentPause{}) } + func() { + defer func() { + if _, ok := recover().(intentPause); !ok { + t.Fatal("response publication boundary was not reached") + } + }() + _, err := WithFileTransaction(t.Context(), c, func(view Store) error { + entry := Entry{Token: &oauth2.Token{AccessToken: "access-only"}} + if err := RecordFileLoginResult(view, "primary", entry); err != nil { + return err + } + return view.Put("primary", entry) + }, nil) + if err != nil { + t.Fatal("login response proof failed") + } + }() + c.beforeIntentCompletion = nil + j := readIntent(t, c) + for nonce, a := range j.Attempts { + if a.Kind == "login" && a.Phase == "response_ready" { + return c, j, nonce + } + } + t.Fatal("response proof missing") + return nil, nil, "" +} + +func TestRefreshIntentNonceAndReplacementOwnership(t *testing.T) { + _, j, nonce := captureLoginReady(t) + a := j.Attempts[nonce] + if !j.canComplete(nonce, a) || !validAttempt(nonce, a) { + t.Fatal("captured owner cannot complete") + } + future, _ := intentNonce() + if validAttempt(future, a) { + t.Fatal("access-only proof was transferable to another nonce") + } + for _, variant := range []string{"reservation", "binding-entry", "owning-entry"} { + t.Run(variant, func(t *testing.T) { + copy, err := cloneRefreshJournal(j) + if err != nil { + t.Fatal("clone failed") + } + candidate := copy.Attempts[nonce] + w := candidate.Writes[0] + g := deref(w.BeforeGeneration) + switch variant { + case "reservation": + old := copy.Attempts[deref(w.BeforeReservation)] + next := *old + copy.Attempts[future] = &next + copy.Generations[g].Reservation = nullable(future) + case "binding-entry": + copy.Bindings[w.Key].Entry = intentHash("different-entry") + case "owning-entry": + copy.Attempts[deref(w.BeforeReservation)].InputEntry = nullable(intentHash("different-entry")) + } + if copy.canComplete(nonce, candidate) { + t.Fatal("completion borrowed future or changed ownership") + } + }) + } +} + +func TestRefreshIntentInvalidEntryAndCapabilityLifetime(t *testing.T) { + if (&FileRefreshStateError{}).Error() == "" { + t.Fatal("zero error has no safe message") + } + c := intentFixture(t) + seedIntent(t, c, "primary", intentToken("0")) + var saved *FileRefreshAttempt + calls := 0 + _, err := WithFileTransaction(t.Context(), c, func(view Store) error { + a, err := BeginFileRefresh(t.Context(), view, "primary", sha256.Sum256([]byte("route"))) + if err != nil { + return err + } + saved = a + calls++ + bad := intentToken("1") + bad.Token.Expiry = time.Date(10000, 1, 1, 0, 0, 0, 0, time.UTC) + return a.RecordResponse(bad) + }, nil) + if !errors.Is(err, ErrRefreshIntentInvalid) || calls != 1 || saved.response != nil { + t.Fatal("invalid entry certified a response") + } + if saved.RecordResponse(intentToken("1")) == nil { + t.Fatal("expired capability accepted response") + } + if !errors.Is(intentRefresh(t, c, "primary", intentToken("1"), &calls), ErrRefreshRecoveryRequired) || calls != 1 { + t.Fatal("encoding failure reopened consumed generation") + } +} + +func TestRefreshIntentWholeFileAndTransitionValidation(t *testing.T) { + c := intentFixture(t) + seedIntent(t, c, "primary", intentToken("0")) + calls := 0 + if err := intentRefresh(t, c, "primary", intentToken("1"), &calls); err != nil { + t.Fatal("completed setup failed") + } + raw, err := os.ReadFile(c.fileLocation + refreshIntentSuffix) + if err != nil { + t.Fatal("read failed") + } + for _, variant := range []string{"partial-length", "partial-frame", "checksum", "trailing"} { + t.Run(variant, func(t *testing.T) { + changed := append([]byte{}, raw...) + switch variant { + case "partial-length": + changed = changed[:9] + case "partial-frame": + changed = changed[:len(changed)-1] + case "checksum": + changed[len(changed)-1] ^= 1 + case "trailing": + changed = append(changed, 0) + } + if loaded, err := decodeRefreshJournal(changed, c.fileLocation); loaded != nil || !errors.Is(err, ErrRefreshIntentInvalid) { + t.Fatal("invalid tail authorized a valid prefix") + } + }) + } + j := readIntent(t, c) + next, err := cloneRefreshJournal(j) + if err != nil { + t.Fatal("clone failed") + } + for _, g := range next.Generations { + if g.Consumed { + g.Consumed = false + break + } + } + if _, err := makeRefreshDelta(j, next, j.sequence+1); !errors.Is(err, ErrRefreshIntentInvalid) { + t.Fatal("checksum-valid semantic rollback was permitted") + } + file, err := os.OpenFile(c.fileLocation+refreshIntentSuffix, os.O_WRONLY, 0) + if err != nil { + t.Fatal("tail fault open failed") + } + _, _ = file.Seek(0, 2) + _, err = file.Write([]byte{0}) + closeErr := file.Close() + if err != nil || closeErr != nil { + t.Fatal("tail fault failed") + } + _, err = WithFileTransaction(t.Context(), c, func(view Store) error { calls++; return nil }, nil) + if !errors.Is(err, ErrRefreshIntentInvalid) || calls != 1 { + t.Fatal("partial file did not block entire transaction") + } +} + +func TestRefreshIntentOriginalContextAndPersistenceOrdering(t *testing.T) { + for _, afterPrepare := range []bool{false, true} { + t.Run(map[bool]string{false: "cancel-before-prepare", true: "cancel-after-prepare"}[afterPrepare], func(t *testing.T) { + c := intentFixture(t) + seedIntent(t, c, "primary", intentToken("0")) + ctx, cancel := context.WithCancel(t.Context()) + defer cancel() + calls := 0 + _, err := WithFileTransaction(ctx, c, func(view Store) error { + if !afterPrepare { + cancel() + } + _, err := BeginFileRefresh(t.Context(), view, "primary", sha256.Sum256([]byte("route"))) + if err != nil { + return err + } + cancel() + return ctx.Err() + }, nil) + if !errors.Is(err, context.Canceled) || calls != 0 { + t.Fatal("original context was reset") + } + if afterPrepare && !errors.Is(intentRefresh(t, c, "primary", intentToken("1"), &calls), ErrRefreshRecoveryRequired) { + t.Fatal("cancellation reopened prepared generation") + } + }) + } + for _, afterCache := range []bool{false, true} { + t.Run(map[bool]string{false: "response-proof-IO", true: "completion-IO"}[afterCache], func(t *testing.T) { + c := intentFixture(t) + seedIntent(t, c, "primary", intentToken("0")) + path := c.fileLocation + refreshIntentSuffix + calls := 0 + t.Cleanup(func() { _ = os.Chmod(path, 0o600) }) + if afterCache { + c.beforeIntentCompletion = func() { + if err := os.Chmod(path, 0o400); err != nil { + t.Fatal("fault setup failed") + } + } + } + _, err := WithFileTransaction(t.Context(), c, func(view Store) error { + a, err := BeginFileRefresh(t.Context(), view, "primary", sha256.Sum256([]byte("route"))) + if err != nil { + return err + } + calls++ + result := intentToken("1") + if err := a.RecordResponse(result); err != nil { + return err + } + if err := view.Put("primary", result); err != nil { + return err + } + if !afterCache { + return os.Chmod(path, 0o400) + } + return nil + }, nil) + c.beforeIntentCompletion = nil + if !errors.Is(err, ErrRefreshIntentIO) || calls != 1 { + t.Fatal("persistence failure incorrectly succeeded") + } + if err := os.Chmod(path, 0o600); err != nil { + t.Fatal("storage repair failed") + } + if afterCache { + _, err = WithFileTransaction(t.Context(), c, func(view Store) error { return nil }, nil) + if err != nil { + t.Fatal("completed cache proof did not reconcile") + } + } else if !errors.Is(intentRefresh(t, c, "primary", intentToken("1"), &calls), ErrRefreshRecoveryRequired) || calls != 1 { + t.Fatal("lost response proof permitted replay") + } + }) + } +} + +func TestRefreshIntentRetirementOwnerAndSameCredentialLogin(t *testing.T) { + for _, ownerDelete := range []bool{true, false} { + t.Run(map[bool]string{true: "owner", false: "copied-alias"}[ownerDelete], func(t *testing.T) { + c := intentFixture(t) + original := intentToken("0") + seedIntent(t, c, "owner", original) + seedIntent(t, c, "copy", original) + _, err := WithFileTransaction(t.Context(), c, func(view Store) error { + _, err := BeginFileRefresh(t.Context(), view, "owner", sha256.Sum256([]byte("route"))) + if err != nil { + return err + } + return errors.New("closed unknown exchange") + }, nil) + if err == nil { + t.Fatal("prepared fixture failed") + } + prior := readIntent(t, c) + g := prior.generation(original) + reservation := deref(prior.Generations[g].Reservation) + key := "copy" + if ownerDelete { + key = "owner" + } + if err := c.Delete(key); err != nil { + t.Fatal("captured deletion failed") + } + current := readIntent(t, c) + if !current.Generations[g].Consumed || current.Generations[g].ActiveEpoch != nil { + t.Fatal("Delete reopened consumed generation") + } + if ownerDelete { + if current.Generations[g].Reservation != nil { + t.Fatal("owner deletion retained its captured reservation") + } + calls := 0 + if !errors.Is(intentRefresh(t, c, "copy", intentToken("1"), &calls), ErrRefreshRecoveryRequired) || calls != 0 { + t.Fatal("Delete authorized an alias") + } + result := intentToken("0") + _, err = WithFileTransaction(t.Context(), c, func(view Store) error { + if err := RecordFileLoginResult(view, "owner", result); err != nil { + return err + } + return view.Put("owner", result) + }, nil) + if err != nil { + t.Fatal("actual fresh login could not open same opaque credential") + } + if !errors.Is(intentRefresh(t, c, "copy", intentToken("1"), &calls), ErrRefreshRecoveryRequired) || calls != 0 { + t.Fatal("fresh login opened an unrelated copied key") + } + if err := intentRefresh(t, c, "owner", intentToken("1"), &calls); err != nil || calls != 1 { + t.Fatal("fresh login did not open only its result epoch") + } + } else { + if deref(current.Generations[g].Reservation) != reservation { + t.Fatal("deleting copied alias cleared original owner") + } + result := intentToken("0") + _, err = WithFileTransaction(t.Context(), c, func(view Store) error { + if err := RecordFileLoginResult(view, "copy", result); err != nil { + return err + } + return view.Put("copy", result) + }, nil) + if !errors.Is(err, ErrRefreshRecoveryRequired) { + t.Fatal("alias login borrowed original owner's reservation") + } + } + }) + } +} + +func TestRefreshIntentRetirementRejectsStaleCapture(t *testing.T) { + c := intentFixture(t) + entry := intentToken("0") + seedIntent(t, c, "owner", entry) + _, _ = WithFileTransaction(t.Context(), c, func(view Store) error { + _, err := BeginFileRefresh(t.Context(), view, "owner", sha256.Sum256([]byte("route"))) + if err != nil { + return err + } + return errors.New("closed stop") + }, nil) + j := readIntent(t, c) + g := j.generation(entry) + reservation := deref(j.Generations[g].Reservation) + owner := j.Attempts[reservation] + proof := refreshRetirement{Key: owner.InputKey, Generation: g, Entry: deref(owner.InputEntry), Reservation: nullable(reservation)} + for _, variant := range []string{"future-reservation", "changed-entry"} { + t.Run(variant, func(t *testing.T) { + copy, err := cloneRefreshJournal(j) + if err != nil { + t.Fatal("clone failed") + } + captured := proof + if variant == "future-reservation" { + future, _ := intentNonce() + next := *copy.Attempts[reservation] + copy.Attempts[future] = &next + copy.Generations[g].Reservation = nullable(future) + } else { + captured.Entry = intentHash("changed-entry") + } + before, err := cloneRefreshJournal(copy) + if err != nil { + t.Fatal("clone failed") + } + if _, err := copy.retireCaptured(captured); !errors.Is(err, ErrRefreshRecoveryRequired) || !sameRefreshState(copy, before) { + t.Fatal("stale retirement changed a future owner") + } + }) + } +} + +type intentAliasKeys struct{} + +func (intentAliasKeys) GetCacheKey() string { return "OWNER" } +func (intentAliasKeys) GetHostCacheKey() string { return "legacy" } +func TestRefreshIntentRetirementIsAtomicAcrossAliases(t *testing.T) { + c := intentFixture(t) + old := intentToken("0") + seedIntent(t, c, "OWNER", old) + seedIntent(t, c, "legacy", old) + wrapped := NewDualWritingStore(c, intentAliasKeys{}) + result := intentToken("1") + _, err := WithFileTransaction(t.Context(), wrapped, func(view Store) error { + a, err := BeginFileRefresh(t.Context(), view, "OWNER", sha256.Sum256([]byte("route"))) + if err != nil { + return err + } + if err := a.RecordResponse(result); err != nil { + return err + } + return view.Put("OWNER", result) + }, nil) + if err != nil { + t.Fatal("wrapped completion failed") + } + _, _ = WithFileTransaction(t.Context(), wrapped, func(view Store) error { + _, err := BeginFileRefresh(t.Context(), view, "OWNER", sha256.Sum256([]byte("route"))) + if err != nil { + return err + } + return errors.New("closed stop") + }, nil) + if err := DeleteWithContext(t.Context(), wrapped, "OWNER"); err != nil { + t.Fatal("owner-first alias batch failed atomic retirement") + } + j := readIntent(t, c) + g := j.generation(result) + if j.Generations[g].Reservation != nil || !j.Generations[g].Consumed { + t.Fatal("batch retirement lost consumed history or ownership") + } +} + +func TestRefreshIntentSmallHistoryAndPriorRecordChecks(t *testing.T) { + c := intentFixture(t) + seedIntent(t, c, "primary", intentToken("0")) + calls := 0 + for _, g := range []string{"1", "2", "3"} { + if err := intentRefresh(t, c, "primary", intentToken(g), &calls); err != nil { + t.Fatal("finite history failed") + } + } + j := readIntent(t, c) + if calls != 3 || j.sequence != 9 || len(j.Attempts) != 3 || len(j.Generations) != 4 { + t.Fatal("finite replay lost state or duplicated transitions") + } + raw, err := os.ReadFile(c.fileLocation + refreshIntentSuffix) + if err != nil { + t.Fatal("history read failed") + } + key := j.key("primary") + binding := j.Bindings[key] + proof := refreshRetirement{Key: key, Generation: binding.Generation, Entry: binding.Entry, Epoch: nullable(binding.Epoch), BindingEntry: nullable(binding.Entry)} + delta := refreshDelta{Sequence: j.sequence + 1, Transition: "retire_bindings", Retirements: []refreshRetirement{proof}, Generations: map[string]generationChange{}, Attempts: map[string]attemptChange{}, Bindings: map[string]bindingChange{key: {Expect: nullable(intentHash("wrong-prior-record"))}}} + payload, err := json.Marshal(delta) + if err != nil { + t.Fatal("fixture encoding failed") + } + frame, _, err := encodeRefreshFrame(payload, j.chain[:], false) + if err != nil { + t.Fatal("fixture framing failed") + } + if decoded, err := decodeRefreshJournal(append(raw, frame...), c.fileLocation); decoded != nil || !errors.Is(err, ErrRefreshIntentInvalid) { + t.Fatal("valid frame checksum bypassed prior-record expectation") + } +} + +func TestRefreshIntentLoginHeaderFailureKeepsCleanup(t *testing.T) { + c := intentFixture(t) + seedIntent(t, c, "primary", intentToken("0")) + seedIntent(t, c, "unrelated", intentToken("9")) + result := intentToken("1") + _, err := WithFileTransaction(t.Context(), c, func(view Store) error { + if err := RecordFileLoginResult(view, "primary", result); err != nil { + return err + } + if err := view.Put("primary", result); err != nil { + return err + } + return os.Mkdir(c.fileLocation+refreshIntentSuffix, 0o700) + }, func(view Store) error { + if err := os.Remove(c.fileLocation + refreshIntentSuffix); err != nil { + return err + } + return view.Delete("primary") + }) + if !errors.Is(err, ErrRefreshIntentIO) { + t.Fatal("header failure incorrectly succeeded") + } + if _, err := c.Lookup("primary"); !errors.Is(err, ErrNotFound) { + t.Fatal("header failure lost stale-token cleanup") + } + if _, err := c.Lookup("unrelated"); err != nil { + t.Fatal("header failure changed unrelated entry") + } +} + +func TestRefreshIntentUnknownVersionAndNullRecordsFailClosed(t *testing.T) { + for _, variant := range []string{"unknown-version", "null-generation", "null-attempt"} { + t.Run(variant, func(t *testing.T) { + c := intentFixture(t) + seedIntent(t, c, "primary", intentToken("0")) + j, err := newRefreshJournal(c.fileLocation) + if err != nil { + t.Fatal("header fixture failed") + } + header := refreshHeader{Version: 1, Scope: j.Scope, Nonce: j.Nonce} + if variant == "unknown-version" { + header.Version = 2 + } + payload, err := json.Marshal(header) + if err != nil { + t.Fatal("header encoding failed") + } + raw, chain, err := encodeRefreshFrame(payload, []byte(refreshIntentMagic), true) + if err != nil { + t.Fatal("header framing failed") + } + if variant != "unknown-version" { + nonce, err := intentNonce() + if err != nil { + t.Fatal("nonce fixture failed") + } + g := j.generation(intentToken("0")) + delta := refreshDelta{Sequence: 1, Transition: "prepare_refresh", Attempt: nullable(nonce), Generations: map[string]generationChange{}, Bindings: map[string]bindingChange{}, Attempts: map[string]attemptChange{}} + if variant == "null-generation" { + delta.Generations[g] = generationChange{} + } else { + delta.Attempts[nonce] = attemptChange{} + } + payload, err = json.Marshal(delta) + if err != nil { + t.Fatal("delta encoding failed") + } + frame, _, err := encodeRefreshFrame(payload, chain[:], false) + if err != nil { + t.Fatal("delta framing failed") + } + raw = append(raw, frame...) + } + if err := os.WriteFile(c.fileLocation+refreshIntentSuffix, raw, 0o600); err != nil { + t.Fatal("invalid-state fixture failed") + } + calls := 0 + _, err = WithFileTransaction(t.Context(), c, func(view Store) error { calls++; return nil }, nil) + if !errors.Is(err, ErrRefreshIntentInvalid) || calls != 0 { + t.Fatal("unknown or null intent authorized callback") + } + }) + } +} diff --git a/libs/auth/storage/refresh_intent_transition.go b/libs/auth/storage/refresh_intent_transition.go new file mode 100644 index 00000000000..c6dc1e65d21 --- /dev/null +++ b/libs/auth/storage/refresh_intent_transition.go @@ -0,0 +1,372 @@ +package storage + +import "reflect" + +// refreshRetirement is the exact leased snapshot of an entry being replaced. +type refreshRetirement struct { + Key string `json:"key_d"` + Generation string `json:"generation"` + Entry string `json:"entry_d"` + Epoch *string `json:"epoch"` + BindingEntry *string `json:"binding_entry_d"` + Reservation *string `json:"reservation"` +} + +func (j *refreshJournal) checkRetirement(p refreshRetirement) (bool, bool, error) { + if !digestValid(p.Key) || !digestValid(p.Generation) || !digestValid(p.Entry) || !optionalValid(p.Epoch) || !optionalValid(p.BindingEntry) || !optionalValid(p.Reservation) || (p.Epoch == nil) != (p.BindingEntry == nil) { + return false, false, invalidIntent(nil) + } + binding := j.Bindings[p.Key] + if binding != nil { + if binding.Generation != p.Generation || binding.Epoch != deref(p.Epoch) || binding.Entry != deref(p.BindingEntry) || binding.Entry != p.Entry { + return false, false, intentError(ErrRefreshRecoveryRequired, nil) + } + } else if p.Epoch != nil { + return false, false, intentError(ErrRefreshRecoveryRequired, nil) + } + gen := j.Generations[p.Generation] + reservation := "" + if gen != nil { + reservation = deref(gen.Reservation) + } + if reservation != deref(p.Reservation) { + return false, false, intentError(ErrRefreshRecoveryRequired, nil) + } + clear := false + if reservation != "" { + owner := j.Attempts[reservation] + if owner == nil || deref(owner.InputGeneration) != p.Generation { + return false, false, invalidIntent(nil) + } + if owner.InputKey == p.Key { + if deref(owner.InputEntry) != p.Entry || deref(owner.InputEpoch) != deref(p.Epoch) { + return false, false, intentError(ErrRefreshRecoveryRequired, nil) + } + clear = true + } + } + return binding != nil, clear, nil +} + +func (j *refreshJournal) retireCapturedSet(proofs []refreshRetirement) ([]refreshRetirement, error) { + type retirement struct { + proof refreshRetirement + drop, clear bool + } + var actions []retirement + var used []refreshRetirement + seen := map[string]bool{} + // Every expectation is checked against the same pre-transition state. + for _, p := range proofs { + if seen[p.Key] { + return nil, invalidIntent(nil) + } + seen[p.Key] = true + drop, clear, err := j.checkRetirement(p) + if err != nil { + return nil, err + } + if drop || clear { + actions = append(actions, retirement{p, drop, clear}) + used = append(used, p) + } + } + for _, a := range actions { + if a.drop { + delete(j.Bindings, a.proof.Key) + } + if a.clear { + reservation := deref(a.proof.Reservation) + j.Attempts[reservation].Phase = "uncertain" + j.Generations[a.proof.Generation].Reservation = nil + } + } + return used, nil +} + +func (j *refreshJournal) retireCaptured(p refreshRetirement) (bool, error) { + used, err := j.retireCapturedSet([]refreshRetirement{p}) + return len(used) != 0, err +} + +// Only records referenced by this transition are copied for semantic validation. +func selectRefreshState(old *refreshJournal, a *refreshAttempt, retire []refreshRetirement, nonce string) *refreshJournal { + view := &refreshJournal{Version: old.Version, Scope: old.Scope, Nonce: old.Nonce, Generations: map[string]*refreshGeneration{}, Bindings: map[string]*refreshBinding{}, Attempts: map[string]*refreshAttempt{}} + takeAttempt := func(key string) { + if value := old.Attempts[key]; value != nil { + copy := *value + view.Attempts[key] = © + } + } + takeGeneration := func(key string) { + if value := old.Generations[key]; value != nil { + copy := *value + view.Generations[key] = © + takeAttempt(deref(value.Reservation)) + } + } + takeBinding := func(key string) { + if value := old.Bindings[key]; value != nil { + copy := *value + view.Bindings[key] = © + } + } + if a != nil { + takeAttempt(nonce) + takeGeneration(deref(a.InputGeneration)) + takeGeneration(deref(a.ResultGeneration)) + takeBinding(a.InputKey) + for _, w := range a.Writes { + takeGeneration(deref(w.BeforeGeneration)) + takeGeneration(deref(w.AfterGeneration)) + takeBinding(w.Key) + } + } + for _, r := range retire { + takeGeneration(r.Generation) + takeBinding(r.Key) + } + return view +} + +func selectedChanges(old, view *refreshJournal, gkeys, bkeys, akeys map[string]bool) (*refreshDelta, error) { + d := &refreshDelta{Generations: map[string]generationChange{}, Bindings: map[string]bindingChange{}, Attempts: map[string]attemptChange{}} + for k := range view.Generations { + gkeys[k] = true + } + for k := range view.Bindings { + bkeys[k] = true + } + for k := range view.Attempts { + akeys[k] = true + } + for k := range gkeys { + if !reflect.DeepEqual(old.Generations[k], view.Generations[k]) { + e, err := expectedRecord("generation", k, old.Generations[k]) + if err != nil { + return nil, err + } + d.Generations[k] = generationChange{Expect: e, Value: view.Generations[k]} + } + } + for k := range bkeys { + if !reflect.DeepEqual(old.Bindings[k], view.Bindings[k]) { + e, err := expectedRecord("binding", k, old.Bindings[k]) + if err != nil { + return nil, err + } + d.Bindings[k] = bindingChange{Expect: e, Value: view.Bindings[k]} + } + } + for k := range akeys { + if !reflect.DeepEqual(old.Attempts[k], view.Attempts[k]) { + e, err := expectedRecord("attempt", k, old.Attempts[k]) + if err != nil { + return nil, err + } + d.Attempts[k] = attemptChange{Expect: e, Value: view.Attempts[k]} + } + } + return d, nil +} + +func selectedKeys(j *refreshJournal) (map[string]bool, map[string]bool, map[string]bool) { + g, b, a := map[string]bool{}, map[string]bool{}, map[string]bool{} + for k := range j.Generations { + g[k] = true + } + for k := range j.Bindings { + b[k] = true + } + for k := range j.Attempts { + a[k] = true + } + return g, b, a +} + +func expectedRefreshChanges(old *refreshJournal, d *refreshDelta) (*refreshDelta, error) { + nonce := deref(d.Attempt) + changed := d.Attempts[nonce].Value + expected := &refreshDelta{Generations: map[string]generationChange{}, Bindings: map[string]bindingChange{}, Attempts: map[string]attemptChange{}} + switch d.Transition { + case "prepare_refresh": + if changed == nil || old.Attempts[nonce] != nil || changed.Kind != "refresh" || changed.Phase != "prepared" { + return nil, invalidIntent(nil) + } + g := deref(changed.InputGeneration) + prior := old.Generations[g] + gen := &refreshGeneration{} + if prior == nil { + if changed.InputEpoch != nil { + return nil, invalidIntent(nil) + } + } else { + *gen = *prior + b := old.Bindings[changed.InputKey] + if prior.Reservation != nil || prior.ActiveEpoch == nil || deref(prior.ActiveEpoch) != deref(changed.InputEpoch) || b == nil || b.Generation != g || b.Epoch != deref(changed.InputEpoch) || b.Entry != deref(changed.InputEntry) { + return nil, invalidIntent(nil) + } + } + gen.Consumed = true + gen.ActiveEpoch = nil + gen.Reservation = nullable(nonce) + e, err := expectedRecord("generation", g, prior) + if err != nil { + return nil, err + } + expected.Generations[g] = generationChange{Expect: e, Value: gen} + expected.Attempts[nonce] = attemptChange{Value: changed} + case "response_ready": + if changed == nil || changed.Phase != "response_ready" { + return nil, invalidIntent(nil) + } + prior := old.Attempts[nonce] + if changed.Kind == "refresh" { + if prior == nil || prior.Phase != "prepared" || !reflect.DeepEqual(skeleton(prior), skeleton(changed)) { + return nil, invalidIntent(nil) + } + } else if prior != nil { + return nil, invalidIntent(nil) + } + if !old.canComplete(nonce, changed) { + return nil, invalidIntent(nil) + } + e, err := expectedRecord("attempt", nonce, prior) + if err != nil { + return nil, err + } + expected.Attempts[nonce] = attemptChange{Expect: e, Value: changed} + case "complete": + prior := old.Attempts[nonce] + if prior == nil || prior.Phase != "response_ready" { + return nil, invalidIntent(nil) + } + view := selectRefreshState(old, prior, nil, nonce) + g, b, a := selectedKeys(view) + if err := view.complete(nonce, view.Attempts[nonce]); err != nil { + return nil, invalidIntent(err) + } + return selectedChanges(old, view, g, b, a) + case "uncertain": + prior := old.Attempts[nonce] + if prior == nil || (prior.Phase != "prepared" && prior.Phase != "response_ready") { + return nil, invalidIntent(nil) + } + copy := *prior + copy.Phase = "uncertain" + e, err := expectedRecord("attempt", nonce, prior) + if err != nil { + return nil, err + } + expected.Attempts[nonce] = attemptChange{Expect: e, Value: ©} + case "retire_bindings": + if d.Attempt != nil || len(d.Retirements) == 0 { + return nil, invalidIntent(nil) + } + view := selectRefreshState(old, nil, d.Retirements, "") + g, b, a := selectedKeys(view) + used, err := view.retireCapturedSet(d.Retirements) + if err != nil || len(used) != len(d.Retirements) { + return nil, invalidIntent(err) + } + return selectedChanges(old, view, g, b, a) + default: + return nil, invalidIntent(nil) + } + return expected, nil +} + +func validateRefreshDelta(old *refreshJournal, d *refreshDelta) error { + if d == nil || d.Generations == nil || d.Bindings == nil || d.Attempts == nil || len(d.Generations)+len(d.Bindings)+len(d.Attempts) == 0 { + return invalidIntent(nil) + } + if d.Transition == "retire_bindings" { + if d.Attempt != nil || len(d.Retirements) == 0 { + return invalidIntent(nil) + } + } else if d.Attempt == nil || !digestValid(*d.Attempt) || len(d.Retirements) != 0 { + return invalidIntent(nil) + } + for k, c := range d.Generations { + if !digestValid(k) || c.Value == nil || !expectationMatches("generation", k, old.Generations[k], c.Expect) || reflect.DeepEqual(old.Generations[k], c.Value) { + return invalidIntent(nil) + } + if old.Generations[k] != nil && old.Generations[k].Consumed && !c.Value.Consumed { + return invalidIntent(nil) + } + } + for k, c := range d.Bindings { + if !digestValid(k) || !expectationMatches("binding", k, old.Bindings[k], c.Expect) || reflect.DeepEqual(old.Bindings[k], c.Value) { + return invalidIntent(nil) + } + } + for k, c := range d.Attempts { + if !validAttempt(k, c.Value) || !expectationMatches("attempt", k, old.Attempts[k], c.Expect) || reflect.DeepEqual(old.Attempts[k], c.Value) { + return invalidIntent(nil) + } + } + expected, err := expectedRefreshChanges(old, d) + if err != nil { + return err + } + if !reflect.DeepEqual(expected.Generations, d.Generations) || !reflect.DeepEqual(expected.Bindings, d.Bindings) || !reflect.DeepEqual(expected.Attempts, d.Attempts) { + return invalidIntent(nil) + } + getAttempt := func(k string) *refreshAttempt { + if c, ok := d.Attempts[k]; ok { + return c.Value + } + return old.Attempts[k] + } + getGeneration := func(k string) *refreshGeneration { + if c, ok := d.Generations[k]; ok { + return c.Value + } + return old.Generations[k] + } + for g, c := range d.Generations { + v := c.Value + if !optionalValid(v.ActiveEpoch) || !optionalValid(v.Reservation) || (v.ActiveEpoch != nil && v.Reservation != nil) { + return invalidIntent(nil) + } + if v.ActiveEpoch != nil { + a := getAttempt(*v.ActiveEpoch) + if a == nil || a.Phase != "completed" || deref(a.ResultGeneration) != g { + return invalidIntent(nil) + } + } + if v.Reservation != nil { + a := getAttempt(*v.Reservation) + if !v.Consumed || a == nil || a.Kind != "refresh" || deref(a.InputGeneration) != g || a.Phase == "completed" { + return invalidIntent(nil) + } + } + } + indexes := map[string]map[string]refreshWrite{} + for k, c := range d.Bindings { + b := c.Value + if b == nil { + continue + } + if !digestValid(b.Generation) || !digestValid(b.Epoch) || !digestValid(b.Entry) || getGeneration(b.Generation) == nil { + return invalidIntent(nil) + } + a := getAttempt(b.Epoch) + if a == nil || a.Phase != "completed" { + return invalidIntent(nil) + } + index := indexes[b.Epoch] + if index == nil { + index = map[string]refreshWrite{} + for _, w := range a.Writes { + index[w.Key] = w + } + indexes[b.Epoch] = index + } + w, ok := index[k] + if !ok || deref(w.AfterGeneration) != b.Generation || deref(w.AfterEpoch) != b.Epoch || w.AfterEntry != b.Entry { + return invalidIntent(nil) + } + } + return nil +} diff --git a/libs/auth/u2m/file_refresh_commit_failure_test.go b/libs/auth/u2m/file_refresh_commit_failure_test.go new file mode 100644 index 00000000000..acbfb60519a --- /dev/null +++ b/libs/auth/u2m/file_refresh_commit_failure_test.go @@ -0,0 +1,165 @@ +package u2m_test + +import ( + "bytes" + "errors" + "io" + "net/http" + "net/url" + "os" + "path/filepath" + "strings" + "testing" + "time" + + "github.com/databricks/cli/libs/auth/storage" + "github.com/databricks/cli/libs/auth/u2m" + "github.com/databricks/databricks-sdk-go/httpclient/fixtures" + "golang.org/x/oauth2" +) + +type fileCommitFailureTransport struct { + path string + backup string + fault bool + armed bool + spent bool + calls []int + successes int + replays int +} + +func (p *fileCommitFailureTransport) RoundTrip(request *http.Request) (*http.Response, error) { + body, err := io.ReadAll(request.Body) + if err != nil { + return nil, errors.New("closed request read failure") + } + request.Body = io.NopCloser(bytes.NewReader(body)) + form, err := url.ParseQuery(string(body)) + if err != nil || form.Get("grant_type") != "refresh_token" { + return nil, errors.New("closed unexpected grant") + } + generation := tokenGeneration(form.Get("refresh_token")) + p.calls = append(p.calls, generation) + if generation != 0 { + return nil, errors.New("closed unexpected generation") + } + status := http.StatusOK + response := `{"access_token":"generation-1","refresh_token":"generation-1","token_type":"Bearer","expires_in":3600}` + if p.spent { + p.replays++ + status = http.StatusBadRequest + response = `{"error":"invalid_grant","error_description":"Refresh token is invalid"}` + } + transport := fixtures.SliceTransport{{ + Method: "POST", Resource: "/oidc/accounts/account/v1/token", + Status: status, Response: response, + ResponseHeaders: map[string][]string{"Content-Type": {"application/json"}}, + }} + reply, err := transport.RoundTrip(request) + if err != nil || status != http.StatusOK { + return reply, err + } + p.spent = true + p.successes++ + if p.fault { + // Reuse the owning transaction fixture's actual atomic-publication fault. + if err := os.Rename(p.path, p.backup); err != nil { + return nil, errors.New("closed fault rename failure") + } + if err := os.Mkdir(p.path, 0o700); err != nil { + return nil, errors.New("closed fault directory failure") + } + p.armed = true + } + return reply, nil +} + +func TestFileRefreshCommitFailure(t *testing.T) { + for _, fault := range []bool{false, true} { + name := "healthy" + if fault { + name = "commit_failure_fresh_instance" + } + t.Run(name, func(t *testing.T) { + cachePath := filepath.Join(t.TempDir(), "tokens.json") + probe := &fileCommitFailureTransport{path: cachePath, backup: cachePath + ".saved", fault: fault} + arg, err := u2m.NewBasicAccountOAuthArgument("https://accounts.databricks.test", "account") + if err != nil { + t.Fatal("closed argument setup failure") + } + initial, err := storage.NewFileStore(t.Context(), storage.WithFileLocation(cachePath)) + if err != nil { + t.Fatal("closed initial FileStore setup failure") + } + if err := initial.Put(arg.GetCacheKey(), storage.Entry{Token: &oauth2.Token{ + AccessToken: "generation-0", RefreshToken: "generation-0", Expiry: time.Now().Add(-time.Hour), + }}); err != nil { + t.Fatal("closed initial generation setup failure") + } + original, err := os.ReadFile(cachePath) + if err != nil { + t.Fatal("closed cache snapshot failure") + } + newAuth := func() *u2m.PersistentAuth { + store, err := storage.NewFileStore(t.Context(), storage.WithFileLocation(cachePath)) + if err != nil { + t.Fatal("closed fresh FileStore setup failure") + } + auth, err := u2m.NewPersistentAuth(t.Context(), u2m.WithTokenStore(store), + u2m.WithOAuthArgument(arg), u2m.WithClientID("synthetic-client"), + u2m.WithOAuthEndpointSupplier(&u2m.BasicOAuthEndpointSupplier{}), + u2m.WithHttpClient(&http.Client{Transport: probe})) + if err != nil { + t.Fatal("closed fresh auth setup failure") + } + t.Cleanup(func() { _ = auth.Close() }) + return auth + } + first := newAuth() + firstToken, firstErr := first.Token() + commitFailed := firstErr != nil && strings.Contains(firstErr.Error(), "cache update") + if probe.successes != 1 || len(probe.calls) != 1 || probe.calls[0] != 0 { + t.Fatal("closed first exchange precondition failure") + } + if fault { + if !probe.armed || !commitFailed || firstToken != nil { + t.Fatal("closed commit-failure precondition failure") + } + if err := os.Remove(cachePath); err != nil { + t.Fatal("closed storage repair remove failure") + } + if err := os.Rename(probe.backup, cachePath); err != nil { + t.Fatal("closed storage repair rename failure") + } + probe.armed = false + restored, err := os.ReadFile(cachePath) + if err != nil || !bytes.Equal(original, restored) { + t.Fatal("closed original cache equality failure") + } + } else if firstErr != nil || firstToken == nil || tokenGeneration(firstToken.AccessToken) != 1 { + t.Fatal("closed healthy first result failure") + } + _ = first.Close() + fresh := newAuth() + freshToken, freshErr := fresh.Token() + cached, err := initial.Lookup(arg.GetCacheKey()) + if err != nil || cached.Token == nil { + t.Fatal("closed final cache lookup failure") + } + cachedGeneration := tokenGeneration(cached.Token.RefreshToken) + freshGeneration := -1 + if freshToken != nil { + freshGeneration = tokenGeneration(freshToken.AccessToken) + } + t.Logf("CLOSED_PROBE fault=%t remote_successes=%d commit_failed=%t calls=%d generations=%v replays=%d cached_generation=%d fresh_failed=%t fresh_generation=%d", + fault, probe.successes, commitFailed, len(probe.calls), probe.calls, probe.replays, cachedGeneration, freshErr != nil, freshGeneration) + if !fault && (freshErr != nil || freshGeneration != 1 || cachedGeneration != 1) { + t.Fatal("healthy cache publication contract failed") + } + if len(probe.calls) != 1 || probe.replays != 0 { + t.Fatalf("fresh instance must not replay spent generation: calls=%d replays=%d", len(probe.calls), probe.replays) + } + }) + } +} diff --git a/libs/auth/u2m/file_refresh_intent_test.go b/libs/auth/u2m/file_refresh_intent_test.go new file mode 100644 index 00000000000..9b27366b3b9 --- /dev/null +++ b/libs/auth/u2m/file_refresh_intent_test.go @@ -0,0 +1,101 @@ +package u2m_test + +import ( + "crypto/sha256" + "errors" + "net/http" + "path/filepath" + "testing" + "time" + + "github.com/databricks/cli/libs/auth/storage" + "github.com/databricks/cli/libs/auth/u2m" + "golang.org/x/oauth2" +) + +func TestFileRefreshIntentPreservesFallbackAndGuardsForce(t *testing.T) { + path := filepath.Join(t.TempDir(), "tokens.json") + store, err := storage.NewFileStore(t.Context(), storage.WithFileLocation(path)) + if err != nil { + t.Fatal("store setup failed") + } + arg, err := u2m.NewBasicAccountOAuthArgument("https://accounts.databricks.test", "account") + if err != nil { + t.Fatal("argument setup failed") + } + old := storage.Entry{Token: &oauth2.Token{AccessToken: "generation-0", RefreshToken: "generation-0", Expiry: time.Now().Add(time.Minute)}} + if err := store.Put(arg.GetCacheKey(), old); err != nil { + t.Fatal("seed failed") + } + _, err = storage.WithFileTransaction(t.Context(), store, func(view storage.Store) error { + if _, err := storage.BeginFileRefresh(t.Context(), view, arg.GetCacheKey(), sha256.Sum256([]byte("route"))); err != nil { + return err + } + return errors.New("closed unknown outcome") + }, nil) + if err == nil { + t.Fatal("prepared setup failed") + } + probe := &fileCommitFailureTransport{path: path} + auth, err := u2m.NewPersistentAuth(t.Context(), u2m.WithTokenStore(store), u2m.WithOAuthArgument(arg), u2m.WithOAuthEndpointSupplier(&u2m.BasicOAuthEndpointSupplier{}), u2m.WithHttpClient(&http.Client{Transport: probe})) + if err != nil { + t.Fatal("auth setup failed") + } + t.Cleanup(func() { _ = auth.Close() }) + token, err := auth.Token() + if err != nil || token == nil || tokenGeneration(token.AccessToken) != 0 || token.RefreshToken != "" { + t.Fatal("still-valid fallback changed") + } + if _, err := auth.ForceRefreshToken(); !errors.Is(err, storage.ErrRefreshRecoveryRequired) { + t.Fatal("force bypassed unknown generation") + } + old.Token.Expiry = time.Now().Add(-time.Hour) + if err := store.Put(arg.GetCacheKey(), old); err != nil { + t.Fatal("expired setup failed") + } + if _, err := auth.Token(); !errors.Is(err, storage.ErrRefreshRecoveryRequired) { + t.Fatal("expired token did not fail closed") + } + if len(probe.calls) != 0 { + t.Fatal("blocked credential reached provider") + } +} + +func TestFileRefreshIntentRepeatedDenialKeepsMappingAndNoReplay(t *testing.T) { + path := filepath.Join(t.TempDir(), "tokens.json") + store, err := storage.NewFileStore(t.Context(), storage.WithFileLocation(path)) + if err != nil { + t.Fatal("store setup failed") + } + arg, err := u2m.NewBasicAccountOAuthArgument("https://accounts.databricks.test", "account") + if err != nil { + t.Fatal("argument setup failed") + } + if err := store.Put(arg.GetCacheKey(), storage.Entry{Token: &oauth2.Token{AccessToken: "generation-0", RefreshToken: "generation-0", Expiry: time.Now().Add(-time.Hour)}}); err != nil { + t.Fatal("seed failed") + } + probe := &fileCommitFailureTransport{path: path, spent: true} + fresh := func() *u2m.PersistentAuth { + current, err := storage.NewFileStore(t.Context(), storage.WithFileLocation(path)) + if err != nil { + t.Fatal("fresh store failed") + } + auth, err := u2m.NewPersistentAuth(t.Context(), u2m.WithTokenStore(current), u2m.WithOAuthArgument(arg), u2m.WithOAuthEndpointSupplier(&u2m.BasicOAuthEndpointSupplier{}), u2m.WithHttpClient(&http.Client{Transport: probe})) + if err != nil { + t.Fatal("fresh auth failed") + } + t.Cleanup(func() { _ = auth.Close() }) + return auth + } + if _, err := fresh().ForceRefreshToken(); err == nil { + t.Fatal("initiating provider denial was lost") + } else if _, ok := errors.AsType[*u2m.InvalidRefreshTokenError](err); !ok { + t.Fatal("initiating provider denial type changed") + } + if _, err := fresh().ForceRefreshToken(); !errors.Is(err, storage.ErrRefreshRecoveryRequired) { + t.Fatal("later client did not retain guarded state") + } + if len(probe.calls) != 1 || probe.calls[0] != 0 { + t.Fatal("repeated denial replayed cached credential") + } +} diff --git a/libs/auth/u2m/file_refresh_process_test.go b/libs/auth/u2m/file_refresh_process_test.go new file mode 100644 index 00000000000..6531f47b3f4 --- /dev/null +++ b/libs/auth/u2m/file_refresh_process_test.go @@ -0,0 +1,308 @@ +package u2m_test + +import ( + "bytes" + "context" + "encoding/json" + "errors" + "fmt" + "io" + "net/http" + "net/url" + "os" + "os/exec" + "path/filepath" + "slices" + "strconv" + "strings" + "sync/atomic" + "testing" + "time" + + "github.com/databricks/cli/libs/auth/storage" + "github.com/databricks/cli/libs/auth/u2m" + "github.com/databricks/databricks-sdk-go/httpclient/fixtures" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "golang.org/x/oauth2" +) + +const fileRefreshActor = "DATABRICKS_TEST_FILE_REFRESH_ACTOR" + +type actorMessage struct { + Kind string + Generation int + Failed bool + Redacted bool +} + +type actorSetup struct { + Path string + Force bool +} + +type gatedFileStore struct { + storage.Store + input *json.Decoder + output *json.Encoder + first atomic.Bool +} + +func (s *gatedFileStore) Lookup(key string) (storage.Entry, error) { + entry, err := s.Store.Lookup(key) + if err != nil || s.first.Swap(true) { + return entry, err + } + if err := s.output.Encode(actorMessage{Kind: "initial-read", Generation: tokenGeneration(entry.Token.AccessToken)}); err != nil { + return storage.Entry{}, err + } + var release actorMessage + if err := s.input.Decode(&release); err != nil { + return storage.Entry{}, err + } + if release.Kind != "release" { + return storage.Entry{}, errors.New("unexpected initial-read control") + } + return entry, nil +} + +// The structural interface keeps this regression usable on the older owner too. +func (s *gatedFileStore) WithFileTransaction(ctx context.Context, update, cleanup func(storage.Store) error) (bool, error) { + owner, ok := s.Store.(interface { + WithFileTransaction(context.Context, func(storage.Store) error, func(storage.Store) error) (bool, error) + }) + if !ok { + return false, nil + } + return owner.WithFileTransaction(ctx, update, cleanup) +} + +type fileRefreshTransport struct { + input *json.Decoder + output *json.Encoder +} + +func (r fileRefreshTransport) RoundTrip(request *http.Request) (*http.Response, error) { + body, err := io.ReadAll(request.Body) + if err != nil { + return nil, err + } + request.Body = io.NopCloser(bytes.NewReader(body)) + form, err := url.ParseQuery(string(body)) + if err != nil || form.Get("grant_type") != "refresh_token" { + return nil, errors.New("unexpected synthetic exchange") + } + if err := r.output.Encode(actorMessage{Kind: "exchange", Generation: tokenGeneration(form.Get("refresh_token"))}); err != nil { + return nil, err + } + var reply actorMessage + if err := r.input.Decode(&reply); err != nil { + return nil, err + } + if reply.Kind != "exchange-result" { + return nil, errors.New("unexpected exchange control") + } + status := http.StatusOK + response := fmt.Sprintf(`{"access_token":"generation-%d","refresh_token":"generation-%d","token_type":"Bearer","expires_in":3600}`, reply.Generation, reply.Generation) + if reply.Failed { + status = http.StatusBadRequest + response = `{"error":"invalid_grant","error_description":"Refresh token is invalid"}` + } + transport := fixtures.SliceTransport{{ + Method: "POST", Resource: "/oidc/accounts/account/v1/token", + Status: status, Response: response, + ResponseHeaders: map[string][]string{"Content-Type": {"application/json"}}, + }} + return transport.RoundTrip(request) +} + +func tokenGeneration(token string) int { + generation, err := strconv.Atoi(strings.TrimPrefix(token, "generation-")) + if err != nil { + return -1 + } + return generation +} + +func runFileRefreshActor(t *testing.T) { + t.Helper() + input := json.NewDecoder(os.Stdin) + output := json.NewEncoder(os.Stdout) + var setup actorSetup + require.NoError(t, input.Decode(&setup), "actor setup") + store, err := storage.NewFileStore(t.Context(), storage.WithFileLocation(setup.Path)) + require.NoError(t, err, "actor file store") + arg, err := u2m.NewBasicAccountOAuthArgument("https://accounts.databricks.test", "account") + require.NoError(t, err) + provider, err := u2m.NewPersistentAuth(t.Context(), + u2m.WithTokenStore(&gatedFileStore{Store: store, input: input, output: output}), + u2m.WithOAuthArgument(arg), u2m.WithClientID("synthetic-client"), + u2m.WithOAuthEndpointSupplier(&u2m.BasicOAuthEndpointSupplier{}), + u2m.WithHttpClient(&http.Client{Transport: fileRefreshTransport{input: input, output: output}}), + ) + require.NoError(t, err) + t.Cleanup(func() { _ = provider.Close() }) + var token *oauth2.Token + if setup.Force { + token, err = provider.ForceRefreshToken() + } else { + token, err = provider.Token() + } + result := actorMessage{Kind: "result", Failed: err != nil} + if token != nil { + result.Generation = tokenGeneration(token.AccessToken) + result.Redacted = token.RefreshToken == "" + } + require.NoError(t, output.Encode(result)) +} + +type refreshChild struct { + command *exec.Cmd + input io.WriteCloser + output *json.Decoder + waited bool +} + +func startFileRefreshActor(t *testing.T, home string, setup actorSetup) *refreshChild { + t.Helper() + require.NoError(t, os.MkdirAll(home, 0o700)) + command := exec.CommandContext(t.Context(), os.Args[0], "-test.run=^TestFileStoreRefreshAcrossProcesses$", "-test.timeout=30s") + command.Env = []string{ + fileRefreshActor + "=1", "HOME=" + home, "USERPROFILE=" + home, + "PATH=" + os.Getenv("PATH"), "TMPDIR=" + os.TempDir(), "TMP=" + os.TempDir(), "TEMP=" + os.TempDir(), + } + input, err := command.StdinPipe() + require.NoError(t, err) + output, err := command.StdoutPipe() + require.NoError(t, err) + command.Stderr = &bytes.Buffer{} + require.NoError(t, command.Start()) + child := &refreshChild{command: command, input: input, output: json.NewDecoder(output)} + t.Cleanup(func() { + _ = input.Close() + if !child.waited { + _ = command.Process.Kill() + _ = command.Wait() + } + }) + require.NoError(t, json.NewEncoder(input).Encode(setup)) + return child +} + +func TestFileStoreRefreshAcrossProcesses(t *testing.T) { + if os.Getenv(fileRefreshActor) == "1" { + runFileRefreshActor(t) + return + } + for _, test := range []struct { + name string + sameHome bool + force bool + expired bool + exchanges int + }{ + {name: "shared-path-different-homes", exchanges: 1}, + {name: "near-expiry-same-home", sameHome: true, exchanges: 1}, + {name: "expired-shared-path", expired: true, exchanges: 1}, + {name: "explicit-force-same-home", sameHome: true, force: true, exchanges: 2}, + } { + t.Run(test.name, func(t *testing.T) { + root := t.TempDir() + path := filepath.Join(root, "shared", "tokens.json") + store, err := storage.NewFileStore(t.Context(), storage.WithFileLocation(path)) + require.NoError(t, err) + arg, err := u2m.NewBasicAccountOAuthArgument("https://accounts.databricks.test", "account") + require.NoError(t, err) + expiry := time.Now().Add(time.Minute) + if test.expired { + expiry = time.Now().Add(-time.Minute) + } else if test.force { + expiry = time.Now().Add(time.Hour) + } + require.NoError(t, store.Put(arg.GetCacheKey(), storage.Entry{Token: &oauth2.Token{ + AccessToken: "generation-0", RefreshToken: "generation-0", Expiry: expiry, + }})) + require.NoError(t, store.Put("unrelated", storage.Entry{Token: &oauth2.Token{AccessToken: "unrelated"}})) + children := make([]*refreshChild, 2) + for actor := range children { + home := filepath.Join(root, fmt.Sprintf("home-%d", actor)) + if test.sameHome { + home = filepath.Join(root, "home-shared") + } + children[actor] = startFileRefreshActor(t, home, actorSetup{Path: path, Force: test.force}) + } + for _, child := range children { + var ready actorMessage + require.NoError(t, child.output.Decode(&ready)) + require.Equal(t, "initial-read", ready.Kind) + require.Zero(t, ready.Generation) + } + type event struct { + actor int + value actorMessage + } + events := make(chan event, 4) + for actor, child := range children { + require.NoError(t, json.NewEncoder(child.input).Encode(actorMessage{Kind: "release"})) + go func() { + for { + var value actorMessage + if err := child.output.Decode(&value); err != nil { + events <- event{actor: actor, value: actorMessage{Kind: "stream-error"}} + return + } + events <- event{actor: actor, value: value} + if value.Kind == "result" { + return + } + } + }() + } + generation, exchanges, replays := 0, 0, 0 + results := make([]actorMessage, 0, 2) + for len(results) < 2 { + event := <-events + switch event.value.Kind { + case "exchange": + exchanges++ + replay := event.value.Generation != generation + if replay { + replays++ + } else { + generation++ + } + require.NoError(t, json.NewEncoder(children[event.actor].input).Encode(actorMessage{Kind: "exchange-result", Generation: generation, Failed: replay})) + case "result": + results = append(results, event.value) + default: + t.Fatal("unexpected actor stage") + } + } + for _, child := range children { + require.NoError(t, child.command.Wait(), "actor completion") + child.waited = true + } + assert.Equal(t, test.exchanges, exchanges) + assert.Zero(t, replays) + returned := make([]int, 0, 2) + for _, result := range results { + assert.False(t, result.Failed) + assert.True(t, result.Redacted) + returned = append(returned, result.Generation) + } + slices.Sort(returned) + if test.force { + assert.Equal(t, []int{1, 2}, returned) + } else { + assert.Equal(t, []int{1, 1}, returned) + } + persisted, err := store.Lookup(arg.GetCacheKey()) + require.NoError(t, err) + assert.Equal(t, test.exchanges, tokenGeneration(persisted.Token.AccessToken)) + assert.Equal(t, test.exchanges, tokenGeneration(persisted.Token.RefreshToken)) + other, err := store.Lookup("unrelated") + require.NoError(t, err) + assert.Equal(t, "unrelated", other.Token.AccessToken) + }) + } +} diff --git a/libs/auth/u2m/file_store_lease_process_test.go b/libs/auth/u2m/file_store_lease_process_test.go new file mode 100644 index 00000000000..b0dd1de8f6e --- /dev/null +++ b/libs/auth/u2m/file_store_lease_process_test.go @@ -0,0 +1,86 @@ +package u2m_test + +import ( + "bufio" + "context" + "encoding/json" + "errors" + "io" + "os" + "os/exec" + "path/filepath" + "testing" + "time" + + "github.com/databricks/cli/libs/auth/storage" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +const fileStoreLeaseActor = "DATABRICKS_TEST_SHARED_FILE_LEASE_ACTOR" + +// Delegate the owner's actual API on both the original and corrected source. +func withActualFileLease(ctx context.Context, store storage.Store, callback func() error) error { + if owner, ok := store.(interface { + WithFileTransaction(context.Context, func(storage.Store) error, func(storage.Store) error) (bool, error) + }); ok { + handled, err := owner.WithFileTransaction(ctx, func(storage.Store) error { return callback() }, nil) + if !handled { + return errors.New("file transaction was not handled") + } + return err + } + if owner, ok := store.(interface { + Lock(context.Context) (func(), error) + }); ok { + release, err := owner.Lock(ctx) + if err != nil { + return err + } + defer release() + return callback() + } + return errors.New("no file lease API") +} + +func TestFileStoreSharedPathLeaseAcrossHomes(t *testing.T) { + if os.Getenv(fileStoreLeaseActor) == "1" { + var path string + require.NoError(t, json.NewDecoder(os.Stdin).Decode(&path)) + store, err := storage.NewFileStore(t.Context(), storage.WithFileLocation(path)) + require.NoError(t, err) + require.NoError(t, withActualFileLease(t.Context(), store, func() error { + if _, err := io.WriteString(os.Stdout, "held\n"); err != nil { + return err + } + _, err := io.Copy(io.Discard, os.Stdin) + return err + })) + return + } + path := filepath.Join(t.TempDir(), "shared", "tokens.json") + store, err := storage.NewFileStore(t.Context(), storage.WithFileLocation(path)) + require.NoError(t, err) + home := t.TempDir() + command := exec.CommandContext(t.Context(), os.Args[0], "-test.run=^TestFileStoreSharedPathLeaseAcrossHomes$", "-test.timeout=30s") + command.Env = []string{fileStoreLeaseActor + "=1", "HOME=" + home, "USERPROFILE=" + home, "PATH=" + os.Getenv("PATH"), "TMPDIR=" + os.TempDir(), "TMP=" + os.TempDir(), "TEMP=" + os.TempDir()} + input, err := command.StdinPipe() + require.NoError(t, err) + output, err := command.StdoutPipe() + require.NoError(t, err) + command.Stderr = io.Discard + require.NoError(t, command.Start()) + t.Cleanup(func() { _ = input.Close(); _ = command.Process.Kill(); _ = command.Wait() }) + require.NoError(t, json.NewEncoder(input).Encode(path)) + ready, err := bufio.NewReader(output).ReadString('\n') + require.NoError(t, err) + require.Equal(t, "held\n", ready) + ctx, cancel := context.WithTimeout(t.Context(), 100*time.Millisecond) + defer cancel() + acquired := false + err = withActualFileLease(ctx, store, func() error { acquired = true; return nil }) + assert.ErrorIs(t, err, context.DeadlineExceeded) + assert.False(t, acquired) + require.NoError(t, input.Close()) + require.NoError(t, command.Wait()) +} diff --git a/libs/auth/u2m/persistent_auth.go b/libs/auth/u2m/persistent_auth.go index c7939da3da6..a8dd27a35a5 100644 --- a/libs/auth/u2m/persistent_auth.go +++ b/libs/auth/u2m/persistent_auth.go @@ -300,7 +300,11 @@ func NewPersistentAuth(ctx context.Context, opts ...PersistentAuthOption) (*Pers // using GetCacheKey(). The returned token may be expired; callers are // responsible for deciding whether and how to refresh it. func (a *PersistentAuth) loadToken() (*oauth2.Token, error) { - e, err := a.store.Lookup(a.oAuthArgument.GetCacheKey()) + return a.loadTokenFrom(a.store) +} + +func (a *PersistentAuth) loadTokenFrom(store storage.Store) (*oauth2.Token, error) { + e, err := store.Lookup(a.oAuthArgument.GetCacheKey()) if err != nil { return nil, fmt.Errorf("cache: %w", err) } @@ -321,7 +325,7 @@ func (a *PersistentAuth) Token() (*oauth2.Token, error) { return nil, err } if needsRefresh(t) { - if refreshedToken, err := a.refresh(t); err == nil { + if refreshedToken, err := a.refresh(t, false); err == nil { t = refreshedToken } else if !t.Valid() { return nil, fmt.Errorf("token refresh: %w", err) @@ -343,7 +347,7 @@ func (a *PersistentAuth) ForceRefreshToken() (*oauth2.Token, error) { if err != nil { return nil, err } - t, err = a.refresh(t) + t, err = a.refresh(t, true) if err != nil { return nil, fmt.Errorf("forced token refresh: %w", err) } @@ -400,32 +404,76 @@ func (a *PersistentAuth) recoverStoreUpdate(old, candidate *oauth2.Token) *oauth return nil } -// refresh refreshes the token for the given OAuthArgument, storing the new -// token in the cache. -// -// This read-refresh-write sequence is not coordinated across processes. -// Because the CLI is stateless, two separate CLI invocations can load the same -// cached refresh token, both attempt a refresh, and race to update the cache. -// This should be fixed in a follow-up by adding cross-process coordination -// around refresh and cache writes. -func (a *PersistentAuth) refresh(oldToken *oauth2.Token) (*oauth2.Token, error) { - // Fail fast with ErrMissingRefreshToken instead of letting the oauth2 - // library attempt to refresh and return a misleading error (e.g. "token - // expired" when the real problem is that the cached token is not - // refresh-capable). - if oldToken.RefreshToken == "" { +// refresh reloads file-backed credentials under their cache transaction. +func (a *PersistentAuth) refresh(oldToken *oauth2.Token, force bool) (*oauth2.Token, error) { + var token *oauth2.Token + handled, err := storage.WithFileTransaction(a.ctx, a.store, func(view storage.Store) error { + current, err := a.loadTokenFrom(view) + if err != nil { + return err + } + if !force && !needsRefresh(current) { + token = current + return nil + } + if current.RefreshToken == "" { + return ErrMissingRefreshToken + } + cfg, err := a.oauth2Config(a.ctx) + if err != nil { + return err + } + attempt, err := storage.BeginFileRefresh(a.ctx, view, a.oAuthArgument.GetCacheKey(), storage.RefreshRequestDigest(cfg.Endpoint.TokenURL, cfg.ClientID)) + if err != nil { + return err + } + token, err = a.exchangeConfigured(current, cfg) + if err != nil { + return err + } + entry := storage.Entry{Token: token} + if err := attempt.RecordResponse(entry); err != nil { + return err + } + return view.Put(a.oAuthArgument.GetCacheKey(), entry) + }, nil) + if handled { + return token, err + } + // Other stores keep their existing refresh and write-recovery behavior. + return a.exchangeAndStore(oldToken, a.store, true) +} + +func (a *PersistentAuth) exchangeAndStore(current *oauth2.Token, store storage.Store, recoverWrite bool) (*oauth2.Token, error) { + if current.RefreshToken == "" { return nil, ErrMissingRefreshToken } cfg, err := a.oauth2Config(a.ctx) if err != nil { return nil, err } + t, err := a.exchangeConfigured(current, cfg) + if err != nil { + return nil, err + } + if err := store.Put(a.oAuthArgument.GetCacheKey(), storage.Entry{Token: t}); err != nil { + if recoverWrite { + if cached := a.recoverStoreUpdate(current, t); cached != nil { + return cached, nil + } + } + return nil, fmt.Errorf("cache update: %w", err) + } + return t, nil +} + +func (a *PersistentAuth) exchangeConfigured(current *oauth2.Token, cfg *oauth2.Config) (*oauth2.Token, error) { ctx := a.setOAuthContext(a.ctx) // Force the oauth2 library to refresh by ensuring the token appears // expired. PersistentAuth owns the refresh decision (including the // proactive buffer), so the oauth2 library should always perform the // refresh when asked. - expired := *oldToken + expired := *current expired.Expiry = time.Now().Add(-time.Minute) t, err := cfg.TokenSource(ctx, &expired).Token() if err != nil { @@ -436,13 +484,7 @@ func (a *PersistentAuth) refresh(oldToken *oauth2.Token) (*oauth2.Token, error) } return nil, oauthError(err) } - err = a.store.Put(a.oAuthArgument.GetCacheKey(), storage.Entry{Token: t}) - if err != nil { - if cached := a.recoverStoreUpdate(oldToken, t); cached != nil { - return cached, nil - } - return nil, fmt.Errorf("cache update: %w", err) - } + return t, nil } diff --git a/libs/databrickscfg/loader_test.go b/libs/databrickscfg/loader_test.go index 9fdf993d394..f13a83f4d92 100644 --- a/libs/databrickscfg/loader_test.go +++ b/libs/databrickscfg/loader_test.go @@ -1,9 +1,11 @@ package databrickscfg import ( + "context" "errors" "os" "path/filepath" + "slices" "testing" "github.com/databricks/databricks-sdk-go/config" @@ -11,6 +13,21 @@ import ( "github.com/stretchr/testify/require" ) +func stubUnavailableLoaderMetadata(t *testing.T, hosts ...string) { + t.Helper() + previous := config.DefaultHostMetadataResolverFactory + config.DefaultHostMetadataResolverFactory = func(*config.Config) config.HostMetadataResolver { + return func(_ context.Context, host string) (*config.HostMetadata, error) { + if !slices.Contains(hosts, host) { + t.Error("unexpected fixture metadata host") + return nil, errors.New("unexpected fixture metadata host") + } + return nil, nil + } + } + t.Cleanup(func() { config.DefaultHostMetadataResolverFactory = previous }) +} + func TestLoaderSkipsEmptyHost(t *testing.T) { cfg := config.Config{ Loaders: []config.Loader{ @@ -24,6 +41,7 @@ func TestLoaderSkipsEmptyHost(t *testing.T) { } func TestLoaderSkipsExistingAuth(t *testing.T) { + stubUnavailableLoaderMetadata(t, "https://foo.test") cfg := config.Config{ Loaders: []config.Loader{ ResolveProfileFromHost, @@ -68,6 +86,7 @@ func TestResolveNonAuthFromEnvSkipsHostAndAuth(t *testing.T) { } func TestProfileAuthLoadersProfileWinsOverSteeringEnv(t *testing.T) { + stubUnavailableLoaderMetadata(t, "https://default.test") // A PAT profile plus auth-steering/routing env vars. The profile owns // host/routing/auth-steering, so these env vars must not shadow it (#5096). // In particular DATABRICKS_AUTH_TYPE=basic must not force basic auth over the @@ -98,6 +117,7 @@ func TestProfileAuthLoadersProfileWinsOverSteeringEnv(t *testing.T) { } func TestProfileAuthLoadersGapFillsHostFromEnv(t *testing.T) { + stubUnavailableLoaderMetadata(t, "https://env.test") // A selected profile that leaves the host empty falls back to the env host, // mirroring the token gap-fill (#5096). Auth-steering env vars stay ignored. dir := t.TempDir() @@ -167,6 +187,7 @@ func TestEnvAlwaysSkipAttrsCoverSDKInternalEnvAttrs(t *testing.T) { } func TestLoaderSkipsExplicitAuthType(t *testing.T) { + stubUnavailableLoaderMetadata(t, "https://default.test") cfg := config.Config{ Loaders: []config.Loader{ ResolveProfileFromHost, @@ -184,6 +205,7 @@ func TestLoaderSkipsExplicitAuthType(t *testing.T) { } func TestLoaderSkipsNonExistingConfigFile(t *testing.T) { + stubUnavailableLoaderMetadata(t, "https://default.test") cfg := config.Config{ Loaders: []config.Loader{ ResolveProfileFromHost, @@ -211,6 +233,7 @@ func TestLoaderErrorsOnInvalidFile(t *testing.T) { } func TestLoaderSkipsNoMatchingHost(t *testing.T) { + stubUnavailableLoaderMetadata(t, "https://noneofthehostsmatch.test") cfg := config.Config{ Loaders: []config.Loader{ ResolveProfileFromHost, @@ -225,6 +248,7 @@ func TestLoaderSkipsNoMatchingHost(t *testing.T) { } func TestLoaderMatchingHost(t *testing.T) { + stubUnavailableLoaderMetadata(t, "https://default.test") cfg := config.Config{ Loaders: []config.Loader{ ResolveProfileFromHost, @@ -240,6 +264,7 @@ func TestLoaderMatchingHost(t *testing.T) { } func TestLoaderMatchingHostWithQuery(t *testing.T) { + stubUnavailableLoaderMetadata(t, "https://query.test") cfg := config.Config{ Loaders: []config.Loader{ ResolveProfileFromHost, @@ -298,6 +323,7 @@ func TestAsMultipleProfilesReturnsFalseForNil(t *testing.T) { } func TestLoaderDisambiguatesByWorkspaceID(t *testing.T) { + stubUnavailableLoaderMetadata(t, "https://spog.databricks.com") cfg := config.Config{ Loaders: []config.Loader{ ResolveProfileFromHost, @@ -314,6 +340,7 @@ func TestLoaderDisambiguatesByWorkspaceID(t *testing.T) { } func TestLoaderDisambiguatesByWorkspaceIDSecondProfile(t *testing.T) { + stubUnavailableLoaderMetadata(t, "https://spog.databricks.com") cfg := config.Config{ Loaders: []config.Loader{ ResolveProfileFromHost, @@ -404,6 +431,7 @@ func writeTestConfigFile(t *testing.T, contents string) string { } func TestStrictProfileLoader(t *testing.T) { + stubUnavailableLoaderMetadata(t, "https://classic.test", "https://custom.test") profiles := writeTestConfigFile(t, strictProfiles) ambiguous := writeTestConfigFile(t, "[DEFAULT]\nhost = https://classic.test\n\n[workspace_classic]\nhost = https://classic.test\n") missing := filepath.Join(t.TempDir(), "missing") @@ -466,6 +494,7 @@ func TestStrictProfileLoader(t *testing.T) { } func TestStrictProfileLoadersIgnoreEnvHost(t *testing.T) { + stubUnavailableLoaderMetadata(t, "https://custom.test") // The env token reaches only the host of the selected profile, even when the // env also sets a host. t.Setenv("DATABRICKS_HOST", "https://env.test")