pixa/internal/handlers/first_bytes_internal_test.go

package handlers

import (
	"bytes"
	"encoding/json"
	"image/color"
	"log/slog"
	"net/http"
	"slices"
	"testing"
	"testing/fstest"
	"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/imageproxy"
	"sneak.berlin/go/pixa/internal/imagerequest"
	"sneak.berlin/go/pixa/internal/imgcache"
	"sneak.berlin/go/pixa/internal/seal"
	"sneak.berlin/go/pixa/internal/signature"
)

// mismatchedSourcePath is the path of the source that
// TestFirstBytesNotMatchingContentType serves on allowlistedHost. The mock
// fetcher names a file's Content-Type from its extension, so the source is
// sent as image/png.
const mismatchedSourcePath = "/images/source.png"

// newRoutesFetchingThrough returns both image routes, and the Handlers behind
// them, whose image service fetches through fetcher, and which have
// allowlistedHost on their allowlist. The handlers and the image service log
// to log.
func newRoutesFetchingThrough(
	t *testing.T, fetcher httpfetcher.Fetcher, log *slog.Logger,
) (*Handlers, http.Handler) {
	t.Helper()

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

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

	keys, err := seal.DeriveKeys(testSigningKey)
	if err != nil {
		t.Fatalf("seal.DeriveKeys() error = %v", err)
	}

	encGen := encurl.NewGenerator(keys.EncryptedURLKey)

	h := &Handlers{
		log:       log,
		imgSvc:    svc,
		encGen:    encGen,
		signer:    signature.New(testSigningKey),
		allowlist: allowlist.New([]string{allowlistedHost}),
	}

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

	return h, r
}

// mismatchedSourcePlainURL returns the /v1/image/ URL of mismatchedSourcePath
// on allowlistedHost at 50x50 as a JPEG.
func mismatchedSourcePlainURL(*testing.T, *Handlers) string {
	return "/v1/image/" + allowlistedHost + mismatchedSourcePath + "/50x50.jpeg"
}

// mismatchedSourceEncURL returns an encrypted /v1/e/ URL, made with h's
// generator, of mismatchedSourcePath on allowlistedHost at 50x50 as a JPEG.
func mismatchedSourceEncURL(t *testing.T, h *Handlers) string {
	t.Helper()

	token, err := h.encGen.Generate(&imagerequest.Request{
		Source: imagerequest.Source{
			Host: allowlistedHost,
			Path: mismatchedSourcePath,
		},
		Transform: imagerequest.Transform{
			Size:   imagerequest.Size{Width: 50, Height: 50},
			Format: imagerequest.FormatJPEG,
		},
	}, time.Time{})
	if err != nil {
		t.Fatalf("Generate() error = %v", err)
	}

	return "/v1/e/" + token + "/img.jpg"
}

// loggedLevels reads the JSON log lines in logs and returns the messages
// logged at each level.
func loggedLevels(t *testing.T, logs *bytes.Buffer) map[string][]string {
	t.Helper()

	t.Logf("logged:\n%s", logs.String())

	messages := make(map[string][]string)

	dec := json.NewDecoder(logs)
	for dec.More() {
		var line struct {
			Level string `json:"level"`
			Msg   string `json:"msg"`
		}

		err := dec.Decode(&line)
		if err != nil {
			t.Fatalf("decoding log line: %v", err)
		}

		messages[line.Level] = append(messages[line.Level], line.Msg)
	}

	return messages
}

// TestFirstBytesNotMatchingContentType checks that both image routes answer
// 415 for an upstream body whose first bytes are not those of the format its
// Content-Type names, and log it as a refused source, not as an error of
// pixa's. The next request for that source must answer 502 without fetching
// it again, as after an upstream error status.
func TestFirstBytesNotMatchingContentType(t *testing.T) {
	t.Parallel()

	bodies := map[string][]byte{
		"JPEG sent as image/png":          generateTestJPEG(t, 10, 10, color.Black),
		"text sent as image/png":          []byte("This is just text, not a PNG"),
		"body too short to name a format": []byte("short"),
	}

	// The source's URL on each image route.
	sourceURLs := map[string]func(*testing.T, *Handlers) string{
		"plain URL":     mismatchedSourcePlainURL,
		"encrypted URL": mismatchedSourceEncURL,
	}

	for bodyName, body := range bodies {
		for urlName, sourceURL := range sourceURLs {
			t.Run(urlName+", "+bodyName, func(t *testing.T) {
				t.Parallel()

				fetcher := &countingFetcher{
					Fetcher: httpfetcher.NewMock(fstest.MapFS{
						allowlistedHost + mismatchedSourcePath: &fstest.MapFile{
							Data: body,
						},
					}),
				}

				var logs bytes.Buffer

				h, routes := newRoutesFetchingThrough(t, fetcher,
					slog.New(slog.NewJSONHandler(&logs, nil)))
				target := sourceURL(t, h)

				checkErrorBody(t, sendGet(t, routes, target),
					http.StatusUnsupportedMediaType, "unsupported source format")

				levels := loggedLevels(t, &logs)
				if len(levels[slog.LevelError.String()]) != 0 {
					t.Errorf("logged at error level: %q",
						levels[slog.LevelError.String()])
				}

				if !slices.Contains(levels[slog.LevelWarn.String()], "source refused") {
					t.Error(`no "source refused" line logged at warning level`)
				}

				checkErrorBody(t, sendGet(t, routes, target),
					http.StatusBadGateway, "upstream error")

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