net/http/httptest: record failed ResponseWriter writes

CL 709335 changed ResponseWriter.Write to return an error
when trying to write to a response with a status code which
doesn't permit a body, such as 304.

Continue to return an error, but still record the write in
ResponseWriter.Body. This maintains the documented property that
"the data in buf is written to rw.Body".

For #75471

Change-Id: I69139797559fe09d6580c5d25b4458f04263c60e
Reviewed-on: https://go-review.googlesource.com/c/go/+/711940
Reviewed-by: Sean Liao <sean@liao.dev>
TryBot-Bypass: Damien Neil <dneil@google.com>
Auto-Submit: Damien Neil <dneil@google.com>
Reviewed-by: Nicholas Husin <nsh@golang.org>
Reviewed-by: Nicholas Husin <husin@google.com>

2 files changed, 16 insertions, 10 deletions

diff --git a/src/net/http/httptest/recorder.go b/src/net/http/httptest/recorder.go
index 7890b5ef6b..4006f4406d 100644
--- a/src/net/http/httptest/recorder.go
+++ b/src/net/http/httptest/recorder.go
@@ -105,28 +105,28 @@ func (rw *ResponseRecorder) writeHeader(b []byte, str string) {
 // Write implements http.ResponseWriter. The data in buf is written to
 // rw.Body, if not nil.
 func (rw *ResponseRecorder) Write(buf []byte) (int, error) {
-	code := rw.Code
-	if !bodyAllowedForStatus(code) {
-		return 0, http.ErrBodyNotAllowed
-	}
+	// Record the write, even if we're going to return an error.
 	rw.writeHeader(buf, "")
 	if rw.Body != nil {
 		rw.Body.Write(buf)
 	}
+	if !bodyAllowedForStatus(rw.Code) {
+		return 0, http.ErrBodyNotAllowed
+	}
 	return len(buf), nil
 }
 
 // WriteString implements [io.StringWriter]. The data in str is written
 // to rw.Body, if not nil.
 func (rw *ResponseRecorder) WriteString(str string) (int, error) {
-	code := rw.Code
-	if !bodyAllowedForStatus(code) {
-		return 0, http.ErrBodyNotAllowed
-	}
+	// Record the write, even if we're going to return an error.
 	rw.writeHeader(nil, str)
 	if rw.Body != nil {
 		rw.Body.WriteString(str)
 	}
+	if !bodyAllowedForStatus(rw.Code) {
+		return 0, http.ErrBodyNotAllowed
+	}
 	return len(str), nil
 }
 
diff --git a/src/net/http/httptest/recorder_test.go b/src/net/http/httptest/recorder_test.go
index abf8e118d6..0996976514 100644
--- a/src/net/http/httptest/recorder_test.go
+++ b/src/net/http/httptest/recorder_test.go
@@ -5,6 +5,7 @@
 package httptest
 
 import (
+	"bytes"
 	"errors"
 	"fmt"
 	"io"
@@ -312,17 +313,22 @@ func TestRecorder(t *testing.T) {
 
 func TestBodyNotAllowed(t *testing.T) {
 	rw := NewRecorder()
+	rw.Body = new(bytes.Buffer)
 	rw.WriteHeader(204)
 
-	_, err := rw.Write([]byte("hello world"))
+	_, err := rw.Write([]byte("hello "))
 	if !errors.Is(err, http.ErrBodyNotAllowed) {
 		t.Errorf("expected BodyNotAllowed for Write after 204, got: %v", err)
 	}
 
-	_, err = rw.WriteString("hello world")
+	_, err = rw.WriteString("world")
 	if !errors.Is(err, http.ErrBodyNotAllowed) {
 		t.Errorf("expected BodyNotAllowed for WriteString after 204, got: %v", err)
 	}
+
+	if got, want := rw.Body.String(), "hello world"; got != want {
+		t.Errorf("got Body=%q, want %q", got, want)
+	}
 }
 
 // issue 39017 - disallow Content-Length values such as "+3"