smallwebwaf/internal/proxy/lookup_test.go

package proxy_test

import (
	"net/http"
	"net/netip"
	"slices"
	"sync"
	"testing"
	"time"

	"sneak.berlin/go/smallwebwaf/internal/alerts"
	"sneak.berlin/go/smallwebwaf/internal/lookup"
	"sneak.berlin/go/smallwebwaf/internal/requestlog"
)

// asnAndCountry is what a lookup gives a client: its AS number, AS name
// and country.
type asnAndCountry struct{ asn, asName, country string }

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

	// The stand-in for GeoJS answers only once released. A request that
	// waited for it would wait an hour, and get no answer within
	// waitLimit.
	geojsURL, asked, release := startHeldGeoJS(t)
	s, _, server := startWithClock(t, geojsURL, map[string]string{
		lookupTimeout:      "1h",
		rateLimitPerMinute: "1",
	})

	// fromDE's second request breaks the limit and bans it, and fromKP
	// comes too. None waits for GeoJS.
	for _, line := range []logLine{
		s.get(fromDE, http.StatusOK, requestlog.ActionForward),
		s.get(fromDE, http.StatusForbidden, requestlog.ActionRateLimited),
		s.get(fromKP, http.StatusOK, requestlog.ActionForward),
	} {
		got := asnAndCountry{line.ASN, line.ASName, line.Country}
		if got != (asnAndCountry{}) {
			t.Errorf("log line has %+v before GeoJS answered, want nothing", got)
		}
	}

	// Once GeoJS answers, each answer reaches the client's history, and
	// fromDE's reaches the notes of its ban.
	release()

	netblock := netip.MustParsePrefix(fromDE + "/32")

	waitUntil(func() bool {
		return historyOf(t, server, fromDE).ASN != "" &&
			historyOf(t, server, fromKP).ASN != "" &&
			server.Ledger.Bans(netblock)[0].Notes.ASN != ""
	})

	de := asnAndCountry{asnDE, asNameDE, "DE"}

	for addr, want := range map[string]asnAndCountry{
		fromDE: de, fromKP: {asnKP, asNameKP, "KP"},
	} {
		h := historyOf(t, server, addr)
		if got := (asnAndCountry{h.ASN, h.ASName, h.Country}); got != want {
			t.Errorf("%s's history has %+v, want %+v", addr, got, want)
		}
	}

	notes := server.Ledger.Bans(netblock)[0].Notes
	if got := (asnAndCountry{notes.ASN, notes.ASName, notes.Country}); got != de {
		t.Errorf("the ban's notes have %+v, want %+v", got, de)
	}

	// GeoJS was asked about each client once, fromKP after fromDE, whose
	// request was under way when fromKP came.
	if got := asked(); !slices.Equal(got, []string{fromDE, fromKP}) {
		t.Errorf("GeoJS was asked about %v, want %s and %s", got, fromDE, fromKP)
	}
}

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

	geojsURL, _ := startGeoJS(t)
	app := startApp(t, func(http.ResponseWriter, *http.Request) {})
	clk := &clock{now: time.Date(2026, 10, 6, 0, 0, 0, 0, time.UTC)}
	addr, out, server, queue := startProxyWithAlerts(t, app.URL, geojsURL, clk.Now,
		map[string]string{
			trustedProxies:     trustLocalhost,
			alertWebhookURL:    "https://alerts.example/smallwebwaf",
			rateLimitPerMinute: "1",
		})
	s := &sender{t: t, addr: addr, out: out}

	// The answer is kept before the requests, so GeoJS is not asked, and
	// gives no answer of its own.
	netblock := netip.MustParsePrefix(fromDE + "/32")
	server.GeoJS.Load([]lookup.Answer{{
		Client: netblock, ASN: asnDE, ASName: asNameDE, Country: "DE",
		Answered: clk.Now(), Used: clk.Now(),
	}})

	want := asnAndCountry{asnDE, asNameDE, "DE"}

	for _, line := range []logLine{
		s.get(fromDE, http.StatusOK, requestlog.ActionForward),
		s.get(fromDE, http.StatusForbidden, requestlog.ActionRateLimited),
	} {
		if got := (asnAndCountry{line.ASN, line.ASName, line.Country}); got != want {
			t.Errorf("log line has %+v, want %+v", got, want)
		}
	}

	h := historyOf(t, server, fromDE)
	if got := (asnAndCountry{h.ASN, h.ASName, h.Country}); got != want ||
		!h.LookedUp.Equal(clk.Now()) {
		t.Errorf("history has %+v, looked up at %s; want %+v, at %s",
			got, h.LookedUp, want, clk.Now())
	}

	notes := server.Ledger.Bans(netblock)[0].Notes
	if got := (asnAndCountry{notes.ASN, notes.ASName, notes.Country}); got != want {
		t.Errorf("the ban's notes have %+v, want %+v", got, want)
	}

	waiting := queue.Snapshot().Waiting[alerts.DestinationWebhook]
	if len(waiting) != 1 {
		t.Fatalf("alerts waiting %+v, want the ban's alone", waiting)
	}

	alert := waiting[0]
	if got := (asnAndCountry{alert.ASN, alert.ASName, alert.Country}); got != want {
		t.Errorf("the ban's alert has %+v, want %+v", got, want)
	}
}

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

	geojsURL, asked := startGeoJS(t)
	s, clk, server := startWithClock(t, geojsURL, map[string]string{lookupSource: "off"})

	// Even an answer kept from before is not used.
	server.GeoJS.Load([]lookup.Answer{{
		Client: netip.MustParsePrefix(fromDE + "/32"), ASN: asnDE, ASName: asNameDE,
		Country: "DE", Answered: clk.Now(), Used: clk.Now(),
	}})

	for _, from := range []string{fromDE, fromKP} {
		line := s.get(from, http.StatusOK, requestlog.ActionForward)

		got := asnAndCountry{line.ASN, line.ASName, line.Country}
		if got != (asnAndCountry{}) {
			t.Errorf("log line has %+v, want nothing", got)
		}
	}

	if h := historyOf(t, server, fromDE); h.ASN != "" || !h.LookedUp.IsZero() {
		t.Errorf("history has %q, looked up at %s, want no lookup", h.ASN, h.LookedUp)
	}

	if len(asked()) != 0 {
		t.Errorf("GeoJS was asked about %v, want nothing", asked())
	}
}

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

	var (
		mu  sync.Mutex
		got [][2][]string // each request's X-Client-ASN and X-Client-Country
	)

	app := startApp(t, func(_ http.ResponseWriter, r *http.Request) {
		mu.Lock()
		defer mu.Unlock()

		got = append(got, [2][]string{
			r.Header.Values("X-Client-Asn"), r.Header.Values("X-Client-Country"),
		})
	})
	geojsURL, _ := startGeoJS(t)
	addr, out, _ := startProxyWithClock(t, app.URL, geojsURL, time.Now, map[string]string{
		trustedProxies:   trustLocalhost,
		addLookupHeaders: "true",
	})
	s := &sender{t: t, addr: addr, out: out}

	// Each client sends headers of its own. fromDE's first request waits
	// for its answer, which the app is passed; unplaced has none to pass,
	// and a client on a private address is not looked up.
	own := "X-Client-ASN: AS1\r\nX-Client-Country: KP"
	for _, from := range []string{fromDE, unplaced, "10.0.0.8"} {
		s.requestWithHeader(from, "/", own, http.StatusOK, requestlog.ActionForward)
	}

	mu.Lock()
	defer mu.Unlock()

	want := [][2][]string{{{asnDE}, {"DE"}}, {nil, nil}, {nil, nil}}
	if !slices.EqualFunc(got, want, func(a, b [2][]string) bool {
		return slices.Equal(a[0], b[0]) && slices.Equal(a[1], b[1])
	}) {
		t.Errorf("the app was passed %v, want %v", got, want)
	}
}

// waitUntil waits until done reports true, for at most waitLimit.
func waitUntil(done func() bool) {
	deadline := time.Now().Add(waitLimit)
	for !done() && time.Now().Before(deadline) {
		time.Sleep(pollInterval)
	}
}