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"
	"github.com/pdfcpu/pdfcpu/pkg/pdfcpu/types"
)

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

const (
	TEXTRACT_SYNC_LIMIT    int64 = 10 * 1024 * 1024
	TEXTRACT_ASYNC_LIMIT   int64 = 500 * 1024 * 1024
	TEXTRACT_PAGE_LIMIT          = 3000
	TEXTRACT_MAX_DIMENSION       = 10000
	TEXTRACT_MIN_DIMENSION       = 50
	TEXTRACT_MIN_DPI             = 72
	PDF_DEFAULT_DPI              = 72
)

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,
		maxDimension: TEXTRACT_MAX_DIMENSION,
		minDimension: TEXTRACT_MIN_DIMENSION,
		minDPI:       TEXTRACT_MIN_DPI,
	}
}

func (s *PDF) isCorrupt(ctx context.Context) *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
	}

	docErr := s.isValidPageCount(pdfCtx)
	if docErr != nil {
		return docErr
	}

	docErr = s.isValidDimensions(pdfCtx)
	if docErr != nil {
		return docErr
	}

	return nil
}

func (s *PDF) isValidPageCount(pdfCtx *model.Context) *InvalidDocumentReason {
	err := pdfCtx.EnsurePageCount()
	if err != nil {
		reason := InvalidDocumentRead
		return &reason
	}

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

	return nil
}

func (s *PDF) isValidDimensions(pdfCtx *model.Context) *InvalidDocumentReason {
	err := pdfCtx.EnsurePageCount()
	if err != nil {
		reason := InvalidDocumentRead
		return &reason
	}

	_, err = s.pdf.Seek(0, io.SeekStart)
	if err != nil {
		reason := InvalidDocumentRead
		return &reason
	}

	pageSet := types.IntSet{}
	for pageNum := 1; pageNum <= pdfCtx.PageCount; pageNum++ {
		pageSet[pageNum] = true
	}

	boxes, err := pdfCtx.PageBoundaries(pageSet)
	if err != nil {
		reason := InvalidDocumentRead
		return &reason
	}

	for _, box := range boxes {
		mediaBox := box.MediaBox()
		if mediaBox == nil {
			continue
		}

		dims := mediaBox.Dimensions()

		if dims.Height > float64(s.maxDimension) ||
			dims.Width > float64(s.maxDimension) {
			reason := InvalidDocumentLargeDimensions
			return &reason
		} else if dims.Height < float64(s.minDimension) ||
			dims.Width < float64(s.minDimension) {
			reason := InvalidDocumentSmallDimensions
			return &reason
		}

		dpiErr := s.isValidDPI(mediaBox, dims)
		if dpiErr != nil {
			return dpiErr
		}
	}

	return nil
}

func (s *PDF) isValidDPI(mediaBox *types.Rectangle, dims types.Dim) *InvalidDocumentReason {
	widthDPI := mediaBox.Width() / dims.ToInches().Width
	heightDPI := mediaBox.Height() / dims.ToInches().Height

	if widthDPI < float64(s.minDPI) || heightDPI < float64(s.minDPI) {
		reason := InvalidDocumentSmallDPI
		return &reason
	}

	return nil
}