From d7622a2d11d7428b70fe6e029f0f53d548d9f2a7 Mon Sep 17 00:00:00 2001
From: Marten Seemann <martenseemann@gmail.com>
Date: Sun, 22 May 2022 19:35:53 +0200
Subject: [PATCH] make it possible to parse a varint at the end of a reader

An io.Reader can read into the buffer and return the io.EOF in the same
call. In fact, that's how the quic.Stream is implemented.
In that case, we shouldn't return an error from quicvarint.Read,
otherwise the caller won't be able to parse a varint sent just before a
stream was closed.
---
 quicvarint/io.go      |  5 ++++-
 quicvarint/io_test.go | 43 +++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 47 insertions(+), 1 deletion(-)

diff --git a/quicvarint/io.go b/quicvarint/io.go
index c4d976b575c..9368d1c58b6 100644
--- a/quicvarint/io.go
+++ b/quicvarint/io.go
@@ -31,7 +31,10 @@ func NewReader(r io.Reader) Reader {
 
 func (r *byteReader) ReadByte() (byte, error) {
 	var b [1]byte
-	_, err := r.Reader.Read(b[:])
+	n, err := r.Reader.Read(b[:])
+	if n == 1 && err == io.EOF {
+		err = nil
+	}
 	return b[0], err
 }
 
diff --git a/quicvarint/io_test.go b/quicvarint/io_test.go
index cc28cd90fb0..054ab8643aa 100644
--- a/quicvarint/io_test.go
+++ b/quicvarint/io_test.go
@@ -22,6 +22,21 @@ func (r *nopWriter) Write(_ []byte) (int, error) {
 	return 0, io.ErrShortBuffer
 }
 
+// eofReader is a reader that returns data and the io.EOF at the same time in the last Read call
+type eofReader struct {
+	Data []byte
+	pos  int
+}
+
+func (r *eofReader) Read(b []byte) (int, error) {
+	n := copy(b, r.Data[r.pos:])
+	r.pos += n
+	if r.pos >= len(r.Data) {
+		return n, io.EOF
+	}
+	return n, nil
+}
+
 var _ io.Writer = &nopWriter{}
 
 var _ = Describe("Varint I/O", func() {
@@ -46,6 +61,34 @@ var _ = Describe("Varint I/O", func() {
 			Expect(err).To(Equal(io.ErrUnexpectedEOF))
 			Expect(val).To(Equal(byte(0)))
 		})
+
+		Context("EOF handling", func() {
+			It("eofReader works correctly", func() {
+				r := &eofReader{Data: []byte("foobar")}
+				b := make([]byte, 3)
+				n, err := r.Read(b)
+				Expect(n).To(Equal(3))
+				Expect(err).ToNot(HaveOccurred())
+				Expect(string(b)).To(Equal("foo"))
+				n, err = r.Read(b)
+				Expect(n).To(Equal(3))
+				Expect(err).To(MatchError(io.EOF))
+				Expect(string(b)).To(Equal("bar"))
+				n, err = r.Read(b)
+				Expect(err).To(MatchError(io.EOF))
+				Expect(n).To(BeZero())
+			})
+
+			It("correctly handles io.EOF", func() {
+				buf := &bytes.Buffer{}
+				Write(buf, 1337)
+
+				r := NewReader(&eofReader{Data: buf.Bytes()})
+				n, err := Read(r)
+				Expect(err).ToNot(HaveOccurred())
+				Expect(n).To(BeEquivalentTo(1337))
+			})
+		})
 	})
 
 	Context("Writer", func() {