query-orchestration/internal/document/clean/pdf.go

package documentclean

import (
	"context"
	"io"

	"github.com/pdfcpu/pdfcpu/pkg/pdfcpu"
	"github.com/pdfcpu/pdfcpu/pkg/pdfcpu/model"
)

type PDF struct {
	pdf      io.ReadSeeker
	config   *model.Configuration
	maxPages int
	maxBytes int64
}

const (
	TEXTRACT_SYNC_LIMIT  int64 = 10 * 1024 * 1024
	TEXTRACT_ASYNC_LIMIT int64 = 500 * 1024 * 1024
	TEXTRACT_PAGE_LIMIT        = 3000
)

func NewPDF(buf io.ReadSeeker) *PDF {
	config := &model.Configuration{
		ValidationMode: model.ValidationRelaxed,
		Reader15:       true,
		Offline:        true,
	}

	return &PDF{
		pdf:      buf,
		config:   config,
		maxPages: TEXTRACT_PAGE_LIMIT,
		maxBytes: TEXTRACT_SYNC_LIMIT,
	}
}

func (s *PDF) isCorrupted(ctx context.Context) *InvalidDocumentReason {
	var reason InvalidDocumentReason

	pdfCtx, err := pdfcpu.ReadWithContext(ctx, s.pdf, s.config)
	if err != nil {
		reason = InvalidDocumentRead
		return &reason
	}

	if int64(pdfCtx.Read.ReadFileSize()) > s.maxBytes {
		reason = InvalidDocumentLargeFile
		return &reason
	}

	err = pdfCtx.EnsurePageCount()
	if err != nil {
		reason = InvalidDocumentRead
		return &reason
	}

	if pdfCtx.PageCount > s.maxPages {
		reason = InvalidDocumentLargePageCount
		return &reason
	}

	return nil
}