Merge pull request #1271 from martin-sucha/fix-range-check-gocql

Fix range checks in readCollectionSize.

AUTHORS

@@ -110,3 +110,4 @@ Alex Lourie <alex@instaclustr.com>; <djay.il@gmail.com>
 Marco Cadetg <cadetg@gmail.com>
 Karl Matthias <karl@matthias.org>
 Thomas Meson <zllak@hycik.org>
+Martin Sucha <martin.sucha@kiwi.com>; <git@mm.ms47.eu>

marshal.go

@@ -1400,11 +1400,17 @@ func marshalList(info TypeInfo, value interface{}) ([]byte, error) {
 	return nil, marshalErrorf("can not marshal %T into %s", value, info)
 }
 
-func readCollectionSize(info CollectionType, data []byte) (size, read int) {
+func readCollectionSize(info CollectionType, data []byte) (size, read int, err error) {
 	if info.proto > protoVersion2 {
+		if len(data) < 4 {
+			return 0, 0, unmarshalErrorf("unmarshal list: unexpected eof")
+		}
 		size = int(data[0])<<24 | int(data[1])<<16 | int(data[2])<<8 | int(data[3])
 		read = 4
 	} else {
+		if len(data) < 2 {
+			return 0, 0, unmarshalErrorf("unmarshal list: unexpected eof")
+		}
 		size = int(data[0])<<8 | int(data[1])
 		read = 2
 	}
@@ -1437,10 +1443,10 @@ func unmarshalList(info TypeInfo, data []byte, value interface{}) error {
 			rv.Set(reflect.Zero(t))
 			return nil
 		}
-		if len(data) < 2 {
-			return unmarshalErrorf("unmarshal list: unexpected eof")
+		n, p, err := readCollectionSize(listInfo, data)
+		if err != nil {
+			return err
 		}
-		n, p := readCollectionSize(listInfo, data)
 		data = data[p:]
 		if k == reflect.Array {
 			if rv.Len() != n {
@@ -1450,10 +1456,10 @@ func unmarshalList(info TypeInfo, data []byte, value interface{}) error {
 			rv.Set(reflect.MakeSlice(t, n, n))
 		}
 		for i := 0; i < n; i++ {
-			if len(data) < 2 {
-				return unmarshalErrorf("unmarshal list: unexpected eof")
+			m, p, err := readCollectionSize(listInfo, data)
+			if err != nil {
+				return err
 			}
-			m, p := readCollectionSize(listInfo, data)
 			data = data[p:]
 			if err := Unmarshal(listInfo.Elem, data[:m], rv.Index(i).Addr().Interface()); err != nil {
 				return err
@@ -1538,16 +1544,16 @@ func unmarshalMap(info TypeInfo, data []byte, value interface{}) error {
 		return nil
 	}
 	rv.Set(reflect.MakeMap(t))
-	if len(data) < 2 {
-		return unmarshalErrorf("unmarshal map: unexpected eof")
+	n, p, err := readCollectionSize(mapInfo, data)
+	if err != nil {
+		return err
 	}
-	n, p := readCollectionSize(mapInfo, data)
 	data = data[p:]
 	for i := 0; i < n; i++ {
-		if len(data) < 2 {
-			return unmarshalErrorf("unmarshal list: unexpected eof")
+		m, p, err := readCollectionSize(mapInfo, data)
+		if err != nil {
+			return err
 		}
-		m, p := readCollectionSize(mapInfo, data)
 		data = data[p:]
 		key := reflect.New(t.Key())
 		if err := Unmarshal(mapInfo.Key, data[:m], key.Interface()); err != nil {
@@ -1555,7 +1561,10 @@ func unmarshalMap(info TypeInfo, data []byte, value interface{}) error {
 		}
 		data = data[m:]
 
-		m, p = readCollectionSize(mapInfo, data)
+		m, p, err = readCollectionSize(mapInfo, data)
+		if err != nil {
+			return err
+		}
 		data = data[p:]
 		val := reflect.New(t.Elem())
 		if err := Unmarshal(mapInfo.Elem, data[:m], val.Interface()); err != nil {

marshal_test.go

@@ -1480,3 +1480,88 @@ func TestMarshalDuration(t *testing.T) {
 		}
 	}
 }
+
+func TestReadCollectionSize(t *testing.T) {
+	listV2 := CollectionType{
+		NativeType: NativeType{proto: 2, typ: TypeList},
+		Elem:       NativeType{proto: 2, typ: TypeVarchar},
+	}
+	listV3 := CollectionType{
+		NativeType: NativeType{proto: 3, typ: TypeList},
+		Elem:       NativeType{proto: 3, typ: TypeVarchar},
+	}
+
+	tests := []struct {
+		name string
+		info CollectionType
+		data []byte
+		isError bool
+		expectedSize int
+	}{
+		{
+			name: "short read 0 proto 2",
+			info: listV2,
+			data: []byte{},
+			isError: true,
+		},
+		{
+			name: "short read 1 proto 2",
+			info: listV2,
+			data: []byte{0x01},
+			isError: true,
+		},
+		{
+			name: "good read proto 2",
+			info: listV2,
+			data: []byte{0x01, 0x38},
+			expectedSize: 0x0138,
+		},
+		{
+			name: "short read 0 proto 3",
+			info: listV3,
+			data: []byte{},
+			isError: true,
+		},
+		{
+			name: "short read 1 proto 3",
+			info: listV3,
+			data: []byte{0x01},
+			isError: true,
+		},
+		{
+			name: "short read 2 proto 3",
+			info: listV3,
+			data: []byte{0x01, 0x38},
+			isError: true,
+		},
+		{
+			name: "short read 3 proto 3",
+			info: listV3,
+			data: []byte{0x01, 0x38, 0x42},
+			isError: true,
+		},
+		{
+			name: "good read proto 3",
+			info: listV3,
+			data: []byte{0x01, 0x38, 0x42, 0x22},
+			expectedSize: 0x01384222,
+		},
+	}
+	for _, test := range tests {
+		t.Run(test.name, func(t *testing.T) {
+			size, _, err := readCollectionSize(test.info, test.data)
+			if test.isError {
+				if err == nil {
+					t.Fatal("Expected error, but it was nil")
+				}
+			} else {
+				if err != nil {
+					t.Fatalf("Expected no error, got %v", err)
+				}
+				if size != test.expectedSize {
+					t.Fatalf("Expected size of %d, but got %d", test.expectedSize, size)
+				}
+			}
+		})
+	}
+}