smallwebwaf/internal/ratelimit/history_test.go

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{}, tableSize)
	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{}, tableSize)
	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}, tableSize)
	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{}, tableSize)

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