sift/internal/fetch/pdf_test.go

package fetch_test

// P5a task 01: a PDF response gets a larger body cap and a longer deadline than a web page, and
// FetchIf revalidates a cached document with a conditional GET.

import (
	"bytes"
	"context"
	"errors"
	"net/http"
	"net/http/httptest"
	"net/netip"
	"strings"
	"sync"
	"testing"
	"time"

	"git.wntrmute.dev/kyle/sift/internal/fetch"
	"git.wntrmute.dev/kyle/sift/internal/guard"
)

func loopbackFetcher() *fetch.Fetcher {
	g := &guard.Guard{AllowPrefixes: []netip.Prefix{netip.MustParsePrefix("127.0.0.0/8")}}
	return &fetch.Fetcher{Guard: g, UA: "sift-test"}
}

func body(prefix string, n int) []byte {
	b := bytes.Repeat([]byte("x"), n)
	copy(b, prefix)
	return b
}

func serve(t *testing.T, ct string, b []byte) string {
	t.Helper()
	ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if ct != "" {
			w.Header().Set("Content-Type", ct)
		}
		_, _ = w.Write(b)
	}))
	t.Cleanup(ts.Close)
	return ts.URL
}

func status(err error) int {
	var fe *fetch.Error
	if errors.As(err, &fe) {
		return fe.Status
	}
	return 0
}

func TestPDFBodyCap(t *testing.T) {
	six := 6 << 20 // over DefaultMaxBody (5 MiB), under DefaultPDFMaxBody (100 MiB)
	for _, tc := range []struct {
		name, ct string
		b        []byte
		want     int // 0 = ok
	}{
		{"pdf by type", "application/pdf", body("%PDF-1.7\n", six), 0},
		{"pdf by sniffing octet-stream", "application/octet-stream", body("%PDF-1.7\n", six), 0},
		{"pdf with no type", "", body("%PDF-1.7\n", six), 0},
		{"html keeps the 5 MiB cap", "text/html", body("<html>", six), 413},
		{"octet-stream that is not a pdf keeps the cap", "application/octet-stream", body("PK\x03\x04", six), 413},
	} {
		f := loopbackFetcher()
		res, err := f.Fetch(context.Background(), serve(t, tc.ct, tc.b))
		if tc.want == 0 {
			if err != nil || res == nil || len(res.Body) != six {
				t.Errorf("%s: err %v, body %d bytes; want the whole %d", tc.name, err, lenBody(res), six)
			}
			continue
		}
		if status(err) != tc.want {
			t.Errorf("%s: err %v (status %d), want %d", tc.name, err, status(err), tc.want)
		}
	}
}

func lenBody(r *fetch.Resource) int {
	if r == nil {
		return -1
	}
	return len(r.Body)
}

func TestPDFMaxBodyIsConfigurable(t *testing.T) {
	f := loopbackFetcher()
	f.PDFMaxBody = 1 << 20
	_, err := f.Fetch(context.Background(), serve(t, "application/pdf", body("%PDF-1.7\n", 2<<20)))
	if status(err) != 413 || !strings.Contains(err.Error(), "too large") {
		t.Errorf("2 MiB pdf over a 1 MiB PDFMaxBody: %v, want 413 too large", err)
	}
	if fetch.DefaultPDFMaxBody != 100<<20 || fetch.DefaultPDFTimeout != 120*time.Second {
		t.Errorf("defaults = %d, %v; want 100 MiB, 120s", fetch.DefaultPDFMaxBody, fetch.DefaultPDFTimeout)
	}
}

// slow sends the headers, then the body in pieces over about 600 ms.
func slow(t *testing.T, ct string, first string) string {
	t.Helper()
	ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("Content-Type", ct)
		w.WriteHeader(200)
		fl, _ := w.(http.Flusher)
		_, _ = w.Write([]byte(first))
		if fl != nil {
			fl.Flush()
		}
		for i := 0; i < 6; i++ {
			time.Sleep(100 * time.Millisecond)
			_, _ = w.Write(bytes.Repeat([]byte("y"), 1024))
			if fl != nil {
				fl.Flush()
			}
		}
	}))
	t.Cleanup(ts.Close)
	return ts.URL
}

// The web-page deadline (Timeout) still applies to HTML; a PDF gets PDFTimeout.
func TestPDFDeadline(t *testing.T) {
	f := loopbackFetcher()
	f.Timeout = 250 * time.Millisecond
	f.PDFTimeout = 3 * time.Second
	if _, err := f.Fetch(context.Background(), slow(t, "application/pdf", "%PDF-1.7\n")); err != nil {
		t.Errorf("slow pdf under PDFTimeout: %v", err)
	}
	if _, err := f.Fetch(context.Background(), slow(t, "application/octet-stream", "%PDF-1.7\n")); err != nil {
		t.Errorf("slow sniffed pdf under PDFTimeout: %v", err)
	}
	start := time.Now()
	if _, err := f.Fetch(context.Background(), slow(t, "text/html", "<html>")); err == nil {
		t.Errorf("slow html past Timeout succeeded; want an error")
	}
	if d := time.Since(start); d > 2*time.Second {
		t.Errorf("slow html took %v to fail; the 250 ms deadline was not applied", d)
	}
}

func TestFetchIf(t *testing.T) {
	var mu sync.Mutex
	var got []string
	ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		mu.Lock()
		got = append(got, r.Header.Get("If-None-Match")+"|"+r.Header.Get("If-Modified-Since"))
		mu.Unlock()
		if r.Header.Get("If-None-Match") == `"v1"` {
			w.WriteHeader(http.StatusNotModified)
			return
		}
		w.Header().Set("Content-Type", "application/pdf")
		w.Header().Set("ETag", `"v2"`)
		w.Header().Set("Last-Modified", "Fri, 25 Sep 2026 12:00:00 GMT")
		_, _ = w.Write([]byte("%PDF-1.7\nnew"))
	}))
	defer ts.Close()
	f := loopbackFetcher()

	res, err := f.FetchIf(context.Background(), ts.URL, `"v1"`, "")
	if err != nil || res == nil || res.Status != http.StatusNotModified || len(res.Body) != 0 {
		t.Fatalf("matching ETag: res %+v err %v; want status 304, empty body, nil error", res, err)
	}
	res, err = f.FetchIf(context.Background(), ts.URL, `"old"`, "Thu, 24 Sep 2026 12:00:00 GMT")
	if err != nil || res.Status != 200 || string(res.Body) != "%PDF-1.7\nnew" || res.Header.Get("ETag") != `"v2"` {
		t.Fatalf("changed document: res %+v err %v; want 200 with the new body and ETag", res, err)
	}
	mu.Lock()
	sent := append([]string(nil), got...)
	mu.Unlock() // never hold mu across another request: the handler takes it
	if len(sent) != 2 || sent[0] != `"v1"|` || sent[1] != `"old"|Thu, 24 Sep 2026 12:00:00 GMT` {
		t.Errorf("validators sent: %q", sent)
	}
	// With no validators FetchIf is a plain Fetch.
	if _, err := f.FetchIf(context.Background(), ts.URL, "", ""); err != nil {
		t.Errorf("FetchIf with no validators: %v", err)
	}
}