From 59a8581684f0fd5a2d68da3d4be9f55699bd9281 Mon Sep 17 00:00:00 2001
From: Michael Stapelberg <stapelberg@golang.org>
Date: Fri, 23 Jun 2023 14:49:52 +0200
Subject: [PATCH] encoding/protodelim: fix handling of io.EOF
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Before this change, when encountering an io.EOF after reading at least
one byte, the zero value byte was — incorrectly — appended to sizeBuf,
and the io.EOF was ignored, resulting in a complete varint (0 has no
continuation bit), which in turn resulted in incorrect unmarshalling.

Change-Id: If06d45039d998eaddf91d0864814bb31d4cb7ae0
Reviewed-on: https://go-review.googlesource.com/c/protobuf/+/505555
Reviewed-by: Lasse Folger <lassefolger@google.com>
Reviewed-by: Damien Neil <dneil@google.com>
---
 encoding/protodelim/protodelim.go      |  6 +++++-
 encoding/protodelim/protodelim_test.go | 17 +++++++++++++++++
 2 files changed, 22 insertions(+), 1 deletion(-)

diff --git a/encoding/protodelim/protodelim.go b/encoding/protodelim/protodelim.go
index 75a50cb60..47be94e4d 100644
--- a/encoding/protodelim/protodelim.go
+++ b/encoding/protodelim/protodelim.go
@@ -96,7 +96,11 @@ func (o UnmarshalOptions) UnmarshalFrom(r Reader, m proto.Message) error {
 	sizeBuf := sizeArr[:0]
 	for i := range sizeArr {
 		b, err := r.ReadByte()
-		if err != nil && (err != io.EOF || i == 0) {
+		if err != nil {
+			// Immediate EOF is unexpected.
+			if err == io.EOF && i != 0 {
+				break
+			}
 			return err
 		}
 		sizeBuf = append(sizeBuf, b)
diff --git a/encoding/protodelim/protodelim_test.go b/encoding/protodelim/protodelim_test.go
index b9f8386ba..08bdea153 100644
--- a/encoding/protodelim/protodelim_test.go
+++ b/encoding/protodelim/protodelim_test.go
@@ -7,6 +7,7 @@ package protodelim_test
 import (
 	"bufio"
 	"bytes"
+	"encoding/binary"
 	"errors"
 	"io"
 	"testing"
@@ -202,3 +203,19 @@ func TestUnmarshalFrom_UnexpectedEOF(t *testing.T) {
 		t.Errorf("protodelim.UnmarshalFrom(size-only buf, _) = %v, want %v", got, want)
 	}
 }
+
+func TestUnmarshalFrom_PrematureHeader(t *testing.T) {
+	var data = []byte{128} // continuation bit set
+	err := protodelim.UnmarshalFrom(bytes.NewReader(data[:]), nil)
+	if got, want := err, io.ErrUnexpectedEOF; !errors.Is(got, want) {
+		t.Errorf("protodelim.UnmarshalFrom(%#v, nil) = %#v; want = %#v", data, got, want)
+	}
+}
+
+func TestUnmarshalFrom_InvalidVarint(t *testing.T) {
+	var data = bytes.Repeat([]byte{128}, 2*binary.MaxVarintLen64) // continuation bit set
+	err := protodelim.UnmarshalFrom(bytes.NewReader(data[:]), nil)
+	if err == nil {
+		t.Errorf("protodelim.UnmarshalFrom unexpectedly did not error on invalid varint")
+	}
+}