pixa/internal/server/image_proxy_internal_test.go

package server

import (
	"bytes"
	"context"
	"database/sql"
	"image"
	"image/color"
	"image/jpeg"
	"image/png"
	"io"
	"net"
	"net/http"
	"net/http/httptest"
	"net/netip"
	"os"
	"path/filepath"
	"sync/atomic"
	"testing"
	"time"

	"github.com/prometheus/client_golang/prometheus"
	"go.uber.org/fx"
	"go.uber.org/fx/fxtest"

	"sneak.berlin/go/pixa/internal/config"
	"sneak.berlin/go/pixa/internal/database"
	"sneak.berlin/go/pixa/internal/globals"
	"sneak.berlin/go/pixa/internal/handlers"
	"sneak.berlin/go/pixa/internal/healthcheck"
	"sneak.berlin/go/pixa/internal/httpfetcher"
	"sneak.berlin/go/pixa/internal/imagerequest"
	"sneak.berlin/go/pixa/internal/logger"
	"sneak.berlin/go/pixa/internal/middleware"
	"sneak.berlin/go/pixa/internal/providers"
)

// upstreamHost is the upstream host of the image URLs below. It is a
// documentation address (RFC 5737), which the fetcher's URL check accepts as
// public; the fetcher's dial function connects it to the test upstream server.
const upstreamHost = "192.0.2.10"

// TestImageProxyFlow requests images through pixa's router, handlers,
// upstream fetcher, image processor, disk cache and database, with only the
// upstream origin replaced by a local test server. The first request for a URL
// is fetched and converted; the second is served from the cache without
// another upstream request. The source and the converted image are then on
// disk, with their rows in the database.
func TestImageProxyFlow(t *testing.T) {
	t.Parallel()

	source := encodeTestPNG(t, 64, 48)

	tests := []struct {
		name          string
		sizeFormat    string // the <size>.<format> part of the image URL
		contentType   string
		decodeConfig  func(io.Reader) (image.Config, error)
		width, height int
	}{
		{"resize and convert to JPEG", "32x24.jpeg", "image/jpeg",
			jpeg.DecodeConfig, 32, 24},
		{"orig", "orig.orig", "image/png", png.DecodeConfig, 64, 48},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			t.Parallel()

			var upstreamRequests atomic.Int32

			upstream := httptest.NewServer(http.HandlerFunc(
				func(w http.ResponseWriter, _ *http.Request) {
					upstreamRequests.Add(1)
					w.Header().Set("Content-Type", "image/png")
					_, _ = w.Write(source)
				}))
			t.Cleanup(upstream.Close)

			s, db, stateDir := startImageProxy(t, upstream)
			target := "/v1/image/" + upstreamHost + "/photo.png/" + tt.sizeFormat

			first := getImage(t, s, target)

			if got := first.Header().Get("X-Pixa-Cache"); got != "MISS" {
				t.Errorf("first X-Pixa-Cache = %q, want MISS", got)
			}

			if got := first.Header().Get("Content-Type"); got != tt.contentType {
				t.Errorf("Content-Type = %q, want %q", got, tt.contentType)
			}

			decoded, err := tt.decodeConfig(bytes.NewReader(first.Body.Bytes()))
			if err != nil {
				t.Fatalf("decoding the image: %v", err)
			}

			if decoded.Width != tt.width || decoded.Height != tt.height {
				t.Errorf("image is %dx%d, want %dx%d",
					decoded.Width, decoded.Height, tt.width, tt.height)
			}

			second := getImage(t, s, target)

			if got := second.Header().Get("X-Pixa-Cache"); got != "HIT" {
				t.Errorf("second X-Pixa-Cache = %q, want HIT", got)
			}

			if !bytes.Equal(second.Body.Bytes(), first.Body.Bytes()) {
				t.Error("the second response is not the image the first served")
			}

			if got := upstreamRequests.Load(); got != 1 {
				t.Errorf("upstream received %d requests, want 1", got)
			}

			checkSourceCached(t, db, stateDir, source)
			checkVariantCached(t, db, stateDir, first.Body.Bytes())
		})
	}
}

