[ZKTRIE-004] Shallow copy can miscalculate the state root hash

[this-is-chainlight]

Summary

Shallow copy can miscalculate the state root hash.

Description

Any time a node is mutated in a way which could change its hash (i.e. SetChild), first copy it. When SetChild is only being used to replace a HashNode with its real node, the hash will not change, so we do not need to copy as long as we don’t accidentally clear the hash. Multiple threads could be doing the replacement concurrently, so we need to be more careful about detecting this case. This is handled by comparing the child hashes instead of checking the node type.

Impact

Low

The statedb calculation could failed.

Recommendation

diff --git a/trie/zk/merkle_tree.go b/trie/zk/merkle_tree.go
index 3c7492bbd..40d9d1abb 100644
--- a/trie/zk/merkle_tree.go
+++ b/trie/zk/merkle_tree.go
@@ -182,6 +182,7 @@ func (t *MerkleTree) addLeaf(
                        log.Error("fail to addLeaf", "err", err, "level", lvl)
                        return nil, err
                }
+               n = n.Copy()
                n.SetChild(path.Get(lvl), newNode) // Update the node to reflect the modified child
                return n, nil
        case *LeafNode:
@@ -250,6 +251,10 @@ func (t *MerkleTree) Delete(key []byte) error {
        for lvl, p := range path {
                switch n := node.(type) {
                case *ParentNode:
+                       n = n.Copy()
+                       if lvl > 0 {
+                               pathNodes[len(pathNodes)-1].SetChild(path.Get(lvl-1), n)
+                       }
                        pathNodes = append(pathNodes, n)
                        node = t.getChild(n, p)
                case *LeafNode:
diff --git a/trie/zk/merkle_tree_node.go b/trie/zk/merkle_tree_node.go
index b89214d56..522345807 100644
--- a/trie/zk/merkle_tree_node.go
+++ b/trie/zk/merkle_tree_node.go
@@ -68,6 +68,10 @@ func newParentNodeFromBlob(blob []byte) (*ParentNode, error) {
        }, nil
 }

+func (n *ParentNode) Copy() *ParentNode {
+       return &ParentNode{childL: n.childL, childR: n.childR, hash: n.hash}
+}
+
 func (n *ParentNode) Hash() *zkt.Hash { return n.hash }

 func (n *ParentNode) CanonicalValue() []byte {
@@ -92,8 +96,8 @@ func (n *ParentNode) SetChild(path byte, child TreeNode) {
        } else {
                n.childL = child
        }
-       if _, ok := oldChild.(*HashNode); ok && child.Hash() != nil && bytes.Equal(oldChild.Hash()[:], child.Hash()[:]) {
-               // This is a case of converting a HashNode to the original TreeNode. Does not clear the hash.
+       if oldChild.Hash() != nil && child.Hash() != nil && bytes.Equal(oldChild.Hash()[:], child.Hash()[:]) {
+               // The child hash has not changed. Does not clear the hash.
                return
        }
        n.hash = nil

Reference

• N/A

Patch

Reported

[this-is-chainlight]

diff --git a/trie/zk/merkle_tree.go b/trie/zk/merkle_tree.go
index 3c7492bbd..40d9d1abb 100644
--- a/trie/zk/merkle_tree.go
+++ b/trie/zk/merkle_tree.go
@@ -182,6 +182,7 @@ func (t *MerkleTree) addLeaf(
                        log.Error("fail to addLeaf", "err", err, "level", lvl)
                        return nil, err
                }
+               n = n.Copy()
                n.SetChild(path.Get(lvl), newNode) // Update the node to reflect the modified child
                return n, nil
        case *LeafNode:
@@ -250,6 +251,10 @@ func (t *MerkleTree) Delete(key []byte) error {
        for lvl, p := range path {
                switch n := node.(type) {
                case *ParentNode:
+                       n = n.Copy()
+                       if lvl > 0 {
+                               pathNodes[len(pathNodes)-1].SetChild(path.Get(lvl-1), n)
+                       }
                        pathNodes = append(pathNodes, n)
                        node = t.getChild(n, p)
                case *LeafNode:
diff --git a/trie/zk/merkle_tree_node.go b/trie/zk/merkle_tree_node.go
index b89214d56..205c6ee2a 100644
--- a/trie/zk/merkle_tree_node.go
+++ b/trie/zk/merkle_tree_node.go
@@ -68,6 +68,10 @@ func newParentNodeFromBlob(blob []byte) (*ParentNode, error) {
        }, nil
 }

+func (n *ParentNode) Copy() *ParentNode {
+       return &ParentNode{childL: n.childL, childR: n.childR, hash: n.hash}
+}
+
 func (n *ParentNode) Hash() *zkt.Hash { return n.hash }

 func (n *ParentNode) CanonicalValue() []byte {
@@ -92,8 +96,8 @@ func (n *ParentNode) SetChild(path byte, child TreeNode) {
        } else {
                n.childL = child
        }
-       if _, ok := oldChild.(*HashNode); ok && child.Hash() != nil && bytes.Equal(oldChild.Hash()[:], child.Hash()[:]) {
-               // This is a case of converting a HashNode to the original TreeNode. Does not clear the hash.
+       if oldChild != nil && oldChild.Hash() != nil && child.Hash() != nil && bytes.Equal(oldChild.Hash()[:], child.Hash()[:]) {
+               // The child hash has not changed. Does not clear the hash.
                return
        }
        n.hash = nil