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") + } +}