pixa/internal/httpfetcher/fetch_internal_test.go

package httpfetcher

import (
	"context"
	"errors"
	"fmt"
	"io"
	"net"
	"net/http"
	"net/http/httptest"
	"slices"
	"strconv"
	"strings"
	"sync"
	"testing"
	"time"

	"sneak.berlin/go/pixa/internal/imagerequest"
)

// testPublicHost is a TEST-NET-1 (RFC 5737) literal. isPrivateIP treats it as
// public, so validateURL and the redirect check accept it with no DNS lookup,
// while the recording dialer routes it to the local httptest server. The
// address is reserved for documentation and is never routed on the network.
const testPublicHost = "192.0.2.10"

// imagePayload is the body served by the fake upstream's image route.
const imagePayload = "fake-jpeg-bytes"

// errUnexpectedDial reports a dial to any host other than testPublicHost, which
// would mean SSRF protection let a forbidden target reach the transport.
var errUnexpectedDial = errors.New("unexpected dial target")

// upstreamSource returns the source at path on the fake public host.
func upstreamSource(path string) imagerequest.Source {
	return imagerequest.Source{Host: testPublicHost, Path: path}
}

// recordingDialer records every address the transport asks it to dial and
// routes connections for testPublicHost to a real local server, so the SSRF
// checks run against a public-looking host while bytes go to httptest.
type recordingDialer struct {
	target string

	mu     sync.Mutex
	dialed []string
}

func (d *recordingDialer) dialContext(
	ctx context.Context,
	network, addr string,
) (net.Conn, error) {
	d.mu.Lock()
	d.dialed = append(d.dialed, addr)
	d.mu.Unlock()

	host, _, err := net.SplitHostPort(addr)
	if err != nil {
		return nil, err
	}

	if host != testPublicHost {
		return nil, fmt.Errorf("%w: %s", errUnexpectedDial, addr)
	}

	var dialer net.Dialer

	return dialer.DialContext(ctx, network, d.target)
}

// dialedAddrs returns a copy of the addresses the dialer was asked to reach.
func (d *recordingDialer) dialedAddrs() []string {
	d.mu.Lock()
	defer d.mu.Unlock()

	return slices.Clone(d.dialed)
}

// startUpstream launches a fake upstream with the routes the fetch tests
// exercise and stops it when the test finishes.
func startUpstream(t *testing.T) *httptest.Server {
	t.Helper()

	mux := http.NewServeMux()
	mux.HandleFunc("/image", func(w http.ResponseWriter, _ *http.Request) {
		w.Header().Set("Content-Type", imagerequest.FormatJPEG.MIMEType())
		_, _ = io.WriteString(w, imagePayload)
	})
	mux.HandleFunc("/status/500", func(w http.ResponseWriter, _ *http.Request) {
		w.WriteHeader(http.StatusInternalServerError)
	})
	mux.HandleFunc("/html", func(w http.ResponseWriter, _ *http.Request) {
		w.Header().Set("Content-Type", "text/html; charset=utf-8")
		_, _ = io.WriteString(w, "<html></html>")
	})
	mux.HandleFunc("/redirect/private", func(w http.ResponseWriter, r *http.Request) {
		http.Redirect(w, r, "http://169.254.169.254/latest/meta-data/", http.StatusFound)
	})
	mux.HandleFunc("/redirect/public", func(w http.ResponseWriter, r *http.Request) {
		http.Redirect(w, r, "/image", http.StatusFound)
	})
	mux.HandleFunc("/redirect/chain", func(w http.ResponseWriter, r *http.Request) {
		http.Redirect(w, r, "/redirect/hop", http.StatusFound)
	})
	mux.HandleFunc("/redirect/hop", func(w http.ResponseWriter, r *http.Request) {
		http.Redirect(w, r, "/image", http.StatusFound)
	})

	srv := httptest.NewServer(mux)
	t.Cleanup(srv.Close)

	return srv
}

