M tree.go → tree.go

@@ -405,12 +405,10 @@ if bytes.Equal(newKey, oldLeafKey) {
 				return nil, nil, nil, ErrKeyAlreadyExists
 			}
 
-			oldLeafKeyFull := make([]byte, t.hashFunction.Len())
-			// if len(oldLeafKey) > t.hashFunction.Len() { // WIP
-			//         return nil, nil, nil,
-			//                 fmt.Errorf("len(oldLeafKey) > hashFunction.Len()")
-			// }
-			copy(oldLeafKeyFull[:], oldLeafKey)
+			oldLeafKeyFull, err := keyPathFromKey(t.maxLevels, oldLeafKey)
+			if err != nil {
+				return nil, nil, nil, err
+			}
 
 			// if currKey is already used, go down until paths diverge
 			oldPath := getPath(t.maxLevels, oldLeafKeyFull)
@@ -841,7 +839,7 @@ if err != nil {
 		return false, err
 	}
 
-	keyPath := make([]byte, len(siblings))
+	keyPath := make([]byte, int(math.Ceil(float64(len(siblings))/float64(8)))) //nolint:gomnd
 	copy(keyPath[:], k)
 
 	key, _, err := newLeafValue(hashFunc, k, v)

M tree_test.go → tree_test.go

@@ -745,6 +745,30 @@ c.Assert(err, qt.IsNil)
 	c.Assert(verif, qt.IsFalse)
 }
 
+func TestKeyLenBiggerThan32(t *testing.T) {
+	c := qt.New(t)
+	maxLevels := 264
+	database, err := badgerdb.New(badgerdb.Options{Path: c.TempDir()})
+	c.Assert(err, qt.IsNil)
+	tree, err := NewTree(database, maxLevels, HashFunctionBlake2b)
+	c.Assert(err, qt.IsNil)
+
+	bLen := 33
+	err = tree.Add(
+		randomBytes(bLen),
+		randomBytes(bLen))
+	c.Assert(err, qt.IsNil)
+
+	// 2nd key that we add, will find a node with len(key)==32 (due the
+	// hash output size, expect that next Add does not give any error, as
+	// internally it will use a keyPath of size corresponent to the
+	// maxLevels size of the tree
+	err = tree.Add(
+		randomBytes(bLen),
+		randomBytes(bLen))
+	c.Assert(err, qt.IsNil)
+}
+
 func BenchmarkAdd(b *testing.B) {
 	bLen := 32 // for both Poseidon & Sha256
 	// prepare inputs