diff --git a/internal/document/clean/pdf.go b/internal/document/clean/pdf.go index 8ead5308..c5e111ba 100644 --- a/internal/document/clean/pdf.go +++ b/internal/document/clean/pdf.go @@ -16,9 +16,9 @@ type PDF struct { } const ( - TEXTRACT_SYNC_LIMIT = 10 * 1024 * 1024 - TEXTRACT_ASYNC_LIMIT = 500 * 1024 * 1024 - TEXTRACT_PAGE_LIMIT = 3000 + TEXTRACT_SYNC_LIMIT int64 = 10 * 1024 * 1024 + TEXTRACT_ASYNC_LIMIT int64 = 500 * 1024 * 1024 + TEXTRACT_PAGE_LIMIT = 3000 ) func NewPDF(buf io.ReadSeeker) *PDF { @@ -50,5 +50,16 @@ func (s *PDF) isCorrupted(ctx context.Context) *InvalidDocumentReason { return &reason } + err = pdfCtx.EnsurePageCount() + if err != nil { + reason = InvalidDocumentRead + return &reason + } + + if pdfCtx.PageCount > s.maxPages { + reason = InvalidDocumentLargePageCount + return &reason + } + return nil } diff --git a/internal/document/clean/pdf_test.go b/internal/document/clean/pdf_test.go index ac67e9dc..2c932686 100644 --- a/internal/document/clean/pdf_test.go +++ b/internal/document/clean/pdf_test.go @@ -31,6 +31,14 @@ func TestIsCorrupt(t *testing.T) { reason := pdf.isCorrupted(ctx) assert.Equal(t, InvalidDocumentLargeFile, *reason) }) + t.Run("large page count", func(t *testing.T) { + ctx := context.Background() + pdf := NewPDF(strings.NewReader(pdfTwoPages)) + pdf.maxPages = 1 + + reason := pdf.isCorrupted(ctx) + assert.Equal(t, InvalidDocumentLargePageCount, *reason) + }) t.Run("version 1.0", func(t *testing.T) { ctx := context.Background() pdf := NewPDF(strings.NewReader(PDF1_0)) @@ -101,7 +109,8 @@ func TestNewPDF(t *testing.T) { pdf := NewPDF(strings.NewReader(pdfHelloWorld)) assert.NotNil(t, pdf.pdf) assert.NotNil(t, pdf.config) - assert.Equal(t, 3000, pdf.maxPages) + assert.Equal(t, TEXTRACT_PAGE_LIMIT, pdf.maxPages) + assert.Equal(t, TEXTRACT_SYNC_LIMIT, pdf.maxBytes) }) } @@ -738,3 +747,56 @@ trailer startxref 485 %%EOF` + +const pdfTwoPages = `%PDF-1.4 +%âãÏÓ +1 0 obj +<> +endobj +2 0 obj +<> +endobj +3 0 obj +<>>>/MediaBox[0 0 612 792]/Contents 6 0 R>> +endobj +4 0 obj +<>>>/MediaBox[0 0 612 792]/Contents 7 0 R>> +endobj +5 0 obj +<> +endobj +6 0 obj +<> +stream +BT +/F1 14 Tf +200 400 Td +(Hello from page one!) Tj +ET +endstream +endobj +7 0 obj +<> +stream +BT +/F1 14 Tf +200 400 Td +(Goodbye from page two!) Tj +ET +endstream +endobj +xref +0 8 +0000000000 65535 f +0000000015 00000 n +0000000060 00000 n +0000000111 00000 n +0000000212 00000 n +0000000313 00000 n +0000000374 00000 n +0000000492 00000 n +trailer +<> +startxref +612 +%%EOF`