diff --git a/encoding/wkb/collection.go b/encoding/wkb/collection.go index ad8535b..21cf7c4 100644 --- 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 index 460aaca..78e86dd 100644 --- 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 index ac76632..03f4bcc 100644 --- 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 index 7b6b6b1..9603e22 100644 --- 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 index a3e66b7..5408a29 100644 --- 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 index 785954a..7b12aa1 100644 --- 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{