internal/rproxy/limits_test.go

@@ -12,7 +12,6 @@ import (
 	"os"
 	"strconv"
 	"strings"
-	"sync"
 	"sync/atomic"
 	"testing"
 	"time"
@@ -383,6 +382,20 @@ type deadlineRecorder struct {
 	// stalled, if set, sends the part of the request body that never arrives. A
 	// read deadline that has passed ends a read of it, as on a connection.
 	stalled *io.PipeWriter
+	// flushError, if set, is what every flush returns, and nothing is flushed.
+	flushError error
+}
+
+// FlushError, which http.ResponseController calls to flush, returns flushError
+// if it is set, and otherwise flushes.
+func (w *deadlineRecorder) FlushError() error {
+	if w.flushError != nil {
+		return w.flushError
+	}
+
+	w.Flush()
+
+	return nil
 }
 
 // SetReadDeadline notes deadline.
@@ -448,6 +461,45 @@ func TestABodyLeftUnreadIsReadOnlyUntilTheRequestTimeLimit(t *testing.T) {
 	}
 }
 
+func TestABodyLeftUnreadIsReadNoMoreOnceTheTimeToSendTheAnswerRunsOut(t *testing.T) {
+	t.Parallel()
+
+	standIn := startWAF(t, allow)
+	app := startApp(t, answerBeforeTheBody)
+	_, _, server := startServer(t, standIn.URL, app.URL, nil)
+
+	// The visitor sends the first part and a little more, and then nothing.
+	stalled, stop := io.Pipe()
+
+	t.Cleanup(func() { _ = stop.Close() })
+
+	req := httptest.NewRequestWithContext(t.Context(), http.MethodPost, "/",
+		io.MultiReader(bytes.NewReader(make([]byte, questionBodyMaxBytes+1)), stalled))
+	req.ContentLength = 2 * questionBodyMaxBytes
+	req.Header.Set(contentLengthHeader, strconv.Itoa(2*questionBodyMaxBytes))
+
+	// Every flush fails, as once the time to send the answer has run out. The
+	// default time limits are far longer than the test, so rproxy records no
+	// refusal.
+	w := &deadlineRecorder{
+		ResponseRecorder: httptest.NewRecorder(),
+		stalled:          stop,
+		flushError:       os.ErrDeadlineExceeded,
+	}
+
+	before := time.Now()
+
+	server.Handler.ServeHTTP(w, req)
+
+	if w.Code != http.StatusOK {
+		t.Errorf("status %d, want the app's 200", w.Code)
+	}
+
+	if w.readDeadline.Before(before) || w.readDeadline.After(time.Now()) {
+		t.Errorf("the reading of the body ends at %s, want now", w.readDeadline)
+	}
+}
+
 func TestAnAnswerRProxyGivesItselfHasTheResponseTimeToBeSent(t *testing.T) {
 	t.Parallel()
 
@@ -918,49 +970,45 @@ func TestAVisitorThatStopsALittleShortHasTheConnectionClosedWhenTheAnswerIsCutOf
 ) {
 	t.Parallel()
 
-	const limit = 300 * time.Millisecond
+	// rproxy passes on a part sent at once well within this limit, even on a busy
+	// host.
+	const limit = 2 * time.Second
 
-	// With two parts, rproxy fails to send the first and then waits for the
-	// second, which the app sends only once rproxy has ended the request to it.
+	// The app sends the parts in atOnce with its status and headers. It sends the
+	// parts in late only once rproxy has ended the request to it, which rproxy
+	// does when it cuts off the answer, so that those never reach the visitor.
 	for _, tc := range []struct {
-		name  string
-		parts []string
+		name         string
+		atOnce, late []string
 	}{
-		{"with the body in one part", []string{greeting}},
-		{"with the body in two parts", []string{greeting[:5], greeting[5:]}},
+		{"with the body in one part", nil, []string{greeting}},
+		{"with the body in two parts", []string{greeting[:5]}, []string{greeting[5:]}},
 	} {
 		t.Run(tc.name, func(t *testing.T) {
 			t.Parallel()
 
-			headersRead := make(chan struct{})
-			// letAppSendBody also runs when the test ends, as the app's server waits
-			// for the app.
-			letAppSendBody := sync.OnceFunc(func() { close(headersRead) })
 			standIn := startWAF(t, allow)
-			// The app sends its body once the visitor has the status and headers and
-			// the response time limit has run out since, so that rproxy cannot send
-			// it. After each part, it reads the request's body, which ends when rproxy
-			// ends the request to it.
 			app := startApp(t, func(w http.ResponseWriter, r *http.Request) {
 				rc := http.NewResponseController(w)
 				_ = rc.EnableFullDuplex()
 
 				w.Header().Set(contentLengthHeader, strconv.Itoa(len(greeting)))
 
+				for _, part := range tc.atOnce {
+					_, _ = io.WriteString(w, part)
+				}
+
 				_ = rc.Flush()
 
-				<-headersRead
-				time.Sleep(limit)
+				// The request's body ends when rproxy ends the request to the app.
+				_, _ = io.Copy(io.Discard, r.Body)
 
-				for _, part := range tc.parts {
+				for _, part := range tc.late {
 					_, _ = io.WriteString(w, part)
 					_ = rc.Flush()
-					_, _ = io.Copy(io.Discard, r.Body)
 				}
 			})
 
-			t.Cleanup(letAppSendBody)
-
 			// The request time limit is far longer than dial's deadline.
 			addr, _ := startRProxy(t, standIn.URL, app.URL, map[string]string{
 				maxRequestTime:  "120000",
@@ -975,15 +1023,14 @@ func TestAVisitorThatStopsALittleShortHasTheConnectionClosedWhenTheAnswerIsCutOf
 				t.Fatalf("read the answer: %v", err)
 			}
 
-			letAppSendBody()
-
 			// If the connection stays open, dial's deadline ends this read with a
 			// timeout.
 			body, err := io.ReadAll(res.Body)
 			_ = res.Body.Close()
 
-			if len(body) != 0 || !errors.Is(err, io.ErrUnexpectedEOF) {
-				t.Errorf("read %q (%v), want the connection closed", body, err)
+			want := strings.Join(tc.atOnce, "")
+			if string(body) != want || !errors.Is(err, io.ErrUnexpectedEOF) {
+				t.Errorf("read %q (%v), want %q and the connection closed", body, err, want)
 			}
 		})
 	}

internal/rproxy/request.go

@@ -272,6 +272,10 @@ func (rq *request) modifyResponse(res *http.Response) error {
 
 	rq.closeAfterAnswer()
 
+	// An answer that arrives before the app has been sent the whole request ends
+	// the request, and SWRP_MAX_RESPONSE_TIME_MS starts.
+	rq.startResponseTimer(requestEnd)
+
 	res.Body = &responseBody{body: limitBody(res.Body, maxBytes), rq: rq}
 	rq.out.afterStatus = rq.setResponseDeadline
 
@@ -386,14 +390,18 @@ func (rq *request) finish(ctx context.Context) {
 
 	// After a cut-off answer, the reading of a body that has not arrived whole
 	// ends now. Go's server would otherwise keep the connection open while it
-	// reads the rest. A flush fails if the time to send the answer ran out, and
-	// every flush after a failed one fails too. ReverseProxy ignores the failure,
-	// but panics if it was still waiting for more of the app's answer; finish runs
-	// either way. Before a status, a flush would send a 200; after a 101,
-	// ReverseProxy has taken over the connection.
+	// reads the rest. modifyResponse sets afterStatus only when it passes the
+	// app's answer on. Any refusal after that has cut the answer off, since a
+	// refusal ends the request to the app. Any answer is cut off when the time to
+	// send it runs out: a flush then fails, and every flush after a failed one
+	// fails too. ReverseProxy ignores the failure, but panics if it was still
+	// waiting for more of the app's answer; finish runs either way. Before a
+	// status, a flush would send a 200; after a 101, ReverseProxy has taken over
+	// the connection.
 	switch {
-	case rq.out.status != 0 && rq.out.status != http.StatusSwitchingProtocols &&
-		errors.Is(rq.rc.Flush(), os.ErrDeadlineExceeded):
+	case rq.out.afterStatus != nil && rq.refused.Load() != 0,
+		rq.out.status != 0 && rq.out.status != http.StatusSwitchingProtocols &&
+			errors.Is(rq.rc.Flush(), os.ErrDeadlineExceeded):
 		rq.stopPassingOnBody(time.Now())
 	case rq.refused.Load() == 0:
 		rq.stopPassingOnBody(rq.requestDeadline())
@@ -553,19 +561,29 @@ func (rq *request) wroteRequest(info httptrace.WroteRequestInfo) {
 	requestEnd := rq.endRequest()
 
 	rq.mu.Lock()
-	defer rq.mu.Unlock()
+	stopTimer(rq.requestTimer)
+	rq.mu.Unlock()
 
-	if rq.timersStopped {
+	rq.startResponseTimer(requestEnd)
+}
+
+// startResponseTimer starts SWRP_MAX_RESPONSE_TIME_MS, which runs from
+// requestEnd, unless it has started already.
+func (rq *request) startResponseTimer(requestEnd time.Time) {
+	limit := rq.h.config.MaxResponseTime
+	if limit == 0 {
 		return
 	}
 
-	stopTimer(rq.requestTimer)
+	rq.mu.Lock()
+	defer rq.mu.Unlock()
 
-	limit := rq.h.config.MaxResponseTime
-	if limit > 0 {
-		rq.responseTimer = time.AfterFunc(time.Until(requestEnd.Add(limit)),
-			rq.responseTimedOut)
+	if rq.timersStopped || rq.responseTimer != nil {
+		return
 	}
+
+	rq.responseTimer = time.AfterFunc(time.Until(requestEnd.Add(limit)),
+		rq.responseTimedOut)
 }
 
 // responseTimedOut is called when SWRP_MAX_RESPONSE_TIME_MS runs out before the