sift/internal/fetch/pdf.go

package fetch

import (
	"bytes"
	"context"
	"io"
	"net/http"
	"strings"
	"time"
)

const (
	// DefaultPDFMaxBody is the body cap for a PDF response: 100 MiB, versus the
	// 5 MiB web cap (DefaultMaxBody).
	DefaultPDFMaxBody = 100 << 20
	// DefaultPDFTimeout is the deadline for a PDF response: 120 s, versus the
	// 20 s web deadline (DefaultTimeout).
	DefaultPDFTimeout = 120 * time.Second
)

// magicPDF marks the start of a PDF, found within the first sniffWindow bytes of
// an octet-stream or a bodyless response.
var magicPDF = []byte("%PDF-")

// sniffWindow is how many leading bytes are inspected to recognise a PDF whose
// Content-Type does not say so.
const sniffWindow = 1024

func (f *Fetcher) pdfMaxBody() int64 {
	if f.PDFMaxBody <= 0 {
		return DefaultPDFMaxBody
	}
	return f.PDFMaxBody
}

func (f *Fetcher) pdfTimeout() time.Duration {
	if f.PDFTimeout <= 0 {
		return DefaultPDFTimeout
	}
	return f.PDFTimeout
}

// clientTimeout is the http.Client.Timeout: it must cover the longest deadline
// this Fetcher can give a request, so a PDF is not cut off at the web deadline.
func (f *Fetcher) clientTimeout() time.Duration {
	t := f.timeout()
	if p := f.pdfTimeout(); p > t {
		return p
	}
	return t
}

// normalizeContentType lower-cases a media type and drops its parameters.
func normalizeContentType(ct string) string {
	ct = strings.ToLower(strings.TrimSpace(ct))
	if i := strings.IndexByte(ct, ';'); i >= 0 {
		ct = strings.TrimSpace(ct[:i])
	}
	return ct
}

// isPDFContentType reports whether a media type is unambiguously a PDF.
func isPDFContentType(ct string) bool {
	ct = normalizeContentType(ct)
	return ct == "application/pdf" || ct == "application/x-pdf"
}

// FetchIf GETs rawURL like Fetch but adds conditional-GET validators: a non-empty
// etag is sent as If-None-Match and a non-empty lastModified as If-Modified-Since.
// A 304 answer returns a Resource with Status 304, an empty Body and a nil error.
// With both empty it is exactly Fetch.
func (f *Fetcher) FetchIf(ctx context.Context, rawURL, etag, lastModified string) (*Resource, error) {
	return f.doFetch(ctx, rawURL, etag, lastModified)
}

