package ratelimit_test import ( "net/netip" "slices" "testing" "time" "sneak.berlin/go/smallwebwaf/internal/ratelimit" ) func TestSnapshotListsTheClientsByAddress(t *testing.T) { t.Parallel() want := []string{"192.0.2.1/32", "203.0.113.9/32", "203.0.113.10/32", "2001:db8::/64"} limiter := ratelimit.New(ratelimit.Limits{}) for _, i := range []int{2, 3, 0, 1} { limiter.Count(netip.MustParsePrefix(want[i]), midnight()) } snapshot := limiter.Snapshot() got := make([]string, 0, len(snapshot)) for _, c := range snapshot { got = append(got, c.Client.String()) } if !slices.Equal(got, want) { t.Errorf("snapshot %v, want %v", got, want) } counted := ratelimit.Buckets{Start: midnight(), Current: 1} if snapshot[0].Minute != counted || snapshot[0].Day != counted { t.Errorf("buckets %+v and %+v, want %+v", snapshot[0].Minute, snapshot[0].Day, counted) } } func TestLoadedCountsCarryOn(t *testing.T) { t.Parallel() client := netip.MustParsePrefix("203.0.113.9/32") start := midnight() before := ratelimit.New(ratelimit.Limits{PerHour: limit}) for range limit { wantCount(t, before, client, start, "") } // Loaded into a new limiter, as across a restart, the client has no // fresh allowance. later := start.Add(time.Minute) after := ratelimit.New(ratelimit.Limits{PerHour: limit}) after.Load(before.Snapshot(), later) wantCount(t, after, client, later, hour) } func TestLoadEmptiesBucketsWhoseTimeHasPassed(t *testing.T) { t.Parallel() client := netip.MustParsePrefix("203.0.113.9/32") start := midnight() limiter := ratelimit.New(ratelimit.Limits{}) limiter.Count(client, start) limiter.AddToHistory(client, start, ratelimit.Request{Forwarded: true}) loaded := func(now time.Time) ratelimit.Client { t.Helper() after := ratelimit.New(ratelimit.Limits{}) after.Load(limiter.Snapshot(), now) return after.Snapshot()[0] } // Two minutes on, the window that ends then covers neither of the // minute's buckets, which are emptied; the hour's and the day's stay, // and so does the history. got := loaded(start.Add(2 * time.Minute)) if got.Minute != (ratelimit.Buckets{}) || got.Hour.Current != 1 || got.Day.Current != 1 || got.History.Requests != 1 { t.Errorf("loaded two minutes on as %+v", got) } // A moment before, the window still covers some of the earlier one. got = loaded(start.Add(2*time.Minute - time.Nanosecond)) if got.Minute.Current != 1 { t.Errorf("loaded just under two minutes on with minute buckets %+v", got.Minute) } } func TestLoadDropsTheLeastRecentlySeenFirst(t *testing.T) { t.Parallel() const maxClients = 20000 // clients.json lists the clients by address. Here each was last seen // a second before the one listed before it, so the last listed is the // one seen longest ago, and the one dropped. clients := make([]ratelimit.Client, maxClients+1) addr := netip.MustParseAddr("10.0.0.0") for i := range clients { clients[i].Client = netip.PrefixFrom(addr, addr.BitLen()) clients[i].History.LastSeen = midnight().Add(-time.Duration(i) * time.Second) addr = addr.Next() } limiter := ratelimit.New(ratelimit.Limits{}) limiter.Load(clients, midnight()) got := limiter.Snapshot() if len(got) != maxClients || got[0].Client != clients[0].Client || got[maxClients-1].Client != clients[maxClients-1].Client { t.Errorf("%d clients kept, from %s to %s; want %d, from %s to %s", len(got), got[0].Client, got[len(got)-1].Client, maxClients, clients[0].Client, clients[maxClients-1].Client) } }