// newServerFetcher builds a fetcher whose transport routes testPublicHost to
// srv, leaving the real SSRF validation and redirect checks in place.
func newServerFetcher(
	t *testing.T,
	srv *httptest.Server,
	cfg *Config,
) (*HTTPFetcher, *recordingDialer) {
	t.Helper()

	if cfg == nil {
		cfg = &Config{
			Timeout:               30 * time.Second,
			MaxResponseSize:       50 << 20,
			MaxConnectionsPerHost: 20,
			MaxConnections:        64,
		}
	}

	cfg.AllowHTTP = true

	f := New(cfg)

	transport, ok := f.client.Transport.(*http.Transport)
	if !ok {
		t.Fatalf("transport is %T, want *http.Transport", f.client.Transport)
	}

	dialer := &recordingDialer{target: srv.Listener.Addr().String()}
	transport.DialContext = dialer.dialContext

	return f, dialer
}

// testContext returns a context cancelled when the test ends, bounding any
// fetch that would otherwise block on a leaked semaphore slot.
func testContext(t *testing.T) context.Context {
	t.Helper()

	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	t.Cleanup(cancel)

	return ctx
}

// fetchImage fetches path from the fake upstream and fails on error.
func fetchImage(t *testing.T, f *HTTPFetcher, path string) *FetchResult {
	t.Helper()

	res, err := f.Fetch(testContext(t), upstreamSource(path), "", "")
	if err != nil {
		t.Fatalf("Fetch(%s) error = %v", path, err)
	}

	return res
}

// fetchExpectError fetches path and fails unless Fetch returns an error.
func fetchExpectError(t *testing.T, f *HTTPFetcher, path string) error {
	t.Helper()

	res, err := f.Fetch(testContext(t), upstreamSource(path), "", "")
	if err == nil {
		_ = res.Content.Close()

		t.Fatalf("Fetch(%s) = nil error, want an error", path)
	}

	return err
}

// fetchBody fetches path and returns the fully read, closed response body.
func fetchBody(t *testing.T, f *HTTPFetcher, path string) string {
	t.Helper()

	res := fetchImage(t, f, path)
	defer func() { _ = res.Content.Close() }()

	data, err := io.ReadAll(res.Content)
	if err != nil {
		t.Fatalf("read body: %v", err)
	}

	return string(data)
}

// semLen reports how many per-host semaphore slots are currently held.
func semLen(f *HTTPFetcher, host string) int {
	f.hostSemMu.Lock()
	defer f.hostSemMu.Unlock()

	sem, ok := f.hostSems[host]
	if !ok {
		return 0
	}

	return len(sem.slots)
}

// hostSemCount reports how many hosts have a semaphore in hostSems.
func hostSemCount(f *HTTPFetcher) int {
	f.hostSemMu.Lock()
	defer f.hostSemMu.Unlock()

	return len(f.hostSems)
}

func TestFetchRedirectToPrivateIPBlocked(t *testing.T) {
	t.Parallel()

	srv := startUpstream(t)
	f, dialer := newServerFetcher(t, srv, nil)

	_, err := f.Fetch(testContext(t), upstreamSource("/redirect/private"), "", "")
	if !errors.Is(err, ErrSSRFBlocked) {
		t.Fatalf("Fetch() error = %v, want ErrSSRFBlocked", err)
	}

	for _, addr := range dialer.dialedAddrs() {
		if strings.Contains(addr, "169.254.169.254") {
			t.Errorf("dialer connected to the private redirect target: %s", addr)
		}
	}
}

func TestFetchRedirectToPublicSucceeds(t *testing.T) {
	t.Parallel()

	srv := startUpstream(t)
	f, _ := newServerFetcher(t, srv, nil)

	if body := fetchBody(t, f, "/redirect/public"); body != imagePayload {
		t.Errorf("body = %q, want %q", body, imagePayload)
	}
}

func TestFetchRedirectChainSucceeds(t *testing.T) {
	t.Parallel()

	srv := startUpstream(t)
	f, _ := newServerFetcher(t, srv, nil)

	if body := fetchBody(t, f, "/redirect/chain"); body != imagePayload {
		t.Errorf("body = %q, want %q", body, imagePayload)
	}
}