// startImageProxy starts the components pixad's fx app builds, from a config
// with a fresh state directory and upstreamHost on the allowlist, and with an
// upstream fetcher that connects every upstream address to upstream. It
// returns the server with its routes, the database and the state directory.
func startImageProxy(
	t *testing.T, upstream *httptest.Server,
) (*Server, *sql.DB, string) {
	t.Helper()

	stateDir := t.TempDir()
	cfg := &config.Config{
		SigningKey:     testSigningKey,
		StateDir:       stateDir,
		DBURL:          "file:" + filepath.Join(stateDir, "state.sqlite3"),
		AllowlistHosts: []string{upstreamHost},
		// The test upstream server has no TLS.
		AllowHTTP: true,
		// A limit of its own, so the cache does not size itself from the
		// host's free disk space.
		CacheMaxBytes:              64 << 20,
		CacheMaxBytesExplicit:      true,
		UpstreamFetchTimeout:       config.DefaultUpstreamFetchTimeout,
		UpstreamMaxResponseSize:    config.DefaultUpstreamMaxResponseSize,
		UpstreamConnectionsPerHost: config.DefaultUpstreamConnectionsPerHost,
		UpstreamConnections:        config.DefaultUpstreamConnections,
		MaxConcurrentProcessing:    2,
		DownstreamTimeout:          config.DefaultDownstreamTimeout,
	}

	dialUpstream := func(ctx context.Context, network, _ string) (net.Conn, error) {
		var dialer net.Dialer

		return dialer.DialContext(ctx, network, upstream.Listener.Addr().String())
	}

	var (
		h  *handlers.Handlers
		mw *middleware.Middleware
		db *database.Database
	)

	app := fxtest.New(t,
		fx.Supply(cfg),
		fx.Provide(
			globals.New,
			logger.New,
			database.New,
			healthcheck.New,
			handlers.New,
			middleware.New,
			func() prometheus.Registerer { return prometheus.NewRegistry() },
		),
		fx.Provide(providers.Constructors()...),
		fx.Decorate(func(fetcherConfig *httpfetcher.Config) *httpfetcher.Config {
			fetcherConfig.DialContext = dialUpstream

			return fetcherConfig
		}),
		fx.Populate(&h, &mw, &db),
	)
	app.RequireStart()
	t.Cleanup(app.RequireStop)

	// Requests go straight to the router, as in newTestServer; the server's
	// own start hook, which listens on a port, is left out.
	s := &Server{config: cfg, mw: mw, h: h}
	s.SetupRoutes()

	return s, db.DB(), stateDir
}

// startImageProxyWithBlockedNetworks starts the components pixad's fx app
// builds, from a config with a fresh state directory and blockedNetworks as
// blocked_networks. Unlike startImageProxy it keeps pixad's dial function,
// which checks every address the fetcher connects to against the built-in
// blocked ranges and blocked_networks. It returns the server with its
// routes.
func startImageProxyWithBlockedNetworks(
	t *testing.T, blockedNetworks []netip.Prefix,
) *Server {
	t.Helper()

	stateDir := t.TempDir()
	cfg := &config.Config{
		SigningKey:      testSigningKey,
		StateDir:        stateDir,
		DBURL:           "file:" + filepath.Join(stateDir, "state.sqlite3"),
		BlockedNetworks: blockedNetworks,
		// A fetch that gets past the checks and finds nothing at the address
		// fails within this time at the latest.
		UpstreamFetchTimeout: 2 * time.Second,
		// With no connection slots, the fetch would fail before connecting.
		UpstreamConnectionsPerHost: config.DefaultUpstreamConnectionsPerHost,
		UpstreamConnections:        config.DefaultUpstreamConnections,
		DownstreamTimeout:          config.DefaultDownstreamTimeout,
	}

	var (
		h  *handlers.Handlers
		mw *middleware.Middleware
	)

	app := fxtest.New(t,
		fx.Supply(cfg),
		fx.Provide(
			globals.New,
			logger.New,
			database.New,
			healthcheck.New,
			handlers.New,
			middleware.New,
			func() prometheus.Registerer { return prometheus.NewRegistry() },
		),
		fx.Provide(providers.Constructors()...),
		fx.Populate(&h, &mw),
	)
	app.RequireStart()
	t.Cleanup(app.RequireStop)

	s := &Server{config: cfg, mw: mw, h: h}
	s.SetupRoutes()

	return s
}

