diff --git a/.gitignore b/.gitignore index 24e8c42..748a86a 100644 --- a/.gitignore +++ b/.gitignore @@ -18,3 +18,7 @@ node_modules/ # ggshield local cache .cache_ggshield + +# Master key — never commit (default lives in ~/.config, but guard the project root) +master.key +*.master.key diff --git a/README.md b/README.md index facacda..8f9dc80 100644 --- a/README.md +++ b/README.md @@ -18,17 +18,28 @@ claude-code ◄──SSE (detokenized)─── proxy ◄──SSE stream── ## How it works 1. **Request interception** — the JSON body of `POST /v1/messages` is scanned by - `ggshield`. Every detected secret is replaced with a unique - `<<>>` placeholder, and the `token → secret` mapping is kept - in memory (`src/store.ts`). The redacted body is forwarded upstream. -2. **Streaming detokenization** — the Anthropic SSE response is piped through a - stateful `Transform` (`src/detokenize.ts`) that scans the byte stream for - placeholders and substitutes the original secret back. It holds back only the - few characters that could still become a token, so a placeholder split across - chunk boundaries (`…<<>>…`) is reassembled correctly. + `ggshield`. Every detected secret is encrypted with **AES-256-GCM** under a + persistent master key and replaced with a self-contained + `<<>>` placeholder, where the hex payload is `IV ++ authTag ++ + ciphertext`. No server-side mapping is kept — the token carries everything + needed to restore it (`src/tokenize.ts`, `src/store.ts`). The redacted body is + forwarded upstream. +2. **Streaming detokenization** — the response is piped through a stateful + `Transform` (`src/detokenize.ts`) that finds placeholders, decrypts the hex + payload, and substitutes the original secret back. For SSE streams it parses + `content_block_delta` events and reassembles the reconstructed text, because + the model emits a long placeholder fragmented across many deltas + (`…"text":"<<>>"…`) with event framing wedged + between the fragments — a raw byte scan can't bridge that. Plain JSON + responses, where the token is contiguous, take a simpler direct byte scan. -Secrets live only in memory for the lifetime of the process — restart the proxy -and the mapping is gone. +The proxy is **stateless**: tokens are cryptographically self-contained, so they +resolve correctly even after a restart or when replaying older chat histories. +The only persistent state is the 256-bit master key, generated on first run and +stored `0600` at `~/.config/claude-tokenization-proxy/master.key` (override with +`GG_MASTER_KEY_FILE`). Anyone with that key file can decrypt captured tokens — +treat it like any other secret, and deleting it makes existing tokens +unrecoverable. ## Requirements @@ -160,8 +171,8 @@ that session. To make it stick: keys/tokens). A secret literally containing a backslash or double-quote could mis-escape inside the response JSON. See the `ponytail:` note in `src/detokenize.ts`. -- The token map is in-memory only; it is not shared across proxy restarts or - multiple proxy processes. +- Tokens are decryptable by anyone holding the master key file. Multiple proxy + instances must share the same `master.key` to resolve each other's tokens. ## License diff --git a/src/detokenize.ts b/src/detokenize.ts index fefbf7c..0e5c5c5 100644 --- a/src/detokenize.ts +++ b/src/detokenize.ts @@ -8,10 +8,10 @@ import { resolveToken as storeResolve } from "./store.ts"; // Token shape: <<>> (hex, because that's what we generate). const PREFIX = "<<>>/; +type Resolver = (token: string) => string | undefined; + // Could `s` (which starts with '<') still grow into a token with more input? // True for any strict prefix of a well-formed token. function isPartialPrefix(s: string): boolean { @@ -23,33 +23,30 @@ function isPartialPrefix(s: string): boolean { } /** - * Stateful SSE detokenizer. Scans the raw upstream byte stream for our - * placeholder tokens and swaps them back to the original secrets on the fly, - * holding back only the few characters that could still become a token so - * tokens split across chunk boundaries are reassembled correctly. + * Core token scanner. Feed it text fragments; it returns the part that is safe + * to emit (with complete tokens swapped for their secrets) and holds back only + * the trailing characters that could still grow into a token. `escape` controls + * whether a resolved secret is JSON-string-escaped on the way out: true when we + * splice it straight into a raw JSON body, false when the caller re-serializes + * the surrounding object itself (SSE deltas). */ -export class DetokenizeStream extends Transform { +class TokenScanner { private buf = ""; - private decoder = new StringDecoder("utf8"); - private resolve: (token: string) => string | undefined; - - constructor(resolve: (token: string) => string | undefined = storeResolve) { - super(); + private resolve: Resolver; + private escape: boolean; + constructor(resolve: Resolver, escape: boolean) { this.resolve = resolve; + this.escape = escape; } private emitFor(token: string): string { const secret = this.resolve(token); - if (secret === undefined) return token; // unknown token: leave it untouched - // ponytail: token appears inside a JSON string in the SSE `data:` line, so - // re-escape. Assumes secrets are escape-neutral (typical API keys/tokens); - // a secret literally containing a backslash or quote would double-escape — - // upgrade to context-aware splicing if that ever bites. - return JSON.stringify(secret).slice(1, -1); + if (secret === undefined) return token; // unknown/corrupt token: leave untouched + return this.escape ? JSON.stringify(secret).slice(1, -1) : secret; } - override _transform(chunk: Buffer, _enc: BufferEncoding, cb: TransformCallback): void { - this.buf += this.decoder.write(chunk); + push(text: string): string { + this.buf += text; let out = ""; while (this.buf.length > 0) { const lt = this.buf.indexOf("<"); @@ -67,18 +64,120 @@ export class DetokenizeStream extends Transform { this.buf = this.buf.slice(complete[0].length); continue; } - if (isPartialPrefix(this.buf)) break; // might finish next chunk; hold it - // This '<' can't begin a token. Emit it and keep scanning past it. - out += "<"; + if (isPartialPrefix(this.buf)) break; // might finish next fragment; hold it + out += "<"; // this '<' can't begin a token this.buf = this.buf.slice(1); } - cb(null, out); + return out; + } + + // Emit whatever is held back verbatim. Anything still buffered here is, by + // construction, an incomplete token, so it can only be passed through as-is. + flush(): string { + const out = this.buf; + this.buf = ""; + return out; + } +} + +/** + * Stateful detokenizer for the upstream response body. Two modes: + * + * - Raw (plain JSON responses): scan the byte stream directly. Tokens are + * contiguous in a single body, so the scanner reassembles them across TCP + * chunk boundaries and re-escapes the spliced-in secret. + * + * - SSE (streaming responses): the model emits a placeholder fragmented across + * many `content_block_delta` events, with JSON/SSE framing wedged between the + * token's characters. Scanning raw bytes can't bridge that, so we parse each + * event, run the *reconstructed* delta text through the scanner (which holds + * a partial token across deltas), and re-serialize. A token can't span a + * content block, so we flush any held tail at `content_block_stop`. + */ +export class DetokenizeStream extends Transform { + private decoder = new StringDecoder("utf8"); + private sse: boolean; + private scanner: TokenScanner; + private sseBuf = ""; // bytes not yet forming a complete `\n\n`-terminated event + private currentIndex = 0; // index of the open content block (for flush deltas) + + constructor(opts: { sse?: boolean; resolve?: Resolver } = {}) { + super(); + this.sse = opts.sse ?? false; + // Raw mode splices into a JSON body (escape); SSE mode re-stringifies the + // delta object itself, so the secret must stay raw. + this.scanner = new TokenScanner(opts.resolve ?? storeResolve, !this.sse); + } + + override _transform(chunk: Buffer, _enc: BufferEncoding, cb: TransformCallback): void { + const text = this.decoder.write(chunk); + cb(null, this.sse ? this.feedSSE(text) : this.scanner.push(text)); } override _flush(cb: TransformCallback): void { - const out = this.buf + this.decoder.end(); - this.buf = ""; - cb(null, out); + const tail = this.decoder.end(); + if (this.sse) { + let out = this.feedSSE(tail); + out += this.flushPending(); // emit any held token tail as a final delta + out += this.sseBuf; // leftover incomplete event, if any + this.sseBuf = ""; + cb(null, out); + } else { + cb(null, this.scanner.push(tail) + this.scanner.flush()); + } + } + + private feedSSE(text: string): string { + this.sseBuf += text; + let out = ""; + let idx: number; + while ((idx = this.sseBuf.indexOf("\n\n")) !== -1) { + const ev = this.sseBuf.slice(0, idx); + this.sseBuf = this.sseBuf.slice(idx + 2); + out += this.handleEvent(ev) + "\n\n"; + } + return out; + } + + private handleEvent(ev: string): string { + const lines = ev.split("\n"); + const di = lines.findIndex((l) => l.startsWith("data:")); + if (di === -1) return ev; // no data line: passthrough + let obj: Record; + try { + obj = JSON.parse(lines[di].slice(lines[di].indexOf(":") + 1).trim()); + } catch { + return ev; // not JSON (e.g. `data: [DONE]`): passthrough untouched + } + + const delta = obj.delta as { type?: string; text?: string } | undefined; + if (obj.type === "content_block_delta" && delta?.type === "text_delta" && typeof delta.text === "string") { + if (typeof obj.index === "number") this.currentIndex = obj.index; + delta.text = this.scanner.push(delta.text); // may hold a partial token across deltas + lines[di] = "data: " + JSON.stringify(obj); + return lines.join("\n"); + } + + // A token never spans a content block, so flush any held tail before the + // block closes. Pings / message_* / content_block_start must NOT flush — + // they can legitimately arrive mid-token. + if (obj.type === "content_block_stop") { + return this.flushPending() + ev; + } + return ev; + } + + // Emit the scanner's held-back tail (an incomplete token) as its own + // synthesized text_delta on the current block, so it stays valid SSE. + private flushPending(): string { + const tail = this.scanner.flush(); + if (!tail) return ""; + const evt = { + type: "content_block_delta", + index: this.currentIndex, + delta: { type: "text_delta", text: tail }, + }; + return "event: content_block_delta\ndata: " + JSON.stringify(evt) + "\n\n"; } } @@ -93,9 +192,9 @@ function runSelfTest(): void { ]); const resolve = (t: string) => map.get(t); - // Feed input split into arbitrary chunks, collect output. - const run = (chunks: string[]): string => { - const s = new DetokenizeStream(resolve); + // ---- raw mode (plain JSON responses) ---------------------------------- + const runRaw = (chunks: string[]): string => { + const s = new DetokenizeStream({ resolve }); let out = ""; s.on("data", (d) => (out += d.toString())); for (const c of chunks) s.write(Buffer.from(c, "utf8")); @@ -103,36 +202,82 @@ function runSelfTest(): void { return out; }; - // 1. Token split across a chunk boundary (the critical case). assert( - run(["password is <<>> done"]) === - "password is SUPER-SECRET-KEY done", - "split token reassembled", + runRaw(["password is <<>> done"]) === "password is SUPER-SECRET-KEY done", + "raw: split token reassembled", + ); + assert(runRaw("x <<>> y".split("")) === "x hunter2 y", "raw: char-by-char token"); + assert(runRaw(["a < b <<< c <<>>!"]) === "see <<>>!", "raw: unknown token kept"); + assert( + runRaw(["<<>><<>>"]) === "SUPER-SECRET-KEYhunter2", + "raw: adjacent tokens", + ); + assert(runRaw(["tail <<>>"]) === " + `event: content_block_delta\ndata: ${JSON.stringify({ type: "content_block_delta", index, delta: { type: "text_delta", text } })}\n\n`; + const stop = (index = 0) => + `event: content_block_stop\ndata: ${JSON.stringify({ type: "content_block_stop", index })}\n\n`; + const ping = () => `event: ping\ndata: ${JSON.stringify({ type: "ping" })}\n\n`; + + const runSSE = (chunks: string[]): string => { + const s = new DetokenizeStream({ sse: true, resolve }); + let out = ""; + s.on("data", (d) => (out += d.toString())); + for (const c of chunks) s.write(Buffer.from(c, "utf8")); + s.end(); + return out; + }; + // Concatenate the text the client would reconstruct from the output deltas. + const clientText = (sse: string): string => { + let text = ""; + for (const block of sse.split("\n\n")) { + const line = block.split("\n").find((l) => l.startsWith("data:")); + if (!line) continue; + try { + const o = JSON.parse(line.slice(line.indexOf(":") + 1).trim()); + if (o?.delta?.type === "text_delta") text += o.delta.text; + } catch { + /* ignore non-JSON */ + } + } + return text; + }; + + // 1. The critical case: a token fragmented across many text deltas (8 chars + // each), as the model actually streams a long placeholder. + const tok = "<<>>"; + const frag = []; + for (let i = 0; i < tok.length; i += 4) frag.push(delta(tok.slice(i, i + 4))); + assert( + clientText(runSSE(["before "].map((t) => delta(t)).concat(frag, [delta(" after"), stop()])).valueOf()) === + "before SUPER-SECRET-KEY after", + "sse: token fragmented across deltas reassembled", ); - // 2. Token split one char at a time. + // 2. A ping arriving mid-token must NOT break reassembly. assert( - run("x <<>> y".split("")) === "x hunter2 y", - "char-by-char token", + clientText(runSSE([delta("k="), delta("<<>>"), stop()])) === "k=hunter2", + "sse: ping mid-token tolerated", ); - // 3. Plain text with stray '<' and '<<<' that is NOT a token flushes promptly. - assert(run(["a < b <<< c <<>>!"]) === "see <<>>!", "unknown token kept"); - - // 5. Two tokens back to back. + // 3. Incomplete token at content_block_stop is flushed verbatim. assert( - run(["<<>><<>>"]) === "SUPER-SECRET-KEYhunter2", - "adjacent tokens", + clientText(runSSE([delta("oops <<>>!"), stop()])) === "hi hunter2!", + "sse: contiguous token in one delta", + ); - // 7. Overlapping '<' runs. - assert(run(["<<<>>"]) === " { (upRes) => { res.writeHead(upRes.statusCode ?? 502, pick(upRes.headers, STRIP_RES)); if (isTarget) { - // Detokenize the response body (SSE or JSON) on the way back. - upRes.pipe(new DetokenizeStream()).pipe(res); + // Detokenize the response on the way back. SSE streams fragment a + // placeholder across many deltas, so the detokenizer needs to parse + // events; plain JSON bodies carry the token contiguously. + const sse = String(upRes.headers["content-type"] ?? "").includes("text/event-stream"); + upRes.pipe(new DetokenizeStream({ sse })).pipe(res); } else { upRes.pipe(res); } diff --git a/src/store.ts b/src/store.ts index cb56cb2..5b8fc4a 100644 --- a/src/store.ts +++ b/src/store.ts @@ -1,12 +1,73 @@ // Copyright (C) 2026 Rootiest // SPDX-License-Identifier: AGPL-3.0-or-later -// In-memory token <-> secret map. ponytail: a Map is enough; secrets live only -// for the lifetime of the proxy process. Swap for SQLite only if you need -// persistence across restarts (you almost certainly don't — secrets shouldn't -// outlive the session). -export const secrets = new Map(); +import { createDecipheriv, randomBytes } from "node:crypto"; +import { mkdirSync, readFileSync, writeFileSync } from "node:fs"; +import { homedir } from "node:os"; +import { dirname, join } from "node:path"; -export function resolveToken(token: string): string | undefined { - return secrets.get(token); +// Cryptographic self-contained tokens: the placeholder carries its own +// ciphertext, so the proxy is stateless. The only persistent state is the +// master key, kept in a 0600 file so tokens survive restarts. +const KEY_PATH = + process.env.GG_MASTER_KEY_FILE || + join(homedir(), ".config", "claude-tokenization-proxy", "master.key"); + +// AES-256-GCM layout inside the hex blob: 12-byte IV, 16-byte auth tag, then +// the ciphertext. detokenize slices on these fixed boundaries. +export const IV_BYTES = 12; +export const TAG_BYTES = 16; +const IV_HEX = IV_BYTES * 2; +const TAG_HEX = TAG_BYTES * 2; + +let masterKey: Buffer | undefined; + +/** + * Load the 256-bit master key from disk, generating and persisting one (0600) + * on first run. Call once at startup before serving traffic. Idempotent. + */ +export function initMasterKey(): Buffer { + if (masterKey) return masterKey; + try { + const key = readFileSync(KEY_PATH); + if (key.length !== 32) { + throw new Error(`master key at ${KEY_PATH} is ${key.length} bytes, expected 32`); + } + masterKey = key; + } catch (err) { + if ((err as NodeJS.ErrnoException).code !== "ENOENT") throw err; + const key = randomBytes(32); + mkdirSync(dirname(KEY_PATH), { recursive: true, mode: 0o700 }); + writeFileSync(KEY_PATH, key, { mode: 0o600 }); + console.warn(`[proxy] generated new master key at ${KEY_PATH}`); + masterKey = key; + } + return masterKey; +} + +export function getMasterKey(): Buffer { + if (!masterKey) throw new Error("master key not initialized — call initMasterKey() first"); + return masterKey; +} + +/** + * Decrypt a full placeholder token back to its secret. Returns undefined for + * anything that isn't one of our tokens or that fails authentication, so + * unknown/corrupt placeholders pass through the stream untouched. + */ +export function resolveToken(token: string): string | undefined { + const m = /^<<>>$/.exec(token); + if (!m) return undefined; + const blob = m[1]; + if (blob.length <= IV_HEX + TAG_HEX) return undefined; + try { + const iv = Buffer.from(blob.slice(0, IV_HEX), "hex"); + const tag = Buffer.from(blob.slice(IV_HEX, IV_HEX + TAG_HEX), "hex"); + const ciphertext = Buffer.from(blob.slice(IV_HEX + TAG_HEX), "hex"); + const decipher = createDecipheriv("aes-256-gcm", getMasterKey(), iv); + decipher.setAuthTag(tag); + return Buffer.concat([decipher.update(ciphertext), decipher.final()]).toString("utf8"); + } catch { + return undefined; + } } diff --git a/src/tokenize.ts b/src/tokenize.ts index be9cee9..957b34f 100644 --- a/src/tokenize.ts +++ b/src/tokenize.ts @@ -2,16 +2,24 @@ // SPDX-License-Identifier: AGPL-3.0-or-later import { spawn } from "node:child_process"; -import { randomBytes } from "node:crypto"; +import { createCipheriv, randomBytes } from "node:crypto"; import { mkdtemp, writeFile, rm } from "node:fs/promises"; import { tmpdir } from "node:os"; import { join } from "node:path"; -import { secrets } from "./store.ts"; +import { getMasterKey, IV_BYTES } from "./store.ts"; let ggshieldWarned = false; -function newToken(): string { - return `<<>>`; +// Encrypt a secret into a self-contained placeholder. The payload is a single +// hex blob (IV ++ auth tag ++ ciphertext) so it matches the stream parser's +// [0-9a-f]+ token boundary and needs no server-side state to resolve later. +function encryptToken(secret: string): string { + const iv = randomBytes(IV_BYTES); + const cipher = createCipheriv("aes-256-gcm", getMasterKey(), iv); + const ciphertext = Buffer.concat([cipher.update(secret, "utf8"), cipher.final()]); + const tag = cipher.getAuthTag(); + const blob = `${iv.toString("hex")}${tag.toString("hex")}${ciphertext.toString("hex")}`; + return `<<>>`; } // Run ggshield over `text` and return the list of plaintext secret strings it @@ -64,9 +72,9 @@ async function scanSecrets(text: string): Promise { } /** - * Scan a raw request body for secrets, replace each with a placeholder token, - * record the token->secret mapping, and return the redacted body. Longest - * secrets first so a secret that contains another is replaced whole. + * Scan a raw request body for secrets and replace each with an encrypted + * self-contained placeholder. Longest secrets first so a secret that contains + * another is replaced whole. Stateless: the token carries its own ciphertext. */ export async function tokenizeBody(body: string): Promise { const found = await scanSecrets(body); @@ -75,9 +83,7 @@ export async function tokenizeBody(body: string): Promise { let redacted = body; for (const secret of found.sort((a, b) => b.length - a.length)) { if (!redacted.includes(secret)) continue; - const token = newToken(); - secrets.set(token, secret); - redacted = redacted.split(secret).join(token); + redacted = redacted.split(secret).join(encryptToken(secret)); } return redacted; }