diff --git a/src/cliproxy/auth/oauth-handler.ts b/src/cliproxy/auth/oauth-handler.ts index f2aece66..fb50989e 100644 --- a/src/cliproxy/auth/oauth-handler.ts +++ b/src/cliproxy/auth/oauth-handler.ts @@ -11,6 +11,7 @@ */ import * as fs from 'fs'; +import * as path from 'path'; import { fail, info, warn, color, ok } from '../../utils/ui'; import { ensureCLIProxyBinary } from '../binary-manager'; import { generateConfig } from '../config-generator'; @@ -46,7 +47,12 @@ import { normalizeKiroIDCFlow, } from './auth-types'; import { isHeadlessEnvironment, killProcessOnPort, showStep } from './environment-detector'; -import { getProviderTokenDir, isAuthenticated, registerAccountFromToken } from './token-manager'; +import { + getProviderTokenDir, + isAuthenticated, + isTokenFileForProvider, + registerAccountFromToken, +} from './token-manager'; import { executeOAuthProcess } from './oauth-process'; import { importKiroToken } from './kiro-import'; import { @@ -71,6 +77,12 @@ interface PasteCallbackStartData { } const PASTE_CALLBACK_AUTH_URL_POLL_INTERVAL_MS = 3000; +const POLLED_AUTH_LOCAL_TOKEN_GRACE_MS = 15 * 1000; + +type ProviderTokenSnapshot = { + file: string; + mtimeMs: number; +}; export async function requestPasteCallbackStart( provider: CLIProxyProvider, @@ -122,6 +134,134 @@ function sleep(ms: number): Promise { return new Promise((resolve) => setTimeout(resolve, ms)); } +function parseAuthUrlState(url: string | null | undefined): string | null { + if (!url) { + return null; + } + + try { + return new URL(url).searchParams.get('state'); + } catch { + return null; + } +} + +function listProviderTokenSnapshots( + provider: CLIProxyProvider, + tokenDir: string +): ProviderTokenSnapshot[] { + if (!fs.existsSync(tokenDir)) { + return []; + } + + return fs + .readdirSync(tokenDir) + .filter((file) => file.endsWith('.json')) + .map((file): ProviderTokenSnapshot | null => { + const filePath = path.join(tokenDir, file); + if (!isTokenFileForProvider(filePath, provider)) { + return null; + } + + return { + file, + mtimeMs: fs.statSync(filePath).mtimeMs, + }; + }) + .filter((snapshot): snapshot is ProviderTokenSnapshot => snapshot !== null) + .sort((left, right) => right.mtimeMs - left.mtimeMs); +} + +export function findNewTokenSnapshotForManualAuth( + provider: CLIProxyProvider, + tokenDir: string, + knownTokenFiles: ProviderTokenSnapshot[], + expectedAccountId?: string +): ProviderTokenSnapshot | null { + const knownTokenMtimes = new Map( + knownTokenFiles.map((snapshot) => [snapshot.file, snapshot.mtimeMs]) + ); + + return ( + listProviderTokenSnapshots(provider, tokenDir).find((snapshot) => { + const knownMtime = knownTokenMtimes.get(snapshot.file); + if (knownMtime === undefined) { + return true; + } + + if (!expectedAccountId) { + return false; + } + + return snapshot.mtimeMs > knownMtime + 1; + }) || null + ); +} + +async function waitForManualCallbackToken( + provider: CLIProxyProvider, + target: ProxyTarget, + tokenDir: string, + oauthState: string | null, + knownTokenFiles: ProviderTokenSnapshot[], + expectedAccountId: string | undefined, + timeoutMs: number, + pollIntervalMs: number = PASTE_CALLBACK_AUTH_URL_POLL_INTERVAL_MS +): Promise<{ tokenSnapshot: ProviderTokenSnapshot | null; error?: string }> { + const deadline = Date.now() + timeoutMs; + let upstreamCompletedAt: number | null = null; + + while (Date.now() < deadline) { + const tokenSnapshot = findNewTokenSnapshotForManualAuth( + provider, + tokenDir, + knownTokenFiles, + expectedAccountId + ); + if (tokenSnapshot) { + return { tokenSnapshot }; + } + + if (oauthState) { + const response = await fetch( + buildProxyUrl( + target, + `/v0/management/get-auth-status?state=${encodeURIComponent(oauthState)}` + ), + { headers: buildManagementHeaders(target) } + ); + + if (response.ok) { + const data = (await response.json()) as { status?: string; error?: string }; + if (data.status === 'error') { + return { + tokenSnapshot: null, + error: data.error || 'Authentication failed while waiting for local token persistence', + }; + } + if (data.status === 'ok' && upstreamCompletedAt === null) { + upstreamCompletedAt = Date.now(); + } + } + } + + if ( + upstreamCompletedAt !== null && + Date.now() - upstreamCompletedAt >= POLLED_AUTH_LOCAL_TOKEN_GRACE_MS + ) { + break; + } + + if (Date.now() + pollIntervalMs >= deadline) { + break; + } + + await sleep(pollIntervalMs); + } + + return { tokenSnapshot: null }; +} + export async function resolvePasteCallbackAuthUrl( target: ProxyTarget, startData: PasteCallbackStartData, @@ -397,6 +537,9 @@ async function handlePasteCallbackMode( return null; } + const oauthState = startData.state || parseAuthUrlState(authUrl); + const knownTokenFiles = listProviderTokenSnapshots(provider, tokenDir); + // Display auth URL in box console.log(''); console.log(' ╔══════════════════════════════════════════════════════════════╗'); @@ -489,15 +632,49 @@ async function handlePasteCallbackMode( return null; } - console.log(ok('Authentication successful!')); + console.log(info('Callback submitted. Waiting for token exchange...')); + const { tokenSnapshot, error: tokenWaitError } = await waitForManualCallbackToken( + provider, + target, + tokenDir, + oauthState, + knownTokenFiles, + expectedAccountId, + OAUTH_STATE_TIMEOUT_MS + ); + + if (tokenWaitError) { + console.log(fail(tokenWaitError)); + warnPossible403Ban(provider, tokenWaitError); + return null; + } + + if (!tokenSnapshot) { + console.log( + fail( + 'Authentication completed upstream, but no new local token was saved for this account. Update CCS/CLIProxy and retry.' + ) + ); + return null; + } + const account = registerAccountFromToken( provider, tokenDir, nickname, verbose, - expectedAccountId + tokenSnapshot.file ); + if (!account) { + console.log( + fail('Authenticated token could not be matched to the requested account. Retry the flow.') + ); + return null; + } + + console.log(ok('Authentication successful!')); + // Account safety: check for cross-provider conflicts if (account?.email) { const conflicts = checkNewAccountConflict(provider, account.email); diff --git a/src/cliproxy/auth/oauth-process.ts b/src/cliproxy/auth/oauth-process.ts index 74b78d8d..bce9491c 100644 --- a/src/cliproxy/auth/oauth-process.ts +++ b/src/cliproxy/auth/oauth-process.ts @@ -198,7 +198,7 @@ export function getKiroBuilderIdSelectionInput(output: string): string | null { } const promptWindow = output.slice(promptMatch.index, promptMatch.index + 600); - const optionMatch = /(?:^|\n)\s*(\d+)\s*[\).:-]?\s*(?:AWS\s+)?Builder ID\b/im.exec(promptWindow); + const optionMatch = /(?:^|\n)\s*(\d+)\s*[\).:-]?\s*.*\bBuilder ID\b/im.exec(promptWindow); if (!optionMatch) { return null; } @@ -206,6 +206,33 @@ export function getKiroBuilderIdSelectionInput(output: string): string | null { return `${optionMatch[1]}\n`; } +export function extractLikelyOAuthAuthorizationUrl(output: string): string | null { + const urls = Array.from(output.matchAll(/https?:\/\/[^\s]+/g), (match) => match[0]); + let selectedUrl: string | null = null; + let selectedScore = 0; + + for (const url of urls) { + try { + const parsed = new URL(url); + let score = 0; + if (parsed.searchParams.has('redirect_uri')) score += 4; + if (parsed.searchParams.has('state')) score += 2; + if (parsed.searchParams.has('code_challenge')) score += 1; + if (parsed.pathname.includes('/authorize')) score += 1; + if (isLoopbackHost(parsed.hostname)) score -= 3; + + if (score >= selectedScore && score > 0) { + selectedUrl = url; + selectedScore = score; + } + } catch { + continue; + } + } + + return selectedUrl; +} + async function promptManualCallbackUrl( displayName: string, state: ProcessState, @@ -332,17 +359,12 @@ async function handleStdout( !state.kiroMethodSelectionHandled && state.accumulatedOutput.includes('Select login method') ) { - state.kiroMethodSelectionHandled = true; const builderIdSelection = getKiroBuilderIdSelectionInput(state.accumulatedOutput); - if (!builderIdSelection) { - console.log(fail('Unable to auto-select Kiro Builder ID from the upstream login menu.')); - console.log(' The upstream Kiro prompt format may have changed.'); - killWithEscalation(authProcess); - return; + if (builderIdSelection) { + state.kiroMethodSelectionHandled = true; + authProcess.stdin?.write(builderIdSelection); + log(`Auto-selected Kiro Builder ID flow (${builderIdSelection.trim()})`); } - - authProcess.stdin?.write(builderIdSelection); - log(`Auto-selected Kiro Builder ID flow (${builderIdSelection.trim()})`); } // Parse project list when available @@ -413,11 +435,11 @@ async function handleStdout( // Display OAuth URL for all modes (enables VS Code terminal URL detection popup) if (!isDeviceCodeFlow && !state.urlDisplayed) { - const urlMatch = output.match(/https?:\/\/[^\s]+/); - if (urlMatch) { + const authUrl = extractLikelyOAuthAuthorizationUrl(state.accumulatedOutput); + if (authUrl) { console.log(''); console.log(info(`${options.oauthConfig.displayName} OAuth URL:`)); - console.log(` ${urlMatch[0]}`); + console.log(` ${authUrl}`); console.log(''); state.urlDisplayed = true; @@ -426,7 +448,7 @@ async function handleStdout( await replayManualCallback( options.oauthConfig, authProcess, - urlMatch[0], + authUrl, options.verbose, state, 10 * 60 * 1000 @@ -442,11 +464,11 @@ function displayUrlFromStderr( state: ProcessState, oauthConfig: ProviderOAuthConfig ): void { - const urlMatch = output.match(/https?:\/\/[^\s]+/); - if (urlMatch) { + const authUrl = extractLikelyOAuthAuthorizationUrl(output); + if (authUrl) { console.log(''); console.log(info(`${oauthConfig.displayName} OAuth URL:`)); - console.log(` ${urlMatch[0]}`); + console.log(` ${authUrl}`); console.log(''); state.urlDisplayed = true; } diff --git a/tests/unit/cliproxy/oauth-handler-paste-callback.test.ts b/tests/unit/cliproxy/oauth-handler-paste-callback.test.ts index 983e94d1..287b95f8 100644 --- a/tests/unit/cliproxy/oauth-handler-paste-callback.test.ts +++ b/tests/unit/cliproxy/oauth-handler-paste-callback.test.ts @@ -1,4 +1,7 @@ import { afterEach, describe, expect, it } from 'bun:test'; +import * as fs from 'fs'; +import * as os from 'os'; +import * as path from 'path'; import type { ProxyTarget } from '../../../src/cliproxy/proxy-target-resolver'; import { getCapturedFetchRequests, mockFetch, restoreFetch } from '../../mocks'; @@ -130,6 +133,56 @@ describe('resolvePasteCallbackAuthUrl', () => { }); }); +describe('findNewTokenSnapshotForManualAuth', () => { + it('detects newly created provider token files', async () => { + const tokenDir = fs.mkdtempSync(path.join(os.tmpdir(), 'ccs-kiro-manual-auth-')); + const existingFile = path.join(tokenDir, 'kiro-existing.json'); + fs.writeFileSync(existingFile, JSON.stringify({ type: 'kiro', email: 'existing@example.com' })); + const existingMtimeMs = fs.statSync(existingFile).mtimeMs; + + const { findNewTokenSnapshotForManualAuth } = await import( + `../../../src/cliproxy/auth/oauth-handler?manual-auth-new-token=${Date.now()}` + ); + + const newFile = path.join(tokenDir, 'kiro-new.json'); + fs.writeFileSync(newFile, JSON.stringify({ type: 'kiro', email: 'new@example.com' })); + + const snapshot = findNewTokenSnapshotForManualAuth( + 'kiro', + tokenDir, + [{ file: 'kiro-existing.json', mtimeMs: existingMtimeMs }] + ); + + expect(snapshot?.file).toBe('kiro-new.json'); + fs.rmSync(tokenDir, { recursive: true, force: true }); + }); + + it('treats a modified existing token as the new token during reauth', async () => { + const tokenDir = fs.mkdtempSync(path.join(os.tmpdir(), 'ccs-kiro-reauth-')); + const tokenFile = path.join(tokenDir, 'kiro-existing.json'); + fs.writeFileSync(tokenFile, JSON.stringify({ type: 'kiro', email: 'existing@example.com' })); + const existingMtimeMs = fs.statSync(tokenFile).mtimeMs; + + const { findNewTokenSnapshotForManualAuth } = await import( + `../../../src/cliproxy/auth/oauth-handler?manual-auth-updated-token=${Date.now()}` + ); + + fs.writeFileSync(tokenFile, JSON.stringify({ type: 'kiro', email: 'existing@example.com', refreshed: true })); + const bumpedTime = new Date(existingMtimeMs + 10_000); + fs.utimesSync(tokenFile, bumpedTime, bumpedTime); + + const snapshot = findNewTokenSnapshotForManualAuth( + 'kiro', + tokenDir, + [{ file: 'kiro-existing.json', mtimeMs: existingMtimeMs }], + 'kiro-existing.json' + ); + + expect(snapshot?.file).toBe('kiro-existing.json'); + fs.rmSync(tokenDir, { recursive: true, force: true }); + }); +}); + describe('getCliAuthNicknameError', () => { it('allows omitted nicknames for no-email providers', async () => { const { getCliAuthNicknameError } = await import( diff --git a/tests/unit/cliproxy/oauth-process-error-parser.test.ts b/tests/unit/cliproxy/oauth-process-error-parser.test.ts index 322f0a32..a9d7f806 100644 --- a/tests/unit/cliproxy/oauth-process-error-parser.test.ts +++ b/tests/unit/cliproxy/oauth-process-error-parser.test.ts @@ -1,6 +1,7 @@ import { describe, expect, it } from 'bun:test'; import { extractLikelyAuthFailureFromStderr, + extractLikelyOAuthAuthorizationUrl, getExpectedLocalCallback, getKiroBuilderIdSelectionInput, validateManualCallbackUrl, @@ -89,9 +90,9 @@ describe('oauth-process manual callback validation', () => { describe('oauth-process Kiro Builder ID menu parsing', () => { it('selects Builder ID when it is the first option', () => { const output = ` -Select login method -1. Builder ID -2. IAM Identity Center +? Select login method: + 1) Use with Builder ID (personal AWS account) + 2) Use with IDC Account (organization SSO) `; expect(getKiroBuilderIdSelectionInput(output)).toBe('1\n'); @@ -117,3 +118,28 @@ Select login method expect(getKiroBuilderIdSelectionInput(output)).toBeNull(); }); }); + +describe('oauth-process OAuth URL extraction', () => { + it('prefers the real auth URL over the IDC start URL banner', () => { + const authUrl = + 'https://oidc.us-east-1.amazonaws.com/authorize?response_type=code&client_id=test-client&redirect_uri=http%3A%2F%2F127.0.0.1%3A9876%2Foauth%2Fcallback&state=test-state&code_challenge=test-challenge&code_challenge_method=S256'; + const output = ` +Using IDC with Start URL: https://d-123.awsapps.com/start +Region: us-east-1 +URL: ${authUrl} +`; + + expect(extractLikelyOAuthAuthorizationUrl(output)).toBe(authUrl); + }); + + it('ignores local callback server URLs when the auth URL is also present', () => { + const authUrl = + 'https://device.sso.us-east-1.amazonaws.com/authorize?response_type=code&client_id=test-client&redirect_uri=http%3A%2F%2F127.0.0.1%3A9876%2Foauth%2Fcallback&state=test-state&code_challenge=test-challenge&code_challenge_method=S256'; + const output = ` +Callback server started, redirect URI: http://127.0.0.1:9876/oauth/callback +URL: ${authUrl} +`; + + expect(extractLikelyOAuthAuthorizationUrl(output)).toBe(authUrl); + }); +});