func TestFetchRejectsNon2xx(t *testing.T) {
	t.Parallel()

	srv := startUpstream(t)
	f, _ := newServerFetcher(t, srv, nil)

	err := fetchExpectError(t, f, "/status/500")
	if !errors.Is(err, ErrUpstreamError) {
		t.Fatalf("Fetch() error = %v, want ErrUpstreamError", err)
	}
}

func TestFetchRejectsDisallowedContentType(t *testing.T) {
	t.Parallel()

	srv := startUpstream(t)
	f, _ := newServerFetcher(t, srv, nil)

	err := fetchExpectError(t, f, "/html")
	if !errors.Is(err, ErrInvalidContentType) {
		t.Fatalf("Fetch() error = %v, want ErrInvalidContentType", err)
	}
}

func TestFetchMaxResponseSizeEnforced(t *testing.T) {
	t.Parallel()

	srv := startUpstream(t)

	cfg := &Config{
		Timeout:               30 * time.Second,
		MaxResponseSize:       8,
		MaxConnectionsPerHost: 20,
		MaxConnections:        64,
	}

	f, _ := newServerFetcher(t, srv, cfg)

	res := fetchImage(t, f, "/image")
	defer func() { _ = res.Content.Close() }()

	data, err := io.ReadAll(res.Content)
	if !errors.Is(err, ErrResponseTooLarge) {
		t.Fatalf("read error = %v, want ErrResponseTooLarge", err)
	}

	if int64(len(data)) > cfg.MaxResponseSize {
		t.Errorf("read %d bytes, exceeds limit %d", len(data), cfg.MaxResponseSize)
	}
}

// fetchWithMaxResponseSize serves handler as the fake upstream and fetches
// its /image route with MaxResponseSize set to maxResponseSize.
func fetchWithMaxResponseSize(
	t *testing.T,
	maxResponseSize int64,
	handler http.HandlerFunc,
) *FetchResult {
	t.Helper()

	srv := httptest.NewServer(handler)
	t.Cleanup(srv.Close)

	cfg := &Config{
		Timeout:               30 * time.Second,
		MaxResponseSize:       maxResponseSize,
		MaxConnectionsPerHost: 20,
		MaxConnections:        64,
	}

	f, _ := newServerFetcher(t, srv, cfg)

	res := fetchImage(t, f, "/image")
	t.Cleanup(func() { _ = res.Content.Close() })

	return res
}

// TestFetchContentLengthBodyAtMaxResponseSize checks that a body sent with a
// Content-Length is read whole at exactly MaxResponseSize bytes and refused
// at one byte more.
func TestFetchContentLengthBodyAtMaxResponseSize(t *testing.T) {
	t.Parallel()

	const maxResponseSize = 8

	for _, tc := range []struct {
		size    int
		wantErr error
	}{
		{maxResponseSize, nil},
		{maxResponseSize + 1, ErrResponseTooLarge},
	} {
		t.Run(strconv.Itoa(tc.size)+" bytes", func(t *testing.T) {
			t.Parallel()

			body := strings.Repeat("x", tc.size)

			res := fetchWithMaxResponseSize(t, maxResponseSize,
				func(w http.ResponseWriter, _ *http.Request) {
					w.Header().Set("Content-Type", imagerequest.FormatJPEG.MIMEType())
					w.Header().Set("Content-Length", strconv.Itoa(tc.size))
					_, _ = io.WriteString(w, body)
				})

			if res.ContentLength != int64(tc.size) {
				t.Fatalf("ContentLength = %d, want %d", res.ContentLength, tc.size)
			}

			data, err := io.ReadAll(res.Content)
			t.Logf("body of %d bytes: read %d, error %v", tc.size, len(data), err)

			if !errors.Is(err, tc.wantErr) {
				t.Fatalf("read error = %v, want %v", err, tc.wantErr)
			}

			if err == nil && string(data) != body {
				t.Errorf("read %q, want %q", data, body)
			}
		})
	}
}

