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 +}