smallwebwaf/internal/ratelimit/ratelimit_test.go

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 TestHitGivesTheLimitAndTheRequestsCounted(t *testing.T) {
	t.Parallel()

	limiter := ratelimit.New(ratelimit.Limits{PerMinute: limit, PerHour: limit})
	client := netip.MustParsePrefix("203.0.113.9/32")
	start := midnight()

	for range limit {
		_, _, over := limiter.Count(client, start)
		if over {
			t.Fatal("a request within the limit is over it")
		}
	}

	// Over both limits; the minute's is named, with the four requests.
	_, hit, over := limiter.Count(client, start)

	want := ratelimit.Hit{
		Kind: ratelimit.KindRequests, Window: minute, Limit: limit, Count: limit + 1,
	}
	if !over || hit != want {
		t.Errorf("request over the limit gives %+v and %t, want %+v and true",
			hit, over, want)
	}
}

func TestEachByteLimitIsBrokenByTheBytesCounted(t *testing.T) {
	t.Parallel()

	const byteLimit = 1000

	for _, tc := range []struct {
		window string
		limits ratelimit.Limits
	}{
		{minute, ratelimit.Limits{BytesPerMinute: byteLimit}},
		{hour, ratelimit.Limits{BytesPerHour: byteLimit}},
		{"day", ratelimit.Limits{BytesPerDay: byteLimit}},
	} {
		t.Run(tc.window, func(t *testing.T) {
			t.Parallel()

			limiter := ratelimit.New(tc.limits)
			client := netip.MustParsePrefix("203.0.113.9/32")

			// 600 bytes are within the limit, 600 more over it.
			_, _, over := limiter.CountBytes(client, midnight(), 600)
			if over {
				t.Fatal("600 bytes are over the limit of 1000")
			}

			_, hit, over := limiter.CountBytes(client, midnight(), 600)

			want := ratelimit.Hit{
				Kind: ratelimit.KindBytes, Window: tc.window, Limit: byteLimit, Count: 1200,
			}
			if !over || hit != want {
				t.Errorf("1200 bytes give %+v and %t, want %+v and true", hit, over, want)
			}
		})
	}
}

func TestALimitIsBrokenOnlyByWhatIsAddedToIt(t *testing.T) {
	t.Parallel()

	limiter := ratelimit.New(ratelimit.Limits{PerMinute: 2, BytesPerMinute: 1000})
	client := netip.MustParsePrefix("203.0.113.9/32")
	other := netip.MustParsePrefix("203.0.113.10/32")
	start := midnight()

	// The third request breaks the rate limit. The bytes of a request
	// counted after it, within the byte limit, do not break it again.
	for range 2 {
		wantCount(t, limiter, client, start, "")
	}

	wantCount(t, limiter, client, start, minute)
	wantBytesCount(t, limiter, client, start, 500, "")
	wantBytesCount(t, limiter, client, start, 600, ratelimit.KindBytes)

	// Bytes over the byte limit do not have the next request break it, nor
	// the rate limit, which that request is within.
	wantBytesCount(t, limiter, other, start, 1200, ratelimit.KindBytes)
	wantCount(t, limiter, other, start, "")
}

func TestCountGivesTheBytesInEachWindow(t *testing.T) {
	t.Parallel()

	limiter := ratelimit.New(ratelimit.Limits{})
	client := netip.MustParsePrefix("203.0.113.9/32")
	start := midnight()

	limiter.CountBytes(client, start, 300)

	// A quarter into the next hour, the minute has only these 100 bytes.
	// The hour still covers three quarters of the bucket before, whose 300
	// bytes count 225, and these: 325. The day covers all 400.
	later := start.Add(time.Hour + time.Hour/4)
	limiter.CountBytes(client, later, 100)

	// A request's counts give the bytes counted so far too.
	counts, _, _ := limiter.Count(client, later)

	want := ratelimit.Counts{
		Minute: 1, Hour: 1, Day: 1, MinuteBytes: 100, HourBytes: 325, DayBytes: 400,
	}
	if counts != want {
		t.Errorf("counts %+v, want %+v", counts, want)
	}
}

