From d5d1b4d8b42a65c82976b080e2d307bc49def48b Mon Sep 17 00:00:00 2001 From: Daniel Salazar Date: Sun, 4 Oct 2026 05:18:31 -0700 Subject: [PATCH] fix(ai-ocr): bound the page count's parsing and close its under-counts - Object-stream headers are capped (/N, /First) and read incrementally, and the inflate budget scales with the file, so a crafted PDF can't spend seconds and gigabytes before the credit check. - An escaped /ObjStm type, an indirect /Count, untyped page leaves and an object redefined to a smaller tree no longer lower the count. - An account with no balance is refused before the file is parsed. --- src/backend/drivers/ai-ocr/OCRDriver.test.ts | 23 +++++- src/backend/drivers/ai-ocr/OCRDriver.ts | 6 ++ src/backend/drivers/ai-ocr/pdfPages.test.ts | 55 +++++++++++++- src/backend/drivers/ai-ocr/pdfPages.ts | 77 +++++++++++++++----- 4 files changed, 140 insertions(+), 21 deletions(-) diff --git a/src/backend/drivers/ai-ocr/OCRDriver.test.ts b/src/backend/drivers/ai-ocr/OCRDriver.test.ts index bf57ecb14..b071dd8fb 100644 --- a/src/backend/drivers/ai-ocr/OCRDriver.test.ts +++ b/src/backend/drivers/ai-ocr/OCRDriver.test.ts @@ -357,7 +357,7 @@ describe('OCRDriver.recognize (aws-textract)', () => { // fsEntry, and the driver's per-region TextractClient cache leaks // across tests. -it('meters one usage line per detected page at the per-page rate from costs.ts', async () => { + it('meters one usage line per detected page at the per-page rate from costs.ts', async () => { const { actor } = await makeUser(); textractSendMock.mockResolvedValueOnce(sampleTextractResponse); @@ -856,6 +856,27 @@ describe('OCRDriver credit pre-flight', () => { expect(await outstandingHolds(actor)).toBe(0); }); + it('rejects a zero-balance actor before estimating pages or holding credits', async () => { + const { actor } = await makeUser(); + const noUsage = vi + .spyOn(server.services.metering, 'hasAnyUsageCached') + .mockResolvedValueOnce(false); + + await expect( + withActor(actor, () => + driver.recognize({ + source: pdfSource(buildPdf(5)), + provider: 'mistral', + }), + ), + ).rejects.toMatchObject({ statusCode: 402 }); + + // The estimate/hold path (which calls hasEnoughCredits) never ran. + expect(hasCreditsSpy).not.toHaveBeenCalled(); + expect(mistralOcrProcessMock).not.toHaveBeenCalled(); + noUsage.mockRestore(); + }); + it('releases the hold when the provider fails', async () => { const { actor } = await makeUser(); mistralOcrProcessMock.mockRejectedValueOnce(new Error('upstream down')); diff --git a/src/backend/drivers/ai-ocr/OCRDriver.ts b/src/backend/drivers/ai-ocr/OCRDriver.ts index 920b9635f..58f53faea 100644 --- a/src/backend/drivers/ai-ocr/OCRDriver.ts +++ b/src/backend/drivers/ai-ocr/OCRDriver.ts @@ -28,6 +28,7 @@ import { Actor } from '../../core/actor.js'; import { Context } from '../../core/context.js'; import { HttpError } from '../../core/http/HttpError.js'; import { mimeFromName } from '../../util/fileSigning.js'; +import { assertActorHasCredits } from '../../services/metering/enforcement.js'; import type { CreditHold } from '../../services/metering/types.js'; import { PuterDriver } from '../types.js'; import { @@ -322,6 +323,11 @@ export class OCRDriver extends PuterDriver { legacyCode: 'internal_error', }); + // Cheap, cached check: reject a zero-balance actor before paying for + // the file load and page estimate, which a real credit check would + // reject anyway once it runs. + await assertActorHasCredits(this.services.metering, actor, this.config); + const loaded = await loadFileInput( this.stores, this.services.fs, diff --git a/src/backend/drivers/ai-ocr/pdfPages.test.ts b/src/backend/drivers/ai-ocr/pdfPages.test.ts index 26a3541c2..30d0ef46c 100644 --- a/src/backend/drivers/ai-ocr/pdfPages.test.ts +++ b/src/backend/drivers/ai-ocr/pdfPages.test.ts @@ -71,8 +71,9 @@ describe('countPdfPages', () => { ).toBe(9); }); - it('lets a later definition of an object replace an earlier one', () => { - // An incremental update that shrank the tree and dropped a page. + it("can't be shrunk by redefining an object number lower", () => { + // An unreferenced redefinition trying to pass off a smaller tree; + // a genuine object number reuse can only raise the count, never lower it. expect( countPdfPages( pdf( @@ -83,7 +84,55 @@ describe('countPdfPages', () => { '3 0 obj\n<< /Type /Annot >>\nendobj\n', ), ), - ).toBe(1); + ).toBe(3); + }); + + it('treats an indirect /Count as unknown rather than reading its reference as the total', () => { + expect( + countPdfPages( + pdf( + '1 0 obj\n<< /Type /Pages /Kids [2 0 R 3 0 R] /Count 3 0 R >>\nendobj\n', + ), + ), + ).toBeNull(); + }); + + it('counts a kid with no /Type as long as it has no /Kids of its own', () => { + expect( + countPdfPages( + pdf( + '1 0 obj\n<< /Type /Pages /Kids [2 0 R 3 0 R] /Count 1 >>\nendobj\n' + + '2 0 obj\n<< /Type /Page /Parent 1 0 R >>\nendobj\n' + + '3 0 obj\n<< /Parent 1 0 R /MediaBox [0 0 612 792] >>\nendobj\n', + ), + ), + ).toBe(2); + }); + + it('decodes a # escaped /ObjStm type before deciding whether to unpack it', () => { + const header = '2 0\n'; + const packedObject = '<< /Type /Page >>'; + const objStm = + `5 0 obj\n<< /Type /ObjS#74m /N 1 /First ${header.length} >>\n` + + `stream\n${header}${packedObject}\nendstream\nendobj\n`; + expect(countPdfPages(pdf(objStm))).toBe(1); + }); + + it("refuses an object stream that claims more objects than it's worth counting", () => { + const objStm = + '1 0 obj\n<< /Type /ObjStm /N 100001 /First 4 >>\nstream\n1 0\nendstream\nendobj\n'; + expect(countPdfPages(pdf(objStm))).toBeNull(); + }); + + it('refuses an oversized /First without scanning the data behind it', () => { + // Without a cap this would run a regex match over ~2MB of digits. + const data = '1 '.repeat(1_000_000); + const objStm = + `1 0 obj\n<< /Type /ObjStm /N 1 /First ${data.length} >>\n` + + `stream\n${data}\nendstream\nendobj\n`; + const start = Date.now(); + expect(countPdfPages(pdf(objStm))).toBeNull(); + expect(Date.now() - start).toBeLessThan(500); }); it('returns null for bytes that are not a PDF', () => { diff --git a/src/backend/drivers/ai-ocr/pdfPages.ts b/src/backend/drivers/ai-ocr/pdfPages.ts index ff6f659e5..a7276d3f6 100644 --- a/src/backend/drivers/ai-ocr/pdfPages.ts +++ b/src/backend/drivers/ai-ocr/pdfPages.ts @@ -19,14 +19,31 @@ import { inflateSync } from 'node:zlib'; -/** Total that object streams may inflate to before the count gives up. */ +/** Hard ceiling on inflated object-stream bytes, regardless of input size. */ const MAX_INFLATED_BYTES = 64 * 1024 * 1024; +/** + * Zlib can't compress past ~1032:1 — ties the worst case to what was actually + * uploaded. + */ +const MAX_INFLATE_RATIO = 1024; +/** + * An object stream claiming more objects or a bigger header than this can't be + * trusted cheaply. + */ +const MAX_OBJSTM_OBJECTS = 100_000; +const MAX_OBJSTM_HEADER_BYTES = 1024 * 1024; // A name ends at whitespace, a delimiter or the end of input. const TYPE_PAGE = /\/Type[\s\0]*\/Page(?![^\s\0/<>[\]()%{}])/; const TYPE_PAGES = /\/Type[\s\0]*\/Pages(?![^\s\0/<>[\]()%{}])/; const TYPE_OBJECT_STREAM = /\/Type[\s\0]*\/ObjStm(?![^\s\0/<>[\]()%{}])/; +const HAS_TYPE = /\/Type[\s\0]*\//; +const HAS_KIDS = /\/Kids[\s\0]*\[/; +const REF = (key: string) => + new RegExp(`\\/${key}[\\s\\0]+\\d{1,10}[\\s\\0]+\\d{1,10}[\\s\\0]+R\\b`); +const HAS_PARENT = REF('Parent'); const COUNT = /\/Count[\s\0]+(\d{1,10})(?![\d.])/; +const COUNT_INDIRECT = REF('Count'); const STREAM_START = />>[\s\0]*stream(?:\r\n|\r|\n)/; const FLATE_FILTER = /\/Filter[\s\0]*(?:\/(?:FlateDecode|Fl)|\[[\s\0]*\/(?:FlateDecode|Fl)[\s\0]*\])/; @@ -61,10 +78,24 @@ const readObjectStream = ( const count = Number(/\/N[\s\0]+(\d+)/.exec(dict)?.[1]); const first = Number(/\/First[\s\0]+(\d+)/.exec(dict)?.[1]); + if (!Number.isInteger(count) || count > MAX_OBJSTM_OBJECTS) return null; + if (!Number.isInteger(first) || first > MAX_OBJSTM_HEADER_BYTES) + return null; const text = content.toString('latin1'); - if (!Number.isInteger(count) || !(first <= text.length)) return null; - // The header is `count` pairs of object number and offset from `first`. - const header = (text.slice(0, first).match(/\d+/g) ?? []).map(Number); + if (first > text.length) return null; + + // The header is `count` pairs of object number and offset from `first`; + // read only that many integers instead of matching the whole prefix. + const header: number[] = []; + const digits = /\d+/g; + let match: RegExpExecArray | null; + while ( + header.length < count * 2 && + (match = digits.exec(text)) && + match.index < first + ) { + header.push(Number(match[0])); + } if (header.length < count * 2) return null; const objects: Array<[number, string]> = []; @@ -85,21 +116,37 @@ const readObjectStream = ( export function countPdfPages(pdf: Buffer, limit = Infinity): number | null { if (!pdf.subarray(0, 1024).includes('%PDF-')) return null; const text = pdf.toString('latin1'); + // Zlib bombs are a fixed multiple of their compressed size; scale the + // ceiling down for small files instead of always allowing the max. + const inflateBudget = Math.min( + MAX_INFLATED_BYTES, + pdf.length * MAX_INFLATE_RATIO, + ); - // Keyed by object number, so a later definition replaces an earlier one - // as it does in an incrementally updated file. + // Keyed by object number. A redefinition can only raise what's counted + // for that id, never lower it — an unreferenced redefinition crafted to + // shrink the tree can't undercount what an earlier one established. const treeCounts = new Map(); const pageObjects = new Set(); const visit = (id: number, rawDict: string): boolean => { const dict = decodeNameEscapes(rawDict); - treeCounts.delete(id); - pageObjects.delete(id); if (TYPE_PAGES.test(dict)) { - const count = Number(COUNT.exec(dict)?.[1] ?? 0); - treeCounts.set(id, count); + // An indirect /Count (`N 0 R`) isn't a count — it's a reference + // whose object number happens to look like one. + const count = COUNT_INDIRECT.test(dict) + ? 0 + : Number(COUNT.exec(dict)?.[1] ?? 0); + treeCounts.set(id, Math.max(count, treeCounts.get(id) ?? 0)); return count >= limit; } - if (TYPE_PAGE.test(dict)) pageObjects.add(id); + // A kid with no /Type is still a page to any real reader as long as + // it isn't an intermediate node (no /Kids); only trust that reading + // when /Type is missing entirely, not just unrecognized. + const isUntypedPage = + !HAS_TYPE.test(dict) && + HAS_PARENT.test(dict) && + !HAS_KIDS.test(dict); + if (TYPE_PAGE.test(dict) || isUntypedPage) pageObjects.add(id); return pageObjects.size >= limit; }; @@ -117,7 +164,7 @@ export function countPdfPages(pdf: Buffer, limit = Infinity): number | null { const body = text.slice(start, end); const stream = STREAM_START.exec(body); const dict = stream ? body.slice(0, stream.index + 2) : body; - if (!TYPE_OBJECT_STREAM.test(dict)) { + if (!TYPE_OBJECT_STREAM.test(decodeNameEscapes(dict))) { if (visit(Number(header[1]), dict)) return limit; continue; } @@ -128,11 +175,7 @@ export function countPdfPages(pdf: Buffer, limit = Infinity): number | null { start + stream.index + stream[0].length, start + dataEnd, ); - const packed = readObjectStream( - dict, - data, - MAX_INFLATED_BYTES - inflated, - ); + const packed = readObjectStream(dict, data, inflateBudget - inflated); if (!packed) return null; inflated += packed.inflated; for (const [id, objectDict] of packed.objects) {