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

README.md

@@ -110,13 +110,16 @@ in `bin/state` unless `SWWAF_STATE_DIR` is set, and the default rule file of
   allowance (see "State files" below).
 - Counts each client's bytes over a minute, an hour and a day, in the same way:
   once a request passed to the app has ended, the body bytes of its answer, of
-  the request, or of both, as `SWWAF_BYTES_COUNT` says. Bytes that take the
-  client over one of the byte limits below break that limit, and ban the client
-  as a broken rate limit does, so that its next request is refused. The byte
-  limits never cut an answer short: the one whose bytes break a limit has
-  already been passed on. They leave out what the rate limits leave out: a
-  client in `SWWAF_ALLOW_NETS` or `SWWAF_RATE_LIMIT_EXEMPT_NETS`, and a request
-  for a path `SWWAF_RATE_LIMIT_EXEMPT_PATHS` exempts.
+  the request, or of both, as `SWWAF_BYTES_COUNT` says. For a WebSocket, or any
+  other upgraded connection, what it carried from the app counts with the
+  answer, and what it carried from the client with the request, once it closes.
+  Bytes that take the client over one of the byte limits below break that limit,
+  and ban the client as a broken rate limit does, so that its next request is
+  refused. The byte limits never cut an answer or an upgraded connection short:
+  the one whose bytes break a limit has already been passed on, or has closed.
+  They leave out what the rate limits leave out: a client in `SWWAF_ALLOW_NETS`
+  or `SWWAF_RATE_LIMIT_EXEMPT_NETS`, and a request for a path
+  `SWWAF_RATE_LIMIT_EXEMPT_PATHS` exempts.
 - Bans a client that breaks a rate limit or a byte limit, as "Bans" in
   [`SPEC.md`](SPEC.md) describes: the first ban lasts an hour, and a limit
   broken again within a day of a ban ending bans for three times as long as that

internal/proxy/bans.go

