sift/internal/pdftext/runner.go

package pdftext

import (
	"bufio"
	"bytes"
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"os"
	"os/exec"
	"strconv"
	"strings"
	"time"
)

const (
	// defaultDeadline bounds a child that never finishes.
	defaultDeadline = 60 * time.Second
	// defaultMaxText caps child output before it is rejected.
	defaultMaxText = 64 << 20 // 64 MiB
	// defaultMaxPages caps how many pages the child extracts.
	defaultMaxPages = 2000
	// stderrCap keeps only the first part of the child's stderr, for messages.
	stderrCap = 4096
)

// Runner executes the extraction child and parses its output into a Doc.
type Runner struct {
	Exe      string        // "" = /proc/self/exe
	Args     []string      // nil = []string{"pdftext"}; -max-pages is appended
	Env      []string      // the child's whole environment; nil = none
	Deadline time.Duration // 0 = 60 s
	MaxText  int64         // bytes of child output; 0 = 64 MiB
	MaxPages int           // 0 = 2000
}

// Run executes the child, writes pdf to its stdin, and parses its output into a
// Doc. The child is killed at the deadline (or when ctx is done, whichever is
// sooner) and always reaped before Run returns.
func (r *Runner) Run(ctx context.Context, pdf []byte) (*Doc, error) {
	deadline := r.Deadline
	if deadline == 0 {
		deadline = defaultDeadline
	}
	maxText := r.MaxText
	if maxText == 0 {
		maxText = defaultMaxText
	}
	maxPages := r.MaxPages
	if maxPages == 0 {
		maxPages = defaultMaxPages
	}

	ctx, cancel := context.WithDeadline(ctx, time.Now().Add(deadline))
	defer cancel()

	cmd := exec.CommandContext(ctx, r.exe(), r.args(maxPages)...)
	if r.Env == nil {
		cmd.Env = []string{}
	} else {
		cmd.Env = r.Env
	}

	stdin, err := cmd.StdinPipe()
	if err != nil {
		return nil, wrapUnreadable(err)
	}
	stdout, err := cmd.StdoutPipe()
	if err != nil {
		return nil, wrapUnreadable(err)
	}
	stderr, err := cmd.StderrPipe()
	if err != nil {
		return nil, wrapUnreadable(err)
	}

	var stdoutBuf, stderrBuf bytes.Buffer
	done := make(chan error, 1)
	go func() {
		// One extra byte so output over the limit is distinguishable from
		// output exactly at it.
		_, e := io.Copy(&stdoutBuf, io.LimitReader(stdout, maxText+1))
		done <- e
	}()
	stderrDone := make(chan struct{})
	go func() {
		io.Copy(&limitedWriter{n: stderrCap, w: &stderrBuf}, stderr)
		close(stderrDone)
	}()
	go func() {
		if _, err := stdin.Write(pdf); err != nil {
			return
		}
		_ = stdin.Close()
	}()

	if err := cmd.Start(); err != nil {
		cancel()
		_ = stdout.Close()
		_ = stderr.Close()
		<-done
		<-stderrDone
		return nil, wrapUnreadable(err)
	}

	readErr := <-done

	// Output over the limit: kill the child at once. Checked before waiting for stderr, which
	// only closes when the child exits — and a child blocked on a full stdout pipe would not,
	// until the deadline.
	if int64(stdoutBuf.Len()) > maxText {
		cancel()
		<-stderrDone
		_ = cmd.Wait()
		return nil, ErrTooMuchText
	}
	<-stderrDone

	waitErr := cmd.Wait()

	// A deadline reached while the child ran is a timeout, even if the kill
	// surfaced as some other error.
	if ctx.Err() == context.DeadlineExceeded {
		return nil, ErrTimeout
	}
	if readErr != nil && !errors.Is(readErr, io.EOF) {
		return nil, wrapUnreadable(fmt.Errorf("read output: %w", readErr))
	}
	if waitErr != nil {
		exitErr, ok := waitErr.(*exec.ExitError)
		if !ok {
			return nil, wrapUnreadable(fmt.Errorf("child: %v", waitErr))
		}
		switch exitErr.ExitCode() {
		case 3:
			return nil, ErrEncrypted
		case 4:
			return nil, ErrUnreadable
		default:
			return nil, wrapUnreadable([]byte(firstLine(stderrBuf.Bytes())))
		}
	}

	doc, err := parseDoc(stdoutBuf.Bytes(), maxPages)
	if err != nil {
		return nil, err
	}
	return doc, nil
}

