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 golang#75471

Change-Id: I69139797559fe09d6580c5d25b4458f04263c60e
Reviewed-on: https://go-review.googlesource.com/c/go/+/711940
Reviewed-by: Sean Liao <[email protected]>
TryBot-Bypass: Damien Neil <[email protected]>
Auto-Submit: Damien Neil <[email protected]>
Reviewed-by: Nicholas Husin <[email protected]>
Reviewed-by: Nicholas Husin <[email protected]>

Author: neild

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
 }
 

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"