From 97db7c9394d56f840fa1bd9a4ce6616f1f141930 Mon Sep 17 00:00:00 2001 From: Mux Date: Sat, 7 Feb 2026 09:45:26 +0100 Subject: [PATCH 01/22] feat: extract shared OAuth utilities from service files MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Create 3 new shared modules for the duplicated OAuth infrastructure found across 5 service files (codexOauthService, copilotOauthService, muxGatewayOauthService, muxGovernorOauthService, mcpOauthService): - oauthUtils.ts: createDeferred, closeServer, escapeHtml, renderOAuthCallbackHtml — verbatim-duplicated utilities - oauthLoopbackServer.ts: startLoopbackServer — shared local HTTP server pattern for receiving OAuth code redirects - oauthFlowManager.ts: OAuthFlowManager class — shared desktop flow lifecycle (register/waitFor/cancel/finish/shutdownAll) No existing files are modified. Follow-up work will wire each service to use these shared modules. --- src/node/utils/oauthFlowManager.ts | 141 +++++++++++++++++++ src/node/utils/oauthLoopbackServer.ts | 191 ++++++++++++++++++++++++++ src/node/utils/oauthUtils.ts | 101 ++++++++++++++ 3 files changed, 433 insertions(+) create mode 100644 src/node/utils/oauthFlowManager.ts create mode 100644 src/node/utils/oauthLoopbackServer.ts create mode 100644 src/node/utils/oauthUtils.ts diff --git a/src/node/utils/oauthFlowManager.ts b/src/node/utils/oauthFlowManager.ts new file mode 100644 index 0000000000..edc02852d9 --- /dev/null +++ b/src/node/utils/oauthFlowManager.ts @@ -0,0 +1,141 @@ +import type http from "node:http"; +import type { Result } from "@/common/types/result"; +import { Err } from "@/common/types/result"; +import { closeServer } from "@/node/utils/oauthUtils"; +import { log } from "@/node/services/log"; + +/** + * Shared desktop OAuth flow lifecycle manager. + * + * Four OAuth services (Gateway, Governor, Codex, MCP) track in-flight desktop + * flows with an identical `Map` + `waitFor`/`cancel`/ + * `finish`/`shutdownAll` pattern. This class extracts that shared lifecycle + * so each service can delegate flow bookkeeping here. + */ + +// --------------------------------------------------------------------------- +// Types +// --------------------------------------------------------------------------- + +export interface OAuthFlowEntry { + /** The loopback HTTP server (null for device-code flows). */ + server: http.Server | null; + /** Deferred that resolves with the final flow result. */ + resultDeferred: { + promise: Promise>; + resolve: (value: Result) => void; + }; + /** Handle for the server-side timeout (set by `waitFor`). */ + timeoutHandle: ReturnType | null; +} + +// --------------------------------------------------------------------------- +// Manager +// --------------------------------------------------------------------------- + +export class OAuthFlowManager { + private readonly flows = new Map(); + + /** Register a new in-flight flow. */ + register(flowId: string, entry: OAuthFlowEntry): void { + this.flows.set(flowId, entry); + } + + /** Get a flow entry by ID, or undefined if not found. */ + get(flowId: string): OAuthFlowEntry | undefined { + return this.flows.get(flowId); + } + + /** Check whether a flow exists. */ + has(flowId: string): boolean { + return this.flows.has(flowId); + } + + /** + * Wait for a flow to complete, with a timeout. + * + * Mirrors the `waitForDesktopFlow` pattern shared across all four services: + * - Sets up a timeout that races against the deferred promise. + * - Stores the timeout handle in the entry for cleanup. + * - On timeout/error, calls `finish` to close the server and clean up. + */ + async waitFor(flowId: string, timeoutMs: number): Promise> { + const flow = this.flows.get(flowId); + if (!flow) { + return Err("OAuth flow not found"); + } + + let timeoutHandle: ReturnType | null = null; + const timeoutPromise = new Promise>((resolve) => { + timeoutHandle = setTimeout(() => { + resolve(Err("Timed out waiting for OAuth callback")); + }, timeoutMs); + }); + + const result = await Promise.race([flow.resultDeferred.promise, timeoutPromise]); + + if (timeoutHandle !== null) { + clearTimeout(timeoutHandle); + } + + if (!result.success) { + // Ensure listener is closed on timeout/errors. + void this.finish(flowId, result); + } + + return result; + } + + /** + * Cancel a flow — resolves the deferred with an error and cleans up. + * + * Mirrors the `cancelDesktopFlow` pattern. + */ + async cancel(flowId: string): Promise { + const flow = this.flows.get(flowId); + if (!flow) return; + + await this.finish(flowId, Err("OAuth flow cancelled")); + } + + /** + * Finish a flow: resolve the deferred, clear the timeout, close the server, + * and remove the entry from the map. + * + * Idempotent — no-op if the flow was already removed. Mirrors the + * `finishDesktopFlow` pattern. + */ + async finish(flowId: string, result: Result): Promise { + const flow = this.flows.get(flowId); + if (!flow) return; + + // Remove from map first to make re-entrant calls no-ops. + this.flows.delete(flowId); + + if (flow.timeoutHandle !== null) { + clearTimeout(flow.timeoutHandle); + } + + try { + flow.resultDeferred.resolve(result); + + if (flow.server) { + // Stop accepting new connections. + await closeServer(flow.server); + } + } catch (error) { + log.debug("Failed to close OAuth callback listener:", error); + } + } + + /** + * Shut down all active flows — resolves each with an error. + * + * Mirrors the `dispose` pattern where services iterate all flows + * and finish them with `Err("App shutting down")`. + */ + async shutdownAll(): Promise { + const flowIds = [...this.flows.keys()]; + await Promise.all(flowIds.map((id) => this.finish(id, Err("App shutting down")))); + } +} diff --git a/src/node/utils/oauthLoopbackServer.ts b/src/node/utils/oauthLoopbackServer.ts new file mode 100644 index 0000000000..2773892d17 --- /dev/null +++ b/src/node/utils/oauthLoopbackServer.ts @@ -0,0 +1,191 @@ +import http from "node:http"; +import type { Result } from "@/common/types/result"; +import { Err, Ok } from "@/common/types/result"; +import { closeServer, createDeferred, renderOAuthCallbackHtml } from "@/node/utils/oauthUtils"; + +/** + * Shared loopback OAuth callback server. + * + * Four OAuth services (Gateway, Governor, Codex, MCP) spin up a local HTTP + * server to receive the authorization code redirect. The pattern is identical + * across all four — this module extracts that into a single reusable utility. + */ + +// --------------------------------------------------------------------------- +// Types +// --------------------------------------------------------------------------- + +export interface LoopbackServerOptions { + /** Port to listen on. 0 = random (default). Codex uses 1455. */ + port?: number; + /** Host to bind to. Default: "127.0.0.1" */ + host?: string; + /** Path to match for the OAuth callback. Default: "/callback" */ + callbackPath?: string; + /** Expected state parameter value for CSRF validation. */ + expectedState: string; + /** Whether to validate that remoteAddress is a loopback address. Default: false. Codex uses true. */ + validateLoopback?: boolean; + /** Custom HTML renderer. If not provided, uses renderOAuthCallbackHtml with generic branding. */ + renderHtml?: (result: { success: boolean; error?: string }) => string; +} + +export interface LoopbackCallbackResult { + code: string; + state: string; +} + +export interface LoopbackServer { + /** The full redirect URI (http://127.0.0.1:{port}{callbackPath}). */ + redirectUri: string; + /** The underlying HTTP server (needed by OAuthFlowManager for cleanup). */ + server: http.Server; + /** Resolves when callback received or resolves with Err on invalid state. */ + result: Promise>; + /** Cancel and close the server. */ + cancel: () => Promise; + /** Close the server. */ + close: () => Promise; +} + +// --------------------------------------------------------------------------- +// Helpers +// --------------------------------------------------------------------------- + +/** + * Check whether an address string is a loopback address. + * Node may normalize IPv4 loopback to an IPv6-mapped address. + * + * Extracted from codexOauthService.ts where validateLoopback is used. + */ +function isLoopbackAddress(address: string | undefined): boolean { + if (!address) return false; + + // Node may normalize IPv4 loopback to an IPv6-mapped address. + if (address === "::ffff:127.0.0.1") { + return true; + } + + return address === "127.0.0.1" || address === "::1"; +} + +// --------------------------------------------------------------------------- +// Main +// --------------------------------------------------------------------------- + +/** + * Start a loopback HTTP server to receive an OAuth authorization code callback. + * + * Pattern extracted from the `http.createServer` blocks in Gateway, Governor, + * Codex, and MCP OAuth services. The server: + * + * 1. Optionally validates the remote address is loopback (Codex). + * 2. Matches only GET requests on `callbackPath`. + * 3. Validates the `state` query parameter against `expectedState`. + * 4. Extracts `code` from the query string. + * 5. Responds with HTML (success or error). + * 6. Resolves the result deferred — the caller then performs token exchange + * and calls `close()`. + * + * The server does NOT close itself after responding — the caller decides when + * to close (matching the existing pattern where services call `closeServer` + * after token exchange). + */ +export async function startLoopbackServer(options: LoopbackServerOptions): Promise { + const port = options.port ?? 0; + const host = options.host ?? "127.0.0.1"; + const callbackPath = options.callbackPath ?? "/callback"; + const validateLoopback = options.validateLoopback ?? false; + + const deferred = createDeferred>(); + + const render = + options.renderHtml ?? + ((r: { success: boolean; error?: string }) => + renderOAuthCallbackHtml({ + title: r.success ? "Login complete" : "Login failed", + message: r.success + ? "You can return to Mux. You may now close this tab." + : (r.error ?? "Unknown error"), + success: r.success, + })); + + const server = http.createServer((req, res) => { + // Optionally reject non-loopback connections (Codex sets validateLoopback: true). + if (validateLoopback && !isLoopbackAddress(req.socket.remoteAddress)) { + res.statusCode = 403; + res.end("Forbidden"); + return; + } + + const reqUrl = req.url ?? "/"; + const url = new URL(reqUrl, "http://localhost"); + + if (req.method !== "GET" || url.pathname !== callbackPath) { + res.statusCode = 404; + res.end("Not found"); + return; + } + + const state = url.searchParams.get("state"); + if (!state || state !== options.expectedState) { + res.statusCode = 400; + res.setHeader("Content-Type", "text/html"); + res.end("