// exe resolves the executable to run, defaulting to this process's own path.
func (r *Runner) exe() string {
	if r.Exe != "" {
		return r.Exe
	}
	if e, err := os.Readlink("/proc/self/exe"); err == nil {
		return e
	}
	return "sift"
}

// args builds the child's argument list: the configured args, or "pdftext",
// with -max-pages appended.
func (r *Runner) args(maxPages int) []string {
	args := []string{"pdftext"}
	if r.Args != nil {
		args = r.Args
	}
	return append(args, "-max-pages", strconv.Itoa(maxPages))
}

// parseDoc turns the child's JSON Lines into a Doc. The first line is the meta;
// each following line is a page's text or its error. It returns ErrNoText when
// no page was extracted or every extracted page's text is blank, and
// ErrUnreadable when the output is not what the child should have produced.
func parseDoc(data []byte, maxPages int) (*Doc, error) {
	sc := bufio.NewScanner(bytes.NewReader(data))
	sc.Buffer(make([]byte, 1<<20), 64<<20)

	var doc Doc
	for sc.Scan() {
		var line struct {
			Meta  *Meta   `json:"meta"`
			Page  int     `json:"page"`
			Text  *string `json:"text"`
			Error string  `json:"error"`
		}
		if err := json.Unmarshal(sc.Bytes(), &line); err != nil {
			return nil, wrapUnreadable(nil)
		}
		if line.Meta != nil {
			doc.Meta = *line.Meta
			continue
		}
		appendPage(&doc, line.Page, line.Text, line.Error)
	}
	if err := sc.Err(); err != nil {
		return nil, wrapUnreadable(nil)
	}

	if len(doc.Text) == 0 {
		return nil, ErrNoText
	}
	for _, t := range doc.Text {
		if strings.TrimSpace(t) != "" {
			return &doc, nil
		}
	}
	return nil, ErrNoText
}

// appendPage records one page's result, growing Text to fit page n.
func appendPage(doc *Doc, n int, text *string, errMsg string) {
	for len(doc.Text) < n {
		doc.Text = append(doc.Text, "")
	}
	if text != nil {
		doc.Text[n-1] = *text
		return
	}
	doc.Text[n-1] = ""
	if errMsg != "" {
		doc.Failed = append(doc.Failed, n)
	}
}

// wrapUnreadable wraps a cause (an error or a chunk of stderr) as ErrUnreadable.
func wrapUnreadable(v any) error {
	switch s := v.(type) {
	case nil:
		return ErrUnreadable
	case error:
		return fmt.Errorf("%w: %v", ErrUnreadable, s)
	case []byte:
		return fmt.Errorf("%w: %s", ErrUnreadable, firstLine(s))
	case string:
		return fmt.Errorf("%w: %s", ErrUnreadable, s)
	}
	return ErrUnreadable
}

// firstLine returns the first line of b, or b when it has no newline.
func firstLine(b []byte) string {
	if i := bytes.IndexByte(b, '\n'); i >= 0 {
		return string(b[:i])
	}
	return string(b)
}

// limitedWriter writes at most n bytes, dropping the rest so a child that
// overflows stderr cannot block the parent.
type limitedWriter struct {
	n int64
	w io.Writer
}

func (l *limitedWriter) Write(p []byte) (int, error) {
	if l.n <= 0 {
		return 0, nil
	}
	if int64(len(p)) > l.n {
		p = p[:l.n]
	}
	n, err := l.w.Write(p)
	l.n -= int64(n)
	if l.n < 0 {
		l.n = 0
	}
	return n, err
}