// TestFetchChunkedBodyAtMaxResponseSize checks that a chunked body is read
// whole at exactly MaxResponseSize bytes and refused at one byte more. The
// upstream sends the end of the body only once the test has read
// MaxResponseSize bytes, so the read after them comes before the end of the
// body has arrived, as it can from a real upstream.
func TestFetchChunkedBodyAtMaxResponseSize(t *testing.T) {
	t.Parallel()

	const maxResponseSize = 8

	for _, tc := range []struct {
		size    int
		wantErr error
	}{
		{maxResponseSize, nil},
		{maxResponseSize + 1, ErrResponseTooLarge},
	} {
		t.Run(strconv.Itoa(tc.size)+" bytes", func(t *testing.T) {
			t.Parallel()

			body := strings.Repeat("x", tc.size)
			endBody := make(chan struct{})

			res := fetchWithMaxResponseSize(t, maxResponseSize,
				func(w http.ResponseWriter, r *http.Request) {
					w.Header().Set("Content-Type", imagerequest.FormatJPEG.MIMEType())
					_, _ = io.WriteString(w, body)
					// With no Content-Length, flushing sends the body
					// chunked; its end is sent when the handler returns.
					_ = http.NewResponseController(w).Flush()

					select {
					case <-endBody:
					case <-r.Context().Done():
					}
				})

			if res.ContentLength != -1 {
				t.Fatalf("ContentLength = %d, want -1 (chunked)", res.ContentLength)
			}

			head := make([]byte, maxResponseSize)

			_, err := io.ReadFull(res.Content, head)
			if err != nil {
				t.Fatalf("read the first %d bytes: %v", maxResponseSize, err)
			}

			close(endBody)

			rest, err := io.ReadAll(res.Content)
			data := string(head) + string(rest)
			t.Logf("body of %d bytes: read %d, error %v", tc.size, len(data), err)

			if !errors.Is(err, tc.wantErr) {
				t.Fatalf("read error = %v, want %v", err, tc.wantErr)
			}

			if err == nil && data != body {
				t.Errorf("read %q, want %q", data, body)
			}
		})
	}
}

func TestFetchSemaphoreReleasedOnError(t *testing.T) {
	t.Parallel()

	srv := startUpstream(t)

	cfg := &Config{
		Timeout:               30 * time.Second,
		MaxResponseSize:       50 << 20,
		MaxConnectionsPerHost: 1,
		MaxConnections:        64,
	}

	f, _ := newServerFetcher(t, srv, cfg)

	err := fetchExpectError(t, f, "/status/500")
	if !errors.Is(err, ErrUpstreamError) {
		t.Fatalf("Fetch() error = %v, want ErrUpstreamError", err)
	}

	if held := semLen(f, testPublicHost); held != 0 {
		t.Fatalf("semaphore slot leaked after error: %d held", held)
	}

	// One slot per host: this fetch proceeds only if the slot was released.
	res := fetchImage(t, f, "/image")
	_ = res.Content.Close()
}

// assertSlotReleasedByClose fetches an image over a one-slot host, hands the
// open result to consume, and asserts the slot is held before and freed after,
// then that a follow-up fetch can still acquire it.
func assertSlotReleasedByClose(
	t *testing.T,
	consume func(*testing.T, *FetchResult),
) {
	t.Helper()

	srv := startUpstream(t)

	cfg := &Config{
		Timeout:               30 * time.Second,
		MaxResponseSize:       50 << 20,
		MaxConnectionsPerHost: 1,
		MaxConnections:        64,
	}

	f, _ := newServerFetcher(t, srv, cfg)

	res := fetchImage(t, f, "/image")
	if held := semLen(f, testPublicHost); held != 1 {
		t.Fatalf("slot not held while body is open: %d held", held)
	}

	consume(t, res)

	if held := semLen(f, testPublicHost); held != 0 {
		t.Fatalf("slot not released after close: %d held", held)
	}

	next := fetchImage(t, f, "/image")
	_ = next.Content.Close()
}

func TestFetchSemaphoreReleasedOnBodyClose(t *testing.T) {
	t.Parallel()

	assertSlotReleasedByClose(t, func(t *testing.T, res *FetchResult) {
		t.Helper()

		_, err := io.ReadAll(res.Content)
		if err != nil {
			t.Fatalf("read body: %v", err)
		}

		err = res.Content.Close()
		if err != nil {
			t.Fatalf("close body: %v", err)
		}
	})
}

