Repository: fxamacker/cbor (go)

Issue:
feature: Add feature to query next element size from stream decoder
**Is your feature request related to a problem? Please describe.**

In Cadence we would like to know the length of the next element that is to be decoded in a stream decoder. 
For example, when `StreamDecoder.NextType` returns `TextStringType`, we call `DecodeString` to decode the string. Now we would like to know the length of the string before the string is decoded.

Currently, there is no such information available.

**Describe the solution you'd like**

Allow a user of the stream decoder to query the length of the next element that is to be decoded in a stream decoder, before the element itself is actually decoded.

This could be exposed similarly to the function returning the next element's type, `NextType`, as e.g. `NextSize`.
It could also be useful to add a function that returns both values in one function call, e.g. `NextTypeAndSize`.

**Describe alternatives you've considered**

*None*

**Additional context**

*None*


Pull request:
Add NextSize() to stream decoder
#### Description

This PR closes #321 and unblocks https://github.com/onflow/cadence/pull/1374

NextSize() returns next CBOR data size for four CBOR types:
- byte string: length (in bytes) of byte string
- text string: length (in bytes) of text string
- array: number of array elements
- map: number of key/value pairs

Error is returned for indef length data and other CBOR types.

Patch:
diff --git a/stream_decode.go b/stream_decode.go
index 0caeb43..38488ad 100644
--- a/stream_decode.go
+++ b/stream_decode.go
@@ -1,6 +1,7 @@
 package cbor
 
 import (
+	"encoding/binary"
 	"errors"
 	"fmt"
 	"io"
@@ -68,7 +69,27 @@ type WrongTypeError struct {
 }
 
 func (e *WrongTypeError) Error() string {
-	return fmt.Sprintf("cannot decode %s to %s", e.ActualType.String(), e.ExpectedType)
+	return fmt.Sprintf("cbor: cannot decode %s to %s", e.ActualType.String(), e.ExpectedType)
+}
+
+// IndefiniteLengthNotSupportedError is returned by NextSize() if next CBOR data
+// is indefinite length.
+type IndefiniteLengthNotSupportedError struct {
+	t Type
+}
+
+func (e *IndefiniteLengthNotSupportedError) Error() string {
+	return fmt.Sprintf("cbor: size is unavailable for indefinite length %s", e.t)
+}
+
+// TypeNotSupportedError is returned by NextSize() if next CBOR type
+// is not supported/implemented by NextSize().
+type TypeNotSupportedError struct {
+	t Type
+}
+
+func (e *TypeNotSupportedError) Error() string {
+	return fmt.Sprintf("cbor: size operation is not supported for %s", e.t)
 }
 
 // StreamDecoder validates complete CBOR data and decodes it in chunks.
@@ -137,6 +158,50 @@ func (sd *StreamDecoder) NextType() (Type, error) {
 	return UndefinedType, errors.New("cbor: unrecognized type")
 }
 
+// NextSize returns the next CBOR data size for four CBOR types.
+// Returned uint64 represents different kind of size depending on
+// the CBOR type:
+// - byte string: length (in bytes) of byte string
+// - text string: length (in bytes) of text string
+// - array: number of array elements
+// - map: number of key/value pairs
+// Error is returned for indef length data and unsupported types.
+// Support for additional types wasn't added for simplicity.
+func (sd *StreamDecoder) NextSize() (uint64, error) {
+	t, err := sd.NextType()
+	if err != nil {
+		return 0, err
+	}
+
+	ai := sd.dec.d.data[sd.dec.d.off] & 0x1f
+
+	switch t {
+
+	case ByteStringType, TextStringType, ArrayType, MapType:
+		off := sd.dec.d.off + 1
+
+		// This function only interprets valid head because
+		// CBOR data is validated in NextType().
+		var size uint64
+		if ai < 24 {
+			size = uint64(ai)
+		} else if ai == 24 {
+			size = uint64(sd.dec.d.data[off])
+		} else if ai == 25 {
+			size = uint64(binary.BigEndian.Uint16(sd.dec.d.data[off : off+2]))
+		} else if ai == 26 {
+			size = uint64(binary.BigEndian.Uint32(sd.dec.d.data[off : off+4]))
+		} else if ai == 27 {
+			size = binary.BigEndian.Uint64(sd.dec.d.data[off : off+8])
+		} else if ai == 31 {
+			return 0, &IndefiniteLengthNotSupportedError{t}
+		}
+		return size, nil
+	}
+
+	return 0, &TypeNotSupportedError{t}
+}
+
 // Skip skips next CBOR data.
 func (sd *StreamDecoder) Skip() error {
 	if err := sd.prepareNext(); err != nil {


Test patch:
diff --git a/stream_decode_test.go b/stream_decode_test.go
index a4a7a2c..49b8fae 100644
--- a/stream_decode_test.go
+++ b/stream_decode_test.go
@@ -53,7 +53,7 @@ func TestStreamDecodeBool(t *testing.T) {
 					t.Errorf("NextType() returned %s, want %s", nt, expectedType)
 				}
 
-				wantErrorMsg := "cannot decode CBOR boolean type to string"
+				wantErrorMsg := "cbor: cannot decode CBOR boolean type to string"
 
 				// DecodeXXX() should return WrongTypeError with type mismatch (data offset is not moved)
 				_, err = sd.sd.DecodeString()
@@ -119,7 +119,7 @@ func TestStreamDecodeNil(t *testing.T) {
 				t.Errorf("NextType() returned %s, want %s", nt, expectedType)
 			}
 
-			wantErrorMsg := "cannot decode CBOR nil type to bool"
+			wantErrorMsg := "cbor: cannot decode CBOR nil type to bool"
 
 			// DecodeXXX() should return WrongTypeError with type mismatch (data offset is not moved)
 			_, err = sd.sd.DecodeBool()
@@ -196,7 +196,7 @@ func TestStreamDecodeUint(t *testing.T) {
 					t.Errorf("NextType() returned %s, want %s", nt, expectedType)
 				}
 
-				wantErrorMsg := "cannot decode CBOR uint type to bytes"
+				wantErrorMsg := "cbor: cannot decode CBOR uint type to bytes"
 
 				// DecodeXXX() should return WrongTypeError with type mismatch (data offset is not moved)
 				_, err = sd.sd.DecodeBytes()
@@ -287,9 +287,9 @@ func TestStreamDecodeInt(t *testing.T) {
 
 				var wantErrorMsg string
 				if tc.expected >= 0 {
-					wantErrorMsg = "cannot decode CBOR uint type to big.Int"
+					wantErrorMsg = "cbor: cannot decode CBOR uint type to big.Int"
 				} else {
-					wantErrorMsg = "cannot decode CBOR int type to big.Int"
+					wantErrorMsg = "cbor: cannot decode CBOR int type to big.Int"
 				}
 
 				// DecodeXXX() should return WrongTypeError with type mismatch (data offset is not moved)
@@ -408,7 +408,7 @@ func TestStreamDecodeBytes(t *testing.T) {
 					t.Errorf("NextType() returned %s, want %s", nt, expectedType)
 				}
 
-				wantErrorMsg := "cannot decode CBOR byte string type to int64"
+				wantErrorMsg := "cbor: cannot decode CBOR byte string type to int64"
 
 				// DecodeXXX() should return WrongTypeError with type mismatch (data offset is not moved)
 				i, err := sd.sd.DecodeInt64()
@@ -534,7 +534,7 @@ func TestStreamDecodeString(t *testing.T) {
 					t.Errorf("NextType() returned %s, want %s", nt, expectedType)
 				}
 
-				wantErrorMsg := "cannot decode CBOR text string type to bytes"
+				wantErrorMsg := "cbor: cannot decode CBOR text string type to bytes"
 
 				// DecodeXXX() should return WrongTypeError with type mismatch (data offset is not moved)
 				i, err := sd.sd.DecodeBytes()
@@ -718,7 +718,7 @@ func TestStreamDecodeBigInt(t *testing.T) {
 					t.Errorf("NextType() returned %s, want %s", nt, expectedType)
 				}
 
-				wantErrorMsg := "cannot decode CBOR bignum type to string"
+				wantErrorMsg := "cbor: cannot decode CBOR bignum type to string"
 
 				// DecodeXXX() should return WrongTypeError with type mismatch (data offset is not moved)
 				i, err := sd.sd.DecodeString()
@@ -795,7 +795,7 @@ func TestStreamDecodeTag(t *testing.T) {
 				t.Errorf("NextType() returned %s, want %s", nt, TagType)
 			}
 
-			wantErrorMsg := "cannot decode CBOR tag type to array"
+			wantErrorMsg := "cbor: cannot decode CBOR tag type to array"
 
 			// DecodeXXX() should return WrongTypeError with type mismatch (data offset is not moved)
 			i, err := sd.sd.DecodeArrayHead()
@@ -825,7 +825,7 @@ func TestStreamDecodeTag(t *testing.T) {
 				t.Errorf("NextType() returned %s, want %s", nt, TextStringType)
 			}
 
-			wantErrorMsg = "cannot decode CBOR text string type to array"
+			wantErrorMsg = "cbor: cannot decode CBOR text string type to array"
 
 			// DecodeXXX() should return WrongTypeError with type mismatch (data offset is not moved)
 			i, err = sd.sd.DecodeArrayHead()
@@ -890,7 +890,7 @@ func TestStreamDecodeEmptyArray(t *testing.T) {
 				t.Errorf("NextType() returned %s, want %s", nt, ArrayType)
 			}
 
-			wantErrorMsg := "cannot decode CBOR array type to tag"
+			wantErrorMsg := "cbor: cannot decode CBOR array type to tag"
 
 			// DecodeXXX() should return WrongTypeError with type mismatch (data offset is not moved)
 			i, err := sd.sd.DecodeTagNumber()
@@ -955,7 +955,7 @@ func TestStreamDecodeArray(t *testing.T) {
 				t.Errorf("NextType() returned %s, want %s", nt, ArrayType)
 			}
 
-			wantErrorMsg := "cannot decode CBOR array type to tag"
+			wantErrorMsg := "cbor: cannot decode CBOR array type to tag"
 
 			// DecodeXXX() should return WrongTypeError with type mismatch (data offset is not moved)
 			i, err := sd.sd.DecodeTagNumber()
@@ -1091,7 +1091,7 @@ func TestStreamDecodeRawBytes(t *testing.T) {
 				t.Errorf("NextType() returned %s, want %s", nt, ArrayType)
 			}
 
-			wantErrorMsg := "cannot decode CBOR array type to nil"
+			wantErrorMsg := "cbor: cannot decode CBOR array type to nil"
 
 			// DecodeXXX() should return WrongTypeError with type mismatch (data offset is not moved)
 			err = sd.sd.DecodeNil()
@@ -1148,7 +1148,7 @@ func TestStreamDecodeRawBytesZeroCopy(t *testing.T) {
 			t.Errorf("NextType() returned %s, want %s", nt, expectedType)
 		}
 
-		wantErrorMsg := "cannot decode CBOR array type to nil"
+		wantErrorMsg := "cbor: cannot decode CBOR array type to nil"
 
 		// DecodeXXX() should return WrongTypeError with type mismatch (data offset is not moved)
 		err = sd.DecodeNil()
@@ -1194,7 +1194,7 @@ func TestStreamDecodeRawBytesZeroCopy(t *testing.T) {
 			t.Errorf("NextType() returned %s, want %s", nt, expectedType)
 		}
 
-		wantErrorMsg := "cannot decode CBOR array type to uint64"
+		wantErrorMsg := "cbor: cannot decode CBOR array type to uint64"
 
 		// DecodeXXX() should return WrongTypeError with type mismatch (data offset is not moved)
 		_, err = sd.DecodeUint64()
@@ -1701,3 +1701,70 @@ func TestStreamDecoderNumBytesDecoded(t *testing.T) {
 		}
 	})
 }
+
+func TestStreamDecodeNextSize(t *testing.T) {
+	testCases := []struct {
+		name           string
+		data           []byte
+		expectedSize   uint64
+		expectedErrMsg string
+	}{
+		{"no data", []byte{}, 0, "EOF"},
+		{"invalid data", []byte{0x1b, 0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07}, 0, "unexpected EOF"},
+		{"invalid data", []byte{0x9d}, 0, "cbor: invalid additional information"},
+		{"false", []byte{0xf4}, 0, "size operation is not supported"},
+		{"true", []byte{0xf5}, 0, "size operation is not supported"},
+		{"nil", []byte{0xf6}, 0, "size operation is not supported"},
+		{"uint", []byte{0x01}, 0, "size operation is not supported"},                                                            // 1
+		{"int", []byte{0x20}, 0, "size operation is not supported"},                                                             // -1
+		{"tag", []byte{0xc2, 0x49, 0x01, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00}, 0, "size operation is not supported"}, // bignum(18446744073709551616)
+		{"empty byte string", []byte{0x40}, 0, ""},                                                                              // []
+		{"byte string", []byte{0x45, 0x01, 0x02, 0x03, 0x04, 0x05}, 5, ""},                                                      // [1, 2, 3, 4, 5]
+		{"indef length byte string", []byte{0x5f, 0x42, 0x01, 0x02, 0x043, 0x03, 0x04, 0x05, 0xff}, 0, "size is unavailable for indefinite length"},
+		{"empty text string", []byte{0x60}, 0, ""},                                                 // ""
+		{"text string", []byte{0x65, 0x68, 0x65, 0x6c, 0x6c, 0x6f}, 5, ""},                         // "hello"
+		{"text string", []byte{0x78, 0x01, 0x61}, 1, ""},                                           // "a"
+		{"text string", []byte{0x79, 0x00, 0x01, 0x61}, 1, ""},                                     // "a"
+		{"text string", []byte{0x7a, 0x00, 0x00, 0x00, 0x01, 0x61}, 1, ""},                         // "a"
+		{"text string", []byte{0x7b, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x01, 0x61}, 1, ""}, // "a"
+		{"indef length text string", []byte{0x7f, 0x65, 0x73, 0x74, 0x72, 0x65, 0x61, 0x64, 0x6d, 0x69, 0x6e, 0x67, 0xff}, 0, "size is unavailable for indefinite length"},
+		{"empty array", []byte{0x80}, 0, ""},                                                                         // []
+		{"array", []byte{0x83, 0x01, 0x02, 0x03}, 3, ""},                                                             // [1, 2, 3]
+		{"indef length array", []byte{0x9f, 0x01, 0x02, 0x03, 0xff}, 0, "size is unavailable for indefinite length"}, // [1, 2, 3]
+	}
+
+	for _, tc := range testCases {
+
+		// For each test case, test 2 StreamDecoders.
+
+		decoders := []struct {
+			name string
+			sd   *StreamDecoder
+		}{
+			{"byte_decoder", NewByteStreamDecoder(tc.data)},
+			{"reader_decoder", NewStreamDecoder(bytes.NewReader(tc.data))},
+		}
+
+		for _, sd := range decoders {
+
+			t.Run(sd.name+" "+tc.name, func(t *testing.T) {
+				// NextSize() peeks at next CBOR data size
+				size, err := sd.sd.NextSize()
+				if size != tc.expectedSize {
+					t.Errorf("NextSize() returned %d, want %d", size, tc.expectedSize)
+				}
+				if err == nil {
+					if tc.expectedErrMsg != "" {
+						t.Errorf("NextSize() didn't return error, want %s", tc.expectedErrMsg)
+					}
+				} else {
+					if tc.expectedErrMsg == "" {
+						t.Errorf("NextSize() returned error %v", err)
+					} else if !strings.Contains(err.Error(), tc.expectedErrMsg) {
+						t.Errorf("NextSize() returned error %v, want %s", err, tc.expectedErrMsg)
+					}
+				}
+			})
+		}
+	}
+}