diff --git a/src/core/auth/chatgpt_oauth.zig b/src/core/auth/chatgpt_oauth.zig index cf072ea9b..df5199d63 100644 --- a/src/core/auth/chatgpt_oauth.zig +++ b/src/core/auth/chatgpt_oauth.zig @@ -701,7 +701,20 @@ fn sessionFromToken(alloc: Allocator, token: *TokenSet, now_ms: i64) !chatgpt_se return session; } -pub fn extractAccountId(alloc: Allocator, token: []const u8) ![]u8 { +const max_codex_residency_bytes: usize = 64; + +pub const RequestIdentity = struct { + account_id: []u8, + residency: ?[]u8, + + pub fn deinit(self: *RequestIdentity, alloc: Allocator) void { + alloc.free(self.account_id); + if (self.residency) |residency| alloc.free(residency); + self.* = undefined; + } +}; + +pub fn extractRequestIdentity(alloc: Allocator, token: []const u8) !RequestIdentity { var parts = std.mem.splitScalar(u8, token, '.'); _ = parts.next() orelse return error.InvalidChatGptAccessToken; const payload = parts.next() orelse return error.InvalidChatGptAccessToken; @@ -721,8 +734,36 @@ pub fn extractAccountId(alloc: Allocator, token: []const u8) ![]u8 { if (parsed.value != .object) return error.InvalidChatGptAccessToken; const claim = parsed.value.object.get(jwt_auth_claim) orelse return error.InvalidChatGptAccessToken; if (claim != .object) return error.InvalidChatGptAccessToken; - return dupeRequiredString(alloc, claim.object, "chatgpt_account_id") catch + + const account_id = dupeRequiredString(alloc, claim.object, "chatgpt_account_id") catch return error.InvalidChatGptAccessToken; + errdefer alloc.free(account_id); + return .{ + .account_id = account_id, + .residency = try dupeResidencyClaim(alloc, claim.object), + }; +} + +pub fn extractAccountId(alloc: Allocator, token: []const u8) ![]u8 { + const identity = try extractRequestIdentity(alloc, token); + if (identity.residency) |residency| alloc.free(residency); + return identity.account_id; +} + +fn dupeResidencyClaim(alloc: Allocator, claim: std.json.ObjectMap) !?[]u8 { + const value = claim.get("chatgpt_compute_residency") orelse return null; + if (value != .string) return null; + const residency = std.mem.trim(u8, value.string, " \t\r\n"); + if (std.mem.eql(u8, residency, "no_constraint") or !validResidencyClaim(residency)) return null; + return try alloc.dupe(u8, residency); +} + +fn validResidencyClaim(residency: []const u8) bool { + if (residency.len == 0 or residency.len > max_codex_residency_bytes) return false; + for (residency) |byte| { + if (byte < 0x21 or byte > 0x7e) return false; + } + return true; } fn dupeRequiredString(alloc: Allocator, object: std.json.ObjectMap, key: []const u8) ![]u8 { @@ -864,6 +905,20 @@ fn writeStdout(text: []const u8) !void { try std.Io.File.stdout().writeStreamingAll(io_mod.getIo(), text); } +fn testAccessToken(alloc: Allocator, auth_claim_json: []const u8) ![]u8 { + const payload = try std.fmt.allocPrint( + alloc, + "{{\"https://api.openai.com/auth\":{s}}}", + .{auth_claim_json}, + ); + defer alloc.free(payload); + const encoded_len = std.base64.url_safe_no_pad.Encoder.calcSize(payload.len); + const encoded = try alloc.alloc(u8, encoded_len); + defer alloc.free(encoded); + _ = std.base64.url_safe_no_pad.Encoder.encode(encoded, payload); + return std.fmt.allocPrint(alloc, "header.{s}.signature", .{encoded}); +} + test "ChatGPT E2E OAuth endpoint overrides accept only loopback HTTP" { try std.testing.expect(isLoopbackHttpUrl("http://127.0.0.1:1234/token")); try std.testing.expect(isLoopbackHttpUrl("http://localhost:1234/token")); @@ -873,14 +928,7 @@ test "ChatGPT E2E OAuth endpoint overrides accept only loopback HTTP" { test "ChatGPT account id is extracted from the namespaced JWT claim" { const alloc = std.testing.allocator; - const payload = - \\{"https://api.openai.com/auth":{"chatgpt_account_id":"acct_test"}} - ; - const encoded_len = std.base64.url_safe_no_pad.Encoder.calcSize(payload.len); - const encoded = try alloc.alloc(u8, encoded_len); - defer alloc.free(encoded); - _ = std.base64.url_safe_no_pad.Encoder.encode(encoded, payload); - const token = try std.fmt.allocPrint(alloc, "header.{s}.signature", .{encoded}); + const token = try testAccessToken(alloc, "{\"chatgpt_account_id\":\"acct_test\"}"); defer alloc.free(token); const account_id = try extractAccountId(alloc, token); @@ -888,6 +936,75 @@ test "ChatGPT account id is extracted from the namespaced JWT claim" { try std.testing.expectEqualStrings("acct_test", account_id); } +test "ChatGPT request identity extracts bounded compute residency" { + const alloc = std.testing.allocator; + const max_residency = "r" ** max_codex_residency_bytes; + const max_claim = try std.fmt.allocPrint( + alloc, + "{{\"chatgpt_account_id\":\"acct_test\",\"chatgpt_compute_residency\":\"{s}\"}}", + .{max_residency}, + ); + defer alloc.free(max_claim); + const cases = [_]struct { + claim: []const u8, + expected_residency: ?[]const u8, + }{ + .{ + .claim = "{\"chatgpt_account_id\":\"acct_test\",\"chatgpt_data_residency\":\"eu\",\"chatgpt_compute_residency\":\"us\"}", + .expected_residency = "us", + }, + .{ + .claim = "{\"chatgpt_account_id\":\"acct_test\",\"chatgpt_compute_residency\":\" eu \"}", + .expected_residency = "eu", + }, + .{ .claim = max_claim, .expected_residency = max_residency }, + .{ + .claim = "{\"chatgpt_account_id\":\"acct_test\",\"chatgpt_data_residency\":\"eu\"}", + .expected_residency = null, + }, + .{ .claim = "{\"chatgpt_account_id\":\"acct_test\"}", .expected_residency = null }, + }; + + for (cases) |case| { + const token = try testAccessToken(alloc, case.claim); + defer alloc.free(token); + var identity = try extractRequestIdentity(alloc, token); + defer identity.deinit(alloc); + try std.testing.expectEqualStrings("acct_test", identity.account_id); + if (case.expected_residency) |expected| { + try std.testing.expectEqualStrings(expected, identity.residency orelse ""); + } else { + try std.testing.expect(identity.residency == null); + } + } +} + +test "ChatGPT request identity omits unconstrained and unsafe compute residency" { + const alloc = std.testing.allocator; + const overlong_residency = "r" ** (max_codex_residency_bytes + 1); + const overlong_claim = try std.fmt.allocPrint( + alloc, + "{{\"chatgpt_account_id\":\"acct_test\",\"chatgpt_compute_residency\":\"{s}\"}}", + .{overlong_residency}, + ); + defer alloc.free(overlong_claim); + const claims = [_][]const u8{ + "{\"chatgpt_account_id\":\"acct_test\",\"chatgpt_compute_residency\":\"no_constraint\"}", + "{\"chatgpt_account_id\":\"acct_test\",\"chatgpt_compute_residency\":\"\"}", + "{\"chatgpt_account_id\":\"acct_test\",\"chatgpt_compute_residency\":\"us\\r\\nx-injected: yes\"}", + "{\"chatgpt_account_id\":\"acct_test\",\"chatgpt_compute_residency\":42}", + overlong_claim, + }; + + for (claims) |claim| { + const token = try testAccessToken(alloc, claim); + defer alloc.free(token); + var identity = try extractRequestIdentity(alloc, token); + defer identity.deinit(alloc); + try std.testing.expect(identity.residency == null); + } +} + test "Codex refresh uses JSON and accepts omitted token rotation and lifetime" { const State = struct { method: ?oauth_transport.Method = null, diff --git a/src/gateway/openai_codex.zig b/src/gateway/openai_codex.zig index ce0f52016..7c0ee4bd9 100644 --- a/src/gateway/openai_codex.zig +++ b/src/gateway/openai_codex.zig @@ -326,8 +326,8 @@ fn streamCompletionCore(alloc: Allocator, request: stream_provider.Request) !str return error.CodexSubscriptionCredentialRequired; } try validateModel(request.model); - const account_id = try chatgpt_oauth.extractAccountId(alloc, request.api_key); - defer alloc.free(account_id); + var request_identity = try chatgpt_oauth.extractRequestIdentity(alloc, request.api_key); + defer request_identity.deinit(alloc); const auth_header = try std.fmt.allocPrint(alloc, "Bearer {s}", .{request.api_key}); defer secret.zeroAndFree(alloc, auth_header); const request_endpoint = if (io_mod.getenv(e2e_endpoint_env)) |override| endpoint: { @@ -338,8 +338,12 @@ fn streamCompletionCore(alloc: Allocator, request: stream_provider.Request) !str var extra_headers_buf: [7]std.http.Header = undefined; var extra_count: usize = 0; - extra_headers_buf[extra_count] = .{ .name = "chatgpt-account-id", .value = account_id }; + extra_headers_buf[extra_count] = .{ .name = "chatgpt-account-id", .value = request_identity.account_id }; extra_count += 1; + if (request_identity.residency) |residency| { + extra_headers_buf[extra_count] = .{ .name = "x-openai-internal-codex-residency", .value = residency }; + extra_count += 1; + } extra_headers_buf[extra_count] = .{ .name = "originator", .value = "fx" }; extra_count += 1; extra_headers_buf[extra_count] = .{ .name = "OpenAI-Beta", .value = "responses=experimental" }; diff --git a/src/gateway/openai_codex_models.zig b/src/gateway/openai_codex_models.zig index b6992e190..38ec29800 100644 --- a/src/gateway/openai_codex_models.zig +++ b/src/gateway/openai_codex_models.zig @@ -63,11 +63,11 @@ fn fetchCatalogForProvider( } const credential = input.access.authorizationCredential() orelse return .{ .failure = .{ .category = .authentication, .http_status = .unauthorized } }; - const account_id = chatgpt_oauth.extractAccountId(alloc, credential) catch |err| { + var request_identity = chatgpt_oauth.extractRequestIdentity(alloc, credential) catch |err| { if (err == error.OutOfMemory) return error.OutOfMemory; return .{ .failure = .{ .category = .authentication, .http_status = .unauthorized } }; }; - defer alloc.free(account_id); + defer request_identity.deinit(alloc); const request_url = modelsUrl(alloc) catch |err| { if (err == error.OutOfMemory) return error.OutOfMemory; return .{ .failure = .{ .category = .runtime } }; @@ -80,7 +80,8 @@ fn fetchCatalogForProvider( .alloc = alloc, .url = request_url, .credential = credential, - .account_id = account_id, + .account_id = request_identity.account_id, + .residency = request_identity.residency, }; var response = gateway_client.runBoundedHttpOperation( FetchResponse, @@ -135,6 +136,7 @@ const FetchOperation = struct { url: []const u8, credential: []const u8, account_id: []const u8, + residency: ?[]const u8, pub fn run(self: *@This()) !FetchResponse { var client: std.http.Client = .{ .allocator = self.alloc, .io = io_mod.getIo() }; @@ -144,6 +146,18 @@ const FetchOperation = struct { const body_buffer = try self.alloc.alloc(u8, max_catalog_bytes + 1); defer secret.zeroAndFree(self.alloc, body_buffer); var response_writer = std.Io.Writer.fixed(body_buffer); + var extra_headers_buf: [4]std.http.Header = undefined; + var extra_count: usize = 0; + extra_headers_buf[extra_count] = .{ .name = "chatgpt-account-id", .value = self.account_id }; + extra_count += 1; + if (self.residency) |residency| { + extra_headers_buf[extra_count] = .{ .name = "x-openai-internal-codex-residency", .value = residency }; + extra_count += 1; + } + extra_headers_buf[extra_count] = .{ .name = "originator", .value = "fx" }; + extra_count += 1; + extra_headers_buf[extra_count] = .{ .name = "accept", .value = "application/json" }; + extra_count += 1; const result = client.fetch(.{ .location = .{ .url = self.url }, .method = .GET, @@ -152,11 +166,7 @@ const FetchOperation = struct { .user_agent = .{ .override = gateway_client.user_agent }, .accept_encoding = .omit, }, - .extra_headers = &.{ - .{ .name = "chatgpt-account-id", .value = self.account_id }, - .{ .name = "originator", .value = "fx" }, - .{ .name = "accept", .value = "application/json" }, - }, + .extra_headers = extra_headers_buf[0..extra_count], .response_writer = &response_writer, .redirect_behavior = .unhandled, }) catch |err| switch (err) { diff --git a/tests/e2e/acp.test.ts b/tests/e2e/acp.test.ts index f157e82ec..8e3b510f8 100644 --- a/tests/e2e/acp.test.ts +++ b/tests/e2e/acp.test.ts @@ -522,9 +522,13 @@ function writeSeededFxAuth(home: string, teamId?: string): void { function acpChatGptAccessToken( accountId = "acct_acp_e2e", signature = "signature", + authClaims: Record = {}, ): string { const payload = Buffer.from(JSON.stringify({ - "https://api.openai.com/auth": { chatgpt_account_id: accountId }, + "https://api.openai.com/auth": { + chatgpt_account_id: accountId, + ...authClaims, + }, })).toString("base64url"); return `header.${payload}.${signature}`; } @@ -577,19 +581,24 @@ function codexLatestToolResult(body: string): { callId: string; output: string } function startAcpFakeCodex(options: { unauthorizedResponses?: number; route?: (body: string) => string | Promise; + authClaims?: Record; } = {}) { - const accessToken = acpChatGptAccessToken("acct_acp_e2e", "stale"); - const refreshedAccessToken = acpChatGptAccessToken("acct_acp_e2e", "fresh"); - const requests: Array<{ path: string; authorization: string | null; body: string }> = []; - const modelRequests: Array<{ path: string; authorization: string | null }> = []; - const tokenRequests: Array<{ path: string; authorization: string | null }> = []; + const accessToken = acpChatGptAccessToken("acct_acp_e2e", "stale", options.authClaims); + const refreshedAccessToken = acpChatGptAccessToken("acct_acp_e2e", "fresh", options.authClaims); + const requests: Array<{ path: string; authorization: string | null; residency: string | null; body: string }> = []; + const modelRequests: Array<{ path: string; authorization: string | null; residency: string | null }> = []; + const tokenRequests: Array<{ path: string; authorization: string | null; residency: string | null }> = []; let unauthorizedResponses = options.unauthorizedResponses ?? 0; const server = Bun.serve({ hostname: "127.0.0.1", port: 0, async fetch(request) { const path = new URL(request.url).pathname; - const recorded = { path, authorization: request.headers.get("authorization") }; + const recorded = { + path, + authorization: request.headers.get("authorization"), + residency: request.headers.get("x-openai-internal-codex-residency"), + }; if (path === "/models") { modelRequests.push(recorded); return Response.json({ models: [ @@ -7462,6 +7471,53 @@ describe("acp: model catalog authentication", () => { }, TIMEOUT, ); + + test( + "session provider changes propagate Codex residency to model discovery", + async () => { + const root = createIsolatedRoot("fx-acp-codex-residency-"); + const gateway = startFakeGateway([]); + const codex = startAcpFakeCodex({ + authClaims: { chatgpt_compute_residency: "us" }, + }); + writeSeededAcpChatGptLogin(root.home, codex.accessToken); + try { + client = await AcpClient.create({ + cwd: root.workspace, + env: { + ...fakeGatewayEnv(root, gateway), + FX_E2E_OPENAI_CODEX_RESPONSES_URL: codex.responsesUrl, + FX_E2E_OPENAI_CODEX_MODELS_URL: codex.modelsUrl, + FX_E2E_CHATGPT_TOKEN_URL: codex.tokenUrl, + }, + }); + await client.request("initialize", { protocolVersion: 1 }, 1); + await client.request("session/new", { mcpServers: [] }, 2); + await client.readLine(); + + const changed = await client.request("session/set_config_option", { + configId: "provider", + value: "codex", + }, 3) as any; + expect(changed.result.configOptions.find((option: any) => option.id === "provider").currentValue) + .toBe("codex"); + expect(codex.modelRequests).toHaveLength(1); + expect(codex.modelRequests[0]!.residency).toBe("us"); + expect(gateway.requests).toHaveLength(0); + expect(gateway.modelRequests).toHaveLength(1); + for (const request of [...gateway.requests, ...gateway.modelRequests]) { + expect(request.headers.has("x-openai-internal-codex-residency")).toBe(false); + } + expect(client.stderr).toBe(""); + } finally { + await client?.close(); + codex.stop(); + gateway.stop(); + rmSync(root.root, { recursive: true, force: true }); + } + }, + TIMEOUT, + ); }); describe.skipIf(!HAS_API_KEY)("acp: model-backed protocol", () => { diff --git a/tests/e2e/tui-auth-source-selection.test.ts b/tests/e2e/tui-auth-source-selection.test.ts index 63b4bcedf..6c8373067 100644 --- a/tests/e2e/tui-auth-source-selection.test.ts +++ b/tests/e2e/tui-auth-source-selection.test.ts @@ -354,9 +354,15 @@ function startFakeOAuth( }; } -function chatgptAccessToken(accountId = "acct_e2e"): string { +function chatgptAccessToken( + accountId = "acct_e2e", + authClaims: Record = {}, +): string { const payload = Buffer.from(JSON.stringify({ - "https://api.openai.com/auth": { chatgpt_account_id: accountId }, + "https://api.openai.com/auth": { + chatgpt_account_id: accountId, + ...authClaims, + }, })).toString("base64url"); return `header.${payload}.signature`; } @@ -704,9 +710,11 @@ function startFakeCodexToolLoop(options: { toolName?: string; toolArguments?: object; finalText?: string; + authClaims?: Record; } = {}) { const bodies: string[] = []; - const accessToken = chatgptAccessToken("acct_tool_loop"); + const requestHeaders: Headers[] = []; + const accessToken = chatgptAccessToken("acct_tool_loop", options.authClaims); const toolName = options.toolName ?? "read_file"; const toolArguments = options.toolArguments ?? { path: "README.md" }; const finalText = options.finalText ?? "CODEX_TOOL_LOOP_OK"; @@ -720,6 +728,7 @@ function startFakeCodexToolLoop(options: { { slug: "gpt-5.4-mini", visibility: "list", supported_in_api: true, supported_reasoning_levels: [{ effort: "low" }], additional_speed_tiers: [], input_modalities: ["text"], context_window: 128000 }, ] }); } + requestHeaders.push(new Headers(request.headers)); bodies.push(await request.text()); if (bodies.length === 1) { return new Response( @@ -741,6 +750,7 @@ function startFakeCodexToolLoop(options: { return { accessToken, bodies, + requestHeaders, responsesUrl: `http://127.0.0.1:${server.port}/responses`, modelsUrl: `http://127.0.0.1:${server.port}/models`, stop() { server.stop(true); }, @@ -2435,11 +2445,13 @@ tmuxTest( ); test( - "ChatGPT tool loops round-trip encrypted reasoning without Gateway leakage", + "ChatGPT tool loops round-trip encrypted reasoning and residency without Gateway leakage", async () => { home = mkdtempSync(join(tmpdir(), "fx-chatgpt-tool-loop-")); gateway = startFakeGateway([]); - const codex = startFakeCodexToolLoop(); + const codex = startFakeCodexToolLoop({ + authClaims: { chatgpt_compute_residency: "us" }, + }); try { writeSeededChatGptLogin(home, codex.accessToken); writeFileSync( @@ -2467,11 +2479,14 @@ test( expect(result.code, `stdout: ${result.stdout}\nstderr: ${result.stderr}`).toBe(0); expect(result.stdout).toContain("CODEX_TOOL_LOOP_OK"); expect(codex.bodies).toHaveLength(2); + expect(codex.requestHeaders).toHaveLength(2); + expect(codex.requestHeaders.every((headers) => + headers.get("x-openai-internal-codex-residency") === "us" + )).toBe(true); expect(codex.bodies[1]).toContain('"encrypted_content":"opaque-tool-loop"'); expect(codex.bodies[1]).toContain('"type":"function_call_output"'); - for (const request of [...gateway.requests, ...gateway.modelRequests]) { - expect(request.headers.get("authorization")).not.toBe(`Bearer ${codex.accessToken}`); - } + expect(gateway.requests).toHaveLength(0); + expect(gateway.modelRequests).toHaveLength(0); } finally { codex.stop(); }