net/http/httptest: check whether response bodies are allowed

Fixes golang#75471

Change-Id: Ie8fc5fae4b2a9285501198d8379bbffe51ee63f7
Reviewed-on: https://go-review.googlesource.com/c/go/+/709335
Reviewed-by: Damien Neil <[email protected]>
Reviewed-by: Michael Pratt <[email protected]>
LUCI-TryBot-Result: Go LUCI <[email protected]>

Author: seankhliao

src/net/http/httptest/recorder.go

@@ -105,6 +105,10 @@ 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
+	}
 	rw.writeHeader(buf, "")
 	if rw.Body != nil {
 		rw.Body.Write(buf)
@@ -115,13 +119,31 @@ func (rw *ResponseRecorder) Write(buf []byte) (int, error) {
 // 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
+	}
 	rw.writeHeader(nil, str)
 	if rw.Body != nil {
 		rw.Body.WriteString(str)
 	}
 	return len(str), nil
 }
 
+// bodyAllowedForStatus reports whether a given response status code
+// permits a body. See RFC 7230, section 3.3.
+func bodyAllowedForStatus(status int) bool {
+	switch {
+	case status >= 100 && status <= 199:
+		return false
+	case status == 204:
+		return false
+	case status == 304:
+		return false
+	}
+	return true
+}
+
 func checkWriteHeaderCode(code int) {
 	// Issue 22880: require valid WriteHeader status codes.
 	// For now we only enforce that it's three digits.

src/net/http/httptest/recorder_test.go

@@ -5,6 +5,7 @@
 package httptest
 
 import (
+	"errors"
 	"fmt"
 	"io"
 	"net/http"
@@ -309,6 +310,21 @@ func TestRecorder(t *testing.T) {
 	}
 }
 
+func TestBodyNotAllowed(t *testing.T) {
+	rw := NewRecorder()
+	rw.WriteHeader(204)
+
+	_, err := rw.Write([]byte("hello world"))
+	if !errors.Is(err, http.ErrBodyNotAllowed) {
+		t.Errorf("expected BodyNotAllowed for Write after 204, got: %v", err)
+	}
+
+	_, err = rw.WriteString("hello world")
+	if !errors.Is(err, http.ErrBodyNotAllowed) {
+		t.Errorf("expected BodyNotAllowed for WriteString after 204, got: %v", err)
+	}
+}
+
 // issue 39017 - disallow Content-Length values such as "+3"
 func TestParseContentLength(t *testing.T) {
 	tests := []struct {