mirror of
https://github.com/simstudioai/sim.git
synced 2026-09-24 15:45:35 +08:00
improvement(file-parsers): bound PDF text extraction (#6425)
* improvement(file-parsers): bound PDF text extraction Extract page text through pdf.js's streaming API with page, character, and wall-clock budgets instead of buffering the whole document, so extraction memory stays bounded regardless of input. Release the document proxy when done, and route output through sanitizeTextForUTF8 like the other parsers. * fix(file-parsers): only flag truncation when PDF text is actually dropped
This commit is contained in:
@@ -0,0 +1,84 @@
|
||||
/**
|
||||
* @vitest-environment node
|
||||
*/
|
||||
import { deflateSync } from 'zlib'
|
||||
import { describe, expect, it } from 'vitest'
|
||||
import { MAX_PDF_TEXT_CHARS, PdfParser } from '@/lib/file-parsers/pdf-parser'
|
||||
|
||||
/**
|
||||
* Builds a single-page PDF that draws 64 characters per repeat from a
|
||||
* FlateDecode content stream, so a few dozen kilobytes of input yields millions
|
||||
* of extracted characters — what made the unbounded extractor exhaust the heap
|
||||
* and abort the process.
|
||||
*
|
||||
* Hand-assembled rather than built with `pdf-lib` because the fixture's whole
|
||||
* point is the compression ratio of the content stream, which `pdf-lib` gives
|
||||
* no way to control.
|
||||
*/
|
||||
function buildTextBombPdf(repeats: number): Buffer {
|
||||
const unit = `BT /F1 12 Tf 10 700 Td (${'A'.repeat(64)}) Tj ET\n`
|
||||
const compressed = deflateSync(Buffer.from(unit.repeat(repeats)))
|
||||
|
||||
const objects = [
|
||||
Buffer.from('<< /Type /Catalog /Pages 2 0 R >>'),
|
||||
Buffer.from('<< /Type /Pages /Kids [3 0 R] /Count 1 >>'),
|
||||
Buffer.from(
|
||||
'<< /Type /Page /Parent 2 0 R /MediaBox [0 0 612 792] /Contents 4 0 R /Resources << /Font << /F1 5 0 R >> >> >>'
|
||||
),
|
||||
Buffer.concat([
|
||||
Buffer.from(`<< /Length ${compressed.length} /Filter /FlateDecode >>\nstream\n`),
|
||||
compressed,
|
||||
Buffer.from('\nendstream'),
|
||||
]),
|
||||
Buffer.from('<< /Type /Font /Subtype /Type1 /BaseFont /Helvetica >>'),
|
||||
]
|
||||
|
||||
const chunks: Buffer[] = [Buffer.from('%PDF-1.4\n')]
|
||||
const offsets: number[] = []
|
||||
let offset = chunks[0].length
|
||||
|
||||
objects.forEach((object, index) => {
|
||||
offsets.push(offset)
|
||||
const chunk = Buffer.concat([
|
||||
Buffer.from(`${index + 1} 0 obj\n`),
|
||||
object,
|
||||
Buffer.from('\nendobj\n'),
|
||||
])
|
||||
chunks.push(chunk)
|
||||
offset += chunk.length
|
||||
})
|
||||
|
||||
const xrefRows = offsets
|
||||
.map((value) => `${value.toString().padStart(10, '0')} 00000 n \n`)
|
||||
.join('')
|
||||
chunks.push(
|
||||
Buffer.from(
|
||||
`xref\n0 ${objects.length + 1}\n0000000000 65535 f \n${xrefRows}` +
|
||||
`trailer\n<< /Size ${objects.length + 1} /Root 1 0 R >>\nstartxref\n${offset}\n%%EOF\n`
|
||||
)
|
||||
)
|
||||
|
||||
return Buffer.concat(chunks)
|
||||
}
|
||||
|
||||
describe('PdfParser', () => {
|
||||
it('bounds extracted text from a compression-bomb PDF instead of exhausting the heap', async () => {
|
||||
const bomb = buildTextBombPdf(200_000)
|
||||
expect(bomb.length).toBeLessThan(200 * 1024)
|
||||
|
||||
const result = await new PdfParser().parseBuffer(bomb)
|
||||
|
||||
expect(result.metadata?.truncated).toBe(true)
|
||||
expect(result.metadata?.warning).toMatch(/parser limit/i)
|
||||
expect(result.content.length).toBeLessThanOrEqual(MAX_PDF_TEXT_CHARS)
|
||||
}, 120_000)
|
||||
|
||||
it('extracts a small PDF in full and does not flag it as truncated', async () => {
|
||||
const result = await new PdfParser().parseBuffer(buildTextBombPdf(3))
|
||||
|
||||
expect(result.metadata?.truncated).toBe(false)
|
||||
expect(result.metadata?.warning).toBeUndefined()
|
||||
expect(result.metadata?.pageCount).toBe(1)
|
||||
expect(result.content).toContain('AAAA')
|
||||
}, 30_000)
|
||||
})
|
||||
@@ -1,9 +1,136 @@
|
||||
import { readFile } from 'fs/promises'
|
||||
import { createLogger } from '@sim/logger'
|
||||
import type { FileParseResult, FileParser } from '@/lib/file-parsers/types'
|
||||
import { sanitizeTextForUTF8 } from '@/lib/file-parsers/utils'
|
||||
|
||||
const logger = createLogger('PdfParser')
|
||||
|
||||
/** Highest page number visited, bounding documents that declare huge page counts. */
|
||||
const MAX_PDF_PAGES = 10_000
|
||||
|
||||
/** Ceiling on extracted characters — roughly 3,000 pages of dense text. */
|
||||
export const MAX_PDF_TEXT_CHARS = 10_000_000
|
||||
|
||||
/** Wall-clock ceiling for extracting text from a whole document. */
|
||||
const PDF_EXTRACTION_TIMEOUT_MS = 60_000
|
||||
|
||||
const PDF_TRUNCATION_WARNING = 'PDF text extraction stopped at a parser limit and is incomplete'
|
||||
|
||||
type PdfDocumentProxy = Awaited<ReturnType<typeof import('unpdf')['getDocumentProxy']>>
|
||||
type PdfPageProxy = Awaited<ReturnType<PdfDocumentProxy['getPage']>>
|
||||
|
||||
interface TextContentChunk {
|
||||
items?: Array<{ str?: unknown; hasEOL?: unknown }>
|
||||
}
|
||||
|
||||
interface PageExtraction {
|
||||
text: string
|
||||
/** Characters consumed from the caller's budget. */
|
||||
used: number
|
||||
/** False when a budget stopped the read before the page was exhausted. */
|
||||
completed: boolean
|
||||
}
|
||||
|
||||
interface BoundedExtraction {
|
||||
text: string
|
||||
/** Page count the document declares, however many pages were actually read. */
|
||||
totalPages: number
|
||||
/** True when a budget stopped extraction before the document was exhausted. */
|
||||
truncated: boolean
|
||||
}
|
||||
|
||||
/**
|
||||
* Reads one page's text through pdf.js's streaming API, stopping once the
|
||||
* character budget or the deadline is spent.
|
||||
*
|
||||
* `extractText`/`getTextContent` buffer a page's entire text content before
|
||||
* resolving, so a page whose compressed content stream expands to hundreds of
|
||||
* megabytes reaches the V8 heap limit and aborts the process — a fatal error no
|
||||
* `try/catch` can intercept, taking every other in-flight request with it.
|
||||
* `streamTextContent` applies backpressure, so cancelling the reader stops the
|
||||
* evaluator rather than letting it run the expansion to completion.
|
||||
*/
|
||||
async function readPageWithinBudget(
|
||||
page: PdfPageProxy,
|
||||
budget: number,
|
||||
deadline: number
|
||||
): Promise<PageExtraction> {
|
||||
const reader = page
|
||||
.streamTextContent()
|
||||
.getReader() as ReadableStreamDefaultReader<TextContentChunk>
|
||||
|
||||
const parts: string[] = []
|
||||
let remaining = budget
|
||||
let completed = false
|
||||
let dropped = false
|
||||
|
||||
try {
|
||||
/**
|
||||
* Loops until content is actually dropped rather than until the budget hits
|
||||
* zero: text that ends exactly on the budget is complete, not truncated.
|
||||
*/
|
||||
while (!dropped && Date.now() <= deadline) {
|
||||
const { value, done } = await reader.read()
|
||||
if (done) {
|
||||
completed = true
|
||||
break
|
||||
}
|
||||
|
||||
for (const item of value?.items ?? []) {
|
||||
if (typeof item?.str !== 'string') continue
|
||||
|
||||
const piece = item.hasEOL === true ? `${item.str}\n` : item.str
|
||||
if (piece.length > remaining) {
|
||||
parts.push(piece.slice(0, remaining))
|
||||
remaining = 0
|
||||
dropped = true
|
||||
break
|
||||
}
|
||||
|
||||
parts.push(piece)
|
||||
remaining -= piece.length
|
||||
}
|
||||
}
|
||||
} finally {
|
||||
if (!completed) {
|
||||
try {
|
||||
await reader.cancel(new Error('PDF text extraction budget exceeded'))
|
||||
} catch {
|
||||
// Cancelling a stream that already failed is not itself an error, and
|
||||
// throwing here would mask whatever ended the read loop.
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
return { text: parts.join(''), used: budget - remaining, completed }
|
||||
}
|
||||
|
||||
async function extractTextWithinBudget(pdf: PdfDocumentProxy): Promise<BoundedExtraction> {
|
||||
const deadline = Date.now() + PDF_EXTRACTION_TIMEOUT_MS
|
||||
const totalPages = pdf.numPages
|
||||
const pageLimit = Math.min(totalPages, MAX_PDF_PAGES)
|
||||
const pageTexts: string[] = []
|
||||
|
||||
let remainingChars = MAX_PDF_TEXT_CHARS
|
||||
let truncated = totalPages > pageLimit
|
||||
|
||||
for (let pageNumber = 1; pageNumber <= pageLimit; pageNumber++) {
|
||||
const page = await pdf.getPage(pageNumber)
|
||||
const { text, used, completed } = await readPageWithinBudget(page, remainingChars, deadline)
|
||||
|
||||
remainingChars -= used
|
||||
pageTexts.push(text)
|
||||
page.cleanup()
|
||||
|
||||
if (!completed) {
|
||||
truncated = true
|
||||
break
|
||||
}
|
||||
}
|
||||
|
||||
return { text: pageTexts.join('\n').replace(/\s+/g, ' '), totalPages, truncated }
|
||||
}
|
||||
|
||||
export class PdfParser implements FileParser {
|
||||
async parseFile(filePath: string): Promise<FileParseResult> {
|
||||
try {
|
||||
@@ -28,24 +155,34 @@ export class PdfParser implements FileParser {
|
||||
try {
|
||||
logger.info('Starting to parse buffer, size:', dataBuffer.length)
|
||||
|
||||
const { extractText, getDocumentProxy } = await import('unpdf')
|
||||
const { getDocumentProxy } = await import('unpdf')
|
||||
|
||||
const uint8Array = new Uint8Array(dataBuffer)
|
||||
|
||||
const pdf = await getDocumentProxy(uint8Array)
|
||||
|
||||
const { totalPages, text } = await extractText(pdf, { mergePages: true })
|
||||
try {
|
||||
const { text, totalPages, truncated } = await extractTextWithinBudget(pdf)
|
||||
|
||||
logger.info('PDF parsed successfully, pages:', totalPages, 'text length:', text.length)
|
||||
logger.info('PDF parsed successfully, pages:', totalPages, 'text length:', text.length)
|
||||
|
||||
const cleanContent = text.replace(/\u0000/g, '')
|
||||
if (truncated) {
|
||||
logger.warn(PDF_TRUNCATION_WARNING, { totalPages, textLength: text.length })
|
||||
}
|
||||
|
||||
return {
|
||||
content: cleanContent,
|
||||
metadata: {
|
||||
pageCount: totalPages,
|
||||
source: 'unpdf',
|
||||
},
|
||||
return {
|
||||
content: sanitizeTextForUTF8(text),
|
||||
metadata: {
|
||||
pageCount: totalPages,
|
||||
source: 'unpdf',
|
||||
truncated,
|
||||
warning: truncated ? PDF_TRUNCATION_WARNING : undefined,
|
||||
},
|
||||
}
|
||||
} finally {
|
||||
// Releases the document-level page, font, and image caches, which the
|
||||
// per-page cleanup() does not touch.
|
||||
await pdf.destroy().catch(() => {})
|
||||
}
|
||||
} catch (error) {
|
||||
logger.error('Error parsing buffer:', error)
|
||||
|
||||
@@ -1,6 +1,8 @@
|
||||
export interface FileParseMetadata {
|
||||
characterCount?: number
|
||||
pageCount?: number
|
||||
/** True when a parser limit stopped extraction before the input was exhausted. */
|
||||
truncated?: boolean
|
||||
extractionMethod?: string
|
||||
warning?: string
|
||||
messages?: unknown[]
|
||||
|
||||
Reference in New Issue
Block a user