func TestResetSetsTheBytesBackToZero(t *testing.T) {
	t.Parallel()

	limiter := ratelimit.New(ratelimit.Limits{BytesPerDay: 1000})
	client := netip.MustParsePrefix("203.0.113.9/32")
	start := midnight()

	wantBytesCount(t, limiter, client, start, 1200, ratelimit.KindBytes)
	limiter.Reset(client)

	// The client has its whole allowance of bytes again.
	wantBytesCount(t, limiter, client, start, 1000, "")
}

func TestCountGivesTheRequestsInEachWindow(t *testing.T) {
	t.Parallel()

	limiter := ratelimit.New(ratelimit.Limits{})
	client := netip.MustParsePrefix("203.0.113.9/32")
	start := midnight()

	for range 3 {
		limiter.Count(client, start)
	}

	// A quarter into the next hour, the minute has only this request. The
	// hour still covers three quarters of the bucket before, with its three
	// requests, which count 2.25, and this one: 3.25. The day covers all
	// four.
	counts, _, _ := limiter.Count(client, start.Add(time.Hour+time.Hour/4))

	want := ratelimit.Counts{Minute: 1, Hour: 3.25, Day: 4}
	if counts != want {
		t.Errorf("counts %+v, want %+v", counts, want)
	}
}

func TestResetSetsTheCountsBackToZero(t *testing.T) {
	t.Parallel()

	limiter := ratelimit.New(ratelimit.Limits{PerMinute: limit, PerDay: limit})
	client := netip.MustParsePrefix("203.0.113.9/32")
	start := midnight()

	for range limit {
		wantCount(t, limiter, client, start, "")
	}

	wantCount(t, limiter, client, start, minute)
	limiter.Reset(client)

	// At the same moment, the client has its whole allowance again.
	for range limit {
		wantCount(t, limiter, client, start, "")
	}

	wantCount(t, limiter, client, start, minute)
}

func TestClientBackAfterAWholeBucketIsWithinTheLimitAtOnce(t *testing.T) {
	t.Parallel()

	limiter := ratelimit.New(ratelimit.Limits{PerHour: limit})
	client := netip.MustParsePrefix("203.0.113.9/32")
	start := midnight()

	for range limit {
		wantCount(t, limiter, client, start, "")
	}

	wantCount(t, limiter, client, start, hour)

	// No request in the whole next bucket, so a quarter into the one after
	// it the window covers none of the four requests: 1 is within the
	// limit. Were they counted as the bucket before, 3 + 1 would be over.
	wantCount(t, limiter, client, start.Add(2*time.Hour+time.Hour/4), "")
}

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 TestRequestCountedLateGoesInTheBucketUnderWay(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, "")
	}

	// A concurrent request dated a moment before the bucket under way, but
	// counted after it began, is counted in it: 3 + 1 is over the limit.
	wantCount(t, limiter, client, start.Add(-time.Millisecond), minute)
}

func TestClockSetBackStartsTheBucketsAfresh(t *testing.T) {
	t.Parallel()

	limiter := ratelimit.New(ratelimit.Limits{PerHour: limit})
	client := netip.MustParsePrefix("203.0.113.9/32")
	start := midnight()

	for range limit {
		wantCount(t, limiter, client, start, "")
	}

	// Half an hour into the next bucket: 3 / 2 + 1 is within the limit.
	wantCount(t, limiter, client, start.Add(time.Hour+time.Hour/2), "")

	// The clock is set back an hour. Counted in the bucket under way, the
	// next request would find the bucket before it at full weight, 3 + 2,
	// over the limit until the clock caught up. The buckets start afresh
	// instead, and the client is refused only past the limit again.
	setBack := start.Add(time.Hour / 2)
	for range limit {
		wantCount(t, limiter, client, setBack, "")
	}

	wantCount(t, limiter, client, setBack, hour)
}

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

	_, hit, _ := limiter.Count(client, now)
	if hit.Window != want {
		t.Errorf("request from %s at %s is over %q, want %q",
			client, now.Format(time.RFC3339), hit.Window, want)
	}
}

// wantBytesCount counts bytes from client at now, and checks the kind of
// the limit they break, "" for none.
func wantBytesCount(
	t *testing.T, limiter *ratelimit.Limiter, client netip.Prefix, now time.Time,
	bytes int64, want string,
) {
	t.Helper()

	_, hit, _ := limiter.CountBytes(client, now, bytes)
	if hit.Kind != want {
		t.Errorf("%d bytes from %s at %s break a limit on %q, want %q",
			bytes, client, now.Format(time.RFC3339), hit.Kind, want)
	}
}