smallwebwaf/internal/servewaf/lookup_test.go

package servewaf_test

import (
	"net/http"
	"net/netip"
	"path/filepath"
	"slices"
	"sync"
	"testing"
	"testing/synctest"
	"time"

	"sneak.berlin/go/smallwebwaf/internal/alerts"
	"sneak.berlin/go/smallwebwaf/internal/lookup"
	"sneak.berlin/go/smallwebwaf/internal/lookup/lookuptest"
	"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 }

// fileSource is the SWWAF_LOOKUP_SOURCE that looks clients up in the
// lookup database.
const fileSource = "file"

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 TestRequestThatWaitedForItsLookupIsCountedWhenItGoesOn(t *testing.T) {
	t.Parallel()

	// The country list has a new client's request wait for its answer,
	// which the stand-in for GeoJS gives only once released.
	geojsURL, asked, release := startHeldGeoJS(t)
	app := startApp(t, func(http.ResponseWriter, *http.Request) {})
	clk := &clock{now: time.Date(2026, 10, 6, 0, 0, 58, 0, time.UTC)}
	addr, out, server := startProxyWithClock(t, app.URL, geojsURL, clk.Now,
		map[string]string{
			trustedProxies:  trustLocalhost,
			lookupTimeout:   "1h",
			deniedCountries: "kp",
		})

	requestFromDE := func() *http.Request {
		req := newRequest(t, http.MethodGet, addr, "/", http.NoBody)
		req.Header.Set(forwardedFor, fromDE)

		return req
	}

	// fromDE's first request arrives two seconds before a minute ends,
	// and waits for GeoJS.
	first, httpClient := requestFromDE(), newClient(t)
	firstEnded := make(chan error, 1)

	go func() {
		res, err := httpClient.Do(first)
		if err == nil {
			err = readAnswer(res).err
		}

		firstEnded <- err
	}()

	waitUntil(func() bool { return len(asked()) == 1 })

	// Half a second into the next minute, the answer is kept, as an edit
	// of lookups.json would keep it, while the first request still waits
	// for GeoJS. The second request goes on at once, and is counted.
	clk.advance(2500 * time.Millisecond)
	server.GeoJS.Load([]lookup.Answer{{
		Client: netip.MustParsePrefix(fromDE + "/32"), ASN: asnDE, ASName: asNameDE,
		Country: "DE", Answered: clk.Now(), Used: clk.Now(),
	}})
	wantStatus(t, do(t, requestFromDE()), http.StatusOK)

	// GeoJS answers, and the first request goes on.
	release()

	err := <-firstEnded
	if err != nil {
		t.Fatalf("first request: %v", err)
	}

	wantLine(t, out.endedLines(t, 2)[1], http.StatusOK, requestlog.ActionForward)

	// The third request's count holds all three: the first was counted
	// when it went on, in the minute under way, and the second's count
	// was kept.
	clk.advance(500 * time.Millisecond)
	wantStatus(t, do(t, requestFromDE()), http.StatusOK)

	if got := out.endedLines(t, 3)[2].Counts.Minute; got != 3 {
		t.Errorf("the third request's count in the minute is %v, want 3", got)
	}
}

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 TestClientsAreLookedUpInTheLookupDatabaseAndGeoJSIsNotAsked(t *testing.T) {
	t.Parallel()

	geojsURL, asked := startGeoJS(t)
	path := filepath.Join(t.TempDir(), "ipinfo_lite.mmdb")
	lookuptest.Write(t, path, map[string]lookuptest.Network{
		fromDE + "/32": {ASN: asnDE, ASName: asNameDE, Country: "DE"},
		fromKP + "/32": {ASN: asnKP, ASName: asNameKP, Country: "KP"},
	})
	s, clk, server := startWithClock(t, geojsURL, map[string]string{
		lookupSource:       fileSource,
		lookupDBPath:       path,
		allowedCountries:   "DE",
		rateLimitPerMinute: "1",
	})

	// fromDE's second request breaks the limit and bans it. The list
	// refuses fromKP, and unplaced, which the file does not hold.
	de := asnAndCountry{asnDE, asNameDE, "DE"}

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

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

	notes := server.Ledger.Bans(netip.MustParsePrefix(fromDE + "/32"))[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)
	}

	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,
		lookupTimeout:    "1h",
		addLookupHeaders: on,
	})
	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.
	for _, from := range []string{fromDE, unplaced, "10.0.0.8"} {
		s.requestWithHeader(from, "/", clientsOwnLookupHeaders,
			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)
	}
}

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

	var (
		mu           sync.Mutex
		asn, country []string
	)

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

		asn, country = 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,
	})
	s := &sender{t: t, addr: addr, out: out}

	s.requestWithHeader(fromDE, "/", clientsOwnLookupHeaders,
		http.StatusOK, requestlog.ActionForward)

	mu.Lock()
	defer mu.Unlock()

	if asn != nil || country != nil {
		t.Errorf("the app was passed X-Client-ASN %v and X-Client-Country %v, want neither",
			asn, country)
	}
}

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

	// The test runs in a synctest bubble, where the time package runs on a
	// clock of the test's own: the wait lasts exactly as long as it should,
	// however slowly the test process runs. Nothing in it may wait on the
	// network, which would keep that clock from moving on: the request is
	// handed to servewaf's handler, and GeoJS is one that never answers.
	synctest.Test(t, func(t *testing.T) {
		// Not the default. The exclusive list needs the answer, and the app
		// is never reached.
		const timeout = 3 * time.Second

		server, out, _ := newServeWAF(t, unansweredGeoJSURL, time.Now,
			map[string]string{
				lookupTimeout:    timeout.String(),
				allowedCountries: "DE",
			})

		began := time.Now()

		serveFromDE(t, server, http.MethodGet, http.NoBody)

		if waited := time.Since(began); waited != timeout {
			t.Errorf("the request waited %s for its answer, want %s", waited, timeout)
		}

		// Without an answer, the client is in no country the list allows.
		wantLine(t, out.requestLine(t), http.StatusForbidden,
			requestlog.ActionCountryDenied)
	})
}

// unansweredGeoJSURL is where a GeoJS that never answers is asked: a
// request to it waits, without the network, until it is abandoned.
// TestMain registers it with Go's default transport, through which GeoJS
// is asked.
const unansweredGeoJSURL = "unanswered://geojs/v1/ip/geo.json"

func TestMain(m *testing.M) {
	transport, _ := http.DefaultTransport.(*http.Transport)
	transport.RegisterProtocol("unanswered", unansweredGeoJS{})
	transport.RegisterProtocol("abuseipdb", abuseIPDBStandIn{})

	m.Run()
}

// unansweredGeoJS is the GeoJS at unansweredGeoJSURL.
type unansweredGeoJS struct{}

// RoundTrip waits until req is abandoned.
func (unansweredGeoJS) RoundTrip(req *http.Request) (*http.Response, error) {
	<-req.Context().Done()

	return nil, req.Context().Err()
}

// clientsOwnLookupHeaders are the X-Client-ASN and X-Client-Country a
// client sends of its own, each twice, in two cases.
const clientsOwnLookupHeaders = "X-Client-ASN: AS1\r\nx-client-asn: AS2\r\n" +
	"X-CLIENT-COUNTRY: KP\r\nx-client-country: CN"

// 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)
	}
}