pixa/internal/handlers/referer_blocklist_internal_test.go

package handlers

import (
	"context"
	"log/slog"
	"net/http"
	"net/http/httptest"
	"sync/atomic"
	"testing"
	"time"

	"github.com/go-chi/chi/v5"
	"github.com/prometheus/client_golang/prometheus"
	"sneak.berlin/go/pixa/internal/allowlist"
	"sneak.berlin/go/pixa/internal/encurl"
	"sneak.berlin/go/pixa/internal/httpfetcher"
	"sneak.berlin/go/pixa/internal/imageprocessor"
	"sneak.berlin/go/pixa/internal/imagerequest"
	"sneak.berlin/go/pixa/internal/imgcache"
	"sneak.berlin/go/pixa/internal/signature"
)

// blockedReferer is a page on leech.example, which newRefererRoutes puts on
// referer_blocklist.
const blockedReferer = "https://leech.example/page.html"

// countingFetcher passes each fetch on to the fetcher it holds and counts it.
type countingFetcher struct {
	httpfetcher.Fetcher

	fetches atomic.Int32
}

// Fetch counts the fetch and passes it on.
func (f *countingFetcher) Fetch(
	ctx context.Context, source imagerequest.Source,
) (*httpfetcher.FetchResult, error) {
	f.fetches.Add(1)

	return f.Fetcher.Fetch(ctx, source)
}

// newRefererRoutes returns both image routes of a Handlers whose
// referer_blocklist is "leech.example" and ".hotlinker.example", the
// Handlers, and the fetcher the routes fetch through. The JPEG at photoPath
// exists on allowlistedHost and on signedHost.
func newRefererRoutes(t *testing.T) (http.Handler, *Handlers, *countingFetcher) {
	t.Helper()

	fetcher := &countingFetcher{
		Fetcher: newPhotoFetcher(t, allowlistedHost, signedHost),
	}

	cache, err := imgcache.NewCache(setupTestDB(t), imgcache.CacheConfig{
		StateDir:    t.TempDir(),
		NegativeTTL: 5 * time.Minute,
		MaxBytes:    1 << 30,
	}, prometheus.NewRegistry())
	if err != nil {
		t.Fatalf("imgcache.NewCache() error = %v", err)
	}

	svc, err := imgcache.NewService(&imgcache.ServiceConfig{
		Cache:   cache,
		Fetcher: fetcher,
		Processor: imageprocessor.New(imageprocessor.Params{
			MaxInputBytes:           50 << 20,
			MaxConcurrentProcessing: 2,
		}),
		MaxResponseSize: 50 << 20,
	})
	if err != nil {
		t.Fatalf("imgcache.NewService() error = %v", err)
	}

	encGen, err := encurl.NewGenerator(testSigningKey)
	if err != nil {
		t.Fatalf("encurl.NewGenerator() error = %v", err)
	}

	h := &Handlers{
		log:       slog.New(slog.DiscardHandler),
		imgSvc:    svc,
		encGen:    encGen,
		signer:    signature.New(testSigningKey),
		allowlist: allowlist.New([]string{allowlistedHost}),
		refererBlocklist: allowlist.New(
			[]string{"leech.example", ".hotlinker.example"}),
	}

	r := chi.NewRouter()
	r.Get("/v1/image/*", h.HandleImage())
	r.Get("/v1/e/{token}/*", h.HandleImageEnc())

	return r, h, fetcher
}

// getWithReferer sends a GET for target to routes with referer as its
// Referer header, or with none when referer is empty, and returns the
// response.
func getWithReferer(
	t *testing.T, routes http.Handler, target, referer string,
) *httptest.ResponseRecorder {
	t.Helper()

	req := httptest.NewRequestWithContext(t.Context(), http.MethodGet, target, nil)
	if referer != "" {
		req.Header.Set("Referer", referer)
	}

	rec := httptest.NewRecorder()

	routes.ServeHTTP(rec, req)
	t.Logf("GET %s with Referer %q: %d", target, referer, rec.Code)

	return rec
}

// TestRefererBlocklist verifies that both image routes refuse a request whose
// Referer names a host on referer_blocklist with 403 and the JSON error,
// without fetching from the upstream host, and serve a request with no
// Referer, one that does not parse, or one naming any other host. Hosts are
// matched as allowlist_hosts matches them.
func TestRefererBlocklist(t *testing.T) {
	t.Parallel()

	cases := []struct {
		name    string
		referer string
		want    int
	}{
		{"no referer", "", http.StatusOK},
		{"unlisted host", "https://unlisted.example/page.html", http.StatusOK},
		{"unparseable", "%zz", http.StatusOK},
		{"listed host", blockedReferer, http.StatusForbidden},
		{"subdomain of listed host", "https://www.leech.example/", http.StatusOK},
		{"subdomain of dot pattern", "https://www.hotlinker.example/a.html",
			http.StatusForbidden},
		{"dot pattern without its dot", "https://hotlinker.example/",
			http.StatusForbidden},
		{"host continuing past dot pattern",
			"https://hotlinker.example.evil.example/", http.StatusOK},
	}

	// The photo's URL on each image route.
	photoURLs := map[string]func(t *testing.T, h *Handlers) string{
		"plain URL": func(t *testing.T, _ *Handlers) string {
			t.Helper()

			return photoURL(allowlistedHost)
		},
		"encrypted URL": encPhotoURL,
	}

	for urlName, photoURLFor := range photoURLs {
		for _, tc := range cases {
			t.Run(urlName+", "+tc.name, func(t *testing.T) {
				t.Parallel()

				routes, h, fetcher := newRefererRoutes(t)

				rec := getWithReferer(t, routes, photoURLFor(t, h), tc.referer)

				if tc.want == http.StatusOK {
					requireServedPhoto(t, rec)

					return
				}

				checkErrorBody(t, rec, http.StatusForbidden, "referer blocked")

				if n := fetcher.fetches.Load(); n != 0 {
					t.Errorf("upstream fetched %d times, want 0", n)
				}
			})
		}
	}
}

// TestBlockedRefererRefusedWhenImageIsCached verifies that a request whose
// Referer is on referer_blocklist is refused even when the image it asks for
// is already cached, so the answer does not depend on the cache.
func TestBlockedRefererRefusedWhenImageIsCached(t *testing.T) {
	t.Parallel()

	routes, h, _ := newRefererRoutes(t)

	for _, target := range []string{photoURL(allowlistedHost), encPhotoURL(t, h)} {
		requireServedPhoto(t, getWithReferer(t, routes, target, ""))

		rec := getWithReferer(t, routes, target, blockedReferer)
		checkErrorBody(t, rec, http.StatusForbidden, "referer blocked")
	}
}