package ratelimit_test import ( "net/netip" "testing" "time" "sneak.berlin/go/smallwebwaf/internal/ratelimit" ) // limit is the limit the tests set. const limit = 3 // The windows, as Count names them. const ( minute = "minute" hour = "hour" ) func TestEachWindowRefusesAtItsLimitAndLetsTheClientBack(t *testing.T) { t.Parallel() for _, tc := range []struct { window string limits ratelimit.Limits length time.Duration }{ {minute, ratelimit.Limits{PerMinute: limit}, time.Minute}, {hour, ratelimit.Limits{PerHour: limit}, time.Hour}, {"day", ratelimit.Limits{PerDay: limit}, 24 * time.Hour}, } { t.Run(tc.window, func(t *testing.T) { t.Parallel() limiter := ratelimit.New(tc.limits) client := netip.MustParsePrefix("203.0.113.9/32") start := midnight() quarter := tc.length / 4 for range limit { wantCount(t, limiter, client, start, "") } wantCount(t, limiter, client, start, tc.window) // A quarter into the next bucket, the window still covers three // quarters of the bucket before, with its four requests: 3 + 1 // is over the limit. wantCount(t, limiter, client, start.Add(tc.length+quarter), tc.window) // Three quarters into it, a quarter: 1 + 2 is within. wantCount(t, limiter, client, start.Add(tc.length+3*quarter), "") }) } } func TestRefusedRequestsCount(t *testing.T) { t.Parallel() limiter := ratelimit.New(ratelimit.Limits{PerMinute: limit, PerHour: 2 * limit}) refused := netip.MustParsePrefix("203.0.113.9/32") within := netip.MustParsePrefix("203.0.113.10/32") start := midnight() for range limit { wantCount(t, limiter, refused, start, "") wantCount(t, limiter, within, start, "") } for range limit { wantCount(t, limiter, refused, start, minute) } // Half a minute into the next bucket the window covers half of the // bucket before: 3 + 1 is over the minute's limit for the client // whose three refused requests count, and 1.5 + 1 within it for the // other. The first is over the hour's limit too, and the shorter // window is named. halfway := start.Add(time.Minute + time.Minute/2) wantCount(t, limiter, refused, halfway, minute) wantCount(t, limiter, within, halfway, "") // The refused requests count in the hour as well: 6 + 1 + 1 is over // its limit, and 3 + 1 + 1 within it. later := start.Add(10 * time.Minute) wantCount(t, limiter, refused, later, hour) wantCount(t, limiter, within, later, "") } func TestKeepsAtMost20000ClientsDroppingTheLeastRecentlySeen(t *testing.T) { t.Parallel() const maxClients = 20000 limiter := ratelimit.New(ratelimit.Limits{PerMinute: 1}) now := midnight() clients := make([]netip.Prefix, maxClients+1) addr := netip.MustParseAddr("10.0.0.0") for i := range clients { clients[i] = netip.PrefixFrom(addr, addr.BitLen()) addr = addr.Next() } for _, client := range clients[:maxClients] { wantCount(t, limiter, client, now, "") } // The first client is seen again: its second request is over the // limit of one, so it is still counted. wantCount(t, limiter, clients[0], now, minute) // One client more drops the least recently seen, the second, which // starts afresh, while the first is kept. wantCount(t, limiter, clients[maxClients], now, "") wantCount(t, limiter, clients[1], now, "") wantCount(t, limiter, clients[0], now, minute) } // midnight is the start of a bucket in every window. func midnight() time.Time { return time.Date(2026, 10, 4, 0, 0, 0, 0, time.UTC) } // wantCount counts a request from client at now, and checks the window // whose limit it goes over, "" for none. func wantCount( t *testing.T, limiter *ratelimit.Limiter, client netip.Prefix, now time.Time, want string, ) { t.Helper() got := limiter.Count(client, now) if got != want { t.Errorf("request from %s at %s is over %q, want %q", client, now.Format(time.RFC3339), got, want) } }