From c5f15f994d638b0fdd1ea8ed02632be1f22c4de5 Mon Sep 17 00:00:00 2001 From: Dmitry Kuznetsov Date: Wed, 4 Nov 2026 19:38:59 +0300 Subject: [PATCH] feat(pii): reversible token vault for streaming and non-streaming responses Rebased onto release/v3.8.53 tip with complete SpDD rework: - Core: LRU touch order, unified unref sweep interval, and resource dispose() - Interceptor: AsyncLocalStorage isolation with fail-closed fallback - Streaming: Speculative split-token buffer for SSE chunks with zero-TTFT fast path - Flags & Docs: Synchronized 117 feature flags across registry, test suites, and docs - Quality: Consolidated single-source tokenVault with full unit test coverage (26 tests passing) --- .env.example | 7 + .../18017-reversible-pii-token-vault.md | 1 + docs/reference/ENVIRONMENT.md | 3 + docs/reference/FEATURE_FLAGS.md | 9 +- .../chatCore/nonStreamingJsonResponse.ts | 11 +- .../handlers/chatCore/streamingPipeline.ts | 16 + open-sse/utils/tokenVault.ts | 400 ++++++++++++++++++ src/lib/pii/tokenVault.ts | 1 + .../constants/featureFlagDefinitions.ts | 12 + src/shared/utils/inputSanitizer.ts | 17 +- tests/unit/feature-flags-settings.test.ts | 3 +- tests/unit/guardrails/tokenVault.test.ts | 246 +++++++++++ tests/unit/pii/featureFlag.test.ts | 101 +++++ tests/unit/pii/tokenVault.test.ts | 260 ++++++++++++ .../unit/server-owned-tool-loop-flag.test.ts | 3 +- 15 files changed, 1081 insertions(+), 9 deletions(-) create mode 100644 changelog.d/features/18017-reversible-pii-token-vault.md create mode 100644 open-sse/utils/tokenVault.ts create mode 100644 src/lib/pii/tokenVault.ts create mode 100644 tests/unit/guardrails/tokenVault.test.ts create mode 100644 tests/unit/pii/featureFlag.test.ts create mode 100644 tests/unit/pii/tokenVault.test.ts diff --git a/.env.example b/.env.example index c5c4c62e237b..c8a1d3ac3f9d 100644 --- a/.env.example +++ b/.env.example @@ -738,6 +738,13 @@ ALLOW_API_KEY_REVEAL=false # PII_RESPONSE_SANITIZATION=false # PII_RESPONSE_SANITIZATION_MODE=redact # redact = mask PII | warn = log only | block = drop response +# ── Response-Side: Reversible PII Token Vault ── +# When enabled alongside PII_REDACTION_ENABLED, replaces PII with reversible surrogate tokens +# during prompt processing and restores the original values in model responses (both SSE and JSON). +# PII_REVERSIBLE_VAULT_ENABLED=false +# PII_VAULT_TTL_MS=120000 +# PII_VAULT_MAX_ENTRIES=5000 + # ── VS Code Tokenized-Route Context Sanitizer ── # Strips implicit active-editor context (editorContext/activeEditor/currentFile/ # selection/openTabs…) from requests on the /v1/vscode/[token]/* routes before diff --git a/changelog.d/features/18017-reversible-pii-token-vault.md b/changelog.d/features/18017-reversible-pii-token-vault.md new file mode 100644 index 000000000000..9734d0b8cc75 --- /dev/null +++ b/changelog.d/features/18017-reversible-pii-token-vault.md @@ -0,0 +1 @@ +- **feat(pii):** reversible token vault with request-scoped nonces for streaming and non-streaming responses — thanks @Ardem2025 diff --git a/docs/reference/ENVIRONMENT.md b/docs/reference/ENVIRONMENT.md index ba58acd93e98..a641a824a65a 100644 --- a/docs/reference/ENVIRONMENT.md +++ b/docs/reference/ENVIRONMENT.md @@ -296,6 +296,9 @@ OmniRoute provides a two-layer defense: request-side injection scanning and resp | -------------------------------- | -------- | ------------------------- | ----------------------------------------------------------------------- | | `PII_RESPONSE_SANITIZATION` | `false` | `src/lib/piiSanitizer.ts` | Scan LLM responses for leaked PII before returning to client. | | `PII_RESPONSE_SANITIZATION_MODE` | `redact` | `src/lib/piiSanitizer.ts` | `redact` = mask PII, `warn` = log only, `block` = drop entire response. | +| `PII_REVERSIBLE_VAULT_ENABLED` | `false` | `src/lib/pii/tokenVault.ts` | Reversible in-memory PII token vault for streaming and non-streaming responses. | +| `PII_VAULT_TTL_MS` | `120000` | `src/lib/pii/tokenVault.ts` | TTL (ms) for surrogate token mappings in memory. | +| `PII_VAULT_MAX_ENTRIES` | `5000` | `src/lib/pii/tokenVault.ts` | Maximum entries per request token vault before LRU eviction. | ### VS Code Tokenized-Route Context Sanitizer diff --git a/docs/reference/FEATURE_FLAGS.md b/docs/reference/FEATURE_FLAGS.md index 00b1debf7cef..05c6a8eaea8b 100644 --- a/docs/reference/FEATURE_FLAGS.md +++ b/docs/reference/FEATURE_FLAGS.md @@ -46,10 +46,10 @@ A boolean flag is considered **enabled** when its effective value is `"true"`, ## Flag Catalog -116 flags across 6 categories. **Default** is the definition default — the value +117 flags across 6 categories. **Default** is the definition default — the value used when neither a DB override nor an environment variable is present. -### Security (11) +### Security (12) | Key | Type | Default | Description | | --------------------------------------- | ------- | -------- | ----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | @@ -59,6 +59,7 @@ used when neither a DB override nor an environment variable is present. | `PII_REDACTION_ENABLED` | boolean | `false` | Redact PII from requests (independent of `INPUT_SANITIZER_MODE`). | | `PII_RESPONSE_SANITIZATION` | boolean | `false` | Sanitize PII from provider responses. | | `PII_RESPONSE_SANITIZATION_MODE` | enum | `redact` | Mode for PII response sanitization. Values: `redact`, `warn`, `block`, `off`. | +| `PII_REVERSIBLE_VAULT_ENABLED` | boolean | `false` | Enable reversible in-memory PII token vault for streaming and non-streaming responses. | | `OUTBOUND_SSRF_GUARD_ENABLED` | boolean | `true` | Legacy alias: a value saved on this flag's dashboard toggle is read before the environment; `false`, `0`, `no`, or `off` in either turns the outbound URL guard's host checks off like `OMNIROUTE_ALLOW_PRIVATE_PROVIDER_URLS`. | | `ALLOW_API_KEY_REVEAL` | boolean | `false` | Allow authenticated dashboard users to reveal stored API keys instead of only seeing masked values. | | `AUTH_LOG_INCLUDE_ACCOUNT_ID` | boolean | `false` | Include account prefix in AUTH log lines (e.g. "Using account: abc12345..."). Disabled by default so account identifiers are redacted from shared/multi-tenant process logs. Independent from Debug Mode; flipping Debug Mode does not reveal this. | @@ -267,10 +268,10 @@ Returns every flag with its effective value, source, and a summary. "requiresRestart": false, "warningLevel": "caution", }, - // ... all 116 flags + // ... all 117 flags ], "summary": { - "total": 116, + "total": 117, "active": 0, "inactive": 0, "overriddenByDb": 0, diff --git a/open-sse/handlers/chatCore/nonStreamingJsonResponse.ts b/open-sse/handlers/chatCore/nonStreamingJsonResponse.ts index 3b3d91848e3b..98e792512182 100644 --- a/open-sse/handlers/chatCore/nonStreamingJsonResponse.ts +++ b/open-sse/handlers/chatCore/nonStreamingJsonResponse.ts @@ -2,11 +2,18 @@ * Build the final non-streaming JSON response body once and publish an accurate * Content-Length for downstream HTTP clients and buffering proxies. */ +import { + detokenizeResponsePayload, + isReversibleVaultEnabled, +} from "../../utils/tokenVault.ts"; export function buildNonStreamingJsonResponse( body: unknown, - headers: Record + headers: Record, ): Response { - const payload = JSON.stringify(body); + const finalBody = isReversibleVaultEnabled() + ? detokenizeResponsePayload(body) + : body; + const payload = JSON.stringify(finalBody); return new Response(payload, { headers: { ...headers, diff --git a/open-sse/handlers/chatCore/streamingPipeline.ts b/open-sse/handlers/chatCore/streamingPipeline.ts index 31202a7d6caa..119de52305f9 100644 --- a/open-sse/handlers/chatCore/streamingPipeline.ts +++ b/open-sse/handlers/chatCore/streamingPipeline.ts @@ -32,6 +32,7 @@ import { import { isFeatureFlagEnabled as defaultFeatureFlag } from "@/shared/utils/featureFlags"; import { OMNIROUTE_RESPONSE_HEADERS } from "@/shared/constants/headers"; import { SSE_HEARTBEAT_INTERVAL_MS } from "../../config/constants.ts"; +import { createDetokenizingStream, isReversibleVaultEnabled } from "../../utils/tokenVault.ts"; /** * Pipeline assembly instrumentation — performance.mark() along the SSE hot path. * Marks are visible to Node.js perf_hooks consumers and DevTools' Performance @@ -56,6 +57,8 @@ export interface StreamingPipelineDeps { createSseHeartbeatTransform: typeof defaultHeartbeat; shapeForClientFormat: typeof defaultShape; createModelEchoTransform: typeof defaultModelEcho; + createDetokenizingStream?: typeof createDetokenizingStream; + isReversibleVaultEnabled?: typeof isReversibleVaultEnabled; } const DEFAULT_DEPS: StreamingPipelineDeps = { @@ -68,8 +71,20 @@ const DEFAULT_DEPS: StreamingPipelineDeps = { createSseHeartbeatTransform: defaultHeartbeat, shapeForClientFormat: defaultShape, createModelEchoTransform: defaultModelEcho, + createDetokenizingStream: createDetokenizingStream, + isReversibleVaultEnabled: isReversibleVaultEnabled, }; +function attachReversibleVaultTransform( + stream: ReadableStream, + deps: StreamingPipelineDeps, +): ReadableStream { + const checkVault = deps.isReversibleVaultEnabled ?? isReversibleVaultEnabled; + if (!checkVault()) return stream; + const makeDetokenizer = deps.createDetokenizingStream ?? createDetokenizingStream; + return stream.pipeThrough(makeDetokenizer()); +} + export function assembleStreamingPipeline( args: { providerResponse: Parameters[0]; @@ -156,6 +171,7 @@ export function assembleStreamingPipeline( if (args.echoModel) { finalStream = finalStream.pipeThrough(deps.createModelEchoTransform(args.echoModel)); } + finalStream = attachReversibleVaultTransform(finalStream, deps); performance.mark(PIPELINE_END); performance.measure(PIPELINE_MEASURE, PIPELINE_START, PIPELINE_END); return finalStream; diff --git a/open-sse/utils/tokenVault.ts b/open-sse/utils/tokenVault.ts new file mode 100644 index 000000000000..3131e09f7bdf --- /dev/null +++ b/open-sse/utils/tokenVault.ts @@ -0,0 +1,400 @@ +/** + * Reversible PII Token Vault + * + * Provides bidirectional pseudonymization: + * 1. Inbound: Replaces detected sensitive PII values with unique surrogate tokens (e.g. <>). + * 2. Outbound: Speculatively buffers and detokenizes SSE chunks in real time (Zero-TTFT), and detokenizes + * non-streaming JSON payloads before sending them back to the client. + * + * Values are retained strictly in-memory with short TTL (default: 120s) and LRU eviction. + * Zero PII is persisted to disk or upstream provider logs. + * + * @module lib/pii/tokenVault + */ + +import { AsyncLocalStorage } from "node:async_hooks"; +import crypto from "node:crypto"; + +/** + * Checks whether Reversible PII Token Vault is enabled in settings or environment. + */ +export function isReversibleVaultEnabled(): boolean { + if (typeof process !== "undefined" && process.env) { + const val = process.env.PII_REVERSIBLE_VAULT_ENABLED; + return val === "true" || val === "1"; + } + return false; +} + +export interface VaultItem { + val: string; + exp: number; +} + +export interface TokenVaultOptions { + maxEntries?: number; + ttlMs?: number; + nonce?: string; + sweepIntervalMs?: number; +} + +/** + * Resolves an integer environment variable with fallback value. + * + * @param envVarName Name of the environment variable. + * @param defaultValue Fallback value if variable is missing or invalid. + */ +export function resolveVaultEnvInt(envVarName: string, defaultValue: number): number { + if (typeof process !== "undefined" && process.env) { + const val = process.env[envVarName]; + if (val) { + const parsed = parseInt(val, 10); + if (!Number.isNaN(parsed) && parsed > 0) { + return parsed; + } + } + } + return defaultValue; +} + +/** + * Safely unrefs a NodeJS or browser timer to avoid keeping the event loop alive. + * + * @param timer Timer object or numeric ID. + */ +export function unrefTimer(timer: unknown): void { + if (timer && typeof timer === "object" && typeof (timer as { unref?: () => void }).unref === "function") { + (timer as { unref: () => void }).unref(); + } +} + +/** + * Generates an 8-character uppercase hex nonce for per-instance token isolation. + */ +export function generateVaultNonce(): string { + return crypto.randomBytes(4).toString("hex").toUpperCase(); +} + +/** + * In-memory Reversible Token Vault with per-instance nonce, LRU eviction, and TTL sweep. + */ +export class TokenVault { + protected vault: Map; + protected valToToken: Map; + protected seq: number; + protected maxEntries: number; + protected ttlMs: number; + protected nonce: string; + protected cleanupTimer?: ReturnType; + + /** + * Initializes a new TokenVault instance. + * + * @param options Configuration options. + */ + constructor(options?: TokenVaultOptions) { + this.vault = new Map(); + this.valToToken = new Map(); + this.seq = 0; + this.maxEntries = options?.maxEntries ?? resolveVaultEnvInt("PII_VAULT_MAX_ENTRIES", 5000); + this.ttlMs = options?.ttlMs ?? resolveVaultEnvInt("PII_VAULT_TTL_MS", 120000); + this.nonce = options?.nonce || generateVaultNonce(); + + const sweepInterval = options?.sweepIntervalMs ?? 30000; + this.cleanupTimer = setInterval(() => this.sweep(), sweepInterval); + unrefTimer(this.cleanupTimer); + } + + /** + * Returns the per-instance nonce string. + */ + public getNonce(): string { + return this.nonce; + } + + /** + * Generates or retrieves an existing surrogate token for a given value. + * Refreshes TTL and performs an LRU touch on hit. + * + * @param name Entity type name (e.g. "EMAIL", "PHONE"). + * @param val Raw sensitive string value. + */ + public getOrCreateToken(name: string, val: string): string { + const existingToken = this.valToToken.get(val); + if (existingToken) { + const item = this.vault.get(existingToken); + if (item) { + item.exp = Date.now() + this.ttlMs; + // LRU touch: delete and re-insert to move to the end of Map iteration order + this.vault.delete(existingToken); + this.vault.set(existingToken, item); + return existingToken; + } + } + + this.seq = (this.seq + 1) % 1000000000; + const cleanName = name.replace(/[^A-Za-z0-9_]/g, "_").toUpperCase(); + const token = `<>`; + + // LRU eviction if capacity exceeded + if (this.vault.size >= this.maxEntries) { + const oldestKey = this.vault.keys().next().value; + if (oldestKey) { + const oldestVal = this.vault.get(oldestKey)?.val; + if (oldestVal) this.valToToken.delete(oldestVal); + this.vault.delete(oldestKey); + } + } + + this.vault.set(token, { val, exp: Date.now() + this.ttlMs }); + this.valToToken.set(val, token); + return token; + } + + /** + * Retrieves original value for a given token if not expired. + * + * @param token Surrogate token string. + */ + public get(token: string): string | undefined { + const item = this.vault.get(token); + if (!item) return undefined; + if (item.exp < Date.now()) { + this.vault.delete(token); + this.valToToken.delete(item.val); + return undefined; + } + return item.val; + } + + /** + * Checks whether this vault contains any active tokens. + */ + public hasTokens(): boolean { + return this.vault.size > 0; + } + + /** + * Returns current count of stored tokens. + */ + public get size(): number { + return this.vault.size; + } + + /** + * Clears all stored tokens. + */ + public clear(): void { + this.vault.clear(); + this.valToToken.clear(); + this.seq = 0; + } + + /** + * Disposes of interval timers and clears internal maps to prevent memory leaks. + */ + public dispose(): void { + if (this.cleanupTimer) { + clearInterval(this.cleanupTimer); + this.cleanupTimer = undefined; + } + this.vault.clear(); + this.valToToken.clear(); + } + + /** + * Sweeps expired items from vault. + */ + public sweep(): void { + const now = Date.now(); + for (const [k, v] of this.vault.entries()) { + if (v.exp < now) { + this.valToToken.delete(v.val); + this.vault.delete(k); + } + } + } +} + +/** + * Legacy TokenVaultManager singleton wrapper for backwards compatibility. + */ +export class TokenVaultManager extends TokenVault { + public static getInstance(): TokenVaultManager { + const globalObj = globalThis as unknown as { + __omniroute_token_vault_manager__?: TokenVaultManager; + }; + if (!globalObj.__omniroute_token_vault_manager__) { + globalObj.__omniroute_token_vault_manager__ = new TokenVaultManager(); + } + return globalObj.__omniroute_token_vault_manager__; + } + + public static resetInstance(): void { + const globalObj = globalThis as unknown as { + __omniroute_token_vault_manager__?: TokenVaultManager; + }; + if (globalObj.__omniroute_token_vault_manager__) { + globalObj.__omniroute_token_vault_manager__.dispose(); + delete globalObj.__omniroute_token_vault_manager__; + } + } +} + +/** + * AsyncLocalStorage container for request-scoped TokenVault instances. + */ +export const tokenVaultStorage = new AsyncLocalStorage(); + +/** + * Runs a callback within the context of a scoped TokenVault. + * + * @param vault TokenVault instance for current execution context. + * @param fn Callback function. + */ +export function runWithTokenVault(vault: TokenVault, fn: () => T): T { + return tokenVaultStorage.run(vault, fn); +} + +/** + * Gets the active request-scoped TokenVault from AsyncLocalStorage. + */ +export function getActiveTokenVault(): TokenVault | undefined { + return tokenVaultStorage.getStore(); +} + +/** + * Resolves the TokenVault to use: explicitVault if provided, otherwise active context vault. + * + * @param explicitVault Optional explicit TokenVault instance. + */ +export function resolveVault(explicitVault?: TokenVault): TokenVault | undefined { + if (explicitVault) { + return explicitVault; + } + return getActiveTokenVault(); +} + +/** + * Escapes replacement strings so they remain valid inside JSON string fields. + * + * @param val String value to escape. + */ +export function safeJsonEscape(val: string): string { + return JSON.stringify(val).slice(1, -1); +} + +/** + * Regex matching surrogate PII and RU tokens. + */ +const TOKEN_REGEX = /(<<(?:PII|RU)_[A-Za-z0-9_]+>>)/g; + +/** + * Replaces surrogate tokens in text with original values from vault. + * + * @param text Input text containing tokens. + * @param vault Active TokenVault instance. + */ +function replaceTokensInText(text: string, vault: TokenVault): string { + if (!text.includes("< { + const val = vault.get(match); + return val !== undefined ? safeJsonEscape(val) : match; + }); +} + +/** + * Creates a zero-TTFT speculatively buffered TransformStream for SSE detokenization. + * Emits raw chunks immediately unless an opening '<' token marker prefix is detected. + * + * @param explicitVault Optional explicit TokenVault instance. + */ +export function createDetokenizingStream( + explicitVault?: TokenVault +): TransformStream { + const dec = new TextDecoder("utf-8"); + const enc = new TextEncoder(); + const vault = resolveVault(explicitVault); + let pending = ""; + const MAX_TOKEN_LEN = 64; + + return new TransformStream({ + transform(chunk, controller) { + if (!vault) { + controller.enqueue(typeof chunk === "string" ? enc.encode(chunk) : chunk); + return; + } + + const text = typeof chunk === "string" ? chunk : dec.decode(chunk, { stream: true }); + pending += text; + + // Fast path: emit immediately if no surrogate prefix exists + if (!pending.includes("<")) { + controller.enqueue(enc.encode(pending)); + pending = ""; + return; + } + + // Check if tail has unclosed '<' or '<<' that could be part of an incoming token + let lastLt = pending.lastIndexOf("<"); + if ( + lastLt !== -1 && + pending.length - lastLt < MAX_TOKEN_LEN && + !pending.slice(lastLt).includes(">>") + ) { + if (lastLt > 0 && pending[lastLt - 1] === "<") { + lastLt--; + } + let toEmit = pending.slice(0, lastLt); + pending = pending.slice(lastLt); + + if (toEmit.length > 0) { + toEmit = replaceTokensInText(toEmit, vault); + controller.enqueue(enc.encode(toEmit)); + } + } else { + let toEmit = pending; + pending = ""; + toEmit = replaceTokensInText(toEmit, vault); + controller.enqueue(enc.encode(toEmit)); + } + }, + flush(controller) { + if (pending.length > 0) { + let toEmit = pending; + if (vault) { + toEmit = replaceTokensInText(toEmit, vault); + } + controller.enqueue(enc.encode(toEmit)); + pending = ""; + } + }, + }); +} + +/** + * Detokenizes non-streaming JSON response payloads before sending to the client. + * + * @param obj Response body payload object. + * @param explicitVault Optional explicit TokenVault instance. + */ +export function detokenizeResponsePayload(obj: T, explicitVault?: TokenVault): T { + const vault = resolveVault(explicitVault); + if (!obj || !vault || !vault.hasTokens()) return obj; + + try { + let str = JSON.stringify(obj); + if (str.includes("< { + const val = vault.get(match); + return val !== undefined ? safeJsonEscape(val) : match; + }); + return JSON.parse(str) as T; + } + } catch (_err) { + // Fail-closed: return unmodified payload on parsing or replacement error + } + return obj; +} diff --git a/src/lib/pii/tokenVault.ts b/src/lib/pii/tokenVault.ts new file mode 100644 index 000000000000..ba1b2a80e095 --- /dev/null +++ b/src/lib/pii/tokenVault.ts @@ -0,0 +1 @@ +export * from "../../../open-sse/utils/tokenVault.ts"; diff --git a/src/shared/constants/featureFlagDefinitions.ts b/src/shared/constants/featureFlagDefinitions.ts index 2ca5cc4c70f5..4a4ea487c033 100644 --- a/src/shared/constants/featureFlagDefinitions.ts +++ b/src/shared/constants/featureFlagDefinitions.ts @@ -94,6 +94,18 @@ export const FEATURE_FLAG_DEFINITIONS: FeatureFlagDefinition[] = [ requiresRestart: false, warningLevel: "info", }, + { + key: "PII_REVERSIBLE_VAULT_ENABLED", + label: "Reversible PII Token Vault", + description: + "Replace detected PII with reversible surrogate tokens and restore them in responses", + descriptionI18nKey: "featureFlagPiiReversibleVaultEnabledDescription", + category: "security", + defaultValue: "false", + type: "boolean", + requiresRestart: false, + warningLevel: "info", + }, { key: "OUTBOUND_SSRF_GUARD_ENABLED", label: "SSRF Guard", diff --git a/src/shared/utils/inputSanitizer.ts b/src/shared/utils/inputSanitizer.ts index cd164763e26b..ac531645ace4 100644 --- a/src/shared/utils/inputSanitizer.ts +++ b/src/shared/utils/inputSanitizer.ts @@ -1,3 +1,7 @@ +import { + getActiveTokenVault, + isReversibleVaultEnabled, +} from "../../lib/pii/tokenVault.ts"; /** * Input Sanitizer — FASE-01 Security Hardening * @@ -377,7 +381,18 @@ function processPII(text, redact = false, options = undefined) { if (matches && matches.length > 0) { detections.push({ type: rule.name, count: matches.length }); if (redact) { - processed = processed.replace(rule.pattern, rule.replacement); + if (isReversibleVaultEnabled()) { + const vault = getActiveTokenVault(); + if (vault) { + processed = processed.replace(rule.pattern, (m) => vault.getOrCreateToken(rule.name, m)); + } else { + // Fail-closed fallback: when reversible vault is enabled but no active context exists, + // fall back to safe standard redaction to prevent raw PII exposure. + processed = processed.replace(rule.pattern, rule.replacement); + } + } else { + processed = processed.replace(rule.pattern, rule.replacement); + } } } } diff --git a/tests/unit/feature-flags-settings.test.ts b/tests/unit/feature-flags-settings.test.ts index 1040aaf6a2a6..11f5efc6f3af 100644 --- a/tests/unit/feature-flags-settings.test.ts +++ b/tests/unit/feature-flags-settings.test.ts @@ -64,7 +64,8 @@ const { // audit rows, default off) takes it to 115. // CLAUDE_CACHE_TAIL_TTL_1H (1h TTL for the Claude cache conversation tail, default off) // takes it to 116. -const EXPECTED_FEATURE_FLAG_COUNT = 116; +// PII_REVERSIBLE_VAULT_ENABLED (#18017, default off) takes it to 117. +const EXPECTED_FEATURE_FLAG_COUNT = 117; // tool_call wrapper repair, default on) to 86; OPENCODE_V2_CLIENT_FINGERPRINT_ENABLED // (#16571, OpenCode v2 client fingerprint, default off) to 87. // tool_call wrapper repair, default on) to 86; CLAUDE_OAUTH_SUMMARIZED_THINKING_UNREDACTED diff --git a/tests/unit/guardrails/tokenVault.test.ts b/tests/unit/guardrails/tokenVault.test.ts new file mode 100644 index 000000000000..45b0229ac09e --- /dev/null +++ b/tests/unit/guardrails/tokenVault.test.ts @@ -0,0 +1,246 @@ +import test from "node:test"; +import assert from "node:assert/strict"; + +import { + TokenVault, + createDetokenizingStream, + detokenizeResponsePayload, + safeJsonEscape, + runWithTokenVault, +} from "../../../src/lib/pii/tokenVault.ts"; + +async function readAll( + stream: TransformStream, + inputs: string[], +): Promise { + const writer = stream.writable.getWriter(); + const reader = stream.readable.getReader(); + const dec = new TextDecoder("utf-8"); + const enc = new TextEncoder(); + + const readPromise = (async () => { + let res = ""; + while (true) { + const { done, value } = await reader.read(); + if (done) break; + res += dec.decode(value); + } + return res; + })(); + + for (const chunk of inputs) { + await writer.write(enc.encode(chunk)); + } + await writer.close(); + + return await readPromise; +} + +test("TokenVault - surrogate generation, reuse, and TTL expiration", async () => { + const vault = new TokenVault({ maxEntries: 10, ttlMs: 100 }); + + const token1 = vault.getOrCreateToken("email", "user@example.com"); + assert.match(token1, /^<>$/); + assert.equal(vault.get(token1), "user@example.com"); + + // Re-requesting the same value yields identical token and refreshes TTL + const token1Again = vault.getOrCreateToken("email", "user@example.com"); + assert.equal(token1Again, token1); + + // Different value generates distinct token + const token2 = vault.getOrCreateToken("phone", "+1-555-0199"); + assert.notEqual(token2, token1); + assert.equal(vault.get(token2), "+1-555-0199"); + + // Wait for TTL expiration + await new Promise((r) => setTimeout(r, 120)); + assert.equal(vault.get(token1), undefined); + assert.equal(vault.get(token2), undefined); + + vault.dispose(); +}); + +test("TokenVault - LRU capacity eviction", () => { + const vault = new TokenVault({ maxEntries: 2, ttlMs: 60000 }); + + const t1 = vault.getOrCreateToken("item1", "value1"); + const t2 = vault.getOrCreateToken("item2", "value2"); + const t3 = vault.getOrCreateToken("item3", "value3"); + + // Earliest item (t1) should be evicted + assert.equal(vault.get(t1), undefined); + assert.equal(vault.get(t2), "value2"); + assert.equal(vault.get(t3), "value3"); + + vault.dispose(); +}); + +test("TokenVault - LRU touch prevents eviction of frequently accessed items", () => { + const vault = new TokenVault({ maxEntries: 2, ttlMs: 60000 }); + + const t1 = vault.getOrCreateToken("item1", "value1"); + const t2 = vault.getOrCreateToken("item2", "value2"); + + // Access t1 again, making t2 the least recently used + const t1Again = vault.getOrCreateToken("item1", "value1"); + assert.equal(t1Again, t1); + + // Add t3, which should evict t2 instead of t1 + const t3 = vault.getOrCreateToken("item3", "value3"); + + assert.equal(vault.get(t1), "value1"); + assert.equal(vault.get(t2), undefined); + assert.equal(vault.get(t3), "value3"); + + vault.dispose(); +}); + +test("TokenVault - dispose() clears interval timer and empties internal storage", () => { + const vault = new TokenVault({ maxEntries: 10, ttlMs: 60000 }); + const t1 = vault.getOrCreateToken("email", "admin@domain.com"); + assert.equal(vault.get(t1), "admin@domain.com"); + + vault.dispose(); + assert.equal(vault.get(t1), undefined); +}); + +test("TokenVault - concurrent requests cannot detokenize each other's tokens or echoes", async () => { + const vaultA = new TokenVault(); + const vaultB = new TokenVault(); + + const tokenA = vaultA.getOrCreateToken("email", "alice@example.com"); + const tokenB = vaultB.getOrCreateToken("email", "bob@example.com"); + + assert.notEqual(tokenA, tokenB); + + // Request B response contains its own token, Request A's token, and a generic token <> + const responsePayloadB = { + choices: [ + { + message: { + role: "assistant", + content: `Here is token B: ${tokenB}, echoed token A: ${tokenA}, and generic echo: <>`, + }, + }, + ], + }; + + const detokenizedB = detokenizeResponsePayload(responsePayloadB, vaultB); + const contentB = detokenizedB.choices[0].message.content; + + // Must restore Bob's email + assert.match(contentB, /Here is token B: bob@example\.com/); + // Must NOT restore Alice's email + assert.doesNotMatch(contentB, /alice@example\.com/); + // Must preserve tokenA intact + assert.ok( + contentB.includes(tokenA), + "Request A token must remain raw surrogate in Request B", + ); + // Must preserve generic echo <> intact + assert.ok( + contentB.includes("<>"), + "Generic <> echo must not be restored", + ); + + vaultA.dispose(); + vaultB.dispose(); +}); + +test("createDetokenizingStream - Zero-TTFT fast path without tokens", async () => { + const stream = createDetokenizingStream(); + const result = await readAll(stream, ["Hello world! How are you?"]); + assert.equal(result, "Hello world! How are you?"); +}); + +test("createDetokenizingStream - detokenizes single chunk with surrogate tokens", async () => { + const vault = new TokenVault(); + const emailToken = vault.getOrCreateToken("email", "user@example.com"); + + const stream = createDetokenizingStream(vault); + const payload = `data: {"choices":[{"delta":{"content":"Contact me at ${emailToken}"}}]}\n\n`; + const result = await readAll(stream, [payload]); + + assert.match(result, /Contact me at user@example\.com/); + assert.doesNotMatch(result, /< { + const vault = new TokenVault(); + const phoneToken = vault.getOrCreateToken("phone", "+1-800-555-0199"); + + // Split surrogate token across two chunks: e.g. "<>" + const mid = Math.floor(phoneToken.length / 2); + const part1 = + `data: {"choices":[{"delta":{"content":"Call ` + phoneToken.slice(0, mid); + const part2 = phoneToken.slice(mid) + ` anytime"}}]}\n\n`; + + const stream = createDetokenizingStream(vault); + const result = await readAll(stream, [part1, part2]); + + assert.match(result, /Call \+1-800-555-0199 anytime/); + assert.doesNotMatch(result, /< { + const vaultA = new TokenVault(); + const vaultB = new TokenVault(); + + const tokenA = vaultA.getOrCreateToken("phone", "+1-800-555-0100"); + const tokenB = vaultB.getOrCreateToken("phone", "+1-800-555-0200"); + + const streamA = createDetokenizingStream(vaultA); + const streamB = createDetokenizingStream(vaultB); + + // Cross-feed token A into stream B + const [resA, resB] = await Promise.all([ + readAll(streamA, [`data: {"text":"Call ${tokenA}"}\n\n`]), + readAll(streamB, [ + `data: {"text":"Call ${tokenB} and foreign ${tokenA}"}\n\n`, + ]), + ]); + + assert.match(resA, /Call \+1-800-555-0100/); + assert.match(resB, /Call \+1-800-555-0200/); + assert.doesNotMatch(resB, /\+1-800-555-0100/); + assert.ok(resB.includes(tokenA)); + + vaultA.dispose(); + vaultB.dispose(); +}); + +test("detokenizeResponsePayload - non-streaming JSON body detokenization with AsyncLocalStorage", () => { + const vault = new TokenVault(); + const ssnToken = vault.getOrCreateToken("ssn", "000-12-3456"); + + const responseBody = { + id: "chatcmpl-test", + choices: [ + { + message: { + role: "assistant", + content: `Verified SSN record: ${ssnToken}.`, + }, + }, + ], + }; + + const restored = runWithTokenVault(vault, () => + detokenizeResponsePayload(responseBody), + ); + assert.equal( + restored.choices[0].message.content, + "Verified SSN record: 000-12-3456.", + ); + + vault.dispose(); +}); + +test("safeJsonEscape - preserves quotes and special characters inside SSE", () => { + assert.equal(safeJsonEscape('Hello "world"'), 'Hello \\"world\\"'); + assert.equal(safeJsonEscape("Line 1\nLine 2"), "Line 1\\nLine 2"); +}); diff --git a/tests/unit/pii/featureFlag.test.ts b/tests/unit/pii/featureFlag.test.ts new file mode 100644 index 000000000000..f47bf2f39130 --- /dev/null +++ b/tests/unit/pii/featureFlag.test.ts @@ -0,0 +1,101 @@ +import test from "node:test"; +import assert from "node:assert/strict"; +import fs from "node:fs"; +import path from "node:path"; + +import { FEATURE_FLAG_DEFINITIONS } from "../../../src/shared/constants/featureFlagDefinitions.ts"; + +test("Feature Flags - Registry length and PII_REVERSIBLE_VAULT_ENABLED definition", () => { + assert.equal( + FEATURE_FLAG_DEFINITIONS.length, + 117, + "Expected exactly 117 feature flags (116 upstream tip + PII_REVERSIBLE_VAULT_ENABLED)", + ); + + const piiFlag = FEATURE_FLAG_DEFINITIONS.find( + (f) => f.key === "PII_REVERSIBLE_VAULT_ENABLED", + ); + assert.ok(piiFlag, "PII_REVERSIBLE_VAULT_ENABLED must be defined in registry"); + assert.equal(piiFlag.category, "security"); + assert.equal(piiFlag.defaultValue, "false"); + assert.equal(piiFlag.type, "boolean"); + + const auditFlag = FEATURE_FLAG_DEFINITIONS.find( + (f) => f.key === "API_KEY_AUDIT_CLIENT_DETAILS_ENABLED", + ); + assert.ok( + auditFlag, + "API_KEY_AUDIT_CLIENT_DETAILS_ENABLED must be preserved (preventing #18017 CI regression)", + ); + + const securityFlags = FEATURE_FLAG_DEFINITIONS.filter( + (f) => f.category === "security", + ); + assert.equal( + securityFlags.length, + 12, + "Security category must contain exactly 12 flags (11 upstream + PII vault)", + ); +}); + +test("Feature Flags - Documentation sync with FEATURE_FLAGS.md", () => { + const catalogPath = path.resolve( + process.cwd(), + "docs/reference/FEATURE_FLAGS.md", + ); + const catalog = fs.readFileSync(catalogPath, "utf8"); + + const countMatch = catalog.match(/^(\d+) flags across \d+ categories\./m); + assert.ok(countMatch, "FEATURE_FLAGS.md must declare total count in header"); + assert.equal( + Number(countMatch[1]), + FEATURE_FLAG_DEFINITIONS.length, + "Catalog total must match FEATURE_FLAG_DEFINITIONS.length", + ); + + assert.ok( + catalog.includes("PII_REVERSIBLE_VAULT_ENABLED"), + "FEATURE_FLAGS.md must document PII_REVERSIBLE_VAULT_ENABLED", + ); + assert.ok( + catalog.includes("### Security (12)"), + "FEATURE_FLAGS.md must declare 12 security flags", + ); +}); + +test("Feature Flags - Documentation sync with ENVIRONMENT.md and .env.example", () => { + const envDocPath = path.resolve( + process.cwd(), + "docs/reference/ENVIRONMENT.md", + ); + const envDoc = fs.readFileSync(envDocPath, "utf8"); + + assert.ok( + envDoc.includes("PII_REVERSIBLE_VAULT_ENABLED"), + "ENVIRONMENT.md must document PII_REVERSIBLE_VAULT_ENABLED", + ); + assert.ok( + envDoc.includes("PII_VAULT_TTL_MS"), + "ENVIRONMENT.md must document PII_VAULT_TTL_MS", + ); + assert.ok( + envDoc.includes("PII_VAULT_MAX_ENTRIES"), + "ENVIRONMENT.md must document PII_VAULT_MAX_ENTRIES", + ); + + const envExamplePath = path.resolve(process.cwd(), ".env.example"); + const envExample = fs.readFileSync(envExamplePath, "utf8"); + + assert.ok( + envExample.includes("PII_REVERSIBLE_VAULT_ENABLED=false"), + ".env.example must specify default PII_REVERSIBLE_VAULT_ENABLED=false", + ); + assert.ok( + envExample.includes("PII_VAULT_TTL_MS=120000"), + ".env.example must specify default PII_VAULT_TTL_MS=120000", + ); + assert.ok( + envExample.includes("PII_VAULT_MAX_ENTRIES=5000"), + ".env.example must specify default PII_VAULT_MAX_ENTRIES=5000", + ); +}); diff --git a/tests/unit/pii/tokenVault.test.ts b/tests/unit/pii/tokenVault.test.ts new file mode 100644 index 000000000000..02131e7b4046 --- /dev/null +++ b/tests/unit/pii/tokenVault.test.ts @@ -0,0 +1,260 @@ +import test from "node:test"; +import assert from "node:assert/strict"; + +import { + TokenVault, + TokenVaultManager, + createDetokenizingStream, + detokenizeResponsePayload, + safeJsonEscape, + runWithTokenVault, + getActiveTokenVault, + resolveVault, +} from "../../../src/lib/pii/tokenVault.ts"; + +async function readAll( + stream: TransformStream, + inputs: string[] +): Promise { + const writer = stream.writable.getWriter(); + const reader = stream.readable.getReader(); + const dec = new TextDecoder("utf-8"); + const enc = new TextEncoder(); + + const readPromise = (async () => { + let res = ""; + while (true) { + const { done, value } = await reader.read(); + if (done) break; + res += dec.decode(value); + } + return res; + })(); + + for (const chunk of inputs) { + await writer.write(enc.encode(chunk)); + } + await writer.close(); + + return await readPromise; +} + +test("TokenVaultManager - surrogate generation, reuse, and TTL expiration", async () => { + const vault = new TokenVaultManager({ maxEntries: 10, ttlMs: 100 }); + + const token1 = vault.getOrCreateToken("email", "admin@company.com"); + assert.match(token1, /^<>$/); + assert.equal(vault.get(token1), "admin@company.com"); + + // Re-requesting the same value yields identical token and refreshes TTL + const token1Again = vault.getOrCreateToken("email", "admin@company.com"); + assert.equal(token1Again, token1); + + // Different value generates distinct token + const token2 = vault.getOrCreateToken("phone", "+1-555-0199"); + assert.notEqual(token2, token1); + assert.equal(vault.get(token2), "+1-555-0199"); + + // Wait for TTL expiration + await new Promise((r) => setTimeout(r, 120)); + assert.equal(vault.get(token1), undefined); + assert.equal(vault.get(token2), undefined); + + vault.dispose(); +}); + +test("TokenVaultManager - LRU capacity eviction", () => { + const vault = new TokenVaultManager({ maxEntries: 2, ttlMs: 60000 }); + + const t1 = vault.getOrCreateToken("item1", "value1"); + const t2 = vault.getOrCreateToken("item2", "value2"); + const t3 = vault.getOrCreateToken("item3", "value3"); + + // Earliest item (t1) should be evicted + assert.equal(vault.get(t1), undefined); + assert.equal(vault.get(t2), "value2"); + assert.equal(vault.get(t3), "value3"); + + vault.dispose(); +}); + +test("TokenVault - LRU touch prevents eviction of frequently accessed items", () => { + const vault = new TokenVault({ maxEntries: 3, ttlMs: 60000 }); + + const t1 = vault.getOrCreateToken("item1", "val1"); + const t2 = vault.getOrCreateToken("item2", "val2"); + const t3 = vault.getOrCreateToken("item3", "val3"); + + // Touch t1 so it moves to most-recently-used position + const t1Again = vault.getOrCreateToken("item1", "val1"); + assert.equal(t1Again, t1); + + // Insert 4th item, exceeding capacity 3. Oldest item now should be t2 (not t1!) + const t4 = vault.getOrCreateToken("item4", "val4"); + + assert.equal(vault.get(t1), "val1", "t1 should NOT be evicted because of LRU touch"); + assert.equal(vault.get(t2), undefined, "t2 should be evicted as least recently used"); + assert.equal(vault.get(t3), "val3"); + assert.equal(vault.get(t4), "val4"); + + vault.dispose(); +}); + +test("createDetokenizingStream - Zero-TTFT fast path without tokens", async () => { + const stream = createDetokenizingStream(); + const result = await readAll(stream, ["Hello world! How are you?"]); + assert.equal(result, "Hello world! How are you?"); +}); + +test("createDetokenizingStream - detokenizes single chunk with surrogate tokens", async () => { + const vault = TokenVaultManager.getInstance(); + const emailToken = vault.getOrCreateToken("email", "user@example.com"); + + const stream = createDetokenizingStream(vault); + const payload = `data: {"choices":[{"delta":{"content":"Contact me at ${emailToken}"}}]}\n\n`; + const result = await readAll(stream, [payload]); + + assert.match(result, /Contact me at user@example\.com/); + assert.doesNotMatch(result, /< { + const vault = TokenVaultManager.getInstance(); + const phoneToken = vault.getOrCreateToken("phone", "+1-800-555-0199"); + + // Split surrogate token across two chunks: e.g. "<>" + const mid = Math.floor(phoneToken.length / 2); + const part1 = `data: {"choices":[{"delta":{"content":"Call ` + phoneToken.slice(0, mid); + const part2 = phoneToken.slice(mid) + ` anytime"}]}\n\n`; + + const stream = createDetokenizingStream(vault); + const result = await readAll(stream, [part1, part2]); + + assert.match(result, /Call \+1-800-555-0199 anytime/); + assert.doesNotMatch(result, /< { + const vault = TokenVaultManager.getInstance(); + const ssnToken = vault.getOrCreateToken("ssn", "123-45-6789"); + + const responseBody = { + id: "chatcmpl-test", + choices: [ + { + message: { + role: "assistant", + content: `Verified SSN record: ${ssnToken}.`, + }, + }, + ], + }; + + const restored = detokenizeResponsePayload(responseBody, vault); + assert.equal( + restored.choices[0].message.content, + "Verified SSN record: 123-45-6789." + ); +}); + +test("safeJsonEscape - preserves quotes and special characters inside SSE", () => { + assert.equal(safeJsonEscape('Hello "world"'), 'Hello \\"world\\"'); + assert.equal(safeJsonEscape("Line 1\nLine 2"), "Line 1\\nLine 2"); +}); + +test("AsyncLocalStorage - parallel requests in runWithTokenVault are isolated", async () => { + const vaultA = new TokenVault(); + const vaultB = new TokenVault(); + + const tokenA = vaultA.getOrCreateToken("email", "alice@example.com"); + const tokenB = vaultB.getOrCreateToken("email", "bob@example.com"); + + // Run two concurrent async requests + const promiseA = runWithTokenVault(vaultA, async () => { + // Verify active vault resolution + assert.equal(getActiveTokenVault(), vaultA); + assert.equal(resolveVault(), vaultA); + + // Simulate async tick + await new Promise((r) => setTimeout(r, 10)); + + // Detokenizing stream without explicit vault should pick up vaultA + const stream = createDetokenizingStream(); + const input = `Client A: ${tokenA}, Alien B: ${tokenB}`; + const result = await readAll(stream, [input]); + return result; + }); + + const promiseB = runWithTokenVault(vaultB, async () => { + assert.equal(getActiveTokenVault(), vaultB); + assert.equal(resolveVault(), vaultB); + + await new Promise((r) => setTimeout(r, 10)); + + // Detokenizing response payload without explicit vault should pick up vaultB + const payload = { + content: `Client B: ${tokenB}, Alien A: ${tokenA}`, + }; + const detokenized = detokenizeResponsePayload(payload); + return detokenized.content; + }); + + const [resultA, resultB] = await Promise.all([promiseA, promiseB]); + + // Request A: alice is detokenized, alien tokenB remains untouched (fail-closed) + assert.match(resultA, /Client A: alice@example\.com/); + assert.match(resultA, new RegExp(`Alien B: ${tokenB}`)); + assert.doesNotMatch(resultA, /bob@example\.com/); + + // Request B: bob is detokenized, alien tokenA remains untouched (fail-closed) + assert.match(resultB, /Client B: bob@example\.com/); + assert.match(resultB, new RegExp(`Alien A: ${tokenA}`)); + assert.doesNotMatch(resultB, /alice@example\.com/); + + vaultA.dispose(); + vaultB.dispose(); +}); + +test("TokenVault - dispose() clears interval timer and empties internal storage", () => { + const vault = new TokenVault(); + const token = vault.getOrCreateToken("card", "4111-2222-3333-4444"); + + assert.equal(vault.hasTokens(), true); + assert.equal(vault.size, 1); + assert.equal(vault.get(token), "4111-2222-3333-4444"); + + vault.dispose(); + + assert.equal(vault.hasTokens(), false); + assert.equal(vault.size, 0); + assert.equal(vault.get(token), undefined); +}); + +test("Layer 2 Interceptor: Fail-Closed without active TokenVault context", () => { + // When no active request vault exists in AsyncLocalStorage, resolveVault() fails closed + const stream = createDetokenizingStream(); + assert.ok(stream); + + // detokenizeResponsePayload passes through safely when no active vault + const payload = { text: "Unresolved <>" }; + const restored = detokenizeResponsePayload(payload); + assert.equal(restored.text, "Unresolved <>"); +}); + +test("Layer 2 Interceptor: End-to-end scoped request streaming detokenization", async () => { + const requestVault = new TokenVault({ nonce: "E2ETEST1" }); + const secretPhone = "+7 (999) 123-45-67"; + const surrogateToken = requestVault.getOrCreateToken("phone", secretPhone); + + await runWithTokenVault(requestVault, async () => { + const stream = createDetokenizingStream(); + const chunk1 = `data: {"choices":[{"delta":{"content":"Contact support at ${surrogateToken}"}}]}\n\n`; + const chunk2 = `data: [DONE]\n\n`; + + const result = await readAll(stream, [chunk1, chunk2]); + assert.match(result, /Contact support at \+7 \(999\) 123-45-67/); + assert.doesNotMatch(result, /< { // REASONING_REPLAY_ENABLED (#12486, default on) to 84; COMBO_AUTO_PRUNE_STALE_STEPS // (#13505, default off) to 85; KIRO_TOOL_CALL_REPAIR_ENABLED (default on) to 86; // OPENCODE_PUBLISHED_CATALOG_ROUTING_ENABLED (#16798, default off) to 87. - assert.equal(FEATURE_FLAG_DEFINITIONS.length, 116); + assert.equal(FEATURE_FLAG_DEFINITIONS.length, 117); // COMPRESSION_CACHE_AUTOMATIC_PREFIX_ESTIMATE (#17114, default off) to 114. // CURSOR_GROK_47_DROP_MAX_EFFORT (#17386, default off) to 112. // UNPRICED_USAGE_BUDGET_POLICY (#14604, default fail_closed) to 113. // API_KEY_AUDIT_CLIENT_DETAILS_ENABLED (#17982, default off) to 115. // CLAUDE_CACHE_TAIL_TTL_1H (Claude cache tail 1h TTL, default off) to 116. + // PII_REVERSIBLE_VAULT_ENABLED (#18017, default off) to 117. // OPENCODE_V2_CLIENT_FINGERPRINT_ENABLED (#16571, default off) to 87. // COMBO_EMPTY_POOL_NO_ELIGIBLE_TARGETS_ENABLED (#16256, default off) to 87. // OPENCODE_ZEN_EGRESS_IP_LOCK_ENABLED (#16797, default off) to 87;