Invalid OAuth state

"); + deferred.resolve(Err("Invalid OAuth state")); + return; + } + + const code = url.searchParams.get("code"); + const error = url.searchParams.get("error"); + const errorDescription = url.searchParams.get("error_description") ?? undefined; + + if (error) { + const errorMessage = errorDescription ? `${error}: ${errorDescription}` : error; + res.setHeader("Content-Type", "text/html"); + res.statusCode = 400; + res.end(render({ success: false, error: errorMessage })); + deferred.resolve(Err(errorMessage)); + return; + } + + if (!code) { + const errorMessage = "Missing authorization code"; + res.setHeader("Content-Type", "text/html"); + res.statusCode = 400; + res.end(render({ success: false, error: errorMessage })); + deferred.resolve(Err(errorMessage)); + return; + } + + res.setHeader("Content-Type", "text/html"); + res.end(render({ success: true })); + deferred.resolve(Ok({ code, state })); + }); + + // Listen on the specified host/port — mirrors the existing + // `server.listen(port, host, () => resolve())` pattern. + await new Promise((resolve, reject) => { + server.once("error", reject); + server.listen(port, host, () => resolve()); + }); + + const address = server.address(); + if (!address || typeof address === "string") { + await closeServer(server); + throw new Error("Failed to determine OAuth callback listener port"); + } + + const redirectUri = `http://127.0.0.1:${address.port}${callbackPath}`; + + return { + redirectUri, + server, + result: deferred.promise, + cancel: async () => { + deferred.resolve(Err("OAuth flow cancelled")); + await closeServer(server); + }, + close: () => closeServer(server), + }; +} diff --git a/src/node/utils/oauthUtils.ts b/src/node/utils/oauthUtils.ts new file mode 100644 index 0000000000..c6185a3c67 --- /dev/null +++ b/src/node/utils/oauthUtils.ts @@ -0,0 +1,101 @@ +import type http from "node:http"; + +/** + * Shared OAuth utility functions extracted from the individual OAuth service files. + * + * These are verbatim-duplicated across codexOauthService, copilotOauthService, + * muxGatewayOauthService, muxGovernorOauthService, and mcpOauthService. + */ + +/** A deferred promise with an externally-accessible `resolve` handle. */ +export interface Deferred { + promise: Promise; + resolve: (value: T) => void; +} + +/** Create a deferred promise that can be resolved externally. */ +export function createDeferred(): Deferred { + let resolve!: (value: T) => void; + const promise = new Promise((res) => { + resolve = res; + }); + return { promise, resolve }; +} + +/** Gracefully close an HTTP server, resolving when all connections are drained. */ +export function closeServer(server: http.Server): Promise { + return new Promise((resolve) => { + server.close(() => resolve()); + }); +} + +/** Escape HTML special characters to prevent XSS in rendered callback pages. */ +export function escapeHtml(input: string): string { + return input + .replaceAll("&", "&") + .replaceAll("<", "<") + .replaceAll(">", ">") + .replaceAll('"', """) + .replaceAll("'", "'"); +} + +export interface RenderOAuthCallbackHtmlOptions { + /** Page and heading title. */ + title: string; + /** Body message shown below the title. */ + message: string; + /** When true, the page auto-closes via `window.close()`. */ + success: boolean; + /** Optional extra content injected into `` (e.g. an external CSS link). */ + extraHead?: string; +} + +/** + * Render the HTML page returned to the browser after an OAuth callback. + * + * All four loopback-based services (Gateway, Governor, Codex, MCP) return an + * HTML page with a title, message, and auto-close script on success. The + * structure mirrors the common pattern found across those services: + * + * - `` with basic inline styling (centered, system font) + * - Title in `

`, message in `