@@ -58,24 +58,33 @@ func (rq *request) limitBroken(now time.Time) bool {
 // response has ended, and notes the client's byte totals for the log line;
 // its requests stay there as the rate limits counted them. The bytes are
 // the response's body bytes, the request's, or both, as SWWAF_BYTES_COUNT
-// says. Only a request passed to the app has them counted, and only one
-// the rate limits counted; in observe mode, not one that enforce mode
-// would have refused. Bytes that take the client over a byte limit break
-// it; the response was passed on whole.
+// says; for an upgraded connection, such as a WebSocket, which has closed
+// by then, what it carried from the app counts with the response's and
+// what it carried from the client with the request's. Only a request
+// passed to the app has them counted, and only one the rate limits
+// counted; in observe mode, not one that enforce mode would have refused.
+// Bytes that take the client over a byte limit break it; the response was
+// passed on whole.
 func (rq *request) countBytes() {
 	if !rq.counted || rq.line.WouldAction != "" {
 		return
 	}
 
+	response, request := rq.out.bytes, rq.requestBytes()
+	if rq.upgraded != nil {
+		response += rq.upgraded.fromApp.Load()
+		request += rq.upgraded.toApp.Load()
+	}
+
 	var bytes int64
 
 	switch rq.h.config.BytesCount {
 	case "response":
-		bytes = rq.out.bytes
+		bytes = response
 	case "request":
-		bytes = rq.requestBytes()
+		bytes = request
 	default: // both
-		bytes = rq.out.bytes + rq.requestBytes()
+		bytes = response + request
 	}
 
 	now := rq.h.now()

internal/proxy/bodies.go

@@ -103,6 +103,48 @@ func (b *responseBody) Close() error {
 	return b.body.Close()
 }
 
+// upgradedConn is the connection to the app once the app has switched
+// protocols, as for a WebSocket. ReverseProxy writes to it what the client
+// sends and reads from it what the app sends, on goroutines of its own,
+// until the connection closes; it counts the bytes each way, for the byte
+// limits.
+type upgradedConn struct {
+	io.ReadWriteCloser
+
+	// fromApp is how many bytes the app has sent, and toApp how many the
+	// client has.
+	fromApp atomic.Int64
+	toApp   atomic.Int64
+}
+
+// Read reads what the app sends.
+func (c *upgradedConn) Read(p []byte) (int, error) {
+	n, err := c.ReadWriteCloser.Read(p)
+	c.fromApp.Add(int64(n))
+
+	return n, err
+}
+
+// Write sends the app what the client sent.
+func (c *upgradedConn) Write(p []byte) (int, error) {
+	n, err := c.ReadWriteCloser.Write(p)
+	c.toApp.Add(int64(n))
+
+	return n, err
+}
+
+// CloseWrite tells the app that the client sends no more, while what the
+// app sends still passes. ReverseProxy calls it once the client has
+// stopped sending, and closes the connection there if it is not supported.
+func (c *upgradedConn) CloseWrite() error {
+	conn, ok := c.ReadWriteCloser.(interface{ CloseWrite() error })
+	if !ok {
+		return http.ErrNotSupported
+	}
+
+	return conn.CloseWrite()
+}
+
 // limitBody returns body, cut off with an *http.MaxBytesError after
 // maxBytes, or unchanged if maxBytes is zero, which is off.
 func limitBody(body io.ReadCloser, maxBytes int64) io.ReadCloser {

internal/proxy/bytelimits_test.go

@@ -1,6 +1,7 @@
 package proxy_test
 
 import (
+	"bufio"
 	"io"
 	"net/http"
 	"net/netip"
@@ -22,6 +23,13 @@ const (
 	bytesCount          = "SWWAF_BYTES_COUNT"
 )
 
+// The values of SWWAF_BYTES_COUNT.
+const (
+	countResponse = "response"
+	countRequest  = "request"
+	countBoth     = "both"
+)
+
 const (
 	// bodyBytes is the size of the body of each request these tests send
 	// with one, and answerBytes that of each answer of the app.
@@ -105,6 +113,59 @@ func TestResponseOverAByteLimitByItselfIsPassedOnWhole(t *testing.T) {
 	s.get(client, http.StatusForbidden, requestlog.ActionBanned)
 }
 
+func TestBytesOfAnAnswerThatBreaksOffAreCounted(t *testing.T) {
+	t.Parallel()
+
+	s, clk, _, _ := startAppWithAlerts(t, breakOff, map[string]string{
+		bytesLimitPerMinute: "50",
+	})
+	expires := requestlog.FormatTime(clk.Now().Add(time.Hour))
+
+	// The 70 bytes passed on before the app broke off are over the limit of
+	// 50, and ban the client for an hour.
+	line, got := s.requestWithBody(http.MethodGet, client, "/", "", "", http.StatusOK,
+		requestlog.ActionUpstreamError)
+	if len(got.body) != answerBytes || line.LimitHit != minuteBytes ||
+		line.BanExpires != expires {
+		t.Errorf("got %d bytes, and the log line has limit_hit %q and ban_expires %q, "+
+			"want %d, %s and %s", len(got.body), line.LimitHit, line.BanExpires,
+			answerBytes, minuteBytes, expires)
+	}
+
+	s.get(client, http.StatusForbidden, requestlog.ActionBanned)
+}
+
+func TestWebSocketBytesAreCountedOnceItCloses(t *testing.T) {
+	t.Parallel()
+
+	for _, tc := range []struct {
+		setting string
+		counted float64
+	}{
+		{countResponse, answerBytes},
+		{countRequest, bodyBytes},
+		{countBoth, bodyBytes + answerBytes},
+	} {
+		t.Run(tc.setting, func(t *testing.T) {
+			t.Parallel()
+
+			s, _, _, _ := startAppWithAlerts(t, answerAfterUpgrade, map[string]string{
+				bytesLimitPerMinute: "29", bytesCount: tc.setting,
+			})
+
+			// The client sends 30 bytes and the app 70, each over the limit
+			// of 29, which bans the client once the WebSocket has closed.
+			line := s.webSocket()
+			if line.LimitHit != minuteBytes || line.Counts.MinuteBytes != tc.counted {
+				t.Errorf("log line has limit_hit %q and minute_bytes %v, want %s and %v",
+					line.LimitHit, line.Counts.MinuteBytes, minuteBytes, tc.counted)
+			}
+
+			s.get(client, http.StatusForbidden, requestlog.ActionBanned)
+		})
+	}
+}
+
 func TestBytesCountSaysWhichBytesCount(t *testing.T) {
 	t.Parallel()
 
@@ -115,9 +176,9 @@ func TestBytesCountSaysWhichBytesCount(t *testing.T) {
 		each     float64
 		breaking int
 	}{
-		{"response", answerBytes, 2},
-		{"request", bodyBytes, 4},
-		{"both", bodyBytes + answerBytes, 1},
+		{countResponse, answerBytes, 2},
+		{countRequest, bodyBytes, 4},
+		{countBoth, bodyBytes + answerBytes, 1},
 	} {
 		t.Run(tc.setting, func(t *testing.T) {
 			t.Parallel()
@@ -336,6 +397,85 @@ func readAndAnswer(w http.ResponseWriter, r *http.Request) {
 	_, _ = io.WriteString(w, strings.Repeat("a", answerBytes))
 }
 
+// breakOff is an app that announces an answer of twice answerBytes, and
+// breaks off after answerBytes.
+func breakOff(w http.ResponseWriter, _ *http.Request) {
+	w.Header().Set("Content-Length", strconv.Itoa(2*answerBytes))
+	_, _ = io.WriteString(w, strings.Repeat("a", answerBytes))
+}
+
+// answerAfterUpgrade is an app that switches protocols, as for a
+// WebSocket, and then answers each line it receives with a line of
+// answerBytes.
+func answerAfterUpgrade(w http.ResponseWriter, _ *http.Request) {
+	conn, buffered, err := http.NewResponseController(w).Hijack()
+	if err != nil {
+		return
+	}
+
+	defer func() {
+		_ = conn.Close()
+	}()
+
+	_, _ = buffered.WriteString("HTTP/1.1 101 Switching Protocols\r\n" +
+		"Connection: Upgrade\r\nUpgrade: websocket\r\n\r\n")
+	_ = buffered.Flush()
+
+	for {
+		_, err := buffered.ReadString('\n')
+		if err != nil {
+			return
+		}
+
+		_, _ = buffered.WriteString(strings.Repeat("a", answerBytes-1) + "\n")
+		_ = buffered.Flush()
+	}
+}
+
+// webSocket opens a WebSocket from client to answerAfterUpgrade, sends a
+// line of bodyBytes on it, reads the answer, and closes it. It checks the
+// answer, and the log line as request does, and returns the log line.
+func (s *sender) webSocket() logLine {
+	s.t.Helper()
+
+	conn := dial(s.t, s.addr)
+	send(s.t, conn, "GET /socket HTTP/1.1\r\nHost: "+appHost+"\r\n"+forwardedFor+
+		": "+client+"\r\nConnection: Upgrade\r\nUpgrade: websocket\r\n\r\n")
+
+	err := conn.SetReadDeadline(time.Now().Add(waitLimit))
+	if err != nil {
+		s.t.Fatalf("set read deadline: %v", err)
+	}
+
+	reader := bufio.NewReader(conn)
+
+	res, err := http.ReadResponse(reader, nil)
+	if err != nil {
+		s.t.Fatalf("read the answer to the upgrade: %v", err)
+	}
+
+	_ = res.Body.Close()
+
+	if res.StatusCode != http.StatusSwitchingProtocols {
+		s.t.Fatalf("status %d, want %d", res.StatusCode, http.StatusSwitchingProtocols)
+	}
+
+	send(s.t, conn, strings.Repeat("u", bodyBytes-1)+"\n")
+
+	got, err := reader.ReadString('\n')
+	if err != nil || len(got) != answerBytes {
+		s.t.Errorf("got %d bytes (%v), want %d", len(got), err, answerBytes)
+	}
+
+	_ = conn.Close()
+
+	line := s.out.requestLines(s.t, s.sent+1)[s.sent]
+	s.sent++
+	wantLine(s.t, line, http.StatusSwitchingProtocols, requestlog.ActionForward)
+
+	return line
+}
+
 // startWithAnswers is startAppWithAlerts in front of readAndAnswer, for a
 // test that looks at neither the server nor the alerts.
 func startWithAnswers(t *testing.T, env map[string]string) (*sender, *clock) {

internal/proxy/request.go

@@ -3,6 +3,7 @@ package proxy
 import (
 	"context"
 	"errors"
+	"io"
 	"net/http"
 	"net/http/httptrace"
 	"net/http/httputil"
@@ -68,6 +69,9 @@ type request struct {
 	refused atomic.Pointer[refusal]
 	// complete is true once the app's whole answer has been passed on.
 	complete bool
+	// upgraded is the connection to the app once the app has switched
+	// protocols, as for a WebSocket, and nil otherwise.
+	upgraded *upgradedConn
 
 	// mu guards what follows. The timeouts run on goroutines of their
 	// own, and the transport starts and stops them, and notes the times
@@ -328,11 +332,18 @@ func (rq *request) modifyResponse(res *http.Response) error {
 	if res.StatusCode == http.StatusSwitchingProtocols {
 		// An upgraded connection, such as a WebSocket, is not cut by the
 		// timeouts. ReverseProxy writes this answer straight to the
-		// connection it takes over, not through rq.out.
+		// connection it takes over, not through rq.out, and then copies
+		// what passes each way through res.Body, the connection to the app.
 		rq.stopTimers()
 		rq.out.status = res.StatusCode
 		rq.line.Websocket = true
 
+		conn, ok := res.Body.(io.ReadWriteCloser)
+		if ok {
+			rq.upgraded = &upgradedConn{ReadWriteCloser: conn}
+			res.Body = rq.upgraded
+		}
+
 		return nil
 	}