encoding/wkb: only limit multi allocations not read length

paulmach · Mar 9, 2019 · c6e407b · c6e407b
1 parent 363a051
commit c6e407b
Show file tree

Hide file tree

Showing 6 changed files with 56 additions and 12 deletions.
diff --git a/encoding/wkb/collection.go b/encoding/wkb/collection.go
@@ -13,12 +13,13 @@ func readCollection(r io.Reader, bom binary.ByteOrder) (orb.Collection, error) {
 		return nil, err
 	}
 
-	if num > maxMultiAlloc {
+	alloc := num
+	if alloc > maxMultiAlloc {
 		// invalid data can come in here and allocate tons of memory.
-		num = maxMultiAlloc
+		alloc = maxMultiAlloc
 	}
+	result := make(orb.Collection, 0, alloc)
 
-	result := make(orb.Collection, 0, num)
 	for i := 0; i < int(num); i++ {
 		geom, err := NewDecoder(r).Decode()
 		if err != nil {

diff --git a/encoding/wkb/collection_test.go b/encoding/wkb/collection_test.go
@@ -30,6 +30,11 @@ var (
 )
 
 func TestCollection(t *testing.T) {
+	large := orb.Collection{}
+	for i := 0; i < maxMultiAlloc+100; i++ {
+		large = append(large, orb.Point{float64(i), float64(-i)})
+	}
+
 	cases := []struct {
 		name     string
 		data     []byte
@@ -40,6 +45,11 @@ func TestCollection(t *testing.T) {
 			data:     testCollectionData,
 			expected: testCollection,
 		},
+		{
+			name:     "large",
+			data:     MustMarshal(large),
+			expected: large,
+		},
 	}
 
 	for _, tc := range cases {

diff --git a/encoding/wkb/line_string.go b/encoding/wkb/line_string.go
@@ -60,12 +60,13 @@ func readMultiLineString(r io.Reader, bom binary.ByteOrder) (orb.MultiLineString
 		return nil, err
 	}
 
-	if num > maxMultiAlloc {
+	alloc := num
+	if alloc > maxMultiAlloc {
 		// invalid data can come in here and allocate tons of memory.
-		num = maxMultiAlloc
+		alloc = maxMultiAlloc
 	}
+	result := make(orb.MultiLineString, 0, alloc)
 
-	result := make(orb.MultiLineString, 0, num)
 	for i := 0; i < int(num); i++ {
 		byteOrder, typ, err := readByteOrderType(r)
 		if err != nil {

diff --git a/encoding/wkb/line_string_test.go b/encoding/wkb/line_string_test.go
@@ -94,6 +94,11 @@ var (
 )
 
 func TestMultiLineString(t *testing.T) {
+	large := orb.MultiLineString{}
+	for i := 0; i < maxMultiAlloc+100; i++ {
+		large = append(large, orb.LineString{})
+	}
+
 	cases := []struct {
 		name     string
 		data     []byte
@@ -109,6 +114,11 @@ func TestMultiLineString(t *testing.T) {
 			data:     testMultiLineStringSingleData,
 			expected: testMultiLineStringSingle,
 		},
+		{
+			name:     "large",
+			data:     MustMarshal(large),
+			expected: large,
+		},
 	}
 
 	for _, tc := range cases {

diff --git a/encoding/wkb/polygon.go b/encoding/wkb/polygon.go
@@ -15,12 +15,13 @@ func readPolygon(r io.Reader, bom binary.ByteOrder) (orb.Polygon, error) {
 		return nil, err
 	}
 
-	if num > maxMultiAlloc {
+	alloc := num
+	if alloc > maxMultiAlloc {
 		// invalid data can come in here and allocate tons of memory.
-		num = maxMultiAlloc
+		alloc = maxMultiAlloc
 	}
+	result := make(orb.Polygon, 0, alloc)
 
-	result := make(orb.Polygon, 0, num)
 	for i := 0; i < int(num); i++ {
 		ls, err := readLineString(r, bom)
 		if err != nil {
@@ -64,12 +65,13 @@ func readMultiPolygon(r io.Reader, bom binary.ByteOrder) (orb.MultiPolygon, erro
 		return nil, err
 	}
 
-	if num > maxMultiAlloc {
+	alloc := num
+	if alloc > maxMultiAlloc {
 		// invalid data can come in here and allocate tons of memory.
-		num = maxMultiAlloc
+		alloc = maxMultiAlloc
 	}
+	result := make(orb.MultiPolygon, 0, alloc)
 
-	result := make(orb.MultiPolygon, 0, num)
 	for i := 0; i < int(num); i++ {
 		byteOrder, typ, err := readByteOrderType(r)
 		if err != nil {

diff --git a/encoding/wkb/polygon_test.go b/encoding/wkb/polygon_test.go
@@ -29,6 +29,11 @@ var (
 )
 
 func TestPolygon(t *testing.T) {
+	large := orb.Polygon{}
+	for i := 0; i < maxMultiAlloc+100; i++ {
+		large = append(large, orb.Ring{})
+	}
+
 	cases := []struct {
 		name     string
 		data     []byte
@@ -39,6 +44,11 @@ func TestPolygon(t *testing.T) {
 			data:     testPolygonData,
 			expected: testPolygon,
 		},
+		{
+			name:     "large",
+			data:     MustMarshal(large),
+			expected: large,
+		},
 		{
 			name: "two ring polygon",
 			data: []byte{
@@ -155,6 +165,11 @@ var (
 )
 
 func TestMultiPolygon(t *testing.T) {
+	large := orb.MultiPolygon{}
+	for i := 0; i < maxMultiAlloc+100; i++ {
+		large = append(large, orb.Polygon{})
+	}
+
 	cases := []struct {
 		name     string
 		data     []byte
@@ -170,6 +185,11 @@ func TestMultiPolygon(t *testing.T) {
 			data:     testMultiPolygonSingleData,
 			expected: testMultiPolygonSingle,
 		},
+		{
+			name:     "large",
+			data:     MustMarshal(large),
+			expected: large,
+		},
 		{
 			name: "three polygons",
 			data: []byte{