` + * - Auto-close ` + +`; +} From 60a8075bdd3da69935733441ccb8330197805dab Mon Sep 17 00:00:00 2001 From: Mux Date: Sat, 7 Feb 2026 09:48:01 +0100 Subject: [PATCH 02/22] refactor(copilot): use shared createDeferred from oauthUtils Replace the inline deferred-promise pattern with the shared createDeferred() utility extracted into oauthUtils. --- src/node/services/copilotOauthService.ts | 7 ++----- 1 file changed, 2 insertions(+), 5 deletions(-) diff --git a/src/node/services/copilotOauthService.ts b/src/node/services/copilotOauthService.ts index 1b95da5817..c725ce843e 100644 --- a/src/node/services/copilotOauthService.ts +++ b/src/node/services/copilotOauthService.ts @@ -4,6 +4,7 @@ import { Err, Ok } from "@/common/types/result"; import type { ProviderService } from "@/node/services/providerService"; import type { WindowService } from "@/node/services/windowService"; import { log } from "@/node/services/log"; +import { createDeferred } from "@/node/utils/oauthUtils"; const GITHUB_COPILOT_CLIENT_ID = "Ov23liCVKFN3jOo9R7HS"; const SCOPE = "read:user"; @@ -71,11 +72,7 @@ export class CopilotOauthService { return Err("Invalid response from GitHub device code endpoint"); } - // Create deferred promise - let resolveResult!: (result: Result) => void; - const resultPromise = new Promise>((resolve) => { - resolveResult = resolve; - }); + const { promise: resultPromise, resolve: resolveResult } = createDeferred>(); const timeout = setTimeout(() => { void this.finishFlow(flowId, Err("Timed out waiting for GitHub authorization")); From f091a6570bc0f34e6a005c12556da5e1dd5d56af Mon Sep 17 00:00:00 2001 From: Mux Date: Sat, 7 Feb 2026 09:51:56 +0100 Subject: [PATCH 03/22] refactor: use shared OAuth utilities in muxGatewayOauthService Replace duplicated OAuth infrastructure with the new shared utilities: - Import createDeferred, renderOAuthCallbackHtml from oauthUtils - Replace http.createServer block with startLoopbackServer - Replace manual desktopFlows Map + lifecycle methods with OAuthFlowManager - Remove local closeServer, createDeferred, escapeHtml helpers - Remove DesktopFlow interface and handleDesktopCallback method - Remove finishDesktopFlow (delegated to OAuthFlowManager.finish) - Simplify dispose() to use shutdownAll() - Preserve Gateway-specific branding via renderHtml + extraHead Public API unchanged. File reduced from 443 to 239 lines (-46%). --- src/node/services/muxGatewayOauthService.ts | 316 ++++---------------- 1 file changed, 56 insertions(+), 260 deletions(-) diff --git a/src/node/services/muxGatewayOauthService.ts b/src/node/services/muxGatewayOauthService.ts index 1b52ebb80e..6e5addd51d 100644 --- a/src/node/services/muxGatewayOauthService.ts +++ b/src/node/services/muxGatewayOauthService.ts @@ -1,5 +1,4 @@ import * as crypto from "crypto"; -import * as http from "http"; import type { Result } from "@/common/types/result"; import { Err, Ok } from "@/common/types/result"; import { @@ -10,44 +9,20 @@ import { import type { ProviderService } from "@/node/services/providerService"; import type { WindowService } from "@/node/services/windowService"; import { log } from "@/node/services/log"; +import { createDeferred, renderOAuthCallbackHtml } from "@/node/utils/oauthUtils"; +import { startLoopbackServer } from "@/node/utils/oauthLoopbackServer"; +import { OAuthFlowManager } from "@/node/utils/oauthFlowManager"; const DEFAULT_DESKTOP_TIMEOUT_MS = 5 * 60 * 1000; const DEFAULT_SERVER_TIMEOUT_MS = 10 * 60 * 1000; -const COMPLETED_DESKTOP_FLOW_TTL_MS = 60 * 1000; - -interface DesktopFlow { - flowId: string; - authorizeUrl: string; - redirectUri: string; - server: http.Server; - timeout: ReturnType; - cleanupTimeout: ReturnType | null; - resultPromise: Promise>; - resolveResult: (result: Result) => void; - settled: boolean; -} interface ServerFlow { state: string; expiresAtMs: number; } -function closeServer(server: http.Server): Promise { - return new Promise((resolve) => { - server.close(() => resolve()); - }); -} - -function createDeferred(): { promise: Promise; resolve: (value: T) => void } { - let resolve!: (value: T) => void; - const promise = new Promise((res) => { - resolve = res; - }); - return { promise, resolve }; -} - export class MuxGatewayOauthService { - private readonly desktopFlows = new Map(); + private readonly desktopFlows = new OAuthFlowManager(); private readonly serverFlows = new Map(); constructor( @@ -59,118 +34,80 @@ export class MuxGatewayOauthService { Result<{ flowId: string; authorizeUrl: string; redirectUri: string }, string> > { const flowId = crypto.randomUUID(); + const resultDeferred = createDeferred>(); - const { promise: resultPromise, resolve: resolveResult } = - createDeferred>(); - - const server = http.createServer((req, res) => { - const reqUrl = req.url ?? "/"; - const url = new URL(reqUrl, "http://localhost"); - - if (req.method !== "GET" || url.pathname !== "/callback") { - res.statusCode = 404; - res.end("Not found"); - return; - } - - const state = url.searchParams.get("state"); - if (!state || state !== flowId) { - res.statusCode = 400; - res.setHeader("Content-Type", "text/html"); - res.end("

Invalid OAuth state

"); - return; - } - - const code = url.searchParams.get("code"); - const error = url.searchParams.get("error"); - const errorDescription = url.searchParams.get("error_description") ?? undefined; - - void this.handleDesktopCallback({ - flowId, - code, - error, - errorDescription, - res, - }); - }); - + let loopback; try { - await new Promise((resolve, reject) => { - server.once("error", reject); - server.listen(0, "127.0.0.1", () => resolve()); + loopback = await startLoopbackServer({ + expectedState: flowId, + renderHtml: (r) => + renderOAuthCallbackHtml({ + title: r.success ? "Login complete" : "Login failed", + message: r.success + ? "You can return to Mux. You may now close this tab." + : (r.error ?? "Unknown error"), + success: r.success, + extraHead: + '\n ', + }), }); } catch (error) { const message = error instanceof Error ? error.message : String(error); return Err(`Failed to start OAuth callback listener: ${message}`); } - const address = server.address(); - if (!address || typeof address === "string") { - return Err("Failed to determine OAuth callback listener port"); - } + const authorizeUrl = buildAuthorizeUrl({ redirectUri: loopback.redirectUri, state: flowId }); - const redirectUri = `http://127.0.0.1:${address.port}/callback`; - const authorizeUrl = buildAuthorizeUrl({ redirectUri, state: flowId }); - - const timeout = setTimeout(() => { - void this.finishDesktopFlow(flowId, Err("Timed out waiting for OAuth callback")); - }, DEFAULT_DESKTOP_TIMEOUT_MS); - - this.desktopFlows.set(flowId, { - flowId, - authorizeUrl, - redirectUri, - server, - timeout, - cleanupTimeout: null, - resultPromise, - resolveResult, - settled: false, + this.desktopFlows.register(flowId, { + server: loopback.server, + resultDeferred, + timeoutHandle: null, }); + // Background task: await loopback callback, do token exchange, finish flow. + // Race against resultDeferred so that if the flow is cancelled/timed out + // externally, this task exits cleanly instead of dangling on loopback.result. + void (async () => { + const callbackOrDone = await Promise.race([ + loopback.result, + resultDeferred.promise.then((): null => null), + ]); + + // Flow was already finished externally (timeout or cancel). + if (callbackOrDone === null) return; + + log.debug(`Mux Gateway OAuth callback received (flowId=${flowId})`); + + let result: Result; + if (callbackOrDone.success) { + result = await this.handleCallbackAndExchange({ + state: callbackOrDone.data.state, + code: callbackOrDone.data.code, + error: null, + }); + } else { + result = Err(`Mux Gateway OAuth error: ${callbackOrDone.error}`); + } + + await this.desktopFlows.finish(flowId, result); + })(); + log.debug(`Mux Gateway OAuth desktop flow started (flowId=${flowId})`); - return Ok({ flowId, authorizeUrl, redirectUri }); + return Ok({ flowId, authorizeUrl, redirectUri: loopback.redirectUri }); } async waitForDesktopFlow( flowId: string, opts?: { timeoutMs?: number } ): Promise> { - const flow = this.desktopFlows.get(flowId); - if (!flow) { - return Err("OAuth flow not found"); - } - - const timeoutMs = opts?.timeoutMs ?? DEFAULT_DESKTOP_TIMEOUT_MS; - - let timeoutHandle: ReturnType | null = null; - const timeoutPromise = new Promise>((resolve) => { - timeoutHandle = setTimeout(() => { - resolve(Err("Timed out waiting for OAuth callback")); - }, timeoutMs); - }); - - const result = await Promise.race([flow.resultPromise, timeoutPromise]); - - if (timeoutHandle !== null) { - clearTimeout(timeoutHandle); - } - - if (!result.success) { - // Ensure listener is closed on timeout/errors. - void this.finishDesktopFlow(flowId, result); - } - - return result; + return this.desktopFlows.waitFor(flowId, opts?.timeoutMs ?? DEFAULT_DESKTOP_TIMEOUT_MS); } async cancelDesktopFlow(flowId: string): Promise { - const flow = this.desktopFlows.get(flowId); - if (!flow) return; - + if (!this.desktopFlows.has(flowId)) return; log.debug(`Mux Gateway OAuth desktop flow cancelled (flowId=${flowId})`); - await this.finishDesktopFlow(flowId, Err("OAuth flow cancelled")); + await this.desktopFlows.cancel(flowId); } startServerFlow(input: { redirectUri: string }): { authorizeUrl: string; state: string } { @@ -228,118 +165,10 @@ export class MuxGatewayOauthService { } async dispose(): Promise { - // Best-effort: cancel all in-flight flows. - const flowIds = [...this.desktopFlows.keys()]; - await Promise.all(flowIds.map((id) => this.finishDesktopFlow(id, Err("App shutting down")))); - - for (const flow of this.desktopFlows.values()) { - clearTimeout(flow.timeout); - if (flow.cleanupTimeout !== null) { - clearTimeout(flow.cleanupTimeout); - } - } - - this.desktopFlows.clear(); + await this.desktopFlows.shutdownAll(); this.serverFlows.clear(); } - private async handleDesktopCallback(input: { - flowId: string; - code: string | null; - error: string | null; - errorDescription?: string; - res: http.ServerResponse; - }): Promise { - const flow = this.desktopFlows.get(input.flowId); - if (!flow || flow.settled) { - input.res.statusCode = 409; - input.res.setHeader("Content-Type", "text/html"); - input.res.end("

OAuth flow already completed

"); - return; - } - - log.debug(`Mux Gateway OAuth callback received (flowId=${input.flowId})`); - - const result = await this.handleCallbackAndExchange({ - state: input.flowId, - code: input.code, - error: input.error, - errorDescription: input.errorDescription, - }); - - const title = result.success ? "Login complete" : "Login failed"; - const description = result.success - ? "You can return to Mux. You may now close this tab." - : escapeHtml(result.error); - - const html = ` - - - - - - - ${title} - - - -
-
-
-
mux
-
-
- -
-
-
-

${title}

-

${description}

