Skip to content
Closed
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
25 changes: 22 additions & 3 deletions src/oauth/index.ts
Original file line number Diff line number Diff line change
Expand Up @@ -322,6 +322,13 @@ export class OAuthReauthIdentityUnverifiedError extends Error {
}
}

class OAuthLoginSupersededError extends Error {
constructor() {
super("OAuth login was superseded before credential persistence");
this.name = "OAuthLoginSupersededError";
}
}

/** Project arbitrary OAuth failures onto the small, stable public error vocabulary. */
export function publicOAuthAuthenticationErrorMessage(error: unknown): string {
if (error instanceof OAuthMutationBusyError) {
Expand Down Expand Up @@ -1096,6 +1103,7 @@ interface RunLoginDeps {
settleKiroLoginTransaction?: typeof settleKiroLoginTransaction;
removeAccount?: typeof removeAccount;
setActiveAccount?: typeof setActiveAccount;
assertCurrentOwner?: () => void;
}

/** Roll back only accounts created by this forced login, preserving concurrent refreshes of others. */
Expand Down Expand Up @@ -1145,6 +1153,7 @@ export async function runLogin(
const cred: OAuthCredentials = rawCred.source ? rawCred : { ...rawCred, source: "oauth" };
const settleKiroTransaction = deps.settleKiroLoginTransaction ?? settleKiroLoginTransaction;
try {
deps.assertCurrentOwner?.();
// Validate the provider row before credential persistence. A namespace claimed during the
// credential write is handled again below before the latest row is re-upserted.
if (provider !== "chatgpt") {
Expand All @@ -1165,10 +1174,13 @@ export async function runLogin(
if (!identityMatches) {
throw new OAuthReauthIdentityMismatchError();
}
await (deps.saveAccountCredential ?? saveAccountCredential)(provider, opts.reauthAccountId, cred);
await (deps.saveAccountCredential ?? saveAccountCredential)(provider, opts.reauthAccountId, cred, {
assertBeforePersist: deps.assertCurrentOwner,
});
} else {
await (deps.saveCredential ?? saveCredential)(provider, cred, {
preserveIdentityless: opts?.forceLogin === true,
assertBeforePersist: deps.assertCurrentOwner,
});
}
if (provider !== "chatgpt") {
Expand Down Expand Up @@ -1235,6 +1247,7 @@ export async function runLogin(
*/
const loginState = new Map<string, { error?: string; done: boolean }>();
const loginAbort = new Map<string, AbortController>();
const kiroLoginSettling = new Set<string>();

/** Pending paste for a login in progress: either a waiter or a stashed early submission. */
interface ManualCodeSlot {
Expand Down Expand Up @@ -1403,13 +1416,14 @@ export async function startLoginFlow(
const def = OAUTH_PROVIDERS[provider];
if (!def) throw new UnsupportedOAuthProviderError(provider);
const existing = loginState.get(provider);
if (existing && !existing.done) {
if ((existing && !existing.done) || (provider === "kiro" && kiroLoginSettling.has(provider))) {
throw new Error(`A login for ${provider} is already in progress`);
}
clearManualCodeSlot(provider);
loginState.set(provider, { done: false });
const abort = new AbortController();
loginAbort.set(provider, abort);
if (provider === "kiro") kiroLoginSettling.add(provider);
return new Promise((resolve, reject) => {
let urlResolved = false;
const ctrl: OAuthController = {
Expand Down Expand Up @@ -1460,7 +1474,10 @@ export async function startLoginFlow(
};
// Background: runLogin persists the credential + provider entry to disk. The lifecycle hook
// lets a long-lived server config adopt that settled state before clients observe done=true.
void runLogin(provider, ctrl, opts).then(
const assertCurrentOwner = (): void => {
if (loginAbort.get(provider) !== abort) throw new OAuthLoginSupersededError();
};
void runLogin(provider, ctrl, opts, { assertCurrentOwner }).then(
() => settle(),
(e: unknown) => settle(e),
).catch((e: unknown) => {
Expand All @@ -1471,6 +1488,8 @@ export async function startLoginFlow(
const msg = publicOAuthAuthenticationErrorMessage(e);
loginState.set(provider, { done: true, error: msg });
if (!urlResolved) reject(e);
}).finally(() => {
if (provider === "kiro") kiroLoginSettling.delete(provider);
});
});
}
16 changes: 11 additions & 5 deletions src/oauth/store.ts
Original file line number Diff line number Diff line change
Expand Up @@ -465,10 +465,11 @@ function serializeMutation<T>(work: () => Promise<T>, retainedValues: readonly u
drainOAuthMutations();
return result;
}
export function mutateStore<T>(fn:(store:AuthStore)=>T|Promise<T>, retainedValues: readonly unknown[] = [], options?: { waitMs?: number }):Promise<T>{return serializeMutation(async()=>{const guard=await createOAuthFileLock({path:getAuthStoreLockPath(),staleAfterMs:30000}).acquire();try{
export function mutateStore<T>(fn:(store:AuthStore)=>T|Promise<T>, retainedValues: readonly unknown[] = [], options?: { waitMs?: number; assertBeforePersist?: () => void }):Promise<T>{return serializeMutation(async()=>{const guard=await createOAuthFileLock({path:getAuthStoreLockPath(),staleAfterMs:30000}).acquire();try{
const { store, hadLegacy } = loadAuthStoreInternal();
if (hadLegacy) backupLegacyOnce();
const result = await fn(store);
options?.assertBeforePersist?.();
persist(store);
return result;
}finally{guard.release();}}, retainedValues, options?.waitMs);
Expand All @@ -491,7 +492,7 @@ export function getCredential(provider: string): OAuthCredentials | null {
export async function saveCredential(
provider: string,
cred: OAuthCredentials,
opts: { preserveIdentityless?: boolean } = {},
opts: { preserveIdentityless?: boolean; assertBeforePersist?: () => void } = {},
): Promise<void> {
const safe = normalizeCredential(cred);
if (!safe) return;
Expand Down Expand Up @@ -542,7 +543,7 @@ export async function saveCredential(
set.accounts.push({ id, credential: safe, addedAt: Date.now() });
set.activeAccountId = id;
}
}, [provider, safe]);
}, [provider, safe], { assertBeforePersist: opts.assertBeforePersist });
}

/**
Expand Down Expand Up @@ -632,15 +633,20 @@ export function getAccountCredential(provider: string, accountId: string): OAuth
}

/** Persist a refreshed credential for a SPECIFIC account without touching activeAccountId. */
export async function saveAccountCredential(provider: string, accountId: string, cred: OAuthCredentials): Promise<void> {
export async function saveAccountCredential(
provider: string,
accountId: string,
cred: OAuthCredentials,
opts: { assertBeforePersist?: () => void } = {},
): Promise<void> {
const safe = normalizeCredential(cred);
if (!safe) return;
await mutateStore(store => {
const account = store[provider]?.accounts.find(a => a.id === accountId);
if (!account) return;
account.credential = safe;
delete account.needsReauth;
}, [provider, accountId, safe]);
}, [provider, accountId, safe], { assertBeforePersist: opts.assertBeforePersist });
}

export async function setActiveAccount(provider: string, accountId: string): Promise<boolean> {
Expand Down
107 changes: 107 additions & 0 deletions tests/oauth-public-surface.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -448,6 +448,113 @@ describe("legacy ChatGPT OAuth public-surface exclusion", () => {
}
});

test("a superseded OAuth flow cannot commit after its replacement owns the provider", async () => {
saveConfig(config());
const originalLogin = OAUTH_PROVIDERS.xai.login;
let loginCalls = 0;
OAUTH_PROVIDERS.xai.login = async (ctrl) => {
loginCalls += 1;
const call = loginCalls;
ctrl.onAuth({ url: `https://auth.example.test/${call}`, deviceCode: `flow-${call}` });
return {
access: `access-${call}`,
refresh: `refresh-${call}`,
accountId: `account-${call}`,
email: `account-${call}@example.test`,
expires: Date.now() + 60_000,
};
};

let releaseHead!: () => void;
let signalHeadStarted!: () => void;
const headStarted = new Promise<void>(resolve => { signalHeadStarted = resolve; });
const headGate = new Promise<void>(resolve => { releaseHead = resolve; });
const blockingMutation = oauthStore.mutateStore(async () => {
signalHeadStarted();
await headGate;
});

const waitForMutationCount = async (minimum: number): Promise<void> => {
for (let attempt = 0; attempt < 200; attempt += 1) {
if (oauthStore.oauthMutationTailSnapshot().active >= minimum) return;
await Bun.sleep(5);
}
throw new Error(`OAuth mutation queue did not reach ${minimum} active rows`);
};

try {
await headStarted;
await startLoginFlow("xai");
await waitForMutationCount(2);
expect(cancelLoginFlow("xai")).toBe(true);

await startLoginFlow("xai");
await waitForMutationCount(3);
releaseHead();
await blockingMutation;

const status = await waitForOAuthDone("xai");
expect(status).toMatchObject({ done: true, loggedIn: true });
expect(getCredential("xai")).toMatchObject({
access: "access-2",
accountId: "account-2",
});
expect(oauthStore.getAccountSet("xai")?.accounts.map(account => account.credential.accountId))
.toEqual(["account-2"]);
} finally {
releaseHead();
await blockingMutation.catch(() => {});
OAUTH_PROVIDERS.xai.login = originalLogin;
clearLoginState("xai");
}
});

test("Kiro does not start a replacement until the canceled external CLI flow settles", async () => {
saveConfig(config());
const originalLogin = OAUTH_PROVIDERS.kiro.login;
let loginCalls = 0;
OAUTH_PROVIDERS.kiro.login = async (ctrl) => {
loginCalls += 1;
const call = loginCalls;
ctrl.onAuth({ url: "", deviceCode: `kiro-flow-${call}` });
if (call === 1) {
await new Promise<never>((_, reject) => {
ctrl.signal.addEventListener("abort", () => reject(new Error("Kiro login cancelled")), { once: true });
});
}
return {
access: "kiro-replacement-access",
refresh: "kiro-replacement-refresh",
accountId: "kiro-replacement-account",
email: "kiro-replacement@example.test",
expires: Date.now() + 60_000,
};
};

try {
await startLoginFlow("kiro");
expect(cancelLoginFlow("kiro")).toBe(true);
await expect(startLoginFlow("kiro")).rejects.toThrow("A login for kiro is already in progress");

let replacement: Awaited<ReturnType<typeof startLoginFlow>> | undefined;
for (let attempt = 0; attempt < 200; attempt += 1) {
try {
replacement = await startLoginFlow("kiro");
break;
} catch (error) {
if (!(error instanceof Error) || !error.message.includes("already in progress")) throw error;
await Bun.sleep(5);
}
}
expect(replacement).toMatchObject({ deviceCode: "kiro-flow-2" });
expect(await waitForOAuthDone("kiro")).toMatchObject({ done: true, loggedIn: true });
expect(loginCalls).toBe(2);
} finally {
OAUTH_PROVIDERS.kiro.login = originalLogin;
clearLoginState("kiro");
}
});
Comment thread
coderabbitai[bot] marked this conversation as resolved.

test("management OAuth safely reconciles live config after a late namespace claim", async () => {
const liveConfig = config();
liveConfig.hostname = "0.0.0.0";
Expand Down
Loading