middleware: fix httpFancyWriter.ReadFrom double-counting bytes with Tee

middleware/wrap_writer.go

@@ -208,8 +208,10 @@ 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).
 		n, err := io.Copy(&f.basicWriter, r)
-		f.basicWriter.bytes += int(n)
 		return n, err
 	}
 	rf := f.basicWriter.ResponseWriter.(io.ReaderFrom)

middleware/wrap_writer_test.go

@@ -4,6 +4,7 @@ import (
 	"bytes"
 	"net/http"
 	"net/http/httptest"
+	"strings"
 	"testing"
 )
 
@@ -84,3 +85,27 @@ func TestBasicWriterDiscardsWritesToOriginalResponseWriter(t *testing.T) {
 		assertEqual(t, 11, wrap.BytesWritten())
 	})
 }
+
+// TestHttpFancyWriterReadFromByteCountWithTee is a regression test for
+// https://github.com/go-chi/chi/issues/1067.
+// httpFancyWriter.ReadFrom was adding n to basicWriter.bytes even when the
+// write went through basicWriter.Write (which already increments the counter),
+// resulting in double-counting the bytes when a Tee writer was set.
+func TestHttpFancyWriterReadFromByteCountWithTee(t *testing.T) {
+	original := &httptest.ResponseRecorder{
+		HeaderMap: make(http.Header),
+		Body:      new(bytes.Buffer),
+	}
+	f := &httpFancyWriter{basicWriter: basicWriter{ResponseWriter: original}}
+
+	var teeBuf bytes.Buffer
+	f.Tee(&teeBuf)
+
+	const input = "hello world"
+	n, err := f.ReadFrom(strings.NewReader(input))
+	assertNoError(t, err)
+	assertEqual(t, int64(len(input)), n)
+	// Before the fix, BytesWritten() returned 22 (double-counted).
+	assertEqual(t, len(input), f.BytesWritten())
+	assertEqual(t, []byte(input), teeBuf.Bytes())
+}