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

Fixes #75471

Change-Id: Ie8fc5fae4b2a9285501198d8379bbffe51ee63f7
Reviewed-on: https://go-review.googlesource.com/c/go/+/709335
Reviewed-by: Damien Neil <dneil@google.com>
Reviewed-by: Michael Pratt <mpratt@google.com>
LUCI-TryBot-Result: Go LUCI <golang-scoped@luci-project-accounts.iam.gserviceaccount.com>

diff --git a/src/net/http/httptest/recorder.go b/src/net/http/httptest/recorder.go
index 17aa70f06760a2e7adf34e879bdbfa465de9fc8b..7890b5ef6bdf318cac5a462814cf2aaf323fa2ca 100644 (file)
--- a/src/net/http/httptest/recorder.go
+++ b/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,6 +119,10 @@ 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)
@@ -122,6 +130,20 @@ func (rw *ResponseRecorder) WriteString(str string) (int, error) {
        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.

diff --git a/src/net/http/httptest/recorder_test.go b/src/net/http/httptest/recorder_test.go
index 4782eced43e6ce61c2dcab57b827048224be7f93..abf8e118d6092e12204d4d58b5a260ec72fd2e95 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 (
+       "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 {