- ${ - result.success - ? '

Mux should now be in the foreground. You can close this tab.

' - : '

You can close this tab.

' - } -
-
-
-
- - - -`; - - input.res.setHeader("Content-Type", "text/html"); - if (!result.success) { - input.res.statusCode = 400; - } - - input.res.end(html); - - await this.finishDesktopFlow(input.flowId, result); - } - private async handleCallbackAndExchange(input: { state: string; code: string | null; @@ -407,37 +236,4 @@ export class MuxGatewayOauthService { } } - private async finishDesktopFlow(flowId: string, result: Result): Promise { - const flow = this.desktopFlows.get(flowId); - if (!flow || flow.settled) return; - - flow.settled = true; - clearTimeout(flow.timeout); - - try { - flow.resolveResult(result); - - // Stop accepting new connections. - await closeServer(flow.server); - } catch (error) { - log.debug("Failed to close OAuth callback listener:", error); - } finally { - // Keep the completed flow around briefly so callers can still await the result. - if (flow.cleanupTimeout !== null) { - clearTimeout(flow.cleanupTimeout); - } - flow.cleanupTimeout = setTimeout(() => { - this.desktopFlows.delete(flowId); - }, COMPLETED_DESKTOP_FLOW_TTL_MS); - } - } -} - -function escapeHtml(input: string): string { - return input - .replaceAll("&", "&") - .replaceAll("<", "<") - .replaceAll(">", ">") - .replaceAll('"', """) - .replaceAll("'", "'"); } From 4b1ad3ae1518caff2265516d0b0ff5908bd1f6eb Mon Sep 17 00:00:00 2001 From: Mux Date: Sat, 7 Feb 2026 09:53:36 +0100 Subject: [PATCH 04/22] refactor: use shared OAuth utilities in muxGovernorOauthService Replace duplicated OAuth infrastructure with the new shared modules: - Import createDeferred, renderOAuthCallbackHtml from oauthUtils - Replace manual http.createServer with startLoopbackServer - Replace desktopFlows Map + lifecycle methods with OAuthFlowManager - Remove local closeServer, escapeHtml, createDeferred functions - Remove DesktopFlow interface, handleDesktopCallback, finishDesktopFlow Governor-specific behavior preserved: - Dynamic governorOrigin from user input for authorize/exchange URLs - Custom 'Enrollment complete/failed' branding via renderHtml option - policyService.refreshNow() callback after successful enrollment - ServerFlow pattern (startServerFlow, handleServerCallbackAndExchange) left unchanged as it doesn't use loopback All public method signatures remain identical. Reduces file from 484 to 293 lines (-40%). --- src/node/services/muxGovernorOauthService.ts | 284 +++---------------- 1 file changed, 46 insertions(+), 238 deletions(-) diff --git a/src/node/services/muxGovernorOauthService.ts b/src/node/services/muxGovernorOauthService.ts index 95282cd47b..13b79534a2 100644 --- a/src/node/services/muxGovernorOauthService.ts +++ b/src/node/services/muxGovernorOauthService.ts @@ -7,7 +7,6 @@ */ import * as crypto from "crypto"; -import * as http from "http"; import type { Result } from "@/common/types/result"; import { Err, Ok } from "@/common/types/result"; import { @@ -20,23 +19,12 @@ import type { Config } from "@/node/config"; import type { PolicyService } from "@/node/services/policyService"; import type { WindowService } from "@/node/services/windowService"; import { log } from "@/node/services/log"; +import { createDeferred, renderOAuthCallbackHtml } from "@/node/utils/oauthUtils"; +import { startLoopbackServer } from "@/node/utils/oauthLoopbackServer"; +import { OAuthFlowManager } from "@/node/utils/oauthFlowManager"; const DEFAULT_DESKTOP_TIMEOUT_MS = 5 * 60 * 1000; const DEFAULT_SERVER_TIMEOUT_MS = 10 * 60 * 1000; -const COMPLETED_DESKTOP_FLOW_TTL_MS = 60 * 1000; - -interface DesktopFlow { - flowId: string; - governorOrigin: string; - authorizeUrl: string; - redirectUri: string; - server: http.Server; - timeout: ReturnType; - cleanupTimeout: ReturnType | null; - resultPromise: Promise>; - resolveResult: (result: Result) => void; - settled: boolean; -} interface ServerFlow { state: string; @@ -44,22 +32,8 @@ interface ServerFlow { expiresAtMs: number; } -function closeServer(server: http.Server): Promise { - return new Promise((resolve) => { - server.close(() => resolve()); - }); -} - -function createDeferred(): { promise: Promise; resolve: (value: T) => void } { - let resolve!: (value: T) => void; - const promise = new Promise((res) => { - resolve = res; - }); - return { promise, resolve }; -} - export class MuxGovernorOauthService { - private readonly desktopFlows = new Map(); + private readonly desktopFlows = new OAuthFlowManager(); private readonly serverFlows = new Map(); constructor( @@ -82,125 +56,76 @@ export class MuxGovernorOauthService { const flowId = crypto.randomUUID(); - const { promise: resultPromise, resolve: resolveResult } = - createDeferred>(); - - const server = http.createServer((req, res) => { - const reqUrl = req.url ?? "/"; - const url = new URL(reqUrl, "http://localhost"); - - if (req.method !== "GET" || url.pathname !== "/callback") { - res.statusCode = 404; - res.end("Not found"); - return; - } - - const state = url.searchParams.get("state"); - if (!state || state !== flowId) { - res.statusCode = 400; - res.setHeader("Content-Type", "text/html"); - res.end("

Invalid OAuth state

"); - return; - } - - const code = url.searchParams.get("code"); - const error = url.searchParams.get("error"); - const errorDescription = url.searchParams.get("error_description") ?? undefined; - - void this.handleDesktopCallback({ - flowId, - governorOrigin, - code, - error, - errorDescription, - res, - }); - }); - + let loopback: Awaited>; try { - await new Promise((resolve, reject) => { - server.once("error", reject); - server.listen(0, "127.0.0.1", () => resolve()); + loopback = await startLoopbackServer({ + expectedState: flowId, + renderHtml: (r) => + renderOAuthCallbackHtml({ + title: r.success ? "Enrollment complete" : "Enrollment failed", + message: r.success + ? "You can return to Mux. You may now close this tab." + : (r.error ?? "Unknown error"), + success: r.success, + }), }); } catch (error) { const message = error instanceof Error ? error.message : String(error); return Err(`Failed to start OAuth callback listener: ${message}`); } - const address = server.address(); - if (!address || typeof address === "string") { - return Err("Failed to determine OAuth callback listener port"); - } - - const redirectUri = `http://127.0.0.1:${address.port}/callback`; const authorizeUrl = buildGovernorAuthorizeUrl({ governorOrigin, - redirectUri, + redirectUri: loopback.redirectUri, state: flowId, }); - const timeout = setTimeout(() => { - void this.finishDesktopFlow(flowId, Err("Timed out waiting for OAuth callback")); - }, DEFAULT_DESKTOP_TIMEOUT_MS); + const resultDeferred = createDeferred>(); - this.desktopFlows.set(flowId, { - flowId, - governorOrigin, - authorizeUrl, - redirectUri, - server, - timeout, - cleanupTimeout: null, - resultPromise, - resolveResult, - settled: false, + this.desktopFlows.register(flowId, { + server: loopback.server, + resultDeferred, + timeoutHandle: null, + }); + + // Wire the loopback result to drive token exchange + finish. + void loopback.result.then(async (callbackResult) => { + if (!callbackResult.success) { + await this.desktopFlows.finish( + flowId, + Err(`Mux Governor OAuth error: ${callbackResult.error}`) + ); + return; + } + + const result = await this.handleCallbackAndExchange({ + state: flowId, + governorOrigin, + code: callbackResult.data.code, + error: null, + }); + + await this.desktopFlows.finish(flowId, result); }); log.debug( `Mux Governor OAuth desktop flow started (flowId=${flowId}, origin=${governorOrigin})` ); - return Ok({ flowId, authorizeUrl, redirectUri }); + return Ok({ flowId, authorizeUrl, redirectUri: loopback.redirectUri }); } async waitForDesktopFlow( flowId: string, opts?: { timeoutMs?: number } ): Promise> { - const flow = this.desktopFlows.get(flowId); - if (!flow) { - return Err("OAuth flow not found"); - } - - const timeoutMs = opts?.timeoutMs ?? DEFAULT_DESKTOP_TIMEOUT_MS; - - let timeoutHandle: ReturnType | null = null; - const timeoutPromise = new Promise>((resolve) => { - timeoutHandle = setTimeout(() => { - resolve(Err("Timed out waiting for OAuth callback")); - }, timeoutMs); - }); - - const result = await Promise.race([flow.resultPromise, timeoutPromise]); - - if (timeoutHandle !== null) { - clearTimeout(timeoutHandle); - } - - if (!result.success) { - // Ensure listener is closed on timeout/errors. - void this.finishDesktopFlow(flowId, result); - } - - return result; + return this.desktopFlows.waitFor(flowId, opts?.timeoutMs ?? DEFAULT_DESKTOP_TIMEOUT_MS); } async cancelDesktopFlow(flowId: string): Promise { - const flow = this.desktopFlows.get(flowId); - if (!flow) return; - + if (!this.desktopFlows.has(flowId)) return; log.debug(`Mux Governor OAuth desktop flow cancelled (flowId=${flowId})`); - await this.finishDesktopFlow(flowId, Err("OAuth flow cancelled")); + await this.desktopFlows.cancel(flowId); } startServerFlow(input: { @@ -278,94 +203,10 @@ export class MuxGovernorOauthService { } async dispose(): Promise { - // Best-effort: cancel all in-flight flows. - const flowIds = [...this.desktopFlows.keys()]; - await Promise.all(flowIds.map((id) => this.finishDesktopFlow(id, Err("App shutting down")))); - - for (const flow of this.desktopFlows.values()) { - clearTimeout(flow.timeout); - if (flow.cleanupTimeout !== null) { - clearTimeout(flow.cleanupTimeout); - } - } - - this.desktopFlows.clear(); + await this.desktopFlows.shutdownAll(); this.serverFlows.clear(); } - private async handleDesktopCallback(input: { - flowId: string; - governorOrigin: string; - code: string | null; - error: string | null; - errorDescription?: string; - res: http.ServerResponse; - }): Promise { - const flow = this.desktopFlows.get(input.flowId); - if (!flow || flow.settled) { - input.res.statusCode = 409; - input.res.setHeader("Content-Type", "text/html"); - input.res.end("

OAuth flow already completed

"); - return; - } - - log.debug(`Mux Governor OAuth callback received (flowId=${input.flowId})`); - - const result = await this.handleCallbackAndExchange({ - state: input.flowId, - governorOrigin: input.governorOrigin, - code: input.code, - error: input.error, - errorDescription: input.errorDescription, - }); - - const title = result.success ? "Enrollment complete" : "Enrollment failed"; - const description = result.success - ? "You can return to Mux. You may now close this tab." - : escapeHtml(result.error); - - const html = ` - - - - - - ${title} - - - -

${title}

-

${description}

- ${ - result.success - ? '

Mux should now be in the foreground. You can close this tab.

' - : '

You can close this tab.

' - } - - -`; - - input.res.setHeader("Content-Type", "text/html"); - if (!result.success) { - input.res.statusCode = 400; - } - - input.res.end(html); - - await this.finishDesktopFlow(input.flowId, result); - } - private async handleCallbackAndExchange(input: { state: string; governorOrigin: string; @@ -448,37 +289,4 @@ export class MuxGovernorOauthService { } } - private async finishDesktopFlow(flowId: string, result: Result): Promise { - const flow = this.desktopFlows.get(flowId); - if (!flow || flow.settled) return; - - flow.settled = true; - clearTimeout(flow.timeout); - - try { - flow.resolveResult(result); - - // Stop accepting new connections. - await closeServer(flow.server); - } catch (error) { - log.debug("Failed to close OAuth callback listener:", error); - } finally { - // Keep the completed flow around briefly so callers can still await the result. - if (flow.cleanupTimeout !== null) { - clearTimeout(flow.cleanupTimeout); - } - flow.cleanupTimeout = setTimeout(() => { - this.desktopFlows.delete(flowId); - }, COMPLETED_DESKTOP_FLOW_TTL_MS); - } - } -} - -function escapeHtml(input: string): string { - return input - .replaceAll("&", "&") - .replaceAll("<", "<") - .replaceAll(">", ">") - .replaceAll('"', """) - .replaceAll("'", "'"); } From 2999625a912b824bc800b57538d977c8425f391b Mon Sep 17 00:00:00 2001 From: Mux Date: Sat, 7 Feb 2026 09:57:41 +0100 Subject: [PATCH 05/22] refactor(mcpOauthService): use shared OAuth utilities from oauthUtils MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Replace local duplicated helpers with imports from @/node/utils/oauthUtils: - createDeferred (~7 lines) - closeServer (~4 lines) - escapeHtml (~8 lines) - Inline HTML template in handleDesktopCallback → renderOAuthCallbackHtml MCP's unique patterns (DesktopFlow lifecycle, two-phase callback handling, state-mismatch flow termination, telemetry) are preserved unchanged. Net reduction: 26 lines (1492 → 1466). --- src/node/services/mcpOauthService.ts | 56 ++++++++-------------------- 1 file changed, 15 insertions(+), 41 deletions(-) diff --git a/src/node/services/mcpOauthService.ts b/src/node/services/mcpOauthService.ts index e7105f23a3..4551898c75 100644 --- a/src/node/services/mcpOauthService.ts +++ b/src/node/services/mcpOauthService.ts @@ -21,6 +21,12 @@ import type { } from "@/common/types/mcpOauth"; import { stripTrailingSlashes } from "@/node/utils/pathUtils"; import { MutexMap } from "@/node/utils/concurrency/mutexMap"; +import { + closeServer, + createDeferred, + escapeHtml, + renderOAuthCallbackHtml, +} from "@/node/utils/oauthUtils"; const DEFAULT_DESKTOP_TIMEOUT_MS = 5 * 60 * 1000; const DEFAULT_SERVER_TIMEOUT_MS = 10 * 60 * 1000; @@ -114,20 +120,6 @@ interface DesktopFlow extends OAuthFlowBase { type ServerFlow = OAuthFlowBase; -function closeServer(server: http.Server): Promise { - return new Promise((resolve) => { - server.close(() => resolve()); - }); -} - -function createDeferred(): { promise: Promise; resolve: (value: T) => void } { - let resolve!: (value: T) => void; - const promise = new Promise((res) => { - resolve = res; - }); - return { promise, resolve }; -} - function isPlainObject(value: unknown): value is Record { return Boolean(value) && typeof value === "object" && !Array.isArray(value); } @@ -1132,29 +1124,20 @@ export class McpOauthService { errorDescription: input.errorDescription, }); - const title = result.success ? "Login complete" : "Login failed"; - const description = result.success - ? "You can return to Mux. You may now close this tab." - : escapeHtml(result.error); - input.res.setHeader("Content-Type", "text/html"); if (!result.success) { input.res.statusCode = 400; } - input.res.end(` - - - - - - ${title} - - -