// getImage sends a GET for target to s and fails unless it answers 200.
func getImage(t *testing.T, s *Server, target string) *httptest.ResponseRecorder {
	t.Helper()

	rec := httptest.NewRecorder()
	s.ServeHTTP(rec, httptest.NewRequestWithContext(
		t.Context(), http.MethodGet, target, nil))
	t.Logf("GET %s: %d, X-Pixa-Cache %s",
		target, rec.Code, rec.Header().Get("X-Pixa-Cache"))

	if rec.Code != http.StatusOK {
		t.Fatalf("GET %s status = %d, want %d; body %s",
			target, rec.Code, http.StatusOK, rec.Body.String())
	}

	return rec
}

// checkSourceCached checks that source is stored in cache/sources under the
// source key of upstreamHost and /photo.png, and recorded under that key in
// source_content.
func checkSourceCached(t *testing.T, db *sql.DB, stateDir string, source []byte) {
	t.Helper()

	key := string(imagerequest.Source{Host: upstreamHost, Path: "/photo.png"}.Key())

	checkFile(t, filepath.Join(stateDir, "cache", "sources", key[0:2], key[2:4], key),
		source)

	var rows int

	err := db.QueryRowContext(t.Context(),
		"SELECT COUNT(*) FROM source_content WHERE source_key = ?", key,
	).Scan(&rows)
	if err != nil || rows != 1 {
		t.Errorf("source_content rows for the source = %d (error %v), want 1",
			rows, err)
	}
}

// checkVariantCached checks that the converted image served is recorded in
// variant_content, and stored under its cache key in cache/variants.
func checkVariantCached(t *testing.T, db *sql.DB, stateDir string, served []byte) {
	t.Helper()

	var cacheKey string

	err := db.QueryRowContext(t.Context(),
		"SELECT cache_key FROM variant_content",
	).Scan(&cacheKey)
	if err != nil {
		t.Fatalf("variant_content row: %v", err)
	}

	checkFile(t, filepath.Join(stateDir, "cache", "variants",
		cacheKey[0:2], cacheKey[2:4], cacheKey), served)
}

// checkFile checks that the file at path holds want.
func checkFile(t *testing.T, path string, want []byte) {
	t.Helper()

	//nolint:gosec // G304: a path under the test's state directory
	got, err := os.ReadFile(path)
	if err != nil {
		t.Errorf("reading %s: %v", path, err)

		return
	}

	if !bytes.Equal(got, want) {
		t.Errorf("%s holds %d bytes that are not the %d expected",
			path, len(got), len(want))
	}
}

// encodeTestPNG returns an opaque width x height PNG of one color.
func encodeTestPNG(t *testing.T, width, height int) []byte {
	t.Helper()

	img := image.NewRGBA(image.Rect(0, 0, width, height))
	for y := range height {
		for x := range width {
			img.Set(x, y, color.RGBA{R: 200, G: 40, B: 40, A: 255})
		}
	}

	var buf bytes.Buffer

	err := png.Encode(&buf, img)
	if err != nil {
		t.Fatalf("encoding the test PNG: %v", err)
	}

	return buf.Bytes()
}