Avoid storing hashing function pointers in storage make MerkleTree structure upgrade-safe (#5080)

Co-authored-by: cairo <cairoeth@protonmail.com>
Co-authored-by: ernestognw <ernestognw@gmail.com>

contracts/utils/structs/MerkleTree.sol

@@ -17,7 +17,7 @@ import {Panic} from "../Panic.sol";
  *
  * * Depth: The number of levels in the tree, it also defines the maximum number of leaves as 2**depth.
  * * Zero value: The value that represents an empty leaf. Used to avoid regular zero values to be part of the tree.
- * * Hashing function: A cryptographic hash function used to produce internal nodes.
+ * * Hashing function: A cryptographic hash function used to produce internal nodes. Defaults to {Hashes-commutativeKeccak256}.
  *
  * _Available since v5.1._
  */
@@ -27,9 +27,6 @@ library MerkleTree {
      *
      * The `sides` and `zero` arrays are set to have a length equal to the depth of the tree during setup.
      *
-     * The hashing function used during initialization to compute the `zeros` values (value of a node at a given depth
-     * for which the subtree is full of zero leaves). This function is kept in the structure for handling insertions.
-     *
      * Struct members have an underscore prefix indicating that they are "private" and should not be read or written to
      * directly. Use the functions provided below instead. Modifying the struct manually may violate assumptions and
      * lead to unexpected behavior.
@@ -44,7 +41,6 @@ library MerkleTree {
         uint256 _nextLeafIndex;
         bytes32[] _sides;
         bytes32[] _zeros;
-        function(bytes32, bytes32) view returns (bytes32) _fnHash;
     }
 
     /**
@@ -53,6 +49,9 @@ library MerkleTree {
      *
      * Calling this function on MerkleTree that was already setup and used will reset it to a blank state.
      *
+     * Once a tree is setup, any push to it must use the same hashing function. This means that values
+     * should be pushed to it using the default {xref-MerkleTree-push-struct-MerkleTree-Bytes32PushTree-bytes32-}[push] function.
+     *
      * IMPORTANT: The zero value should be carefully chosen since it will be stored in the tree representing
      * empty leaves. It should be a value that is not expected to be part of the tree.
      */
@@ -61,7 +60,10 @@ library MerkleTree {
     }
 
     /**
-     * @dev Same as {setup}, but allows to specify a custom hashing function.
+     * @dev Same as {xref-MerkleTree-setup-struct-MerkleTree-Bytes32PushTree-uint8-bytes32-}[setup], but allows to specify a custom hashing function.
+     *
+     * Once a tree is setup, any push to it must use the same hashing function. This means that values
+     * should be pushed to it using the custom push function, which should be the same one as used during the setup.
      *
      * IMPORTANT: Providing a custom hashing function is a security-sensitive operation since it may
      * compromise the soundness of the tree. Consider using functions from {Hashes}.
@@ -85,7 +87,6 @@ library MerkleTree {
 
         // Set the first root
         self._nextLeafIndex = 0;
-        self._fnHash = fnHash;
 
         return currentZero;
     }
@@ -96,11 +97,32 @@ library MerkleTree {
      *
      * Hashing the leaf before calling this function is recommended as a protection against
      * second pre-image attacks.
+     *
+     * This variant uses {Hashes-commutativeKeccak256} to hash internal nodes. It should only be used on merkle trees
+     * that were setup using the same (default) hashing function (i.e. by calling
+     * {xref-MerkleTree-setup-struct-MerkleTree-Bytes32PushTree-uint8-bytes32-}[the default setup] function).
      */
     function push(Bytes32PushTree storage self, bytes32 leaf) internal returns (uint256 index, bytes32 newRoot) {
+        return push(self, leaf, Hashes.commutativeKeccak256);
+    }
+
+    /**
+     * @dev Insert a new leaf in the tree, and compute the new root. Returns the position of the inserted leaf in the
+     * tree, and the resulting root.
+     *
+     * Hashing the leaf before calling this function is recommended as a protection against
+     * second pre-image attacks.
+     *
+     * This variant uses a custom hashing function to hash internal nodes. It should only be called with the same
+     * function as the one used during the initial setup of the merkle tree.
+     */
+    function push(
+        Bytes32PushTree storage self,
+        bytes32 leaf,
+        function(bytes32, bytes32) view returns (bytes32) fnHash
+    ) internal returns (uint256 index, bytes32 newRoot) {
         // Cache read
         uint256 levels = self._zeros.length;
-        function(bytes32, bytes32) view returns (bytes32) fnHash = self._fnHash;
 
         // Get leaf index
         index = self._nextLeafIndex++;
@@ -123,7 +145,7 @@ library MerkleTree {
             }
 
             // Compute the current node hash by using the hash function
-            // with either the its sibling (side) or the zero value for that level.
+            // with either its sibling (side) or the zero value for that level.
             currentLevelHash = fnHash(
                 isLeft ? currentLevelHash : Arrays.unsafeAccess(self._sides, i).value,
                 isLeft ? Arrays.unsafeAccess(self._zeros, i).value : currentLevelHash

docs/modules/ROOT/pages/utilities.adoc

@@ -157,10 +157,10 @@ Building an on-chain Merkle Tree allow developers to keep track of the history o
 
 The Merkle Tree does not keep track of the roots purposely, so that developers can choose their tracking mechanism. Setting up and using an Merkle Tree in Solidity is as simple as follows:
 
+NOTE: Functions are exposed without access control for demonstration purposes
+
 [source,solidity]
 ----
-// NOTE: Functions are exposed without access control for demonstration purposes
-
 using MerkleTree for MerkleTree.Bytes32PushTree;
 MerkleTree.Bytes32PushTree private _tree;
 
@@ -174,6 +174,31 @@ function push(bytes32 leaf) public /* onlyOwner */ {
 }
 ----
 
+The library also supports custom hashing functions, which can be passed as an extra parameter to the xref:api:utils.adoc#MerkleTree-push-struct-MerkleTree-Bytes32PushTree-bytes32-[`push`] and xref:api:utils.adoc#MerkleTree-setup-struct-MerkleTree-Bytes32PushTree-uint8-bytes32-[`setup`] functions.
+
+Using custom hashing functions is a sensitive operation. After setup, it requires to keep using the same hashing function for every new valued pushed to the tree to avoid corrupting the tree. For this reason, it's a good practice to keep your hashing function static in your implementation contract as follows:
+
+[source,solidity]
+----
+using MerkleTree for MerkleTree.Bytes32PushTree;
+MerkleTree.Bytes32PushTree private _tree;
+
+function setup(uint8 _depth, bytes32 _zero) public /* onlyOwner */ {
+    root = _tree.setup(_depth, _zero, _hashFn);
+}
+
+function push(bytes32 leaf) public /* onlyOwner */ {
+    (uint256 leafIndex, bytes32 currentRoot) = _tree.push(leaf, _hashFn);
+    // Store the new root.
+}
+
+function _hashFn(bytes32 a, bytes32 b) internal view returns(bytes32) {
+    // Custom hash function implementation
+    // Kept as an internal implementation detail to 
+    // guarantee the same function is always used
+}
+---- 
+
 [[misc]]
 == Misc