From 9d5de12933f2f3afb3f6befe12fe4f9c152650b7 Mon Sep 17 00:00:00 2001
From: Marten Seemann <martenseemann@gmail.com>
Date: Thu, 26 May 2022 19:52:50 +0200
Subject: [PATCH] make it possible to parse a varint at the end of a reader
 (#3428)

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 c4d976b5..9368d1c5 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 cc28cd90..054ab864 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() {