Honor Discard() in httpFancyWriter.ReadFrom (#1110)

Author: DucMinhNe

middleware/wrap_writer.go

@@ -207,10 +207,12 @@ func (f *http2FancyWriter) Push(target string, opts *http.PushOptions) error {
 }
 
 func (f *httpFancyWriter) ReadFrom(r io.Reader) (int64, error) {
-	if f.basicWriter.tee != nil {
-		// Route through basicWriter.Write so that data is also written to the
-		// tee writer. basicWriter.Write already increments basicWriter.bytes,
-		// so we must NOT add n again here (that would double-count).
+	if f.basicWriter.tee != nil || f.basicWriter.discard {
+		// Route through basicWriter.Write so that the tee and discard semantics
+		// are honored (the fast ReaderFrom path below would bypass both, writing
+		// straight to the original ResponseWriter). basicWriter.Write already
+		// increments basicWriter.bytes, so we must NOT add n again here (that
+		// would double-count).
 		n, err := io.Copy(&f.basicWriter, r)
 		return n, err
 	}

middleware/wrap_writer_test.go

@@ -2,12 +2,26 @@ package middleware
 
 import (
 	"bytes"
+	"io"
 	"net/http"
 	"net/http/httptest"
 	"strings"
 	"testing"
 )
 
+// readerFromRecorder is a ResponseRecorder that also satisfies io.ReaderFrom,
+// mirroring the real net/http response writer (which implements ReadFrom and
+// is what httpFancyWriter.ReadFrom fast-paths to). Bytes copied via ReadFrom
+// are recorded to the body so a test can observe whether they reached the
+// original writer.
+type readerFromRecorder struct {
+	*httptest.ResponseRecorder
+}
+
+func (r *readerFromRecorder) ReadFrom(src io.Reader) (int64, error) {
+	return io.Copy(r.ResponseRecorder.Body, src)
+}
+
 func TestHttpFancyWriterRemembersWroteHeaderWhenFlushed(t *testing.T) {
 	f := &httpFancyWriter{basicWriter: basicWriter{ResponseWriter: httptest.NewRecorder()}}
 	f.Flush()
@@ -109,3 +123,25 @@ func TestHttpFancyWriterReadFromByteCountWithTee(t *testing.T) {
 	assertEqual(t, len(input), f.BytesWritten())
 	assertEqual(t, []byte(input), teeBuf.Bytes())
 }
+
+// TestHttpFancyWriterReadFromHonorsDiscard verifies that Discard() suppresses
+// the body even when it is written via ReadFrom/io.Copy. Before the fix, the
+// non-tee ReadFrom path streamed straight to the original ResponseWriter's
+// ReaderFrom, bypassing the discard flag entirely.
+func TestHttpFancyWriterReadFromHonorsDiscard(t *testing.T) {
+	original := &readerFromRecorder{&httptest.ResponseRecorder{
+		HeaderMap: make(http.Header),
+		Body:      new(bytes.Buffer),
+	}}
+	f := &httpFancyWriter{basicWriter: basicWriter{ResponseWriter: original}}
+	f.Discard()
+
+	const input = "hello world"
+	n, err := f.ReadFrom(strings.NewReader(input))
+	assertNoError(t, err)
+	assertEqual(t, int64(len(input)), n)
+	// The data must NOT reach the original writer once Discard() is set.
+	assertEqual(t, 0, original.Body.Len())
+	// BytesWritten still reflects the bytes consumed.
+	assertEqual(t, len(input), f.BytesWritten())
+}