${title}

-

${description}

- -`); + input.res.end( + renderOAuthCallbackHtml({ + title: result.success ? "Login complete" : "Login failed", + message: result.success + ? "You can return to Mux. You may now close this tab." + : escapeHtml(result.error), + success: result.success, + }) + ); await this.finishDesktopFlow(input.flowId, result); } @@ -1481,12 +1464,3 @@ export class McpOauthService { }); } } - -function escapeHtml(input: string): string { - return input - .replaceAll("&", "&") - .replaceAll("<", "<") - .replaceAll(">", ">") - .replaceAll('"', """) - .replaceAll("'", "'"); -} From aabe338d035960dfe4c0ce1c23f5d841251fd21a Mon Sep 17 00:00:00 2001 From: Mux Date: Sat, 7 Feb 2026 10:03:39 +0100 Subject: [PATCH 06/22] refactor: use shared OAuth utilities in codexOauthService MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Replace local createDeferred, closeServer, escapeHtml, isLoopbackAddress with imports from shared oauthUtils/oauthLoopbackServer - Replace DesktopFlow interface + Map with OAuthFlowManager - Replace manual http.createServer with startLoopbackServer - Remove handleDesktopCallback (logic moved to background task in startDesktopFlow) - Remove finishDesktopFlow (delegated to OAuthFlowManager.finish) - Simplify dispose() with shutdownAll() - Keep DeviceFlow and all device flow methods unchanged 931 → 749 lines (-182) --- src/node/services/codexOauthService.ts | 282 +++++-------------------- 1 file changed, 50 insertions(+), 232 deletions(-) diff --git a/src/node/services/codexOauthService.ts b/src/node/services/codexOauthService.ts index a0e2fcebc3..a81c593e75 100644 --- a/src/node/services/codexOauthService.ts +++ b/src/node/services/codexOauthService.ts @@ -1,5 +1,4 @@ import * as crypto from "crypto"; -import * as http from "http"; import type { Result } from "@/common/types/result"; import { Err, Ok } from "@/common/types/result"; import { @@ -24,26 +23,14 @@ import { parseCodexOauthAuth, type CodexOauthAuth, } from "@/node/utils/codexOauthAuth"; +import { createDeferred } from "@/node/utils/oauthUtils"; +import { startLoopbackServer } from "@/node/utils/oauthLoopbackServer"; +import { OAuthFlowManager } from "@/node/utils/oauthFlowManager"; const DEFAULT_DESKTOP_TIMEOUT_MS = 5 * 60 * 1000; const DEFAULT_DEVICE_TIMEOUT_MS = 15 * 60 * 1000; const COMPLETED_FLOW_TTL_MS = 60 * 1000; -interface DesktopFlow { - flowId: string; - authorizeUrl: string; - redirectUri: string; - codeVerifier: string; - - server: http.Server; - timeout: ReturnType; - cleanupTimeout: ReturnType | null; - - resultPromise: Promise>; - resolveResult: (result: Result) => void; - settled: boolean; -} - interface DeviceFlow { flowId: string; deviceAuthId: string; @@ -63,20 +50,6 @@ interface DeviceFlow { settled: boolean; } -function closeServer(server: http.Server): Promise { - return new Promise((resolve) => { - server.close(() => resolve()); - }); -} - -function createDeferred(): { promise: Promise; resolve: (value: T) => void } { - let resolve!: (value: T) => void; - const promise = new Promise((res) => { - resolve = res; - }); - return { promise, resolve }; -} - function sha256Base64Url(value: string): string { return crypto.createHash("sha256").update(value).digest().toString("base64url"); } @@ -89,17 +62,6 @@ function isPlainObject(value: unknown): value is Record { return Boolean(value) && typeof value === "object" && !Array.isArray(value); } -function isLoopbackAddress(address: string | undefined): boolean { - if (!address) return false; - - // Node may normalize IPv4 loopback to an IPv6-mapped address. - if (address === "::ffff:127.0.0.1") { - return true; - } - - return address === "127.0.0.1" || address === "::1"; -} - function parseOptionalNumber(value: unknown): number | null { if (typeof value === "number" && Number.isFinite(value)) { return value; @@ -133,7 +95,7 @@ function isInvalidGrantError(errorText: string): boolean { } export class CodexOauthService { - private readonly desktopFlows = new Map(); + private readonly desktopFlows = new OAuthFlowManager(); private readonly deviceFlows = new Map(); private readonly refreshMutex = new AsyncMutex(); @@ -159,81 +121,64 @@ export class CodexOauthService { const codeVerifier = randomBase64Url(); const codeChallenge = sha256Base64Url(codeVerifier); - - const { promise: resultPromise, resolve: resolveResult } = - createDeferred>(); - const redirectUri = CODEX_OAUTH_BROWSER_REDIRECT_URI; - const server = http.createServer((req, res) => { - if (!isLoopbackAddress(req.socket.remoteAddress)) { - res.statusCode = 403; - res.end("Forbidden"); - return; - } - - const reqUrl = req.url ?? "/"; - const url = new URL(reqUrl, "http://localhost"); - - if (req.method !== "GET" || url.pathname !== "/auth/callback") { - res.statusCode = 404; - res.end("Not found"); - return; - } - - const state = url.searchParams.get("state"); - if (!state || state !== flowId) { - res.statusCode = 400; - res.setHeader("Content-Type", "text/html"); - res.end("

Invalid OAuth state

"); - return; - } - - const code = url.searchParams.get("code"); - const error = url.searchParams.get("error"); - const errorDescription = url.searchParams.get("error_description") ?? undefined; - - void this.handleDesktopCallback({ - flowId, - code, - error, - errorDescription, - res, - }); - }); - + let loopback: Awaited>; try { - await new Promise((resolve, reject) => { - server.once("error", reject); - server.listen(1455, "localhost", () => resolve()); + loopback = await startLoopbackServer({ + port: 1455, + host: "localhost", + callbackPath: "/auth/callback", + validateLoopback: true, + expectedState: flowId, }); } catch (error) { const message = error instanceof Error ? error.message : String(error); return Err(`Failed to start OAuth callback listener: ${message}`); } + const resultDeferred = createDeferred>(); + + this.desktopFlows.register(flowId, { + server: loopback.server, + resultDeferred, + timeoutHandle: null, + }); + const authorizeUrl = buildCodexAuthorizeUrl({ redirectUri, state: flowId, codeChallenge, }); - const timeout = setTimeout(() => { - void this.finishDesktopFlow(flowId, Err("Timed out waiting for OAuth callback")); - }, DEFAULT_DESKTOP_TIMEOUT_MS); + // Background task: wait for the loopback callback, exchange code for tokens, + // then finish the flow. Races against resultDeferred (which resolves on + // cancel/timeout) so the task exits cleanly if the flow is cancelled. + void (async () => { + const callbackResult = await Promise.race([ + loopback.result, + resultDeferred.promise.then(() => null), + ]); - this.desktopFlows.set(flowId, { - flowId, - authorizeUrl, - redirectUri, - codeVerifier, - server, - timeout, - cleanupTimeout: null, - resultPromise, - resolveResult, - settled: false, - }); + // null means the flow was finished externally (cancel/timeout). + if (!callbackResult) return; + + if (!callbackResult.success) { + await this.desktopFlows.finish(flowId, Err(callbackResult.error)); + return; + } + + const exchangeResult = await this.handleDesktopCallbackAndExchange({ + flowId, + redirectUri, + codeVerifier, + code: callbackResult.data.code, + error: null, + errorDescription: undefined, + }); + + await this.desktopFlows.finish(flowId, exchangeResult); + })(); log.debug(`[Codex OAuth] Desktop flow started (flowId=${flowId})`); @@ -244,40 +189,12 @@ export class CodexOauthService { flowId: string, opts?: { timeoutMs?: number } ): Promise> { - const flow = this.desktopFlows.get(flowId); - if (!flow) { - return Err("OAuth flow not found"); - } - - const timeoutMs = opts?.timeoutMs ?? DEFAULT_DESKTOP_TIMEOUT_MS; - - let timeoutHandle: ReturnType | null = null; - const timeoutPromise = new Promise>((resolve) => { - timeoutHandle = setTimeout(() => { - resolve(Err("Timed out waiting for OAuth callback")); - }, timeoutMs); - }); - - const result = await Promise.race([flow.resultPromise, timeoutPromise]); - - if (timeoutHandle !== null) { - clearTimeout(timeoutHandle); - } - - if (!result.success) { - // Ensure listener is closed on timeout/errors. - void this.finishDesktopFlow(flowId, result); - } - - return result; + return this.desktopFlows.waitFor(flowId, opts?.timeoutMs ?? DEFAULT_DESKTOP_TIMEOUT_MS); } async cancelDesktopFlow(flowId: string): Promise { - const flow = this.desktopFlows.get(flowId); - if (!flow) return; - log.debug(`[Codex OAuth] Desktop flow cancelled (flowId=${flowId})`); - await this.finishDesktopFlow(flowId, Err("OAuth flow cancelled")); + await this.desktopFlows.cancel(flowId); } async startDeviceFlow(): Promise< @@ -414,19 +331,11 @@ export class CodexOauthService { } async dispose(): Promise { - const desktopIds = [...this.desktopFlows.keys()]; - await Promise.all(desktopIds.map((id) => this.finishDesktopFlow(id, Err("App shutting down")))); + await this.desktopFlows.shutdownAll(); const deviceIds = [...this.deviceFlows.keys()]; await Promise.all(deviceIds.map((id) => this.finishDeviceFlow(id, Err("App shutting down")))); - for (const flow of this.desktopFlows.values()) { - clearTimeout(flow.timeout); - if (flow.cleanupTimeout !== null) { - clearTimeout(flow.cleanupTimeout); - } - } - for (const flow of this.deviceFlows.values()) { clearTimeout(flow.timeout); if (flow.cleanupTimeout !== null) { @@ -434,7 +343,6 @@ export class CodexOauthService { } } - this.desktopFlows.clear(); this.deviceFlows.clear(); } @@ -458,67 +366,6 @@ export class CodexOauthService { return result; } - private async handleDesktopCallback(input: { - flowId: string; - code: string | null; - error: string | null; - errorDescription?: string; - res: http.ServerResponse; - }): Promise { - const flow = this.desktopFlows.get(input.flowId); - if (!flow || flow.settled) { - input.res.statusCode = 409; - input.res.setHeader("Content-Type", "text/html"); - input.res.end("

OAuth flow already completed

"); - return; - } - - log.debug(`[Codex OAuth] Desktop callback received (flowId=${input.flowId})`); - - const result = await this.handleDesktopCallbackAndExchange({ - flowId: input.flowId, - redirectUri: flow.redirectUri, - codeVerifier: flow.codeVerifier, - code: input.code, - error: input.error, - errorDescription: input.errorDescription, - }); - - const title = result.success ? "Login complete" : "Login failed"; - const description = result.success - ? "You can return to Mux. You may now close this tab." - : escapeHtml(result.error); - - input.res.setHeader("Content-Type", "text/html"); - if (!result.success) { - input.res.statusCode = 400; - } - - input.res.end(` - - - - - - ${title} - - -

