package lookup_test import ( "encoding/json" "log/slog" "net/http" "net/http/httptest" "net/netip" "slices" "strings" "sync" "testing" "time" "sneak.berlin/go/smallwebwaf/internal/lookup" ) const ( // germany is where the stand-in for GeoJS places every address but // unplaced. germany = "DE" // unplaced is the address it cannot place. unplaced = "192.0.2.1" // timeout is how long a new client waits for its answer. timeout = time.Second // waitLimit bounds how long a test waits for what should happen. waitLimit = 10 * time.Second // pollInterval is how often a test looks again. pollInterval = 10 * time.Millisecond // week is how long an answer is kept. week = 7 * 24 * time.Hour ) func TestKeptAnswerIsUsedFor7DaysThenAskedAgain(t *testing.T) { t.Parallel() geojs, clock, g := start(t) placed := netip.MustParsePrefix("203.0.113.9/32") notPlaced := netip.MustParsePrefix(unplaced + "/32") wantCountry(t, g, placed, germany) wantCountry(t, g, notPlaced, "") wantRequests(t, geojs, 2) // An answer without a country is kept too. clock.advance(week - time.Second) wantCountry(t, g, placed, germany) wantCountry(t, g, notPlaced, "") wantRequests(t, geojs, 2) clock.advance(time.Second) wantCountry(t, g, placed, germany) wantRequests(t, geojs, 3) wantAsked(t, geojs, 2, "203.0.113.9") } func TestNewClientWaitsAtMostOneSecondThenCountsAsNotFound(t *testing.T) { t.Parallel() geojs, clock, g := start(t) client := netip.MustParsePrefix("203.0.113.9/32") // GeoJS does not answer: the request to it is abandoned. geojs.set(hanging) began := time.Now() wantCountry(t, g, client, "") took := time.Since(began) if took < timeout || took > timeout+waitLimit/2 { t.Errorf("waited %s for the answer, want %s", took, timeout) } // Its next request does not wait. began = time.Now() wantCountry(t, g, client, "") took = time.Since(began) if took > timeout/2 { t.Errorf("waited %s again, want no wait", took) } // Once GeoJS answers, the client is asked about again in the // background, and has its country. geojs.set(answering) deadline := time.Now().Add(waitLimit) for g.Country(t.Context(), client) != germany { if time.Now().After(deadline) { t.Fatalf("no answer after %s", waitLimit) } clock.advance(time.Minute) time.Sleep(pollInterval) } wantRequests(t, geojs, 2) wantAsked(t, geojs, 1, "203.0.113.9") } func TestWaitingClientsAreAskedAboutInOneRequest(t *testing.T) { t.Parallel() geojs, clock, g := start(t) // GeoJS fails, and is then left alone for a second, while three more // clients come. An IPv6 client is a /64, and GeoJS is asked about its // first address. geojs.set(failing) wantCountry(t, g, netip.MustParsePrefix("203.0.113.1/32"), "") wantCountry(t, g, netip.MustParsePrefix("203.0.113.2/32"), "") wantCountry(t, g, netip.MustParsePrefix("2001:db8:1:2::/64"), "") wantRequests(t, geojs, 1) geojs.set(answering) clock.advance(time.Second) wantCountry(t, g, netip.MustParsePrefix("203.0.113.3/32"), germany) wantRequests(t, geojs, 2) wantAsked(t, geojs, 1, "203.0.113.1", "203.0.113.2", "2001:db8:1:2::", "203.0.113.3") } func TestKeptAnswersUnaffectedWhileGeoJSFailsAndAskedAgainWithBackoff(t *testing.T) { t.Parallel() geojs, clock, g := start(t) clients := newClients() kept := clients() wantCountry(t, g, kept, germany) geojs.set(failing) wantCountry(t, g, kept, germany) wantRequests(t, geojs, 1) // Each failure leaves GeoJS alone twice as long as the one before, up // to five minutes. New clients meanwhile count as not found. requests := 1 for _, delay := range []time.Duration{ time.Second, 2 * time.Second, 4 * time.Second, 8 * time.Second, 16 * time.Second, 32 * time.Second, 64 * time.Second, 128 * time.Second, 256 * time.Second, 5 * time.Minute, 5 * time.Minute, } { wantCountry(t, g, clients(), "") requests++ wantRequests(t, geojs, requests) clock.advance(delay - time.Millisecond) wantCountry(t, g, clients(), "") wantRequests(t, geojs, requests) clock.advance(time.Millisecond) } // Once GeoJS answers again, it is asked about every client waiting. geojs.set(answering) wantCountry(t, g, clients(), germany) wantRequests(t, geojs, requests+1) asked := waitForRequests(t, geojs, requests+1) if len(asked[requests]) != 23 { t.Errorf("GeoJS was asked about %d clients, want 23", len(asked[requests])) } } func TestAtMost200AddressesInOneRequest(t *testing.T) { t.Parallel() geojs, clock, g := start(t) clients := newClients() first := clients() // 201 clients wait while GeoJS is left alone after a failure. geojs.set(failing) wantCountry(t, g, first, "") for range 200 { wantCountry(t, g, clients(), "") } // The first one's next request has GeoJS asked again. geojs.set(answering) clock.advance(time.Second) wantCountry(t, g, first, "") asked := waitForRequests(t, geojs, 3) if len(asked[1]) != 200 || len(asked[2]) != 1 { t.Errorf("GeoJS was asked about %d and then %d clients, want 200 and 1", len(asked[1]), len(asked[2])) } } func TestAtMost10000ClientsWait(t *testing.T) { t.Parallel() geojs, clock, g := start(t) clients := newClients() first := clients() // 10,000 clients wait while GeoJS is left alone after a failure, and // one more cannot join them. geojs.set(failing) wantCountry(t, g, first, "") for range 9999 { wantCountry(t, g, clients(), "") } extra := clients() wantCountry(t, g, extra, "") // The first one's next request has GeoJS asked about the 10,000, 200 // at a time, and not about the one more. geojs.set(answering) clock.advance(time.Second) wantCountry(t, g, first, "") asked := waitForRequests(t, geojs, 51) for i, request := range asked { if slices.Contains(request, extra.Addr().String()) { t.Errorf("request %d asked about %s", i, extra.Addr()) } } // With room among those waiting, it is asked about. wantCountry(t, g, extra, germany) } // How the stand-in for GeoJS answers. const ( answering = iota failing // with 503 hanging // not at all, until the request is abandoned ) // standIn is a stand-in for GeoJS. It notes the addresses each request // asks about. type standIn struct { server *httptest.Server mu sync.Mutex answers int requests [][]string } // ServeHTTP answers a request about the addresses in its ip parameter. func (s *standIn) ServeHTTP(w http.ResponseWriter, r *http.Request) { addrs := strings.Split(r.URL.Query().Get("ip"), ",") s.mu.Lock() s.requests = append(s.requests, addrs) answers := s.answers s.mu.Unlock() switch answers { case failing: w.WriteHeader(http.StatusServiceUnavailable) return case hanging: <-r.Context().Done() return } list := make([]map[string]string, 0, len(addrs)) for _, addr := range addrs { country := germany if addr == unplaced { country = "" } list = append(list, map[string]string{"ip": addr, "country": country}) } err := json.NewEncoder(w).Encode(list) if err != nil { http.Error(w, err.Error(), http.StatusInternalServerError) } } // set sets how the stand-in answers. func (s *standIn) set(answers int) { s.mu.Lock() defer s.mu.Unlock() s.answers = answers } // asked returns the addresses each request has asked about so far. func (s *standIn) asked() [][]string { s.mu.Lock() defer s.mu.Unlock() return slices.Clone(s.requests) } // testClock is a clock the test sets. type testClock struct { mu sync.Mutex now time.Time } // Now tells the time. func (c *testClock) Now() time.Time { c.mu.Lock() defer c.mu.Unlock() return c.now } // advance moves the clock on by d. func (c *testClock) advance(d time.Duration) { c.mu.Lock() defer c.mu.Unlock() c.now = c.now.Add(d) } // start starts a stand-in for GeoJS that answers, and returns it, a // clock, and a GeoJS asking it by that clock. func start(t *testing.T) (*standIn, *testClock, *lookup.GeoJS) { t.Helper() geojs := &standIn{} geojs.server = httptest.NewServer(geojs) t.Cleanup(geojs.server.Close) clock := &testClock{now: time.Date(2026, 10, 4, 0, 0, 0, 0, time.UTC)} g := lookup.New(lookup.Params{ URL: geojs.server.URL, Now: clock.Now, ProcessLog: slog.New(slog.DiscardHandler), }) return geojs, clock, g } // newClients returns what returns a new IPv4 client each time it is // called. func newClients() func() netip.Prefix { addr := netip.MustParseAddr("10.0.0.0") return func() netip.Prefix { addr = addr.Next() return netip.PrefixFrom(addr, addr.BitLen()) } } // wantCountry checks the country g gives client. func wantCountry(t *testing.T, g *lookup.GeoJS, client netip.Prefix, want string) { t.Helper() got := g.Country(t.Context(), client) if got != want { t.Errorf("%s is in %q, want %q", client, got, want) } } // wantRequests checks how many requests GeoJS has had. func wantRequests(t *testing.T, geojs *standIn, want int) { t.Helper() got := len(geojs.asked()) if got != want { t.Errorf("GeoJS had %d requests, want %d", got, want) } } // wantAsked checks the addresses request i asked about, in any order. func wantAsked(t *testing.T, geojs *standIn, i int, want ...string) { t.Helper() asked := geojs.asked() if len(asked) <= i { t.Fatalf("GeoJS had %d requests, want more than %d", len(asked), i) } got := slices.Sorted(slices.Values(asked[i])) slices.Sort(want) if !slices.Equal(got, want) { t.Errorf("request %d asked about %v, want %v", i, got, want) } } // waitForRequests waits for GeoJS to have had count requests, and returns // the addresses each asked about. func waitForRequests(t *testing.T, geojs *standIn, count int) [][]string { t.Helper() deadline := time.Now().Add(waitLimit) for time.Now().Before(deadline) { asked := geojs.asked() if len(asked) >= count { return asked } time.Sleep(pollInterval) } t.Fatalf("fewer than %d requests to GeoJS after %s", count, waitLimit) return nil }