Maven/internal/webfetch/webfetch_test.go

package webfetch

import (
	"bytes"
	"context"
	"errors"
	"fmt"
	"io"
	"net"
	"net/http"
	"net/http/httptest"
	"strings"
	"testing"
	"time"
)

// The limits in this package are the reason a crawler is allowed to exist on
// this box at all, so each one has a test that fails loudly if it is removed.

func TestPrivateAddressesAreRefused(t *testing.T) {
	// The wireguard range (10.42.0.0/24), the LAN (192.168.1.0/24) and the
	// cloud metadata address are the three that matter here; the rest come
	// along for free.
	for _, s := range []string{
		"127.0.0.1", "127.1.2.3", "10.42.0.7", "10.0.0.5", "192.168.1.104",
		"172.16.4.4", "169.254.169.254", "100.64.1.1", "0.0.0.0",
		"::1", "fc00::1", "fd12:3456::1", "fe80::1",
	} {
		if !IsPrivateIP(net.ParseIP(s)) {
			t.Errorf("IsPrivateIP(%s) = false, want true", s)
		}
	}
	for _, s := range []string{"8.8.8.8", "1.1.1.1", "93.184.216.34", "2606:2800:220:1::1"} {
		if IsPrivateIP(net.ParseIP(s)) {
			t.Errorf("IsPrivateIP(%s) = true, want false", s)
		}
	}
}

func TestGetRefusesPrivateLiteral(t *testing.T) {
	f := New(Config{})
	for _, u := range []string{
		"http://127.0.0.1:8034/search",
		"http://10.42.0.1/",
		"http://192.168.1.104/dash",
		"http://[::1]:9100/mcp",
	} {
		if _, err := f.Get(context.Background(), u); !errors.Is(err, ErrPrivate) {
			t.Errorf("Get(%s) error = %v, want ErrPrivate", u, err)
		}
	}
}

// A hostname that resolves into private space must fail too — that is the
// rebinding case, and it is why the check lives in the dialer.
func TestGetRefusesPrivateResolution(t *testing.T) {
	f := New(Config{})
	if _, err := f.Get(context.Background(), "http://localhost:8034/"); !errors.Is(err, ErrPrivate) {
		t.Fatalf("Get(localhost) error = %v, want ErrPrivate", err)
	}
}

func TestGetRefusesNonHTTPSchemes(t *testing.T) {
	f := New(Config{})
	for _, u := range []string{"file:///etc/passwd", "ftp://example.com/x", "gopher://example.com"} {
		if _, err := f.Get(context.Background(), u); !errors.Is(err, ErrScheme) {
			t.Errorf("Get(%s) error = %v, want ErrScheme", u, err)
		}
	}
}

// testFetcher — a fetcher pointed at an httptest server, which necessarily
// listens on loopback. AllowPrivate is the test-only escape hatch.
func testFetcher(t *testing.T, cfg Config) *Fetcher {
	t.Helper()
	cfg.AllowPrivate = true
	if cfg.HostInterval == 0 {
		cfg.HostInterval = time.Nanosecond
	}
	return New(cfg)
}

func TestAllowAndDenyLists(t *testing.T) {
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.Write([]byte("ok"))
	}))
	defer srv.Close()

	f := testFetcher(t, Config{AllowHosts: []string{"example.com"}})
	if _, err := f.Get(context.Background(), srv.URL); !errors.Is(err, ErrBlocked) {
		t.Fatalf("off-allowlist host: error = %v, want ErrBlocked", err)
	}
	f = testFetcher(t, Config{DenyHosts: []string{"127.0.0.1"}})
	if _, err := f.Get(context.Background(), srv.URL); !errors.Is(err, ErrBlocked) {
		t.Fatalf("denied host: error = %v, want ErrBlocked", err)
	}
	f = testFetcher(t, Config{AllowHosts: []string{"127.0.0.1"}})
	if _, err := f.Get(context.Background(), srv.URL); err != nil {
		t.Fatalf("allowlisted host: %v", err)
	}
}

func TestHostMatchesSubdomains(t *testing.T) {
	pats := []string{"example.com", "*.news.org"}
	for _, h := range []string{"example.com", "news.example.com", "a.b.example.com", "news.org", "feeds.news.org"} {
		if !HostMatches(h, pats) {
			t.Errorf("HostMatches(%q) = false, want true", h)
		}
	}
	for _, h := range []string{"notexample.com", "example.com.evil.net", "org"} {
		if HostMatches(h, pats) {
			t.Errorf("HostMatches(%q) = true, want false", h)
		}
	}
}

