From 60530bcb07bc0140dfaa2f7547188f44111de20c Mon Sep 17 00:00:00 2001 From: Harry Yao Date: Sat, 7 Nov 2026 10:30:37 +0000 Subject: [PATCH 1/3] Serialize file-backed OAuth cache transactions across writers Reload and recheck the canonical file store under one operation-context lease, then atomically publish the staged map. Coordinate login, logout, pending login and legacy mirror writers while preserving main's secure backend and Store API. Build on the OS locking foundation from Vivek1106-04 in related PR #6759; lock_unix.go and lock_windows.go are copied byte-identically. This is a file-backed alternative proposed against main from an owned branch. Co-authored-by: Vivek1106-04 Signed-off-by: Harry Yao --- .../file-backed-oauth-cache-transactions.md | 1 + acceptance/bundle/deploy/mlops-stacks/script | 2 +- .../non-interactive/nominal/output.txt | 1 + acceptance/cmd/completion/script | 10 +- bundle/bundle_test.go | 5 + bundle/config/workspace_test.go | 24 ++ cmd/auth/auth_test.go | 16 +- cmd/auth/login.go | 11 + cmd/auth/login_device_code.go | 4 +- cmd/auth/login_device_code_test.go | 3 + cmd/auth/login_test.go | 88 ++++ cmd/auth/logout.go | 84 ++-- cmd/auth/logout_test.go | 97 +++++ cmd/auth/token_test.go | 12 + libs/auth/arguments_test.go | 20 + libs/auth/credentials_test.go | 5 + libs/auth/error_test.go | 1 + libs/auth/storage/dual_writing_store.go | 27 ++ libs/auth/storage/file_transaction.go | 95 +++++ libs/auth/storage/file_transaction_test.go | 387 ++++++++++++++++++ libs/auth/storage/filestore.go | 85 ++-- libs/auth/storage/lock.go | 64 +++ libs/auth/storage/lock_unix.go | 29 ++ libs/auth/storage/lock_windows.go | 40 ++ libs/auth/storage/not_found_hint.go | 15 + libs/auth/u2m/file_refresh_process_test.go | 308 ++++++++++++++ .../auth/u2m/file_store_lease_process_test.go | 86 ++++ libs/auth/u2m/persistent_auth.go | 61 ++- libs/databrickscfg/loader_test.go | 29 ++ 29 files changed, 1521 insertions(+), 89 deletions(-) create mode 100644 .nextchanges/cli/file-backed-oauth-cache-transactions.md create mode 100644 libs/auth/storage/file_transaction.go create mode 100644 libs/auth/storage/file_transaction_test.go create mode 100644 libs/auth/storage/lock.go create mode 100644 libs/auth/storage/lock_unix.go create mode 100644 libs/auth/storage/lock_windows.go create mode 100644 libs/auth/u2m/file_refresh_process_test.go create mode 100644 libs/auth/u2m/file_store_lease_process_test.go 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..b791e177984 --- /dev/null +++ b/.nextchanges/cli/file-backed-oauth-cache-transactions.md @@ -0,0 +1 @@ +* Fix concurrent file-backed OAuth refreshes and cache updates overwriting newer credentials or unrelated entries. 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/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/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..5db8de1dd0e 100644 --- a/cmd/auth/login.go +++ b/cmd/auth/login.go @@ -126,6 +126,17 @@ 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 { return view.Put(key, storage.Entry{Token: token}) }, + 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_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_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_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..0e45ab7bca5 --- /dev/null +++ b/libs/auth/storage/file_transaction.go @@ -0,0 +1,95 @@ +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 +} + +func (t *fileTransaction) Lookup(key string) (Entry, error) { + 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 { + 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 _, 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..05db1d785e8 100644 --- a/libs/auth/storage/filestore.go +++ b/libs/auth/storage/filestore.go @@ -99,17 +99,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 +111,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,17 +119,40 @@ 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 fmt.Errorf("load: %w", err) + 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) + } + staged := &fileTransaction{file: cloneStoreFile(file)} + if err := update(staged); err != nil { + return true, err } - if _, ok := f.Tokens[key]; !ok { - return nil + if !staged.dirty { + return true, nil } - delete(f.Tokens, key) - return c.write(f) + if err := c.write(staged.file); err != nil { + err = fmt.Errorf("cache update: %w", err) + if cleanup == nil { + return true, err + } + rollback := &fileTransaction{file: cloneStoreFile(file)} + cleanupErr := cleanup(rollback) + if cleanupErr == nil && rollback.dirty { + cleanupErr = c.write(rollback.file) + } + return true, errors.Join(err, cleanupErr) + } + return true, nil } // write marshals f and atomically replaces the store file. @@ -165,14 +178,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/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..c321b769efc 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,20 +404,33 @@ 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 + } + token, err = a.exchangeAndStore(current, view, false) + return err + }, 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) @@ -425,7 +442,7 @@ func (a *PersistentAuth) refresh(oldToken *oauth2.Token) (*oauth2.Token, error) // 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,10 +453,12 @@ 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}) + err = store.Put(a.oAuthArgument.GetCacheKey(), storage.Entry{Token: t}) if err != nil { - if cached := a.recoverStoreUpdate(oldToken, t); cached != nil { - return cached, nil + if recoverWrite { + if cached := a.recoverStoreUpdate(current, t); cached != nil { + return cached, nil + } } return nil, fmt.Errorf("cache update: %w", err) } 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") From d4cda6b9a57ff3858fddb6d13326da94e40b4e7a Mon Sep 17 00:00:00 2001 From: Harry Yao Date: Sat, 7 Nov 2026 10:58:50 +0000 Subject: [PATCH 2/3] Link file-backed OAuth cache changelog to introducing PR Signed-off-by: Harry Yao --- .nextchanges/cli/file-backed-oauth-cache-transactions.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.nextchanges/cli/file-backed-oauth-cache-transactions.md b/.nextchanges/cli/file-backed-oauth-cache-transactions.md index b791e177984..dff1f571539 100644 --- a/.nextchanges/cli/file-backed-oauth-cache-transactions.md +++ b/.nextchanges/cli/file-backed-oauth-cache-transactions.md @@ -1 +1 @@ -* Fix concurrent file-backed OAuth refreshes and cache updates overwriting newer credentials or unrelated entries. +* Fix concurrent file-backed OAuth refreshes and cache updates overwriting newer credentials or unrelated entries. ([#7396](https://github.com/databricks/cli/pull/7396)) From 7ab1fc01fb128c9181fc7937256a10f9c8806d64 Mon Sep 17 00:00:00 2001 From: Harry Yao Date: Sat, 7 Nov 2026 14:13:01 +0000 Subject: [PATCH 3/3] Guard uncertain file-backed OAuth refresh generations Persist consumed generations before exchange using a synced append-only companion under the existing canonical-cache lease. Bind successful responses to their exact cache publication and written-key epochs, preserving aliases, access-only login, valid-access fallback and explicit force safeguards. Keep consumed tombstones and fail closed on incomplete metadata. Preserve initiating provider denial handling and render guarded recovery errors through the existing actionable/JSON paths. Signed-off-by: Harry Yao --- .../file-backed-oauth-cache-transactions.md | 2 +- .../docker/invalid-refresh-token/output.txt | 8 +- .../output.txt | 2 +- cmd/auth/login.go | 8 +- cmd/auth/login_intent_test.go | 70 ++ cmd/auth/token.go | 2 +- libs/auth/error.go | 20 +- libs/auth/error_intent_test.go | 58 ++ libs/auth/storage/file_transaction.go | 25 +- libs/auth/storage/filestore.go | 50 +- libs/auth/storage/refresh_intent.go | 734 ++++++++++++++++++ libs/auth/storage/refresh_intent_file.go | 636 +++++++++++++++ .../storage/refresh_intent_process_test.go | 140 ++++ libs/auth/storage/refresh_intent_test.go | 654 ++++++++++++++++ .../auth/storage/refresh_intent_transition.go | 372 +++++++++ .../u2m/file_refresh_commit_failure_test.go | 165 ++++ libs/auth/u2m/file_refresh_intent_test.go | 101 +++ libs/auth/u2m/persistent_auth.go | 45 +- 18 files changed, 3051 insertions(+), 41 deletions(-) create mode 100644 cmd/auth/login_intent_test.go create mode 100644 libs/auth/error_intent_test.go create mode 100644 libs/auth/storage/refresh_intent.go create mode 100644 libs/auth/storage/refresh_intent_file.go create mode 100644 libs/auth/storage/refresh_intent_process_test.go create mode 100644 libs/auth/storage/refresh_intent_test.go create mode 100644 libs/auth/storage/refresh_intent_transition.go create mode 100644 libs/auth/u2m/file_refresh_commit_failure_test.go create mode 100644 libs/auth/u2m/file_refresh_intent_test.go diff --git a/.nextchanges/cli/file-backed-oauth-cache-transactions.md b/.nextchanges/cli/file-backed-oauth-cache-transactions.md index dff1f571539..ba143e26d69 100644 --- a/.nextchanges/cli/file-backed-oauth-cache-transactions.md +++ b/.nextchanges/cli/file-backed-oauth-cache-transactions.md @@ -1 +1 @@ -* Fix concurrent file-backed OAuth refreshes and cache updates overwriting newer credentials or unrelated entries. ([#7396](https://github.com/databricks/cli/pull/7396)) +* 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/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/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/cmd/auth/login.go b/cmd/auth/login.go index 5db8de1dd0e..f8f02ad325b 100644 --- a/cmd/auth/login.go +++ b/cmd/auth/login.go @@ -127,7 +127,13 @@ func storeLoginToken(ctx context.Context, tokenStore storage.Store, mode storage tokenStore = storage.WrapForOAuthArgument(ctx, tokenStore, mode, arg) key := arg.GetCacheKey() handled, err := storage.WithFileTransaction(ctx, tokenStore, - func(view storage.Store) error { return view.Put(key, storage.Entry{Token: token}) }, + 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 { 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/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/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/storage/file_transaction.go b/libs/auth/storage/file_transaction.go index 0e45ab7bca5..5f94b4dfe9b 100644 --- a/libs/auth/storage/file_transaction.go +++ b/libs/auth/storage/file_transaction.go @@ -50,11 +50,24 @@ func DeleteWithContext(ctx context.Context, store Store, key string) error { } type fileTransaction struct { - file *tokenStoreFile - dirty bool + 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 @@ -63,6 +76,10 @@ func (t *fileTransaction) Lookup(key string) (Entry, error) { } 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) @@ -73,6 +90,10 @@ func (t *fileTransaction) Put(key string, entry Entry) error { } 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 diff --git a/libs/auth/storage/filestore.go b/libs/auth/storage/filestore.go index 05db1d785e8..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 @@ -133,34 +135,44 @@ func (c *fileStore) WithFileTransaction(ctx context.Context, update, cleanup fun if err != nil { return true, fmt.Errorf("load: %w", err) } - staged := &fileTransaction{file: cloneStoreFile(file)} - if err := update(staged); err != nil { + 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 true, err } - if !staged.dirty { - return true, nil + if err := c.reconcileRefreshJournal(ctx, journal, raw); err != nil { + return true, err } - if err := c.write(staged.file); err != nil { - err = fmt.Errorf("cache update: %w", err) - if cleanup == nil { - return true, err - } - rollback := &fileTransaction{file: cloneStoreFile(file)} - cleanupErr := cleanup(rollback) - if cleanupErr == nil && rollback.dirty { - cleanupErr = c.write(rollback.file) - } - return true, errors.Join(err, cleanupErr) + staged := newFileTransaction(ctx, c, file, raw, journal) + defer func() { staged.live = false }() + if err := update(staged); err != nil { + return true, err } - return true, nil + 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) } 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/persistent_auth.go b/libs/auth/u2m/persistent_auth.go index c321b769efc..a8dd27a35a5 100644 --- a/libs/auth/u2m/persistent_auth.go +++ b/libs/auth/u2m/persistent_auth.go @@ -419,8 +419,23 @@ func (a *PersistentAuth) refresh(oldToken *oauth2.Token, force bool) (*oauth2.To if current.RefreshToken == "" { return ErrMissingRefreshToken } - token, err = a.exchangeAndStore(current, view, false) - return err + 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 @@ -437,6 +452,22 @@ func (a *PersistentAuth) exchangeAndStore(current *oauth2.Token, store storage.S 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 @@ -453,15 +484,7 @@ func (a *PersistentAuth) exchangeAndStore(current *oauth2.Token, store storage.S } return nil, oauthError(err) } - err = store.Put(a.oAuthArgument.GetCacheKey(), storage.Entry{Token: t}) - if 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 }