func TestFetchSemaphoreReleasedOnPartialReadClose(t *testing.T) {
	t.Parallel()

	assertSlotReleasedByClose(t, func(t *testing.T, res *FetchResult) {
		t.Helper()

		buf := make([]byte, 1)

		_, err := res.Content.Read(buf)
		if err != nil {
			t.Fatalf("partial read: %v", err)
		}

		err = res.Content.Close()
		if err != nil {
			t.Fatalf("close body: %v", err)
		}
	})
}

// The dial-time re-resolution in dialSSRFSafe is what closes the DNS
// rebinding window: even if validateURL saw a public answer earlier, the
// dialer independently re-checks the address it is about to connect to. A full
// rebinding simulation (a resolver returning public, then private) would mean
// replacing the global net.DefaultResolver with a fake DNS server, which is
// heavyweight and unsafe to mutate under parallel -race tests. The property is
// proven directly here instead: the dialer rejects a private target outright,
// which is exactly the check that fires when a validated host later resolves
// to a private address.
func TestSSRFSafeDialerBlocksPrivateTarget(t *testing.T) {
	t.Parallel()

	for _, addr := range []string{
		"169.254.169.254:80", // link-local (cloud metadata)
		"127.0.0.1:80",       // loopback
		"10.0.0.5:80",        // RFC 1918 private
	} {
		t.Run(addr, func(t *testing.T) {
			t.Parallel()

			_, err := dialSSRFSafe(context.Background(), "tcp", addr, nil)
			if !errors.Is(err, ErrSSRFBlocked) {
				t.Errorf("dialSSRFSafe(%q) = %v, want ErrSSRFBlocked", addr, err)
			}
		})
	}
}

func TestSSRFSafeDialerAllowsPublicTarget(t *testing.T) {
	t.Parallel()

	// A cancelled context makes the dial fail immediately without touching the
	// network; the point is only that a public literal is not SSRF-blocked.
	ctx, cancel := context.WithCancel(context.Background())
	cancel()

	_, err := dialSSRFSafe(ctx, "tcp", testPublicHost+":80", nil)
	if err == nil {
		t.Fatal("expected a dial error for an unreachable public target")
	}

	if errors.Is(err, ErrSSRFBlocked) {
		t.Errorf("public target was SSRF-blocked: %v", err)
	}
}

// TestFetchSchemeFollowsAllowHTTP checks that the fetcher fetches a source
// over https, on port 443, unless AllowHTTP is set, and then over http, on
// port 80. The test server speaks plain HTTP, so only the fetch over http
// succeeds.
func TestFetchSchemeFollowsAllowHTTP(t *testing.T) {
	t.Parallel()

	tests := []struct {
		allowHTTP bool
		wantAddr  string
	}{
		{false, net.JoinHostPort(testPublicHost, "443")},
		{true, net.JoinHostPort(testPublicHost, "80")},
	}

	for _, tt := range tests {
		t.Run(fmt.Sprintf("AllowHTTP %v", tt.allowHTTP), func(t *testing.T) {
			t.Parallel()

			srv := startUpstream(t)
			dialer := &recordingDialer{target: srv.Listener.Addr().String()}

			f := New(&Config{
				Timeout:               30 * time.Second,
				MaxResponseSize:       50 << 20,
				AllowHTTP:             tt.allowHTTP,
				MaxConnectionsPerHost: 20,
				MaxConnections:        64,
				DialContext:           dialer.dialContext,
			})

			res, err := f.Fetch(testContext(t), upstreamSource("/image"), "", "")
			if err == nil {
				_ = res.Content.Close()
			}

			dialed := dialer.dialedAddrs()
			t.Logf("Fetch(): error %v, dialed %v", err, dialed)

			if served := err == nil; served != tt.allowHTTP {
				t.Errorf("served = %v, want %v", served, tt.allowHTTP)
			}

			if len(dialed) != 1 || dialed[0] != tt.wantAddr {
				t.Errorf("dialed %v, want only %s", dialed, tt.wantAddr)
			}
		})
	}
}