func TestSizeCap(t *testing.T) {
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.Write([]byte(strings.Repeat("x", 5000)))
	}))
	defer srv.Close()

	f := testFetcher(t, Config{MaxBytes: 100})
	if _, err := f.Get(context.Background(), srv.URL); !errors.Is(err, ErrTooLarge) {
		t.Fatalf("error = %v, want ErrTooLarge", err)
	}
	f = testFetcher(t, Config{MaxBytes: 6000})
	resp, err := f.Get(context.Background(), srv.URL)
	if err != nil {
		t.Fatalf("under the cap: %v", err)
	}
	if len(resp.Body) != 5000 {
		t.Fatalf("body = %d bytes, want 5000", len(resp.Body))
	}
}

func TestRedirectCap(t *testing.T) {
	var srv *httptest.Server
	srv = httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		http.Redirect(w, r, srv.URL+"/again", http.StatusFound)
	}))
	defer srv.Close()

	f := testFetcher(t, Config{MaxRedirects: 2})
	if _, err := f.Get(context.Background(), srv.URL); !errors.Is(err, ErrRedirects) {
		t.Fatalf("error = %v, want ErrRedirects", err)
	}
}

// A redirect off the allowlist is the interesting redirect: the first hop is
// permitted, the second must not be.
func TestRedirectRecheckedAgainstDenylist(t *testing.T) {
	target := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.Write([]byte("secret"))
	}))
	defer target.Close()
	hop := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		http.Redirect(w, r, target.URL, http.StatusFound)
	}))
	defer hop.Close()

	// Reach the hop under the name "localhost" and allow only that name; the
	// redirect lands on the same box under its literal address, which the
	// allowlist does not cover. Without the CheckRedirect hook this fetch
	// succeeds and returns "secret".
	f := testFetcher(t, Config{AllowHosts: []string{"localhost"}})
	viaName := strings.Replace(hop.URL, "127.0.0.1", "localhost", 1)
	if _, err := f.Get(context.Background(), viaName); !errors.Is(err, ErrBlocked) {
		t.Fatalf("error = %v, want ErrBlocked", err)
	}
}

func TestPerHostRateLimit(t *testing.T) {
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.Write([]byte("ok"))
	}))
	defer srv.Close()

	f := testFetcher(t, Config{HostInterval: 60 * time.Millisecond})
	start := time.Now()
	for i := 0; i < 3; i++ {
		if _, err := f.Get(context.Background(), srv.URL); err != nil {
			t.Fatalf("request %d: %v", i, err)
		}
	}
	if elapsed := time.Since(start); elapsed < 120*time.Millisecond {
		t.Fatalf("three requests took %s, want at least 120ms of spacing", elapsed)
	}
}

func TestRateLimitHonoursContext(t *testing.T) {
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {}))
	defer srv.Close()

	f := testFetcher(t, Config{HostInterval: 10 * time.Second})
	if _, err := f.Get(context.Background(), srv.URL); err != nil {
		t.Fatalf("first request: %v", err)
	}
	ctx, cancel := context.WithTimeout(context.Background(), 20*time.Millisecond)
	defer cancel()
	if _, err := f.Get(ctx, srv.URL); !errors.Is(err, context.DeadlineExceeded) {
		t.Fatalf("error = %v, want DeadlineExceeded", err)
	}
}

func TestNon2xxIsAnError(t *testing.T) {
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		http.Error(w, "nope", http.StatusInternalServerError)
	}))
	defer srv.Close()
	f := testFetcher(t, Config{})
	if _, err := f.Get(context.Background(), srv.URL); !errors.Is(err, ErrStatus) {
		t.Fatalf("error = %v, want ErrStatus", err)
	}
}

func TestNon2xxCarriesTheCodeAndBeatsTheSizeCap(t *testing.T) {
	// A big error page used to be read in full and reported as ErrTooLarge,
	// which names the size and hides the 503. The status is checked first now,
	// and the code survives for a caller that has to tell 5xx from 404.
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.WriteHeader(http.StatusServiceUnavailable)
		w.Write(bytes.Repeat([]byte("x"), 5000))
	}))
	defer srv.Close()
	f := testFetcher(t, Config{MaxBytes: 100})
	_, err := f.Get(context.Background(), srv.URL)
	if !errors.Is(err, ErrStatus) || errors.Is(err, ErrTooLarge) {
		t.Fatalf("error = %v, want ErrStatus and not ErrTooLarge", err)
	}
	var se *StatusError
	if !errors.As(err, &se) || se.Code != http.StatusServiceUnavailable {
		t.Fatalf("error = %v, want a StatusError carrying 503", err)
	}
}

func TestUserAgentIsSent(t *testing.T) {
	got := make(chan string, 1)
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		got <- r.Header.Get("User-Agent")
	}))
	defer srv.Close()
	f := testFetcher(t, Config{UserAgent: "Maven/test"})
	if _, err := f.Get(context.Background(), srv.URL); err != nil {
		t.Fatal(err)
	}
	if ua := <-got; ua != "Maven/test" {
		t.Fatalf("user-agent = %q", ua)
	}
}

