From 93854178e3844c073d63a5c8f05eaf32adc9c040 Mon Sep 17 00:00:00 2001 From: audanye-sudo Date: Fri, 17 Jul 2026 00:56:43 +0800 Subject: [PATCH] feat(auth): support exact multi-account profile identities Store DingTalk credentials in corpId:userId identity slots while retaining organization and legacy mirrors for forward compatibility. Resolve organization, account, friendly-name, current, previous, and deletion selectors without silently choosing among ambiguous accounts. Make identity tokens the source of truth, serialize migration and refresh reads, reject unsafe mirror recovery, and sweep orphan token entries during reset. Persist token source and client ID for exact remote revocation, require user identity before first-login persistence, and add cross-platform regression coverage for migration, deletion, refresh, and keychain failures. --- internal/auth/auth_extra_test.go | 37 + internal/auth/cli_auth_check_test.go | 4 + internal/auth/coverage_edges_test.go | 293 ++++- internal/auth/device_flow.go | 29 +- internal/auth/keychain_store.go | 71 ++ internal/auth/oauth_helpers.go | 4 + internal/auth/oauth_provider.go | 53 +- internal/auth/oauth_provider_coverage_test.go | 5 + internal/auth/portable_store_test.go | 61 +- internal/auth/profiles.go | 1027 ++++++++++++++--- internal/auth/token.go | 281 ++++- internal/auth/token_preflight_test.go | 64 +- internal/auth/token_test.go | 897 +++++++++++++- internal/keychain/auth_entries_file.go | 14 + internal/keychain/auth_entries_windows.go | 4 + internal/keychain/keychain.go | 7 + internal/keychain/keychain_test.go | 30 + internal/keychain/keychain_windows.go | 28 + 18 files changed, 2590 insertions(+), 319 deletions(-) diff --git a/internal/auth/auth_extra_test.go b/internal/auth/auth_extra_test.go index 5d702d32..d64d9ee8 100644 --- a/internal/auth/auth_extra_test.go +++ b/internal/auth/auth_extra_test.go @@ -306,6 +306,43 @@ func TestRevokeTokenRemote(t *testing.T) { // Can't easily test since LogoutURL is a const. Just test that it doesn't panic with real URL. } +func TestRevokeTokenRemoteForDataUsesExactTokenMetadata(t *testing.T) { + configDir := t.TempDir() + var got struct { + ClientID string `json:"clientId"` + AccessToken string `json:"accessToken"` + } + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.URL.Path != MCPRevokeTokenPath { + t.Errorf("revoke path = %q, want %q", r.URL.Path, MCPRevokeTokenPath) + } + if err := json.NewDecoder(r.Body).Decode(&got); err != nil { + t.Errorf("decode revoke body: %v", err) + } + w.WriteHeader(http.StatusOK) + })) + defer srv.Close() + + t.Setenv("DWS_CONFIG_DIR", configDir) + if err := os.WriteFile(filepath.Join(configDir, "mcp_url"), []byte(srv.URL), 0o600); err != nil { + t.Fatalf("write mcp_url: %v", err) + } + SetClientID("wrong-global-client") + t.Cleanup(func() { SetClientID("") }) + + data := &TokenData{ + AccessToken: "exact-account-token", + ClientID: "exact-account-client", + Source: "mcp", + } + if err := RevokeTokenRemoteForData(t.Context(), data); err != nil { + t.Fatalf("RevokeTokenRemoteForData() error = %v", err) + } + if got.ClientID != data.ClientID || got.AccessToken != data.AccessToken { + t.Fatalf("revoke body = %#v, want exact token metadata", got) + } +} + // ─── oauth_helpers.go ────────────────────────────────────────────────── type tokenResponse struct { diff --git a/internal/auth/cli_auth_check_test.go b/internal/auth/cli_auth_check_test.go index 420af07e..6ed68644 100644 --- a/internal/auth/cli_auth_check_test.go +++ b/internal/auth/cli_auth_check_test.go @@ -495,6 +495,10 @@ func TestDeviceFlow_LoginOnce_CLIAuthEnabled_Success(t *testing.T) { logger: newDeviceFlowTestLogger(), Output: io.Discard, httpClient: srv.Client(), + IdentityEnricher: func(_ context.Context, data *TokenData) error { + data.UserID = "user123" + return nil + }, } token, err := provider.loginOnce(context.Background(), 1) diff --git a/internal/auth/coverage_edges_test.go b/internal/auth/coverage_edges_test.go index 1fa26b34..8154cf52 100644 --- a/internal/auth/coverage_edges_test.go +++ b/internal/auth/coverage_edges_test.go @@ -396,7 +396,8 @@ func TestCrossPlatformCoverageDeviceFlowLoginRetry(t *testing.T) { return } _ = json.NewEncoder(w).Encode(map[string]any{ - "accessToken": "device-access", "refreshToken": "refresh", "expiresIn": 7200, "corpId": "corp-device", + "accessToken": "device-access", "refreshToken": "refresh", "expiresIn": 7200, + "corpId": "corp-device", "userId": "user-device", }) case strings.HasSuffix(r.URL.Path, CLIAuthEnabledPath): _ = json.NewEncoder(w).Encode(CLIAuthStatus{Success: true, Result: &CLIAuthResult{CLIAuthEnabled: true}}) @@ -493,7 +494,7 @@ func TestCrossPlatformCoverageProfilesLifecycleEdges(t *testing.T) { } normalizeProfilesConfig(cfg) if len(cfg.Profiles) != 1 || cfg.Profiles[0].Name != "Acme" || - cfg.PrimaryProfile != "corp-a" || cfg.CurrentProfile != "corp-a" || cfg.PreviousProfile != "" { + cfg.PrimaryProfile != "" || cfg.CurrentProfile != "" || cfg.PreviousProfile != "" { t.Fatalf("normalized profiles = %#v", cfg) } @@ -515,6 +516,16 @@ func TestCrossPlatformCoverageProfilesLifecycleEdges(t *testing.T) { if err := UpsertProfileFromTokenWithCurrent(dir, second, true); err != nil { t.Fatal(err) } + if err := SaveTokenDataKeychainForIdentity(first.CorpID, first.UserID, first); err != nil { + t.Fatal(err) + } + if err := SaveTokenDataKeychainForCorpID(second.CorpID, second); err != nil { + t.Fatal(err) + } + t.Cleanup(func() { + _ = DeleteTokenDataKeychainForIdentity(first.CorpID, first.UserID) + _ = DeleteTokenDataKeychainForCorpID(second.CorpID) + }) profilesForAmbiguity, err := LoadProfiles(dir) if err != nil { t.Fatal(err) @@ -575,7 +586,7 @@ func TestCrossPlatformCoverageProfilesLifecycleEdges(t *testing.T) { {CorpID: "two", Name: "Acme-12345678"}, {CorpID: "three", Name: "Acme-12345678-2"}, }} - if got := chooseProfileName(nameCfg, &TokenData{CorpID: "corp-12345678", CorpName: "Acme"}); got != "Acme-12345678-3" { + if got := chooseProfileName(nameCfg, &TokenData{CorpID: "corp-12345678", CorpName: "Acme"}); got != "Acme-2" { t.Fatalf("collision profile name = %q", got) } if chooseProfileName(&ProfilesConfig{}, &TokenData{}) != "profile" { @@ -585,12 +596,7 @@ func TestCrossPlatformCoverageProfilesLifecycleEdges(t *testing.T) { !shouldRefreshProfileName(&Profile{}, first) { t.Fatal("profile refresh-name decisions failed") } - if findProfile(nil, "x") != nil || findProfile(cfg, "") != nil || - profileIndexByCorpID(nil, "x") != -1 || firstProfileCorpID(nil) != "" { - t.Fatal("nil profile helper failed") - } - if shortCorpID("short") != "short" || shortCorpID("corp-12345678") != "12345678" || - timeOrRFC3339(time.Time{}) != "" { + if shortCorpID("short") != "short" || shortCorpID("corp-12345678") != "12345678" { t.Fatal("profile formatting helpers failed") } } @@ -1454,18 +1460,26 @@ func TestCrossPlatformCoverageTokenStorageAndRevocationCoverageEdges(t *testing. oldRemove := tokenRemove oldGlob := tokenGlob oldSaveCorp := tokenSaveKeychainForCorpID + oldSaveIdentity := tokenSaveKeychainForIdentity oldSaveLegacy := tokenSaveKeychain oldLoadCorp := tokenLoadKeychainForCorpID + oldLoadIdentity := tokenLoadKeychainIdentity oldLoadLegacy := tokenLoadKeychain oldExists := tokenKeychainExists oldDeleteCorp := tokenDeleteKeychainForCorpID + oldDeleteIdentity := tokenDeleteKeychainIdentity oldDeleteLegacy := tokenDeleteKeychain + oldRemoveAuthEntries := tokenRemoveAuthTokenEntries oldLoadSecure := tokenLoadSecure oldDeleteSecure := tokenDeleteSecure + oldEnsureProfiles := profilesEnsureMigration oldResolve := tokenResolveProfile + oldResolveDeletion := tokenResolveDeletion + oldResolveSelection := tokenResolveSelection oldUpsert := tokenUpsertProfile oldRemoveProfile := tokenRemoveProfile oldSync := tokenSyncLegacyMirror + oldSyncOrganization := tokenSyncOrganizationMirror oldLoadProfiles := tokenLoadProfiles oldWriteMarker := tokenWriteMarker oldDeleteMarker := tokenDeleteMarker @@ -1474,6 +1488,7 @@ func TestCrossPlatformCoverageTokenStorageAndRevocationCoverageEdges(t *testing. oldDefaultDir := tokenDefaultConfigDir oldLoadData := tokenLoadData oldRevokeURL := tokenRevokeURL + oldMCPBaseURL := tokenMCPBaseURL oldLogoutURL := tokenLogoutURL oldLogoutContinue := tokenLogoutContinueURL oldLogoutClient := tokenLogoutHTTPClient @@ -1488,18 +1503,26 @@ func TestCrossPlatformCoverageTokenStorageAndRevocationCoverageEdges(t *testing. tokenRemove = oldRemove tokenGlob = oldGlob tokenSaveKeychainForCorpID = oldSaveCorp + tokenSaveKeychainForIdentity = oldSaveIdentity tokenSaveKeychain = oldSaveLegacy tokenLoadKeychainForCorpID = oldLoadCorp + tokenLoadKeychainIdentity = oldLoadIdentity tokenLoadKeychain = oldLoadLegacy tokenKeychainExists = oldExists tokenDeleteKeychainForCorpID = oldDeleteCorp + tokenDeleteKeychainIdentity = oldDeleteIdentity tokenDeleteKeychain = oldDeleteLegacy + tokenRemoveAuthTokenEntries = oldRemoveAuthEntries tokenLoadSecure = oldLoadSecure tokenDeleteSecure = oldDeleteSecure + profilesEnsureMigration = oldEnsureProfiles tokenResolveProfile = oldResolve + tokenResolveDeletion = oldResolveDeletion + tokenResolveSelection = oldResolveSelection tokenUpsertProfile = oldUpsert tokenRemoveProfile = oldRemoveProfile tokenSyncLegacyMirror = oldSync + tokenSyncOrganizationMirror = oldSyncOrganization tokenLoadProfiles = oldLoadProfiles tokenWriteMarker = oldWriteMarker tokenDeleteMarker = oldDeleteMarker @@ -1508,6 +1531,7 @@ func TestCrossPlatformCoverageTokenStorageAndRevocationCoverageEdges(t *testing. tokenDefaultConfigDir = oldDefaultDir tokenLoadData = oldLoadData tokenRevokeURL = oldRevokeURL + tokenMCPBaseURL = oldMCPBaseURL tokenLogoutURL = oldLogoutURL tokenLogoutContinueURL = oldLogoutContinue tokenLogoutHTTPClient = oldLogoutClient @@ -1653,16 +1677,36 @@ func TestCrossPlatformCoverageTokenStorageAndRevocationCoverageEdges(t *testing. t.Run("delete branches", func(t *testing.T) { dir := t.TempDir() fail := errors.New("fail") - selected := &Profile{CorpID: "corp"} - tokenResolveProfile = func(string, string) (*Profile, error) { return nil, fail } + selected := &Profile{CorpID: "corp", UserID: "user"} + profilesEnsureMigration = func(string) error { return fail } + if err := deleteTokenDataForProfileLocked(dir, ""); !errors.Is(err, fail) { + t.Fatalf("delete migration error = %v", err) + } + profilesEnsureMigration = func(string) error { return nil } + tokenLoadProfiles = func(string) (*ProfilesConfig, error) { return nil, fail } + if err := deleteTokenDataForProfileLocked(dir, ""); !errors.Is(err, fail) { + t.Fatalf("delete profiles load error = %v", err) + } + cfg := &ProfilesConfig{ + CurrentProfile: ProfileSelector(*selected), + OrgCurrentProfiles: map[string]string{"corp": ProfileSelector(*selected)}, + Profiles: []Profile{*selected}, + } + tokenLoadProfiles = func(string) (*ProfilesConfig, error) { return cfg, nil } + tokenResolveDeletion = func(*ProfilesConfig, string) (*Profile, bool, error) { return nil, false, fail } if err := deleteTokenDataForProfileLocked(dir, ""); !errors.Is(err, fail) { t.Fatalf("delete resolution error = %v", err) } - tokenResolveProfile = func(string, string) (*Profile, error) { return selected, nil } - tokenDeleteKeychainForCorpID = func(string) error { return fail } + tokenResolveDeletion = func(*ProfilesConfig, string) (*Profile, bool, error) { return selected, true, 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 } + tokenDeleteKeychainForCorpID = func(string) error { return fail } if err := deleteTokenDataForProfileLocked(dir, ""); !errors.Is(err, fail) { t.Fatalf("delete corp keychain error = %v", err) } @@ -1672,6 +1716,7 @@ func TestCrossPlatformCoverageTokenStorageAndRevocationCoverageEdges(t *testing. t.Fatalf("remove profile error = %v", err) } tokenRemoveProfile = func(string, string) (*Profile, error) { return selected, nil } + cfg.OrgCurrentProfiles = nil tokenSyncLegacyMirror = func(string) error { return fail } if err := deleteTokenDataForProfileLocked(dir, ""); !errors.Is(err, fail) { t.Fatalf("sync mirror error = %v", err) @@ -1682,7 +1727,8 @@ func TestCrossPlatformCoverageTokenStorageAndRevocationCoverageEdges(t *testing. t.Fatalf("secure cleanup error = %v", err) } - tokenResolveProfile = func(string, string) (*Profile, error) { return nil, nil } + cfg.CurrentProfile = "" + cfg.Profiles = nil tokenDeleteKeychain = func() error { return fail } tokenDeleteMarker = func(string) error { return nil } if err := deleteTokenDataForProfileLocked(dir, ""); !errors.Is(err, fail) { @@ -1710,22 +1756,17 @@ func TestCrossPlatformCoverageTokenStorageAndRevocationCoverageEdges(t *testing. t.Run("delete all", func(t *testing.T) { fail := errors.New("fail") base := func() string { - tokenLoadProfiles = func(string) (*ProfilesConfig, error) { return &ProfilesConfig{}, nil } - tokenDeleteKeychainForCorpID = func(string) error { return nil } + tokenRemoveAuthTokenEntries = func(string) error { return nil } tokenRemove = func(string) error { return os.ErrNotExist } tokenGlob = func(string) ([]string, error) { return nil, nil } - tokenDeleteKeychain = func() error { return nil } tokenDeleteSecure = func(string) error { return nil } tokenDeleteMarker = func(string) error { return nil } return t.TempDir() } dir := base() - tokenLoadProfiles = func(string) (*ProfilesConfig, error) { - return &ProfilesConfig{Profiles: []Profile{{CorpID: "corp"}}}, nil - } - tokenDeleteKeychainForCorpID = func(string) error { return fail } + tokenRemoveAuthTokenEntries = func(string) error { return fail } if err := DeleteAllTokenData(dir); !errors.Is(err, fail) { - t.Fatalf("delete-all corp error = %v", err) + t.Fatalf("delete-all auth namespace error = %v", err) } dir = base() tokenRemove = func(string) error { return fail } @@ -1745,11 +1786,6 @@ func TestCrossPlatformCoverageTokenStorageAndRevocationCoverageEdges(t *testing. t.Fatalf("delete-all quarantine error = %v", err) } dir = base() - tokenDeleteKeychain = func() error { return fail } - if err := DeleteAllTokenData(dir); !errors.Is(err, fail) { - t.Fatalf("delete-all legacy keychain error = %v", err) - } - dir = base() tokenDeleteSecure = func(string) error { return fail } if err := DeleteAllTokenData(dir); !errors.Is(err, fail) { t.Fatalf("delete-all secure error = %v", err) @@ -1765,6 +1801,9 @@ func TestCrossPlatformCoverageTokenStorageAndRevocationCoverageEdges(t *testing. resetClientIDFromMCP() SetClientID("client") fail := errors.New("fail") + tokenLoadData = func(string) (*TokenData, error) { + return &TokenData{AccessToken: "access", ClientID: "client", Source: "direct"}, nil + } tokenParseURL = func(string) (*url.URL, error) { return nil, fail } if err := RevokeTokenRemote(context.Background()); !errors.Is(err, fail) { t.Fatalf("logout parse error = %v", err) @@ -1797,16 +1836,14 @@ func TestCrossPlatformCoverageTokenStorageAndRevocationCoverageEdges(t *testing. } SetClientIDFromMCP("mcp-client") - tokenRevokeURL = func() string { return "" } - if err := RevokeTokenRemote(context.Background()); err != nil { - t.Fatal(err) - } - tokenRevokeURL = func() string { return "https://revoke.test" } tokenLoadData = func(string) (*TokenData, error) { return nil, fail } if err := RevokeTokenRemote(context.Background()); err != nil { t.Fatal("missing token revoke should be a no-op") } - tokenLoadData = func(string) (*TokenData, error) { return &TokenData{AccessToken: "access"}, nil } + tokenMCPBaseURL = func() string { return "https://revoke.test" } + tokenLoadData = func(string) (*TokenData, error) { + return &TokenData{AccessToken: "access", ClientID: "mcp-client", Source: "mcp"}, nil + } tokenJSONMarshal = func(any) ([]byte, error) { return nil, fail } if err := RevokeTokenRemote(context.Background()); !errors.Is(err, fail) { t.Fatalf("revoke marshal error = %v", err) @@ -2207,8 +2244,8 @@ func TestCrossPlatformCoverageProfilesCoverageEdges(t *testing.T) { } primaryOnly := &ProfilesConfig{PrimaryProfile: "primary", Profiles: []Profile{{Name: "primary", CorpID: "primary"}}} profilesLoad = func(string) (*ProfilesConfig, error) { return primaryOnly, nil } - if got, err := ResolveProfile(dir, ""); err != nil || got == nil || got.CorpID != "primary" { - t.Fatalf("primary profile resolution = %#v %v", got, err) + if got, err := ResolveProfile(dir, ""); err != nil || got != nil { + t.Fatalf("primary-only profile resolution = %#v %v", got, err) } profilesLoad = func(string) (*ProfilesConfig, error) { return &ProfilesConfig{}, nil } if got, err := ResolveProfile(dir, ""); err != nil || got != nil { @@ -2222,8 +2259,8 @@ func TestCrossPlatformCoverageProfilesCoverageEdges(t *testing.T) { } cfgNoCurrent := &ProfilesConfig{PrimaryProfile: "primary", Profiles: []Profile{{CorpID: "primary"}}} profilesLoad = func(string) (*ProfilesConfig, error) { return cfgNoCurrent, nil } - if got, err := resolveProfileForLoad(dir, ""); err != nil || got == nil || got.CorpID != "primary" { - t.Fatalf("profile-for-load primary fallback = %#v %v", got, err) + if got, err := resolveProfileForLoad(dir, ""); err != nil || got != nil { + t.Fatalf("profile-for-load primary-only result = %#v %v", got, err) } profilesLoad = func(string) (*ProfilesConfig, error) { return cfg, nil } @@ -2242,6 +2279,10 @@ func TestCrossPlatformCoverageProfilesCoverageEdges(t *testing.T) { Profiles: []Profile{{Name: "primary", CorpID: "primary"}, {Name: "current", CorpID: "current"}, {Name: "previous", CorpID: "previous"}}, } profilesSave = func(string, *ProfilesConfig) error { return nil } + profilesLoadCorp = func(corpID string) (*TokenData, error) { + return &TokenData{CorpID: corpID, AccessToken: "x"}, nil + } + profilesSaveCorp = func(string, *TokenData) error { return nil } profilesSyncLegacyMirror = func(string) error { return fail } if _, err := setCurrentProfileLocked(dir, "current"); !errors.Is(err, fail) { t.Fatalf("set-current mirror error = %v", err) @@ -2269,9 +2310,13 @@ func TestCrossPlatformCoverageProfilesCoverageEdges(t *testing.T) { deletedLegacy, deletedMarker := false, false profilesDeleteLegacy = func() error { deletedLegacy = true; return fail } profilesDeleteMarker = func(string) error { deletedMarker = true; return fail } - if err := syncLegacyTokenMirrorLocked(dir); err != nil || !deletedLegacy || !deletedMarker { + if err := syncLegacyTokenMirrorLocked(dir); !errors.Is(err, fail) || !deletedLegacy || deletedMarker { t.Fatalf("empty mirror cleanup = %v %v %v", err, deletedLegacy, deletedMarker) } + profilesDeleteLegacy = func() error { return nil } + if err := syncLegacyTokenMirrorLocked(dir); !errors.Is(err, fail) || !deletedMarker { + t.Fatalf("empty mirror marker cleanup = %v %v", err, deletedMarker) + } normalizeProfilesConfig(nil) } @@ -2911,3 +2956,175 @@ func TestCrossPlatformCoverageDeviceFlowHighLevelCoverageEdges(t *testing.T) { } dfPrintBox(io.Discard, []string{"short"}) } + +func TestCrossPlatformCoverageMultiAccountSelectorAndIdentityLoadEdges(t *testing.T) { + cfg := &ProfilesConfig{ + Profiles: []Profile{ + {Name: "local-a", CorpID: "corp-a", CorpName: "Shared", UserID: "u1", UserName: "Alice"}, + {Name: "duplicate", CorpID: "corp-a", CorpName: "Shared", UserID: "u2", UserName: "Alice"}, + {Name: "duplicate", CorpID: "corp-b", CorpName: "Shared", UserID: "u3", UserName: "Bob"}, + {Name: "solo", CorpID: "corp-c", CorpName: "Solo", UserID: "u4", UserName: "Carol"}, + }, + } + + for _, tc := range []struct { + selector string + wantErr bool + }{ + {"", true}, + {"missing:u1", true}, + {"Shared:u1", true}, + {"corp-a:u1", false}, + {"corp-a:Alice", true}, + {"corp-a:missing", true}, + {"corp-a", true}, + {"corp-c", false}, + {"Solo", false}, + {"local-a", false}, + {"duplicate", true}, + {"missing", true}, + } { + _, _, err := resolveProfileSelection("", cfg, tc.selector) + if (err != nil) != tc.wantErr { + t.Fatalf("resolveProfileSelection(%q) error = %v", tc.selector, err) + } + } + if _, _, err := resolveProfileSelection("", nil, "x"); err == nil { + t.Fatal("nil profile config selection succeeded") + } + cfg.OrgCurrentProfiles = map[string]string{"corp-a": "corp-a:u2"} + if got, _, err := resolveProfileSelection("", cfg, "corp-a"); err != nil || got.UserID != "u2" { + t.Fatalf("organization current selection = %#v %v", got, err) + } + + for _, tc := range []struct { + selector string + wantErr bool + exact bool + }{ + {"", true, false}, + {"corp-a:u1", false, true}, + {"corp-a", false, false}, + {"Solo", false, false}, + {"local-a", false, true}, + {"duplicate", true, true}, + {"missing", true, false}, + } { + _, exact, err := resolveProfileDeletionSelection(cfg, tc.selector) + if (err != nil) != tc.wantErr || (!tc.wantErr && exact != tc.exact) { + t.Fatalf("resolveProfileDeletionSelection(%q) = exact %v, %v", tc.selector, exact, err) + } + } + if _, _, err := resolveProfileDeletionSelection(nil, "x"); err == nil { + t.Fatal("nil deletion config succeeded") + } + + for _, selector := range []string{"", "corp-a", "missing", "Solo", "Shared"} { + _, _ = resolveOrganizationCorpID(cfg, selector) + } + if got, err := resolveOrganizationCorpID(nil, "x"); err != nil || got != "" { + t.Fatalf("nil organization resolution = %q %v", got, err) + } + if _, _, err := resolveOrganizationDefault(cfg, "missing", "missing", nil); err == nil { + t.Fatal("empty organization default succeeded") + } + if got, _, err := resolveOrganizationDefault(cfg, "corp-c", "Solo", profilesForCorpID(cfg, "corp-c")); err != nil || got.UserID != "u4" { + t.Fatalf("single account organization default = %#v %v", got, err) + } + if _, _, err := resolveOrganizationDefault(&ProfilesConfig{}, "corp-a", "corp-a", profilesForCorpID(cfg, "corp-a")); err == nil { + t.Fatal("ambiguous organization default succeeded") + } + if got := profileSelectorCandidates([]*Profile{nil, &cfg.Profiles[1], &cfg.Profiles[0]}); len(got) != 2 || got[0] != "corp-a:u1" { + t.Fatalf("selector candidates = %#v", got) + } + for selector, want := range map[string]bool{ + "": false, + "corp-a:u1": true, + "Shared:Bob": true, + "corp-c": true, + "local-a": true, + "definitely-missing": false, + } { + if got := profileSelectorReferenceExists(cfg, selector); got != want { + t.Fatalf("profileSelectorReferenceExists(%q) = %v", selector, got) + } + } + if profileSelectorReferenceExists(nil, "corp-a") { + t.Fatal("nil profile selector reference exists") + } + + oldAcquire := profilesAcquireDualLock + oldEnsure := profilesEnsureMigration + oldLoad := profilesLoad + t.Cleanup(func() { + profilesAcquireDualLock = oldAcquire + profilesEnsureMigration = oldEnsure + profilesLoad = oldLoad + }) + profilesAcquireDualLock = func(context.Context, string) (*DualLock, error) { return &DualLock{}, nil } + profilesEnsureMigration = func(string) error { return nil } + profilesLoad = func(string) (*ProfilesConfig, error) { return cfg, nil } + if got, exact, err := ResolveProfileDeletionScope("cfg", "corp-a:u1"); err != nil || !exact || got.UserID != "u1" { + t.Fatalf("deletion scope = %#v %v %v", got, exact, err) + } + profilesEnsureMigration = func(string) error { return errors.New("migration") } + if _, _, err := ResolveProfileDeletionScope("cfg", "corp-a"); err == nil { + t.Fatal("deletion scope migration failure succeeded") + } + profilesEnsureMigration = func(string) error { return nil } + profilesLoad = func(string) (*ProfilesConfig, error) { return nil, errors.New("load") } + if _, _, err := ResolveProfileDeletionScope("cfg", "corp-a"); err == nil { + t.Fatal("deletion scope load failure succeeded") + } + + oldLoadIdentity := profilesLoadIdentity + oldLoadCorp := profilesLoadCorp + oldSaveIdentity := profilesSaveIdentity + t.Cleanup(func() { + profilesLoadIdentity = oldLoadIdentity + profilesLoadCorp = oldLoadCorp + profilesSaveIdentity = oldSaveIdentity + }) + profile := Profile{CorpID: "corp-a", UserID: "u1"} + profilesLoadIdentity = func(string, string) (*TokenData, error) { return &TokenData{AccessToken: "identity"}, nil } + if got, err := loadTokenForProfileIdentity(profile); err != nil || got.AccessToken != "identity" { + t.Fatalf("identity token load = %#v %v", got, err) + } + fail := errors.New("fail") + profilesLoadIdentity = func(string, string) (*TokenData, error) { return nil, fail } + if _, err := loadTokenForProfileIdentity(profile); !errors.Is(err, fail) { + t.Fatalf("identity load failure = %v", err) + } + profilesLoadIdentity = func(string, string) (*TokenData, error) { return nil, ErrTokenDataNotFound } + profilesLoadCorp = func(string) (*TokenData, error) { return nil, ErrTokenDataNotFound } + if _, err := loadTokenForProfileIdentity(profile); !errors.Is(err, ErrTokenDataNotFound) { + t.Fatalf("missing organization mirror = %v", err) + } + profilesLoadCorp = func(string) (*TokenData, error) { return nil, fail } + if _, err := loadTokenForProfileIdentity(profile); !errors.Is(err, fail) { + t.Fatalf("organization mirror failure = %v", err) + } + for _, data := range []*TokenData{ + {CorpID: "corp-a"}, + {CorpID: "corp-a", UserID: "other"}, + } { + profilesLoadCorp = func(string) (*TokenData, error) { return data, nil } + if _, err := loadTokenForProfileIdentity(profile); err == nil { + t.Fatalf("invalid organization mirror %#v succeeded", data) + } + } + profilesLoadCorp = func(string) (*TokenData, error) { + return &TokenData{CorpID: "corp-a", UserID: "u1", AccessToken: "mirror"}, nil + } + profilesSaveIdentity = func(string, string, *TokenData) error { return fail } + if _, err := loadTokenForProfileIdentity(profile); !errors.Is(err, fail) { + t.Fatalf("identity repair save failure = %v", err) + } + profilesSaveIdentity = func(string, string, *TokenData) error { return nil } + if got, err := loadTokenForProfileIdentity(profile); err != nil || got.AccessToken != "mirror" { + t.Fatalf("identity repair = %#v %v", got, err) + } + if _, err := loadTokenForProfileIdentity(Profile{CorpID: "corp-a"}); err != nil { + t.Fatalf("organization-only token load = %v", err) + } +} diff --git a/internal/auth/device_flow.go b/internal/auth/device_flow.go index 02e96332..644a9aba 100644 --- a/internal/auth/device_flow.go +++ b/internal/auth/device_flow.go @@ -75,15 +75,16 @@ var ( ) type DeviceFlowProvider struct { - configDir string - clientID string - scope string - baseURL string - terminalBaseURL string - logger *slog.Logger - Output io.Writer - httpClient *http.Client - NoBrowser bool + configDir string + clientID string + scope string + baseURL string + terminalBaseURL string + logger *slog.Logger + Output io.Writer + httpClient *http.Client + NoBrowser bool + IdentityEnricher func(context.Context, *TokenData) error } func NewDeviceFlowProvider(configDir string, logger *slog.Logger) *DeviceFlowProvider { @@ -267,9 +268,10 @@ func (p *DeviceFlowProvider) loginOnce(ctx context.Context, attempt int) (*Token dfPrintStep(p.output(), 3, i18n.T("使用授权码换取 Access Token..."), 0) oauthProvider := &OAuthProvider{ - configDir: p.configDir, - clientID: p.clientID, - logger: p.logger, + configDir: p.configDir, + clientID: p.clientID, + logger: p.logger, + IdentityEnricher: p.IdentityEnricher, } tokenData, err := deviceExchangeCode(oauthProvider, ctx, tokenResult.AuthCode) if err != nil { @@ -352,6 +354,9 @@ func (p *DeviceFlowProvider) loginOnce(ctx context.Context, attempt int) (*Token // Save token data with associated client ID for refresh tokenData.ClientID = p.clientID + if err := oauthProvider.prepareLoginToken(ctx, tokenData); err != nil { + return nil, fmt.Errorf("%s: %w", i18n.T("保存 token 失败"), err) + } if err := deviceSaveToken(p.configDir, tokenData); err != nil { return nil, fmt.Errorf("%s: %w", i18n.T("保存 token 失败"), err) } diff --git a/internal/auth/keychain_store.go b/internal/auth/keychain_store.go index b03611eb..839780e3 100644 --- a/internal/auth/keychain_store.go +++ b/internal/auth/keychain_store.go @@ -14,6 +14,7 @@ package auth import ( + "crypto/sha256" "encoding/json" "errors" "fmt" @@ -52,6 +53,14 @@ func TokenAccountForCorpID(corpID string) string { return keychain.AccountToken + ":" + strings.TrimSpace(corpID) } +// TokenAccountForIdentity returns the stable keychain account used for one +// DingTalk identity. The hash avoids collisions caused by delimiter escaping +// or keychain/file-name restrictions. +func TokenAccountForIdentity(corpID, userID string) string { + identity := strings.TrimSpace(corpID) + "\x00" + strings.TrimSpace(userID) + return fmt.Sprintf("%s:id:%x", keychain.AccountToken, sha256.Sum256([]byte(identity))) +} + // SaveTokenDataKeychainForCorpID saves TokenData to a corp-scoped keychain slot. func SaveTokenDataKeychainForCorpID(corpID string, data *TokenData) error { corpID = strings.TrimSpace(corpID) @@ -61,6 +70,16 @@ func SaveTokenDataKeychainForCorpID(corpID string, data *TokenData) error { return saveTokenDataKeychainAccount(TokenAccountForCorpID(corpID), data) } +// SaveTokenDataKeychainForIdentity saves TokenData to an identity-scoped slot. +func SaveTokenDataKeychainForIdentity(corpID, userID string, data *TokenData) error { + corpID = strings.TrimSpace(corpID) + userID = strings.TrimSpace(userID) + if corpID == "" || userID == "" { + return fmt.Errorf("corpId and userId are required for identity token storage") + } + return saveTokenDataKeychainAccount(TokenAccountForIdentity(corpID, userID), data) +} + func saveTokenDataKeychainAccount(account string, data *TokenData) error { jsonData, err := authKeychainMarshal(data, "", " ") if err != nil { @@ -93,6 +112,16 @@ func LoadTokenDataKeychainForCorpID(corpID string) (*TokenData, error) { return loadTokenDataKeychainAccount(TokenAccountForCorpID(corpID)) } +// LoadTokenDataKeychainForIdentity loads TokenData from an identity-scoped slot. +func LoadTokenDataKeychainForIdentity(corpID, userID string) (*TokenData, error) { + corpID = strings.TrimSpace(corpID) + userID = strings.TrimSpace(userID) + if corpID == "" || userID == "" { + return nil, fmt.Errorf("corpId and userId are required for identity token storage") + } + return loadTokenDataKeychainAccount(TokenAccountForIdentity(corpID, userID)) +} + func loadTokenDataKeychainAccount(account string) (*TokenData, error) { jsonStr, err := authKeychainGet(keychain.Service, account) if err != nil { @@ -137,6 +166,19 @@ func preflightTokenPersistence(configDir string) error { ) } } + for _, profile := range cfg.Profiles { + corpID := strings.TrimSpace(profile.CorpID) + userID := strings.TrimSpace(profile.UserID) + if corpID == "" || userID == "" { + continue + } + if _, err := LoadTokenDataKeychainForIdentity(corpID, userID); err != nil && !errors.Is(err, ErrTokenDataNotFound) { + return fmt.Errorf( + "identity token slot %q is unreadable; remove only this account with `dws auth logout --profile %q`, or use `dws auth reset` only when discarding all local profiles: %w", + TokenAccountForIdentity(corpID, userID), ProfileSelector(profile), err, + ) + } + } if err := authValidateEntries(keychain.Service); err != nil { return fmt.Errorf( "auth token ciphertext inventory is unreadable; on macOS first try `env -u DWS_DISABLE_KEYCHAIN dws auth migrate-keychain --to file-dek --dry-run`; if the ciphertext is damaged, use `dws auth reset` only when discarding all local profiles: %w", @@ -161,6 +203,15 @@ func preflightTokenRefreshPersistence(data *TokenData) error { return nil } corpID := strings.TrimSpace(data.CorpID) + userID := strings.TrimSpace(data.UserID) + if userID != "" { + if _, err := LoadTokenDataKeychainForIdentity(corpID, userID); err != nil && !errors.Is(err, ErrTokenDataNotFound) { + return fmt.Errorf("identity token slot %q is unreadable: %w", TokenAccountForIdentity(corpID, userID), err) + } + } + if _, _, exact := ParseIdentitySelector(RuntimeProfile()); exact { + return nil + } if _, err := LoadTokenDataKeychainForCorpID(corpID); err != nil && !errors.Is(err, ErrTokenDataNotFound) { return fmt.Errorf("profile token slot %q is unreadable: %w", TokenAccountForCorpID(corpID), err) } @@ -181,6 +232,16 @@ func DeleteTokenDataKeychainForCorpID(corpID string) error { return authKeychainRemove(keychain.Service, TokenAccountForCorpID(corpID)) } +// DeleteTokenDataKeychainForIdentity removes one identity-scoped token. +func DeleteTokenDataKeychainForIdentity(corpID, userID string) error { + corpID = strings.TrimSpace(corpID) + userID = strings.TrimSpace(userID) + if corpID == "" || userID == "" { + return fmt.Errorf("corpId and userId are required for identity token storage") + } + return authKeychainRemove(keychain.Service, TokenAccountForIdentity(corpID, userID)) +} + // TokenDataExistsKeychain checks if token data exists in keychain. func TokenDataExistsKeychain() bool { return authKeychainExists(keychain.Service, keychain.AccountToken) @@ -195,6 +256,16 @@ func TokenDataExistsKeychainForCorpID(corpID string) bool { return authKeychainExists(keychain.Service, TokenAccountForCorpID(corpID)) } +// TokenDataExistsKeychainForIdentity checks if an identity-scoped token exists. +func TokenDataExistsKeychainForIdentity(corpID, userID string) bool { + corpID = strings.TrimSpace(corpID) + userID = strings.TrimSpace(userID) + if corpID == "" || userID == "" { + return false + } + return authKeychainExists(keychain.Service, TokenAccountForIdentity(corpID, userID)) +} + // EnsureMigration performs one-time migration from legacy .data to keychain. // This should be called early in the auth flow (e.g., during GetAccessToken). // The migration is idempotent and thread-safe. diff --git a/internal/auth/oauth_helpers.go b/internal/auth/oauth_helpers.go index be3658f3..bac901c9 100644 --- a/internal/auth/oauth_helpers.go +++ b/internal/auth/oauth_helpers.go @@ -293,6 +293,8 @@ func (p *OAuthProvider) parseMCPTokenResponse(body []byte) (*TokenData, error) { CorpName string `json:"corpName"` CorpNameSnake string `json:"corp_name"` OrgName string `json:"orgName"` + UserID string `json:"userId"` + UserName string `json:"userName"` // Error fields (when request fails) ErrorCode string `json:"errorCode,omitempty"` ErrorMsg string `json:"errorMsg,omitempty"` @@ -320,6 +322,8 @@ func (p *OAuthProvider) parseMCPTokenResponse(body []byte) (*TokenData, error) { RefreshExpAt: now.Add(config.DefaultRefreshTokenLifetime), CorpID: resp.CorpID, CorpName: firstNonEmpty(resp.CorpName, resp.CorpNameSnake, resp.OrgName), + UserID: resp.UserID, + UserName: resp.UserName, } if resp.PersistentCode != "" { data.PersistentCode = resp.PersistentCode diff --git a/internal/auth/oauth_provider.go b/internal/auth/oauth_provider.go index 7b51d4fc..9fa4bcaf 100644 --- a/internal/auth/oauth_provider.go +++ b/internal/auth/oauth_provider.go @@ -44,6 +44,7 @@ var ( oauthPollInterval = 5 * time.Second oauthSuccessPause = 2 * time.Second oauthLoadToken = LoadTokenData + oauthLoadTokenLocked = loadTokenDataForProfileLocked oauthAcquireLock = AcquireDualLock oauthMarkProfile = MarkProfileStatus oauthFetchClientID = FetchClientIDFromMCP @@ -73,6 +74,9 @@ type OAuthProvider struct { httpClient *http.Client NoBrowser bool TargetCorpID string + // IdentityEnricher resolves userId/userName/corpName while the freshly + // exchanged access token is still only in memory. + IdentityEnricher func(context.Context, *TokenData) error } // NewOAuthProvider creates a new OAuth provider. @@ -555,7 +559,7 @@ continueLogin: // Save token data with associated client ID for refresh tokenData.ClientID = p.clientID - if err := oauthSaveToken(p.configDir, tokenData); err != nil { + if err := p.persistLoginToken(ctx, tokenData); err != nil { return nil, fmt.Errorf("%s: %w", i18n.T("保存 token 失败"), err) } @@ -595,12 +599,13 @@ func (p *OAuthProvider) GetAccessToken(ctx context.Context) (string, error) { if rErr == nil { return refreshed.AccessToken, nil } - _ = oauthMarkProfile(p.configDir, data.CorpID, ProfileStatusExpired) + _ = oauthMarkProfile(p.configDir, TokenProfileSelector(data), ProfileStatusExpired) if p.logger != nil { p.logger.Warn(i18n.T("refresh_token 刷新失败"), "error", rErr) } + return "", fmt.Errorf("%s: %w", i18n.T("refresh_token 刷新失败"), rErr) } else { - _ = oauthMarkProfile(p.configDir, data.CorpID, ProfileStatusExpired) + _ = oauthMarkProfile(p.configDir, TokenProfileSelector(data), ProfileStatusExpired) } return "", errors.New(i18n.T("所有凭证已失效，请运行 dws auth login 重新登录")) @@ -635,7 +640,7 @@ func (p *OAuthProvider) lockedRefresh(ctx context.Context) (*TokenData, error) { // Double-check: re-load from disk — another goroutine/process may have refreshed // while we were waiting for the lock. - data, err := oauthLoadToken(p.configDir) + data, err := oauthLoadTokenLocked(p.configDir, RuntimeProfile()) if err != nil { return nil, err } @@ -672,14 +677,50 @@ func (p *OAuthProvider) ExchangeAuthCode(ctx context.Context, authCode, uid stri return nil, fmt.Errorf("%s: %w", i18n.T("换取 token 失败"), err) } if uid != "" { - tokenData.UserID = uid + tokenData.UserID = strings.TrimSpace(uid) + if err := p.persistKnownLoginToken(tokenData); err != nil { + return nil, fmt.Errorf("%s: %w", i18n.T("保存 token 失败"), err) + } + return tokenData, nil } - if err := oauthSaveToken(p.configDir, tokenData); err != nil { + if err := p.persistLoginToken(ctx, tokenData); err != nil { return nil, fmt.Errorf("%s: %w", i18n.T("保存 token 失败"), err) } return tokenData, nil } +func (p *OAuthProvider) persistLoginToken(ctx context.Context, tokenData *TokenData) error { + if err := p.prepareLoginToken(ctx, tokenData); err != nil { + return err + } + return oauthSaveToken(p.configDir, tokenData) +} + +func (p *OAuthProvider) prepareLoginToken(ctx context.Context, tokenData *TokenData) error { + if tokenData == nil { + return fmt.Errorf("token data is empty") + } + if p != nil && p.IdentityEnricher != nil { + if err := p.IdentityEnricher(ctx, tokenData); err != nil { + return fmt.Errorf("resolve login identity: %w", err) + } + } + if strings.TrimSpace(tokenData.CorpID) != "" && strings.TrimSpace(tokenData.UserID) == "" { + return fmt.Errorf("resolve login identity: userId is required for corpId %q", tokenData.CorpID) + } + return nil +} + +func (p *OAuthProvider) persistKnownLoginToken(tokenData *TokenData) error { + if tokenData == nil { + return fmt.Errorf("token data is empty") + } + if strings.TrimSpace(tokenData.CorpID) != "" && strings.TrimSpace(tokenData.UserID) == "" { + return fmt.Errorf("resolve login identity: userId is required for corpId %q", tokenData.CorpID) + } + return oauthSaveToken(p.configDir, tokenData) +} + // Logout clears all stored credentials. func (p *OAuthProvider) Logout() error { return DeleteTokenData(p.configDir) diff --git a/internal/auth/oauth_provider_coverage_test.go b/internal/auth/oauth_provider_coverage_test.go index 08b04d71..1be87470 100644 --- a/internal/auth/oauth_provider_coverage_test.go +++ b/internal/auth/oauth_provider_coverage_test.go @@ -641,6 +641,7 @@ func TestCrossPlatformCoverageOAuthRefreshAndParsingEdges(t *testing.T) { func TestCrossPlatformCoverageOAuthProviderHighLevelEdges(t *testing.T) { oldLoad := oauthLoadToken + oldLoadLocked := oauthLoadTokenLocked oldAcquire := oauthAcquireLock oldMark := oauthMarkProfile oldFetch := oauthFetchClientID @@ -651,6 +652,7 @@ func TestCrossPlatformCoverageOAuthProviderHighLevelEdges(t *testing.T) { oldSaveLocked := oauthSaveTokenLocked t.Cleanup(func() { oauthLoadToken = oldLoad + oauthLoadTokenLocked = oldLoadLocked oauthAcquireLock = oldAcquire oauthMarkProfile = oldMark oauthFetchClientID = oldFetch @@ -665,6 +667,9 @@ func TestCrossPlatformCoverageOAuthProviderHighLevelEdges(t *testing.T) { }) fail := errors.New("fail") p := &OAuthProvider{configDir: t.TempDir(), logger: slog.Default(), Output: io.Discard} + oauthLoadTokenLocked = func(configDir, _ string) (*TokenData, error) { + return oauthLoadToken(configDir) + } valid := &TokenData{AccessToken: "valid", ExpiresAt: time.Now().Add(time.Hour)} oauthLoadToken = func(string) (*TokenData, error) { return valid, nil } diff --git a/internal/auth/portable_store_test.go b/internal/auth/portable_store_test.go index 9fc98ec5..eb7ac9d5 100644 --- a/internal/auth/portable_store_test.go +++ b/internal/auth/portable_store_test.go @@ -288,7 +288,7 @@ func TestPortableAuthBundleRoundTripPreservesProfiles(t *testing.T) { if err != nil { t.Fatalf("LoadProfiles() after import error = %v", err) } - if cfg.PrimaryProfile != "corp_a" || cfg.CurrentProfile != "corp_b" || cfg.PreviousProfile != "corp_a" { + if cfg.PrimaryProfile != "" || cfg.CurrentProfile != "corp_b" || cfg.PreviousProfile != "corp_a" { t.Fatalf("profiles after import = %#v", cfg) } if len(cfg.Profiles) != 2 { @@ -311,6 +311,65 @@ func TestPortableAuthBundleRoundTripPreservesProfiles(t *testing.T) { } } +func TestPortableAuthBundleRoundTripPreservesSameCorpAccounts(t *testing.T) { + requirePortableFileBackend(t) + t.Setenv(keychain.DisableKeychainEnv, "1") + SetRuntimeProfile("") + t.Cleanup(func() { SetRuntimeProfile("") }) + + sourceKeychain := filepath.Join(t.TempDir(), "source-keychain") + t.Setenv(keychain.StorageDirEnv, sourceKeychain) + sourceConfig := filepath.Join(t.TempDir(), ".dws") + + first := &TokenData{ + AccessToken: "access-first", + RefreshToken: "refresh-first", + ExpiresAt: time.Now().Add(time.Hour), + RefreshExpAt: time.Now().Add(30 * 24 * time.Hour), + CorpID: "corp_same", + CorpName: "Same Org", + UserID: "user_1", + UserName: "账号一", + } + second := *first + second.AccessToken = "access-second" + second.RefreshToken = "refresh-second" + second.UserID = "user_2" + second.UserName = "账号二" + if err := SaveTokenData(sourceConfig, first); err != nil { + t.Fatalf("SaveTokenData(first) error = %v", err) + } + if err := SaveTokenData(sourceConfig, &second); err != nil { + t.Fatalf("SaveTokenData(second) error = %v", err) + } + + var bundle bytes.Buffer + if err := ExportPortableAuthBundle(sourceConfig, &bundle); err != nil { + t.Fatalf("ExportPortableAuthBundle() error = %v", err) + } + + targetKeychain := filepath.Join(t.TempDir(), "target-keychain") + t.Setenv(keychain.StorageDirEnv, targetKeychain) + targetConfig := filepath.Join(t.TempDir(), ".dws") + if _, err := ImportPortableAuthBundle(targetConfig, bytes.NewReader(bundle.Bytes())); err != nil { + t.Fatalf("ImportPortableAuthBundle() error = %v", err) + } + + for selector, wantToken := range map[string]string{ + "corp_same:user_1": "access-first", + "corp_same:user_2": "access-second", + "corp_same": "access-second", + } { + loaded, err := LoadTokenDataForProfile(targetConfig, selector) + if err != nil { + t.Fatalf("LoadTokenDataForProfile(%s) error = %v", selector, err) + } + if loaded.AccessToken != wantToken { + t.Fatalf("profile %s token = %q, want %q", selector, loaded.AccessToken, wantToken) + } + } +} + func requirePortableFileBackend(t *testing.T) { t.Helper() if runtime.GOOS == "windows" { diff --git a/internal/auth/profiles.go b/internal/auth/profiles.go index 79dd671c..0badcf10 100644 --- a/internal/auth/profiles.go +++ b/internal/auth/profiles.go @@ -16,9 +16,11 @@ package auth import ( "context" "encoding/json" + "errors" "fmt" "os" "path/filepath" + "sort" "strings" "sync" "time" @@ -43,8 +45,10 @@ var ( profilesTokenExists = TokenDataExistsKeychain profilesLoadLegacy = LoadTokenDataKeychain profilesSaveCorp = SaveTokenDataKeychainForCorpID + profilesSaveIdentity = SaveTokenDataKeychainForIdentity profilesTokenExistsCorp = TokenDataExistsKeychainForCorpID profilesLoadCorp = LoadTokenDataKeychainForCorpID + profilesLoadIdentity = LoadTokenDataKeychainForIdentity profilesSaveLegacy = SaveTokenDataKeychain profilesWriteMarker = WriteTokenMarker profilesDeleteLegacy = DeleteTokenDataKeychain @@ -68,7 +72,10 @@ func withProfilesLock(configDir string, fn func() error) error { return fn() } -const profilesJSONFile = "profiles.json" +const ( + profilesJSONFile = "profiles.json" + profilesVersion = 2 +) const ( ProfileStatusActive = "active" @@ -78,11 +85,12 @@ const ( // ProfilesConfig stores non-sensitive profile metadata. Token material stays in keychain. type ProfilesConfig struct { - Version int `json:"version"` - PrimaryProfile string `json:"primaryProfile,omitempty"` - CurrentProfile string `json:"currentProfile,omitempty"` - PreviousProfile string `json:"previousProfile,omitempty"` - Profiles []Profile `json:"profiles,omitempty"` + Version int `json:"version"` + PrimaryProfile string `json:"primaryProfile,omitempty"` + CurrentProfile string `json:"currentProfile,omitempty"` + PreviousProfile string `json:"previousProfile,omitempty"` + OrgCurrentProfiles map[string]string `json:"orgCurrentProfiles,omitempty"` + Profiles []Profile `json:"profiles,omitempty"` } // Profile is a logged-in DingTalk organization identity. @@ -193,20 +201,141 @@ func ensureProfilesMigrationLocked(configDir string) error { if err != nil { return err } - if len(cfg.Profiles) > 0 { - 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. + if cfg.Version >= profilesVersion { + return nil + } + if !profilesTokenExists() { + return nil + } + data, loadErr := profilesLoadLegacy() + if loadErr != nil || data == nil || strings.TrimSpace(data.CorpID) == "" { + return nil + } + if err := profilesSaveCorp(data.CorpID, data); err != nil { + return err + } + if strings.TrimSpace(data.UserID) != "" { + if err := profilesSaveIdentity(data.CorpID, data.UserID, data); err != nil { + return err + } + } + return upsertProfileFromToken(configDir, cfg, data, false) } - if !profilesTokenExists() { - return nil + + legacySelectionState := cfg.Version < profilesVersion + changed := false + if cfg.OrgCurrentProfiles == nil { + cfg.OrgCurrentProfiles = make(map[string]string) } - data, err := profilesLoadLegacy() - if err != nil || data == nil || strings.TrimSpace(data.CorpID) == "" { - return nil + orgTokens := make(map[string]*TokenData) + for i := range cfg.Profiles { + p := &cfg.Profiles[i] + corpID := strings.TrimSpace(p.CorpID) + if corpID == "" { + continue + } + orgToken, loaded := orgTokens[corpID] + if !loaded { + token, loadErr := profilesLoadCorp(corpID) + if loadErr == nil { + orgToken = token + } + orgTokens[corpID] = orgToken + } + if orgToken == nil { + continue + } + if strings.TrimSpace(p.UserID) == "" && strings.TrimSpace(orgToken.UserID) != "" { + if existing := findExactProfile(cfg, corpID, orgToken.UserID); existing != nil && existing != p { + p.CorpID = "" + changed = true + continue + } + p.UserID = strings.TrimSpace(orgToken.UserID) + if p.UserName == "" { + p.UserName = strings.TrimSpace(orgToken.UserName) + } + changed = true + } + if strings.TrimSpace(p.UserID) == "" || strings.TrimSpace(orgToken.UserID) != strings.TrimSpace(p.UserID) { + continue + } + _, identityErr := profilesLoadIdentity(corpID, p.UserID) + if errors.Is(identityErr, ErrTokenDataNotFound) { + if err := profilesSaveIdentity(corpID, p.UserID, orgToken); err != nil { + return err + } + } else if identityErr != nil { + return identityErr + } } - if err := profilesSaveCorp(data.CorpID, data); err != nil { - return err + for _, corpID := range uniqueProfileCorpIDs(cfg) { + if exactProfileSelectorForCorp(cfg, corpID, cfg.OrgCurrentProfiles[corpID]) != "" { + continue + } + if _, exists := cfg.OrgCurrentProfiles[corpID]; exists { + delete(cfg.OrgCurrentProfiles, corpID) + changed = true + } + if !legacySelectionState { + continue + } + + if orgToken := orgTokens[corpID]; orgToken != nil { + if p := findExactProfile(cfg, corpID, orgToken.UserID); p != nil { + cfg.OrgCurrentProfiles[corpID] = ProfileSelector(*p) + changed = true + continue + } + } + for _, pointer := range []string{cfg.CurrentProfile, cfg.PreviousProfile, cfg.PrimaryProfile} { + if exact := exactProfileSelectorForCorp(cfg, corpID, pointer); exact != "" { + cfg.OrgCurrentProfiles[corpID] = exact + changed = true + break + } + } + if cfg.OrgCurrentProfiles[corpID] != "" { + continue + } + profiles := profilesForCorpID(cfg, corpID) + if len(profiles) == 1 { + cfg.OrgCurrentProfiles[corpID] = ProfileSelector(*profiles[0]) + changed = true + } } - return upsertProfileFromToken(configDir, cfg, data, false) + if exact := canonicalStoredSelector(cfg, cfg.CurrentProfile); exact != "" && exact != cfg.CurrentProfile { + cfg.CurrentProfile = exact + changed = true + } + if legacySelectionState && cfg.CurrentProfile == "" { + if exact := canonicalStoredSelector(cfg, cfg.PrimaryProfile); exact != "" { + cfg.CurrentProfile = exact + changed = true + } else if len(cfg.Profiles) == 1 { + cfg.CurrentProfile = ProfileSelector(cfg.Profiles[0]) + changed = true + } + } + if exact := canonicalStoredSelector(cfg, cfg.PreviousProfile); exact != "" && exact != cfg.PreviousProfile { + cfg.PreviousProfile = exact + changed = true + } + if cfg.PreviousProfile != "" && cfg.PreviousProfile == cfg.CurrentProfile { + cfg.PreviousProfile = "" + changed = true + } + if cfg.Version < profilesVersion { + cfg.Version = profilesVersion + changed = true + } + if changed { + return profilesSave(configDir, cfg) + } + return nil } // UpsertProfileFromToken updates profiles.json after a successful login or refresh. @@ -239,35 +368,38 @@ func upsertProfileFromToken(configDir string, cfg *ProfilesConfig, data *TokenDa return nil } normalizeProfilesConfig(cfg) - now := time.Now().Format(time.RFC3339) - idx := profileIndexByCorpID(cfg, corpID) + cfg.Version = profilesVersion + userID := strings.TrimSpace(data.UserID) + var previousCurrent *Profile + if makeCurrent && strings.TrimSpace(cfg.CurrentProfile) != "" { + previousCurrent, _, _ = resolveProfileSelection(configDir, cfg, cfg.CurrentProfile) + } + idx := profileIndexByIdentity(cfg, corpID, userID) + if idx < 0 && userID != "" { + idx = legacyProfileIndexByCorpID(cfg, corpID) + } if idx < 0 { profile := Profile{ - Name: chooseProfileName(cfg, data), - CorpID: corpID, - CorpName: strings.TrimSpace(data.CorpName), - UserID: strings.TrimSpace(data.UserID), - UserName: strings.TrimSpace(data.UserName), - ClientID: strings.TrimSpace(data.ClientID), - Status: ProfileStatusActive, - ExpiresAt: timeOrRFC3339(data.ExpiresAt), - RefreshExpAt: timeOrRFC3339(data.RefreshExpAt), - LastLoginAt: now, - LastUsedAt: now, - UpdatedAt: now, + Name: chooseProfileName(cfg, data), + CorpID: corpID, + CorpName: strings.TrimSpace(data.CorpName), + UserID: userID, + UserName: strings.TrimSpace(data.UserName), + ClientID: strings.TrimSpace(data.ClientID), + Status: ProfileStatusActive, } cfg.Profiles = append(cfg.Profiles, profile) } else { p := &cfg.Profiles[idx] + if userID != "" { + p.UserID = userID + } if shouldRefreshProfileName(p, data) { p.Name = chooseProfileName(cfg, data) } if v := strings.TrimSpace(data.CorpName); v != "" { p.CorpName = v } - if v := strings.TrimSpace(data.UserID); v != "" { - p.UserID = v - } if v := strings.TrimSpace(data.UserName); v != "" { p.UserName = v } @@ -275,81 +407,153 @@ func upsertProfileFromToken(configDir string, cfg *ProfilesConfig, data *TokenDa p.ClientID = v } p.Status = ProfileStatusActive - p.ExpiresAt = timeOrRFC3339(data.ExpiresAt) - p.RefreshExpAt = timeOrRFC3339(data.RefreshExpAt) - p.LastLoginAt = now - p.LastUsedAt = now - p.UpdatedAt = now } - if cfg.PrimaryProfile == "" { - cfg.PrimaryProfile = corpID - } - if makeCurrent && cfg.CurrentProfile != corpID { - if cfg.CurrentProfile != "" { - cfg.PreviousProfile = cfg.CurrentProfile + if makeCurrent { + newSelector := profileSelector(corpID, userID) + if previousCurrent != nil && ProfileSelector(*previousCurrent) != newSelector { + cfg.PreviousProfile = ProfileSelector(*previousCurrent) } - cfg.CurrentProfile = corpID + cfg.CurrentProfile = newSelector + setOrgCurrentProfile(cfg, corpID, newSelector) } if cfg.CurrentProfile == "" { - cfg.CurrentProfile = corpID + cfg.CurrentProfile = profileSelector(corpID, userID) + setOrgCurrentProfile(cfg, corpID, cfg.CurrentProfile) } return profilesSave(configDir, cfg) } -// ResolveProfile returns a profile selected by name/corpId or by current/primary fallback. -func ResolveProfile(configDir, selector string) (*Profile, error) { - if err := profilesEnsureMigration(configDir); err != nil { - return nil, err - } - cfg, err := profilesLoad(configDir) - if err != nil { - return nil, err - } - selector = strings.TrimSpace(selector) - if selector != "" { - p := findProfile(cfg, selector) - if p == nil { - return nil, fmt.Errorf("profile %q not found", selector) - } - return p, nil - } - if p := findProfile(cfg, cfg.CurrentProfile); p != nil { - return p, nil - } - if p := findProfile(cfg, cfg.PrimaryProfile); p != nil { - return p, nil - } - return nil, nil +// ProfileSelector returns the exact identity selector for a profile when its +// userId is known, otherwise it returns the historical corpId selector. +func ProfileSelector(profile Profile) string { + return profileSelector(profile.CorpID, profile.UserID) } -func resolveProfileForLoad(configDir, selector string) (*Profile, error) { +// TokenProfileSelector returns the exact identity selector for token data when +// its userId is known, otherwise it returns the historical corpId selector. +func TokenProfileSelector(data *TokenData) string { + if data == nil { + return "" + } + return profileSelector(data.CorpID, data.UserID) +} + +func profileSelector(corpID, userID string) string { + corpID = strings.TrimSpace(corpID) + userID = strings.TrimSpace(userID) + if corpID == "" { + return "" + } + if userID == "" { + return corpID + } + return corpID + ":" + userID +} + +// ParseIdentitySelector splits corpId:userId selectors. +func ParseIdentitySelector(selector string) (corpID, userID string, ok bool) { + selector = strings.TrimSpace(selector) + idx := strings.LastIndex(selector, ":") + if idx <= 0 || idx >= len(selector)-1 { + return "", "", false + } + corpID = strings.TrimSpace(selector[:idx]) + userID = strings.TrimSpace(selector[idx+1:]) + if corpID == "" || userID == "" { + return "", "", false + } + return corpID, userID, true +} + +// ResolveProfile returns a profile selected by name/corpId/identity or by +// current fallback. +func ResolveProfile(configDir, selector string) (*Profile, error) { + p, _, err := ResolveProfileWithScope(configDir, selector) + return p, err +} + +// ResolveProfileWithScope resolves a selector and reports whether it targets +// one identity (compound selector or local profile name) rather than an +// organization as a whole. +func ResolveProfileWithScope(configDir, selector string) (*Profile, bool, error) { + var ( + result *Profile + exact bool + ) + err := withProfilesLock(configDir, func() error { + var resolveErr error + result, exact, resolveErr = resolveProfileWithScopeLocked(configDir, selector) + return resolveErr + }) + return result, exact, err +} + +func resolveProfileWithScopeLocked(configDir, selector string) (*Profile, bool, error) { if err := profilesEnsureMigration(configDir); err != nil { - return nil, err + return nil, false, err } cfg, err := profilesLoad(configDir) if err != nil { - return nil, err + return nil, false, err } selector = strings.TrimSpace(selector) if selector != "" { - p := findProfile(cfg, selector) - if p == nil { - return nil, fmt.Errorf("profile %q not found", selector) + return resolveProfileSelection(configDir, cfg, selector) + } + if strings.TrimSpace(cfg.CurrentProfile) != "" { + return resolveProfileSelection(configDir, cfg, cfg.CurrentProfile) + } + return nil, false, nil +} + +// ResolveProfileDeletionScope resolves selectors for destructive profile +// removal. Organization selectors intentionally resolve to the whole +// organization even when it has multiple accounts and no current account. +func ResolveProfileDeletionScope(configDir, selector string) (*Profile, bool, error) { + var ( + result *Profile + exact bool + ) + err := withProfilesLock(configDir, func() error { + if err := profilesEnsureMigration(configDir); err != nil { + return err } - return p, nil - } - for _, candidate := range []string{cfg.CurrentProfile, cfg.PrimaryProfile} { - if p := findProfile(cfg, candidate); p != nil && profilesTokenExistsCorp(p.CorpID) { - return p, nil + cfg, err := profilesLoad(configDir) + if err != nil { + return err } + var resolveErr error + result, exact, resolveErr = resolveProfileDeletionSelection(cfg, selector) + return resolveErr + }) + return result, exact, err +} + +func resolveProfileForLoadLocked(configDir, selector string) (*Profile, *ProfilesConfig, error) { + if err := profilesEnsureMigration(configDir); err != nil { + return nil, nil, err } - if p := findProfile(cfg, cfg.CurrentProfile); p != nil { - return p, nil + cfg, err := profilesLoad(configDir) + if err != nil { + return nil, nil, err } - if p := findProfile(cfg, cfg.PrimaryProfile); p != nil { - return p, nil + selector = strings.TrimSpace(selector) + if selector != "" { + p, _, resolveErr := resolveProfileSelection(configDir, cfg, selector) + return p, cfg, resolveErr } - return nil, nil + if strings.TrimSpace(cfg.CurrentProfile) != "" { + p, _, resolveErr := resolveProfileSelection(configDir, cfg, cfg.CurrentProfile) + return p, cfg, resolveErr + } + return nil, cfg, nil +} + +// resolveProfileForLoad is retained as a focused test seam around the locked +// resolver. Production callers already hold the auth lock. +func resolveProfileForLoad(configDir, selector string) (*Profile, error) { + profile, _, err := resolveProfileForLoadLocked(configDir, selector) + return profile, err } // SetCurrentProfile persists the selected current profile. @@ -371,24 +575,32 @@ func setCurrentProfileLocked(configDir, selector string) (*Profile, error) { if err != nil { return nil, err } - p := findProfile(cfg, selector) - if p == nil { - return nil, fmt.Errorf("profile %q not found", strings.TrimSpace(selector)) + p, _, err := resolveProfileSelection(configDir, cfg, selector) + if err != nil { + return nil, err } - if cfg.CurrentProfile != p.CorpID { - if cfg.CurrentProfile != "" { - cfg.PreviousProfile = cfg.CurrentProfile + var previousCurrent *Profile + if strings.TrimSpace(cfg.CurrentProfile) != "" { + previousCurrent, _, _ = resolveProfileSelection(configDir, cfg, cfg.CurrentProfile) + } + storedSelector := ProfileSelector(*p) + if cfg.CurrentProfile != storedSelector { + if previousCurrent != nil { + cfg.PreviousProfile = ProfileSelector(*previousCurrent) } - cfg.CurrentProfile = p.CorpID + cfg.CurrentProfile = storedSelector } - touchProfile(cfg, p.CorpID) + setOrgCurrentProfile(cfg, p.CorpID, storedSelector) if err := profilesSave(configDir, cfg); err != nil { return nil, err } + if err := syncOrganizationTokenMirrorForProfile(*p); err != nil { + return nil, err + } if err := profilesSyncLegacyMirror(configDir); err != nil { return nil, err } - return findProfile(cfg, p.CorpID), nil + return findExactProfile(cfg, p.CorpID, p.UserID), nil } // UsePreviousProfile toggles currentProfile and previousProfile. @@ -414,25 +626,40 @@ func usePreviousProfileLocked(configDir string) (*Profile, error) { if prev == "" { return nil, fmt.Errorf("previous profile is empty") } - p := findProfile(cfg, prev) - if p == nil { - return nil, fmt.Errorf("previous profile %q not found", prev) + p, _, err := resolveProfileSelection(configDir, cfg, prev) + if err != nil { + return nil, fmt.Errorf("resolve previous profile %q: %w", prev, err) } - cfg.PreviousProfile, cfg.CurrentProfile = cfg.CurrentProfile, p.CorpID - touchProfile(cfg, p.CorpID) + var current *Profile + if strings.TrimSpace(cfg.CurrentProfile) != "" { + current, _, _ = resolveProfileSelection(configDir, cfg, cfg.CurrentProfile) + } + cfg.CurrentProfile = ProfileSelector(*p) + if current != nil { + cfg.PreviousProfile = ProfileSelector(*current) + } else { + cfg.PreviousProfile = "" + } + setOrgCurrentProfile(cfg, p.CorpID, ProfileSelector(*p)) if err := profilesSave(configDir, cfg); err != nil { return nil, err } + if err := syncOrganizationTokenMirrorForProfile(*p); err != nil { + return nil, err + } if err := profilesSyncLegacyMirror(configDir); err != nil { return nil, err } - return findProfile(cfg, p.CorpID), nil + return findExactProfile(cfg, p.CorpID, p.UserID), nil } // RemoveProfile removes a profile from metadata and returns the removed profile. func RemoveProfile(configDir, selector string) (*Profile, error) { var result *Profile err := withProfilesLock(configDir, func() error { + if err := ensureProfilesMigrationLocked(configDir); err != nil { + return err + } p, e := removeProfileLocked(configDir, selector) result = p return e @@ -445,34 +672,68 @@ func removeProfileLocked(configDir, selector string) (*Profile, error) { if err != nil { return nil, err } - p := findProfile(cfg, selector) - if p == nil { - return nil, fmt.Errorf("profile %q not found", strings.TrimSpace(selector)) + p, exact, err := resolveProfileDeletionSelection(cfg, selector) + if err != nil { + return nil, err } removed := *p kept := cfg.Profiles[:0] for _, profile := range cfg.Profiles { - if profile.CorpID != removed.CorpID { + remove := profile.CorpID == removed.CorpID + if exact { + remove = sameProfileIdentity(profile.CorpID, profile.UserID, removed.CorpID, removed.UserID) + } + if !remove { kept = append(kept, profile) } } cfg.Profiles = kept - if cfg.PrimaryProfile == removed.CorpID { - cfg.PrimaryProfile = firstProfileCorpID(cfg) - } - if cfg.CurrentProfile == removed.CorpID { - cfg.CurrentProfile = cfg.PrimaryProfile - if cfg.CurrentProfile == "" { - cfg.CurrentProfile = firstProfileCorpID(cfg) + replacementSelector := "" + if exact && len(profilesForCorpID(cfg, removed.CorpID)) > 0 { + remaining := profilesForCorpID(cfg, removed.CorpID) + if len(remaining) == 1 { + replacementSelector = ProfileSelector(*remaining[0]) } } - if cfg.PreviousProfile == removed.CorpID { + if exact { + if selectorMatchesIdentity(cfg.OrgCurrentProfiles[removed.CorpID], removed) { + if replacementSelector == "" { + delete(cfg.OrgCurrentProfiles, removed.CorpID) + } else { + cfg.OrgCurrentProfiles[removed.CorpID] = replacementSelector + } + } + } else { + delete(cfg.OrgCurrentProfiles, removed.CorpID) + } + pointers := []*string{&cfg.PrimaryProfile, &cfg.CurrentProfile, &cfg.PreviousProfile} + for _, pointer := range pointers { + if exact { + if selectorMatchesIdentity(*pointer, removed) { + *pointer = replacementSelector + } + continue + } + if selectorTargetsCorp(*pointer, removed.CorpID) { + *pointer = "" + } + } + if cfg.CurrentProfile == "" { + if previous := canonicalStoredSelector(cfg, cfg.PreviousProfile); previous != "" { + cfg.CurrentProfile = previous + cfg.PreviousProfile = "" + } else if len(cfg.Profiles) == 1 { + cfg.CurrentProfile = ProfileSelector(cfg.Profiles[0]) + } + } + if cfg.PreviousProfile != "" && cfg.PreviousProfile == cfg.CurrentProfile { cfg.PreviousProfile = "" } if len(cfg.Profiles) == 0 { cfg.PrimaryProfile = "" cfg.CurrentProfile = "" cfg.PreviousProfile = "" + cfg.OrgCurrentProfiles = nil } if err := profilesSave(configDir, cfg); err != nil { return nil, err @@ -480,23 +741,38 @@ func removeProfileLocked(configDir, selector string) (*Profile, error) { return &removed, nil } -// MarkProfileStatus updates a profile status if it exists. -func MarkProfileStatus(configDir, corpID, status string) error { - if strings.TrimSpace(corpID) == "" { +func selectorMatchesIdentity(selector string, profile Profile) bool { + corpID, userID, exact := ParseIdentitySelector(selector) + return exact && sameProfileIdentity(corpID, userID, profile.CorpID, profile.UserID) +} + +func selectorTargetsCorp(selector, corpID string) bool { + selector = strings.TrimSpace(selector) + corpID = strings.TrimSpace(corpID) + if selector == corpID { + return true + } + selectedCorpID, _, exact := ParseIdentitySelector(selector) + return exact && selectedCorpID == corpID +} + +// MarkProfileStatus updates a selected profile status if it exists. +func MarkProfileStatus(configDir, selector, status string) error { + if strings.TrimSpace(selector) == "" { return nil } return withProfilesLock(configDir, func() error { - return markProfileStatusLocked(configDir, corpID, status) + return markProfileStatusLocked(configDir, selector, status) }) } -func markProfileStatusLocked(configDir, corpID, status string) error { +func markProfileStatusLocked(configDir, selector, status string) error { cfg, err := profilesLoad(configDir) if err != nil { return err } - p := findProfile(cfg, corpID) - if p == nil { + p, _, resolveErr := resolveProfileSelection(configDir, cfg, selector) + if resolveErr != nil { return nil } p.Status = strings.TrimSpace(status) @@ -516,17 +792,17 @@ func syncLegacyTokenMirrorLocked(configDir string) error { if err != nil { return err } - hadReadError := false - for _, candidate := range []string{cfg.CurrentProfile, cfg.PrimaryProfile} { - p := findProfile(cfg, candidate) - if p == nil { - continue + current := strings.TrimSpace(cfg.CurrentProfile) + if current != "" { + p, _, resolveErr := resolveProfileSelection(configDir, cfg, current) + if resolveErr != nil { + return resolveErr } - data, loadErr := profilesLoadCorp(p.CorpID) + data, loadErr := loadTokenForProfileIdentity(*p) if loadErr != nil { - // Transient keychain read failure: do NOT touch the existing mirror. - hadReadError = true - continue + // 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 { @@ -535,35 +811,73 @@ func syncLegacyTokenMirrorLocked(configDir string) error { return profilesWriteMarker(configDir) } } - if hadReadError { - // Keep the existing legacy mirror untouched rather than wiping a host - // app's login state just because keychain was momentarily unavailable. - return nil + // No current profile (or its token is confirmed absent): clear the mirror. + if err := profilesDeleteLegacy(); err != nil { + return err } - // All candidate profiles confirmed absent (no token): clear the mirror. - _ = profilesDeleteLegacy() - _ = profilesDeleteMarker(configDir) - return nil + return profilesDeleteMarker(configDir) +} + +func syncOrganizationTokenMirrorForProfile(profile Profile) error { + data, err := loadTokenForProfileIdentity(profile) + if err != nil { + return err + } + return profilesSaveCorp(profile.CorpID, data) +} + +func loadTokenForProfileIdentity(profile Profile) (*TokenData, error) { + if strings.TrimSpace(profile.UserID) != "" { + data, err := profilesLoadIdentity(profile.CorpID, profile.UserID) + if err == nil { + return data, nil + } + if !errors.Is(err, ErrTokenDataNotFound) { + return nil, err + } + orgData, orgErr := profilesLoadCorp(profile.CorpID) + if orgErr != nil { + if errors.Is(orgErr, ErrTokenDataNotFound) { + return nil, err + } + return nil, orgErr + } + if strings.TrimSpace(orgData.UserID) == "" { + return nil, fmt.Errorf("organization token mirror for corpId %q has no userId; cannot use it for profile %q", profile.CorpID, ProfileSelector(profile)) + } + if strings.TrimSpace(orgData.UserID) != strings.TrimSpace(profile.UserID) { + return nil, err + } + if saveErr := profilesSaveIdentity(profile.CorpID, profile.UserID, orgData); saveErr != nil { + return nil, saveErr + } + return orgData, nil + } + return profilesLoadCorp(profile.CorpID) } func normalizeProfilesConfig(cfg *ProfilesConfig) { if cfg == nil { return } - cfg.Version = 1 + if cfg.Version <= 0 { + cfg.Version = 1 + } seen := make(map[string]bool, len(cfg.Profiles)) profiles := cfg.Profiles[:0] for _, p := range cfg.Profiles { p.CorpID = strings.TrimSpace(p.CorpID) - if p.CorpID == "" || seen[p.CorpID] { + p.UserID = strings.TrimSpace(p.UserID) + identity := profileIdentityKey(p.CorpID, p.UserID) + if p.CorpID == "" || seen[identity] { continue } - seen[p.CorpID] = true + seen[identity] = true p.Name = strings.TrimSpace(p.Name) if p.Name == "" { p.Name = p.CorpID } - if corpName := strings.TrimSpace(p.CorpName); p.Name == p.CorpID && corpName != "" && !profileNameTakenByOtherCorp(cfg, corpName, p.CorpID) { + if corpName := strings.TrimSpace(p.CorpName); p.Name == p.CorpID && corpName != "" && !profileNameTakenByOtherIdentity(cfg, corpName, p.CorpID, p.UserID) { p.Name = corpName } if p.Status == "" { @@ -572,21 +886,29 @@ func normalizeProfilesConfig(cfg *ProfilesConfig) { profiles = append(profiles, p) } cfg.Profiles = profiles - if cfg.PrimaryProfile != "" && findProfile(cfg, cfg.PrimaryProfile) == nil { + if len(cfg.OrgCurrentProfiles) > 0 { + normalized := make(map[string]string, len(cfg.OrgCurrentProfiles)) + for corpID, selector := range cfg.OrgCurrentProfiles { + corpID = strings.TrimSpace(corpID) + if exact := exactProfileSelectorForCorp(cfg, corpID, selector); exact != "" { + normalized[corpID] = exact + } + } + if len(normalized) == 0 { + cfg.OrgCurrentProfiles = nil + } else { + cfg.OrgCurrentProfiles = normalized + } + } + if cfg.PrimaryProfile != "" && !profileSelectorReferenceExists(cfg, cfg.PrimaryProfile) { cfg.PrimaryProfile = "" } - if cfg.CurrentProfile != "" && findProfile(cfg, cfg.CurrentProfile) == nil { + if cfg.CurrentProfile != "" && !profileSelectorReferenceExists(cfg, cfg.CurrentProfile) { cfg.CurrentProfile = "" } - if cfg.PreviousProfile != "" && findProfile(cfg, cfg.PreviousProfile) == nil { + if cfg.PreviousProfile != "" && !profileSelectorReferenceExists(cfg, cfg.PreviousProfile) { cfg.PreviousProfile = "" } - if cfg.PrimaryProfile == "" { - cfg.PrimaryProfile = firstProfileCorpID(cfg) - } - if cfg.CurrentProfile == "" { - cfg.CurrentProfile = cfg.PrimaryProfile - } } func chooseProfileName(cfg *ProfilesConfig, data *TokenData) string { @@ -597,17 +919,27 @@ func chooseProfileName(cfg *ProfilesConfig, data *TokenData) string { if base == "" { base = "profile" } - if !profileNameTakenByOtherCorp(cfg, base, data.CorpID) { + corpID := strings.TrimSpace(data.CorpID) + userID := strings.TrimSpace(data.UserID) + if !profileNameTakenByOtherIdentity(cfg, base, corpID, userID) { return base } - suffix := shortCorpID(data.CorpID) - name := base + "-" + suffix - if !profileNameTakenByOtherCorp(cfg, name, data.CorpID) { - return name + for _, suffix := range []string{ + strings.TrimSpace(data.UserName), + shortProfileID(userID), + shortCorpID(corpID), + } { + if suffix == "" { + continue + } + name := base + "-" + suffix + if !profileNameTakenByOtherIdentity(cfg, name, corpID, userID) { + return name + } } for i := 2; ; i++ { - candidate := fmt.Sprintf("%s-%s-%d", base, suffix, i) - if !profileNameTakenByOtherCorp(cfg, candidate, data.CorpID) { + candidate := fmt.Sprintf("%s-%d", base, i) + if !profileNameTakenByOtherIdentity(cfg, candidate, corpID, userID) { return candidate } } @@ -624,72 +956,387 @@ func shouldRefreshProfileName(p *Profile, data *TokenData) bool { return strings.TrimSpace(data.CorpName) != "" && name == strings.TrimSpace(p.CorpID) } -func profileNameTakenByOtherCorp(cfg *ProfilesConfig, name, corpID string) bool { +func profileNameTakenByOtherIdentity(cfg *ProfilesConfig, name, corpID, userID string) bool { name = strings.TrimSpace(name) corpID = strings.TrimSpace(corpID) + userID = strings.TrimSpace(userID) for _, p := range cfg.Profiles { - if p.CorpID != corpID && p.Name == name { + if strings.TrimSpace(p.Name) == name && !sameProfileIdentity(p.CorpID, p.UserID, corpID, userID) { return true } } return false } -func findProfile(cfg *ProfilesConfig, selector string) *Profile { +// resolveProfileSelection resolves a user-facing selector to one exact identity. +// The bool reports whether the selector targets one identity (compound selector +// or local profile name) rather than an organization as a whole. +func resolveProfileSelection(_ string, cfg *ProfilesConfig, selector string) (*Profile, bool, error) { if cfg == nil { - return nil + return nil, false, fmt.Errorf("profile %q not found", strings.TrimSpace(selector)) } selector = strings.TrimSpace(selector) if selector == "" { - return nil + return nil, false, fmt.Errorf("profile selector is empty") } - var corpNameMatch *Profile - for i := range cfg.Profiles { - if cfg.Profiles[i].CorpID == selector || cfg.Profiles[i].Name == selector { - return &cfg.Profiles[i] + + if organization, account, compound := ParseIdentitySelector(selector); compound { + corpID, err := resolveOrganizationCorpID(cfg, organization) + if err != nil { + return nil, true, err } - if strings.TrimSpace(cfg.Profiles[i].CorpName) == selector { - if corpNameMatch != nil { - return nil + if corpID == "" { + return nil, true, fmt.Errorf("organization %q not found", organization) + } + if p := findExactProfile(cfg, corpID, account); p != nil { + return p, true, nil + } + var matches []*Profile + for _, p := range profilesForCorpID(cfg, corpID) { + if strings.TrimSpace(p.UserName) == account { + matches = append(matches, p) } - corpNameMatch = &cfg.Profiles[i] + } + switch len(matches) { + case 1: + return matches[0], true, nil + case 0: + return nil, true, fmt.Errorf("account %q not found in organization %q", account, organization) + default: + return nil, true, fmt.Errorf( + "account name %q is ambiguous in organization %q; use one of: %s", + account, + organization, + strings.Join(profileSelectorCandidates(matches), ", "), + ) } } - return corpNameMatch + + if profiles := profilesForCorpID(cfg, selector); len(profiles) > 0 { + return resolveOrganizationDefault(cfg, selector, selector, profiles) + } + + orgIDs := make(map[string]struct{}) + for i := range cfg.Profiles { + if strings.TrimSpace(cfg.Profiles[i].CorpName) != selector { + continue + } + orgIDs[strings.TrimSpace(cfg.Profiles[i].CorpID)] = struct{}{} + } + if len(orgIDs) == 1 { + for corpID := range orgIDs { + return resolveOrganizationDefault(cfg, corpID, selector, profilesForCorpID(cfg, corpID)) + } + } + if len(orgIDs) > 1 { + candidates := make([]string, 0, len(orgIDs)) + for corpID := range orgIDs { + candidates = append(candidates, corpID) + } + sort.Strings(candidates) + return nil, false, fmt.Errorf( + "organization name %q is ambiguous; use one of: %s", + selector, + strings.Join(candidates, ", "), + ) + } + + var nameMatches []*Profile + for i := range cfg.Profiles { + if strings.TrimSpace(cfg.Profiles[i].Name) != selector { + continue + } + nameMatches = append(nameMatches, &cfg.Profiles[i]) + } + switch len(nameMatches) { + case 1: + return nameMatches[0], true, nil + case 0: + return nil, false, fmt.Errorf("profile %q not found", selector) + default: + return nil, true, fmt.Errorf( + "profile name %q is ambiguous; use one of: %s", + selector, + strings.Join(profileSelectorCandidates(nameMatches), ", "), + ) + } } -func profileIndexByCorpID(cfg *ProfilesConfig, corpID string) int { +func resolveProfileDeletionSelection(cfg *ProfilesConfig, selector string) (*Profile, bool, error) { + if cfg == nil { + return nil, false, fmt.Errorf("profile %q not found", strings.TrimSpace(selector)) + } + selector = strings.TrimSpace(selector) + if selector == "" { + return nil, false, fmt.Errorf("profile selector is empty") + } + if _, _, compound := ParseIdentitySelector(selector); compound { + return resolveProfileSelection("", cfg, selector) + } + if profiles := profilesForCorpID(cfg, selector); len(profiles) > 0 { + return profiles[0], false, nil + } + corpID, err := resolveOrganizationCorpID(cfg, selector) + if err != nil { + return nil, false, err + } + if corpID != "" { + profiles := profilesForCorpID(cfg, corpID) + if len(profiles) > 0 { + return profiles[0], false, nil + } + } + + var nameMatches []*Profile + for i := range cfg.Profiles { + if strings.TrimSpace(cfg.Profiles[i].Name) == selector { + nameMatches = append(nameMatches, &cfg.Profiles[i]) + } + } + switch len(nameMatches) { + case 1: + return nameMatches[0], true, nil + case 0: + return nil, false, fmt.Errorf("profile %q not found", selector) + default: + return nil, true, fmt.Errorf( + "profile name %q is ambiguous; use one of: %s", + selector, + strings.Join(profileSelectorCandidates(nameMatches), ", "), + ) + } +} + +func resolveOrganizationCorpID(cfg *ProfilesConfig, selector string) (string, error) { + if cfg == nil { + return "", nil + } + selector = strings.TrimSpace(selector) + if selector == "" { + return "", nil + } + if len(profilesForCorpID(cfg, selector)) > 0 { + return selector, nil + } + orgIDs := make(map[string]struct{}) + for i := range cfg.Profiles { + if strings.TrimSpace(cfg.Profiles[i].CorpName) != selector { + continue + } + orgIDs[strings.TrimSpace(cfg.Profiles[i].CorpID)] = struct{}{} + } + if len(orgIDs) == 0 { + return "", nil + } + if len(orgIDs) == 1 { + for corpID := range orgIDs { + return corpID, nil + } + } + candidates := make([]string, 0, len(orgIDs)) + for corpID := range orgIDs { + candidates = append(candidates, corpID) + } + sort.Strings(candidates) + return "", fmt.Errorf( + "organization name %q is ambiguous; use one of: %s", + selector, + strings.Join(candidates, ", "), + ) +} + +func resolveOrganizationDefault(cfg *ProfilesConfig, corpID, displaySelector string, profiles []*Profile) (*Profile, bool, error) { + if len(profiles) == 0 { + return nil, false, fmt.Errorf("organization %q not found", displaySelector) + } + if exact := exactProfileSelectorForCorp(cfg, corpID, cfg.OrgCurrentProfiles[corpID]); exact != "" { + selectedCorpID, userID, _ := ParseIdentitySelector(exact) + if p := findExactProfile(cfg, selectedCorpID, userID); p != nil { + return p, false, nil + } + } + if len(profiles) == 1 { + return profiles[0], false, nil + } + return nil, false, fmt.Errorf( + "organization %q has multiple accounts and no current account; use one of: %s", + displaySelector, + strings.Join(profileSelectorCandidates(profiles), ", "), + ) +} + +func profileSelectorCandidates(profiles []*Profile) []string { + candidates := make([]string, 0, len(profiles)) + for _, p := range profiles { + if p == nil { + continue + } + candidates = append(candidates, ProfileSelector(*p)) + } + sort.Strings(candidates) + return candidates +} + +func exactProfileSelectorForCorp(cfg *ProfilesConfig, corpID, selector string) string { + selectedCorpID, userID, exact := ParseIdentitySelector(selector) + if !exact || strings.TrimSpace(selectedCorpID) != strings.TrimSpace(corpID) { + return "" + } + if p := findExactProfile(cfg, selectedCorpID, userID); p != nil { + return ProfileSelector(*p) + } + return "" +} + +func canonicalStoredSelector(cfg *ProfilesConfig, selector string) string { + selector = strings.TrimSpace(selector) + if selector == "" { + return "" + } + if corpID, userID, exact := ParseIdentitySelector(selector); exact { + if p := findExactProfile(cfg, corpID, userID); p != nil { + return ProfileSelector(*p) + } + return "" + } + if profiles := profilesForCorpID(cfg, selector); len(profiles) > 0 { + if exact := exactProfileSelectorForCorp(cfg, selector, cfg.OrgCurrentProfiles[selector]); exact != "" { + return exact + } + if len(profiles) == 1 { + return ProfileSelector(*profiles[0]) + } + return "" + } + p, _, err := resolveProfileSelection("", cfg, selector) + if err != nil || p == nil { + return "" + } + return ProfileSelector(*p) +} + +func setOrgCurrentProfile(cfg *ProfilesConfig, corpID, selector string) { + if cfg == nil { + return + } + exact := exactProfileSelectorForCorp(cfg, corpID, selector) + if exact == "" { + return + } + if cfg.OrgCurrentProfiles == nil { + cfg.OrgCurrentProfiles = make(map[string]string) + } + cfg.OrgCurrentProfiles[strings.TrimSpace(corpID)] = exact +} + +func uniqueProfileCorpIDs(cfg *ProfilesConfig) []string { + if cfg == nil { + return nil + } + seen := make(map[string]bool) + result := make([]string, 0) + for _, p := range cfg.Profiles { + corpID := strings.TrimSpace(p.CorpID) + if corpID == "" || seen[corpID] { + continue + } + seen[corpID] = true + result = append(result, corpID) + } + return result +} + +func profileIndexByIdentity(cfg *ProfilesConfig, corpID, userID string) int { if cfg == nil { return -1 } for i := range cfg.Profiles { - if cfg.Profiles[i].CorpID == corpID { + if sameProfileIdentity(cfg.Profiles[i].CorpID, cfg.Profiles[i].UserID, corpID, userID) { return i } } return -1 } -func firstProfileCorpID(cfg *ProfilesConfig) string { - if cfg == nil || len(cfg.Profiles) == 0 { - return "" +func legacyProfileIndexByCorpID(cfg *ProfilesConfig, corpID string) int { + if cfg == nil { + return -1 } - return cfg.Profiles[0].CorpID + match := -1 + for i := range cfg.Profiles { + if strings.TrimSpace(cfg.Profiles[i].CorpID) != strings.TrimSpace(corpID) || strings.TrimSpace(cfg.Profiles[i].UserID) != "" { + continue + } + if match >= 0 { + return -1 + } + match = i + } + return match } -func touchProfile(cfg *ProfilesConfig, corpID string) { - if p := findProfile(cfg, corpID); p != nil { - now := time.Now().Format(time.RFC3339) - p.LastUsedAt = now - p.UpdatedAt = now +func findExactProfile(cfg *ProfilesConfig, corpID, userID string) *Profile { + idx := profileIndexByIdentity(cfg, corpID, userID) + if idx < 0 { + return nil } + return &cfg.Profiles[idx] } -func timeOrRFC3339(t time.Time) string { - if t.IsZero() { - return "" +func profilesForCorpID(cfg *ProfilesConfig, corpID string) []*Profile { + if cfg == nil { + return nil } - return t.Format(time.RFC3339) + corpID = strings.TrimSpace(corpID) + result := make([]*Profile, 0) + for i := range cfg.Profiles { + if strings.TrimSpace(cfg.Profiles[i].CorpID) == corpID { + result = append(result, &cfg.Profiles[i]) + } + } + return result +} + +func profileSelectorReferenceExists(cfg *ProfilesConfig, selector string) bool { + if cfg == nil { + return false + } + selector = strings.TrimSpace(selector) + if selector == "" { + return false + } + if corpID, userID, exact := ParseIdentitySelector(selector); exact { + if findExactProfile(cfg, corpID, userID) != nil { + return true + } + for i := range cfg.Profiles { + p := &cfg.Profiles[i] + orgMatches := strings.TrimSpace(p.CorpID) == corpID || strings.TrimSpace(p.CorpName) == corpID + accountMatches := strings.TrimSpace(p.UserID) == userID || strings.TrimSpace(p.UserName) == userID + if orgMatches && accountMatches { + return true + } + } + return false + } + if len(profilesForCorpID(cfg, selector)) > 0 { + return true + } + for i := range cfg.Profiles { + if strings.TrimSpace(cfg.Profiles[i].CorpName) == selector || + strings.TrimSpace(cfg.Profiles[i].Name) == selector { + return true + } + } + return false +} + +func profileIdentityKey(corpID, userID string) string { + return strings.TrimSpace(corpID) + "\x00" + strings.TrimSpace(userID) +} + +func sameProfileIdentity(leftCorpID, leftUserID, rightCorpID, rightUserID string) bool { + return profileIdentityKey(leftCorpID, leftUserID) == profileIdentityKey(rightCorpID, rightUserID) } func shortCorpID(corpID string) string { @@ -699,3 +1346,11 @@ func shortCorpID(corpID string) string { } return corpID[len(corpID)-8:] } + +func shortProfileID(value string) string { + value = strings.TrimSpace(value) + if len(value) <= 8 { + return value + } + return value[len(value)-8:] +} diff --git a/internal/auth/token.go b/internal/auth/token.go index 3e41e816..16bcf664 100644 --- a/internal/auth/token.go +++ b/internal/auth/token.go @@ -28,6 +28,7 @@ import ( "github.com/google/uuid" + "github.com/DingTalk-Real-AI/dingtalk-workspace-cli/internal/keychain" "github.com/DingTalk-Real-AI/dingtalk-workspace-cli/pkg/edition" ) @@ -40,30 +41,41 @@ var ( tokenRemove = os.Remove tokenGlob = filepath.Glob tokenSaveKeychainForCorpID = SaveTokenDataKeychainForCorpID + tokenSaveKeychainForIdentity = SaveTokenDataKeychainForIdentity tokenSaveKeychain = SaveTokenDataKeychain tokenLoadKeychainForCorpID = LoadTokenDataKeychainForCorpID + tokenLoadKeychainIdentity = LoadTokenDataKeychainForIdentity tokenLoadKeychain = LoadTokenDataKeychain tokenKeychainExists = TokenDataExistsKeychain tokenDeleteKeychainForCorpID = DeleteTokenDataKeychainForCorpID + tokenDeleteKeychainIdentity = DeleteTokenDataKeychainForIdentity tokenDeleteKeychain = DeleteTokenDataKeychain + tokenRemoveAuthTokenEntries = keychain.RemoveAuthTokenEntries tokenLoadSecure = LoadSecureTokenData tokenDeleteSecure = DeleteSecureData - tokenResolveProfile = resolveProfileForLoad - tokenUpsertProfile = upsertProfileFromTokenWithCurrentLocked - tokenRemoveProfile = removeProfileLocked - tokenSyncLegacyMirror = syncLegacyTokenMirrorLocked - tokenLoadProfiles = LoadProfiles - tokenWriteMarker = WriteTokenMarker - tokenDeleteMarker = DeleteTokenMarker - tokenParseURL = url.Parse - tokenNewRequest = http.NewRequestWithContext - tokenDefaultConfigDir = getDefaultConfigDir - tokenLoadData = LoadTokenData - tokenRevokeURL = GetRevokeTokenURL - tokenLogoutURL = LogoutURL - tokenLogoutContinueURL = LogoutContinueURL - tokenLogoutHTTPClient = &http.Client{Timeout: 10 * time.Second, CheckRedirect: func(*http.Request, []*http.Request) error { return http.ErrUseLastResponse }} - tokenRevokeHTTPClient = &http.Client{Timeout: 10 * time.Second} + tokenResolveProfile = func(configDir, selector string) (*Profile, error) { + profile, _, err := resolveProfileForLoadLocked(configDir, selector) + return profile, err + } + tokenResolveDeletion = resolveProfileDeletionSelection + tokenResolveSelection = resolveProfileSelection + tokenUpsertProfile = upsertProfileFromTokenWithCurrentLocked + tokenRemoveProfile = removeProfileLocked + tokenSyncLegacyMirror = syncLegacyTokenMirrorLocked + tokenSyncOrganizationMirror = syncOrganizationTokenMirrorForProfile + tokenLoadProfiles = LoadProfiles + tokenWriteMarker = WriteTokenMarker + tokenDeleteMarker = DeleteTokenMarker + tokenParseURL = url.Parse + tokenNewRequest = http.NewRequestWithContext + tokenDefaultConfigDir = getDefaultConfigDir + tokenLoadData = LoadTokenData + tokenRevokeURL = GetRevokeTokenURL + tokenMCPBaseURL = GetMCPBaseURL + tokenLogoutURL = LogoutURL + tokenLogoutContinueURL = LogoutContinueURL + tokenLogoutHTTPClient = &http.Client{Timeout: 10 * time.Second, CheckRedirect: func(*http.Request, []*http.Request) error { return http.ErrUseLastResponse }} + tokenRevokeHTTPClient = &http.Client{Timeout: 10 * time.Second} ) // TokenData holds the OAuth token set persisted to disk. @@ -158,13 +170,40 @@ func saveTokenDataLocked(configDir string, data *TokenData) error { return saveTokenViaHook(h, configDir, data) } if data != nil && strings.TrimSpace(data.CorpID) != "" { - if err := tokenSaveKeychainForCorpID(data.CorpID, data); err != nil { - return err + corpID := strings.TrimSpace(data.CorpID) + userID := strings.TrimSpace(data.UserID) + if userID != "" { + if err := tokenSaveKeychainForIdentity(corpID, userID, data); err != nil { + return 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) + } + } } - makeCurrent := strings.TrimSpace(RuntimeProfile()) == "" + runtimeSelector := strings.TrimSpace(RuntimeProfile()) + makeCurrent := runtimeSelector == "" if err := tokenUpsertProfile(configDir, data, makeCurrent); err != nil { return 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 + } + } if makeCurrent { if err := tokenSaveKeychain(data); err != nil { return err @@ -213,13 +252,23 @@ func LoadTokenDataForProfile(configDir, profile string) (*TokenData, error) { return &td, nil } + var result *TokenData + err := withProfilesLock(configDir, func() error { + var loadErr error + result, loadErr = loadTokenDataForProfileLocked(configDir, profile) + return loadErr + }) + return result, err +} + +func loadTokenDataForProfileLocked(configDir, profile string) (*TokenData, error) { // Default: keychain with legacy .data migration selected, err := tokenResolveProfile(configDir, profile) if err != nil { return nil, err } if selected != nil { - data, err := tokenLoadKeychainForCorpID(selected.CorpID) + data, err := tokenLoadProfileIdentity(*selected) if err == nil { return data, nil } @@ -231,13 +280,21 @@ func LoadTokenDataForProfile(configDir, profile string) (*TokenData, error) { // the SAME org; otherwise surface the error instead of silently acting // as a different organization (the legacy mirror may have drifted). if legacy, lerr := tokenLoadKeychain(); lerr == nil && legacy != nil && - strings.TrimSpace(legacy.CorpID) == strings.TrimSpace(selected.CorpID) { + strings.TrimSpace(legacy.CorpID) == strings.TrimSpace(selected.CorpID) && + (strings.TrimSpace(selected.UserID) == "" || strings.TrimSpace(legacy.UserID) == strings.TrimSpace(selected.UserID)) { return legacy, nil } else if lerr != nil && !errors.Is(lerr, ErrTokenDataNotFound) { return nil, lerr } return nil, err } + cfg, err := tokenLoadProfiles(configDir) + if err != nil { + return nil, err + } + if cfg != nil && cfg.Version >= profilesVersion { + return nil, ErrTokenDataNotFound + } if tokenKeychainExists() { return tokenLoadKeychain() } @@ -253,6 +310,36 @@ func LoadTokenDataForProfile(configDir, profile string) (*TokenData, error) { return data, nil } +func tokenLoadProfileIdentity(profile Profile) (*TokenData, error) { + if strings.TrimSpace(profile.UserID) == "" { + return tokenLoadKeychainForCorpID(profile.CorpID) + } + data, err := tokenLoadKeychainIdentity(profile.CorpID, profile.UserID) + if err == nil { + return data, nil + } + if !errors.Is(err, ErrTokenDataNotFound) { + return nil, err + } + orgData, orgErr := tokenLoadKeychainForCorpID(profile.CorpID) + if orgErr != nil { + if errors.Is(orgErr, ErrTokenDataNotFound) { + return nil, err + } + return nil, orgErr + } + if strings.TrimSpace(orgData.UserID) == "" { + return nil, fmt.Errorf("organization token mirror for corpId %q has no userId; cannot use it for profile %q", profile.CorpID, ProfileSelector(profile)) + } + if strings.TrimSpace(orgData.UserID) != strings.TrimSpace(profile.UserID) { + return nil, err + } + if saveErr := tokenSaveKeychainForIdentity(profile.CorpID, profile.UserID, orgData); saveErr != nil { + return nil, saveErr + } + return orgData, nil +} + // DeleteTokenData removes token data. When an edition hook (DeleteToken) is // registered, it delegates entirely to the hook; otherwise it falls back // to keychain + legacy cleanup. @@ -275,25 +362,79 @@ func DeleteTokenDataForProfile(configDir, profile string) error { } func deleteTokenDataForProfileLocked(configDir, profile string) error { - selected, err := tokenResolveProfile(configDir, profile) + if err := profilesEnsureMigration(configDir); err != nil { + return err + } + cfg, err := tokenLoadProfiles(configDir) if err != nil { return err } + effectiveSelector := strings.TrimSpace(profile) + if effectiveSelector == "" { + effectiveSelector = strings.TrimSpace(cfg.CurrentProfile) + } + var selected *Profile + exact := false + if effectiveSelector != "" { + selected, exact, err = tokenResolveDeletion(cfg, effectiveSelector) + if err != nil { + return err + } + } if selected != nil { - keychainErr := tokenDeleteKeychainForCorpID(selected.CorpID) - _, removeErr := tokenRemoveProfile(configDir, selected.CorpID) - legacyErr := tokenSyncLegacyMirror(configDir) - secureErr := tokenDeleteSecure(configDir) - if keychainErr != nil { - return keychainErr + removed := *selected + if exact { + 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 removeErr != nil { - return removeErr + if err := tokenSyncLegacyMirror(configDir); err != nil { + return err } - if legacyErr != nil { - return legacyErr - } - return secureErr + return tokenDeleteSecure(configDir) } keychainErr := tokenDeleteKeychain() @@ -315,14 +456,10 @@ func DeleteAllTokenData(configDir string) error { } return withProfilesLock(configDir, func() error { var firstErr error - // Best-effort: even if profiles.json is unreadable, still clear every - // other slot so the user can always self-heal via auth reset / logout. - if cfg, err := tokenLoadProfiles(configDir); err == nil { - for _, profile := range cfg.Profiles { - if e := tokenDeleteKeychainForCorpID(profile.CorpID); e != nil && firstErr == nil { - firstErr = e - } - } + // Sweep the complete auth-token namespace so orphan identity slots that + // are not present in profiles.json cannot survive reset/logout --all. + if e := tokenRemoveAuthTokenEntries(keychain.Service); e != nil { + firstErr = e } if e := tokenRemove(ProfilesPath(configDir)); e != nil && !os.IsNotExist(e) && firstErr == nil { firstErr = e @@ -335,15 +472,19 @@ func DeleteAllTokenData(configDir string) error { } } } - if e := tokenDeleteKeychain(); e != nil && firstErr == nil { - firstErr = e - } if e := tokenDeleteSecure(configDir); e != nil && firstErr == nil { firstErr = e } if e := tokenDeleteMarker(configDir); e != nil && firstErr == nil { firstErr = e } + if firstErr != nil { + // Preserve an explicit v2 empty registry so any stale mirror that + // could not be removed is never imported on a later read. + if e := profilesSave(configDir, &ProfilesConfig{Version: profilesVersion}); e != nil { + return fmt.Errorf("%v; save logged-out profile tombstone: %w", firstErr, e) + } + } return firstErr }) } @@ -353,18 +494,42 @@ func DeleteAllTokenData(configDir string) error { // This should be called before deleting local token data. // The function is best-effort: errors are returned but callers may choose to ignore them. func RevokeTokenRemote(ctx context.Context) error { - // Use MCP revoke endpoint when clientID is from MCP - if IsClientIDFromMCP() { - return revokeTokenViaMCP(ctx) + tokenData, err := tokenLoadData(tokenDefaultConfigDir()) + if err != nil || tokenData == nil { + return nil } - // Direct mode: use DingTalk logout endpoint + // Historical token records may not have Source. Preserve the legacy + // process-wide MCP decision only for those records. + if strings.TrimSpace(tokenData.Source) == "" && IsClientIDFromMCP() { + copy := *tokenData + copy.Source = "mcp" + tokenData = © + } + return RevokeTokenRemoteForData(ctx, tokenData) +} + +// RevokeTokenRemoteForData revokes the supplied account token using the +// credential source and client ID persisted with that exact identity. +func RevokeTokenRemoteForData(ctx context.Context, tokenData *TokenData) error { + if tokenData == nil { + return nil + } + clientID := strings.TrimSpace(tokenData.ClientID) + if clientID == "" { + clientID = ClientID() + } + if strings.EqualFold(strings.TrimSpace(tokenData.Source), "mcp") { + return revokeTokenViaMCP(ctx, tokenData, clientID) + } + + // Direct mode: use DingTalk logout endpoint. logoutURL, err := tokenParseURL(tokenLogoutURL) if err != nil { return fmt.Errorf("parsing logout URL: %w", err) } q := logoutURL.Query() - q.Set("client_id", ClientID()) + q.Set("client_id", clientID) q.Set("continue", tokenLogoutContinueURL) logoutURL.RawQuery = q.Encode() @@ -388,20 +553,10 @@ func RevokeTokenRemote(ctx context.Context) error { } // revokeTokenViaMCP revokes token via MCP endpoint. -func revokeTokenViaMCP(ctx context.Context) error { - revokeURL := tokenRevokeURL() - if revokeURL == "" { - return nil // No revoke endpoint available - } - - // Load current token to get accessToken - tokenData, err := tokenLoadData(tokenDefaultConfigDir()) - if err != nil || tokenData == nil { - return nil // No token to revoke - } - +func revokeTokenViaMCP(ctx context.Context, tokenData *TokenData, clientID string) error { + revokeURL := tokenMCPBaseURL() + MCPRevokeTokenPath body := map[string]string{ - "clientId": ClientID(), + "clientId": clientID, "accessToken": tokenData.AccessToken, } bodyBytes, err := tokenJSONMarshal(body) diff --git a/internal/auth/token_preflight_test.go b/internal/auth/token_preflight_test.go index 4c7e72ce..5d495c78 100644 --- a/internal/auth/token_preflight_test.go +++ b/internal/auth/token_preflight_test.go @@ -92,7 +92,7 @@ func TestLoadTokenDataFallsBackToLegacyOnlyWhenCurrentSlotIsMissing(t *testing.T } } -func TestLoadTokenDataDoesNotHideUnreadableCurrentSlotWithLegacyFallback(t *testing.T) { +func TestLoadTokenDataUsesIdentitySlotWhenOrganizationMirrorIsUnreadable(t *testing.T) { cleanupKeychain(t) t.Setenv(keychain.DisableKeychainEnv, "1") configDir := t.TempDir() @@ -106,11 +106,11 @@ func TestLoadTokenDataDoesNotHideUnreadableCurrentSlotWithLegacyFallback(t *test } loaded, err := LoadTokenData(configDir) - if err == nil { - t.Fatalf("LoadTokenData() = %#v, nil; want unreadable profile error", loaded) + if err != nil { + t.Fatalf("LoadTokenData() error = %v", err) } - if loaded != nil { - t.Fatalf("LoadTokenData() data = %#v, want nil", loaded) + if loaded == nil || loaded.AccessToken != data.AccessToken || loaded.UserID != data.UserID { + t.Fatalf("LoadTokenData() = %#v, want identity token %#v", loaded, data) } } @@ -150,6 +150,26 @@ func TestPreflightTokenPersistenceRejectsUnreadableProfileSlot(t *testing.T) { } } +func TestExactProfileRefreshIgnoresUnreadableOrgMirror(t *testing.T) { + cleanupKeychain(t) + t.Setenv(keychain.DisableKeychainEnv, "1") + configDir := t.TempDir() + data := testToken("at_exact_refresh", "corp_exact", "Exact Org") + data.UserID = "user_exact" + if err := SaveTokenData(configDir, data); err != nil { + t.Fatalf("SaveTokenData() error = %v", err) + } + if err := os.WriteFile(profileCiphertextPathForTest(data.CorpID), []byte("corrupt ciphertext"), 0o600); err != nil { + t.Fatalf("WriteFile(profile ciphertext) error = %v", err) + } + + SetRuntimeProfile("corp_exact:user_exact") + defer SetRuntimeProfile("") + if err := preflightTokenRefreshPersistence(data); err != nil { + t.Fatalf("preflightTokenRefreshPersistence(exact) error = %v", err) + } +} + func TestExchangeAuthCodePreflightsOrphanProfileCiphertextBeforeHTTP(t *testing.T) { cleanupKeychain(t) t.Setenv(keychain.DisableKeychainEnv, "1") @@ -374,3 +394,37 @@ func TestExchangeAuthCodeAllowsFirstLogin(t *testing.T) { t.Fatalf("HTTP calls = %d, want 1", got) } } + +func TestExchangeAuthCodeExplicitUIDSkipsIdentityOverride(t *testing.T) { + cleanupKeychain(t) + t.Setenv(keychain.DisableKeychainEnv, "1") + setPreflightTestCredentials(t) + configDir := t.TempDir() + + var identityCalls atomic.Int32 + provider := NewOAuthProvider(configDir, nil) + provider.IdentityEnricher = func(context.Context, *TokenData) error { + identityCalls.Add(1) + return errors.New("identity lookup should not run for explicit uid") + } + provider.httpClient = &http.Client{Transport: preflightRoundTripFunc(func(*http.Request) (*http.Response, error) { + return &http.Response{ + StatusCode: http.StatusOK, + Header: make(http.Header), + Body: io.NopCloser(strings.NewReader( + `{"accessToken":"new-access","refreshToken":"new-refresh","expiresIn":7200,"corpId":"corp_new"}`, + )), + }, nil + })} + + data, err := provider.ExchangeAuthCode(context.Background(), "new-code", "explicit-user") + if err != nil { + t.Fatalf("ExchangeAuthCode() error = %v", err) + } + if data.UserID != "explicit-user" { + t.Fatalf("ExchangeAuthCode() userId = %q, want explicit-user", data.UserID) + } + if got := identityCalls.Load(); got != 0 { + t.Fatalf("IdentityEnricher calls = %d, want 0", got) + } +} diff --git a/internal/auth/token_test.go b/internal/auth/token_test.go index 4651b1b8..be6c5adf 100644 --- a/internal/auth/token_test.go +++ b/internal/auth/token_test.go @@ -14,7 +14,10 @@ package auth import ( + "context" + "errors" "os" + "strings" "testing" "time" @@ -147,7 +150,9 @@ func TestMultiProfileSaveLoadAndSwitch(t *testing.T) { if err != nil { t.Fatalf("LoadProfiles() error = %v", err) } - if cfg.PrimaryProfile != "corp_a" || cfg.CurrentProfile != "corp_b" || cfg.PreviousProfile != "corp_a" { + if cfg.PrimaryProfile != "" || + cfg.CurrentProfile != "corp_b:user_corp_b" || + cfg.PreviousProfile != "corp_a:user_corp_a" { t.Fatalf("profile pointers = primary %q current %q previous %q", cfg.PrimaryProfile, cfg.CurrentProfile, cfg.PreviousProfile) } @@ -212,8 +217,8 @@ func TestRuntimeProfileOverrideDoesNotMutateCurrent(t *testing.T) { if err != nil { t.Fatalf("LoadProfiles() error = %v", err) } - if cfg.CurrentProfile != "corp_a" { - t.Fatalf("current profile = %q, want corp_a", cfg.CurrentProfile) + if cfg.CurrentProfile != "corp_a:user_corp_a" { + t.Fatalf("current profile = %q, want corp_a:user_corp_a", cfg.CurrentProfile) } loadedB, err := LoadTokenDataForProfile(configDir, "corp_b") if err != nil { @@ -258,12 +263,33 @@ func TestDeleteProfilePreservesOtherProfiles(t *testing.T) { if err != nil { t.Fatalf("LoadProfiles() error = %v", err) } - if len(cfg.Profiles) != 1 || cfg.CurrentProfile != "corp_a" { + if len(cfg.Profiles) != 1 || + cfg.CurrentProfile != "corp_a:user_corp_a" || + cfg.PreviousProfile != "" { t.Fatalf("profiles after delete = %#v", cfg) } } -func TestUpsertProfileFromTokenOverwritesSameCorp(t *testing.T) { +func TestDeleteMissingProfileDoesNotDeleteCurrentToken(t *testing.T) { + cleanupKeychain(t) + configDir := t.TempDir() + if err := SaveTokenData(configDir, testToken("at_current", "corp_current", "Current Org")); err != nil { + t.Fatalf("SaveTokenData() error = %v", err) + } + + if err := DeleteTokenDataForProfile(configDir, "corp_missing"); err == nil { + t.Fatal("DeleteTokenDataForProfile(missing) error = nil, want not found") + } + loaded, err := LoadTokenData(configDir) + if err != nil { + t.Fatalf("LoadTokenData() after missing delete error = %v", err) + } + if loaded.AccessToken != "at_current" { + t.Fatalf("current token after missing delete = %q, want at_current", loaded.AccessToken) + } +} + +func TestUpsertProfileFromTokenUpdatesSameIdentity(t *testing.T) { cleanupKeychain(t) configDir := t.TempDir() @@ -272,7 +298,7 @@ func TestUpsertProfileFromTokenOverwritesSameCorp(t *testing.T) { t.Fatalf("SaveTokenData(first) error = %v", err) } second := testToken("at_second", "corp_same", "新组织名") - second.UserID = "user_updated" + second.UserID = first.UserID second.UserName = "Updated User" second.ClientID = "client_updated" if err := SaveTokenData(configDir, second); err != nil { @@ -290,7 +316,7 @@ func TestUpsertProfileFromTokenOverwritesSameCorp(t *testing.T) { if profile.CorpName != "新组织名" { t.Fatalf("corpName = %q, want 新组织名", profile.CorpName) } - if profile.UserID != "user_updated" || profile.UserName != "Updated User" || profile.ClientID != "client_updated" { + if profile.UserID != first.UserID || profile.UserName != "Updated User" || profile.ClientID != "client_updated" { t.Fatalf("profile metadata was not overwritten: %#v", profile) } loaded, err := LoadTokenDataForProfile(configDir, "corp_same") @@ -302,6 +328,606 @@ func TestUpsertProfileFromTokenOverwritesSameCorp(t *testing.T) { } } +func TestSameCorpDifferentUsersAreStoredSeparately(t *testing.T) { + cleanupKeychain(t) + configDir := t.TempDir() + + first := testToken("at_first", "corp_same", "同一组织") + first.UserID = "user_1" + first.UserName = "账号一" + second := testToken("at_second", "corp_same", "同一组织") + second.UserID = "user_2" + second.UserName = "账号二" + + 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) + } + + cfg, err := LoadProfiles(configDir) + if err != nil { + t.Fatalf("LoadProfiles() error = %v", err) + } + if len(cfg.Profiles) != 2 { + t.Fatalf("profiles len = %d, want 2: %#v", len(cfg.Profiles), cfg.Profiles) + } + + loadedFirst, err := LoadTokenDataForProfile(configDir, "corp_same:user_1") + if err != nil { + t.Fatalf("LoadTokenDataForProfile(first identity) error = %v", err) + } + if loadedFirst.AccessToken != "at_first" { + t.Fatalf("first access token = %q, want at_first", loadedFirst.AccessToken) + } + loadedSecond, err := LoadTokenDataForProfile(configDir, "corp_same:user_2") + if err != nil { + t.Fatalf("LoadTokenDataForProfile(second identity) error = %v", err) + } + if loadedSecond.AccessToken != "at_second" { + t.Fatalf("second access token = %q, want at_second", loadedSecond.AccessToken) + } + + loadedOrg, err := LoadTokenDataForProfile(configDir, "corp_same") + if err != nil { + t.Fatalf("LoadTokenDataForProfile(org) error = %v", err) + } + if loadedOrg.UserID != "user_2" { + t.Fatalf("org current user = %q, want user_2", loadedOrg.UserID) + } +} + +func TestMigrationBuildsOrgCurrentFromMatchingCorpMirror(t *testing.T) { + cleanupKeychain(t) + configDir := t.TempDir() + + first := testToken("at_first", "corp_same", "同一组织") + first.UserID = "user_1" + second := testToken("at_second", "corp_same", "同一组织") + second.UserID = "user_2" + if err := SaveTokenDataKeychainForIdentity(first.CorpID, first.UserID, first); err != nil { + t.Fatalf("SaveTokenDataKeychainForIdentity(first) error = %v", err) + } + if err := SaveTokenDataKeychainForIdentity(second.CorpID, second.UserID, second); err != nil { + t.Fatalf("SaveTokenDataKeychainForIdentity(second) error = %v", err) + } + if err := SaveTokenDataKeychainForCorpID(second.CorpID, second); err != nil { + t.Fatalf("SaveTokenDataKeychainForCorpID(second) error = %v", err) + } + if err := SaveProfiles(configDir, &ProfilesConfig{ + Version: 1, + CurrentProfile: "corp_same", + Profiles: []Profile{ + {Name: "账号一", CorpID: "corp_same", CorpName: "同一组织", UserID: "user_1"}, + {Name: "账号二", CorpID: "corp_same", CorpName: "同一组织", UserID: "user_2"}, + }, + }); err != nil { + t.Fatalf("SaveProfiles() error = %v", err) + } + + if err := EnsureProfilesMigration(configDir); err != nil { + t.Fatalf("EnsureProfilesMigration() error = %v", err) + } + cfg, err := LoadProfiles(configDir) + if err != nil { + t.Fatalf("LoadProfiles() error = %v", err) + } + if got := cfg.OrgCurrentProfiles["corp_same"]; got != "corp_same:user_2" { + t.Fatalf("org current = %q, want corp_same:user_2", got) + } + if cfg.CurrentProfile != "corp_same:user_2" { + t.Fatalf("current profile = %q, want corp_same:user_2", cfg.CurrentProfile) + } +} + +func TestMigrationDoesNotGuessOrgCurrentForMultipleAccounts(t *testing.T) { + cleanupKeychain(t) + configDir := t.TempDir() + + for _, userID := range []string{"user_1", "user_2"} { + token := testToken("at_"+userID, "corp_same", "同一组织") + token.UserID = userID + if err := SaveTokenDataKeychainForIdentity(token.CorpID, token.UserID, token); err != nil { + t.Fatalf("SaveTokenDataKeychainForIdentity(%s) error = %v", userID, err) + } + } + raw := `{ + "version": 1, + "currentProfile": "corp_same", + "profiles": [ + {"name":"账号一","corpId":"corp_same","corpName":"同一组织","userId":"user_1"}, + {"name":"账号二","corpId":"corp_same","corpName":"同一组织","userId":"user_2"} + ] +}` + if err := os.MkdirAll(configDir, 0o700); err != nil { + t.Fatalf("MkdirAll() error = %v", err) + } + if err := os.WriteFile(ProfilesPath(configDir), []byte(raw), 0o600); err != nil { + t.Fatalf("WriteFile(profiles.json) error = %v", err) + } + + if err := EnsureProfilesMigration(configDir); err != nil { + t.Fatalf("EnsureProfilesMigration() error = %v", err) + } + cfg, err := LoadProfiles(configDir) + if err != nil { + t.Fatalf("LoadProfiles() error = %v", err) + } + if got := cfg.OrgCurrentProfiles["corp_same"]; got != "" { + t.Fatalf("org current = %q, want empty", got) + } + if _, err := LoadTokenDataForProfile(configDir, "corp_same"); err == nil { + t.Fatal("LoadTokenDataForProfile(corp) error = nil, want ambiguous default failure") + } +} + +func TestLoginStoresExactCurrentAndPreviousSelectors(t *testing.T) { + cleanupKeychain(t) + configDir := t.TempDir() + + first := testToken("at_first", "corp_first", "第一组织") + first.UserID = "user_1" + second := testToken("at_second", "corp_second", "第二组织") + second.UserID = "user_2" + 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) + } + + cfg, err := LoadProfiles(configDir) + if err != nil { + t.Fatalf("LoadProfiles() error = %v", err) + } + if cfg.CurrentProfile != "corp_second:user_2" { + t.Fatalf("current profile = %q, want corp_second:user_2", cfg.CurrentProfile) + } + if cfg.PreviousProfile != "corp_first:user_1" { + t.Fatalf("previous profile = %q, want corp_first:user_1", cfg.PreviousProfile) + } + if got := cfg.OrgCurrentProfiles["corp_first"]; got != "corp_first:user_1" { + t.Fatalf("first org current = %q, want corp_first:user_1", got) + } + if got := cfg.OrgCurrentProfiles["corp_second"]; got != "corp_second:user_2" { + t.Fatalf("second org current = %q, want corp_second:user_2", got) + } +} + +func TestProfileSelectorSupportsOrganizationAndAccountNames(t *testing.T) { + cleanupKeychain(t) + configDir := t.TempDir() + + first := testToken("at_first", "corp_same", "钉钉") + first.UserID = "user_1" + first.UserName = "孙博文" + second := testToken("at_second", "corp_same", "钉钉") + second.UserID = "user_2" + second.UserName = "钉三多" + 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) + } + + for _, selector := range []string{ + "corp_same:user_1", + "corp_same:孙博文", + "钉钉:user_1", + "钉钉:孙博文", + } { + loaded, err := LoadTokenDataForProfile(configDir, selector) + if err != nil { + t.Fatalf("LoadTokenDataForProfile(%q) error = %v", selector, err) + } + if loaded.UserID != "user_1" || loaded.AccessToken != "at_first" { + t.Fatalf("LoadTokenDataForProfile(%q) = %#v, want first account", selector, loaded) + } + } +} + +func TestProfileSelectorRejectsDuplicateAccountName(t *testing.T) { + cleanupKeychain(t) + configDir := t.TempDir() + + for _, userID := range []string{"user_1", "user_2"} { + token := testToken("at_"+userID, "corp_same", "钉钉") + token.UserID = userID + token.UserName = "同名用户" + if err := SaveTokenData(configDir, token); err != nil { + t.Fatalf("SaveTokenData(%s) error = %v", userID, err) + } + } + + _, err := LoadTokenDataForProfile(configDir, "钉钉:同名用户") + if err == nil { + t.Fatal("LoadTokenDataForProfile(duplicate userName) error = nil") + } + for _, want := range []string{"ambiguous", "corp_same:user_1", "corp_same:user_2"} { + if !strings.Contains(err.Error(), want) { + t.Fatalf("error %q does not contain %q", err, want) + } + } +} + +func TestProfileSelectorRejectsDuplicateOrganizationName(t *testing.T) { + cleanupKeychain(t) + configDir := t.TempDir() + + for _, corpID := range []string{"corp_1", "corp_2"} { + token := testToken("at_"+corpID, corpID, "同名组织") + token.UserID = "user_1" + token.UserName = "同一用户" + if err := SaveTokenData(configDir, token); err != nil { + t.Fatalf("SaveTokenData(%s) error = %v", corpID, err) + } + } + + _, err := LoadTokenDataForProfile(configDir, "同名组织:同一用户") + if err == nil { + t.Fatal("LoadTokenDataForProfile(duplicate corpName) error = nil") + } + for _, want := range []string{"ambiguous", "corp_1", "corp_2"} { + if !strings.Contains(err.Error(), want) { + t.Fatalf("error %q does not contain %q", err, want) + } + } +} + +func TestExactProfileSwitchPreservesPreviousIdentityAndUpdatesOrgCurrent(t *testing.T) { + cleanupKeychain(t) + configDir := t.TempDir() + + first := testToken("at_first", "corp_same", "同一组织") + first.UserID = "user_1" + second := testToken("at_second", "corp_same", "同一组织") + second.UserID = "user_2" + 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) + } + + selected, err := SetCurrentProfile(configDir, "corp_same:user_1") + if err != nil { + t.Fatalf("SetCurrentProfile(exact) error = %v", err) + } + if selected.UserID != "user_1" { + t.Fatalf("selected user = %q, want user_1", selected.UserID) + } + cfg, err := LoadProfiles(configDir) + if err != nil { + t.Fatalf("LoadProfiles() error = %v", err) + } + if cfg.CurrentProfile != "corp_same:user_1" { + t.Fatalf("current profile = %q, want corp_same:user_1", cfg.CurrentProfile) + } + if cfg.PreviousProfile != "corp_same:user_2" { + t.Fatalf("previous profile = %q, want corp_same:user_2", cfg.PreviousProfile) + } + loadedOrg, err := LoadTokenDataForProfile(configDir, "corp_same") + if err != nil { + t.Fatalf("LoadTokenDataForProfile(org) error = %v", err) + } + if loadedOrg.UserID != "user_1" { + t.Fatalf("org current user = %q, want user_1", loadedOrg.UserID) + } + previous, err := UsePreviousProfile(configDir) + if err != nil { + t.Fatalf("UsePreviousProfile() error = %v", err) + } + if previous.UserID != "user_2" { + t.Fatalf("previous user = %q, want user_2", previous.UserID) + } + cfg, err = LoadProfiles(configDir) + if err != nil { + t.Fatalf("LoadProfiles() after previous error = %v", err) + } + if cfg.CurrentProfile != "corp_same:user_2" || cfg.PreviousProfile != "corp_same:user_1" { + t.Fatalf("profile pointers after previous = %#v", cfg) + } +} + +func TestExactRuntimeRefreshDoesNotChangeOrgCurrent(t *testing.T) { + cleanupKeychain(t) + configDir := t.TempDir() + + first := testToken("at_first", "corp_same", "同一组织") + first.UserID = "user_1" + second := testToken("at_second", "corp_same", "同一组织") + second.UserID = "user_2" + 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) + } + + refreshed := *first + refreshed.AccessToken = "at_first_refreshed" + SetRuntimeProfile("corp_same:user_1") + if err := SaveTokenData(configDir, &refreshed); err != nil { + t.Fatalf("SaveTokenData(exact refresh) error = %v", err) + } + SetRuntimeProfile("") + + loadedOrg, err := LoadTokenDataForProfile(configDir, "corp_same") + if err != nil { + t.Fatalf("LoadTokenDataForProfile(org) error = %v", err) + } + if loadedOrg.UserID != "user_2" || loadedOrg.AccessToken != "at_second" { + t.Fatalf("org token changed after exact refresh: %#v", loadedOrg) + } + loadedFirst, err := LoadTokenDataForProfile(configDir, "corp_same:user_1") + if err != nil { + t.Fatalf("LoadTokenDataForProfile(first) error = %v", err) + } + if loadedFirst.AccessToken != "at_first_refreshed" { + t.Fatalf("exact refreshed token = %q, want at_first_refreshed", loadedFirst.AccessToken) + } +} + +func TestDeleteExactProfilePreservesOtherAccountInSameCorp(t *testing.T) { + cleanupKeychain(t) + configDir := t.TempDir() + + first := testToken("at_first", "corp_same", "同一组织") + first.UserID = "user_1" + second := testToken("at_second", "corp_same", "同一组织") + second.UserID = "user_2" + 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 := DeleteTokenDataForProfile(configDir, "corp_same:user_2"); err != nil { + t.Fatalf("DeleteTokenDataForProfile(exact current) error = %v", err) + } + if _, err := LoadTokenDataForProfile(configDir, "corp_same:user_2"); err == nil { + t.Fatal("deleted exact profile is still loadable") + } + loadedFirst, err := LoadTokenDataForProfile(configDir, "corp_same:user_1") + if err != nil { + t.Fatalf("LoadTokenDataForProfile(remaining) error = %v", err) + } + if loadedFirst.AccessToken != "at_first" { + t.Fatalf("remaining access token = %q, want at_first", loadedFirst.AccessToken) + } + loadedOrg, err := LoadTokenDataForProfile(configDir, "corp_same") + if err != nil { + t.Fatalf("LoadTokenDataForProfile(org) error = %v", err) + } + if loadedOrg.UserID != "user_1" { + t.Fatalf("org current user after delete = %q, want user_1", loadedOrg.UserID) + } +} + +func TestDeleteOrgCurrentLeavesMultipleAccountsUnselected(t *testing.T) { + cleanupKeychain(t) + configDir := t.TempDir() + + for _, userID := range []string{"user_1", "user_2", "user_3"} { + token := testToken("at_"+userID, "corp_same", "同一组织") + token.UserID = userID + token.UserName = "账号" + userID + if err := SaveTokenData(configDir, token); err != nil { + t.Fatalf("SaveTokenData(%s) error = %v", userID, err) + } + } + + if err := DeleteTokenDataForProfile(configDir, "corp_same:user_3"); err != nil { + t.Fatalf("DeleteTokenDataForProfile(current) error = %v", err) + } + cfg, err := LoadProfiles(configDir) + if err != nil { + t.Fatalf("LoadProfiles() error = %v", err) + } + if got := cfg.OrgCurrentProfiles["corp_same"]; got != "" { + t.Fatalf("org current = %q, want empty", got) + } + if _, err := LoadTokenDataForProfile(configDir, "corp_same"); err == nil { + t.Fatal("LoadTokenDataForProfile(corp) error = nil, want explicit account requirement") + } + + if err := DeleteTokenDataForProfile(configDir, "corp_same"); err != nil { + t.Fatalf("DeleteTokenDataForProfile(unselected org) error = %v", err) + } + cfg, err = LoadProfiles(configDir) + if err != nil { + t.Fatalf("LoadProfiles() after organization delete error = %v", err) + } + if len(cfg.Profiles) != 0 { + t.Fatalf("profiles after unselected organization delete = %#v, want empty", cfg.Profiles) + } +} + +func TestDeleteOrgProfileRemovesAllAccountsInCorp(t *testing.T) { + cleanupKeychain(t) + configDir := t.TempDir() + + first := testToken("at_first", "corp_same", "同一组织") + first.UserID = "user_1" + second := testToken("at_second", "corp_same", "同一组织") + second.UserID = "user_2" + 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 := DeleteTokenDataForProfile(configDir, "corp_same"); err != nil { + t.Fatalf("DeleteTokenDataForProfile(org) error = %v", err) + } + if TokenDataExistsKeychainForIdentity("corp_same", "user_1") || + TokenDataExistsKeychainForIdentity("corp_same", "user_2") || + TokenDataExistsKeychainForCorpID("corp_same") { + t.Fatal("organization logout left account token slots behind") + } + cfg, err := LoadProfiles(configDir) + if err != nil { + t.Fatalf("LoadProfiles() error = %v", err) + } + if len(cfg.Profiles) != 0 { + t.Fatalf("profiles after org logout = %#v, want empty", cfg.Profiles) + } +} + +func TestOAuthLoginEnrichesIdentityBeforePersistingSameCorpAccount(t *testing.T) { + cleanupKeychain(t) + configDir := t.TempDir() + + existing := testToken("at_existing", "corp_same", "同一组织") + existing.UserID = "user_1" + if err := SaveTokenData(configDir, existing); err != nil { + t.Fatalf("SaveTokenData(existing) error = %v", err) + } + + incoming := testToken("at_incoming", "corp_same", "同一组织") + incoming.UserID = "" + incoming.UserName = "" + provider := NewOAuthProvider(configDir, nil) + provider.IdentityEnricher = func(_ context.Context, data *TokenData) error { + loaded, err := LoadTokenDataForProfile(configDir, "corp_same") + if err != nil { + t.Fatalf("LoadTokenDataForProfile(existing during enrichment) error = %v", err) + } + if loaded.UserID != "user_1" || loaded.AccessToken != "at_existing" { + t.Fatalf("existing account was changed before identity enrichment: %#v", loaded) + } + data.UserID = "user_2" + data.UserName = "账号二" + return nil + } + + if err := provider.persistLoginToken(context.Background(), incoming); err != nil { + t.Fatalf("persistLoginToken() error = %v", err) + } + cfg, err := LoadProfiles(configDir) + if err != nil { + t.Fatalf("LoadProfiles() error = %v", err) + } + if len(cfg.Profiles) != 2 { + t.Fatalf("profiles len = %d, want 2: %#v", len(cfg.Profiles), cfg.Profiles) + } +} + +func TestOAuthLoginIdentityFailureLeavesExistingSameCorpAccountUntouched(t *testing.T) { + cleanupKeychain(t) + configDir := t.TempDir() + + existing := testToken("at_existing", "corp_same", "同一组织") + existing.UserID = "user_1" + if err := SaveTokenData(configDir, existing); err != nil { + t.Fatalf("SaveTokenData(existing) error = %v", err) + } + + incoming := testToken("at_incoming", "corp_same", "同一组织") + incoming.UserID = "" + provider := NewOAuthProvider(configDir, nil) + provider.IdentityEnricher = func(context.Context, *TokenData) error { + return errors.New("identity lookup failed") + } + if err := provider.persistLoginToken(context.Background(), incoming); err == nil { + t.Fatal("persistLoginToken() error = nil, want identity lookup failure") + } + + loaded, err := LoadTokenDataForProfile(configDir, "corp_same") + if err != nil { + t.Fatalf("LoadTokenDataForProfile(existing) error = %v", err) + } + if loaded.UserID != "user_1" || loaded.AccessToken != "at_existing" { + t.Fatalf("existing account changed after identity failure: %#v", loaded) + } + cfg, err := LoadProfiles(configDir) + if err != nil { + t.Fatalf("LoadProfiles() error = %v", err) + } + if len(cfg.Profiles) != 1 { + t.Fatalf("profiles len = %d, want 1", len(cfg.Profiles)) + } +} + +func TestOAuthLoginMissingIdentityCannotOverwriteLegacySameCorpAccount(t *testing.T) { + cleanupKeychain(t) + configDir := t.TempDir() + + existing := testToken("at_existing", "corp_same", "同一组织") + existing.UserID = "" + existing.UserName = "" + if err := SaveTokenData(configDir, existing); err != nil { + t.Fatalf("SaveTokenData(existing legacy) error = %v", err) + } + + incoming := testToken("at_incoming", "corp_same", "同一组织") + incoming.UserID = "" + incoming.UserName = "" + provider := NewOAuthProvider(configDir, nil) + provider.IdentityEnricher = func(context.Context, *TokenData) error { return nil } + if err := provider.persistLoginToken(context.Background(), incoming); err == nil { + t.Fatal("persistLoginToken() error = nil, want missing identity failure") + } + + loaded, err := LoadTokenDataForProfile(configDir, "corp_same") + if err != nil { + t.Fatalf("LoadTokenDataForProfile(existing) error = %v", err) + } + if loaded.AccessToken != "at_existing" { + t.Fatalf("legacy account changed after missing identity: %#v", loaded) + } +} + +func TestLoadProfilesKeepsDifferentIdentitiesInSameCorp(t *testing.T) { + configDir := t.TempDir() + raw := `{ + "version": 1, + "primaryProfile": "corp_same:user_1", + "currentProfile": "corp_same:user_2", + "previousProfile": "corp_same:user_1", + "profiles": [ + { + "name": "账号一", + "corpId": "corp_same", + "userId": "user_1" + }, + { + "name": "账号二", + "corpId": "corp_same", + "userId": "user_2" + }, + { + "name": "账号二重复记录", + "corpId": "corp_same", + "userId": "user_2" + } + ] +}` + if err := os.MkdirAll(configDir, 0o700); err != nil { + t.Fatalf("MkdirAll() error = %v", err) + } + if err := os.WriteFile(ProfilesPath(configDir), []byte(raw), 0o600); err != nil { + t.Fatalf("WriteFile(profiles.json) error = %v", err) + } + + cfg, err := LoadProfiles(configDir) + if err != nil { + t.Fatalf("LoadProfiles() error = %v", err) + } + if len(cfg.Profiles) != 2 { + t.Fatalf("profiles len = %d, want 2: %#v", len(cfg.Profiles), cfg.Profiles) + } + if cfg.PrimaryProfile != "corp_same:user_1" || + cfg.CurrentProfile != "corp_same:user_2" || + cfg.PreviousProfile != "corp_same:user_1" { + t.Fatalf("identity pointers were not preserved: %#v", cfg) + } +} + func TestUpsertProfileFromTokenPromotesCorpIDNameToCorpName(t *testing.T) { cleanupKeychain(t) configDir := t.TempDir() @@ -387,12 +1013,50 @@ func TestLegacyKeychainMigrationInitializesProfile(t *testing.T) { if err != nil { t.Fatalf("LoadProfiles() error = %v", err) } - if cfg.PrimaryProfile != "corp_legacy" || cfg.CurrentProfile != "corp_legacy" { + if cfg.PrimaryProfile != "" || + cfg.CurrentProfile != "corp_legacy:user_corp_legacy" || + cfg.OrgCurrentProfiles["corp_legacy"] != "corp_legacy:user_corp_legacy" { t.Fatalf("profile pointers after migration = %#v", cfg) } if !TokenDataExistsKeychainForCorpID("corp_legacy") { t.Fatal("corp-scoped token should exist after migration") } + if !TokenDataExistsKeychainForIdentity("corp_legacy", legacy.UserID) { + t.Fatal("identity-scoped token should exist after migration") + } +} + +func TestLegacyCorpScopedTokenMigratesToIdentitySlot(t *testing.T) { + cleanupKeychain(t) + configDir := t.TempDir() + legacy := testToken("at_legacy_corp", "corp_legacy", "Legacy Org") + if err := SaveTokenDataKeychainForCorpID(legacy.CorpID, legacy); err != nil { + t.Fatalf("SaveTokenDataKeychainForCorpID() error = %v", err) + } + if err := SaveProfiles(configDir, &ProfilesConfig{ + Version: 1, + PrimaryProfile: legacy.CorpID, + CurrentProfile: legacy.CorpID, + Profiles: []Profile{{ + Name: legacy.CorpName, + CorpID: legacy.CorpID, + CorpName: legacy.CorpName, + UserID: legacy.UserID, + UserName: legacy.UserName, + }}, + }); err != nil { + t.Fatalf("SaveProfiles() error = %v", err) + } + + if err := EnsureProfilesMigration(configDir); err != nil { + t.Fatalf("EnsureProfilesMigration() error = %v", err) + } + if !TokenDataExistsKeychainForIdentity(legacy.CorpID, legacy.UserID) { + t.Fatal("identity-scoped token should exist after corp-scoped migration") + } + if !TokenDataExistsKeychainForCorpID(legacy.CorpID) { + t.Fatal("legacy corp-scoped token should remain after migration") + } } func TestTokenDataExistsKeychain(t *testing.T) { @@ -420,6 +1084,223 @@ func TestTokenDataExistsKeychain(t *testing.T) { } } +func TestProfileReadPathsWaitForAuthLock(t *testing.T) { + cleanupKeychain(t) + configDir := t.TempDir() + + for _, tc := range []struct { + name string + read func() (*Profile, *TokenData, error) + }{ + { + name: "resolve profile", + read: func() (*Profile, *TokenData, error) { + profile, err := ResolveProfile(configDir, "corp_locked:user_locked") + return profile, nil, err + }, + }, + { + name: "load token", + read: func() (*Profile, *TokenData, error) { + data, err := LoadTokenDataForProfile(configDir, "corp_locked:user_locked") + return nil, data, err + }, + }, + } { + t.Run(tc.name, func(t *testing.T) { + _ = os.Remove(ProfilesPath(configDir)) + _ = DeleteTokenDataKeychainForIdentity("corp_locked", "user_locked") + + lock, err := AcquireDualLock(context.Background(), configDir) + if err != nil { + t.Fatalf("AcquireDualLock() error = %v", err) + } + + type result struct { + profile *Profile + data *TokenData + err error + } + done := make(chan result, 1) + go func() { + profile, data, err := tc.read() + done <- result{profile: profile, data: data, err: err} + }() + + select { + case got := <-done: + lock.Release() + t.Fatalf("read completed before auth lock release: %+v", got) + case <-time.After(100 * time.Millisecond): + } + + data := testToken("at_locked", "corp_locked", "Locked Org") + data.UserID = "user_locked" + if err := SaveTokenDataKeychainForIdentity(data.CorpID, data.UserID, data); err != nil { + lock.Release() + t.Fatalf("SaveTokenDataKeychainForIdentity() error = %v", err) + } + if err := SaveProfiles(configDir, &ProfilesConfig{ + Version: profilesVersion, + CurrentProfile: "corp_locked:user_locked", + OrgCurrentProfiles: map[string]string{ + "corp_locked": "corp_locked:user_locked", + }, + Profiles: []Profile{{ + Name: "Locked Org", + CorpID: "corp_locked", + CorpName: "Locked Org", + UserID: "user_locked", + }}, + }); err != nil { + lock.Release() + t.Fatalf("SaveProfiles() error = %v", err) + } + lock.Release() + + select { + case got := <-done: + if got.err != nil { + t.Fatalf("read after auth lock release error = %v", got.err) + } + if got.profile != nil && ProfileSelector(*got.profile) != "corp_locked:user_locked" { + t.Fatalf("ResolveProfile() = %#v", got.profile) + } + if got.data != nil && got.data.AccessToken != "at_locked" { + t.Fatalf("LoadTokenDataForProfile() = %#v", got.data) + } + case <-time.After(2 * time.Second): + t.Fatal("read remained blocked after auth lock release") + } + }) + } +} + +func TestV2EmptyProfilesDoesNotRestoreLegacyToken(t *testing.T) { + cleanupKeychain(t) + configDir := t.TempDir() + legacy := testToken("at_legacy_stale", "corp_stale", "Stale Org") + if err := SaveTokenDataKeychain(legacy); err != nil { + t.Fatalf("SaveTokenDataKeychain() error = %v", err) + } + if err := SaveProfiles(configDir, &ProfilesConfig{Version: profilesVersion}); err != nil { + t.Fatalf("SaveProfiles() error = %v", err) + } + + if _, err := LoadTokenData(configDir); !errors.Is(err, ErrTokenDataNotFound) { + t.Fatalf("LoadTokenData() error = %v, want ErrTokenDataNotFound", err) + } + cfg, err := LoadProfiles(configDir) + if err != nil { + t.Fatalf("LoadProfiles() error = %v", err) + } + if len(cfg.Profiles) != 0 { + t.Fatalf("v2 empty config restored stale profile: %#v", cfg.Profiles) + } +} + +func TestIdentityLoadRejectsOrganizationMirrorWithoutUserID(t *testing.T) { + cleanupKeychain(t) + configDir := t.TempDir() + if err := SaveProfiles(configDir, &ProfilesConfig{ + Version: profilesVersion, + CurrentProfile: "corp_same:user_1", + OrgCurrentProfiles: map[string]string{ + "corp_same": "corp_same:user_1", + }, + Profiles: []Profile{{ + Name: "账号一", + CorpID: "corp_same", + UserID: "user_1", + }}, + }); err != nil { + t.Fatalf("SaveProfiles() error = %v", err) + } + legacy := testToken("at_unknown_identity", "corp_same", "同一组织") + legacy.UserID = "" + if err := SaveTokenDataKeychainForCorpID("corp_same", legacy); err != nil { + t.Fatalf("SaveTokenDataKeychainForCorpID() error = %v", err) + } + + _, err := LoadTokenDataForProfile(configDir, "corp_same:user_1") + if err == nil || !strings.Contains(err.Error(), "has no userId") { + t.Fatalf("LoadTokenDataForProfile() error = %v, want missing mirror identity error", err) + } + if TokenDataExistsKeychainForIdentity("corp_same", "user_1") { + t.Fatal("organization mirror without userId was persisted as an identity token") + } +} + +func TestProfilesMigrationDoesNotOverwriteUnreadableIdentityToken(t *testing.T) { + cleanupKeychain(t) + configDir := t.TempDir() + data := testToken("at_org_mirror", "corp_same", "同一组织") + data.UserID = "user_1" + if err := SaveTokenDataKeychainForCorpID(data.CorpID, data); err != nil { + t.Fatalf("SaveTokenDataKeychainForCorpID() error = %v", err) + } + if err := keychain.Set(keychain.Service, TokenAccountForIdentity(data.CorpID, data.UserID), "{unreadable"); err != nil { + t.Fatalf("write unreadable identity slot: %v", err) + } + if err := SaveProfiles(configDir, &ProfilesConfig{ + Version: profilesVersion, + Profiles: []Profile{{ + Name: "账号一", + CorpID: data.CorpID, + UserID: data.UserID, + }}, + }); err != nil { + t.Fatalf("SaveProfiles() error = %v", err) + } + + err := EnsureProfilesMigration(configDir) + if err == nil || !strings.Contains(err.Error(), "parse token data") { + t.Fatalf("EnsureProfilesMigration() error = %v, want unreadable identity error", err) + } + raw, getErr := keychain.Get(keychain.Service, TokenAccountForIdentity(data.CorpID, data.UserID)) + if getErr != nil { + t.Fatalf("read identity slot after migration: %v", getErr) + } + if raw != "{unreadable" { + t.Fatalf("migration overwrote unreadable identity slot: %q", raw) + } +} + +func TestOAuthLoginRequiresIdentityForFirstProfile(t *testing.T) { + cleanupKeychain(t) + configDir := t.TempDir() + incoming := testToken("at_missing_identity", "corp_first", "First Org") + incoming.UserID = "" + incoming.UserName = "" + + provider := NewOAuthProvider(configDir, nil) + provider.IdentityEnricher = func(context.Context, *TokenData) error { return nil } + if err := provider.persistLoginToken(context.Background(), incoming); err == nil || + !strings.Contains(err.Error(), "userId is required") { + t.Fatalf("persistLoginToken() error = %v, want required userId", err) + } + if TokenDataExistsKeychain() || TokenDataExistsKeychainForCorpID("corp_first") { + t.Fatal("missing-identity login persisted token data") + } +} + +func TestDeleteAllTokenDataRemovesOrphanIdentitySlots(t *testing.T) { + cleanupKeychain(t) + configDir := t.TempDir() + orphan := testToken("at_orphan", "corp_orphan", "Orphan Org") + orphan.UserID = "user_orphan" + if err := SaveTokenDataKeychainForIdentity(orphan.CorpID, orphan.UserID, orphan); err != nil { + t.Fatalf("SaveTokenDataKeychainForIdentity() error = %v", err) + } + + if err := DeleteAllTokenData(configDir); err != nil { + t.Fatalf("DeleteAllTokenData() error = %v", err) + } + if TokenDataExistsKeychainForIdentity(orphan.CorpID, orphan.UserID) { + t.Fatal("DeleteAllTokenData() left orphan identity token behind") + } +} + func testToken(accessToken, corpID, corpName string) *TokenData { now := time.Now().UTC() return &TokenData{ diff --git a/internal/keychain/auth_entries_file.go b/internal/keychain/auth_entries_file.go index 30d6b2ff..87ffba24 100644 --- a/internal/keychain/auth_entries_file.go +++ b/internal/keychain/auth_entries_file.go @@ -61,3 +61,17 @@ func isAuthTokenCiphertextFile(name string) bool { return name == legacyName || (strings.HasPrefix(name, strings.TrimSuffix(legacyName, ".enc")+"_") && strings.HasSuffix(name, ".enc")) } + +func platformRemoveAuthTokenEntries(service string) error { + paths, err := authTokenCiphertextPaths(service) + if err != nil { + return err + } + var firstErr error + for _, path := range paths { + if err := os.Remove(path); err != nil && !os.IsNotExist(err) && firstErr == nil { + firstErr = fmt.Errorf("remove keychain entry %q: %w", filepath.Base(path), err) + } + } + return firstErr +} diff --git a/internal/keychain/auth_entries_windows.go b/internal/keychain/auth_entries_windows.go index 07b7930f..c41bf669 100644 --- a/internal/keychain/auth_entries_windows.go +++ b/internal/keychain/auth_entries_windows.go @@ -20,3 +20,7 @@ package keychain func platformValidateAuthTokenEntries(string) error { return nil } + +func platformRemoveAuthTokenEntries(service string) error { + return registryRemoveAuthTokenEntries(service) +} diff --git a/internal/keychain/keychain.go b/internal/keychain/keychain.go index 40d3ac14..7cbb7133 100644 --- a/internal/keychain/keychain.go +++ b/internal/keychain/keychain.go @@ -166,6 +166,13 @@ func ValidateAuthTokenEntries(service string) error { return platformValidateAuthTokenEntries(service) } +// RemoveAuthTokenEntries deletes the legacy token plus every organization and +// identity scoped auth-token entry for service. Other keychain accounts are +// preserved. +func RemoveAuthTokenEntries(service string) error { + return platformRemoveAuthTokenEntries(service) +} + // Exists checks if an entry exists in the keychain. func Exists(service, account string) bool { val, err := Get(service, account) diff --git a/internal/keychain/keychain_test.go b/internal/keychain/keychain_test.go index ba17194a..16ca3e12 100644 --- a/internal/keychain/keychain_test.go +++ b/internal/keychain/keychain_test.go @@ -146,6 +146,36 @@ func TestKeychainOverwrite(t *testing.T) { } } +func TestRemoveAuthTokenEntriesPreservesOtherAccounts(t *testing.T) { + service := "test-service-" + t.Name() + authAccounts := []string{ + AccountToken, + AccountToken + ":corp-a", + AccountToken + ":id:0123456789abcdef", + } + for _, account := range authAccounts { + if err := Set(service, account, "secret"); err != nil { + t.Fatalf("Set(%q) error = %v", account, err) + } + } + const unrelated = "app-secret:client-a" + if err := Set(service, unrelated, "preserve"); err != nil { + t.Fatalf("Set(unrelated) error = %v", err) + } + + if err := RemoveAuthTokenEntries(service); err != nil { + t.Fatalf("RemoveAuthTokenEntries() error = %v", err) + } + for _, account := range authAccounts { + if Exists(service, account) { + t.Fatalf("auth account %q still exists", account) + } + } + if got, err := Get(service, unrelated); err != nil || got != "preserve" { + t.Fatalf("unrelated account = %q, %v; want preserved", got, err) + } +} + func TestMigrationNoLegacyData(t *testing.T) { t.Parallel() diff --git a/internal/keychain/keychain_windows.go b/internal/keychain/keychain_windows.go index 0ee89e28..4b73099c 100644 --- a/internal/keychain/keychain_windows.go +++ b/internal/keychain/keychain_windows.go @@ -204,3 +204,31 @@ func registryRemove(service, account string) error { _ = k.DeleteValue(valueNameForAccount(account)) return nil } + +func registryRemoveAuthTokenEntries(service string) error { + keyPath := registryPathForService(service) + k, err := registry.OpenKey(registry.CURRENT_USER, keyPath, registry.QUERY_VALUE|registry.SET_VALUE) + if err != nil { + return nil + } + defer k.Close() + + names, err := k.ReadValueNames(-1) + if err != nil { + return fmt.Errorf("registry list values failed: %w", err) + } + for _, name := range names { + accountBytes, decodeErr := base64.RawURLEncoding.DecodeString(name) + if decodeErr != nil { + continue + } + account := string(accountBytes) + if account != AccountToken && !strings.HasPrefix(account, AccountToken+":") { + continue + } + if err := k.DeleteValue(name); err != nil { + return fmt.Errorf("registry delete auth token value failed: %w", err) + } + } + return nil +}