clean, fetch: recognise PDFs; 100 MB / 120 s for PDFs; FetchIf revalidates
Recognise PDFs in clean.Kind (declared types and %PDF- sniffing on octet-stream/bodyless bodies), give PDF responses the 100 MB body cap and 120 s deadline while web pages keep 5 MB / 20 s, and add FetchIf, which sends If-None-Match / If-Modified-Since and maps a 304 to a Resource with Status 304 and an empty Body. Fetch and FetchIf share one doFetch path for guard, client, limits and error mapping. Co-Authored-By: ornith-1.5-35b-a3b <ornith-1.5-35b-a3b@llama.cpp.invalid>
This commit is contained in:
ornith-1.5-35b-a3b
@@ -0,0 +1,233 @@
|
||||
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,*/*;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
|
||||
}
|
||||
Reference in New Issue
Block a user