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 <golang-scoped@luci-project-accounts.iam.gserviceaccount.com>
Reviewed-by: Nicholas Husin <husin@google.com>
Reviewed-by: Damien Neil <dneil@google.com>

1 files changed, 61 insertions, 46 deletions

diff --git a/src/net/http/transfer.go b/src/net/http/transfer.go
index 675551287f..9b972aaa44 100644
--- a/src/net/http/transfer.go
+++ b/src/net/http/transfer.go
@@ -809,17 +809,17 @@ func fixTrailer(header Header, chunked bool) (Header, error) {
 // Close ensures that the body has been fully read
 // and then reads the trailer if necessary.
 type body struct {
-	src          io.Reader
-	hdr          any           // non-nil (Response or Request) value means read trailer
-	r            *bufio.Reader // underlying wire-format reader for the trailer
-	closing      bool          // is the connection to be closed after reading body?
-	doEarlyClose bool          // whether Close should stop early
+	src     io.Reader
+	hdr     any           // non-nil (Response or Request) value means read trailer
+	r       *bufio.Reader // underlying wire-format reader for the trailer
+	closing bool          // is the connection to be closed after reading body?
 
-	mu         sync.Mutex // guards following, and calls to Read and Close
-	sawEOF     bool
-	closed     bool
-	earlyClose bool   // Close called and we didn't read to the end of src
-	onHitEOF   func() // if non-nil, func to call when EOF is Read
+	mu                sync.Mutex // guards following, and calls to Read and Close
+	sawEOF            bool
+	closed            bool
+	incompleteDiscard bool   // if true, we failed to discard the body content completely
+	inRequestHandler  bool   // used so Close() calls from within request handlers do not discard body
+	onHitEOF          func() // if non-nil, func to call when EOF is Read
 }
 
 // ErrBodyReadAfterClose is returned when reading a [Request] or [Response]
@@ -968,51 +968,69 @@ func (b *body) unreadDataSizeLocked() int64 {
 	return -1
 }
 
+// shouldDiscardBodyLocked determines whether a body needs to discard its
+// content or not.
+// b.mu must be held.
+func (b *body) shouldDiscardBodyLocked() bool {
+	// Already saw EOF. No point in discarding the body.
+	if b.sawEOF {
+		return false
+	}
+	// No trailer and closing the connection next. No point in discarding the
+	// body since it will not be reusable.
+	if b.hdr == nil && b.closing {
+		return false
+	}
+	return true
+}
+
+// tryDiscardBody attempts to discard the body content. If the body cannot be
+// discarded completely, b.incompleteDiscard will be set to true.
+func (b *body) tryDiscardBody() {
+	b.mu.Lock()
+	defer b.mu.Unlock()
+	if !b.shouldDiscardBodyLocked() {
+		return
+	}
+	// Read up to maxPostHandlerReadBytes bytes of the body, looking for EOF
+	// (and trailers), so we can re-use this connection.
+	if lr, ok := b.src.(*io.LimitedReader); ok && lr.N > maxPostHandlerReadBytes {
+		// There was a declared Content-Length, and we have more bytes
+		// remaining than our maxPostHandlerReadBytes tolerance. So, give up.
+		b.incompleteDiscard = true
+		return
+	}
+	// Consume the body, which will also lead to us reading the trailer headers
+	// after the body, if present.
+	n, err := io.CopyN(io.Discard, bodyLocked{b}, maxPostHandlerReadBytes)
+	if err == io.EOF {
+		err = nil
+	}
+	if n == maxPostHandlerReadBytes || err != nil {
+		b.incompleteDiscard = true
+	}
+}
+
 func (b *body) Close() error {
 	b.mu.Lock()
 	defer b.mu.Unlock()
 	if b.closed {
 		return nil
 	}
-	var err error
-	switch {
-	case b.sawEOF:
-		// Already saw EOF, so no need going to look for it.
-	case b.hdr == nil && b.closing:
-		// no trailer and closing the connection next.
-		// no point in reading to EOF.
-	case b.doEarlyClose:
-		// Read up to maxPostHandlerReadBytes bytes of the body, looking
-		// for EOF (and trailers), so we can re-use this connection.
-		if lr, ok := b.src.(*io.LimitedReader); ok && lr.N > maxPostHandlerReadBytes {
-			// There was a declared Content-Length, and we have more bytes remaining
-			// than our maxPostHandlerReadBytes tolerance. So, give up.
-			b.earlyClose = true
-		} else {
-			var n int64
-			// Consume the body, or, which will also lead to us reading
-			// the trailer headers after the body, if present.
-			n, err = io.CopyN(io.Discard, bodyLocked{b}, maxPostHandlerReadBytes)
-			if err == io.EOF {
-				err = nil
-			}
-			if n == maxPostHandlerReadBytes {
-				b.earlyClose = true
-			}
-		}
-	default:
-		// Fully consume the body, which will also lead to us reading
-		// the trailer headers after the body, if present.
-		_, err = io.Copy(io.Discard, bodyLocked{b})
-	}
 	b.closed = true
+	if !b.shouldDiscardBodyLocked() || b.inRequestHandler {
+		return nil
+	}
+	// Fully consume the body, which will also lead to us reading
+	// the trailer headers after the body, if present.
+	_, err := io.Copy(io.Discard, bodyLocked{b})
 	return err
 }
 
-func (b *body) didEarlyClose() bool {
+func (b *body) didIncompleteDiscard() bool {
 	b.mu.Lock()
 	defer b.mu.Unlock()
-	return b.earlyClose
+	return b.incompleteDiscard
 }
 
 // bodyRemains reports whether future Read calls might
@@ -1036,9 +1054,6 @@ type bodyLocked struct {
 }
 
 func (bl bodyLocked) Read(p []byte) (n int, err error) {
-	if bl.b.closed {
-		return 0, ErrBodyReadAfterClose
-	}
 	return bl.b.readLocked(p)
 }