smallwebwaf/internal/ratelimit/snapshot_test.go

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.CountBytes(client, start, 5)
	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, of requests and of bytes, 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)
	}

	if got.MinuteBytes != (ratelimit.Buckets{}) || got.HourBytes.Current != 5 ||
		got.DayBytes.Current != 5 {
		t.Errorf("loaded two minutes on with buckets of bytes %+v, %+v and %+v",
			got.MinuteBytes, got.HourBytes, got.DayBytes)
	}

	// 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 || got.MinuteBytes.Current != 5 {
		t.Errorf("loaded just under two minutes on with minute buckets %+v and %+v",
			got.Minute, got.MinuteBytes)
	}
}

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