net/http: do not discard body content when closing it within request handlers

(*body).Close() internally tries to discard the content of a request
body up to 256 KB. We rely on this behavior to allow connection re-use,
by calling (*body).Close() when our request handler exits.

Unfortunately, this causes an unfortunate side-effect where we would
prematurely try to discard a body content when (*body).Close() is called
from within a request handler.

There should not be a good reason for (*body).Close() to do this when
called from within a request handler. As such, this CL modifies
(*body).Close() to not discard body contents when called from within a
request handler. Note that when a request handler exits, it will still
try to discard the body content for connection re-use.

For #75933

Change-Id: I71d2431a540579184066dd35d3da49d6c85c3daf
Reviewed-on: https://go-review.googlesource.com/c/go/+/720380
LUCI-TryBot-Result: Go LUCI <[email protected]>
Reviewed-by: Nicholas Husin <[email protected]>
Reviewed-by: Damien Neil <[email protected]>

Author: nicholashusin

src/net/http/serve_test.go

@@ -2150,118 +2150,137 @@ func TestServerUnreadRequestBodyLarge(t *testing.T) {
 	}
 }
 
-type handlerBodyCloseTest struct {
+type bodyDiscardTest struct {
 	bodySize     int
 	bodyChunked  bool
 	reqConnClose bool
 
-	wantEOFSearch bool // should Handler's Body.Close do Reads, looking for EOF?
-	wantNextReq   bool // should it find the next request on the same conn?
+	shouldDiscardBody bool // should the handler discard body after it exits?
 }
 
-func (t handlerBodyCloseTest) connectionHeader() string {
+func (t bodyDiscardTest) connectionHeader() string {
 	if t.reqConnClose {
 		return "Connection: close\r\n"
 	}
 	return ""
 }
 
-var handlerBodyCloseTests = [...]handlerBodyCloseTest{
-	// Small enough to slurp past to the next request +
-	// has Content-Length.
+var bodyDiscardTests = [...]bodyDiscardTest{
+	// Have:
+	// - Small body.
+	// - Content-Length defined.
+	// Should:
+	// - Discard remaining body.
 	0: {
-		bodySize:      20 << 10,
-		bodyChunked:   false,
-		reqConnClose:  false,
-		wantEOFSearch: true,
-		wantNextReq:   true,
+		bodySize:          20 << 10,
+		bodyChunked:       false,
+		reqConnClose:      false,
+		shouldDiscardBody: true,
 	},
 
-	// Small enough to slurp past to the next request +
-	// is chunked.
+	// Have:
+	// - Small body.
+	// - Chunked (no Content-Length defined).
+	// Should:
+	// - Discard remaining body.
 	1: {
-		bodySize:      20 << 10,
-		bodyChunked:   true,
-		reqConnClose:  false,
-		wantEOFSearch: true,
-		wantNextReq:   true,
+		bodySize:          20 << 10,
+		bodyChunked:       true,
+		reqConnClose:      false,
+		shouldDiscardBody: true,
 	},
 
-	// Small enough to slurp past to the next request +
-	// has Content-Length +
-	// declares Connection: close (so pointless to read more).
+	// Have:
+	// - Small body.
+	// - Content-Length defined.
+	// - Connection: close.
+	// Should:
+	// - Not discard remaining body (no point as Connection: close).
 	2: {
-		bodySize:      20 << 10,
-		bodyChunked:   false,
-		reqConnClose:  true,
-		wantEOFSearch: false,
-		wantNextReq:   false,
+		bodySize:          20 << 10,
+		bodyChunked:       false,
+		reqConnClose:      true,
+		shouldDiscardBody: false,
 	},
 
-	// Small enough to slurp past to the next request +
-	// declares Connection: close,
-	// but chunked, so it might have trailers.
-	// TODO: maybe skip this search if no trailers were declared
-	// in the headers.
+	// Have:
+	// - Small body.
+	// - Chunked (no Content-Length defined).
+	// - Connection: close.
+	// Should:
+	// - Discard remaining body (chunked, so it might have trailers).
+	//
+	// TODO: maybe skip this if no trailers were declared in the headers.
 	3: {
-		bodySize:      20 << 10,
-		bodyChunked:   true,
-		reqConnClose:  true,
-		wantEOFSearch: true,
-		wantNextReq:   false,
+		bodySize:          20 << 10,
+		bodyChunked:       true,
+		reqConnClose:      true,
+		shouldDiscardBody: true,
 	},
 
-	// Big with Content-Length, so give up immediately if we know it's too big.
+	// Have:
+	// - Large body.
+	// - Content-Length defined.
+	// Should:
+	// - Not discard remaining body (we know it is too large from Content-Length).
 	4: {
-		bodySize:      1 << 20,
-		bodyChunked:   false, // has a Content-Length
-		reqConnClose:  false,
-		wantEOFSearch: false,
-		wantNextReq:   false,
+		bodySize:          1 << 20,
+		bodyChunked:       false,
+		reqConnClose:      false,
+		shouldDiscardBody: false,
 	},
 
-	// Big chunked, so read a bit before giving up.
+	// Have:
+	// - Large body.
+	// - Chunked (no Content-Length defined).
+	// Should:
+	// - Discard remaining body (chunked, so we try up to a limit before giving up).
 	5: {
-		bodySize:      1 << 20,
-		bodyChunked:   true,
-		reqConnClose:  false,
-		wantEOFSearch: true,
-		wantNextReq:   false,
+		bodySize:          1 << 20,
+		bodyChunked:       true,
+		reqConnClose:      false,
+		shouldDiscardBody: true,
 	},
 
-	// Big with Connection: close, but chunked, so search for trailers.
-	// TODO: maybe skip this search if no trailers were declared
-	// in the headers.
+	// Have:
+	// - Large body.
+	// - Content-Length defined.
+	// - Connection: close.
+	// Should:
+	// - Not discard remaining body (Connection: Close, and Content-Length is too large).
 	6: {
-		bodySize:      1 << 20,
-		bodyChunked:   true,
-		reqConnClose:  true,
-		wantEOFSearch: true,
-		wantNextReq:   false,
+		bodySize:          1 << 20,
+		bodyChunked:       false,
+		reqConnClose:      true,
+		shouldDiscardBody: false,
 	},
-
-	// Big with Connection: close, so don't do any reads on Close.
-	// With Content-Length.
+	// Have:
+	// - Large body.
+	// - Chunked (no Content-Length defined).
+	// - Connection: close.
+	// Should:
+	// - Discard remaining body (chunked, so it might have trailers).
+	//
+	// TODO: maybe skip this if no trailers were declared in the headers.
 	7: {
-		bodySize:      1 << 20,
-		bodyChunked:   false,
-		reqConnClose:  true,
-		wantEOFSearch: false,
-		wantNextReq:   false,
+		bodySize:          1 << 20,
+		bodyChunked:       true,
+		reqConnClose:      true,
+		shouldDiscardBody: true,
 	},
 }
 
-func TestHandlerBodyClose(t *testing.T) {
+func TestBodyDiscard(t *testing.T) {
 	setParallel(t)
 	if testing.Short() && testenv.Builder() == "" {
 		t.Skip("skipping in -short mode")
 	}
-	for i, tt := range handlerBodyCloseTests {
-		testHandlerBodyClose(t, i, tt)
+	for i, tt := range bodyDiscardTests {
+		testBodyDiscard(t, i, tt)
 	}
 }
 
-func testHandlerBodyClose(t *testing.T, i int, tt handlerBodyCloseTest) {
+func testBodyDiscard(t *testing.T, i int, tt bodyDiscardTest) {
 	conn := new(testConn)
 	body := strings.Repeat("x", tt.bodySize)
 	if tt.bodyChunked {
@@ -2275,12 +2294,12 @@ func testHandlerBodyClose(t *testing.T, i int, tt handlerBodyCloseTest) {
 		cw.Close()
 		conn.readBuf.WriteString("\r\n")
 	} else {
-		conn.readBuf.Write([]byte(fmt.Sprintf(
+		conn.readBuf.Write(fmt.Appendf(nil,
 			"POST / HTTP/1.1\r\n"+
 				"Host: test\r\n"+
 				tt.connectionHeader()+
 				"Content-Length: %d\r\n"+
-				"\r\n", len(body))))
+				"\r\n", len(body)))
 		conn.readBuf.Write([]byte(body))
 	}
 	if !tt.reqConnClose {
@@ -2295,26 +2314,23 @@ func testHandlerBodyClose(t *testing.T, i int, tt handlerBodyCloseTest) {
 	}
 
 	ls := &oneConnListener{conn}
-	var numReqs int
-	var size0, size1 int
+	var initialSize, closedSize, exitSize int
 	go Serve(ls, HandlerFunc(func(rw ResponseWriter, req *Request) {
-		numReqs++
-		if numReqs == 1 {
-			size0 = readBufLen()
-			req.Body.Close()
-			size1 = readBufLen()
-		}
+		initialSize = readBufLen()
+		req.Body.Close()
+		closedSize = readBufLen()
 	}))
 	<-conn.closec
-	if numReqs < 1 || numReqs > 2 {
-		t.Fatalf("%d. bug in test. unexpected number of requests = %d", i, numReqs)
+	exitSize = readBufLen()
+
+	if initialSize != closedSize {
+		t.Errorf("%d. Close() within request handler should be a no-op, but body size went from %d to %d", i, initialSize, closedSize)
 	}
-	didSearch := size0 != size1
-	if didSearch != tt.wantEOFSearch {
-		t.Errorf("%d. did EOF search = %v; want %v (size went from %d to %d)", i, didSearch, !didSearch, size0, size1)
+	if tt.shouldDiscardBody && closedSize <= exitSize {
+		t.Errorf("%d. want body content to be discarded upon request handler exit, but size went from %d to %d", i, closedSize, exitSize)
 	}
-	if tt.wantNextReq && numReqs != 2 {
-		t.Errorf("%d. numReq = %d; want 2", i, numReqs)
+	if !tt.shouldDiscardBody && closedSize != exitSize {
+		t.Errorf("%d. want body content to not be discarded upon request handler exit, but size went from %d to %d", i, closedSize, exitSize)
 	}
 }
 

src/net/http/server.go

@@ -1077,9 +1077,6 @@ func (c *conn) readRequest(ctx context.Context) (w *response, err error) {
 	req.ctx = ctx
 	req.RemoteAddr = c.remoteAddr
 	req.TLS = c.tlsState
-	if body, ok := req.Body.(*body); ok {
-		body.doEarlyClose = true
-	}
 
 	// Adjust the read deadline if necessary.
 	if !hdrDeadline.Equal(wholeReqDeadline) {
@@ -1711,9 +1708,11 @@ func (w *response) finishRequest() {
 
 	w.conn.r.abortPendingRead()
 
-	// Close the body (regardless of w.closeAfterReply) so we can
-	// re-use its bufio.Reader later safely.
-	w.reqBody.Close()
+	// Try to discard the body (regardless of w.closeAfterReply), so we can
+	// potentially reuse it in the same connection.
+	if b, ok := w.reqBody.(*body); ok {
+		b.tryDiscardBody()
+	}
 
 	if w.req.MultipartForm != nil {
 		w.req.MultipartForm.RemoveAll()
@@ -1741,16 +1740,16 @@ func (w *response) shouldReuseConnection() bool {
 		return false
 	}
 
-	if w.closedRequestBodyEarly() {
+	if w.didIncompleteDiscard() {
 		return false
 	}
 
 	return true
 }
 
-func (w *response) closedRequestBodyEarly() bool {
+func (w *response) didIncompleteDiscard() bool {
 	body, ok := w.req.Body.(*body)
-	return ok && body.didEarlyClose()
+	return ok && body.didIncompleteDiscard()
 }
 
 func (w *response) Flush() {
@@ -2106,6 +2105,18 @@ func (c *conn) serve(ctx context.Context) {
 		// But we're not going to implement HTTP pipelining because it
 		// was never deployed in the wild and the answer is HTTP/2.
 		inFlightResponse = w
+		// Ensure that Close() invocations within request handlers do not
+		// discard the body.
+		if b, ok := w.reqBody.(*body); ok {
+			b.mu.Lock()
+			b.inRequestHandler = true
+			b.mu.Unlock()
+			defer func() {
+				b.mu.Lock()
+				b.inRequestHandler = false
+				b.mu.Unlock()
+			}()
+		}
 		serverHandler{c.server}.ServeHTTP(w, w.req)
 		inFlightResponse = nil
 		w.cancelCtx()
@@ -2116,7 +2127,7 @@ func (c *conn) serve(ctx context.Context) {
 		w.finishRequest()
 		c.rwc.SetWriteDeadline(time.Time{})
 		if !w.shouldReuseConnection() {
-			if w.requestBodyLimitHit || w.closedRequestBodyEarly() {
+			if w.requestBodyLimitHit || w.didIncompleteDiscard() {
 				c.closeWriteAndWait()
 			}
 			return