sift/internal/pdftext/pdftext.go

// Package pdftext extracts text, title and outline from a PDF. Extraction runs
// in a child process (see Runner) so a page that hangs or panics cannot take
// the parent down; Extract is the child's own work, usable directly in tests.
package pdftext

import (
	"encoding/json"
	"errors"
	"flag"
	"fmt"
	"io"
	"os"
	"runtime/debug"
	"strings"
	"unicode/utf16"

	"github.com/go-pdfkit/extract"
	"github.com/go-pdfkit/reader"
)

// The errors Extract and Runner can return.
var (
	ErrEncrypted   = errors.New("encrypted PDF (password required)")
	ErrUnreadable  = errors.New("unreadable PDF")
	ErrNoText      = errors.New("no text layer (scanned PDF?)")
	ErrTimeout     = errors.New("pdf extraction timed out")
	ErrTooMuchText = errors.New("pdf text over the output limit")
)

// OutlineItem is one bookmark: its title, the 1-based page it points at, and
// its nesting level (1 = top level).
type OutlineItem struct {
	Title string `json:"title"`
	Page  int    `json:"page"`  // 1-based
	Level int    `json:"level"` // 1 = top level
}

// Meta is the document-level summary written as the first output line.
type Meta struct {
	Pages   int           `json:"pages"`   // pages in the document (not only those extracted)
	Title   string        `json:"title"`   // /Info /Title, "" when absent
	Outline []OutlineItem `json:"outline"` // never nil: [] when the document has none
}

// Doc is an extracted document: Text[i] is page i+1's text, for the pages
// extracted (len(Text) ≤ Pages). Failed lists pages whose extraction failed
// (their Text is "").
type Doc struct {
	Meta
	Text   []string
	Failed []int
}

// metaLine is the first output line: the document summary.
type metaLine struct {
	Meta Meta `json:"meta"`
}

// pageOK and pageErr are the per-page lines. A page that extracts yields pageOK
// even when its text is empty (an image-only page); one that fails yields
// pageErr and extraction continues.
type pageOK struct {
	Page int    `json:"page"`
	Text string `json:"text"`
}

type pageErr struct {
	Page  int    `json:"page"`
	Error string `json:"error"`
}

// Extract writes one {"meta":{…}} line, then one {"page":n,"text":"…"} line per
// page from 1 up to min(Pages, maxPages); a page that fails yields
// {"page":n,"error":"…"} and extraction goes on. It returns ErrEncrypted when
// the document needs a user password, ErrUnreadable when it cannot be opened or
// has no pages, and never panics.
func Extract(pdf []byte, maxPages int, w io.Writer) error {
	doc, err := reader.Open(pdf)
	if err != nil {
		if errors.Is(err, reader.ErrWrongPassword) {
			return ErrEncrypted
		}
		return ErrUnreadable
	}
	if doc.PageCount() == 0 {
		return ErrUnreadable
	}

	meta := Meta{
		Pages:   doc.PageCount(),
		Title:   infoTitle(doc),
		Outline: walkOutline(doc),
	}
	if meta.Outline == nil {
		meta.Outline = []OutlineItem{}
	}
	if err := writeJSONLine(w, metaLine{Meta: meta}); err != nil {
		return err
	}

	n := doc.PageCount()
	if maxPages > 0 && n > maxPages {
		n = maxPages
	}
	for i := 1; i <= n; i++ {
		func() {
			defer func() {
				if r := recover(); r != nil {
					_ = writeJSONLine(w, pageErr{Page: i, Error: fmt.Sprintf("panic: %v", r)})
				}
			}()
			text, err := extract.Text(doc, i)
			if err != nil {
				_ = writeJSONLine(w, pageErr{Page: i, Error: err.Error()})
				return
			}
			_ = writeJSONLine(w, pageOK{Page: i, Text: text})
		}()
	}
	return nil
}

// writeJSONLine marshals v and writes it as one JSON line.
func writeJSONLine(w io.Writer, v any) error {
	b, err := json.Marshal(v)
	if err != nil {
		return err
	}
	if _, err := w.Write(b); err != nil {
		return err
	}
	_, err = w.Write([]byte{'\n'})
	return err
}

// infoTitle returns the document title from the trailer's /Info /Title, decoded
// as a PDF text string (rule 3), "" when absent.
func infoTitle(doc *reader.Document) string {
	trailer := doc.Trailer()
	if trailer == nil {
		return ""
	}
	info, err := doc.Resolve(trailer.Get("Info"))
	if err != nil || info == nil {
		return ""
	}
	d, ok := reader.ToDict(info)
	if !ok {
		return ""
	}
	return textString(d.Get("Title"))
}

// textString decodes a reader.String (raw bytes) as a PDF text string per rule 3.
func textString(o reader.Object) string {
	s, ok := reader.ToString(o)
	if !ok {
		return ""
	}
	return decodePDFString(s)
}

// decodePDFString follows rule 3: bytes starting FE FF are UTF-16BE (the rest is
// decoded as such); otherwise each byte is treated as a Latin-1 code point.
func decodePDFString(b []byte) string {
	if len(b) >= 2 && b[0] == 0xFE && b[1] == 0xFF {
		return strings.TrimSpace(string(utf16.Decode(utf16BE(b[2:]))))
	}
	return strings.TrimSpace(string(b))
}

// utf16BE packs big-endian byte pairs into uint16s for utf16.Decode.
func utf16BE(b []byte) []uint16 {
	n := len(b) / 2
	u := make([]uint16, n)
	for i := 0; i < 2*n; i += 2 {
		u[i/2] = uint16(b[i])<<8 | uint16(b[i+1])
	}
	return u
}

// ChildMain is `sift pdftext`: flags (-max-pages N, default 2000), the PDF on
// stdin (read with a 200 MiB cap), Extract to stdout. Exit codes: 0 ok, 3
// ErrEncrypted, 4 ErrUnreadable, 1 anything else. It runs with a soft memory
// limit, reads no environment variable, and opens no network connection.
func ChildMain(args []string, stdin io.Reader, stdout io.Writer) int {
	debug.SetMemoryLimit(768 << 20)

	fs := flag.NewFlagSet("pdftext", flag.ContinueOnError)
	maxPages := fs.Int("max-pages", 2000, "maximum number of pages to extract")
	if err := fs.Parse(args); err != nil {
		return 1
	}

	pdf, err := io.ReadAll(io.LimitReader(stdin, 200<<20))
	if err != nil {
		fmt.Fprintln(os.Stderr, "pdftext: read input:", err)
		return 4
	}
	if len(pdf) == 0 {
		fmt.Fprintln(os.Stderr, "pdftext: empty input")
		return 4
	}

	if err := Extract(pdf, *maxPages, stdout); err != nil {
		switch {
		case errors.Is(err, ErrEncrypted):
			return 3
		case errors.Is(err, ErrUnreadable):
			return 4
		default:
			return 1
		}
	}
	return 0
}