Skip to content

Commit e0516b2

Browse files
ndycodendycode
authored
Merge pull request #338 from ndycode/audit/pr2-runtime-recovery-safety
fix runtime recovery token restore safety
2 parents 16af48b + d0bb32d commit e0516b2

6 files changed

Lines changed: 407 additions & 40 deletions

File tree

‎index.ts‎

Lines changed: 23 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -254,6 +254,7 @@ import {
254254
saveFlaggedAccounts,
255255
setStorageBackupEnabled,
256256
setStoragePath,
257+
withAccountAndFlaggedStorageTransaction,
257258
withAccountStorageTransaction,
258259
} from "./lib/storage.js";
259260
import {
@@ -398,6 +399,26 @@ export const OpenAIOAuthPlugin: Plugin = async ({ client }: PluginInput) => {
398399
findMatchingAccountIndex,
399400
modelFamilies: MODEL_FAMILIES,
400401
});
402+
const persistAccountPoolAndFlagged = async (
403+
results: TokenSuccessWithAccount[],
404+
flaggedStorage: Parameters<typeof saveFlaggedAccounts>[0],
405+
replaceAll: boolean = false,
406+
): Promise<void> =>
407+
withAccountAndFlaggedStorageTransaction(async (loadedStorage, persist) => {
408+
await persistAccountPoolResults({
409+
results,
410+
replaceAll,
411+
modelFamilies: MODEL_FAMILIES,
412+
withAccountStorageTransaction: async (handler) =>
413+
handler(loadedStorage, async (storage) => {
414+
await persist(storage, flaggedStorage);
415+
}),
416+
findMatchingAccountIndex,
417+
extractAccountId,
418+
extractAccountEmail,
419+
sanitizeEmail,
420+
});
421+
});
401422

