From b80e73d7de6484b9901b3f4d6e5360d3d33641f5 Mon Sep 17 00:00:00 2001
From: Atsushi Watanabe <atsushi.w@ieee.org>
Date: Fri, 20 Dec 2019 22:04:06 +0900
Subject: [PATCH] Fix short unknown data size read (#113)

---
 unmarshal.go  |  2 +-
 value.go      |  7 +++++++
 value_test.go | 38 +++++++++++++++++++++++++++++++++++++-
 3 files changed, 45 insertions(+), 2 deletions(-)

diff --git a/unmarshal.go b/unmarshal.go
index 460ce76..d619dec 100644
--- a/unmarshal.go
+++ b/unmarshal.go
@@ -103,7 +103,7 @@ func readElement(r0 io.Reader, n int64, vo reflect.Value, depth int, pos uint64,
 			return nil, ErrUnknownElement
 		}
 
-		size, nb, err := readVInt(r)
+		size, nb, err := readDataSize(r)
 		headerSize += uint64(nb)
 		if err != nil {
 			return nil, err
diff --git a/value.go b/value.go
index 9135fce..6a71cbe 100644
--- a/value.go
+++ b/value.go
@@ -50,6 +50,13 @@ var perTypeReader = map[DataType]func(io.Reader, uint64) (interface{}, error){
 	DataTypeBlock:  readBlock,
 }
 
+func readDataSize(r io.Reader) (uint64, int, error) {
+	v, n, err := readVInt(r)
+	if v == (uint64(0xFFFFFFFFFFFFFFFF) >> uint(64-n*7)) {
+		return SizeUnknown, n, err
+	}
+	return v, n, err
+}
 func readVInt(r io.Reader) (uint64, int, error) {
 	var bs [1]byte
 	bytesRead, err := io.ReadFull(r, bs[:])
diff --git a/value_test.go b/value_test.go
index 1cc3013..f8f23cd 100644
--- a/value_test.go
+++ b/value_test.go
@@ -32,7 +32,7 @@ func TestDataSize(t *testing.T) {
 	}
 
 	for n, c := range testCases {
-		t.Run("Decode "+n, func(t *testing.T) {
+		t.Run("DecodeVInt "+n, func(t *testing.T) {
 			r, _, err := readVInt(bytes.NewBuffer(c.b))
 			if err != nil {
 				t.Fatalf("Failed to readVInt: %v", err)
@@ -42,6 +42,17 @@ func TestDataSize(t *testing.T) {
 			}
 		})
 	}
+	for n, c := range testCases {
+		t.Run("DecodeDataSize "+n, func(t *testing.T) {
+			r, _, err := readDataSize(bytes.NewBuffer(c.b))
+			if err != nil {
+				t.Fatalf("Failed to readDataSize: %v", err)
+			}
+			if r != c.i {
+				t.Errorf("Unexpected readVInt result, expected: %d, got: %d", c.i, r)
+			}
+		})
+	}
 	for n, c := range testCases {
 		t.Run("Encode "+n, func(t *testing.T) {
 			b := encodeDataSize(c.i, 0)
@@ -52,6 +63,31 @@ func TestDataSize(t *testing.T) {
 	}
 }
 
+func TestDataSize_Unknown(t *testing.T) {
+	testCases := map[string][]byte{
+		"1 byte":  []byte{0xFF},
+		"2 bytes": []byte{0x7F, 0xFF},
+		"3 bytes": []byte{0x3F, 0xFF, 0xFF},
+		"4 bytes": []byte{0x1F, 0xFF, 0xFF, 0xFF},
+		"5 bytes": []byte{0x0F, 0xFF, 0xFF, 0xFF, 0xFF},
+		"6 bytes": []byte{0x07, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF},
+		"7 bytes": []byte{0x03, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF},
+		"8 bytes": []byte{0x01, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff},
+	}
+
+	for n, b := range testCases {
+		t.Run("DecodeDataSize "+n, func(t *testing.T) {
+			r, _, err := readDataSize(bytes.NewBuffer(b))
+			if err != nil {
+				t.Fatalf("Failed to readDataSize: %v", err)
+			}
+			if r != SizeUnknown {
+				t.Errorf("Unexpected readDataSize result, expected: %d, got: %d", SizeUnknown, r)
+			}
+		})
+	}
+}
+
 func TestElementID(t *testing.T) {
 	testCases := map[string]struct {
 		b []byte