// doFetch is the shared body of Fetch and FetchIf: one guard check, one client,
// one set of limits and error mappings, differing only in request headers.
func (f *Fetcher) doFetch(ctx context.Context, rawURL, etag, lastModified string) (*Resource, error) {
	if f == nil || f.Guard == nil {
		return nil, errf(500, "", "internal error")
	}

	// A PDF may take longer than a web page. Give the request the PDF deadline
	// and cancel it at the ordinary (web) deadline unless the response is
	// recognised as a PDF sooner, by headers or by the first bytes sniffed.
	pdfDeadline := f.pdfTimeout()
	ctx, cancel := context.WithTimeout(ctx, pdfDeadline)
	var webTimer *time.Timer
	if web := f.timeout(); web < pdfDeadline {
		webTimer = time.AfterFunc(web, cancel)
	}

	if _, _, err := f.Guard.CheckURL(ctx, rawURL); err != nil {
		cancel()
		return nil, mapGuard(err)
	}

	req, err := http.NewRequestWithContext(ctx, http.MethodGet, rawURL, nil)
	if err != nil {
		cancel()
		return nil, errf(400, "", "invalid url")
	}
	req.Header.Set("User-Agent", f.ua())
	req.Header.Set("Accept", "text/html,application/xhtml+xml,application/json;q=0.9,text/plain;q=0.8,application/pdf;q=0.8,*/*;q=0.1")
	if etag != "" {
		req.Header.Set("If-None-Match", etag)
	}
	if lastModified != "" {
		req.Header.Set("If-Modified-Since", lastModified)
	}
	// no cookies, no auth, no Referer

	resp, err := f.httpClient().Do(req)
	if err != nil {
		cancel()
		return nil, mapDial(err)
	}
	defer resp.Body.Close()

	// The headers are in; a declared PDF type no longer needs the web deadline.
	ct := resp.Header.Get("Content-Type")
	pdfByHeader := isPDFContentType(ct)
	if pdfByHeader && webTimer != nil {
		webTimer.Stop()
	}

	body, err := f.readBody(resp, ct, pdfByHeader, func() {
		// Sniffing an octet-stream or bodyless response as a PDF: keep the PDF
		// deadline instead of cancelling at the web deadline.
		if webTimer != nil {
			webTimer.Stop()
		}
	})
	if err != nil {
		cancel()
		return nil, err
	}

	final := rawURL
	if resp.Request != nil && resp.Request.URL != nil {
		final = resp.Request.URL.String()
	}
	res := &Resource{
		URL:         rawURL,
		FinalURL:    final,
		Status:      resp.StatusCode,
		ContentType: ct,
		Header:      resp.Header.Clone(),
		Body:        body,
		FetchedAt:   time.Now().UTC(),
	}

	switch {
	case resp.StatusCode == http.StatusOK:
		cancel()
		return res, nil
	case resp.StatusCode == 403 || resp.StatusCode == 429 || resp.StatusCode == 451:
		cancel()
		return res, errf(422, "browser", "upstream %d (likely bot-blocked); hint: browser", resp.StatusCode)
	case resp.StatusCode == 503 && looksLikeChallenge(body, resp.Header):
		cancel()
		return res, errf(422, "browser", "upstream 503 (likely bot-blocked); hint: browser")
	case resp.StatusCode == http.StatusNotModified:
		cancel()
		_, _ = io.Copy(io.Discard, resp.Body)
		return &Resource{
			URL:         rawURL,
			FinalURL:    final,
			Status:      resp.StatusCode,
			ContentType: ct,
			Header:      resp.Header.Clone(),
		}, nil
	case resp.StatusCode >= 400:
		cancel()
		return res, errf(422, "", "upstream %d", resp.StatusCode)
	default:
		cancel()
		return res, nil
	}
}

// readBody reads resp.Body under the cap that applies and reports whether the
// response is a PDF (by declared type or by sniffing an octet-stream/bodyless
// body). onPDF, when non-nil, is called once the body is sniffed as a PDF so the
// caller can drop the web deadline.
func (f *Fetcher) readBody(resp *http.Response, ct string, pdfByHeader bool, onPDF func()) ([]byte, error) {
	norm := normalizeContentType(ct)
	if !pdfByHeader && norm != "application/octet-stream" && norm != "binary/octet-stream" && norm != "" {
		// A known non-PDF type: the web cap, no sniffing.
		body, err := io.ReadAll(io.LimitReader(resp.Body, f.maxBody()+1))
		if err != nil {
			return nil, mapDial(err)
		}
		if int64(len(body)) > f.maxBody() {
			return nil, errf(413, "", "too large (>%d MB decompressed)", f.maxBody()>>20)
		}
		return body, nil
	}

	// PDF by declared type, or a type that may hide one: sniff the leading bytes
	// to pick the cap, then read under it.
	pdf := pdfByHeader
	cap := f.maxBody()
	if pdf {
		cap = f.pdfMaxBody()
	}
	var body []byte
	if !pdfByHeader {
		head, err := io.ReadAll(io.LimitReader(resp.Body, sniffWindow))
		if err != nil {
			return nil, mapDial(err)
		}
		if bytes.Contains(head, magicPDF) {
			pdf = true
			cap = f.pdfMaxBody()
			if onPDF != nil {
				onPDF()
			}
		}
		body = head
	}
	if remaining := cap + 1 - int64(len(body)); remaining > 0 {
		more, err := io.ReadAll(io.LimitReader(resp.Body, remaining))
		if err != nil {
			return nil, mapDial(err)
		}
		body = append(body, more...)
	}
	if int64(len(body)) > cap {
		return nil, errf(413, "", "too large (>%d MB decompressed)", cap>>20)
	}
	return body, nil
}