${title}

-

${description}

- - -`); - - await this.finishDesktopFlow(input.flowId, result); - } - private async handleDesktopCallbackAndExchange(input: { flowId: string; redirectUri: string; @@ -845,28 +692,6 @@ export class CodexOauthService { } } - private async finishDesktopFlow(flowId: string, result: Result): Promise { - const flow = this.desktopFlows.get(flowId); - if (!flow || flow.settled) return; - - flow.settled = true; - clearTimeout(flow.timeout); - - try { - flow.resolveResult(result); - await closeServer(flow.server); - } catch (error) { - log.debug("[Codex OAuth] Failed to close OAuth callback listener:", error); - } finally { - if (flow.cleanupTimeout !== null) { - clearTimeout(flow.cleanupTimeout); - } - flow.cleanupTimeout = setTimeout(() => { - this.desktopFlows.delete(flowId); - }, COMPLETED_FLOW_TTL_MS); - } - } - private finishDeviceFlow(flowId: string, result: Result): Promise { const flow = this.deviceFlows.get(flowId); if (!flow || flow.settled) { @@ -921,11 +746,4 @@ async function sleepWithAbort(ms: number, signal: AbortSignal): Promise { }); } -function escapeHtml(input: string): string { - return input - .replaceAll("&", "&") - .replaceAll("<", "<") - .replaceAll(">", ">") - .replaceAll('"', """) - .replaceAll("'", "'"); -} + From 60b6e3f572faec47649bc3e87dc265dc41d58dca Mon Sep 17 00:00:00 2001 From: Mux Date: Sat, 7 Feb 2026 10:10:28 +0100 Subject: [PATCH 07/22] test: add unit tests for shared OAuth utility modules Add 3 test files covering the shared OAuth utilities: - oauthUtils.test.ts (21 tests): createDeferred, closeServer, escapeHtml, renderOAuthCallbackHtml (XSS escaping, auto-close, extraHead) - oauthLoopbackServer.test.ts (9 tests): server start, valid callback, state mismatch, error params, missing code, 404, cancel, custom path/renderer - oauthFlowManager.test.ts (13 tests): register/get/has, waitFor with resolve/timeout/missing, cancel, finish (idempotent, null server, timeout cleanup), shutdownAll Uses real HTTP requests for loopback server tests and mock http.Server for flow manager tests. --- src/node/utils/oauthFlowManager.test.ts | 220 +++++++++++++++++++++ src/node/utils/oauthLoopbackServer.test.ts | 171 ++++++++++++++++ src/node/utils/oauthUtils.test.ts | 195 ++++++++++++++++++ 3 files changed, 586 insertions(+) create mode 100644 src/node/utils/oauthFlowManager.test.ts create mode 100644 src/node/utils/oauthLoopbackServer.test.ts create mode 100644 src/node/utils/oauthUtils.test.ts diff --git a/src/node/utils/oauthFlowManager.test.ts b/src/node/utils/oauthFlowManager.test.ts new file mode 100644 index 0000000000..7813fc4178 --- /dev/null +++ b/src/node/utils/oauthFlowManager.test.ts @@ -0,0 +1,220 @@ +import type http from "node:http"; +import { describe, it, expect, beforeEach } from "bun:test"; +import type { Result } from "@/common/types/result"; +import { Ok, Err } from "@/common/types/result"; +import { createDeferred } from "@/node/utils/oauthUtils"; +import type { OAuthFlowEntry } from "./oauthFlowManager"; +import { OAuthFlowManager } from "./oauthFlowManager"; + +// --------------------------------------------------------------------------- +// Mock http.Server +// --------------------------------------------------------------------------- + +/** + * Minimal mock that satisfies the `http.Server` contract used by + * OAuthFlowManager: only `close()` is called (via `closeServer`). + */ +function createMockServer(): http.Server { + const mock = { + close: (cb?: (err?: Error) => void) => { + if (cb) cb(); + return mock; + }, + } as unknown as http.Server; + return mock; +} + +/** Create a fresh OAuthFlowEntry backed by a mock server. */ +function createFlowEntry(server?: http.Server | null): OAuthFlowEntry { + return { + server: server === undefined ? createMockServer() : server, + resultDeferred: createDeferred>(), + timeoutHandle: null, + }; +} + +// --------------------------------------------------------------------------- +// Tests +// --------------------------------------------------------------------------- + +describe("OAuthFlowManager", () => { + let manager: OAuthFlowManager; + + beforeEach(() => { + manager = new OAuthFlowManager(); + }); + + // ----------------------------------------------------------------------- + // register / get / has + // ----------------------------------------------------------------------- + + describe("register / get / has", () => { + it("registers and retrieves a flow entry", () => { + const entry = createFlowEntry(); + manager.register("f1", entry); + + expect(manager.has("f1")).toBe(true); + expect(manager.get("f1")).toBe(entry); + }); + + it("returns undefined for unregistered flows", () => { + expect(manager.has("nope")).toBe(false); + expect(manager.get("nope")).toBeUndefined(); + }); + }); + + // ----------------------------------------------------------------------- + // waitFor + // ----------------------------------------------------------------------- + + describe("waitFor", () => { + it("resolves when the deferred resolves with Ok", async () => { + const entry = createFlowEntry(); + manager.register("f1", entry); + + // Resolve the deferred in the background. + setTimeout(() => entry.resultDeferred.resolve(Ok(undefined)), 5); + + const result = await manager.waitFor("f1", 5_000); + expect(result.success).toBe(true); + }); + + it("times out when the deferred does not resolve in time", async () => { + const entry = createFlowEntry(); + manager.register("f1", entry); + + // Use a very short timeout so the test is fast. + const result = await manager.waitFor("f1", 20); + + expect(result.success).toBe(false); + if (!result.success) { + expect(result.error).toContain("Timed out"); + } + }); + + it("returns Err when flow ID is not found", async () => { + const result = await manager.waitFor("missing", 1_000); + expect(result.success).toBe(false); + if (!result.success) { + expect(result.error).toContain("not found"); + } + }); + }); + + // ----------------------------------------------------------------------- + // cancel + // ----------------------------------------------------------------------- + + describe("cancel", () => { + it("resolves the deferred with Err and removes the flow", async () => { + const entry = createFlowEntry(); + manager.register("f1", entry); + + await manager.cancel("f1"); + + const result = await entry.resultDeferred.promise; + expect(result.success).toBe(false); + if (!result.success) { + expect(result.error).toContain("cancelled"); + } + expect(manager.has("f1")).toBe(false); + }); + + it("is a no-op for non-existent flows", async () => { + // Should not throw. + await manager.cancel("nope"); + }); + }); + + // ----------------------------------------------------------------------- + // finish + // ----------------------------------------------------------------------- + + describe("finish", () => { + it("resolves the deferred, removes the flow, and closes the server", async () => { + let serverClosed = false; + const mockServer = { + close: (cb?: (err?: Error) => void) => { + serverClosed = true; + if (cb) cb(); + return mockServer; + }, + } as unknown as http.Server; + + const entry = createFlowEntry(mockServer); + manager.register("f1", entry); + + await manager.finish("f1", Ok(undefined)); + + const result = await entry.resultDeferred.promise; + expect(result.success).toBe(true); + expect(manager.has("f1")).toBe(false); + expect(serverClosed).toBe(true); + }); + + it("is idempotent — second call is a no-op", async () => { + const entry = createFlowEntry(); + manager.register("f1", entry); + + await manager.finish("f1", Ok(undefined)); + // Second call should not throw. + await manager.finish("f1", Err("should be ignored")); + + const result = await entry.resultDeferred.promise; + // Should still be the first result. + expect(result.success).toBe(true); + }); + + it("works when server is null (device-code flows)", async () => { + const entry = createFlowEntry(null); + manager.register("f1", entry); + + await manager.finish("f1", Ok(undefined)); + + const result = await entry.resultDeferred.promise; + expect(result.success).toBe(true); + expect(manager.has("f1")).toBe(false); + }); + + it("clears the timeout handle", async () => { + const entry = createFlowEntry(); + // Simulate a stored timeout handle. + entry.timeoutHandle = setTimeout(() => {}, 60_000); + manager.register("f1", entry); + + await manager.finish("f1", Ok(undefined)); + + expect(manager.has("f1")).toBe(false); + }); + }); + + // ----------------------------------------------------------------------- + // shutdownAll + // ----------------------------------------------------------------------- + + describe("shutdownAll", () => { + it("finishes all registered flows", async () => { + const entry1 = createFlowEntry(); + const entry2 = createFlowEntry(); + manager.register("f1", entry1); + manager.register("f2", entry2); + + await manager.shutdownAll(); + + expect(manager.has("f1")).toBe(false); + expect(manager.has("f2")).toBe(false); + + const r1 = await entry1.resultDeferred.promise; + const r2 = await entry2.resultDeferred.promise; + expect(r1.success).toBe(false); + expect(r2.success).toBe(false); + if (!r1.success) expect(r1.error).toContain("shutting down"); + if (!r2.success) expect(r2.error).toContain("shutting down"); + }); + + it("is a no-op when there are no flows", async () => { + // Should not throw. + await manager.shutdownAll(); + }); + }); +}); diff --git a/src/node/utils/oauthLoopbackServer.test.ts b/src/node/utils/oauthLoopbackServer.test.ts new file mode 100644 index 0000000000..d247f29299 --- /dev/null +++ b/src/node/utils/oauthLoopbackServer.test.ts @@ -0,0 +1,171 @@ +import http from "node:http"; +import { describe, it, expect, afterEach } from "bun:test"; +import type { LoopbackServer } from "./oauthLoopbackServer"; +import { startLoopbackServer } from "./oauthLoopbackServer"; + +// --------------------------------------------------------------------------- +// Helpers +// --------------------------------------------------------------------------- + +/** Extract the port from a redirectUri like http://127.0.0.1:12345/callback */ +function portFromUri(uri: string): number { + return new URL(uri).port ? Number(new URL(uri).port) : 80; +} + +/** Simple GET helper that returns { status, body }. */ +async function httpGet(url: string): Promise<{ status: number; body: string }> { + return new Promise((resolve, reject) => { + http.get(url, (res) => { + let body = ""; + res.on("data", (chunk: Buffer) => { + body += chunk.toString(); + }); + res.on("end", () => resolve({ status: res.statusCode ?? 0, body })); + }).on("error", reject); + }); +} + +// --------------------------------------------------------------------------- +// Tests +// --------------------------------------------------------------------------- + +describe("startLoopbackServer", () => { + let loopback: LoopbackServer | undefined; + + afterEach(async () => { + // Ensure the server is always cleaned up. + if (loopback?.server.listening) { + await loopback.close(); + } + loopback = undefined; + }); + + it("starts a server and provides a redirectUri with the listening port", async () => { + loopback = await startLoopbackServer({ expectedState: "s1" }); + + expect(loopback.redirectUri).toMatch(/^http:\/\/127\.0\.0\.1:\d+\/callback$/); + const port = portFromUri(loopback.redirectUri); + expect(port).toBeGreaterThan(0); + expect(loopback.server.listening).toBe(true); + }); + + it("resolves with Ok({code, state}) on a valid callback", async () => { + loopback = await startLoopbackServer({ expectedState: "state123" }); + + const callbackUrl = `${loopback.redirectUri}?state=state123&code=authcode456`; + const res = await httpGet(callbackUrl); + + expect(res.status).toBe(200); + expect(res.body).toContain(""); + + const result = await loopback.result; + expect(result.success).toBe(true); + if (result.success) { + expect(result.data.code).toBe("authcode456"); + expect(result.data.state).toBe("state123"); + } + }); + + it("resolves with Err on state mismatch", async () => { + loopback = await startLoopbackServer({ expectedState: "good" }); + + const callbackUrl = `${loopback.redirectUri}?state=bad&code=c`; + const res = await httpGet(callbackUrl); + + expect(res.status).toBe(400); + + const result = await loopback.result; + expect(result.success).toBe(false); + if (!result.success) { + expect(result.error).toContain("Invalid OAuth state"); + } + }); + + it("resolves with Err when provider returns an error param", async () => { + loopback = await startLoopbackServer({ expectedState: "s1" }); + + const callbackUrl = `${loopback.redirectUri}?state=s1&error=access_denied&error_description=User+denied`; + const res = await httpGet(callbackUrl); + + expect(res.status).toBe(400); + + const result = await loopback.result; + expect(result.success).toBe(false); + if (!result.success) { + expect(result.error).toContain("access_denied"); + expect(result.error).toContain("User denied"); + } + }); + + it("resolves with Err when code is missing", async () => { + loopback = await startLoopbackServer({ expectedState: "s1" }); + + const callbackUrl = `${loopback.redirectUri}?state=s1`; + const res = await httpGet(callbackUrl); + + expect(res.status).toBe(400); + + const result = await loopback.result; + expect(result.success).toBe(false); + if (!result.success) { + expect(result.error).toContain("Missing authorization code"); + } + }); + + it("returns 404 for the wrong path", async () => { + loopback = await startLoopbackServer({ expectedState: "s1" }); + + const port = portFromUri(loopback.redirectUri); + const res = await httpGet(`http://127.0.0.1:${port}/wrong`); + + expect(res.status).toBe(404); + expect(res.body).toContain("Not found"); + }); + + it("cancel resolves result with Err and closes the server", async () => { + loopback = await startLoopbackServer({ expectedState: "s1" }); + expect(loopback.server.listening).toBe(true); + + await loopback.cancel(); + + const result = await loopback.result; + expect(result.success).toBe(false); + if (!result.success) { + expect(result.error).toContain("cancelled"); + } + expect(loopback.server.listening).toBe(false); + }); + + it("uses a custom callbackPath", async () => { + loopback = await startLoopbackServer({ + expectedState: "s1", + callbackPath: "/oauth/done", + }); + + expect(loopback.redirectUri).toContain("/oauth/done"); + + const callbackUrl = `${loopback.redirectUri}?state=s1&code=c1`; + const res = await httpGet(callbackUrl); + expect(res.status).toBe(200); + + const result = await loopback.result; + expect(result.success).toBe(true); + }); + + it("uses a custom renderHtml function", async () => { + const customHtml = "Custom!"; + loopback = await startLoopbackServer({ + expectedState: "s1", + renderHtml: () => customHtml, + }); + + const callbackUrl = `${loopback.redirectUri}?state=s1&code=c1`; + const res = await httpGet(callbackUrl); + + expect(res.status).toBe(200); + expect(res.body).toBe(customHtml); + + const result = await loopback.result; + expect(result.success).toBe(true); + }); +}); diff --git a/src/node/utils/oauthUtils.test.ts b/src/node/utils/oauthUtils.test.ts new file mode 100644 index 0000000000..2490d86791 --- /dev/null +++ b/src/node/utils/oauthUtils.test.ts @@ -0,0 +1,195 @@ +import http from "node:http"; +import { describe, it, expect, afterEach } from "bun:test"; +import { createDeferred, closeServer, escapeHtml, renderOAuthCallbackHtml } from "./oauthUtils"; + +// --------------------------------------------------------------------------- +// createDeferred +// --------------------------------------------------------------------------- + +describe("createDeferred", () => { + it("resolves with the value passed to resolve()", async () => { + const d = createDeferred(); + d.resolve("hello"); + expect(await d.promise).toBe("hello"); + }); + + it("works with numeric types", async () => { + const d = createDeferred(); + d.resolve(42); + expect(await d.promise).toBe(42); + }); + + it("works with object types", async () => { + const d = createDeferred<{ ok: boolean }>(); + d.resolve({ ok: true }); + expect(await d.promise).toEqual({ ok: true }); + }); + + it("can be resolved asynchronously", async () => { + const d = createDeferred(); + setTimeout(() => d.resolve("async"), 5); + expect(await d.promise).toBe("async"); + }); +}); + +// --------------------------------------------------------------------------- +// closeServer +// --------------------------------------------------------------------------- + +describe("closeServer", () => { + let server: http.Server | undefined; + + afterEach(() => { + // Safety net in case a test fails before closing. + if (server?.listening) { + server.close(); + } + server = undefined; + }); + + it("closes a listening HTTP server", async () => { + server = http.createServer(); + await new Promise((resolve) => server!.listen(0, "127.0.0.1", resolve)); + expect(server.listening).toBe(true); + + await closeServer(server); + expect(server.listening).toBe(false); + }); +}); + +// --------------------------------------------------------------------------- +// escapeHtml +// --------------------------------------------------------------------------- + +describe("escapeHtml", () => { + it("escapes ampersand", () => { + expect(escapeHtml("a&b")).toBe("a&b"); + }); + + it("escapes less-than", () => { + expect(escapeHtml("a { + expect(escapeHtml("a>b")).toBe("a>b"); + }); + + it("escapes double quote", () => { + expect(escapeHtml('a"b')).toBe("a"b"); + }); + + it("escapes single quote", () => { + expect(escapeHtml("a'b")).toBe("a'b"); + }); + + it("escapes all special chars together", () => { + expect(escapeHtml(`&<>"'`)).toBe("&<>"'"); + }); + + it("returns plain strings unchanged", () => { + expect(escapeHtml("hello world")).toBe("hello world"); + }); +}); + +// --------------------------------------------------------------------------- +// renderOAuthCallbackHtml +// --------------------------------------------------------------------------- + +describe("renderOAuthCallbackHtml", () => { + it("renders a valid HTML page with the given title and message", () => { + const html = renderOAuthCallbackHtml({ + title: "Test Title", + message: "Test message", + success: true, + }); + expect(html).toContain(""); + expect(html).toContain("Test Title"); + expect(html).toContain("

Test Title

"); + expect(html).toContain("

Test message

"); + }); + + it("includes an auto-close script when success is true", () => { + const html = renderOAuthCallbackHtml({ + title: "Done", + message: "All good", + success: true, + }); + expect(html).toContain("window.close()"); + expect(html).toContain("const ok = true;"); + }); + + it("does NOT auto-close when success is false", () => { + const html = renderOAuthCallbackHtml({ + title: "Failed", + message: "Something went wrong", + success: false, + }); + expect(html).toContain("const ok = false;"); + // The script tag is present but guarded by `if (!ok) return;` + expect(html).toContain("if (!ok) return;"); + }); + + it("escapes title to prevent XSS", () => { + const html = renderOAuthCallbackHtml({ + title: '', + message: "ok", + success: true, + }); + expect(html).not.toContain("