only allow writing a response body for status codes that allow bodies

fixes #386

h2quic/response_writer.go

@@ -21,6 +21,7 @@ type responseWriter struct {
 	headerStreamMutex *sync.Mutex
 
 	header        http.Header
+	status        int // status code passed to WriteHeader
 	headerWritten bool
 }
 
@@ -43,6 +44,7 @@ func (w *responseWriter) WriteHeader(status int) {
 		return
 	}
 	w.headerWritten = true
+	w.status = status
 
 	var headers bytes.Buffer
 	enc := hpack.NewEncoder(&headers)
@@ -72,6 +74,9 @@ func (w *responseWriter) Write(p []byte) (int, error) {
 	if !w.headerWritten {
 		w.WriteHeader(200)
 	}
+	if !bodyAllowedForStatus(w.status) {
+		return 0, http.ErrBodyNotAllowed
+	}
 	return w.dataStream.Write(p)
 }
 
@@ -79,3 +84,18 @@ func (w *responseWriter) Flush() {}
 
 // test that we implement http.Flusher
 var _ http.Flusher = &responseWriter{}
+
+// copied from http2/http2.go
+// bodyAllowedForStatus reports whether a given response status code
+// permits a body. See RFC 2616, section 4.4.
+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
+}

h2quic/response_writer_test.go

@@ -92,4 +92,12 @@ var _ = Describe("Response Writer", func() {
 		w.WriteHeader(500)
 		Expect(headerStream.Bytes()).To(Equal([]byte{0x0, 0x0, 0x1, 0x1, 0x4, 0x0, 0x0, 0x0, 0x5, 0x88})) // 0x88 is 200
 	})
+
+	It("doesn't allow writes if the status code doesn't allow a body", func() {
+		w.WriteHeader(304)
+		n, err := w.Write([]byte("foobar"))
+		Expect(n).To(BeZero())
+		Expect(err).To(MatchError(http.ErrBodyNotAllowed))
+		Expect(dataStream.Bytes()).To(HaveLen(0))
+	})
 })