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()
 	}
-	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,110 +48,119 @@ var _ = Describe("Body", func() {
 	for _, bt := range []bodyType{bodyTypeRequest, bodyTypeResponse} {
 		bodyType := bt
 
-		BeforeEach(func() {
-			cb := &closingBuffer{Buffer: buf}
-			switch bodyType {
-			case bodyTypeRequest:
-				rb = newRequestBody(cb)
-			case bodyTypeResponse:
-				rb = newResponseBody(cb)
+		Context(fmt.Sprintf("using a %s body", bodyType), func() {
+			BeforeEach(func() {
+				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(str)
+				case bodyTypeResponse:
+					rb = newResponseBody(str)
+				}
+			})
+
+			It("reads DATA frames in a single run", func() {
+				buf.Write(getDataFrame([]byte("foobar")))
+				b := make([]byte, 6)
+				n, err := rb.Read(b)
+				Expect(err).ToNot(HaveOccurred())
+				Expect(n).To(Equal(6))
+				Expect(b).To(Equal([]byte("foobar")))
+			})
+
+			It("reads DATA frames in multiple runs", func() {
+				buf.Write(getDataFrame([]byte("foobar")))
+				b := make([]byte, 3)
+				n, err := rb.Read(b)
+				Expect(err).ToNot(HaveOccurred())
+				Expect(n).To(Equal(3))
+				Expect(b).To(Equal([]byte("foo")))
+				n, err = rb.Read(b)
+				Expect(err).ToNot(HaveOccurred())
+				Expect(n).To(Equal(3))
+				Expect(b).To(Equal([]byte("bar")))
+			})
+
+			It("reads DATA frames into too large buffers", func() {
+				buf.Write(getDataFrame([]byte("foobar")))
+				b := make([]byte, 10)
+				n, err := rb.Read(b)
+				Expect(err).ToNot(HaveOccurred())
+				Expect(n).To(Equal(6))
+				Expect(b[:n]).To(Equal([]byte("foobar")))
+			})
+
+			It("reads DATA frames into too large buffers, in multiple runs", func() {
+				buf.Write(getDataFrame([]byte("foobar")))
+				b := make([]byte, 4)
+				n, err := rb.Read(b)
+				Expect(err).ToNot(HaveOccurred())
+				Expect(n).To(Equal(4))
+				Expect(b).To(Equal([]byte("foob")))
+				n, err = rb.Read(b)
+				Expect(err).ToNot(HaveOccurred())
+				Expect(n).To(Equal(2))
+				Expect(b[:n]).To(Equal([]byte("ar")))
+			})
+
+			It("reads multiple DATA frames", func() {
+				buf.Write(getDataFrame([]byte("foo")))
+				buf.Write(getDataFrame([]byte("bar")))
+				b := make([]byte, 6)
+				n, err := rb.Read(b)
+				Expect(err).ToNot(HaveOccurred())
+				Expect(n).To(Equal(3))
+				Expect(b[:n]).To(Equal([]byte("foo")))
+				n, err = rb.Read(b)
+				Expect(err).ToNot(HaveOccurred())
+				Expect(n).To(Equal(3))
+				Expect(b[:n]).To(Equal([]byte("bar")))
+			})
+
+			It("skips HEADERS frames", func() {
+				buf.Write(getDataFrame([]byte("foo")))
+				(&headersFrame{Length: 10}).Write(buf)
+				buf.Write(make([]byte, 10))
+				buf.Write(getDataFrame([]byte("bar")))
+				b := make([]byte, 6)
+				n, err := io.ReadFull(rb, b)
+				Expect(err).ToNot(HaveOccurred())
+				Expect(n).To(Equal(6))
+				Expect(b).To(Equal([]byte("foobar")))
+			})
+
+			It("errors when it can't parse the frame", func() {
+				buf.Write([]byte("invalid"))
+				_, err := rb.Read([]byte{0})
+				Expect(err).To(HaveOccurred())
+			})
+
+			It("errors on unexpected frames", func() {
+				(&settingsFrame{}).Write(buf)
+				_, 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())
+				})
+			}
+
+			if bodyType == bodyTypeResponse {
+				It("closes responses", func() {
+					str.EXPECT().CancelRead(quic.ErrorCode(errorRequestCanceled))
+					Expect(rb.Close()).To(Succeed())
+				})
 			}
 		})
-
-		It("reads DATA frames in a single run", func() {
-			buf.Write(getDataFrame([]byte("foobar")))
-			b := make([]byte, 6)
-			n, err := rb.Read(b)
-			Expect(err).ToNot(HaveOccurred())
-			Expect(n).To(Equal(6))
-			Expect(b).To(Equal([]byte("foobar")))
-		})
-
-		It("reads DATA frames in multiple runs", func() {
-			buf.Write(getDataFrame([]byte("foobar")))
-			b := make([]byte, 3)
-			n, err := rb.Read(b)
-			Expect(err).ToNot(HaveOccurred())
-			Expect(n).To(Equal(3))
-			Expect(b).To(Equal([]byte("foo")))
-			n, err = rb.Read(b)
-			Expect(err).ToNot(HaveOccurred())
-			Expect(n).To(Equal(3))
-			Expect(b).To(Equal([]byte("bar")))
-		})
-
-		It("reads DATA frames into too large buffers", func() {
-			buf.Write(getDataFrame([]byte("foobar")))
-			b := make([]byte, 10)
-			n, err := rb.Read(b)
-			Expect(err).ToNot(HaveOccurred())
-			Expect(n).To(Equal(6))
-			Expect(b[:n]).To(Equal([]byte("foobar")))
-		})
-
-		It("reads DATA frames into too large buffers, in multiple runs", func() {
-			buf.Write(getDataFrame([]byte("foobar")))
-			b := make([]byte, 4)
-			n, err := rb.Read(b)
-			Expect(err).ToNot(HaveOccurred())
-			Expect(n).To(Equal(4))
-			Expect(b).To(Equal([]byte("foob")))
-			n, err = rb.Read(b)
-			Expect(err).ToNot(HaveOccurred())
-			Expect(n).To(Equal(2))
-			Expect(b[:n]).To(Equal([]byte("ar")))
-		})
-
-		It("reads multiple DATA frames", func() {
-			buf.Write(getDataFrame([]byte("foo")))
-			buf.Write(getDataFrame([]byte("bar")))
-			b := make([]byte, 6)
-			n, err := rb.Read(b)
-			Expect(err).ToNot(HaveOccurred())
-			Expect(n).To(Equal(3))
-			Expect(b[:n]).To(Equal([]byte("foo")))
-			n, err = rb.Read(b)
-			Expect(err).ToNot(HaveOccurred())
-			Expect(n).To(Equal(3))
-			Expect(b[:n]).To(Equal([]byte("bar")))
-		})
-
-		It("skips HEADERS frames", func() {
-			buf.Write(getDataFrame([]byte("foo")))
-			(&headersFrame{Length: 10}).Write(buf)
-			buf.Write(make([]byte, 10))
-			buf.Write(getDataFrame([]byte("bar")))
-			b := make([]byte, 6)
-			n, err := io.ReadFull(rb, b)
-			Expect(err).ToNot(HaveOccurred())
-			Expect(n).To(Equal(6))
-			Expect(b).To(Equal([]byte("foobar")))
-		})
-
-		It("errors when it can't parse the frame", func() {
-			buf.Write([]byte("invalid"))
-			_, err := rb.Read([]byte{0})
-			Expect(err).To(HaveOccurred())
-		})
-
-		It("errors on unexpected frames", func() {
-			(&settingsFrame{}).Write(buf)
-			_, err := rb.Read([]byte{0})
-			Expect(err).To(MatchError("unexpected frame"))
-		})
 	}
-
-	It("closes requests", func() {
-		cb := &closingBuffer{Buffer: buf}
-		rb := newRequestBody(cb)
-		Expect(rb.Close()).To(Succeed())
-		Expect(cb.closed).To(BeFalse())
-	})
-
-	It("closes responses", func() {
-		cb := &closingBuffer{Buffer: buf}
-		rb := newResponseBody(cb)
-		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())
-	})
-})