func TestPostSendsBodyAndHeaders(t *testing.T) {
	type seen struct {
		method, ctype, accept, ua, custom string
		body                              []byte
	}
	ch := make(chan seen, 1)
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		b, _ := io.ReadAll(r.Body)
		ch <- seen{r.Method, r.Header.Get("Content-Type"), r.Header.Get("Accept"),
			r.Header.Get("User-Agent"), r.Header.Get("X-Thing"), b}
		w.Header().Set("Mcp-Session-Id", "sess-9")
		w.Header().Set("Content-Type", "application/json")
		_, _ = w.Write([]byte(`{"ok":true}`))
	}))
	defer srv.Close()
	f := testFetcher(t, Config{UserAgent: "Maven/test"})
	resp, err := f.Post(context.Background(), srv.URL, "application/json",
		[]byte(`{"jsonrpc":"2.0"}`), map[string]string{"Accept": "text/event-stream", "X-Thing": "1"})
	if err != nil {
		t.Fatal(err)
	}
	if string(resp.Body) != `{"ok":true}` {
		t.Fatalf("body = %q", resp.Body)
	}
	if resp.Header["Mcp-Session-Id"] != "sess-9" {
		t.Fatalf("response headers not surfaced: %+v", resp.Header)
	}
	s := <-ch
	if s.method != http.MethodPost {
		t.Fatalf("method = %s", s.method)
	}
	if string(s.body) != `{"jsonrpc":"2.0"}` {
		t.Fatalf("request body = %q", s.body)
	}
	if s.ctype != "application/json" {
		t.Fatalf("content-type = %q", s.ctype)
	}
	if s.accept != "text/event-stream" || s.custom != "1" {
		t.Fatalf("caller headers dropped: %+v", s)
	}
	if s.ua != "Maven/test" {
		t.Fatalf("user-agent = %q — a caller must not be able to override it", s.ua)
	}
}

// The whole point of routing MCP through webfetch: a POST is guarded exactly
// like a GET. A body does not buy a caller a way onto the LAN.
func TestPostRefusesPrivateAddress(t *testing.T) {
	f := New(Config{}) // no AllowPrivate
	_, err := f.Post(context.Background(), "http://127.0.0.1:9100/mcp", "application/json", []byte(`{}`), nil)
	if !errors.Is(err, ErrPrivate) {
		t.Fatalf("error = %v, want ErrPrivate", err)
	}
}

func TestPostRefusesNonHTTPScheme(t *testing.T) {
	f := New(Config{})
	if _, err := f.Post(context.Background(), "file:///etc/passwd", "application/json", nil, nil); !errors.Is(err, ErrScheme) {
		t.Fatalf("error = %v, want ErrScheme", err)
	}
}

func TestPostObeysDenylist(t *testing.T) {
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {}))
	defer srv.Close()
	f := testFetcher(t, Config{DenyHosts: []string{"127.0.0.1"}})
	if _, err := f.Post(context.Background(), srv.URL, "application/json", []byte(`{}`), nil); !errors.Is(err, ErrBlocked) {
		t.Fatalf("error = %v, want ErrBlocked", err)
	}
}

// f.last used to hold one entry per host ever dialed, for the life of the
// process. A host whose last dial is older than HostInterval cannot delay
// anything, so it is dropped once the map is worth walking.
func TestHostRateMapIsPruned(t *testing.T) {
	f := New(Config{HostInterval: time.Minute, AllowPrivate: true})
	stale := time.Now().Add(-time.Hour)
	for i := 0; i < pruneHostsAbove*2; i++ {
		f.last[fmt.Sprintf("h%d.example", i)] = stale
	}

	// One real turn is what triggers the sweep.
	if err := f.waitTurn(context.Background(), "fresh.example"); err != nil {
		t.Fatal(err)
	}
	if len(f.last) != 1 {
		t.Fatalf("len(f.last) = %d after the sweep, want 1 (only the host just dialed)", len(f.last))
	}
	if _, ok := f.last["fresh.example"]; !ok {
		t.Fatal("the host just dialed was pruned, so its own rate limit is lost")
	}

	// A host inside the interval is kept: pruning must not hand out a free turn.
	f.last["recent.example"] = time.Now()
	for i := 0; i < pruneHostsAbove*2; i++ {
		f.last[fmt.Sprintf("g%d.example", i)] = stale
	}
	if err := f.waitTurn(context.Background(), "other.example"); err != nil {
		t.Fatal(err)
	}
	if _, ok := f.last["recent.example"]; !ok {
		t.Fatal("a host dialed inside HostInterval was pruned")
	}
}