From 650401d3580cc2cbc0e67f464ad55131c20d3818 Mon Sep 17 00:00:00 2001 From: kartikshukla17 Date: Thu, 30 Jul 2026 02:00:41 +0530 Subject: [PATCH] fix: bound knowledge-base ingestion --- sdk/typescript/README.md | 4 + sdk/typescript/src/knowledge-base.ts | 283 +++++++++++++++--- .../tests-ts/knowledge-base.test.ts | 106 ++++++- 3 files changed, 348 insertions(+), 45 deletions(-) diff --git a/sdk/typescript/README.md b/sdk/typescript/README.md index 8d069014..bfb970c3 100644 --- a/sdk/typescript/README.md +++ b/sdk/typescript/README.md @@ -175,6 +175,10 @@ to Repeat `--knowledge-base PATH` for multiple files or directories. Directories are searched recursively for Markdown, text, PDF, and Word (`.docx`) files. +Knowledge bases are limited to 128 documents and 4,096 discovered entries across +16 directory levels. Each input and extracted document is limited to 8 MiB, +aggregate input and extracted text are each limited to 32 MiB, and PDFs are +limited to 512 pages. On macOS/Linux, an existing output directory must be private to the current user (`chmod 700`). diff --git a/sdk/typescript/src/knowledge-base.ts b/sdk/typescript/src/knowledge-base.ts index ad82a3c5..02712a5b 100644 --- a/sdk/typescript/src/knowledge-base.ts +++ b/sdk/typescript/src/knowledge-base.ts @@ -2,8 +2,8 @@ import { constants } from "node:fs"; import { lstat, mkdtemp, - readFile, - readdir, + open, + opendir, realpath, rm, writeFile, @@ -19,6 +19,23 @@ const SUPPORTED_EXTENSIONS = new Set([ ".pdf", ".docx", ]); +const MAX_DOCUMENTS = 128; +const MAX_DIRECTORY_DEPTH = 16; +const MAX_DISCOVERY_ENTRIES = 4_096; +const MAX_DOCUMENT_BYTES = 8 * 1024 * 1024; +const MAX_INPUT_BYTES = 32 * 1024 * 1024; +const MAX_EXTRACTED_DOCUMENT_BYTES = 8 * 1024 * 1024; +const MAX_EXTRACTED_BYTES = 32 * 1024 * 1024; +const MAX_PDF_PAGES = 512; +const READ_CHUNK_BYTES = 64 * 1024; + +interface DiscoveryState { + documents: Set; + entries: number; + inputBytes: number; +} + +class KnowledgeBaseLimitError extends Error {} export interface PreparedKnowledgeBase { path: string; @@ -31,7 +48,11 @@ export async function prepareKnowledgeBase( signal?: AbortSignal, ): Promise { const sources = new Set(); - const documents = new Set(); + const discovery: DiscoveryState = { + documents: new Set(), + entries: 0, + inputBytes: 0, + }; for (const requested of paths) { signal?.throwIfAborted(); @@ -49,46 +70,61 @@ export async function prepareKnowledgeBase( } const source = await realpath(path); - const selected = metadata.isDirectory() ? await discover(source) : [source]; - if (selected.length === 0) { + signal?.throwIfAborted(); + if (sources.has(source)) continue; + let selected = false; + if (metadata.isDirectory()) { + selected = await discover(source, 0, discovery, signal); + } else { + if (!SUPPORTED_EXTENSIONS.has(extname(source).toLowerCase())) { + throw new Error(`Unsupported knowledge base document: ${source}`); + } + await addDocument(source, await lstat(source), discovery, signal); + selected = true; + } + if (!selected) { throw new Error( `Knowledge base directory contains no supported documents: ${path}`, ); } - for (const document of selected) { - if (!SUPPORTED_EXTENSIONS.has(extname(document).toLowerCase())) { - throw new Error(`Unsupported knowledge base document: ${document}`); - } - documents.add(document); - } sources.add(source); } + signal?.throwIfAborted(); const path = await mkdtemp(join(tmpdir(), "codex-security-knowledge-")); try { let index = 0; - for (const document of documents) { + let inputBytes = 0; + let extractedBytes = 0; + for (const document of discovery.documents) { signal?.throwIfAborted(); - const metadata = await lstat(document); - if (process.platform !== "win32" && (metadata.mode & 0o444) === 0) { - throw new Error(`Knowledge base document is not readable: ${document}`); - } - const bytes = await readFile(document, { - flag: constants.O_RDONLY | (constants.O_NOFOLLOW ?? 0), - signal, - }); + const bytes = await readDocument(document, inputBytes, signal); + inputBytes += bytes.byteLength; const extension = extname(document).toLowerCase(); const text = extension === ".pdf" - ? await extractPdf(document, bytes) + ? await extractPdf(document, bytes, signal) : extension === ".docx" - ? extractDocx(document, bytes) + ? extractDocx(document, bytes, signal) : decodeText(document, bytes); + signal?.throwIfAborted(); if ((extension === ".pdf" || extension === ".docx") && !text.trim()) { throw new Error( `Knowledge base document contains no extractable text: ${document}`, ); } + const textBytes = Buffer.byteLength(text, "utf8"); + if (textBytes > MAX_EXTRACTED_DOCUMENT_BYTES) { + throw new KnowledgeBaseLimitError( + `Knowledge base document exceeds the ${MAX_EXTRACTED_DOCUMENT_BYTES}-byte extracted-text limit: ${document}`, + ); + } + if (extractedBytes + textBytes > MAX_EXTRACTED_BYTES) { + throw new KnowledgeBaseLimitError( + `Knowledge base extracted text exceeds the ${MAX_EXTRACTED_BYTES}-byte aggregate limit.`, + ); + } + extractedBytes += textBytes; await writeFile( join(path, `${index++}-${basename(document)}.txt`), text, @@ -111,22 +147,135 @@ export async function prepareKnowledgeBase( }; } -async function discover(directory: string): Promise { - const documents: string[] = []; - const entries = await readdir(directory, { withFileTypes: true }); - for (const entry of entries) { +async function discover( + directory: string, + depth: number, + state: DiscoveryState, + signal?: AbortSignal, +): Promise { + signal?.throwIfAborted(); + if (depth > MAX_DIRECTORY_DEPTH) { + throw new KnowledgeBaseLimitError( + `Knowledge base directory exceeds the ${MAX_DIRECTORY_DEPTH}-level nesting limit: ${directory}`, + ); + } + let selected = false; + const entries = await opendir(directory); + for await (const entry of entries) { + signal?.throwIfAborted(); + state.entries += 1; + if (state.entries > MAX_DISCOVERY_ENTRIES) { + throw new KnowledgeBaseLimitError( + `Knowledge base discovery exceeds the ${MAX_DISCOVERY_ENTRIES}-entry limit.`, + ); + } const path = join(directory, entry.name); if (entry.isSymbolicLink()) continue; if (entry.isDirectory()) { - documents.push(...(await discover(path))); + if (await discover(path, depth + 1, state, signal)) selected = true; } else if ( entry.isFile() && SUPPORTED_EXTENSIONS.has(extname(path).toLowerCase()) ) { - documents.push(path); + const metadata = await lstat(path); + signal?.throwIfAborted(); + if (metadata.isSymbolicLink() || !metadata.isFile()) continue; + selected = true; + await addDocument(path, metadata, state, signal); } } - return documents; + signal?.throwIfAborted(); + return selected; +} + +async function addDocument( + path: string, + metadata: { size: number }, + state: DiscoveryState, + signal?: AbortSignal, +): Promise { + signal?.throwIfAborted(); + if (state.documents.has(path)) return; + if (state.documents.size >= MAX_DOCUMENTS) { + throw new KnowledgeBaseLimitError( + `Knowledge base contains more than ${MAX_DOCUMENTS} documents.`, + ); + } + if (metadata.size > MAX_DOCUMENT_BYTES) { + throw new KnowledgeBaseLimitError( + `Knowledge base document exceeds the ${MAX_DOCUMENT_BYTES}-byte input limit: ${path}`, + ); + } + if (state.inputBytes + metadata.size > MAX_INPUT_BYTES) { + throw new KnowledgeBaseLimitError( + `Knowledge base input exceeds the ${MAX_INPUT_BYTES}-byte aggregate limit.`, + ); + } + state.documents.add(path); + state.inputBytes += metadata.size; +} + +async function readDocument( + path: string, + consumedBytes: number, + signal?: AbortSignal, +): Promise { + signal?.throwIfAborted(); + const file = await open( + path, + constants.O_RDONLY | (constants.O_NOFOLLOW ?? 0), + ); + try { + const metadata = await file.stat(); + signal?.throwIfAborted(); + if (!metadata.isFile()) { + throw new Error(`Knowledge base document is not a file: ${path}`); + } + if (process.platform !== "win32" && (metadata.mode & 0o444) === 0) { + throw new Error(`Knowledge base document is not readable: ${path}`); + } + if (metadata.size > MAX_DOCUMENT_BYTES) { + throw new KnowledgeBaseLimitError( + `Knowledge base document exceeds the ${MAX_DOCUMENT_BYTES}-byte input limit: ${path}`, + ); + } + if (consumedBytes + metadata.size > MAX_INPUT_BYTES) { + throw new KnowledgeBaseLimitError( + `Knowledge base input exceeds the ${MAX_INPUT_BYTES}-byte aggregate limit.`, + ); + } + + const maximum = Math.min( + MAX_DOCUMENT_BYTES, + MAX_INPUT_BYTES - consumedBytes, + ); + const chunks: Buffer[] = []; + let length = 0; + while (length <= maximum) { + signal?.throwIfAborted(); + const chunk = Buffer.allocUnsafe( + Math.min(READ_CHUNK_BYTES, maximum + 1 - length), + ); + const { bytesRead } = await file.read(chunk, 0, chunk.byteLength, null); + signal?.throwIfAborted(); + if (bytesRead === 0) break; + chunks.push(chunk.subarray(0, bytesRead)); + length += bytesRead; + } + if (length > MAX_DOCUMENT_BYTES) { + throw new KnowledgeBaseLimitError( + `Knowledge base document exceeds the ${MAX_DOCUMENT_BYTES}-byte input limit: ${path}`, + ); + } + if (consumedBytes + length > MAX_INPUT_BYTES) { + throw new KnowledgeBaseLimitError( + `Knowledge base input exceeds the ${MAX_INPUT_BYTES}-byte aggregate limit.`, + ); + } + return Buffer.concat(chunks, length); + } finally { + await file.close(); + } } function decodeText(path: string, bytes: Uint8Array): string { @@ -139,49 +288,95 @@ function decodeText(path: string, bytes: Uint8Array): string { } } -async function extractPdf(path: string, bytes: Uint8Array): Promise { +async function extractPdf( + path: string, + bytes: Uint8Array, + signal?: AbortSignal, +): Promise { + signal?.throwIfAborted(); try { const { getDocument, VerbosityLevel } = await import( "pdfjs-dist/legacy/build/pdf.mjs" ); - const document = await getDocument({ + const loadingTask = getDocument({ data: new Uint8Array(bytes), isEvalSupported: false, stopAtErrors: true, verbosity: VerbosityLevel.ERRORS, - }).promise; + }); + let document: Awaited | undefined; + let destroying: Promise | null = null; + const destroy = (): Promise => + (destroying ??= + document === undefined ? loadingTask.destroy() : document.destroy()); + const onAbort = (): void => { + void destroy().catch(() => {}); + }; + signal?.addEventListener("abort", onAbort, { once: true }); + if (signal?.aborted === true) onAbort(); try { + document = await loadingTask.promise; + signal?.throwIfAborted(); + if (document.numPages > MAX_PDF_PAGES) { + throw new KnowledgeBaseLimitError( + `Knowledge base PDF exceeds the ${MAX_PDF_PAGES}-page limit: ${path}`, + ); + } const pages: string[] = []; + let extractedBytes = 0; for (let number = 1; number <= document.numPages; number++) { + signal?.throwIfAborted(); const content = await (await document.getPage(number)).getTextContent(); - pages.push( - content.items - .map((item) => ("str" in item ? item.str : "")) - .join(" "), - ); + signal?.throwIfAborted(); + const page = content.items + .map((item) => ("str" in item ? item.str : "")) + .join(" "); + const pageBytes = + Buffer.byteLength(page, "utf8") + (pages.length === 0 ? 0 : 1); + if (extractedBytes + pageBytes > MAX_EXTRACTED_DOCUMENT_BYTES) { + throw new KnowledgeBaseLimitError( + `Knowledge base document exceeds the ${MAX_EXTRACTED_DOCUMENT_BYTES}-byte extracted-text limit: ${path}`, + ); + } + pages.push(page); + extractedBytes += pageBytes; } return pages.join("\n"); } finally { - await document.destroy(); + signal?.removeEventListener("abort", onAbort); + await destroy().catch((error: unknown) => { + signal?.throwIfAborted(); + throw error; + }); } } catch (error) { + signal?.throwIfAborted(); + if (error instanceof KnowledgeBaseLimitError) throw error; throw new Error(`Cannot extract text from knowledge base PDF: ${path}`, { cause: error, }); } } -function extractDocx(path: string, bytes: Uint8Array): string { +function extractDocx( + path: string, + bytes: Uint8Array, + signal?: AbortSignal, +): string { + signal?.throwIfAborted(); try { const files = unzipSync(bytes, { filter: (file) => { if (file.name !== "word/document.xml") return false; - if (file.originalSize > 25 * 1024 * 1024) { - throw new Error("DOCX document text exceeds 25 MB."); + if (file.originalSize > MAX_EXTRACTED_DOCUMENT_BYTES) { + throw new KnowledgeBaseLimitError( + `Knowledge base document exceeds the ${MAX_EXTRACTED_DOCUMENT_BYTES}-byte extracted-text limit: ${path}`, + ); } return true; }, }); + signal?.throwIfAborted(); const document = files["word/document.xml"]; if (document === undefined) throw new Error("Missing word/document.xml."); const xml = decodeText(path, document); @@ -190,13 +385,17 @@ function extractDocx(path: string, bytes: Uint8Array): string { ) { throw new Error("Malformed word/document.xml."); } - return decodeXml( + const text = decodeXml( xml .replace(/<\/(?:\w+:)?p\s*>/gu, "\n") .replace(/<(?:\w+:)?tab\b[^>]*\/>/gu, "\t") .replace(/<[^>]+>/gu, ""), ); + signal?.throwIfAborted(); + return text; } catch (error) { + signal?.throwIfAborted(); + if (error instanceof KnowledgeBaseLimitError) throw error; throw new Error(`Cannot extract text from knowledge base DOCX: ${path}`, { cause: error, }); diff --git a/sdk/typescript/tests-ts/knowledge-base.test.ts b/sdk/typescript/tests-ts/knowledge-base.test.ts index 79ca83c4..9bd67020 100644 --- a/sdk/typescript/tests-ts/knowledge-base.test.ts +++ b/sdk/typescript/tests-ts/knowledge-base.test.ts @@ -8,6 +8,7 @@ import { rm, stat, symlink, + truncate, writeFile, } from "node:fs/promises"; import { tmpdir } from "node:os"; @@ -49,13 +50,19 @@ function docx(text: string): Uint8Array { }); } -function pdf(text: string): Uint8Array { +function pdf(text: string, pages = 1): Uint8Array { const escaped = text.replace(/[\\()]/gu, "\\$&"); const stream = `BT /F1 12 Tf 72 720 Td (${escaped}) Tj ET`; + const pageObjects = Array.from({ length: pages }, (_, index) => index + 3); + const font = pages + 3; + const content = pages + 4; const objects = [ "<< /Type /Catalog /Pages 2 0 R >>", - "<< /Type /Pages /Kids [3 0 R] /Count 1 >>", - "<< /Type /Page /Parent 2 0 R /MediaBox [0 0 612 792] /Resources << /Font << /F1 4 0 R >> >> /Contents 5 0 R >>", + `<< /Type /Pages /Kids [${pageObjects.map((number) => `${number} 0 R`).join(" ")}] /Count ${pages} >>`, + ...pageObjects.map( + () => + `<< /Type /Page /Parent 2 0 R /MediaBox [0 0 612 792] /Resources << /Font << /F1 ${font} 0 R >> >> /Contents ${content} 0 R >>`, + ), "<< /Type /Font /Subtype /Type1 /BaseFont /Helvetica >>", `<< /Length ${Buffer.byteLength(stream)} >>\nstream\n${stream}\nendstream`, ]; @@ -187,6 +194,99 @@ describe("scan knowledge bases", () => { ); }); + test("bounds document count, nesting depth, and individual input size", async () => { + const countRoot = await temporaryDirectory(); + const documents = Array.from({ length: 129 }, (_, index) => + join(countRoot, `${index}.md`), + ); + await Promise.all(documents.map((path) => writeFile(path, "scope"))); + await expect(prepareKnowledgeBase(documents)).rejects.toThrow( + "more than 128 documents", + ); + + const depthRoot = await temporaryDirectory(); + let nested = depthRoot; + for (let depth = 0; depth < 17; depth += 1) { + nested = join(nested, "nested"); + await mkdir(nested); + } + await writeFile(join(nested, "scope.md"), "scope"); + await expect(prepareKnowledgeBase([depthRoot])).rejects.toThrow( + "16-level nesting limit", + ); + + const sizeRoot = await temporaryDirectory(); + const oversized = join(sizeRoot, "oversized.md"); + await writeFile(oversized, ""); + await truncate(oversized, 8 * 1024 * 1024 + 1); + await expect(prepareKnowledgeBase([oversized])).rejects.toThrow( + "8388608-byte input limit", + ); + }); + + test("bounds aggregate input and extracted text", async () => { + const inputRoot = await temporaryDirectory(); + const inputs = Array.from({ length: 5 }, (_, index) => + join(inputRoot, `${index}.md`), + ); + for (const [index, path] of inputs.entries()) { + await writeFile(path, ""); + await truncate(path, index === inputs.length - 1 ? 1 : 8 * 1024 * 1024); + } + await expect(prepareKnowledgeBase(inputs)).rejects.toThrow( + "33554432-byte aggregate limit", + ); + + const documentOutputRoot = await temporaryDirectory(); + const oversizedOutput = join(documentOutputRoot, "oversized.docx"); + await writeFile(oversizedOutput, docx("x".repeat(8 * 1024 * 1024))); + await expect(prepareKnowledgeBase([oversizedOutput])).rejects.toThrow( + "8388608-byte extracted-text limit", + ); + + const outputRoot = await temporaryDirectory(); + const compressedText = "x".repeat(7 * 1024 * 1024); + for (let index = 0; index < 5; index += 1) { + await writeFile(join(outputRoot, `${index}.docx`), docx(compressedText)); + } + await expect(prepareKnowledgeBase([outputRoot])).rejects.toThrow( + "extracted text exceeds the 33554432-byte aggregate limit", + ); + }); + + test("limits PDF page extraction", async () => { + const root = await temporaryDirectory(); + const oversized = join(root, "oversized.pdf"); + await writeFile(oversized, pdf("scope", 513)); + + await expect(prepareKnowledgeBase([oversized])).rejects.toThrow( + "512-page limit", + ); + }); + + test("observes cancellation while discovering directories", async () => { + const root = await temporaryDirectory(); + const nested = join(root, "one", "two"); + await mkdir(nested, { recursive: true }); + await writeFile(join(nested, "scope.md"), "scope"); + const controller = new AbortController(); + const reason = new DOMException("cancel discovery", "AbortError"); + const throwIfAborted = controller.signal.throwIfAborted.bind( + controller.signal, + ); + let checks = 0; + controller.signal.throwIfAborted = (): void => { + checks += 1; + if (checks === 5) controller.abort(reason); + throwIfAborted(); + }; + + await expect(prepareKnowledgeBase([root], controller.signal)).rejects.toBe( + reason, + ); + expect(checks).toBe(5); + }); + testPosix("does not follow symbolic links", async () => { const root = await temporaryDirectory(); const source = join(root, "scope.md");