diff --git a/cli/__tests__/provider-material-redaction.test.mjs b/cli/__tests__/provider-material-redaction.test.mjs index 1508ef30a6..833533dc2f 100644 --- a/cli/__tests__/provider-material-redaction.test.mjs +++ b/cli/__tests__/provider-material-redaction.test.mjs @@ -12,6 +12,130 @@ test.each([1, 2, 7, 50])('provider material never reaches model output with chun expect(Buffer.concat(output).toString()).toBe('before [redacted] between [redacted] after 中文'); }); +const collect = async (redactor, input, size) => { + const output = []; + redactor.on('data', (chunk) => output.push(chunk)); + const finished = new Promise((resolve) => redactor.on('end', resolve)); + for (let index = 0; index < input.length; index += size) redactor.write(input.subarray(index, index + size)); + redactor.end(); await finished; + return Buffer.concat(output).toString(); +}; + +// The redactor as it was before it released early: hold the last +// (key length - 1) bytes, always. The current one must match it byte for byte. +const { Transform } = await import('node:stream'); +const holdKeyLength = (material) => { + const secret = Buffer.from(material); + let pending = Buffer.alloc(0); + return new Transform({ + transform(chunk, _encoding, done) { + pending = Buffer.concat([pending, Buffer.from(chunk)]); + while (pending.length >= secret.length) { + const index = pending.indexOf(secret); + if (index >= 0) { + this.push(pending.subarray(0, index)); this.push('[redacted]'); + pending = pending.subarray(index + secret.length); + } else { + const count = pending.length - secret.length + 1; + this.push(pending.subarray(0, count)); pending = pending.subarray(count); + } + } + done(); + }, + flush(done) { this.push(pending); pending = Buffer.alloc(0); done(); }, + }); +}; + +// A seeded generator, so a failing stream can be replayed. +const seeded = (seed) => () => { + seed = (seed + 0x6d2b79f5) | 0; + let t = Math.imul(seed ^ (seed >>> 15), 1 | seed); + t = (t + Math.imul(t ^ (t >>> 7), 61 | t)) ^ t; + return ((t ^ (t >>> 14)) >>> 0) / 4294967296; +}; + +test.each(['', Buffer.alloc(0)])('refuses empty material instead of scanning forever (%p)', (material) => { + expect(() => redactProviderMaterial(material)).toThrow(TypeError); +}); + +test('releases bytes that cannot start the key at once, and holds only a possible key prefix', async () => { + const material = 'sk-PRIVATE-provider-key'; + const output = []; + const stream = redactProviderMaterial(material); + stream.on('data', (chunk) => output.push(chunk.toString())); + stream.write('data: {"token":"hi"}\n\n'); + await new Promise((resolve) => setImmediate(resolve)); + expect(output.join('')).toBe('data: {"token":"hi"}\n\n'); + stream.write('then sk-PRI'); + await new Promise((resolve) => setImmediate(resolve)); + expect(output.join('')).toBe('data: {"token":"hi"}\n\nthen '); + stream.write('VATE-provider-key done'); + stream.end(); + await new Promise((resolve) => stream.on('end', resolve)); + expect(output.join('')).toBe('data: {"token":"hi"}\n\nthen [redacted] done'); +}); + +test('redacts a key split at every byte boundary', async () => { + const material = 'sk-PRIVATE-provider-key'; + for (let cut = 1; cut < material.length; cut += 1) { + const stream = redactProviderMaterial(material); + const output = []; + stream.on('data', (chunk) => output.push(chunk)); + const finished = new Promise((resolve) => stream.on('end', resolve)); + stream.write(`a ${material.slice(0, cut)}`); stream.end(`${material.slice(cut)} b`); + await finished; // eslint-disable-line no-await-in-loop + expect(Buffer.concat(output).toString()).toBe('a [redacted] b'); + } +}); + +describe('matches the hold-key-length redactor byte for byte', () => { + const KEY = `sk-ant-api03-${'ab'.repeat(46)}`; + const cases = [`x${KEY}y`, `${KEY}${KEY}`, `sk-ant-${KEY}`, `${KEY.slice(0, -1)}!`, + `aaa${KEY.slice(0, 5)}`, 'plain text no key', KEY]; + test.each(cases.flatMap((input) => [1, 2, 7, 50, 1000].map((size) => [input.slice(0, 24), size, input])))( + '%s at chunk size %i', + async (_label, size, input) => { + const got = await collect(redactProviderMaterial(KEY), Buffer.from(input), size); + expect(got).toBe(await collect(holdKeyLength(KEY), Buffer.from(input), size)); + expect(got).not.toContain(KEY); + }, + ); + + test('on 2,000 random streams of key fragments', async () => { + const random = seeded(2753); + for (let n = 0; n < 2000; n += 1) { + const parts = []; + for (let k = 0; k < 6; k += 1) { + parts.push(random() < 0.3 ? KEY.slice(0, Math.floor(random() * KEY.length) + 1) + : random().toString(36).slice(2)); + } + const input = Buffer.from(parts.join('')); + const size = 1 + Math.floor(random() * 60); + /* eslint-disable no-await-in-loop */ + expect([n, await collect(redactProviderMaterial(KEY), input, size)]) + .toEqual([n, await collect(holdKeyLength(KEY), input, size)]); + /* eslint-enable no-await-in-loop */ + } + }); + + // Keys over a two-letter alphabet overlap themselves, the hard case for + // deciding which held bytes could still start a match. + test('on 2,000 random streams with self-overlapping keys', async () => { + const random = seeded(4000); + const word = (length) => Array.from({ length }, () => (random() < 0.5 ? 'a' : 'b')).join(''); + for (let n = 0; n < 2000; n += 1) { + const key = word(1 + Math.floor(random() * 12)); + const input = Buffer.from(word(Math.floor(random() * 80))); + const size = 1 + Math.floor(random() * 9); + /* eslint-disable no-await-in-loop */ + const got = await collect(redactProviderMaterial(key), input, size); + expect([n, key, got]).toEqual([n, key, await collect(holdKeyLength(key), input, size)]); + expect(got.replace(/\[redacted\]/g, '|')).not.toContain(key); + /* eslint-enable no-await-in-loop */ + } + }); +}); + test('the live parent proxy resolves memory-only material and closes at spawn exit', async () => { const { createServer } = await import('node:http'); const { mkdtemp, readFile, rm } = await import('node:fs/promises'); diff --git a/cli/src/lib/provider-proxy.js b/cli/src/lib/provider-proxy.js index 15333e5371..e1ce311e73 100644 --- a/cli/src/lib/provider-proxy.js +++ b/cli/src/lib/provider-proxy.js @@ -11,6 +11,8 @@ import { isAbsolute, relative, sep } from 'path'; // material across arbitrary stream chunk boundaries before any model sees it. export const redactProviderMaterial = (material) => { const secret = Buffer.from(material); + // An empty key matches everywhere and the scan below would never advance. + if (secret.length === 0) throw new TypeError('provider material to redact is empty'); let pending = Buffer.alloc(0); return new Transform({ transform(chunk, _encoding, done) { @@ -25,6 +27,13 @@ export const redactProviderMaterial = (material) => { this.push(pending.subarray(0, count)); pending = pending.subarray(count); } } + // Hold back only the longest tail that could still start the key, so a + // streamed token is not delayed until a key's length of bytes follows it. + let keep = Math.min(pending.length, secret.length - 1); + while (keep > 0 && !pending.subarray(pending.length - keep).equals(secret.subarray(0, keep))) keep -= 1; + if (pending.length > keep) { + this.push(pending.subarray(0, pending.length - keep)); pending = pending.subarray(pending.length - keep); + } done(); }, flush(done) { this.push(pending); pending = Buffer.alloc(0); done(); },