io: allow to restrict slice size in ReadVarBytes

pkg/io/binaryReader.go

@@ -168,8 +168,16 @@ func (r *BinReader) ReadVarUint() uint64 {
 
 // ReadVarBytes reads the next set of bytes from the underlying reader.
 // ReadVarUInt() is used to determine how large that slice is
-func (r *BinReader) ReadVarBytes() []byte {
+func (r *BinReader) ReadVarBytes(maxSize ...int) []byte {
 	n := r.ReadVarUint()
+	ms := maxArraySize
+	if len(maxSize) != 0 {
+		ms = maxSize[0]
+	}
+	if n > uint64(ms) {
+		r.Err = fmt.Errorf("byte-slice is too big (%d)", n)
+		return nil
+	}
 	b := make([]byte, n)
 	r.ReadBytes(b)
 	return b

pkg/io/binaryrw_test.go

@@ -143,6 +143,35 @@ func TestBufBinWriter_Len(t *testing.T) {
 	require.Equal(t, 1, bw.Len())
 }
 
+func TestBinReader_ReadVarBytes(t *testing.T) {
+	buf := make([]byte, 11)
+	for i := range buf {
+		buf[i] = byte(i)
+	}
+	w := NewBufBinWriter()
+	w.WriteVarBytes(buf)
+	require.NoError(t, w.Err)
+	data := w.Bytes()
+
+	t.Run("NoArguments", func(t *testing.T) {
+		r := NewBinReaderFromBuf(data)
+		actual := r.ReadVarBytes()
+		require.NoError(t, r.Err)
+		require.Equal(t, buf, actual)
+	})
+	t.Run("Good", func(t *testing.T) {
+		r := NewBinReaderFromBuf(data)
+		actual := r.ReadVarBytes(11)
+		require.NoError(t, r.Err)
+		require.Equal(t, buf, actual)
+	})
+	t.Run("Bad", func(t *testing.T) {
+		r := NewBinReaderFromBuf(data)
+		r.ReadVarBytes(10)
+		require.Error(t, r.Err)
+	})
+}
+
 func TestWriterErrHandling(t *testing.T) {
 	var badio = &badRW{}
 	bw := NewBinWriterFromIO(badio)