From 5224d9c527c04bd1c54f212aacddbfb243aa3ac9 Mon Sep 17 00:00:00 2001 From: audanye-sudo Date: Fri, 17 Jul 2026 02:11:45 +0800 Subject: [PATCH] fix(auth): harden multi-account profile compatibility Preserve manual-token defaults across explicit profile refreshes and selective logout while keeping legacy marker behavior compatible. Make profile login, refresh, switch, and logout writes rollback-safe; reject unsupported future profile versions before remote side effects; and prevent cross-profile token fallback. Forward token overrides through usage recording, use the newly authenticated identity for post-login authorization, distinguish unavailable profile state, and propagate Windows registry deletion failures. --- internal/app/access_token_resolve.go | 6 + internal/app/app_small_coverage_test.go | 6 + internal/app/auth_command.go | 22 +- .../app/auth_command_full_coverage_test.go | 38 +- internal/app/auth_command_test.go | 65 +- internal/app/profile_command.go | 8 +- internal/app/profile_command_test.go | 21 + internal/app/usage_recording.go | 16 + internal/app/usage_recording_test.go | 30 +- internal/auth/coverage_edges_test.go | 41 +- internal/auth/keychain_store.go | 23 +- internal/auth/oauth_provider.go | 2 +- internal/auth/profiles.go | 200 ++++-- internal/auth/remaining_coverage_test.go | 4 +- internal/auth/token.go | 568 ++++++++++++++--- internal/auth/token_preflight_test.go | 87 ++- internal/auth/token_test.go | 569 ++++++++++++++++++ internal/keychain/keychain_windows.go | 34 +- internal/keychain/keychain_windows_test.go | 48 ++ 19 files changed, 1656 insertions(+), 132 deletions(-) create mode 100644 internal/keychain/keychain_windows_test.go diff --git a/internal/app/access_token_resolve.go b/internal/app/access_token_resolve.go index 570c089b..e9a24e84 100644 --- a/internal/app/access_token_resolve.go +++ b/internal/app/access_token_resolve.go @@ -56,6 +56,12 @@ func resolveAccessTokenFromDir(ctx context.Context, configDir string) (string, e if tokenErr != nil && errors.Is(tokenErr, authpkg.ErrTokenDecryption) { return "", tokenErr } + if strings.TrimSpace(authpkg.RuntimeProfile()) != "" { + if tokenErr != nil { + return "", tokenErr + } + return "", nil + } manager := newLegacyTokenManager(configDir) if leg, _, err := manager.GetToken(); err == nil && strings.TrimSpace(leg) != "" { return strings.TrimSpace(leg), nil diff --git a/internal/app/app_small_coverage_test.go b/internal/app/app_small_coverage_test.go index 6ec27f9d..ee6e313e 100644 --- a/internal/app/app_small_coverage_test.go +++ b/internal/app/app_small_coverage_test.go @@ -217,6 +217,12 @@ func TestCrossPlatformCoverageConfigAndTokenSeamsCoverage(t *testing.T) { if got, err := resolveAccessTokenFromDir(context.Background(), "unused"); err != nil || got != "legacy" { t.Fatalf("legacy token = %q, %v", got, err) } + authpkg.SetRuntimeProfile("corp:user") + t.Cleanup(func() { authpkg.SetRuntimeProfile("") }) + if got, err := resolveAccessTokenFromDir(context.Background(), "unused"); got != "" || err == nil || err.Error() != "missing" { + t.Fatalf("explicit profile fallback = token %q error %v, want profile error", got, err) + } + authpkg.SetRuntimeProfile("") newAccessTokenProvider = func(string) accessTokenGetter { return fakeAccessTokenGetter{err: errors.New("load")} } newLegacyTokenManager = func(string) legacyTokenGetter { return fakeLegacyTokenGetter{err: errors.New("missing")} } edition.Override(&edition.Hooks{}) diff --git a/internal/app/auth_command.go b/internal/app/auth_command.go index ca27057c..04ae6751 100644 --- a/internal/app/auth_command.go +++ b/internal/app/auth_command.go @@ -178,6 +178,8 @@ func newAuthLoginCommand(patCaller edition.ToolCaller) *cobra.Command { if !recommendAuthMode { return nil } + restoreProfile := replaceRuntimeProfile(authpkg.TokenProfileSelector(tokenData)) + defer restoreProfile() recommendScopeMode := pat.LoginRecommendScopeRecommended var initialPlan *pat.LoginRecommendPlan if postLoginTUIMode { @@ -647,8 +649,10 @@ func logoutOneProfile(_ *cobra.Command, ctx context.Context, configDir, selector if selected == nil { return apperrors.NewValidation(fmt.Sprintf("profile %q not found", selector)) } + stableSelector := selected.CorpID if exact { - if data, loadErr := authLoadTokenForProfile(configDir, authpkg.ProfileSelector(*selected)); loadErr == nil { + stableSelector = authpkg.ProfileSelector(*selected) + if data, loadErr := authLoadTokenForProfile(configDir, stableSelector); loadErr == nil { _ = authRevokeTokenForData(ctx, data) } } else if cfg, loadErr := authLoadProfiles(configDir); loadErr == nil { @@ -661,7 +665,7 @@ func logoutOneProfile(_ *cobra.Command, ctx context.Context, configDir, selector } } } - if err := authDeleteProfileToken(configDir, selector); err != nil { + if err := authDeleteProfileToken(configDir, stableSelector); err != nil { if strings.Contains(err.Error(), "not found") { return apperrors.NewValidation(err.Error()) } @@ -705,6 +709,14 @@ func pushRuntimeProfile(selector string) func() { } } +func replaceRuntimeProfile(selector string) func() { + previous := authpkg.RuntimeProfile() + authpkg.SetRuntimeProfile(strings.TrimSpace(selector)) + return func() { + authpkg.SetRuntimeProfile(previous) + } +} + func newAuthExportCommand() *cobra.Command { return newAuthExportCommandWithSupport(authpkg.PortableExportSupportError) } @@ -1248,13 +1260,12 @@ func enrichAuthLoginProfileFromContact(ctx context.Context, _ string, caller edi return nil } - args := map[string]any{"profile": corpID} var ( result *edition.ToolResult err error ) if tokenCaller, ok := caller.(tokenOverrideToolCaller); ok && strings.TrimSpace(data.AccessToken) != "" { - result, err = tokenCaller.CallToolWithToken(ctx, data.AccessToken, "contact", "get_current_user_profile", args) + result, err = tokenCaller.CallToolWithToken(ctx, data.AccessToken, "contact", "get_current_user_profile", nil) } else { if strings.TrimSpace(data.UserID) == "" { return fmt.Errorf("login identity lookup requires an in-memory token override") @@ -1262,6 +1273,9 @@ func enrichAuthLoginProfileFromContact(ctx context.Context, _ string, caller edi return nil } if err != nil { + if strings.TrimSpace(data.UserID) != "" { + return nil + } return err } identity, ok := contactProfileIdentityFromToolResult(result) diff --git a/internal/app/auth_command_full_coverage_test.go b/internal/app/auth_command_full_coverage_test.go index 98fa1760..8bd0f67f 100644 --- a/internal/app/auth_command_full_coverage_test.go +++ b/internal/app/auth_command_full_coverage_test.go @@ -355,6 +355,14 @@ func TestCrossPlatformCoverageAuthCoverageContactEnrichment(t *testing.T) { if err := enrichAuthLoginProfileFromContact(ctx, "cfg", &authCoverageCaller{err: errors.New("call")}, &authpkg.TokenData{CorpID: "ding"}); err == nil { t.Fatal("caller error should propagate") } + if err := enrichAuthLoginProfileFromContact( + ctx, + "cfg", + &authCoverageCaller{err: errors.New("call")}, + &authpkg.TokenData{CorpID: "ding", UserID: "known", AccessToken: "token"}, + ); err != nil { + t.Fatalf("optional contact metadata failure with known userId = %v", err) + } for _, text := range []string{"", "not-json", `{"result":[]}`, `{"result":[{"orgEmployeeModel":{}}]}`} { caller := &authCoverageCaller{result: &edition.ToolResult{Content: []edition.ContentBlock{{Text: text}}}} if err := enrichAuthLoginProfileFromContact(ctx, "cfg", caller, &authpkg.TokenData{CorpID: "ding", AccessToken: "token"}); err != nil { @@ -526,14 +534,40 @@ func TestCrossPlatformCoverageAuthCoverageStatusAndLogout(t *testing.T) { return &authpkg.TokenData{CorpID: "ding", UserID: "user"}, nil } authRevokeTokenForData = func(context.Context, *authpkg.TokenData) error { return errors.New("ignored") } - authDeleteProfileToken = func(string, string) error { return errors.New("delete") } + var deletedSelector string + authDeleteProfileToken = func(_ string, selector string) error { + deletedSelector = selector + return errors.New("delete") + } if err := logoutOneProfile(nil, context.Background(), "cfg", "x"); err == nil { t.Fatal("profile delete should fail") } - authDeleteProfileToken = func(string, string) error { return nil } + if deletedSelector != "ding:user" { + t.Fatalf("exact deletion selector = %q, want stable identity selector", deletedSelector) + } + authDeleteProfileToken = func(_ string, selector string) error { + deletedSelector = selector + return nil + } if err := logoutOneProfile(nil, context.Background(), "cfg", "x"); err != nil { t.Fatal(err) } + if deletedSelector != "ding:user" { + t.Fatalf("exact deletion selector = %q, want stable identity selector", deletedSelector) + } + + authResolveProfileDeletion = func(string, string) (*authpkg.Profile, bool, error) { + return &authpkg.Profile{CorpID: "ding", UserID: "user"}, false, nil + } + authLoadProfiles = func(string) (*authpkg.ProfilesConfig, error) { + return &authpkg.ProfilesConfig{Profiles: []authpkg.Profile{{CorpID: "ding", UserID: "user"}}}, nil + } + if err := logoutOneProfile(nil, context.Background(), "cfg", "organization-name"); err != nil { + t.Fatal(err) + } + if deletedSelector != "ding" { + t.Fatalf("organization deletion selector = %q, want stable corpId", deletedSelector) + } authEnsureProfilesMigration = func(string) error { return errors.New("migrate") } if err := logoutAllProfiles(nil, context.Background(), "cfg"); err == nil { diff --git a/internal/app/auth_command_test.go b/internal/app/auth_command_test.go index a457a4c8..7e03f6c8 100644 --- a/internal/app/auth_command_test.go +++ b/internal/app/auth_command_test.go @@ -1087,6 +1087,12 @@ func TestAuthLoginRecommendSkipsPostLoginTUI(t *testing.T) { `{"success":true,"data":{"items":[{"scope":"calendar.event:read","productCode":"calendar","productName":"日历"}],"selectedScopes":["calendar.event:read"]}}`, `{"success":true,"data":{"grantedScopes":["calendar.event:read"]}}`, }} + authpkg.SetRuntimeProfile("corp_old:user_old") + t.Cleanup(func() { authpkg.SetRuntimeProfile("") }) + var authorizationProfiles []string + fake.beforeCall = func(string) { + authorizationProfiles = append(authorizationProfiles, authpkg.RuntimeProfile()) + } cmd := newAuthLoginCommand(fake) var out bytes.Buffer cmd.SetOut(&out) @@ -1105,6 +1111,61 @@ func TestAuthLoginRecommendSkipsPostLoginTUI(t *testing.T) { if got := fake.args[0]["recommend"]; got != true { t.Fatalf("--recommend plan recommend = %#v, want true", got) } + for _, profile := range authorizationProfiles { + if profile != "" { + t.Fatalf("manual token post-login profile = %q, want empty runtime selector", profile) + } + } + if got := authpkg.RuntimeProfile(); got != "corp_old:user_old" { + t.Fatalf("runtime profile after authorization = %q, want restored selector", got) + } +} + +func TestAuthLoginRecommendUsesNewExactIdentity(t *testing.T) { + t.Setenv(keychain.DisableKeychainEnv, "1") + t.Setenv(keychain.StorageDirEnv, t.TempDir()) + t.Setenv("DWS_CONFIG_DIR", t.TempDir()) + + oldOAuthLogin := authOAuthLogin + oldInteractive := authLoginInteractiveTerminal + t.Cleanup(func() { + authOAuthLogin = oldOAuthLogin + authLoginInteractiveTerminal = oldInteractive + authpkg.SetRuntimeProfile("") + }) + authLoginInteractiveTerminal = func() bool { return false } + authOAuthLogin = func(*authpkg.OAuthProvider, context.Context, bool) (*authpkg.TokenData, error) { + return &authpkg.TokenData{ + AccessToken: "new-token", + CorpID: "corp_same", + UserID: "user_new", + ExpiresAt: time.Now().Add(time.Hour), + }, nil + } + fake := &authLoginRecommendSequenceCaller{responses: []string{ + `{"success":true,"data":{"items":[],"selectedScopes":[]}}`, + }} + var authorizationProfiles []string + fake.beforeCall = func(string) { + authorizationProfiles = append(authorizationProfiles, authpkg.RuntimeProfile()) + } + authpkg.SetRuntimeProfile("corp_same:user_old") + + cmd := newAuthLoginCommand(fake) + cmd.SetOut(io.Discard) + cmd.SetErr(io.Discard) + cmd.SetArgs([]string{"--recommend"}) + if err := cmd.Execute(); err != nil { + t.Fatalf("auth login --recommend error = %v", err) + } + for _, profile := range authorizationProfiles { + if profile != "corp_same:user_new" { + t.Fatalf("post-login authorization profile = %q, want new exact identity", profile) + } + } + if got := authpkg.RuntimeProfile(); got != "corp_same:user_old" { + t.Fatalf("runtime profile after authorization = %q, want restored old identity", got) + } } func TestAuthLoginDefaultTUIModeSkipsSelectorWhenAllGranted(t *testing.T) { @@ -1377,8 +1438,8 @@ func TestEnrichAuthLoginProfileFromContactBeforePersist(t *testing.T) { if len(fake.tools) != 1 || fake.tools[0] != "get_current_user_profile" { t.Fatalf("tool calls = %v, want get_current_user_profile", fake.tools) } - if got := fake.args[0]["profile"]; got != "ding32fff839a3e0105d" { - t.Fatalf("contact profile arg = %#v, want ding32fff839a3e0105d", got) + if len(fake.args[0]) != 0 { + t.Fatalf("contact profile args = %#v, want no arguments", fake.args[0]) } if len(fake.tokens) != 1 || fake.tokens[0] != "access-token" { t.Fatalf("token overrides = %v, want access-token", fake.tokens) diff --git a/internal/app/profile_command.go b/internal/app/profile_command.go index c4bfb493..e0315e5a 100644 --- a/internal/app/profile_command.go +++ b/internal/app/profile_command.go @@ -146,6 +146,7 @@ var ( profileSwitchTUIRunner = runProfileSwitchTUI profileEnsureProfilesMigration = authpkg.EnsureProfilesMigration profileLoadProfiles = authpkg.LoadProfiles + profileLoadTokenData = authpkg.LoadTokenDataForProfile profileUsePrevious = authpkg.UsePreviousProfile profileSetCurrent = authpkg.SetCurrentProfile profileRunTeaProgram = (*tea.Program).Run @@ -743,10 +744,13 @@ func profileViewFromProfile( } func loadProfileTokenState(configDir string, profile authpkg.Profile) *profileTokenState { - data, err := authpkg.LoadTokenDataForProfile(configDir, authpkg.ProfileSelector(profile)) - if err != nil || data == nil { + data, err := profileLoadTokenData(configDir, authpkg.ProfileSelector(profile)) + if errors.Is(err, authpkg.ErrTokenDataNotFound) || (err == nil && data == nil) { return &profileTokenState{Status: authpkg.ProfileStatusRevoked} } + if err != nil { + return &profileTokenState{Status: authpkg.ProfileStatusUnavailable} + } status := authpkg.ProfileStatusExpired if data.IsAccessTokenValid() { status = authpkg.ProfileStatusActive diff --git a/internal/app/profile_command_test.go b/internal/app/profile_command_test.go index 0c496087..de0950b1 100644 --- a/internal/app/profile_command_test.go +++ b/internal/app/profile_command_test.go @@ -16,6 +16,7 @@ package app import ( "bytes" "encoding/json" + "errors" "fmt" "strings" "testing" @@ -188,6 +189,26 @@ func TestProfileListUsesRealIdentityTokenState(t *testing.T) { } } +func TestProfileListDistinguishesMissingAndUnavailableTokenState(t *testing.T) { + originalLoad := profileLoadTokenData + t.Cleanup(func() { profileLoadTokenData = originalLoad }) + profile := authpkg.Profile{CorpID: "corp", UserID: "user"} + + profileLoadTokenData = func(string, string) (*authpkg.TokenData, error) { + return nil, authpkg.ErrTokenDataNotFound + } + if state := loadProfileTokenState("cfg", profile); state.Status != authpkg.ProfileStatusRevoked { + t.Fatalf("missing token status = %q, want revoked", state.Status) + } + + profileLoadTokenData = func(string, string) (*authpkg.TokenData, error) { + return nil, errors.New("keychain unavailable") + } + if state := loadProfileTokenState("cfg", profile); state.Status != authpkg.ProfileStatusUnavailable { + t.Fatalf("unavailable token status = %q, want unavailable", state.Status) + } +} + func TestProfileListCurrentFlagsUseStoredExactSelectors(t *testing.T) { first := authLogoutTestToken("corp_same") first.UserID = "user_1" diff --git a/internal/app/usage_recording.go b/internal/app/usage_recording.go index 4ef9f19e..b6ead5df 100644 --- a/internal/app/usage_recording.go +++ b/internal/app/usage_recording.go @@ -15,6 +15,7 @@ package app import ( "context" + "fmt" "github.com/DingTalk-Real-AI/dingtalk-workspace-cli/internal/shortcut/usage" "github.com/DingTalk-Real-AI/dingtalk-workspace-cli/pkg/edition" @@ -38,6 +39,21 @@ func (r recordingToolCaller) CallTool(ctx context.Context, product, tool string, return res, err } +func (r recordingToolCaller) CallToolWithToken( + ctx context.Context, + token, product, tool string, + args map[string]any, +) (*edition.ToolResult, error) { + inner, ok := r.inner.(tokenOverrideToolCaller) + if !ok { + return nil, fmt.Errorf("ToolCaller token override is not configured") + } + recordedArgs := cloneToolArgs(args) + res, err := inner.CallToolWithToken(ctx, token, product, tool, args) + usage.Append(product, tool, recordedArgs, err == nil, r.inner.DryRun()) + return res, err +} + func (r recordingToolCaller) Format() string { return r.inner.Format() } func (r recordingToolCaller) DryRun() bool { return r.inner.DryRun() } func (r recordingToolCaller) Fields() string { return r.inner.Fields() } diff --git a/internal/app/usage_recording_test.go b/internal/app/usage_recording_test.go index 5cc6cbfe..3b9705d4 100644 --- a/internal/app/usage_recording_test.go +++ b/internal/app/usage_recording_test.go @@ -25,6 +25,7 @@ import ( type crossPlatformCoverageCaller struct { args map[string]any + token string dryRun bool } @@ -33,6 +34,16 @@ func (c *crossPlatformCoverageCaller) CallTool(_ context.Context, _, _ string, a return &edition.ToolResult{}, nil } +func (c *crossPlatformCoverageCaller) CallToolWithToken( + _ context.Context, + token, _, _ string, + args map[string]any, +) (*edition.ToolResult, error) { + c.token = token + c.args = args + return &edition.ToolResult{}, nil +} + func (*crossPlatformCoverageCaller) Format() string { return "json" } func (c *crossPlatformCoverageCaller) DryRun() bool { return c.dryRun } func (*crossPlatformCoverageCaller) Fields() string { return "id,name" } @@ -88,6 +99,23 @@ func TestCrossPlatformCoverageRecordingToolCaller(t *testing.T) { t.Fatal("sensitive text must not be recorded") } + tokenCaller, ok := caller.(tokenOverrideToolCaller) + if !ok { + t.Fatal("recording caller dropped token override support") + } + if _, err := tokenCaller.CallToolWithToken( + context.Background(), + "temporary-token", + "contact", + "get_current_user_profile", + nil, + ); err != nil { + t.Fatal(err) + } + if inner.token != "temporary-token" || len(inner.args) != 0 { + t.Fatalf("token override forwarding = token %q args %#v", inner.token, inner.args) + } + inner.dryRun = true if _, err := caller.CallTool(context.Background(), "chat", "send_message", args); err != nil { t.Fatal(err) @@ -96,7 +124,7 @@ func TestCrossPlatformCoverageRecordingToolCaller(t *testing.T) { if err != nil { t.Fatal(err) } - if len(records) != 1 { + if len(records) != 2 { t.Fatalf("dry-run call must not be recorded: %#v", records) } } diff --git a/internal/auth/coverage_edges_test.go b/internal/auth/coverage_edges_test.go index 8154cf52..7bb6f4b2 100644 --- a/internal/auth/coverage_edges_test.go +++ b/internal/auth/coverage_edges_test.go @@ -1455,6 +1455,7 @@ func TestCrossPlatformCoverageTokenStorageAndRevocationCoverageEdges(t *testing. oldMarshalIndent := tokenJSONMarshalIndent oldMarshal := tokenJSONMarshal oldMkdir := tokenMkdirAll + oldRead := tokenReadFile oldWrite := tokenWriteFile oldRename := tokenRename oldRemove := tokenRemove @@ -1481,7 +1482,9 @@ func TestCrossPlatformCoverageTokenStorageAndRevocationCoverageEdges(t *testing. oldSync := tokenSyncLegacyMirror oldSyncOrganization := tokenSyncOrganizationMirror oldLoadProfiles := tokenLoadProfiles + oldSaveProfiles := tokenSaveProfiles oldWriteMarker := tokenWriteMarker + oldWriteManualMarker := tokenWriteManualMarker oldDeleteMarker := tokenDeleteMarker oldParseURL := tokenParseURL oldNewRequest := tokenNewRequest @@ -1498,6 +1501,7 @@ func TestCrossPlatformCoverageTokenStorageAndRevocationCoverageEdges(t *testing. tokenJSONMarshalIndent = oldMarshalIndent tokenJSONMarshal = oldMarshal tokenMkdirAll = oldMkdir + tokenReadFile = oldRead tokenWriteFile = oldWrite tokenRename = oldRename tokenRemove = oldRemove @@ -1524,7 +1528,9 @@ func TestCrossPlatformCoverageTokenStorageAndRevocationCoverageEdges(t *testing. tokenSyncLegacyMirror = oldSync tokenSyncOrganizationMirror = oldSyncOrganization tokenLoadProfiles = oldLoadProfiles + tokenSaveProfiles = oldSaveProfiles tokenWriteMarker = oldWriteMarker + tokenWriteManualMarker = oldWriteManualMarker tokenDeleteMarker = oldDeleteMarker tokenParseURL = oldParseURL tokenNewRequest = oldNewRequest @@ -1596,10 +1602,12 @@ func TestCrossPlatformCoverageTokenStorageAndRevocationCoverageEdges(t *testing. t.Fatalf("corp marker error = %v", err) } SetRuntimeProfile("") + tokenWriteManualMarker = func(string) error { return fail } if err := saveTokenDataLocked(dir, &TokenData{}); !errors.Is(err, fail) { t.Fatalf("legacy marker error = %v", err) } tokenWriteMarker = oldWriteMarker + tokenWriteManualMarker = oldWriteManualMarker edition.Override(&edition.Hooks{Name: "coverage", SaveToken: func(string, []byte) error { return nil }}) if err := saveTokenDataLocked(dir, data); err != nil { @@ -1654,6 +1662,9 @@ func TestCrossPlatformCoverageTokenStorageAndRevocationCoverageEdges(t *testing. t.Fatalf("failed migration load = %#v %v", got, err) } tokenSaveKeychain = func(*TokenData) error { return nil } + tokenLoadKeychain = func() (*TokenData, error) { return nil, ErrTokenDataNotFound } + tokenReadFile = func(string) ([]byte, error) { return nil, os.ErrNotExist } + tokenWriteManualMarker = func(string) error { return nil } deleted := false tokenDeleteSecure = func(string) error { deleted = true; return nil } if _, err := LoadTokenDataForProfile(dir, ""); err != nil || !deleted { @@ -1698,19 +1709,43 @@ func TestCrossPlatformCoverageTokenStorageAndRevocationCoverageEdges(t *testing. t.Fatalf("delete resolution error = %v", err) } tokenResolveDeletion = func(*ProfilesConfig, string) (*Profile, bool, error) { return selected, true, nil } + tokenLoadKeychainIdentity = func(string, string) (*TokenData, error) { + return &TokenData{CorpID: "corp", UserID: "user"}, nil + } + tokenLoadKeychainForCorpID = func(string) (*TokenData, error) { + return &TokenData{CorpID: "corp", UserID: "user"}, nil + } + tokenLoadKeychain = func() (*TokenData, error) { + return &TokenData{CorpID: "corp", UserID: "user"}, nil + } + tokenReadFile = func(string) ([]byte, error) { return nil, os.ErrNotExist } + tokenSaveProfiles = func(string, *ProfilesConfig) error { return nil } + tokenSaveKeychainForIdentity = func(string, string, *TokenData) error { return nil } + tokenSaveKeychainForCorpID = func(string, *TokenData) error { return nil } + tokenSaveKeychain = func(*TokenData) error { return nil } + tokenWriteMarker = func(string) error { return nil } + tokenDeleteMarker = func(string) error { return nil } + tokenRemoveProfile = func(string, string) (*Profile, error) { return selected, nil } + tokenSyncLegacyMirror = func(string) error { return nil } + tokenSyncOrganizationMirror = func(Profile) error { return nil } + tokenDeleteSecure = func(string) error { return nil } + tokenDeleteKeychainForCorpID = func(string) error { return nil } tokenDeleteKeychainIdentity = func(string, string) error { return fail } if err := deleteTokenDataForProfileLocked(dir, ""); !errors.Is(err, fail) { t.Fatalf("delete identity keychain error = %v", err) } tokenDeleteKeychainIdentity = func(string, string) error { return nil } - tokenRemoveProfile = func(string, string) (*Profile, error) { return selected, nil } - tokenSyncLegacyMirror = func(string) error { return nil } - tokenDeleteSecure = func(string) error { return nil } + cfg.OrgCurrentProfiles = map[string]string{"corp": ProfileSelector(*selected)} + tokenRemoveProfile = func(string, string) (*Profile, error) { + cfg.OrgCurrentProfiles = nil + return selected, nil + } tokenDeleteKeychainForCorpID = func(string) error { return fail } if err := deleteTokenDataForProfileLocked(dir, ""); !errors.Is(err, fail) { t.Fatalf("delete corp keychain error = %v", err) } tokenDeleteKeychainForCorpID = func(string) error { return nil } + cfg.OrgCurrentProfiles = map[string]string{"corp": ProfileSelector(*selected)} tokenRemoveProfile = func(string, string) (*Profile, error) { return nil, fail } if err := deleteTokenDataForProfileLocked(dir, ""); !errors.Is(err, fail) { t.Fatalf("remove profile error = %v", err) diff --git a/internal/auth/keychain_store.go b/internal/auth/keychain_store.go index 839780e3..16f33142 100644 --- a/internal/auth/keychain_store.go +++ b/internal/auth/keychain_store.go @@ -156,6 +156,9 @@ func preflightTokenPersistence(configDir string) error { if err != nil { return fmt.Errorf("load token profiles: %w", err) } + if err := ensureProfilesWritable(cfg); err != nil { + return err + } for _, profile := range cfg.Profiles { // LoadProfiles normalizes away blank and duplicate corp IDs. corpID := profile.CorpID @@ -191,10 +194,17 @@ func preflightTokenPersistence(configDir string) error { // preflightTokenRefreshPersistence checks only the slots a refresh can write. // An unrelated broken profile must not prevent the current profile from using // its still-valid credentials. -func preflightTokenRefreshPersistence(data *TokenData) error { +func preflightTokenRefreshPersistence(configDir string, data *TokenData) error { if h := edition.Get(); h.SaveToken != nil { return nil } + cfg, err := LoadProfiles(configDir) + if err != nil { + return err + } + if err := ensureProfilesWritable(cfg); err != nil { + return err + } if _, err := LoadTokenDataKeychain(); err != nil && !errors.Is(err, ErrTokenDataNotFound) { return fmt.Errorf("legacy token slot %q is unreadable: %w", keychain.AccountToken, err) @@ -209,11 +219,16 @@ func preflightTokenRefreshPersistence(data *TokenData) error { return fmt.Errorf("identity token slot %q is unreadable: %w", TokenAccountForIdentity(corpID, userID), err) } } + checkOrganizationMirror := true if _, _, exact := ParseIdentitySelector(RuntimeProfile()); exact { - return nil + checkOrganizationMirror = + exactProfileSelectorForCorp(cfg, corpID, cfg.OrgCurrentProfiles[corpID]) == + profileSelector(corpID, userID) } - if _, err := LoadTokenDataKeychainForCorpID(corpID); err != nil && !errors.Is(err, ErrTokenDataNotFound) { - return fmt.Errorf("profile token slot %q is unreadable: %w", TokenAccountForCorpID(corpID), err) + if checkOrganizationMirror { + if _, err := LoadTokenDataKeychainForCorpID(corpID); err != nil && !errors.Is(err, ErrTokenDataNotFound) { + return fmt.Errorf("profile token slot %q is unreadable: %w", TokenAccountForCorpID(corpID), err) + } } return nil } diff --git a/internal/auth/oauth_provider.go b/internal/auth/oauth_provider.go index 9fa4bcaf..a5c85e1a 100644 --- a/internal/auth/oauth_provider.go +++ b/internal/auth/oauth_provider.go @@ -659,7 +659,7 @@ func (p *OAuthProvider) lockedRefresh(ctx context.Context) (*TokenData, error) { if !data.IsRefreshTokenValid() { return nil, fmt.Errorf("refresh_token 已过期") } - if err := preflightTokenRefreshPersistence(data); err != nil { + if err := preflightTokenRefreshPersistence(p.configDir, data); err != nil { return nil, fmt.Errorf("%s: %w", i18n.T("本地登录态无法安全更新"), err) } diff --git a/internal/auth/profiles.go b/internal/auth/profiles.go index 0badcf10..dab1954f 100644 --- a/internal/auth/profiles.go +++ b/internal/auth/profiles.go @@ -31,28 +31,30 @@ import ( ) var ( - profilesAcquireDualLock = AcquireDualLock - profilesReadFile = os.ReadFile - profilesRename = os.Rename - profilesMkdirAll = os.MkdirAll - profilesMarshalIndent = json.MarshalIndent - profilesWriteFile = os.WriteFile - profilesRemove = os.Remove - profilesLoad = LoadProfiles - profilesSave = SaveProfiles - profilesEnsureMigration = ensureProfilesMigrationLocked - profilesSyncLegacyMirror = syncLegacyTokenMirrorLocked - profilesTokenExists = TokenDataExistsKeychain - profilesLoadLegacy = LoadTokenDataKeychain - profilesSaveCorp = SaveTokenDataKeychainForCorpID - profilesSaveIdentity = SaveTokenDataKeychainForIdentity - profilesTokenExistsCorp = TokenDataExistsKeychainForCorpID - profilesLoadCorp = LoadTokenDataKeychainForCorpID - profilesLoadIdentity = LoadTokenDataKeychainForIdentity - profilesSaveLegacy = SaveTokenDataKeychain - profilesWriteMarker = WriteTokenMarker - profilesDeleteLegacy = DeleteTokenDataKeychain - profilesDeleteMarker = DeleteTokenMarker + profilesAcquireDualLock = AcquireDualLock + profilesReadFile = os.ReadFile + profilesRename = os.Rename + profilesMkdirAll = os.MkdirAll + profilesMarshalIndent = json.MarshalIndent + profilesWriteFile = os.WriteFile + profilesRemove = os.Remove + profilesLoad = LoadProfiles + profilesSave = SaveProfiles + profilesEnsureMigration = ensureProfilesMigrationLocked + profilesSyncLegacyMirror = syncLegacyTokenMirrorLocked + profilesTokenExists = TokenDataExistsKeychain + profilesLoadLegacy = LoadTokenDataKeychain + profilesSaveCorp = SaveTokenDataKeychainForCorpID + profilesSaveIdentity = SaveTokenDataKeychainForIdentity + profilesTokenExistsCorp = TokenDataExistsKeychainForCorpID + profilesLoadCorp = LoadTokenDataKeychainForCorpID + profilesDeleteCorp = DeleteTokenDataKeychainForCorpID + profilesLoadIdentity = LoadTokenDataKeychainForIdentity + profilesSaveLegacy = SaveTokenDataKeychain + profilesWriteMarker = WriteTokenMarker + profilesWriteManualMarker = WriteManualTokenMarker + profilesDeleteLegacy = DeleteTokenDataKeychain + profilesDeleteMarker = DeleteTokenMarker ) // withProfilesLock runs fn while holding the auth dual-layer lock (process + @@ -78,9 +80,10 @@ const ( ) const ( - ProfileStatusActive = "active" - ProfileStatusExpired = "expired" - ProfileStatusRevoked = "revoked" + ProfileStatusActive = "active" + ProfileStatusExpired = "expired" + ProfileStatusRevoked = "revoked" + ProfileStatusUnavailable = "unavailable" ) // ProfilesConfig stores non-sensitive profile metadata. Token material stays in keychain. @@ -162,6 +165,9 @@ func SaveProfiles(configDir string, cfg *ProfilesConfig) error { if cfg == nil { cfg = &ProfilesConfig{} } + if err := ensureProfilesWritable(cfg); err != nil { + return err + } normalizeProfilesConfig(cfg) if err := profilesMkdirAll(configDir, config.DirPerm); err != nil { return fmt.Errorf("create config dir: %w", err) @@ -201,6 +207,9 @@ func ensureProfilesMigrationLocked(configDir string) error { if err != nil { return err } + if cfg.Version > profilesVersion { + return nil + } if len(cfg.Profiles) == 0 { // Version 2 with no profiles is an intentional logged-out tombstone. // Never resurrect a stale legacy mirror after logout/reset. @@ -240,9 +249,13 @@ func ensureProfilesMigrationLocked(configDir string) error { orgToken, loaded := orgTokens[corpID] if !loaded { token, loadErr := profilesLoadCorp(corpID) - if loadErr == nil { - orgToken = token + if legacySelectionState && loadErr != nil && !errors.Is(loadErr, ErrTokenDataNotFound) { + return loadErr } + if loadErr != nil { + token = nil + } + orgToken = token orgTokens[corpID] = orgToken } if orgToken == nil { @@ -367,6 +380,9 @@ func upsertProfileFromToken(configDir string, cfg *ProfilesConfig, data *TokenDa if corpID == "" { return nil } + if err := ensureProfilesWritable(cfg); err != nil { + return err + } normalizeProfilesConfig(cfg) cfg.Version = profilesVersion userID := strings.TrimSpace(data.UserID) @@ -453,7 +469,7 @@ func profileSelector(corpID, userID string) string { // ParseIdentitySelector splits corpId:userId selectors. func ParseIdentitySelector(selector string) (corpID, userID string, ok bool) { selector = strings.TrimSpace(selector) - idx := strings.LastIndex(selector, ":") + idx := strings.Index(selector, ":") if idx <= 0 || idx >= len(selector)-1 { return "", "", false } @@ -522,6 +538,9 @@ func ResolveProfileDeletionScope(configDir, selector string) (*Profile, bool, er if err != nil { return err } + if err := ensureProfilesWritable(cfg); err != nil { + return err + } var resolveErr error result, exact, resolveErr = resolveProfileDeletionSelection(cfg, selector) return resolveErr @@ -575,10 +594,18 @@ func setCurrentProfileLocked(configDir, selector string) (*Profile, error) { if err != nil { return nil, err } + if err := ensureProfilesWritable(cfg); err != nil { + return nil, err + } p, _, err := resolveProfileSelection(configDir, cfg, selector) if err != nil { return nil, err } + originalCfg := cloneProfilesConfig(cfg) + mirrors, err := snapshotProfileSelectionMirrors(configDir, p.CorpID) + if err != nil { + return nil, err + } var previousCurrent *Profile if strings.TrimSpace(cfg.CurrentProfile) != "" { previousCurrent, _, _ = resolveProfileSelection(configDir, cfg, cfg.CurrentProfile) @@ -595,10 +622,10 @@ func setCurrentProfileLocked(configDir, selector string) (*Profile, error) { return nil, err } if err := syncOrganizationTokenMirrorForProfile(*p); err != nil { - return nil, err + return nil, rollbackProfileSelection(configDir, originalCfg, p.CorpID, mirrors, err) } if err := profilesSyncLegacyMirror(configDir); err != nil { - return nil, err + return nil, rollbackProfileSelection(configDir, originalCfg, p.CorpID, mirrors, err) } return findExactProfile(cfg, p.CorpID, p.UserID), nil } @@ -622,6 +649,9 @@ func usePreviousProfileLocked(configDir string) (*Profile, error) { if err != nil { return nil, err } + if err := ensureProfilesWritable(cfg); err != nil { + return nil, err + } prev := strings.TrimSpace(cfg.PreviousProfile) if prev == "" { return nil, fmt.Errorf("previous profile is empty") @@ -630,6 +660,11 @@ func usePreviousProfileLocked(configDir string) (*Profile, error) { if err != nil { return nil, fmt.Errorf("resolve previous profile %q: %w", prev, err) } + originalCfg := cloneProfilesConfig(cfg) + mirrors, err := snapshotProfileSelectionMirrors(configDir, p.CorpID) + if err != nil { + return nil, err + } var current *Profile if strings.TrimSpace(cfg.CurrentProfile) != "" { current, _, _ = resolveProfileSelection(configDir, cfg, cfg.CurrentProfile) @@ -645,10 +680,10 @@ func usePreviousProfileLocked(configDir string) (*Profile, error) { return nil, err } if err := syncOrganizationTokenMirrorForProfile(*p); err != nil { - return nil, err + return nil, rollbackProfileSelection(configDir, originalCfg, p.CorpID, mirrors, err) } if err := profilesSyncLegacyMirror(configDir); err != nil { - return nil, err + return nil, rollbackProfileSelection(configDir, originalCfg, p.CorpID, mirrors, err) } return findExactProfile(cfg, p.CorpID, p.UserID), nil } @@ -672,6 +707,9 @@ func removeProfileLocked(configDir, selector string) (*Profile, error) { if err != nil { return nil, err } + if err := ensureProfilesWritable(cfg); err != nil { + return nil, err + } p, exact, err := resolveProfileDeletionSelection(cfg, selector) if err != nil { return nil, err @@ -771,6 +809,9 @@ func markProfileStatusLocked(configDir, selector, status string) error { if err != nil { return err } + if err := ensureProfilesWritable(cfg); err != nil { + return err + } p, _, resolveErr := resolveProfileSelection(configDir, cfg, selector) if resolveErr != nil { return nil @@ -780,6 +821,90 @@ func markProfileStatusLocked(configDir, selector, status string) error { return profilesSave(configDir, cfg) } +func ensureProfilesWritable(cfg *ProfilesConfig) error { + if cfg != nil && cfg.Version > profilesVersion { + return fmt.Errorf( + "profiles.json version %d is newer than supported version %d; upgrade dws before changing profiles", + cfg.Version, + profilesVersion, + ) + } + return nil +} + +type profileSelectionMirrorSnapshot struct { + organization tokenSlotSnapshot + legacy tokenSlotSnapshot + marker tokenMarkerSnapshot +} + +func snapshotProfileSelectionMirrors(configDir, corpID string) (profileSelectionMirrorSnapshot, error) { + organization, err := snapshotTokenSlot(func() (*TokenData, error) { + return profilesLoadCorp(corpID) + }) + if err != nil { + return profileSelectionMirrorSnapshot{}, err + } + legacy, err := snapshotTokenSlot(profilesLoadLegacy) + if err != nil { + return profileSelectionMirrorSnapshot{}, err + } + marker, err := snapshotTokenMarker(configDir) + if err != nil { + return profileSelectionMirrorSnapshot{}, err + } + return profileSelectionMirrorSnapshot{ + organization: organization, + legacy: legacy, + marker: marker, + }, nil +} + +func rollbackProfileSelection( + configDir string, + cfg *ProfilesConfig, + corpID string, + mirrors profileSelectionMirrorSnapshot, + operationErr error, +) error { + var rollbackErr error + if err := profilesSave(configDir, cloneProfilesConfig(cfg)); err != nil { + rollbackErr = errors.Join(rollbackErr, err) + } + if mirrors.organization.exists { + if err := profilesSaveCorp(corpID, mirrors.organization.token); err != nil { + rollbackErr = errors.Join(rollbackErr, err) + } + } else if err := profilesDeleteCorp(corpID); err != nil { + rollbackErr = errors.Join(rollbackErr, err) + } + if mirrors.legacy.exists { + if err := profilesSaveLegacy(mirrors.legacy.token); err != nil { + rollbackErr = errors.Join(rollbackErr, err) + } + } else if err := profilesDeleteLegacy(); err != nil { + rollbackErr = errors.Join(rollbackErr, err) + } + switch { + case !mirrors.marker.exists: + if err := profilesDeleteMarker(configDir); err != nil { + rollbackErr = errors.Join(rollbackErr, err) + } + case mirrors.marker.manual: + if err := profilesWriteManualMarker(configDir); err != nil { + rollbackErr = errors.Join(rollbackErr, err) + } + default: + if err := profilesWriteMarker(configDir); err != nil { + rollbackErr = errors.Join(rollbackErr, err) + } + } + if rollbackErr != nil { + return errors.Join(operationErr, fmt.Errorf("rollback profile selection: %w", rollbackErr)) + } + return operationErr +} + // SyncLegacyTokenMirror mirrors the current profile token into legacy auth-token. func SyncLegacyTokenMirror(configDir string) error { return withProfilesLock(configDir, func() error { @@ -800,9 +925,14 @@ func syncLegacyTokenMirrorLocked(configDir string) error { } data, loadErr := loadTokenForProfileIdentity(*p) if loadErr != nil { - // Keep the existing legacy mirror untouched rather than wiping a host - // app's login state just because keychain was momentarily unavailable. - return nil + if errors.Is(loadErr, ErrTokenDataNotFound) { + // The current identity and its organization mirror are both + // confirmed absent. Clear the stale global compatibility mirror. + } else { + // Keep the existing legacy mirror untouched rather than wiping a host + // app's login state just because keychain was momentarily unavailable. + return nil + } } if data != nil { if err := profilesSaveLegacy(data); err != nil { diff --git a/internal/auth/remaining_coverage_test.go b/internal/auth/remaining_coverage_test.go index 828d635f..60b2c8b1 100644 --- a/internal/auth/remaining_coverage_test.go +++ b/internal/auth/remaining_coverage_test.go @@ -46,7 +46,7 @@ func TestCrossPlatformCoverageTokenPersistencePreflightRemainingEdges(t *testing if err := preflightTokenPersistence(t.TempDir()); err != nil { t.Fatal(err) } - if err := preflightTokenRefreshPersistence(&TokenData{CorpID: "corp"}); err != nil { + if err := preflightTokenRefreshPersistence(t.TempDir(), &TokenData{CorpID: "corp"}); err != nil { t.Fatal(err) } @@ -79,7 +79,7 @@ func TestCrossPlatformCoverageTokenPersistencePreflightRemainingEdges(t *testing if err := preflightTokenPersistence(dir); !errors.Is(err, fail) { t.Fatalf("profile slot error = %v", err) } - if err := preflightTokenRefreshPersistence(&TokenData{CorpID: "corp"}); !errors.Is(err, fail) { + if err := preflightTokenRefreshPersistence(t.TempDir(), &TokenData{CorpID: "corp"}); !errors.Is(err, fail) { t.Fatalf("refresh profile slot error = %v", err) } diff --git a/internal/auth/token.go b/internal/auth/token.go index 16bcf664..1aa44462 100644 --- a/internal/auth/token.go +++ b/internal/auth/token.go @@ -36,6 +36,7 @@ var ( tokenJSONMarshalIndent = json.MarshalIndent tokenJSONMarshal = json.Marshal tokenMkdirAll = os.MkdirAll + tokenReadFile = os.ReadFile tokenWriteFile = os.WriteFile tokenRename = os.Rename tokenRemove = os.Remove @@ -64,7 +65,9 @@ var ( tokenSyncLegacyMirror = syncLegacyTokenMirrorLocked tokenSyncOrganizationMirror = syncOrganizationTokenMirrorForProfile tokenLoadProfiles = LoadProfiles + tokenSaveProfiles = SaveProfiles tokenWriteMarker = WriteTokenMarker + tokenWriteManualMarker = WriteManualTokenMarker tokenDeleteMarker = DeleteTokenMarker tokenParseURL = url.Parse tokenNewRequest = http.NewRequestWithContext @@ -121,14 +124,29 @@ const tokenJSONFile = "token.json" // TokenMarker is a lightweight file the host application reads to detect // whether the CLI has a valid token without accessing the keychain. type TokenMarker struct { - UpdatedAt string `json:"updated_at"` + UpdatedAt string `json:"updated_at"` + ManualToken bool `json:"manual_token,omitempty"` } // WriteTokenMarker writes a token.json marker containing only an updated_at // timestamp. The host application uses this file's presence and mtime to // decide whether it needs to trigger a new auth exchange. func WriteTokenMarker(configDir string) error { - marker := TokenMarker{UpdatedAt: time.Now().Format(time.RFC3339)} + return writeTokenMarker(configDir, false) +} + +// WriteManualTokenMarker marks the legacy global keychain slot as an explicit +// `auth login --token` credential. The additive field keeps older hosts, which +// only inspect token.json presence and mtime, fully compatible. +func WriteManualTokenMarker(configDir string) error { + return writeTokenMarker(configDir, true) +} + +func writeTokenMarker(configDir string, manual bool) error { + marker := TokenMarker{ + UpdatedAt: time.Now().Format(time.RFC3339), + ManualToken: manual, + } data, _ := tokenJSONMarshalIndent(marker, "", " ") if err := tokenMkdirAll(configDir, 0o700); err != nil { return err @@ -140,6 +158,23 @@ func WriteTokenMarker(configDir string) error { return tokenRename(tmp, filepath.Join(configDir, tokenJSONFile)) } +func manualTokenMarkerActive(configDir string) (bool, error) { + data, err := tokenReadFile(filepath.Join(configDir, tokenJSONFile)) + if err != nil { + if os.IsNotExist(err) { + return false, nil + } + return false, err + } + var marker TokenMarker + if err := json.Unmarshal(data, &marker); err != nil { + // Historical hosts only require the marker's presence. A malformed old + // marker must not make profile authentication unusable. + return false, nil + } + return marker.ManualToken, nil +} + // DeleteTokenMarker removes the token.json marker file. func DeleteTokenMarker(configDir string) error { if err := tokenRemove(filepath.Join(configDir, tokenJSONFile)); err != nil && !os.IsNotExist(err) { @@ -172,51 +207,98 @@ func saveTokenDataLocked(configDir string, data *TokenData) error { if data != nil && strings.TrimSpace(data.CorpID) != "" { corpID := strings.TrimSpace(data.CorpID) userID := strings.TrimSpace(data.UserID) + cfg, err := tokenLoadProfiles(configDir) + if err != nil { + return err + } + if err := ensureProfilesWritable(cfg); err != nil { + return err + } + runtimeSelector := strings.TrimSpace(RuntimeProfile()) + makeCurrent := runtimeSelector == "" + exactSelector := profileSelector(corpID, userID) + mirrorOrg := makeCurrent || + exactProfileSelectorForCorp(cfg, corpID, cfg.OrgCurrentProfiles[corpID]) == exactSelector + snapshot, err := snapshotTokenPersistence(configDir, cfg, corpID, userID, mirrorOrg) + if err != nil { + return err + } + preserveManualDefault := !makeCurrent && + snapshot.marker.known && + snapshot.marker.exists && + snapshot.marker.manual + rollback := func(operationErr error) error { + if rollbackErr := restoreTokenPersistence(configDir, snapshot); rollbackErr != nil { + return errors.Join(operationErr, fmt.Errorf("rollback token persistence: %w", rollbackErr)) + } + return operationErr + } if userID != "" { if err := tokenSaveKeychainForIdentity(corpID, userID, data); err != nil { - return err + return rollback(err) } } else { - cfg, err := tokenLoadProfiles(configDir) - if err != nil { - return err - } for _, profile := range cfg.Profiles { if strings.TrimSpace(profile.CorpID) == corpID && strings.TrimSpace(profile.UserID) != "" { return fmt.Errorf("cannot store profile for corpId %q without userId because account identities already exist", corpID) } } } - runtimeSelector := strings.TrimSpace(RuntimeProfile()) - makeCurrent := runtimeSelector == "" if err := tokenUpsertProfile(configDir, data, makeCurrent); err != nil { - return err + return rollback(err) } - cfg, err := tokenLoadProfiles(configDir) - if err != nil { - return err - } - exactSelector := profileSelector(corpID, userID) - mirrorOrg := makeCurrent || - exactProfileSelectorForCorp(cfg, corpID, cfg.OrgCurrentProfiles[corpID]) == exactSelector if mirrorOrg { if err := tokenSaveKeychainForCorpID(corpID, data); err != nil { - return err + return rollback(err) } } if makeCurrent { if err := tokenSaveKeychain(data); err != nil { - return err + return rollback(err) + } + } else if !preserveManualDefault { + if err := tokenSyncLegacyMirror(configDir); err != nil { + return rollback(err) } - } else if err := tokenSyncLegacyMirror(configDir); err != nil { - return err } - return tokenWriteMarker(configDir) + if preserveManualDefault { + if err := tokenWriteManualMarker(configDir); err != nil { + return rollback(err) + } + } else if err := tokenWriteMarker(configDir); err != nil { + return rollback(err) + } + return nil + } + legacySnapshot, err := snapshotTokenSlot(tokenLoadKeychain) + if err != nil { + return err + } + markerSnapshot, err := snapshotTokenMarker(configDir) + if err != nil { + return err } if err := tokenSaveKeychain(data); err != nil { return err } - return tokenWriteMarker(configDir) + if err := tokenWriteManualMarker(configDir); err != nil { + var rollbackErr error + if restoreErr := restoreTokenSlot( + legacySnapshot, + tokenSaveKeychain, + tokenDeleteKeychain, + ); restoreErr != nil { + rollbackErr = errors.Join(rollbackErr, restoreErr) + } + if restoreErr := restoreTokenMarker(configDir, markerSnapshot); restoreErr != nil { + rollbackErr = errors.Join(rollbackErr, restoreErr) + } + if rollbackErr != nil { + return errors.Join(err, fmt.Errorf("rollback manual token persistence: %w", rollbackErr)) + } + return err + } + return nil } func saveTokenViaHook(h *edition.Hooks, configDir string, data *TokenData) error { @@ -263,6 +345,21 @@ func LoadTokenDataForProfile(configDir, profile string) (*TokenData, error) { func loadTokenDataForProfileLocked(configDir, profile string) (*TokenData, error) { // Default: keychain with legacy .data migration + if strings.TrimSpace(profile) == "" { + manual, err := manualTokenMarkerActive(configDir) + if err != nil { + return nil, err + } + if manual { + data, loadErr := tokenLoadKeychain() + if loadErr == nil && data != nil && strings.TrimSpace(data.CorpID) == "" { + return data, nil + } + if loadErr != nil && !errors.Is(loadErr, ErrTokenDataNotFound) { + return nil, loadErr + } + } + } selected, err := tokenResolveProfile(configDir, profile) if err != nil { return nil, err @@ -362,6 +459,15 @@ func DeleteTokenDataForProfile(configDir, profile string) error { } func deleteTokenDataForProfileLocked(configDir, profile string) error { + if strings.TrimSpace(profile) == "" { + manual, err := manualTokenMarkerActive(configDir) + if err != nil { + return err + } + if manual { + return deleteManualTokenDataLocked(configDir) + } + } if err := profilesEnsureMigration(configDir); err != nil { return err } @@ -383,58 +489,84 @@ func deleteTokenDataForProfileLocked(configDir, profile string) error { } if selected != nil { removed := *selected + originalCfg := cloneProfilesConfig(cfg) + identitySnapshots, err := snapshotDeletionIdentities(cfg, removed, exact) + if err != nil { + return err + } + orgSnapshot := snapshotTokenSlotForDeletion(func() (*TokenData, error) { + return tokenLoadKeychainForCorpID(removed.CorpID) + }) + legacySnapshot := snapshotTokenSlotForDeletion(tokenLoadKeychain) + markerSnapshot := snapshotTokenMarkerForDeletion(configDir) + + // Clean the deprecated secure-store copy before changing the profile + // transaction. A cleanup failure therefore leaves all current metadata + // and keychain slots untouched. + if err := tokenDeleteSecure(configDir); err != nil { + return err + } + + removeSelector := removed.CorpID + orgCurrent := false if exact { - orgCurrent := exactProfileSelectorForCorp( + removeSelector = ProfileSelector(removed) + orgCurrent = exactProfileSelectorForCorp( cfg, removed.CorpID, cfg.OrgCurrentProfiles[removed.CorpID], ) == ProfileSelector(removed) - if strings.TrimSpace(removed.UserID) != "" { - if err := tokenDeleteKeychainIdentity(removed.CorpID, removed.UserID); err != nil { - return err - } - } - if _, err := tokenRemoveProfile(configDir, ProfileSelector(removed)); err != nil { - return err - } - if orgCurrent { - if err := tokenDeleteKeychainForCorpID(removed.CorpID); err != nil { - return err - } - updated, loadErr := tokenLoadProfiles(configDir) - if loadErr != nil { - return loadErr - } - if replacementSelector := updated.OrgCurrentProfiles[removed.CorpID]; replacementSelector != "" { - replacement, _, resolveErr := tokenResolveSelection(configDir, updated, replacementSelector) - if resolveErr != nil { - return resolveErr - } - if err := tokenSyncOrganizationMirror(*replacement); err != nil { - return err - } - } - } - } else { - for _, candidate := range cfg.Profiles { - if strings.TrimSpace(candidate.CorpID) != strings.TrimSpace(removed.CorpID) || strings.TrimSpace(candidate.UserID) == "" { - continue - } - if err := tokenDeleteKeychainIdentity(candidate.CorpID, candidate.UserID); err != nil { - return err - } - } - if err := tokenDeleteKeychainForCorpID(removed.CorpID); err != nil { - return err - } - if _, err := tokenRemoveProfile(configDir, removed.CorpID); err != nil { - return err - } } - if err := tokenSyncLegacyMirror(configDir); err != nil { + if _, err := tokenRemoveProfile(configDir, removeSelector); err != nil { return err } - return tokenDeleteSecure(configDir) + rollback := func(operationErr error) error { + if rollbackErr := restoreProfileDeletion( + configDir, + originalCfg, + identitySnapshots, + removed.CorpID, + orgSnapshot, + legacySnapshot, + markerSnapshot, + ); rollbackErr != nil { + return errors.Join(operationErr, fmt.Errorf("rollback profile deletion: %w", rollbackErr)) + } + return operationErr + } + + if !exact || orgCurrent { + updated, loadErr := tokenLoadProfiles(configDir) + if loadErr != nil { + return rollback(loadErr) + } + replacementSelector := updated.OrgCurrentProfiles[removed.CorpID] + if exact && replacementSelector != "" { + replacement, _, resolveErr := tokenResolveSelection(configDir, updated, replacementSelector) + if resolveErr != nil { + return rollback(resolveErr) + } + if err := tokenSyncOrganizationMirror(*replacement); err != nil { + return rollback(err) + } + } else if err := tokenDeleteKeychainForCorpID(removed.CorpID); err != nil { + return rollback(err) + } + } + preserveManualDefault := markerSnapshot.known && + markerSnapshot.exists && + markerSnapshot.manual + if !preserveManualDefault { + if err := tokenSyncLegacyMirror(configDir); err != nil { + return rollback(err) + } + } + for _, snapshot := range identitySnapshots { + if err := tokenDeleteKeychainIdentity(snapshot.profile.CorpID, snapshot.profile.UserID); err != nil { + return rollback(err) + } + } + return nil } keychainErr := tokenDeleteKeychain() @@ -449,6 +581,308 @@ func deleteTokenDataForProfileLocked(configDir, profile string) error { return markerErr } +func deleteManualTokenDataLocked(configDir string) error { + legacySnapshot := snapshotTokenSlotForDeletion(tokenLoadKeychain) + if err := tokenDeleteSecure(configDir); err != nil { + return err + } + if err := tokenDeleteKeychain(); err != nil { + return err + } + if err := tokenDeleteMarker(configDir); err != nil { + if legacySnapshot.known { + if rollbackErr := restoreTokenSlot( + legacySnapshot, + tokenSaveKeychain, + tokenDeleteKeychain, + ); rollbackErr != nil { + return errors.Join(err, fmt.Errorf("rollback manual token deletion: %w", rollbackErr)) + } + } + return err + } + return nil +} + +type deletionIdentitySnapshot struct { + profile Profile + token *TokenData +} + +type tokenSlotSnapshot struct { + token *TokenData + known bool + exists bool +} + +type tokenMarkerSnapshot struct { + known bool + exists bool + manual bool +} + +type tokenPersistenceSnapshot struct { + profiles *ProfilesConfig + corpID string + userID string + identity tokenSlotSnapshot + org tokenSlotSnapshot + legacy tokenSlotSnapshot + marker tokenMarkerSnapshot +} + +func cloneProfilesConfig(cfg *ProfilesConfig) *ProfilesConfig { + if cfg == nil { + return nil + } + cloned := *cfg + cloned.Profiles = append([]Profile(nil), cfg.Profiles...) + for i := range cloned.Profiles { + cloned.Profiles[i].AuthorizedDomains = append( + []string(nil), + cloned.Profiles[i].AuthorizedDomains..., + ) + } + if cfg.OrgCurrentProfiles != nil { + cloned.OrgCurrentProfiles = make(map[string]string, len(cfg.OrgCurrentProfiles)) + for corpID, selector := range cfg.OrgCurrentProfiles { + cloned.OrgCurrentProfiles[corpID] = selector + } + } + return &cloned +} + +func snapshotDeletionIdentities(cfg *ProfilesConfig, removed Profile, exact bool) ([]deletionIdentitySnapshot, error) { + var profiles []Profile + if exact { + profiles = []Profile{removed} + } else { + for _, candidate := range cfg.Profiles { + if strings.TrimSpace(candidate.CorpID) == strings.TrimSpace(removed.CorpID) { + profiles = append(profiles, candidate) + } + } + } + snapshots := make([]deletionIdentitySnapshot, 0, len(profiles)) + for _, candidate := range profiles { + if strings.TrimSpace(candidate.UserID) == "" { + continue + } + data, err := tokenLoadKeychainIdentity(candidate.CorpID, candidate.UserID) + if err != nil { + if errors.Is(err, ErrTokenDataNotFound) { + snapshots = append(snapshots, deletionIdentitySnapshot{profile: candidate}) + continue + } + // A damaged target slot must remain removable. It cannot be restored + // during rollback, but every readable slot in the same transaction + // still is. + snapshots = append(snapshots, deletionIdentitySnapshot{profile: candidate}) + continue + } + snapshots = append(snapshots, deletionIdentitySnapshot{profile: candidate, token: data}) + } + return snapshots, nil +} + +func snapshotTokenSlot(load func() (*TokenData, error)) (tokenSlotSnapshot, error) { + data, err := load() + if err != nil { + if errors.Is(err, ErrTokenDataNotFound) { + return tokenSlotSnapshot{known: true}, nil + } + return tokenSlotSnapshot{}, err + } + return tokenSlotSnapshot{token: data, known: true, exists: data != nil}, nil +} + +func snapshotTokenSlotForDeletion(load func() (*TokenData, error)) tokenSlotSnapshot { + snapshot, err := snapshotTokenSlot(load) + if err != nil { + return tokenSlotSnapshot{} + } + return snapshot +} + +func snapshotTokenMarker(configDir string) (tokenMarkerSnapshot, error) { + data, err := tokenReadFile(filepath.Join(configDir, tokenJSONFile)) + if err != nil { + if os.IsNotExist(err) { + return tokenMarkerSnapshot{known: true}, nil + } + return tokenMarkerSnapshot{}, err + } + var marker TokenMarker + if err := json.Unmarshal(data, &marker); err != nil { + return tokenMarkerSnapshot{known: true, exists: true}, nil + } + return tokenMarkerSnapshot{known: true, exists: true, manual: marker.ManualToken}, nil +} + +func snapshotTokenMarkerForDeletion(configDir string) tokenMarkerSnapshot { + snapshot, err := snapshotTokenMarker(configDir) + if err != nil { + return tokenMarkerSnapshot{} + } + return snapshot +} + +func restoreProfileDeletion( + configDir string, + cfg *ProfilesConfig, + identities []deletionIdentitySnapshot, + corpID string, + org tokenSlotSnapshot, + legacy tokenSlotSnapshot, + marker tokenMarkerSnapshot, +) error { + var rollbackErr error + for _, snapshot := range identities { + if snapshot.token == nil { + continue + } + if err := tokenSaveKeychainForIdentity( + snapshot.profile.CorpID, + snapshot.profile.UserID, + snapshot.token, + ); err != nil { + rollbackErr = errors.Join(rollbackErr, err) + } + } + if err := tokenSaveProfiles(configDir, cloneProfilesConfig(cfg)); err != nil { + rollbackErr = errors.Join(rollbackErr, err) + } + if org.known { + if org.exists { + if err := tokenSaveKeychainForCorpID(corpID, org.token); err != nil { + rollbackErr = errors.Join(rollbackErr, err) + } + } else if err := tokenDeleteKeychainForCorpID(corpID); err != nil { + rollbackErr = errors.Join(rollbackErr, err) + } + } + if legacy.known { + if legacy.exists { + if err := tokenSaveKeychain(legacy.token); err != nil { + rollbackErr = errors.Join(rollbackErr, err) + } + } else if err := tokenDeleteKeychain(); err != nil { + rollbackErr = errors.Join(rollbackErr, err) + } + } + if marker.known { + if err := restoreTokenMarker(configDir, marker); err != nil { + rollbackErr = errors.Join(rollbackErr, err) + } + } + return rollbackErr +} + +func snapshotTokenPersistence( + configDir string, + cfg *ProfilesConfig, + corpID, userID string, + includeOrganization bool, +) (tokenPersistenceSnapshot, error) { + snapshot := tokenPersistenceSnapshot{ + profiles: cloneProfilesConfig(cfg), + corpID: corpID, + userID: userID, + } + var err error + if strings.TrimSpace(userID) != "" { + snapshot.identity, err = snapshotTokenSlot(func() (*TokenData, error) { + return tokenLoadKeychainIdentity(corpID, userID) + }) + if err != nil { + return tokenPersistenceSnapshot{}, err + } + } + if includeOrganization { + snapshot.org, err = snapshotTokenSlot(func() (*TokenData, error) { + return tokenLoadKeychainForCorpID(corpID) + }) + if err != nil { + return tokenPersistenceSnapshot{}, err + } + } + snapshot.legacy, err = snapshotTokenSlot(tokenLoadKeychain) + if err != nil { + return tokenPersistenceSnapshot{}, err + } + snapshot.marker, err = snapshotTokenMarker(configDir) + if err != nil { + return tokenPersistenceSnapshot{}, err + } + return snapshot, nil +} + +func restoreTokenPersistence(configDir string, snapshot tokenPersistenceSnapshot) error { + var rollbackErr error + if err := tokenSaveProfiles(configDir, cloneProfilesConfig(snapshot.profiles)); err != nil { + rollbackErr = errors.Join(rollbackErr, err) + } + if strings.TrimSpace(snapshot.userID) != "" && snapshot.identity.known { + if err := restoreTokenSlot( + snapshot.identity, + func(data *TokenData) error { + return tokenSaveKeychainForIdentity(snapshot.corpID, snapshot.userID, data) + }, + func() error { + return tokenDeleteKeychainIdentity(snapshot.corpID, snapshot.userID) + }, + ); err != nil { + rollbackErr = errors.Join(rollbackErr, err) + } + } + if snapshot.org.known { + if err := restoreTokenSlot( + snapshot.org, + func(data *TokenData) error { + return tokenSaveKeychainForCorpID(snapshot.corpID, data) + }, + func() error { + return tokenDeleteKeychainForCorpID(snapshot.corpID) + }, + ); err != nil { + rollbackErr = errors.Join(rollbackErr, err) + } + } + if snapshot.legacy.known { + if err := restoreTokenSlot(snapshot.legacy, tokenSaveKeychain, tokenDeleteKeychain); err != nil { + rollbackErr = errors.Join(rollbackErr, err) + } + } + if snapshot.marker.known { + if err := restoreTokenMarker(configDir, snapshot.marker); err != nil { + rollbackErr = errors.Join(rollbackErr, err) + } + } + return rollbackErr +} + +func restoreTokenSlot( + snapshot tokenSlotSnapshot, + save func(*TokenData) error, + remove func() error, +) error { + if snapshot.exists { + return save(snapshot.token) + } + return remove() +} + +func restoreTokenMarker(configDir string, marker tokenMarkerSnapshot) error { + switch { + case !marker.exists: + return tokenDeleteMarker(configDir) + case marker.manual: + return tokenWriteManualMarker(configDir) + default: + return tokenWriteMarker(configDir) + } +} + // DeleteAllTokenData removes all profile-scoped and legacy token data. func DeleteAllTokenData(configDir string) error { if h := edition.Get(); h.DeleteToken != nil { diff --git a/internal/auth/token_preflight_test.go b/internal/auth/token_preflight_test.go index 5d495c78..059fd309 100644 --- a/internal/auth/token_preflight_test.go +++ b/internal/auth/token_preflight_test.go @@ -18,6 +18,7 @@ package auth import ( "bytes" "context" + "encoding/json" "errors" "io" "net/http" @@ -150,7 +151,7 @@ func TestPreflightTokenPersistenceRejectsUnreadableProfileSlot(t *testing.T) { } } -func TestExactProfileRefreshIgnoresUnreadableOrgMirror(t *testing.T) { +func TestExactOrgCurrentRefreshRejectsUnreadableOrgMirror(t *testing.T) { cleanupKeychain(t) t.Setenv(keychain.DisableKeychainEnv, "1") configDir := t.TempDir() @@ -165,8 +166,47 @@ func TestExactProfileRefreshIgnoresUnreadableOrgMirror(t *testing.T) { SetRuntimeProfile("corp_exact:user_exact") defer SetRuntimeProfile("") - if err := preflightTokenRefreshPersistence(data); err != nil { - t.Fatalf("preflightTokenRefreshPersistence(exact) error = %v", err) + if err := preflightTokenRefreshPersistence(configDir, data); err == nil || + !strings.Contains(err.Error(), "profile token slot") { + t.Fatalf("preflightTokenRefreshPersistence(exact current) error = %v, want unreadable org mirror", err) + } +} + +func TestExactNonOrgCurrentRefreshIgnoresUnreadableOrgMirror(t *testing.T) { + cleanupKeychain(t) + t.Setenv(keychain.DisableKeychainEnv, "1") + configDir := t.TempDir() + first := testToken("at_first", "corp_exact", "Exact Org") + first.UserID = "user_first" + second := testToken("at_second", "corp_exact", "Exact Org") + second.UserID = "user_second" + if err := SaveTokenData(configDir, first); err != nil { + t.Fatalf("SaveTokenData(first) error = %v", err) + } + if err := SaveTokenData(configDir, second); err != nil { + t.Fatalf("SaveTokenData(second) error = %v", err) + } + if err := os.WriteFile(profileCiphertextPathForTest(first.CorpID), []byte("corrupt ciphertext"), 0o600); err != nil { + t.Fatalf("WriteFile(profile ciphertext) error = %v", err) + } + + SetRuntimeProfile("corp_exact:user_first") + defer SetRuntimeProfile("") + if err := preflightTokenRefreshPersistence(configDir, first); err != nil { + t.Fatalf("preflightTokenRefreshPersistence(exact non-current) error = %v", err) + } + updated := *first + updated.AccessToken = "at_first_refreshed" + updated.RefreshToken = "rt_first_refreshed" + if err := SaveTokenData(configDir, &updated); err != nil { + t.Fatalf("SaveTokenData(exact non-current refresh) error = %v", err) + } + loaded, err := LoadTokenDataForProfile(configDir, "corp_exact:user_first") + if err != nil { + t.Fatalf("LoadTokenDataForProfile(refreshed) error = %v", err) + } + if loaded.AccessToken != updated.AccessToken || loaded.RefreshToken != updated.RefreshToken { + t.Fatalf("refreshed exact token = %#v, want %#v", loaded, updated) } } @@ -252,7 +292,7 @@ func TestRefreshPreflightIgnoresUnreadableUnrelatedProfile(t *testing.T) { t.Fatalf("WriteFile(A profile ciphertext) error = %v", err) } - if err := preflightTokenRefreshPersistence(dataB); err != nil { + if err := preflightTokenRefreshPersistence(configDir, dataB); err != nil { t.Fatalf("preflightTokenRefreshPersistence(B) error = %v", err) } loaded, err := NewOAuthProvider(configDir, nil).Login(context.Background(), false) @@ -363,6 +403,45 @@ func TestLockedRefreshPreflightsLegacyMirrorBeforeHTTP(t *testing.T) { } } +func TestLockedRefreshRejectsFutureProfilesVersionBeforeHTTP(t *testing.T) { + cleanupKeychain(t) + t.Setenv(keychain.DisableKeychainEnv, "1") + setPreflightTestCredentials(t) + configDir := t.TempDir() + data := testToken("at_future_refresh", "corp_future", "Future Org") + data.ExpiresAt = time.Now().Add(-time.Hour) + data.RefreshExpAt = time.Now().Add(time.Hour) + if err := SaveTokenData(configDir, data); err != nil { + t.Fatalf("SaveTokenData() error = %v", err) + } + cfg, err := LoadProfiles(configDir) + if err != nil { + t.Fatalf("LoadProfiles() error = %v", err) + } + cfg.Version = profilesVersion + 1 + raw, err := json.Marshal(cfg) + if err != nil { + t.Fatalf("json.Marshal() error = %v", err) + } + if err := os.WriteFile(ProfilesPath(configDir), raw, 0o600); err != nil { + t.Fatalf("write future profiles: %v", err) + } + + var calls atomic.Int32 + provider := NewOAuthProvider(configDir, nil) + provider.httpClient = &http.Client{Transport: preflightRoundTripFunc(func(*http.Request) (*http.Response, error) { + calls.Add(1) + return nil, errors.New("unexpected refresh request") + })} + _, err = provider.Login(context.Background(), false) + if err == nil || !strings.Contains(err.Error(), "newer than supported") { + t.Fatalf("Login() error = %v, want future profiles rejection", err) + } + if got := calls.Load(); got != 0 { + t.Fatalf("refresh HTTP calls = %d, want 0", got) + } +} + func TestExchangeAuthCodeAllowsFirstLogin(t *testing.T) { cleanupKeychain(t) t.Setenv(keychain.DisableKeychainEnv, "1") diff --git a/internal/auth/token_test.go b/internal/auth/token_test.go index be6c5adf..0e345f0b 100644 --- a/internal/auth/token_test.go +++ b/internal/auth/token_test.go @@ -17,6 +17,7 @@ import ( "context" "errors" "os" + "path/filepath" "strings" "testing" "time" @@ -133,6 +134,141 @@ func TestTokenOverwrite(t *testing.T) { } } +func TestManualTokenRemainsLoadableWithV2Profiles(t *testing.T) { + cleanupKeychain(t) + configDir := t.TempDir() + + profileToken := testToken("profile-token", "corp_same", "同一组织") + profileToken.UserID = "user_1" + if err := SaveTokenData(configDir, profileToken); err != nil { + t.Fatalf("SaveTokenData(profile) error = %v", err) + } + + manualToken := &TokenData{ + AccessToken: "manual-token", + ExpiresAt: time.Now().Add(time.Hour), + } + if err := SaveTokenData(configDir, manualToken); err != nil { + t.Fatalf("SaveTokenData(manual) error = %v", err) + } + + loaded, err := LoadTokenData(configDir) + if err != nil { + t.Fatalf("LoadTokenData() error = %v", err) + } + if loaded.AccessToken != manualToken.AccessToken { + t.Fatalf("default token = %q, want manual token", loaded.AccessToken) + } + + explicit, err := LoadTokenDataForProfile(configDir, ProfileSelector(Profile{ + CorpID: profileToken.CorpID, + UserID: profileToken.UserID, + })) + if err != nil { + t.Fatalf("LoadTokenDataForProfile() error = %v", err) + } + if explicit.AccessToken != profileToken.AccessToken { + t.Fatalf("explicit profile token = %q, want %q", explicit.AccessToken, profileToken.AccessToken) + } + + refreshed := *profileToken + refreshed.AccessToken = "profile-token-refreshed" + SetRuntimeProfile(ProfileSelector(Profile{CorpID: profileToken.CorpID, UserID: profileToken.UserID})) + if err := SaveTokenData(configDir, &refreshed); err != nil { + t.Fatalf("SaveTokenData(explicit refresh) error = %v", err) + } + SetRuntimeProfile("") + loaded, err = LoadTokenData(configDir) + if err != nil { + t.Fatalf("LoadTokenData() after explicit refresh error = %v", err) + } + if loaded.AccessToken != manualToken.AccessToken { + t.Fatalf("default token after explicit refresh = %q, want manual token", loaded.AccessToken) + } + explicit, err = LoadTokenDataForProfile( + configDir, + ProfileSelector(Profile{CorpID: profileToken.CorpID, UserID: profileToken.UserID}), + ) + if err != nil { + t.Fatalf("LoadTokenDataForProfile() after refresh error = %v", err) + } + if explicit.AccessToken != refreshed.AccessToken { + t.Fatalf("explicit refreshed token = %q, want %q", explicit.AccessToken, refreshed.AccessToken) + } + + if err := DeleteTokenDataForProfile( + configDir, + ProfileSelector(Profile{CorpID: profileToken.CorpID, UserID: profileToken.UserID}), + ); err != nil { + t.Fatalf("DeleteTokenDataForProfile() error = %v", err) + } + loaded, err = LoadTokenData(configDir) + if err != nil { + t.Fatalf("LoadTokenData() after selective logout error = %v", err) + } + if loaded.AccessToken != manualToken.AccessToken { + t.Fatalf("default token after selective logout = %q, want manual token", loaded.AccessToken) + } +} + +func TestManualTokenRemainsLoadableWithEmptyV2Tombstone(t *testing.T) { + cleanupKeychain(t) + configDir := t.TempDir() + if err := SaveProfiles(configDir, &ProfilesConfig{Version: profilesVersion}); err != nil { + t.Fatalf("SaveProfiles() error = %v", err) + } + + manualToken := &TokenData{ + AccessToken: "manual-token", + ExpiresAt: time.Now().Add(time.Hour), + } + if err := SaveTokenData(configDir, manualToken); err != nil { + t.Fatalf("SaveTokenData() error = %v", err) + } + loaded, err := LoadTokenData(configDir) + if err != nil { + t.Fatalf("LoadTokenData() error = %v", err) + } + if loaded.AccessToken != manualToken.AccessToken { + t.Fatalf("loaded token = %q, want %q", loaded.AccessToken, manualToken.AccessToken) + } +} + +func TestDefaultDeleteRemovesManualTokenWithoutDeletingProfiles(t *testing.T) { + cleanupKeychain(t) + configDir := t.TempDir() + profileToken := testToken("profile-token", "corp_same", "同一组织") + profileToken.UserID = "user_1" + if err := SaveTokenData(configDir, profileToken); err != nil { + t.Fatalf("SaveTokenData(profile) error = %v", err) + } + manualToken := &TokenData{ + AccessToken: "manual-token", + ExpiresAt: time.Now().Add(time.Hour), + } + if err := SaveTokenData(configDir, manualToken); err != nil { + t.Fatalf("SaveTokenData(manual) error = %v", err) + } + + if err := DeleteTokenData(configDir); err != nil { + t.Fatalf("DeleteTokenData() error = %v", err) + } + cfg, err := LoadProfiles(configDir) + if err != nil { + t.Fatalf("LoadProfiles() error = %v", err) + } + if len(cfg.Profiles) != 1 || cfg.CurrentProfile != "corp_same:user_1" { + t.Fatalf("profiles after default manual delete = %#v", cfg) + } + loaded, err := LoadTokenData(configDir) + if err != nil { + t.Fatalf("LoadTokenData() after manual delete error = %v", err) + } + if loaded.AccessToken != profileToken.AccessToken { + t.Fatalf("default token after manual delete = %q, want profile token", loaded.AccessToken) + } +} + func TestMultiProfileSaveLoadAndSwitch(t *testing.T) { cleanupKeychain(t) configDir := t.TempDir() @@ -236,6 +372,210 @@ func TestRuntimeProfileOverrideDoesNotMutateCurrent(t *testing.T) { } } +func TestSaveTokenDataRollsBackPartialProfilePersistence(t *testing.T) { + for _, failurePoint := range []string{"organization", "legacy", "marker"} { + t.Run(failurePoint, func(t *testing.T) { + cleanupKeychain(t) + configDir := t.TempDir() + existing := testToken("at_existing", "corp_existing", "现有组织") + if err := SaveTokenData(configDir, existing); err != nil { + t.Fatalf("SaveTokenData(existing) error = %v", err) + } + before, err := LoadProfiles(configDir) + if err != nil { + t.Fatalf("LoadProfiles() error = %v", err) + } + + incoming := testToken("at_incoming", "corp_incoming", "新组织") + failure := errors.New("persist " + failurePoint) + originalSaveOrg := tokenSaveKeychainForCorpID + originalSaveLegacy := tokenSaveKeychain + originalWriteMarker := tokenWriteMarker + switch failurePoint { + case "organization": + tokenSaveKeychainForCorpID = func(corpID string, data *TokenData) error { + if corpID == incoming.CorpID { + return failure + } + return originalSaveOrg(corpID, data) + } + case "legacy": + tokenSaveKeychain = func(data *TokenData) error { + if data != nil && data.CorpID == incoming.CorpID { + return failure + } + return originalSaveLegacy(data) + } + case "marker": + tokenWriteMarker = func(string) error { return failure } + } + + err = SaveTokenData(configDir, incoming) + tokenSaveKeychainForCorpID = originalSaveOrg + tokenSaveKeychain = originalSaveLegacy + tokenWriteMarker = originalWriteMarker + if !errors.Is(err, failure) { + t.Fatalf("SaveTokenData() error = %v, want %v", err, failure) + } + + after, err := LoadProfiles(configDir) + if err != nil { + t.Fatalf("LoadProfiles() after rollback error = %v", err) + } + if len(after.Profiles) != len(before.Profiles) || + after.CurrentProfile != before.CurrentProfile || + after.PreviousProfile != before.PreviousProfile { + t.Fatalf("profiles after rollback = %#v, want %#v", after, before) + } + if TokenDataExistsKeychainForIdentity(incoming.CorpID, incoming.UserID) || + TokenDataExistsKeychainForCorpID(incoming.CorpID) { + t.Fatal("failed login left incoming token slots behind") + } + legacy, err := LoadTokenDataKeychain() + if err != nil { + t.Fatalf("LoadTokenDataKeychain() after rollback error = %v", err) + } + if legacy.CorpID != existing.CorpID || legacy.UserID != existing.UserID { + t.Fatalf("legacy mirror after rollback = %#v, want existing identity", legacy) + } + }) + } +} + +func TestSetCurrentProfileRollsBackWhenLegacyMirrorSyncFails(t *testing.T) { + cleanupKeychain(t) + configDir := t.TempDir() + for _, userID := range []string{"user_1", "user_2"} { + data := testToken("at_"+userID, "corp_same", "同一组织") + data.UserID = userID + if err := SaveTokenData(configDir, data); err != nil { + t.Fatalf("SaveTokenData(%s) error = %v", userID, err) + } + } + + before, err := LoadProfiles(configDir) + if err != nil { + t.Fatalf("LoadProfiles() error = %v", err) + } + originalSaveLegacy := profilesSaveLegacy + failure := errors.New("save legacy mirror") + profilesSaveLegacy = func(*TokenData) error { return failure } + _, err = SetCurrentProfile(configDir, "corp_same:user_1") + profilesSaveLegacy = originalSaveLegacy + if !errors.Is(err, failure) { + t.Fatalf("SetCurrentProfile() error = %v, want %v", err, failure) + } + assertProfileSwitchRolledBack(t, configDir, before) +} + +func TestUsePreviousProfileRollsBackWhenLegacyMirrorSyncFails(t *testing.T) { + cleanupKeychain(t) + configDir := t.TempDir() + for _, userID := range []string{"user_1", "user_2"} { + data := testToken("at_"+userID, "corp_same", "同一组织") + data.UserID = userID + if err := SaveTokenData(configDir, data); err != nil { + t.Fatalf("SaveTokenData(%s) error = %v", userID, err) + } + } + + before, err := LoadProfiles(configDir) + if err != nil { + t.Fatalf("LoadProfiles() error = %v", err) + } + originalSaveLegacy := profilesSaveLegacy + failure := errors.New("save legacy mirror") + profilesSaveLegacy = func(*TokenData) error { return failure } + _, err = UsePreviousProfile(configDir) + profilesSaveLegacy = originalSaveLegacy + if !errors.Is(err, failure) { + t.Fatalf("UsePreviousProfile() error = %v, want %v", err, failure) + } + assertProfileSwitchRolledBack(t, configDir, before) +} + +func assertProfileSwitchRolledBack(t *testing.T, configDir string, before *ProfilesConfig) { + t.Helper() + after, err := LoadProfiles(configDir) + if err != nil { + t.Fatalf("LoadProfiles() after rollback error = %v", err) + } + if after.CurrentProfile != before.CurrentProfile || + after.PreviousProfile != before.PreviousProfile || + after.OrgCurrentProfiles["corp_same"] != before.OrgCurrentProfiles["corp_same"] { + t.Fatalf("profile selection after rollback = %#v, want %#v", after, before) + } + _, currentUserID, _ := ParseIdentitySelector(before.CurrentProfile) + org, err := LoadTokenDataKeychainForCorpID("corp_same") + if err != nil { + t.Fatalf("LoadTokenDataKeychainForCorpID() error = %v", err) + } + if org.UserID != currentUserID { + t.Fatalf("organization mirror user = %q, want %q", org.UserID, currentUserID) + } + legacy, err := LoadTokenDataKeychain() + if err != nil { + t.Fatalf("LoadTokenDataKeychain() error = %v", err) + } + if legacy.UserID != currentUserID { + t.Fatalf("legacy mirror user = %q, want %q", legacy.UserID, currentUserID) + } +} + +func TestFutureProfilesVersionIsNotDowngradedOrOverwritten(t *testing.T) { + cleanupKeychain(t) + configDir := t.TempDir() + future := []byte(`{ + "version": 3, + "currentProfile": "corp_future:user_1", + "orgCurrentProfiles": {"corp_future": "corp_future:user_1"}, + "futureField": {"preserve": true}, + "profiles": [{ + "name": "未来账号", + "corpId": "corp_future", + "userId": "user_1", + "futureProfileField": "preserve" + }] +} +`) + if err := os.WriteFile(ProfilesPath(configDir), future, 0o600); err != nil { + t.Fatalf("write future profiles fixture: %v", err) + } + + incoming := testToken("at_new", "corp_new", "新组织") + if err := preflightTokenPersistence(configDir); err == nil || + !strings.Contains(err.Error(), "newer than supported") { + t.Fatalf("preflightTokenPersistence() error = %v, want future version rejection", err) + } + if err := preflightTokenRefreshPersistence(configDir, incoming); err == nil || + !strings.Contains(err.Error(), "newer than supported") { + t.Fatalf("preflightTokenRefreshPersistence() error = %v, want future version rejection", err) + } + if err := SaveTokenData(configDir, incoming); err == nil || + !strings.Contains(err.Error(), "newer than supported") { + t.Fatalf("SaveTokenData() error = %v, want future version rejection", err) + } + if TokenDataExistsKeychainForIdentity(incoming.CorpID, incoming.UserID) { + t.Fatal("future version rejection left a new identity token behind") + } + if _, err := SetCurrentProfile(configDir, "corp_future:user_1"); err == nil || + !strings.Contains(err.Error(), "newer than supported") { + t.Fatalf("SetCurrentProfile() error = %v, want future version rejection", err) + } + if _, _, err := ResolveProfileDeletionScope(configDir, "corp_future:user_1"); err == nil || + !strings.Contains(err.Error(), "newer than supported") { + t.Fatalf("ResolveProfileDeletionScope() error = %v, want future version rejection", err) + } + + after, err := os.ReadFile(ProfilesPath(configDir)) + if err != nil { + t.Fatalf("read future profiles fixture: %v", err) + } + if string(after) != string(future) { + t.Fatalf("future profiles file was rewritten:\n%s", after) + } +} + func TestDeleteProfilePreservesOtherProfiles(t *testing.T) { cleanupKeychain(t) configDir := t.TempDir() @@ -778,6 +1118,167 @@ func TestDeleteOrgProfileRemovesAllAccountsInCorp(t *testing.T) { } } +func TestDeleteExactProfileRollsBackWhenOrganizationMirrorDeleteFails(t *testing.T) { + cleanupKeychain(t) + configDir := t.TempDir() + data := testToken("at_user_1", "corp_same", "同一组织") + data.UserID = "user_1" + if err := SaveTokenData(configDir, data); err != nil { + t.Fatalf("SaveTokenData() error = %v", err) + } + + originalDeleteCorp := tokenDeleteKeychainForCorpID + failure := errors.New("delete organization mirror") + tokenDeleteKeychainForCorpID = func(string) error { return failure } + err := DeleteTokenDataForProfile(configDir, "corp_same:user_1") + tokenDeleteKeychainForCorpID = originalDeleteCorp + if !errors.Is(err, failure) { + t.Fatalf("DeleteTokenDataForProfile() error = %v, want %v", err, failure) + } + + assertProfileDeletionRolledBack(t, configDir, []string{"corp_same:user_1"}, "corp_same:user_1") +} + +func TestDeleteExactProfileRollsBackWhenIdentityDeleteFails(t *testing.T) { + cleanupKeychain(t) + configDir := t.TempDir() + for _, userID := range []string{"user_1", "user_2"} { + data := testToken("at_"+userID, "corp_same", "同一组织") + data.UserID = userID + if err := SaveTokenData(configDir, data); err != nil { + t.Fatalf("SaveTokenData(%s) error = %v", userID, err) + } + } + + originalDeleteIdentity := tokenDeleteKeychainIdentity + failure := errors.New("delete identity") + tokenDeleteKeychainIdentity = func(corpID, userID string) error { + if corpID == "corp_same" && userID == "user_2" { + return failure + } + return originalDeleteIdentity(corpID, userID) + } + err := DeleteTokenDataForProfile(configDir, "corp_same:user_2") + tokenDeleteKeychainIdentity = originalDeleteIdentity + if !errors.Is(err, failure) { + t.Fatalf("DeleteTokenDataForProfile() error = %v, want %v", err, failure) + } + + assertProfileDeletionRolledBack( + t, + configDir, + []string{"corp_same:user_1", "corp_same:user_2"}, + "corp_same:user_2", + ) +} + +func TestDeleteOrganizationRollsBackPreviouslyDeletedIdentities(t *testing.T) { + cleanupKeychain(t) + configDir := t.TempDir() + for _, userID := range []string{"user_1", "user_2"} { + data := testToken("at_"+userID, "corp_same", "同一组织") + data.UserID = userID + if err := SaveTokenData(configDir, data); err != nil { + t.Fatalf("SaveTokenData(%s) error = %v", userID, err) + } + } + + originalDeleteIdentity := tokenDeleteKeychainIdentity + deleteCalls := 0 + failure := errors.New("delete second identity") + tokenDeleteKeychainIdentity = func(corpID, userID string) error { + deleteCalls++ + if deleteCalls == 2 { + return failure + } + return originalDeleteIdentity(corpID, userID) + } + err := DeleteTokenDataForProfile(configDir, "corp_same") + tokenDeleteKeychainIdentity = originalDeleteIdentity + if !errors.Is(err, failure) { + t.Fatalf("DeleteTokenDataForProfile() error = %v, want %v", err, failure) + } + + assertProfileDeletionRolledBack( + t, + configDir, + []string{"corp_same:user_1", "corp_same:user_2"}, + "corp_same:user_2", + ) +} + +func TestDeleteProfileAllowsUnreadableTargetTokenSlots(t *testing.T) { + cleanupKeychain(t) + configDir := t.TempDir() + data := testToken("at_user_1", "corp_same", "同一组织") + data.UserID = "user_1" + if err := SaveTokenData(configDir, data); err != nil { + t.Fatalf("SaveTokenData() error = %v", err) + } + for _, account := range []string{ + TokenAccountForIdentity(data.CorpID, data.UserID), + TokenAccountForCorpID(data.CorpID), + keychain.AccountToken, + } { + if err := keychain.Set(keychain.Service, account, "{unreadable"); err != nil { + t.Fatalf("write unreadable token slot %q: %v", account, err) + } + } + + if err := DeleteTokenDataForProfile(configDir, ProfileSelector(Profile{ + CorpID: data.CorpID, + UserID: data.UserID, + })); err != nil { + t.Fatalf("DeleteTokenDataForProfile() error = %v", err) + } + cfg, err := LoadProfiles(configDir) + if err != nil { + t.Fatalf("LoadProfiles() error = %v", err) + } + if len(cfg.Profiles) != 0 { + t.Fatalf("profiles after unreadable target deletion = %#v, want empty", cfg.Profiles) + } + if TokenDataExistsKeychainForIdentity(data.CorpID, data.UserID) || + TokenDataExistsKeychainForCorpID(data.CorpID) || + TokenDataExistsKeychain() { + t.Fatal("unreadable target token slots still exist after logout") + } +} + +func assertProfileDeletionRolledBack(t *testing.T, configDir string, selectors []string, current string) { + t.Helper() + cfg, err := LoadProfiles(configDir) + if err != nil { + t.Fatalf("LoadProfiles() error = %v", err) + } + if len(cfg.Profiles) != len(selectors) { + t.Fatalf("profiles after rollback = %#v, want %d profiles", cfg.Profiles, len(selectors)) + } + if cfg.CurrentProfile != current { + t.Fatalf("current profile after rollback = %q, want %q", cfg.CurrentProfile, current) + } + for _, selector := range selectors { + if _, err := LoadTokenDataForProfile(configDir, selector); err != nil { + t.Fatalf("LoadTokenDataForProfile(%q) after rollback error = %v", selector, err) + } + } + legacy, err := LoadTokenDataKeychain() + if err != nil { + t.Fatalf("LoadTokenDataKeychain() after rollback error = %v", err) + } + _, currentUserID, _ := ParseIdentitySelector(current) + if legacy.UserID != currentUserID { + t.Fatalf("legacy mirror user after rollback = %q, want %q", legacy.UserID, currentUserID) + } + org, err := LoadTokenDataKeychainForCorpID("corp_same") + if err != nil { + t.Fatalf("LoadTokenDataKeychainForCorpID() after rollback error = %v", err) + } + if org.UserID != currentUserID { + t.Fatalf("organization mirror user after rollback = %q, want %q", org.UserID, currentUserID) + } +} + func TestOAuthLoginEnrichesIdentityBeforePersistingSameCorpAccount(t *testing.T) { cleanupKeychain(t) configDir := t.TempDir() @@ -1266,6 +1767,74 @@ func TestProfilesMigrationDoesNotOverwriteUnreadableIdentityToken(t *testing.T) } } +func TestProfilesMigrationDoesNotIgnoreUnreadableOrganizationMirror(t *testing.T) { + cleanupKeychain(t) + configDir := t.TempDir() + if err := SaveProfiles(configDir, &ProfilesConfig{ + Version: 1, + Profiles: []Profile{{ + Name: "账号一", + CorpID: "corp_same", + UserID: "user_1", + }}, + }); err != nil { + t.Fatalf("SaveProfiles() error = %v", err) + } + if err := keychain.Set(keychain.Service, TokenAccountForCorpID("corp_same"), "{unreadable"); err != nil { + t.Fatalf("write unreadable organization mirror: %v", err) + } + + err := EnsureProfilesMigration(configDir) + if err == nil || !strings.Contains(err.Error(), "parse token data") { + t.Fatalf("EnsureProfilesMigration() error = %v, want unreadable mirror error", err) + } + cfg, loadErr := LoadProfiles(configDir) + if loadErr != nil { + t.Fatalf("LoadProfiles() error = %v", loadErr) + } + if cfg.Version != 1 { + t.Fatalf("profiles version after failed migration = %d, want 1", cfg.Version) + } +} + +func TestProfileSelectorRoundTripsUserIDContainingColon(t *testing.T) { + profile := Profile{CorpID: "corp", UserID: "user:with:colon"} + selector := ProfileSelector(profile) + if selector != "corp:user:with:colon" { + t.Fatalf("ProfileSelector() = %q", selector) + } + corpID, userID, exact := ParseIdentitySelector(selector) + if !exact || corpID != profile.CorpID || userID != profile.UserID { + t.Fatalf("ParseIdentitySelector(%q) = %q, %q, %v", selector, corpID, userID, exact) + } +} + +func TestMissingCurrentIdentityClearsLegacyMirror(t *testing.T) { + cleanupKeychain(t) + configDir := t.TempDir() + data := testToken("at_user_1", "corp_same", "同一组织") + data.UserID = "user_1" + if err := SaveTokenData(configDir, data); err != nil { + t.Fatalf("SaveTokenData() error = %v", err) + } + if err := DeleteTokenDataKeychainForIdentity(data.CorpID, data.UserID); err != nil { + t.Fatalf("DeleteTokenDataKeychainForIdentity() error = %v", err) + } + if err := DeleteTokenDataKeychainForCorpID(data.CorpID); err != nil { + t.Fatalf("DeleteTokenDataKeychainForCorpID() error = %v", err) + } + + if err := SyncLegacyTokenMirror(configDir); err != nil { + t.Fatalf("SyncLegacyTokenMirror() error = %v", err) + } + if TokenDataExistsKeychain() { + t.Fatal("legacy mirror still exists after current identity token was confirmed missing") + } + if _, err := os.Stat(filepath.Join(configDir, tokenJSONFile)); !os.IsNotExist(err) { + t.Fatalf("token marker stat error = %v, want not exist", err) + } +} + func TestOAuthLoginRequiresIdentityForFirstProfile(t *testing.T) { cleanupKeychain(t) configDir := t.TempDir() diff --git a/internal/keychain/keychain_windows.go b/internal/keychain/keychain_windows.go index 4b73099c..69e66af6 100644 --- a/internal/keychain/keychain_windows.go +++ b/internal/keychain/keychain_windows.go @@ -17,6 +17,7 @@ package keychain import ( "encoding/base64" + "errors" "fmt" "os" "path/filepath" @@ -28,6 +29,14 @@ import ( "golang.org/x/sys/windows/registry" ) +var registryDeleteValue = func(key registry.Key, name string) error { + return key.DeleteValue(name) +} + +var registryOpenDeleteKey = func(path string, access uint32) (registry.Key, error) { + return registry.OpenKey(registry.CURRENT_USER, path, access) +} + // --------------------------------------------------------------------------- // Windows backend: DPAPI + HKCU registry // --------------------------------------------------------------------------- @@ -196,20 +205,35 @@ func registrySet(service, account string, protected []byte) error { func registryRemove(service, account string) error { keyPath := registryPathForService(service) - k, err := registry.OpenKey(registry.CURRENT_USER, keyPath, registry.SET_VALUE) + k, err := registryOpenDeleteKey(keyPath, registry.SET_VALUE) if err != nil { - return nil + if errors.Is(err, windows.ERROR_FILE_NOT_FOUND) { + return nil + } + return fmt.Errorf("registry open for delete failed: %w", err) } defer k.Close() - _ = k.DeleteValue(valueNameForAccount(account)) + return deleteRegistryValue(k, valueNameForAccount(account)) +} + +func deleteRegistryValue(key registry.Key, name string) error { + if err := registryDeleteValue(key, name); err != nil { + if errors.Is(err, windows.ERROR_FILE_NOT_FOUND) { + return nil + } + return fmt.Errorf("registry delete failed: %w", err) + } return nil } func registryRemoveAuthTokenEntries(service string) error { keyPath := registryPathForService(service) - k, err := registry.OpenKey(registry.CURRENT_USER, keyPath, registry.QUERY_VALUE|registry.SET_VALUE) + k, err := registryOpenDeleteKey(keyPath, registry.QUERY_VALUE|registry.SET_VALUE) if err != nil { - return nil + if errors.Is(err, windows.ERROR_FILE_NOT_FOUND) { + return nil + } + return fmt.Errorf("registry open for auth token cleanup failed: %w", err) } defer k.Close() diff --git a/internal/keychain/keychain_windows_test.go b/internal/keychain/keychain_windows_test.go new file mode 100644 index 00000000..9d9f4b80 --- /dev/null +++ b/internal/keychain/keychain_windows_test.go @@ -0,0 +1,48 @@ +// Copyright 2026 Alibaba Group +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +//go:build windows + +package keychain + +import ( + "errors" + "testing" + + "golang.org/x/sys/windows" + "golang.org/x/sys/windows/registry" +) + +func TestDeleteRegistryValuePropagatesFailure(t *testing.T) { + originalDelete := registryDeleteValue + failure := errors.New("delete failed") + registryDeleteValue = func(registry.Key, string) error { return failure } + t.Cleanup(func() { registryDeleteValue = originalDelete }) + + if err := deleteRegistryValue(0, "auth-token"); !errors.Is(err, failure) { + t.Fatalf("deleteRegistryValue() error = %v, want %v", err, failure) + } +} + +func TestRegistryRemoveAuthTokenEntriesPropagatesOpenFailure(t *testing.T) { + originalOpen := registryOpenDeleteKey + failure := windows.ERROR_ACCESS_DENIED + registryOpenDeleteKey = func(string, uint32) (registry.Key, error) { + return 0, failure + } + t.Cleanup(func() { registryOpenDeleteKey = originalOpen }) + + if err := registryRemoveAuthTokenEntries("service"); !errors.Is(err, failure) { + t.Fatalf("registryRemoveAuthTokenEntries() error = %v, want %v", err, failure) + } +}