402423
const accountManagerCacheInvalidationDeps = {
403424
setCachedAccountManager: (value: unknown) => {
@@ -2598,6 +2619,8 @@ export const OpenAIOAuthPlugin: Plugin = async ({ client }: PluginInput) => {
25982619
queuedRefresh,
25992620
resolveTokenSuccessAccount,
26002621
persistAccounts: persistAccountPool,
2622+
persistAccountsAndFlagged:
2623+
persistAccountPoolAndFlagged,
26012624
invalidateAccountManagerCache: () =>
26022625
invalidateRuntimeAccountManagerCache(
26032626
accountManagerCacheInvalidationDeps,

‎lib/runtime/account-check.ts‎

Lines changed: 23 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -65,6 +65,13 @@ export async function runRuntimeAccountCheck(
6565
version: 1;
6666
accounts: FlaggedAccountMetadataV1[];
6767
}) => Promise<void>;
68+
persistAccountAndFlaggedStorage?: (
69+
accountStorage: AccountStorageV3,
70+
flaggedStorage: {
71+
version: 1;
72+
accounts: FlaggedAccountMetadataV1[];
73+
},
74+
) => Promise<void>;
6875
now?: () => number;
6976
showLine: (message: string) => void;
7077
},
@@ -145,14 +152,6 @@ export async function runRuntimeAccountCheck(
145152
) {
146153
accessToken = cached.accessToken;
147154
authDetail = "OK (Codex CLI cache)";
148-
149-
if (
150-
cached.refreshToken &&
151-
cached.refreshToken !== account.refreshToken
152-
) {
153-
account.refreshToken = cached.refreshToken;
154-
state.storageChanged = true;
155-
}
156155
if (
157156
cached.accessToken &&
158157
cached.accessToken !== account.accessToken
@@ -318,12 +317,23 @@ export async function runRuntimeAccountCheck(
318317
state.storageChanged = true;
319318
}
320319

321-
if (state.flaggedChanged) {
322-
await deps.saveFlaggedAccounts(state.flaggedStorage);
323-
}
324-
if (state.storageChanged) {
325-
await deps.saveAccounts(workingStorage);
320+
const persistAccountAndFlaggedStorage =
321+
deps.persistAccountAndFlaggedStorage;
322+
const canPersistTogether =
323+
state.flaggedChanged &&
324+
state.storageChanged &&
325+
typeof persistAccountAndFlaggedStorage === "function";
326+
if (canPersistTogether) {
327+
await persistAccountAndFlaggedStorage(workingStorage, state.flaggedStorage);
326328
deps.invalidateAccountManagerCache();
329+
} else {
330+
if (state.flaggedChanged) {
331+
await deps.saveFlaggedAccounts(state.flaggedStorage);
332+
}
333+
if (state.storageChanged) {
334+
await deps.saveAccounts(workingStorage);
335+
deps.invalidateAccountManagerCache();
336+
}
327337
}
328338

329339
deps.showLine("");

‎lib/runtime/verify-flagged.ts‎

Lines changed: 41 additions & 26 deletions
Original file line numberDiff line numberDiff line change
@@ -43,6 +43,14 @@ export async function verifyRuntimeFlaggedAccounts(deps: {
4343
results: Array<TokenSuccessWithAccount<SuccessfulAccountTokens>>,
4444
replaceAll?: boolean,
4545
) => Promise<void>;
46+
persistAccountsAndFlagged?: (
47+
results: Array<TokenSuccessWithAccount<SuccessfulAccountTokens>>,
48+
flaggedStorage: {
49+
version: 1;
50+
accounts: FlaggedAccountMetadataV1[];
51+
},
52+
replaceAll?: boolean,
53+
) => Promise<void>;
4654
invalidateAccountManagerCache: () => void;
4755
saveFlaggedAccounts: (storage: {
4856
version: 1;
@@ -78,26 +86,28 @@ export async function verifyRuntimeFlaggedAccounts(deps: {
7886
const refreshToken =
7987
typeof cached.refreshToken === "string" && cached.refreshToken.trim()
8088
? cached.refreshToken.trim()
81-
: flagged.refreshToken;
82-
const resolved = deps.resolveTokenSuccessAccount({
83-
type: "success",
84-
access: cached.accessToken,
85-
refresh: refreshToken,
86-
expires: cached.expiresAt,
87-
multiAccount: true,
88-
});
89-
if (!resolved.accountIdOverride && flagged.accountId) {
90-
resolved.accountIdOverride = flagged.accountId;
91-
resolved.accountIdSource = flagged.accountIdSource ?? "manual";
89+
: undefined;
90+
if (refreshToken) {
91+
const resolved = deps.resolveTokenSuccessAccount({
92+
type: "success",
93+
access: cached.accessToken,
94+
refresh: refreshToken,
95+
expires: cached.expiresAt,
96+
multiAccount: true,
97+
});
98+
if (!resolved.accountIdOverride && flagged.accountId) {
99+
resolved.accountIdOverride = flagged.accountId;
100+
resolved.accountIdSource = flagged.accountIdSource ?? "manual";
101+
}
102+
if (!resolved.accountLabel && flagged.accountLabel) {
103+
resolved.accountLabel = flagged.accountLabel;
104+
}
105+
state.restored.push(resolved);
106+
deps.showLine(
107+
`[${i + 1}/${flaggedStorage.accounts.length}] ${label}: RESTORED (Codex CLI cache)`,
108+
);
109+
continue;
92110
}
93-
if (!resolved.accountLabel && flagged.accountLabel) {
94-
resolved.accountLabel = flagged.accountLabel;
95-
}
96-
state.restored.push(resolved);
97-
deps.showLine(
98-
`[${i + 1}/${flaggedStorage.accounts.length}] ${label}: RESTORED (Codex CLI cache)`,
99-
);
100-
continue;
101111
}
102112

103113
const refreshResult = await deps.queuedRefresh(flagged.refreshToken);
@@ -134,16 +144,21 @@ export async function verifyRuntimeFlaggedAccounts(deps: {
134144
}
135145
}
136146

137-
if (state.restored.length > 0) {
138-
await deps.persistAccounts(state.restored, false);
147+
const nextFlaggedStorage = {
148+
version: 1 as const,
149+
accounts: state.remaining,
150+
};
151+
if (state.restored.length > 0 && deps.persistAccountsAndFlagged) {
152+
await deps.persistAccountsAndFlagged(state.restored, nextFlaggedStorage, false);
139153
deps.invalidateAccountManagerCache();
154+
} else {
155+
if (state.restored.length > 0) {
156+
await deps.persistAccounts(state.restored, false);
157+
deps.invalidateAccountManagerCache();
158+
}
159+
await deps.saveFlaggedAccounts(nextFlaggedStorage);
140160
}
141161

142-
await deps.saveFlaggedAccounts({
143-
version: 1,
144-
accounts: state.remaining,
145-
});
146-
147162
deps.showLine("");
148163
deps.showLine(
149164
`Results: ${state.restored.length} restored, ${state.remaining.length} still flagged`,

‎test/index.test.ts‎

Lines changed: 80 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -156,6 +156,7 @@ vi.mock("../lib/logger.js", () => ({
156156
logInfo: vi.fn(),
157157
logWarn: vi.fn(),
158158
logError: vi.fn(),
159+
maskEmail: vi.fn((email: string) => email),
159160
setCorrelationId: vi.fn(() => "test-correlation-id"),
160161
clearCorrelationId: vi.fn(),
161162
createLogger: vi.fn(() => ({
@@ -333,6 +334,7 @@ const withAccountStorageTransactionMock = vi.fn(
333334
return nextRun;
334335
},
335336
);
337+
const withAccountAndFlaggedStorageTransactionMock = vi.fn();
336338

337339
const syncCodexCliSelectionMock = vi.fn(async (_index: number) => {});
338340

@@ -343,6 +345,8 @@ vi.mock("../lib/storage.js", async () => {
343345
getStoragePath: () => "/mock/path/accounts.json",
344346
loadAccounts: loadAccountsMock,
345347
saveAccounts: saveAccountsMock,
348+
withAccountAndFlaggedStorageTransaction:
349+
withAccountAndFlaggedStorageTransactionMock,
346350
withAccountStorageTransaction: withAccountStorageTransactionMock,
347351
clearAccounts: clearAccountsMock,
348352
setStoragePath: vi.fn(),
@@ -776,6 +780,82 @@ describe("OpenAIOAuthPlugin", () => {
776780
expect.stringContaining("aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa"),
777781
);
778782
});
783+
784+
it("uses combined flagged persistence when verify-flagged restores from the login menu", async () => {
785+
const cliModule = await import("../lib/cli.js");
786+
const refreshQueueModule = await import("../lib/refresh-queue.js");
787+
const storageModule = await import("../lib/storage.js");
788+
const now = Date.now();
789+
const flaggedStorage = {
790+
version: 1 as const,
791+
accounts: [
792+
{
793+
refreshToken: "flagged-refresh",
794+
accountId: "flagged-account",
795+
email: "flagged@example.com",
796+
addedAt: now - 1_000,
797+
lastUsed: now - 1_000,
798+
flaggedAt: now - 5_000,
799+
},
800+
],
801+
};
802+
let loadFlaggedCallCount = 0;
803+
804+
vi.mocked(cliModule.promptLoginMode)
805+
.mockResolvedValueOnce({ mode: "verify-flagged" } as never)
806+
.mockResolvedValueOnce({ mode: "cancel" } as never);
807+
vi.mocked(refreshQueueModule.queuedRefresh).mockResolvedValueOnce({
808+
type: "success",
809+
access: "restored-access",
810+
refresh: "restored-refresh",
811+
expires: now + 60_000,
812+
});
813+
vi.mocked(storageModule.loadFlaggedAccounts).mockImplementation(async () => {
814+
loadFlaggedCallCount += 1;
815+
return loadFlaggedCallCount <= 2
816+
? flaggedStorage
817+
: { version: 1, accounts: [] };
818+
});
819+
withAccountAndFlaggedStorageTransactionMock.mockImplementationOnce(
820+
async (handler) =>
821+
handler(
822+
{
823+
version: 3,
824+
accounts: mockStorage.accounts.map((account) =>
825+
cloneMockAccount(account),
826+
),
827+
activeIndex: mockStorage.activeIndex,
828+
activeIndexByFamily: { ...mockStorage.activeIndexByFamily },
829+
},
830+
async (storage, nextFlaggedStorage) => {
831+
await saveAccountsMock(storage);
832+
await vi.mocked(storageModule.saveFlaggedAccounts)(
833+
nextFlaggedStorage,
834+
);
835+
},
836+
flaggedStorage,
837+
),
838+
);
839+
840+
const autoMethod = plugin.auth.methods[0] as unknown as {
841+
authorize: () => Promise<unknown>;
842+
};
843+
await autoMethod.authorize();
844+
845+
expect(withAccountAndFlaggedStorageTransactionMock).toHaveBeenCalledTimes(1);
846+
expect(mockStorage.accounts).toEqual(
847+
expect.arrayContaining([
848+
expect.objectContaining({
849+
refreshToken: "restored-refresh",
850+
accessToken: "restored-access",
851+
}),
852+
]),
853+
);
854+
expect(vi.mocked(storageModule.saveFlaggedAccounts)).toHaveBeenCalledWith({
855+
version: 1,
856+
accounts: [],
857+
});
858+
});
779859
});
780860

781861
describe("event handler", () => {

0 commit comments

Comments
 (0)