package ratelimit_test import ( "net/netip" "testing" "time" "sneak.berlin/go/smallwebwaf/internal/ratelimit" ) func TestHistoryKeepsEveryRequest(t *testing.T) { t.Parallel() limiter := ratelimit.New(ratelimit.Limits{}) client := netip.MustParsePrefix("203.0.113.9/32") start := midnight() for i, r := range []ratelimit.Request{ {Forwarded: true, Status: 200, RequestBytes: 10, ResponseBytes: 100}, {Forwarded: true, Status: 101}, {Forwarded: true, Status: 304, RequestBytes: 5}, {Refused: true, Status: 403, ResponseBytes: 10, BrokeLimit: true}, {Forwarded: true, Status: 502, ResponseBytes: 12}, // Closed without an answer: refused, and no response. {Refused: true, Status: 0}, // Answered 404 at smallwebwaf's own endpoints: neither forwarded // nor refused. {Status: 404}, } { limiter.AddToHistory(client, start.Add(time.Duration(i)*time.Minute), r) } want := ratelimit.History{ FirstSeen: start, LastSeen: start.Add(6 * time.Minute), Requests: 7, Forwarded: 4, Refused: 2, RequestBytes: 15, ResponseBytes: 122, Responses: ratelimit.Responses{ Status1xx: 1, Status2xx: 1, Status3xx: 1, Status4xx: 2, Status5xx: 1, }, Offences: ratelimit.Offences{Limit: 1}, } got := historyOf(t, limiter, client) if got != want { t.Errorf("history\n%+v\nwant\n%+v", got, want) } } func TestLookupReachesTheHistoryOfAClientInTheTable(t *testing.T) { t.Parallel() limiter := ratelimit.New(ratelimit.Limits{}) client := netip.MustParsePrefix("203.0.113.9/32") other := netip.MustParsePrefix("198.51.100.7/32") start := midnight() limiter.AddToHistory(client, start, ratelimit.Request{Forwarded: true}) limiter.AddLookup(client, start, "AS64496", "Example Net", "DE") // A later answer replaces it, and one for a client the table does not // hold adds no client. limiter.AddLookup(client, start.Add(time.Hour), "AS64497", "Other Net", "FR") limiter.AddLookup(other, start, "AS64496", "Example Net", "DE") want := ratelimit.History{ FirstSeen: start, LastSeen: start, ASN: "AS64497", ASName: "Other Net", Country: "FR", LookedUp: start.Add(time.Hour), Requests: 1, Forwarded: 1, } got := historyOf(t, limiter, client) if got != want { t.Errorf("history\n%+v\nwant\n%+v", got, want) } if clients := limiter.Snapshot(); len(clients) != 1 { t.Errorf("the table holds %+v, want %s alone", clients, client) } } func TestResetKeepsTheHistory(t *testing.T) { t.Parallel() limiter := ratelimit.New(ratelimit.Limits{PerMinute: limit}) client := netip.MustParsePrefix("203.0.113.9/32") start := midnight() for range limit { wantCount(t, limiter, client, start, "") limiter.AddToHistory(client, start, ratelimit.Request{Forwarded: true}) } limiter.Reset(client) if got := historyOf(t, limiter, client).Requests; got != limit { t.Errorf("the history counts %d requests, want %d", got, limit) } } func TestRequestsAddsUpTheClientsInsideTheNetblock(t *testing.T) { t.Parallel() limiter := ratelimit.New(ratelimit.Limits{}) for client, requests := range map[string]int{ "198.51.100.9/32": 2, "198.51.100.10/32": 3, "192.0.2.1/32": 5, "2001:db8:5::/64": 7, } { for range requests { limiter.AddToHistory(netip.MustParsePrefix(client), midnight(), ratelimit.Request{}) } } for netblock, want := range map[string]int64{ "198.51.100.9/32": 2, "198.51.100.0/24": 5, "2001:db8:5::/64": 7, "203.0.113.0/24": 0, } { got := limiter.Requests(netip.MustParsePrefix(netblock)) if got != want { t.Errorf("%s has sent %d requests, want %d", netblock, got, want) } } } // historyOf returns client's history. func historyOf( t *testing.T, limiter *ratelimit.Limiter, client netip.Prefix, ) ratelimit.History { t.Helper() for _, c := range limiter.Snapshot() { if c.Client == client { return c.History } } t.Fatalf("%s is not in the table", client) return ratelimit.History{} }