From 394ae237207e83709ed5594c7f178da83ae6c385 Mon Sep 17 00:00:00 2001 From: MumuTW <42820974+MumuTW@users.noreply.github.com> Date: Fri, 28 Aug 2026 20:31:08 +0800 Subject: [PATCH] fix(sse): harden session affinity key extraction (#11746) Hardens session-affinity key extraction: no more JSON.stringify on arbitrary request objects, recognizes bounded text from Responses/chat/Anthropic/Gemini/common string-root shapes, enforces a shared 4096-char processing budget, and rejects oversized explicit session IDs before trim/regex/hash work. 128/128 focused affinity/failover tests passing. Closes #11744. Thanks! --- .../11746-session-affinity-key-hardening.md | 3 + src/sse/services/auth.ts | 78 +--- src/sse/services/sessionAffinityPin.ts | 178 ++++++++ .../unit/session-affinity-edge-cases.test.ts | 408 ++++++++++++++++++ 4 files changed, 591 insertions(+), 76 deletions(-) create mode 100644 changelog.d/fixes/11746-session-affinity-key-hardening.md create mode 100644 tests/unit/session-affinity-edge-cases.test.ts diff --git a/changelog.d/fixes/11746-session-affinity-key-hardening.md b/changelog.d/fixes/11746-session-affinity-key-hardening.md new file mode 100644 index 0000000000..7165f81245 --- /dev/null +++ b/changelog.d/fixes/11746-session-affinity-key-hardening.md @@ -0,0 +1,3 @@ +- **fix(routing):** harden session-affinity key derivation against oversized explicit IDs and + empty structural-payload collisions, while recognizing Gemini and common text request shapes + ([#11746](https://github.com/diegosouzapw/OmniRoute/pull/11746)) — thanks @MumuTW diff --git a/src/sse/services/auth.ts b/src/sse/services/auth.ts index 1f754c21e9..a5d93bd39b 100644 --- a/src/sse/services/auth.ts +++ b/src/sse/services/auth.ts @@ -1,4 +1,4 @@ -import { randomUUID, createHash } from "crypto"; +import { randomUUID } from "crypto"; import { nodeTypeFromId } from "@/lib/db/providerNodeSelect"; import { extractGoogApiKeyHeader } from "./googApiKeyAuth.ts"; import { describeUpstreamFailure } from "@/shared/utils/upstreamError"; @@ -218,81 +218,6 @@ function toStringOrNull(value: unknown): string | null { function toBooleanOrDefault(value: unknown, fallback: boolean): boolean { return typeof value === "boolean" ? value : fallback; } -function normalizeSessionKey(value: unknown, prefix: string): string | null { - if (typeof value !== "string" || value.trim().length === 0) return null; - const trimmed = value.trim(); - if (trimmed.length <= 180 && /^[A-Za-z0-9._:-]+$/.test(trimmed)) { - return `${prefix}:${trimmed}`; - } - return `${prefix}:sha256:${createHash("sha256").update(trimmed).digest("hex")}`; -} -function extractTextForSessionHash(value: unknown): string | null { - if (typeof value === "string") return value; - if (Array.isArray(value)) { - const parts = value - .map((item) => { - if (typeof item === "string") return item; - const record = asRecord(item); - if (typeof record.text === "string") return record.text; - if (typeof record.content === "string") return record.content; - return null; - }) - .filter(Boolean) as string[]; - return parts.length > 0 ? parts.join("\n") : JSON.stringify(value); - } - if (value && typeof value === "object") return JSON.stringify(value); - return null; -} -function getFirstInputText(body: unknown): string | null { - const record = asRecord(body); - if (record.input !== undefined) { - if (typeof record.input === "string") return record.input; - if (Array.isArray(record.input)) { - for (const item of record.input) { - const itemRecord = asRecord(item); - const text = extractTextForSessionHash(itemRecord.content ?? item); - if (text && text.trim().length > 0) return text; - } - } - const text = extractTextForSessionHash(record.input); - if (text && text.trim().length > 0) return text; - } - - if (Array.isArray(record.messages)) { - const userMessage = record.messages.find((message) => asRecord(message).role === "user"); - const firstMessage = userMessage ?? record.messages[0]; - const text = extractTextForSessionHash(asRecord(firstMessage).content ?? firstMessage); - if (text && text.trim().length > 0) return text; - } - - return null; -} -export function extractSessionAffinityKey( - body: unknown, - headers?: Headers | { get?: (name: string) => string | null } | null -): string | null { - const headerKey = normalizeSessionKey( - readHeaderValue(headers, "x-codex-session-id") ?? - readHeaderValue(headers, "x-session-id") ?? - readHeaderValue(headers, "x-omniroute-session"), - "header" - ); - if (headerKey) return headerKey; - - const record = asRecord(body); - const metadata = asRecord(record.metadata); - const explicitKey = - normalizeSessionKey(metadata.session_id, "metadata") ?? - normalizeSessionKey(metadata.sessionId, "metadata") ?? - normalizeSessionKey(record.conversation_id, "conversation") ?? - normalizeSessionKey(record.session_id, "session") ?? - normalizeSessionKey(record.prompt_cache_key, "prompt-cache"); - if (explicitKey) return explicitKey; - - const inputText = getFirstInputText(body); - if (!inputText || inputText.trim().length === 0) return null; - return `input:sha256:${createHash("sha256").update(inputText.slice(0, 4096)).digest("hex")}`; -} function getCodexLimitPolicy(providerSpecificData: JsonRecord): { use5h: boolean; useWeekly: boolean; @@ -1023,6 +948,7 @@ export { fisherYatesShuffle, getNextFromDeckSync as getNextFromDeck }; // Re-export readHeaderValue and AuthRequestHeaders from headerReader.ts for // backwards compat with existing imports (e.g. googApiKeyAuth.ts). export { readHeaderValue, type AuthRequestHeaders } from "./headerReader.ts"; +export { extractSessionAffinityKey } from "./sessionAffinityPin"; const PROVIDER_SEARCH_PAIRS: string[][] = [ ["nvidia", "nvidia_nim"], ["kimi-coding", "kimi-coding-apikey"], diff --git a/src/sse/services/sessionAffinityPin.ts b/src/sse/services/sessionAffinityPin.ts index f7c5009df6..fad970f134 100644 --- a/src/sse/services/sessionAffinityPin.ts +++ b/src/sse/services/sessionAffinityPin.ts @@ -26,6 +26,7 @@ * check wrapping `evaluateQuotaLimitPolicy` — are injected as callbacks. */ +import { createHash } from "crypto"; import { getSessionAccountAffinity, upsertSessionAccountAffinity, @@ -42,6 +43,183 @@ import { } from "@omniroute/open-sse/services/accountFallback.ts"; import { isComboPerModelTimeoutAbort } from "@omniroute/open-sse/services/combo/comboAbortReasons.ts"; import * as log from "../utils/logger"; +import { readHeaderValue } from "./headerReader.ts"; + +export { readHeaderValue } from "./headerReader.ts"; + +function asRecord(value: unknown): Record { + return value && typeof value === "object" && !Array.isArray(value) + ? (value as Record) + : {}; +} + +/** Maximum accepted session-key input length. Longer identifiers are rejected before processing. */ +const SESSION_KEY_MAX_INPUT_LEN = 4096; + +function normalizeSessionKey(value: unknown, prefix: string): string | null { + if (typeof value !== "string" || value.length > SESSION_KEY_MAX_INPUT_LEN) return null; + const trimmed = value.trim(); + if (trimmed.length === 0) return null; + if (trimmed.length <= 180 && /^[A-Za-z0-9._:-]+$/.test(trimmed)) { + return `${prefix}:${trimmed}`; + } + return `${prefix}:sha256:${createHash("sha256").update(trimmed).digest("hex")}`; +} + +/** Upper bound on text extracted for session hashing (matches the slice in extractSessionAffinityKey). */ +const SESSION_HASH_TEXT_LIMIT = 4096; + +function extractBoundedNonEmptyText( + value: unknown, + limit = SESSION_HASH_TEXT_LIMIT +): string | null { + if (typeof value !== "string" || limit <= 0) return null; + const bounded = value.slice(0, limit); + return bounded.trim().length > 0 ? bounded : null; +} + +/** + * Extracts human-readable text from a value for session-affinity hashing. + * + * Design constraints (post-adversarial-review): + * - NEVER calls `JSON.stringify` on arbitrary objects — avoids synchronous + * Event Loop blocking on huge payloads (e.g. 50MB multimodal base64). + * - NEVER returns structural tokens like `"[]"` or `"{}"` — avoids global + * session-key collisions across unrelated users with empty payloads. + * - Only extracts values from known text fields (`.text`, `.content`) or + * raw strings. Payloads without recognisable text content get `null`, + * which correctly signals "no input-based session affinity" — callers + * should use explicit session IDs instead. + */ +function extractTextForSessionHash(value: unknown): string | null { + if (typeof value === "string") return extractBoundedNonEmptyText(value); + + if (Array.isArray(value)) { + const parts: string[] = []; + let totalLen = 0; + for (const item of value) { + if (totalLen >= SESSION_HASH_TEXT_LIMIT) break; + let candidate: unknown = null; + if (typeof item === "string") { + candidate = item; + } else { + const record = asRecord(item); + if (typeof record.text === "string") candidate = record.text; + else if (typeof record.content === "string") candidate = record.content; + } + const separatorLength = parts.length > 0 ? 1 : 0; + const text = extractBoundedNonEmptyText( + candidate, + SESSION_HASH_TEXT_LIMIT - totalLen - separatorLength + ); + if (text) { + parts.push(text); + totalLen += separatorLength + text.length; + } + } + return parts.length > 0 ? parts.join("\n") : null; + } + + if (value && typeof value === "object") { + const record = value as Record; + // Known text fields — covers OpenAI, Anthropic, and most providers + const directText = + extractBoundedNonEmptyText(record.text) ?? + extractBoundedNonEmptyText(record.content) ?? + extractBoundedNonEmptyText(record.prompt); + if (directText) return directText; + // Gemini format: { parts: [{ text: "..." }, ...] } + if (Array.isArray(record.parts)) { + const partsText = extractTextForSessionHash(record.parts); + if (partsText) return partsText; + } + // No recognisable text field — return null rather than risking a + // potentially huge JSON.stringify on arbitrary payload shapes. + return null; + } + + return null; +} + +function getFirstInputText(body: unknown): string | null { + const record = asRecord(body); + + // Codex / Responses API: { input: "..." | [...] } + if (record.input !== undefined) { + if (typeof record.input === "string") return extractBoundedNonEmptyText(record.input); + if (Array.isArray(record.input)) { + for (const item of record.input) { + const itemRecord = asRecord(item); + const text = extractTextForSessionHash(itemRecord.content ?? item); + if (text) return text; + } + } + const text = extractTextForSessionHash(record.input); + if (text) return text; + } + + // OpenAI Chat / Anthropic Messages: { messages: [...] } + if (Array.isArray(record.messages)) { + const userMessage = record.messages.find((message) => asRecord(message).role === "user"); + const firstMessage = userMessage ?? record.messages[0]; + const text = extractTextForSessionHash(asRecord(firstMessage).content ?? firstMessage); + if (text) return text; + } + + // Google Gemini: { contents: [{ role: "user", parts: [{ text: "..." }] }] } + if (Array.isArray(record.contents)) { + const userContent = record.contents.find((c) => asRecord(c).role === "user"); + const firstContent = userContent ?? record.contents[0]; + const text = extractTextForSessionHash(asRecord(firstContent).parts ?? firstContent); + if (text) return text; + } + + // OpenAI Legacy Completions / Anthropic /v1/complete / Ollama: { prompt: "..." } + const prompt = extractBoundedNonEmptyText(record.prompt); + if (prompt) return prompt; + + // Other common root-level text fields + const query = extractBoundedNonEmptyText(record.query); + if (query) return query; + const instruction = extractBoundedNonEmptyText(record.instruction); + if (instruction) return instruction; + + return null; +} + +/** + * Derives the stable connection-affinity key for a request. + * + * @param body - Parsed request body containing explicit session identifiers or recognized text. + * @param headers - Optional request headers that may carry an explicit session identifier. + * @returns A namespaced affinity key, or `null` when no safe key can be derived. + */ +export function extractSessionAffinityKey( + body: unknown, + headers?: Headers | { get?: (name: string) => string | null } | null +): string | null { + const headerKey = normalizeSessionKey( + readHeaderValue(headers, "x-codex-session-id") ?? + readHeaderValue(headers, "x-session-id") ?? + readHeaderValue(headers, "x-omniroute-session"), + "header" + ); + if (headerKey) return headerKey; + + const record = asRecord(body); + const metadata = asRecord(record.metadata); + const explicitKey = + normalizeSessionKey(metadata.session_id, "metadata") ?? + normalizeSessionKey(metadata.sessionId, "metadata") ?? + normalizeSessionKey(record.conversation_id, "conversation") ?? + normalizeSessionKey(record.session_id, "session") ?? + normalizeSessionKey(record.prompt_cache_key, "prompt-cache"); + if (explicitKey) return explicitKey; + + const inputText = getFirstInputText(body); + if (!inputText) return null; + return `input:sha256:${createHash("sha256").update(inputText).digest("hex")}`; +} /** Minimal structural view of a provider connection this module reads. */ export interface AffinityPinConnection { diff --git a/tests/unit/session-affinity-edge-cases.test.ts b/tests/unit/session-affinity-edge-cases.test.ts new file mode 100644 index 0000000000..3537f6f9f0 --- /dev/null +++ b/tests/unit/session-affinity-edge-cases.test.ts @@ -0,0 +1,408 @@ +/** + * Edge-case unit tests for `extractSessionAffinityKey`, `getFirstInputText`, + * and `extractTextForSessionHash` — focused on the functions extracted from + * `auth.ts` into `sessionAffinityPin.ts`. + * + * Covers: + * - Array payloads passed as `body` (asRecord Array.isArray guard) + * - Multimodal payloads with base64 images (DoS protection via recognized text fields) + * - Deeply nested message structures + * - Empty / null / undefined edge cases + * - Large payload truncation (SESSION_HASH_TEXT_LIMIT = 4096) + * - Google Gemini contents/parts format + * - OpenAI Legacy / Anthropic / Ollama prompt field + * - normalizeSessionKey DoS protection + */ +import test from "node:test"; +import assert from "node:assert/strict"; +import { createHash } from "node:crypto"; + +const { extractSessionAffinityKey, readHeaderValue } = + await import("../../src/sse/services/sessionAffinityPin.ts"); + +// --------------------------------------------------------------------------- +// Helpers +// --------------------------------------------------------------------------- +function sha256hex(input: string): string { + return createHash("sha256").update(input).digest("hex"); +} + +// --------------------------------------------------------------------------- +// 1. asRecord Array guard — body is an array, not an object +// --------------------------------------------------------------------------- + +test("extractSessionAffinityKey returns null for a bare array body with no extractable text", () => { + // A bare array like `[1, 2, 3]` should be treated as non-record and produce + // null — not crash or misinterpret array indices as keys. + const result = extractSessionAffinityKey([1, 2, 3]); + assert.equal(result, null); +}); + +test("extractSessionAffinityKey returns null for an empty array body", () => { + assert.equal(extractSessionAffinityKey([]), null); +}); + +test("extractSessionAffinityKey returns null for null body", () => { + assert.equal(extractSessionAffinityKey(null), null); +}); + +test("extractSessionAffinityKey returns null for undefined body", () => { + assert.equal(extractSessionAffinityKey(undefined), null); +}); + +// --------------------------------------------------------------------------- +// 2. Standard message-based payloads (OpenAI-style) +// --------------------------------------------------------------------------- + +test("extractSessionAffinityKey hashes the first user message content", () => { + const body = { + messages: [ + { role: "system", content: "You are helpful." }, + { role: "user", content: "Hello world" }, + ], + }; + const result = extractSessionAffinityKey(body); + const expected = `input:sha256:${sha256hex("Hello world")}`; + assert.equal(result, expected); +}); + +test("extractSessionAffinityKey falls back to first message if no user role", () => { + const body = { + messages: [{ role: "assistant", content: "I am here" }], + }; + const result = extractSessionAffinityKey(body); + const expected = `input:sha256:${sha256hex("I am here")}`; + assert.equal(result, expected); +}); + +// --------------------------------------------------------------------------- +// 3. Input-based payloads (Codex-style) +// --------------------------------------------------------------------------- + +test("extractSessionAffinityKey handles string input directly", () => { + const body = { input: "simple text input" }; + const result = extractSessionAffinityKey(body); + const expected = `input:sha256:${sha256hex("simple text input")}`; + assert.equal(result, expected); +}); + +test("extractSessionAffinityKey handles array input with content parts", () => { + const body = { + input: [{ content: "first part" }, { content: "second part" }], + }; + const result = extractSessionAffinityKey(body); + // getFirstInputText iterates input array, extracts content from first item + const expected = `input:sha256:${sha256hex("first part")}`; + assert.equal(result, expected); +}); + +// --------------------------------------------------------------------------- +// 4. Multimodal payloads — base64 images (DoS protection) +// --------------------------------------------------------------------------- + +test("extractSessionAffinityKey ignores large non-text multimodal content", () => { + // Simulate a multimodal message with a 1MB base64 image + const largeBase64 = "A".repeat(1_000_000); + const body = { + messages: [ + { + role: "user", + content: [ + { type: "text", text: "Describe this image" }, + { type: "image_url", image_url: { url: `data:image/png;base64,${largeBase64}` } }, + ], + }, + ], + }; + + const result = extractSessionAffinityKey(body); + + // Should produce a valid hash + assert.ok(result !== null); + assert.ok(result!.startsWith("input:sha256:")); + + // Should extract the text part "Describe this image" from the content array + const expected = `input:sha256:${sha256hex("Describe this image")}`; + assert.equal(result, expected); +}); + +test("extractSessionAffinityKey returns null for multimodal content with ONLY non-text parts", () => { + // When content is an array but has no text/content string fields, + // we return null rather than JSON.stringifying (prevents DoS + hash collisions) + const body = { + messages: [ + { + role: "user", + content: [{ type: "image_url", image_url: { url: "data:image/png;base64,abc123" } }], + }, + ], + }; + const result = extractSessionAffinityKey(body); + // No extractable text → no input-based session affinity + // (users should use explicit session IDs for image-only requests) + assert.equal(result, null); +}); + +// --------------------------------------------------------------------------- +// 5. Large string payloads — truncation at 4096 chars +// --------------------------------------------------------------------------- + +test("extractSessionAffinityKey produces same hash regardless of text length beyond 4096", () => { + const base = "x".repeat(4096); + const body1 = { input: base + "AAAA" }; + const body2 = { input: base + "BBBB" }; + + const result1 = extractSessionAffinityKey(body1); + const result2 = extractSessionAffinityKey(body2); + + // Both should hash only the first 4096 chars (the slice in extractSessionAffinityKey) + assert.equal(result1, result2); + const expected = `input:sha256:${sha256hex(base)}`; + assert.equal(result1, expected); +}); + +// --------------------------------------------------------------------------- +// 6. Explicit session keys take priority over input hashing +// --------------------------------------------------------------------------- + +test("extractSessionAffinityKey prefers header session ID over body content", () => { + const body = { input: "some text" }; + const headers = new Headers({ "x-codex-session-id": "my-session-123" }); + const result = extractSessionAffinityKey(body, headers); + assert.equal(result, "header:my-session-123"); +}); + +test("extractSessionAffinityKey prefers metadata.session_id over input hashing", () => { + const body = { + metadata: { session_id: "meta-sess-42" }, + input: "some text", + }; + const result = extractSessionAffinityKey(body); + assert.equal(result, "metadata:meta-sess-42"); +}); + +// --------------------------------------------------------------------------- +// 7. Empty / whitespace-only content +// --------------------------------------------------------------------------- + +test("extractSessionAffinityKey returns null for whitespace-only input", () => { + assert.equal(extractSessionAffinityKey({ input: " " }), null); +}); + +test("extractSessionAffinityKey bounds direct input before scanning whitespace", () => { + assert.equal(extractSessionAffinityKey({ input: `${" ".repeat(4096)}A` }), null); +}); + +test("extractSessionAffinityKey bounds recognized text fields before scanning whitespace", () => { + const lateText = `${" ".repeat(4096)}A`; + const bodies = [ + { input: { text: lateText } }, + { messages: [{ role: "user", content: [{ type: "text", text: lateText }] }] }, + { prompt: lateText }, + { query: lateText }, + { instruction: lateText }, + ]; + + for (const body of bodies) { + assert.equal(extractSessionAffinityKey(body), null); + } +}); + +test("extractSessionAffinityKey returns null for empty string input", () => { + assert.equal(extractSessionAffinityKey({ input: "" }), null); +}); + +test("extractSessionAffinityKey returns null for messages with empty content", () => { + const body = { messages: [{ role: "user", content: "" }] }; + assert.equal(extractSessionAffinityKey(body), null); +}); + +// --------------------------------------------------------------------------- +// 8. readHeaderValue edge cases +// --------------------------------------------------------------------------- + +test("readHeaderValue handles record-style headers with array values", () => { + const headers = { "x-session-id": ["first-val", "second-val"] }; + assert.equal(readHeaderValue(headers, "x-session-id"), "first-val"); +}); + +test("readHeaderValue returns null for empty string header value", () => { + const headers = new Headers({ "x-session-id": "" }); + assert.equal(readHeaderValue(headers, "x-session-id"), null); +}); + +test("readHeaderValue returns null for whitespace-only header value", () => { + const headers = new Headers({ "x-session-id": " " }); + assert.equal(readHeaderValue(headers, "x-session-id"), null); +}); + +// --------------------------------------------------------------------------- +// 9. Object input — text field extraction (no JSON.stringify) +// --------------------------------------------------------------------------- + +test("extractSessionAffinityKey returns null for object input without text fields", () => { + // Objects without recognisable .text/.content/.prompt fields return null — + // no JSON.stringify, no DoS, no hash collision + const body = { + input: { nested: { deep: "value" } }, + }; + assert.equal(extractSessionAffinityKey(body), null); +}); + +test("extractSessionAffinityKey extracts .text field from object input", () => { + const body = { + input: { text: "meaningful content", metadata: { irrelevant: true } }, + }; + const result = extractSessionAffinityKey(body); + const expected = `input:sha256:${sha256hex("meaningful content")}`; + assert.equal(result, expected); +}); + +test("extractSessionAffinityKey extracts .content field from object input", () => { + const body = { + input: { content: "another text field", type: "document" }, + }; + const result = extractSessionAffinityKey(body); + const expected = `input:sha256:${sha256hex("another text field")}`; + assert.equal(result, expected); +}); + +test("extractSessionAffinityKey does not enumerate arbitrary object input", () => { + const input = new Proxy( + {}, + { + ownKeys() { + throw new Error("arbitrary input must not be serialized or enumerated"); + }, + } + ); + + assert.equal(extractSessionAffinityKey({ input }), null); +}); + +// --------------------------------------------------------------------------- +// 10. Empty payload collision prevention (adversarial review finding) +// --------------------------------------------------------------------------- + +test("extractSessionAffinityKey returns null for { input: [] } — prevents hash collision", () => { + // Empty array input must not produce a shared session key + assert.equal(extractSessionAffinityKey({ input: [] }), null); +}); + +test("extractSessionAffinityKey returns null for { input: {} } — prevents hash collision", () => { + // Empty object input must not produce a shared session key + assert.equal(extractSessionAffinityKey({ input: {} }), null); +}); + +test("extractSessionAffinityKey returns null for message without content field", () => { + // { role: "user" } with no content field must not produce a shared structural hash + const body = { messages: [{ role: "user" }] }; + assert.equal(extractSessionAffinityKey(body), null); +}); + +// --------------------------------------------------------------------------- +// 11. OpenAI Legacy / Anthropic /v1/complete / Ollama — prompt field +// --------------------------------------------------------------------------- + +test("extractSessionAffinityKey extracts root-level prompt field", () => { + const body = { prompt: "Once upon a time" }; + const expected = `input:sha256:${sha256hex("Once upon a time")}`; + assert.equal(extractSessionAffinityKey(body), expected); +}); + +test("extractSessionAffinityKey ignores empty prompt field", () => { + assert.equal(extractSessionAffinityKey({ prompt: "" }), null); + assert.equal(extractSessionAffinityKey({ prompt: " " }), null); +}); + +// --------------------------------------------------------------------------- +// 12. Google Gemini — contents / parts format +// --------------------------------------------------------------------------- + +test("extractSessionAffinityKey extracts text from Gemini contents format", () => { + const body = { + contents: [ + { + role: "user", + parts: [{ text: "Explain quantum computing" }], + }, + ], + }; + const expected = `input:sha256:${sha256hex("Explain quantum computing")}`; + assert.equal(extractSessionAffinityKey(body), expected); +}); + +test("extractSessionAffinityKey extracts text from Gemini multi-part content", () => { + const body = { + contents: [ + { + role: "user", + parts: [ + { text: "First part" }, + { inlineData: { mimeType: "image/png", data: "abc123" } }, + { text: "Second part" }, + ], + }, + ], + }; + const expected = `input:sha256:${sha256hex("First part\nSecond part")}`; + assert.equal(extractSessionAffinityKey(body), expected); +}); + +test("extractSessionAffinityKey returns null for Gemini image-only parts", () => { + const body = { + contents: [ + { + role: "user", + parts: [{ inlineData: { mimeType: "image/png", data: "abc123" } }], + }, + ], + }; + assert.equal(extractSessionAffinityKey(body), null); +}); + +test("extractSessionAffinityKey prefers messages over contents", () => { + // If both messages and contents exist, messages takes priority + const body = { + messages: [{ role: "user", content: "From messages" }], + contents: [{ role: "user", parts: [{ text: "From contents" }] }], + }; + const expected = `input:sha256:${sha256hex("From messages")}`; + assert.equal(extractSessionAffinityKey(body), expected); +}); + +// --------------------------------------------------------------------------- +// 13. Other root-level text fields — query, instruction +// --------------------------------------------------------------------------- + +test("extractSessionAffinityKey extracts root-level query field", () => { + const body = { query: "search for something" }; + const expected = `input:sha256:${sha256hex("search for something")}`; + assert.equal(extractSessionAffinityKey(body), expected); +}); + +test("extractSessionAffinityKey extracts root-level instruction field", () => { + const body = { instruction: "translate to French" }; + const expected = `input:sha256:${sha256hex("translate to French")}`; + assert.equal(extractSessionAffinityKey(body), expected); +}); + +// --------------------------------------------------------------------------- +// 14. normalizeSessionKey DoS protection (adversarial review R3) +// --------------------------------------------------------------------------- + +test("extractSessionAffinityKey rejects overlong session IDs instead of truncating them", () => { + const sharedPrefix = "A".repeat(4096); + + assert.equal(extractSessionAffinityKey({ session_id: `${sharedPrefix}X` }), null); + assert.equal(extractSessionAffinityKey({ session_id: `${sharedPrefix}Y` }), null); +}); + +test("extractSessionAffinityKey rejects an overlong whitespace-padded session ID", () => { + const whitespacePaddedSessionId = `${" ".repeat(4096)}A`; + + assert.equal( + extractSessionAffinityKey({ metadata: { session_id: whitespacePaddedSessionId } }), + null + ); +});