rproxy: read the rest of a Connection: close visitor's body after an early answer (closes #373)

Go's server reads up to 256 KiB of a body left unread once the handler
returns, but not a body with a declared length from a visitor that sent
Connection: close. It closes that connection with the rest unread, and
the visitor, still sending, gets a reset. rproxy now reads that rest
itself, until the deadline already set on the connection, through a
MaxBytesReader made with Go's own ResponseWriter. With more than 256 KiB
left, that has Go's server tell the visitor it sends no more and wait
500 ms before it closes, as for any other visitor.

Model: opus-5-5

internal/rproxy/limits_test.go

@@ -93,10 +93,15 @@ func TestABodyAnnouncedOverItsLimitIsRefusedOnceWAFAllowsIt(t *testing.T) {
 				},
 			)
 
-			got := visit(t, newRequest(t, http.MethodPost, addr, "/",
-				bytes.NewReader(make([]byte, tc.size))))
-			if got.status != http.StatusRequestEntityTooLarge {
-				t.Errorf("answered %+v, want 413", got)
+			// The visitor sends no Connection: close, so Go's server reads the rest
+			// of the body before it closes the connection. Closed with the rest
+			// unread, the connection is reset while the visitor still sends.
+			conn := dial(t, addr)
+			send(t, conn, post+announcedLength(tc.size)+
+				strings.Repeat("x", tc.size))
+
+			if status := readStatus(t, conn); status != http.StatusRequestEntityTooLarge {
+				t.Errorf("status %d, want 413", status)
 			}
 
 			if asked := standIn.question(t); len(asked.body) != tc.asked {
@@ -131,6 +136,73 @@ func TestADenyOutranksTheBodySizeLimit(t *testing.T) {
 	}
 }
 
+func TestAVisitorSendingConnectionCloseReadsAnAnswerBeforeTheRestOfItsBody(
+	t *testing.T,
+) {
+	t.Parallel()
+
+	for _, tc := range []struct {
+		name   string
+		waf    http.HandlerFunc
+		rest   int
+		status int
+	}{
+		{"over the size limit", allow, 65 << 10, http.StatusRequestEntityTooLarge},
+		{"denied", deny, 65 << 10, http.StatusForbidden},
+		// rproxy reads 256 KiB of this rest. What it leaves unread must fit in the
+		// connection's buffers. Otherwise the visitor is still sending when the
+		// connection closes, and gets a reset.
+		{
+			"with more than 256 KiB left", allow, 272 << 10,
+			http.StatusRequestEntityTooLarge,
+		},
+	} {
+		t.Run(tc.name, func(t *testing.T) {
+			t.Parallel()
+
+			standIn := startWAF(t, tc.waf)
+			app := startApp(t, func(http.ResponseWriter, *http.Request) {
+				t.Error("the app was asked")
+			})
+			addr, _ := startRProxy(t, standIn.URL, app.URL,
+				map[string]string{maxRequestSize: "128K"})
+
+			conn := dial(t, addr)
+			send(t, conn, post+"Connection: close\r\n"+
+				announcedLength(questionBodyMaxBytes+tc.rest)+
+				strings.Repeat("x", questionBodyMaxBytes))
+
+			reader := bufio.NewReader(conn)
+
+			res, err := http.ReadResponse(reader, nil)
+			if err != nil {
+				t.Fatalf("read the answer: %v", err)
+			}
+
+			_, _ = io.Copy(io.Discard, res.Body)
+			_ = res.Body.Close()
+
+			if res.StatusCode != tc.status {
+				t.Errorf("status %d, want %d", res.StatusCode, tc.status)
+			}
+
+			// The rest goes in small parts, so that a reset fails a later one.
+			for sent := 0; sent < tc.rest; sent += 1 << 10 {
+				send(t, conn, strings.Repeat("x", 1<<10))
+			}
+
+			got, err := io.ReadAll(reader)
+			if err != nil || len(got) != 0 {
+				t.Errorf(
+					"after the answer, read %q and %v, want the end",
+					got,
+					err,
+				)
+			}
+		})
+	}
+}
+
 func TestABodyNotAnnouncedOverItsLimitIsRefused(t *testing.T) {
 	t.Parallel()
 

internal/rproxy/request.go

@@ -38,6 +38,10 @@ const plainText = "text/plain; charset=utf-8"
 // and the visitor then gets a closed connection and no answer.
 const minSendTime = 10 * time.Second
 
+// maxBodyReadAfterAnswer is the most of a body left unread that Go's server
+// reads once the handler returns, for a visitor without Connection: close.
+const maxBodyReadAfterAnswer = 256 << 10
+
 // errResponseTooLate refuses an app's answer that arrives once
 // SWW_RPROXY_MAX_RESPONSE_TIME_MS has run out.
 var errResponseTooLate = errors.New(
@@ -338,8 +342,8 @@ func (rq *request) answer(status int) {
 	}
 
 	// A visitor found too slow is read no more; any other may go on sending until
-	// its time is up, so that Go's server can read the rest of the body and end
-	// the request cleanly.
+	// its time is up, so that the rest of the body is read and the request ends
+	// cleanly: by Go's server, or by readRestOfBody.
 	deadline := rq.requestDeadline()
 	if status == http.StatusRequestTimeout {
 		deadline = time.Now()
@@ -523,6 +527,31 @@ func (rq *request) stopReadingBody(deadline time.Time) {
 	_ = rq.rc.SetReadDeadline(deadline)
 }
 
+// readRestOfBody reads the rest of the body once the visitor has been
+// answered, when the visitor sent Connection: close and the body has a declared
+// length. Go's server closes such a connection with the rest unread. The
+// visitor, still sending, gets a reset and may never read the answer. For any
+// other visitor, and for a chunked body, Go's server reads up to 256 KiB of the
+// rest itself. When more is left, it tells the visitor that it sends no more,
+// then waits 500 ms before it closes the connection. readRestOfBody does the
+// same, and its read ends at the deadline already set on the connection.
+// MaxBytesReader, made with Go's own ResponseWriter, has Go's server end the
+// connection that way when more is left.
+func (rq *request) readRestOfBody() {
+	if !rq.in.Close || rq.in.ContentLength <= 0 || rq.body.received.Load() ||
+		rq.out.status == 0 ||
+		rq.out.status == http.StatusSwitchingProtocols {
+		return
+	}
+
+	rest := http.MaxBytesReader(
+		rq.out.ResponseWriter,
+		rq.in.Body,
+		maxBodyReadAfterAnswer,
+	)
+	_, _ = io.Copy(io.Discard, rest)
+}
+
 // stopPassingOnBody ends the transport's reading of a visitor's body that has
 // not arrived whole, and lets Go's server read the rest until deadline. Go's
 // server reads it once the handler returns, but with no deadline at all if

internal/rproxy/rproxy.go

@@ -169,6 +169,7 @@ func newWAFTransport() *http.Transport {
 // it writes the request's log line and sends its report.
 func (h *handler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 	rq := h.newRequest(w, r)
+	defer rq.readRestOfBody()
 	defer rq.finish(r.Context())
 
 	if r.Method == http.MethodGet && r.URL.Path == waf.HealthPath {