minecraft/nbt: Bound array and list allocations by the input

minecraft/nbt/decode.go

@@ -254,17 +254,11 @@ func (d *Decoder) unmarshalTag(val reflect.Value, t tagType, tagName string) err
 		if err != nil {
 			return err
 		}
-		if err = d.checkRemaining(int(length), "ByteArray"); err != nil {
+		b, err := d.r.readArrayBytes(length, 1, "ByteArray")
+		if err != nil {
 			return err
 		}
-		if length < 0 {
-			return BufferOverrunError{Op: "ByteArray"}
-		}
-		b := make([]byte, length)
-		if _, err := d.r.Read(b); err != nil {
-			return BufferOverrunError{Op: "ByteArray"}
-		}
-		value := reflect.New(reflect.ArrayOf(int(length), byteType)).Elem()
+		value := reflect.New(reflect.ArrayOf(len(b), byteType)).Elem()
 		reflect.Copy(value, reflect.ValueOf(b))
 
 		switch {
@@ -348,13 +342,10 @@ func (d *Decoder) unmarshalTag(val reflect.Value, t tagType, tagName string) err
 				val.Set(reflect.MakeSlice(sliceType, int(length), int(length)))
 				break
 			}
-			if err = d.checkRemaining(int(length), "ByteSlice"); err != nil {
+			b, err := d.r.readArrayBytes(length, 1, "ByteSlice")
+			if err != nil {
 				return err
 			}
-			b := make([]byte, length)
-			if _, err := d.r.Read(b); err != nil {
-				return BufferOverrunError{Op: "ByteSlice"}
-			}
 			switch {
 			case k == reflect.Slice && val.Type().Elem().Kind() == reflect.Uint8, isAny(val):
 				val.Set(reflect.ValueOf(b))
@@ -388,11 +379,18 @@ func (d *Decoder) unmarshalTag(val reflect.Value, t tagType, tagName string) err
 			if err != nil {
 				return err
 			}
-			if length < 0 {
-				return BufferOverrunError{Op: "Slice"}
+			c, err := d.r.sliceCap(length, 1, "Slice")
+			if err != nil {
+				return err
 			}
-			v := reflect.MakeSlice(sliceType, int(length), int(length))
+			// Addressable so the slice can grow in place when the capacity was bounded below length.
+			v := reflect.New(sliceType).Elem()
+			v.Set(reflect.MakeSlice(sliceType, 0, c))
 			for i := 0; i < int(length); i++ {
+				if i == v.Cap() {
+					v.Grow(min(int(length)-i, v.Cap()))
+				}
+				v.SetLen(i + 1)
 				if err := d.unmarshalTag(v.Index(i), listType, ""); err != nil {
 					// An error occurred during the decoding of one of the elements of the TAG_List, meaning it
 					// either had an invalid type or the NBT was invalid.
@@ -545,13 +543,6 @@ func (d *Decoder) tag() (t tagType, tagName string, err error) {
 	return t, tagName, err
 }
 
-func (d *Decoder) checkRemaining(length int, op string) error {
-	if remaining, ok := d.r.Reader.(interface{ Len() int }); ok && length > remaining.Len() {
-		return BufferOverrunError{Op: op}
-	}
-	return nil
-}
-
 // isAny checks if a reflect.Value has the type `any` or `interface{}`.
 func isAny(v reflect.Value) bool {
 	return v.Kind() == reflect.Interface && v.NumMethod() == 0

minecraft/nbt/encoding.go

@@ -191,15 +191,17 @@ func (e networkLittleEndian) Int32Slice(r *offsetReader) ([]int32, error) {
 	if err != nil {
 		return nil, BufferOverrunError{Op: "Int32Slice"}
 	}
-	if n < 0 {
-		return nil, BufferOverrunError{Op: "Int32Slice"}
+	c, err := r.sliceCap(n, 1, "Int32Slice")
+	if err != nil {
+		return nil, err
 	}
-	m := make([]int32, n)
+	m := make([]int32, 0, c)
 	for i := int32(0); i < n; i++ {
-		m[i], err = e.Int32(r)
+		v, err := e.Int32(r)
 		if err != nil {
 			return nil, BufferOverrunError{Op: "Int32Slice"}
 		}
+		m = append(m, v)
 	}
 	return m, nil
 }
@@ -210,15 +212,17 @@ func (e networkLittleEndian) Int64Slice(r *offsetReader) ([]int64, error) {
 	if err != nil {
 		return nil, BufferOverrunError{Op: "Int64Slice"}
 	}
-	if n < 0 {
-		return nil, BufferOverrunError{Op: "Int64Slice"}
+	c, err := r.sliceCap(n, 1, "Int64Slice")
+	if err != nil {
+		return nil, err
 	}
-	m := make([]int64, n)
+	m := make([]int64, 0, c)
 	for i := int32(0); i < n; i++ {
-		m[i], err = e.Int64(r)
+		v, err := e.Int64(r)
 		if err != nil {
 			return nil, BufferOverrunError{Op: "Int64Slice"}
 		}
+		m = append(m, v)
 	}
 	return m, nil
 }

minecraft/nbt/reader.go

@@ -81,3 +81,19 @@ func (b *offsetReader) readArrayBytes(count int32, width int, op string) ([]byte
 	}
 	return data, nil
 }
+
+// sliceCap bounds the capacity preallocated for count variable-width elements of at least minWidth bytes
+// each. A reader that reports its length rejects more elements than it can hold; otherwise the capacity
+// starts small and the slice grows with the elements actually read.
+func (b *offsetReader) sliceCap(count int32, minWidth int, op string) (int, error) {
+	if count < 0 {
+		return 0, BufferOverrunError{Op: op}
+	}
+	if remaining, ok := b.Reader.(interface{ Len() int }); ok {
+		if int64(count)*int64(minWidth) > int64(remaining.Len()) {
+			return 0, BufferOverrunError{Op: op}
+		}
+		return int(count), nil
+	}
+	return min(int(count), 256), nil
+}