fix closing of http.Response and http.Request bodies

http3/body.go

@@ -3,11 +3,13 @@ package http3
 import (
 	"errors"
 	"io"
+
+	"github.com/lucas-clemente/quic-go"
 )
 
 // The body of a http.Request or http.Response.
 type body struct {
-	str io.ReadCloser
+	str quic.Stream
 
 	isRequest bool
 
@@ -16,14 +18,14 @@ type body struct {
 
 var _ io.ReadCloser = &body{}
 
-func newRequestBody(str io.ReadCloser) *body {
+func newRequestBody(str quic.Stream) *body {
 	return &body{
 		str:       str,
 		isRequest: true,
 	}
 }
 
-func newResponseBody(str io.ReadCloser) *body {
+func newResponseBody(str quic.Stream) *body {
 	return &body{str: str}
 }
 
@@ -62,7 +64,8 @@ func (r *body) Read(b []byte) (int, error) {
 func (r *body) Close() error {
 	// quic.Stream.Close() closes the write side, not the read side
 	if r.isRequest {
-		return nil
-	}
 		return r.str.Close()
 	}
+	r.str.CancelRead(quic.ErrorCode(errorRequestCanceled))
+	return nil
+}

http3/body_test.go

@@ -2,20 +2,17 @@ package http3
 
 import (
 	"bytes"
+	"fmt"
 	"io"
 
+	"github.com/golang/mock/gomock"
+	"github.com/lucas-clemente/quic-go"
+	mockquic "github.com/lucas-clemente/quic-go/internal/mocks/quic"
+
 	. "github.com/onsi/ginkgo"
 	. "github.com/onsi/gomega"
 )
 
-type closingBuffer struct {
-	*bytes.Buffer
-
-	closed bool
-}
-
-func (b *closingBuffer) Close() error { b.closed = true; return nil }
-
 type bodyType uint8
 
 const (
@@ -23,9 +20,19 @@ const (
 	bodyTypeResponse
 )
 
+func (t bodyType) String() string {
+	if t == bodyTypeRequest {
+		return "request"
+	}
+	return "response"
+}
+
 var _ = Describe("Body", func() {
-	var rb *body
-	var buf *bytes.Buffer
+	var (
+		rb  *body
+		str *mockquic.MockStream
+		buf *bytes.Buffer
+	)
 
 	getDataFrame := func(data []byte) []byte {
 		b := &bytes.Buffer{}
@@ -41,13 +48,21 @@ var _ = Describe("Body", func() {
 	for _, bt := range []bodyType{bodyTypeRequest, bodyTypeResponse} {
 		bodyType := bt
 
+		Context(fmt.Sprintf("using a %s body", bodyType), func() {
 			BeforeEach(func() {
-			cb := &closingBuffer{Buffer: buf}
+				str = mockquic.NewMockStream(mockCtrl)
+				str.EXPECT().Write(gomock.Any()).DoAndReturn(func(b []byte) (int, error) {
+					return buf.Write(b)
+				}).AnyTimes()
+				str.EXPECT().Read(gomock.Any()).DoAndReturn(func(b []byte) (int, error) {
+					return buf.Read(b)
+				}).AnyTimes()
+
 				switch bodyType {
 				case bodyTypeRequest:
-				rb = newRequestBody(cb)
+					rb = newRequestBody(str)
 				case bodyTypeResponse:
-				rb = newResponseBody(cb)
+					rb = newResponseBody(str)
 				}
 			})
 
@@ -132,19 +147,20 @@ var _ = Describe("Body", func() {
 				_, err := rb.Read([]byte{0})
 				Expect(err).To(MatchError("unexpected frame"))
 			})
+
+			if bodyType == bodyTypeRequest {
+				It("closes requests", func() {
+					str.EXPECT().Close()
+					Expect(rb.Close()).To(Succeed())
+				})
 			}
 
-	It("closes requests", func() {
-		cb := &closingBuffer{Buffer: buf}
-		rb := newRequestBody(cb)
-		Expect(rb.Close()).To(Succeed())
-		Expect(cb.closed).To(BeFalse())
-	})
-
+			if bodyType == bodyTypeResponse {
 				It("closes responses", func() {
-		cb := &closingBuffer{Buffer: buf}
-		rb := newResponseBody(cb)
+					str.EXPECT().CancelRead(quic.ErrorCode(errorRequestCanceled))
 					Expect(rb.Close()).To(Succeed())
-		Expect(cb.closed).To(BeTrue())
 				})
+			}
+		})
+	}
 })

http3/client.go

@@ -187,7 +187,7 @@ func (c *client) RoundTrip(req *http.Request) (*http.Response, error) {
 			res.Header.Add(hf.Name, hf.Value)
 		}
 	}
-	respBody := newResponseBody(&responseBody{str})
+	respBody := newResponseBody(str)
 	if requestGzip && res.Header.Get("Content-Encoding") == "gzip" {
 		res.Header.Del("Content-Encoding")
 		res.Header.Del("Content-Length")

http3/response_body.go

@@ -1,18 +0,0 @@
-package http3
-
-import (
-	"io"
-
-	quic "github.com/lucas-clemente/quic-go"
-)
-
-type responseBody struct {
-	quic.Stream
-}
-
-var _ io.ReadCloser = &responseBody{}
-
-func (rb *responseBody) Close() error {
-	rb.Stream.CancelRead(0)
-	return nil
-}

http3/response_body_test.go

@@ -1,26 +0,0 @@
-package http3
-
-import (
-	"github.com/golang/mock/gomock"
-	mockquic "github.com/lucas-clemente/quic-go/internal/mocks/quic"
-
-	. "github.com/onsi/ginkgo"
-	. "github.com/onsi/gomega"
-)
-
-var _ = Describe("Response Body", func() {
-	var (
-		stream *mockquic.MockStream
-		body   *responseBody
-	)
-
-	BeforeEach(func() {
-		stream = mockquic.NewMockStream(mockCtrl)
-		body = &responseBody{stream}
-	})
-
-	It("calls CancelRead when closing", func() {
-		stream.EXPECT().CancelRead(gomock.Any())
-		Expect(body.Close()).To(Succeed())
-	})
-})