Byte limits per client over a minute, an hour and a day (closes #20)

SWWAF_BYTES_LIMIT_PER_MINUTE, _PER_HOUR and _PER_DAY (10G, 20G, 50G)
and SWWAF_BYTES_COUNT (both). A request's bytes are counted once its
answer has ended, for a request passed to the app that the rate limits
count; what a WebSocket carries each way, once it closes. Bytes over a
limit ban the client as a broken rate limit does, and cut nothing
short. clients.json keeps the byte buckets, the log line's counts carry
the byte totals, ban notes say what the limit is on, and the limit hits
metric is labelled by kind.

Judgement call: limit_hit names a byte window minute_bytes, hour_bytes
or day_bytes, as counts names the byte totals.
Judgement call: in observe mode, the bytes of a request enforce mode
would have refused are not counted.

Model: opus-5-5

internal/ratelimit/ratelimit.go

@@ -1,9 +1,9 @@
 // Package ratelimit keeps the table of clients: each client's requests
-// counted over a minute, an hour and a day, as the "Counting method"
-// section of SPEC.md describes, which tell when a request takes the client
-// over a rate limit, and each client's history since it was first seen.
-// At most 20,000 clients are kept, in memory, and written to clients.json
-// and read from it by the state package.
+// and bytes counted over a minute, an hour and a day, as the "Counting
+// method" section of SPEC.md describes, which tell when a request takes
+// the client over a rate limit or a byte limit, and each client's history
+// since it was first seen. At most 20,000 clients are kept, in memory, and
+// written to clients.json and read from it by the state package.
 package ratelimit
 
 import (
@@ -23,19 +23,30 @@ const maxClients = 20000
 
 const day = 24 * time.Hour
 
+// The kinds of limits, as the metrics name them.
+const (
+	// KindRequests is a rate limit, on a client's requests.
+	KindRequests = "requests"
+	// KindBytes is a byte limit, on a client's bytes.
+	KindBytes = "bytes"
+)
+
 // Limits are the most requests a client may make in a minute, an hour and
-// a day. Zero is no limit.
+// a day, and the most bytes. Zero is no limit.
 type Limits struct {
-	PerMinute int64
-	PerHour   int64
-	PerDay    int64
+	PerMinute      int64
+	PerHour        int64
+	PerDay         int64
+	BytesPerMinute int64
+	BytesPerHour   int64
+	BytesPerDay    int64
 }
 
-// Limiter counts each client's requests against the limits, and keeps
-// its history. It is safe for concurrent use.
+// Limiter counts each client's requests and bytes against the limits, and
+// keeps its history. It is safe for concurrent use.
 type Limiter struct {
 	// windows are the minute, the hour and the day, in the order of
-	// Client.buckets.
+	// Client.buckets and Client.byteBuckets.
 	windows [3]window
 
 	mu      sync.Mutex
@@ -43,17 +54,23 @@ type Limiter struct {
 }
 
 // Client is a client in the table, as clients.json holds it: its buckets
-// in each window, and its history.
+// of requests and of bytes in each window, and its history.
+//
+//nolint:tagliatelle // the state files use snake_case, as the request log does
 type Client struct {
-	Client  netip.Prefix `json:"client"`
-	Minute  Buckets      `json:"minute"`
-	Hour    Buckets      `json:"hour"`
-	Day     Buckets      `json:"day"`
-	History History      `json:"history"`
+	Client      netip.Prefix `json:"client"`
+	Minute      Buckets      `json:"minute"`
+	Hour        Buckets      `json:"hour"`
+	Day         Buckets      `json:"day"`
+	MinuteBytes Buckets      `json:"minute_bytes"`
+	HourBytes   Buckets      `json:"hour_bytes"`
+	DayBytes    Buckets      `json:"day_bytes"`
+	History     History      `json:"history"`
 }
 
-// Buckets are a client's two buckets in one window: the requests in the
-// bucket under way, which began at Start, and in the bucket before it.
+// Buckets are a client's two buckets in one window: the requests, or the
+// bytes, in the bucket under way, which began at Start, and in the bucket
+// before it.
 type Buckets struct {
 	Start    time.Time `json:"start"`
 	Current  int64     `json:"current"`
@@ -101,7 +118,7 @@ type Responses struct {
 
 // Offences are a client's offences, by kind.
 type Offences struct {
-	// Limit is its requests that broke a rate limit.
+	// Limit is its requests that broke a rate limit or a byte limit.
 	Limit int64 `json:"limit"`
 }
 
@@ -119,7 +136,8 @@ type Request struct {
 	// and of its response.
 	RequestBytes  int64
 	ResponseBytes int64
-	// BrokeLimit is true for a request that broke a rate limit.
+	// BrokeLimit is true for a request that broke a rate limit or a byte
+	// limit.
 	BrokeLimit bool
 }
 
@@ -132,62 +150,71 @@ func New(limits Limits) *Limiter {
 
 	return &Limiter{
 		windows: [3]window{
-			{name: "minute", length: time.Minute, limit: limits.PerMinute},
-			{name: "hour", length: time.Hour, limit: limits.PerHour},
-			{name: "day", length: day, limit: limits.PerDay},
+			{
+				name: "minute", length: time.Minute,
+				limit: limits.PerMinute, byteLimit: limits.BytesPerMinute,
+			},
+			{
+				name: "hour", length: time.Hour,
+				limit: limits.PerHour, byteLimit: limits.BytesPerHour,
+			},
+			{
+				name: "day", length: day,
+				limit: limits.PerDay, byteLimit: limits.BytesPerDay,
+			},
 		},
 		clients: clients,
 	}
 }
 
-// Hit is a request that takes a client over a rate limit.
+// Hit is a request that takes a client over a rate limit, or whose bytes
+// take it over a byte limit.
 type Hit struct {
+	// Kind is KindRequests for a rate limit, KindBytes for a byte limit.
+	Kind string
 	// Window is "minute", "hour" or "day".
 	Window string
 	// Limit is the window's limit.
 	Limit int64
-	// Requests is the client's requests counted in the window, this one
-	// included.
-	Requests float64
+	// Count is the client's requests, or bytes, counted in the window,
+	// this request's included.
+	Count float64
 }
 
-// Counts are a client's requests in the minute, the hour and the day that
-// end at a request, that request included.
+// Counts are a client's requests and bytes in the minute, the hour and
+// the day that end at a request, that request's included.
+//
+//nolint:tagliatelle // SPEC.md's request log names its fields in snake_case
 type Counts struct {
-	Minute float64 `json:"minute"`
-	Hour   float64 `json:"hour"`
-	Day    float64 `json:"day"`
+	Minute      float64 `json:"minute"`
+	Hour        float64 `json:"hour"`
+	Day         float64 `json:"day"`
+	MinuteBytes float64 `json:"minute_bytes"`
+	HourBytes   float64 `json:"hour_bytes"`
+	DayBytes    float64 `json:"day_bytes"`
 }
 
 // Count counts a request from client at now, in every window, whether or
-// not it is refused, and returns the client's requests in each window. It
-// reports whether the request takes the client over a limit, and the
-// window whose limit it goes over, the shortest if it is over several.
+// not it is refused, and returns the client's counts in each window. It
+// reports whether the request takes the client over a rate limit, and the
+// hit: the window whose limit it goes over, the shortest if it is over
+// several.
 func (l *Limiter) Count(client netip.Prefix, now time.Time) (Counts, Hit, bool) {
-	l.mu.Lock()
-	defer l.mu.Unlock()
-
-	var (
-		requests [3]float64
-		hit      Hit
-	)
-
-	for i, b := range l.get(client).buckets() {
-		w := l.windows[i]
-
-		requests[i] = b.add(now, w.length)
-		if hit.Window == "" && w.limit > 0 && requests[i] > float64(w.limit) {
-			hit = Hit{Window: w.name, Limit: w.limit, Requests: requests[i]}
-		}
-	}
-
-	counts := Counts{Minute: requests[0], Hour: requests[1], Day: requests[2]}
-
-	return counts, hit, hit.Window != ""
+	return l.count(client, now, 1, 0)
 }
 
-// Reset sets client's counts in every window back to zero. Its history
-// keeps its totals.
+// CountBytes counts bytes, those of a request from client that has ended,
+// at now, in every window, and returns the client's counts in each window.
+// It reports whether the bytes take the client over a byte limit, and the
+// hit, as Count does.
+func (l *Limiter) CountBytes(
+	client netip.Prefix, now time.Time, bytes int64,
+) (Counts, Hit, bool) {
+	return l.count(client, now, 0, bytes)
+}
+
+// Reset sets client's counts of requests and of bytes in every window
+// back to zero. Its history keeps its totals.
 func (l *Limiter) Reset(client netip.Prefix) {
 	l.mu.Lock()
 	defer l.mu.Unlock()
@@ -195,6 +222,7 @@ func (l *Limiter) Reset(client netip.Prefix) {
 	c, seen := l.clients.Peek(client)
 	if seen {
 		c.Minute, c.Hour, c.Day = Buckets{}, Buckets{}, Buckets{}
+		c.MinuteBytes, c.HourBytes, c.DayBytes = Buckets{}, Buckets{}, Buckets{}
 	}
 }
 
@@ -328,12 +356,13 @@ func (l *Limiter) Load(clients []Client, now time.Time) {
 	l.clients.Purge()
 
 	for _, c := range clients {
-		for i, b := range c.buckets() {
-			// The window that ends at now covers neither bucket once it
-			// begins after the bucket under way has ended.
-			length := l.windows[i].length
-			if !now.Add(-length).Before(b.Start.Add(length)) {
-				*b = Buckets{}
+		for i, w := range l.windows {
+			for _, b := range []*Buckets{c.buckets()[i], c.byteBuckets()[i]} {
+				// The window that ends at now covers neither bucket once it
+				// begins after the bucket under way has ended.
+				if !now.Add(-w.length).Before(b.Start.Add(w.length)) {
+					*b = Buckets{}
+				}
 			}
 		}
 
@@ -341,6 +370,49 @@ func (l *Limiter) Load(clients []Client, now time.Time) {
 	}
 }
 
+// count adds requests and bytes from client at now to its buckets in
+// every window, and returns its counts. A limit is broken only by what is
+// added to it, so that a request whose bytes are counted after another of
+// the client's requests broke a rate limit does not break it too.
+func (l *Limiter) count(
+	client netip.Prefix, now time.Time, requests, bytes int64,
+) (Counts, Hit, bool) {
+	l.mu.Lock()
+	defer l.mu.Unlock()
+
+	c := l.get(client)
+	requestBuckets, byteBuckets := c.buckets(), c.byteBuckets()
+
+	var (
+		requestCounts, byteCounts [3]float64
+		hit                       Hit
+	)
+
+	for i, w := range l.windows {
+		requestCounts[i] = requestBuckets[i].add(now, w.length, requests)
+		byteCounts[i] = byteBuckets[i].add(now, w.length, bytes)
+
+		switch {
+		case hit.Window != "":
+		case requests > 0 && w.limit > 0 && requestCounts[i] > float64(w.limit):
+			hit = Hit{
+				Kind: KindRequests, Window: w.name, Limit: w.limit, Count: requestCounts[i],
+			}
+		case bytes > 0 && w.byteLimit > 0 && byteCounts[i] > float64(w.byteLimit):
+			hit = Hit{
+				Kind: KindBytes, Window: w.name, Limit: w.byteLimit, Count: byteCounts[i],
+			}
+		}
+	}
+
+	counts := Counts{
+		Minute: requestCounts[0], Hour: requestCounts[1], Day: requestCounts[2],
+		MinuteBytes: byteCounts[0], HourBytes: byteCounts[1], DayBytes: byteCounts[2],
+	}
+
+	return counts, hit, hit.Window != ""
+}
+
 // get returns client's entry in the table, a new one if it has none, and
 // makes it the most recently seen.
 func (l *Limiter) get(client netip.Prefix) *Client {
@@ -353,30 +425,39 @@ func (l *Limiter) get(client netip.Prefix) *Client {
 	return c
 }
 
-// buckets returns c's buckets in the minute, the hour and the day.
+// buckets returns c's buckets of requests in the minute, the hour and the
+// day.
 func (c *Client) buckets() [3]*Buckets {
 	return [3]*Buckets{&c.Minute, &c.Hour, &c.Day}
 }
 
-// window is a length of time over which requests are counted, and the
-// most requests a client may make in it.
-type window struct {
-	name   string
-	length time.Duration
-	limit  int64
+// byteBuckets returns c's buckets of bytes in the minute, the hour and the
+// day.
+func (c *Client) byteBuckets() [3]*Buckets {
+	return [3]*Buckets{&c.MinuteBytes, &c.HourBytes, &c.DayBytes}
 }
 
-// add counts a request at now in a window of length, and returns the
-// client's requests in the window that ends at now: those in the bucket
-// under way, and those in the bucket before it weighted by how much of
-// that bucket the window still covers.
+// window is a length of time over which requests and bytes are counted,
+// and the most requests and the most bytes a client may have in it.
+type window struct {
+	name      string
+	length    time.Duration
+	limit     int64
+	byteLimit int64
+}
+
+// add counts n requests, or n bytes, at now in a window of length, and
+// returns the client's count in the window that ends at now: what is in
+// the bucket under way, and what is in the bucket before it weighted by
+// how much of that bucket the window still covers. With n zero it counts
+// nothing, and returns the count.
 //
 // Concurrent requests can be counted out of order, so now can be a moment
 // before the bucket under way began; such a request is counted in that
 // bucket. A request dated more than a second before it means the clock
 // was set back, and the buckets start afresh: otherwise the bucket before
 // would keep its full weight until the clock caught up.
-func (b *Buckets) add(now time.Time, length time.Duration) float64 {
+func (b *Buckets) add(now time.Time, length time.Duration, n int64) float64 {
 	if now.Before(b.Start.Add(-time.Second)) {
 		*b = Buckets{}
 	}
@@ -393,7 +474,7 @@ func (b *Buckets) add(now time.Time, length time.Duration) float64 {
 		b.Current = 0
 	}
 
-	b.Current++
+	b.Current += n
 
 	elapsed := max(now.Sub(b.Start), 0)
 	covered := 1 - float64(elapsed)/float64(length)

internal/ratelimit/ratelimit_test.go

@@ -71,13 +71,116 @@ func TestHitGivesTheLimitAndTheRequestsCounted(t *testing.T) {
 	// Over both limits; the minute's is named, with the four requests.
 	_, hit, over := limiter.Count(client, start)
 
-	want := ratelimit.Hit{Window: minute, Limit: limit, Requests: limit + 1}
+	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()
 
@@ -267,3 +370,18 @@ func wantCount(
 			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)
+	}
+}

internal/ratelimit/snapshot_test.go

@@ -64,6 +64,7 @@ func TestLoadEmptiesBucketsWhoseTimeHasPassed(t *testing.T) {
 
 	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 {
@@ -76,19 +77,25 @@ func TestLoadEmptiesBucketsWhoseTimeHasPassed(t *testing.T) {
 	}
 
 	// Two minutes on, the window that ends then covers neither of the
-	// minute's buckets, which are emptied; the hour's and the day's stay,
-	// and so does the history.
+	// 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 {
-		t.Errorf("loaded just under two minutes on with minute buckets %+v",
-			got.Minute)
+	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)
 	}
 }