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>

diff --git a/src/net/http/httptest/recorder.go b/src/net/http/httptest/recorder.go
index 7890b5ef6bdf318cac5a462814cf2aaf323fa2ca..4006f4406d27a52f837effe2efaa81075c6e8323 100644 (file)
--- 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 abf8e118d6092e12204d4d58b5a260ec72fd2e95..099697651465adb1e7131b3c7a10e3b326a69b6d 100644 (file)
--- 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"