M tree.go → tree.go

@@ -39,6 +39,9 @@ PrefixValueIntermediate = 2
 
 	// nChars is used to crop the Graphviz nodes labels
 	nChars = 4
+
+	maxKeyLen           = 255
+	maxSiblingsFieldLen = 65535
 )
 
 var (
@@ -243,16 +246,21 @@ // the same leafs.
 func (t *Tree) loadVT(rTx db.ReadTx) (vt, error) {
 	vt := newVT(t.maxLevels, t.hashFunction)
 	vt.params.dbg = t.dbg
-	err := t.IterateWithTx(rTx, nil, func(k, v []byte) {
+	var callbackErr error
+	err := t.IterateWithStopWithTx(rTx, nil, func(_ int, k, v []byte) bool {
 		if v[0] != PrefixValueLeaf {
-			return
+			return false
 		}
 		leafK, leafV := ReadLeafValue(v)
 		if err := vt.add(0, leafK, leafV); err != nil {
-			// TODO instead of panic, return this error
-			panic(err)
+			callbackErr = err
+			return true
 		}
+		return false
 	})
+	if callbackErr != nil {
+		return vt, callbackErr
+	}
 
 	return vt, err
 }
@@ -304,6 +312,9 @@ }
 
 func (t *Tree) add(wTx db.WriteTx, root []byte, fromLvl int, k, v []byte) ([]byte, error) {
 	keyPath := make([]byte, t.hashFunction.Len())
+	if len(k) > len(keyPath) {
+		return nil, fmt.Errorf("len(k) > len(keyPath)")
+	}
 	copy(keyPath[:], k)
 
 	path := getPath(t.maxLevels, keyPath)
@@ -376,6 +387,10 @@ return nil, nil, nil, ErrKeyAlreadyExists
 			}
 
 			oldLeafKeyFull := make([]byte, t.hashFunction.Len())
+			if len(oldLeafKey) > len(oldLeafKeyFull) {
+				return nil, nil, nil,
+					fmt.Errorf("len(oldLeafKey) > len(oldLeafKeyFull)")
+			}
 			copy(oldLeafKeyFull[:], oldLeafKey)
 
 			// if currKey is already used, go down until paths diverge
@@ -479,6 +494,9 @@ return nil, nil, err
 	}
 	var leafValue []byte
 	leafValue = append(leafValue, byte(PrefixValueLeaf))
+	if len(k) > maxKeyLen {
+		return nil, nil, fmt.Errorf("len(k) > %v", maxKeyLen)
+	}
 	leafValue = append(leafValue, byte(len(k)))
 	leafValue = append(leafValue, k...)
 	leafValue = append(leafValue, v...)
@@ -514,6 +532,9 @@ // [ type of node | length of left key | left key | right key ]
 func newIntermediate(hashFunc HashFunction, l, r []byte) ([]byte, []byte, error) {
 	b := make([]byte, PrefixValueLen+hashFunc.Len()*2)
 	b[0] = PrefixValueIntermediate
+	if len(l) > maxKeyLen {
+		return nil, nil, fmt.Errorf("len(k) > %v", maxKeyLen)
+	}
 	b[1] = byte(len(l))
 	copy(b[PrefixValueLen:PrefixValueLen+hashFunc.Len()], l)
 	copy(b[PrefixValueLen+hashFunc.Len():], r)
@@ -575,6 +596,9 @@
 	var err error
 
 	keyPath := make([]byte, t.hashFunction.Len())
+	if len(k) > len(keyPath) {
+		return fmt.Errorf("len(k) > len(keyPath)")
+	}
 	copy(keyPath[:], k)
 	path := getPath(t.maxLevels, keyPath)
 
@@ -632,6 +656,9 @@ // GenProofWithTx does the same than the GenProof method, but allowing to pass
 // the db.ReadTx that is used.
 func (t *Tree) GenProofWithTx(rTx db.ReadTx, k []byte) ([]byte, []byte, []byte, bool, error) {
 	keyPath := make([]byte, t.hashFunction.Len())
+	if len(k) > len(keyPath) {
+		return nil, nil, nil, false, fmt.Errorf("len(k) > len(keyPath)")
+	}
 	copy(keyPath[:], k)
 
 	root, err := t.RootWithTx(rTx)
@@ -647,7 +674,10 @@ if err != nil {
 		return nil, nil, nil, false, err
 	}
 
-	s := PackSiblings(t.hashFunction, siblings)
+	s, err := PackSiblings(t.hashFunction, siblings)
+	if err != nil {
+		return nil, nil, nil, false, err
+	}
 
 	leafK, leafV := ReadLeafValue(value)
 	if !bytes.Equal(k, leafK) {
@@ -665,7 +695,7 @@ // Where the bitmap indicates if the sibling is 0 or a value from the siblings
 // array. And S is the size of the output of the hash function used for the
 // Tree. The 2 2-byte that define the full length and bitmap length, are
 // encoded in little-endian.
-func PackSiblings(hashFunc HashFunction, siblings [][]byte) []byte {
+func PackSiblings(hashFunc HashFunction, siblings [][]byte) ([]byte, error) {
 	var b []byte
 	var bitmap []bool
 	emptySibling := make([]byte, hashFunc.Len())
@@ -683,11 +713,17 @@ l := len(bitmapBytes)
 
 	fullLen := 4 + l + len(b) //nolint:gomnd
 	res := make([]byte, fullLen)
+	if fullLen > maxSiblingsFieldLen {
+		return nil, fmt.Errorf("fullLen > %v", maxSiblingsFieldLen)
+	}
 	binary.LittleEndian.PutUint16(res[0:2], uint16(fullLen)) // set full length
-	binary.LittleEndian.PutUint16(res[2:4], uint16(l))       // set the bitmapBytes length
+	if l > maxSiblingsFieldLen {
+		return nil, fmt.Errorf("l > %v", maxSiblingsFieldLen)
+	}
+	binary.LittleEndian.PutUint16(res[2:4], uint16(l)) // set the bitmapBytes length
 	copy(res[4:4+l], bitmapBytes)
 	copy(res[4+l:], b)
-	return res
+	return res, nil
 }
 
 // UnpackSiblings unpacks the siblings from a byte array.
@@ -1038,18 +1074,31 @@
 	// WARNING current encoding only supports key & values of 255 bytes each
 	// (due using only 1 byte for the length headers).
 	var b []byte
-	err := t.Iterate(fromRoot, func(k, v []byte) {
+	var callbackErr error
+	err := t.IterateWithStop(fromRoot, func(_ int, k, v []byte) bool {
 		if v[0] != PrefixValueLeaf {
-			return
+			return false
 		}
 		leafK, leafV := ReadLeafValue(v)
 		kv := make([]byte, 2+len(leafK)+len(leafV))
+		if len(leafK) > maxKeyLen {
+			callbackErr = fmt.Errorf("len(leafK) > %v", maxKeyLen)
+			return true
+		}
 		kv[0] = byte(len(leafK))
+		if len(leafV) > maxKeyLen {
+			callbackErr = fmt.Errorf("len(leafV) > %v", maxKeyLen)
+			return true
+		}
 		kv[1] = byte(len(leafV))
 		copy(kv[2:2+len(leafK)], leafK)
 		copy(kv[2+len(leafK):], leafV)
 		b = append(b, kv...)
+		return false
 	})
+	if callbackErr != nil {
+		return nil, callbackErr
+	}
 	return b, err
 }