[storage/mmr] hash leaf count in last to avoid malleability concern (#3392)

Author: roberto-bayardo

storage/conformance.toml

@@ -72,7 +72,7 @@ hash = "506b5a02d1a7c3e249f514b889e6099c27bb0723ebc711514263f0c71f0676ef"
 
 ["commonware_storage::merkle::mmr::conformance::MmrRootStability"]
 n_cases = 200
-hash = "d9de2d71a41b1c656744b4ff9b1bd8e3a1b7d9da65de654b7551912431eb02d7"
+hash = "4078ce1f5e242f8edb1d41575d51e166af620404036124f8094fc74d4b401047"
 
 ["commonware_storage::merkle::proof::tests::conformance::CodecConformance<Proof<Sha256Digest>>"]
 n_cases = 65536

storage/src/merkle/mmr/hasher.rs

@@ -43,18 +43,23 @@ pub trait Hasher: Clone + Send + Sync {
         self.hash([acc.as_ref(), peak.as_ref()])
     }
 
-    /// Computes the root for an MMR given its size and an iterator over the digests of its peaks in
-    /// decreasing order of height.
+    /// Computes the root for an MMR given its leaf count and peak digests in canonical order.
+    ///
+    /// The root digest is computed as `Hash(leaves || fold(peak_digests))`, where `fold` is
+    /// defined as `fold(acc, peak) = Hash(acc || peak)`. The `peak_digests` are assumed to be
+    /// in canonical order.
     fn root<'a>(
         &mut self,
         leaves: Location,
         peak_digests: impl IntoIterator<Item = &'a Self::Digest>,
     ) -> Self::Digest {
-        let mut acc = self.digest(&leaves.to_be_bytes());
-        for digest in peak_digests {
-            acc = self.fold(&acc, digest);
-        }
-        acc
+        let mut iter = peak_digests.into_iter();
+        let Some(first) = iter.next() else {
+            return self.digest(&leaves.to_be_bytes());
+        };
+        let acc = iter.fold(*first, |acc, digest| self.fold(&acc, digest));
+
+        self.hash([leaves.to_be_bytes().as_slice(), acc.as_ref()])
     }
 }
 

storage/src/merkle/mmr/verification.rs

@@ -92,17 +92,18 @@ impl<D: Digest> ProofStore<D> {
 
         let mut digests: Vec<D> = Vec::new();
         if !bp.fold_prefix.is_empty() {
-            let mut acc = self
-                .fold_acc
-                .unwrap_or_else(|| hasher.digest(&leaves.to_be_bytes()));
+            // Start from the stored fold accumulator (which does not include the leaf count).
+            let mut acc = self.fold_acc;
             // Fold in peaks beyond those already covered by the stored accumulator.
             for &pos in bp.fold_prefix.iter().skip(self.num_fold_peaks) {
                 match self.digests.get(&pos) {
-                    Some(d) => acc = hasher.fold(&acc, d),
+                    Some(d) => {
+                        acc = Some(acc.map_or(*d, |a| hasher.fold(&a, d)));
+                    }
                     None => return Err(Error::ElementPruned(pos)),
                 }
             }
-            digests.push(acc);
+            digests.push(acc.expect("fold_prefix is non-empty so acc must be set"));
         }
 
         for &pos in &bp.fetch_nodes {
@@ -186,14 +187,12 @@ pub async fn historical_range_proof<D: Digest, H: Hasher<Digest = D>, S: Storage
 
     let mut digests: Vec<D> = Vec::new();
     if !bp.fold_prefix.is_empty() {
-        let mut acc = hasher.digest(&leaves.to_be_bytes());
         let node_futures = bp.fold_prefix.iter().map(|&pos| mmr.get_node(pos));
         let results = try_join_all(node_futures).await?;
-        for (i, result) in results.into_iter().enumerate() {
-            match result {
-                Some(d) => acc = hasher.fold(&acc, &d),
-                None => return Err(Error::ElementPruned(bp.fold_prefix[i])),
-            }
+        let mut acc = results[0].ok_or(Error::ElementPruned(bp.fold_prefix[0]))?;
+        for (i, &result) in results.iter().enumerate().skip(1) {
+            let d = result.ok_or(Error::ElementPruned(bp.fold_prefix[i]))?;
+            acc = hasher.fold(&acc, &d);
         }
         digests.push(acc);
     }

storage/src/merkle/proof.rs

@@ -236,9 +236,11 @@ impl<D: Digest> Proof<D> {
                 if bp.fold_prefix.is_empty() { 0 } else { 1 } + bp.fetch_nodes.len(),
             );
             if !bp.fold_prefix.is_empty() {
-                // Fold prefix peaks into accumulator
-                let mut acc = hasher.digest(&self.leaves.to_be_bytes());
-                for &pos in &bp.fold_prefix {
+                // Fold prefix peaks into accumulator (without the leaf count).
+                let mut acc = *node_digests
+                    .get(&bp.fold_prefix[0])
+                    .expect("must exist by construction");
+                for &pos in &bp.fold_prefix[1..] {
                     let d = node_digests.get(&pos).expect("must exist by construction");
                     acc = hasher.fold(&acc, d);
                 }
@@ -350,13 +352,17 @@ impl<D: Digest> Proof<D> {
             .zip(pinned_nodes.iter().copied())
             .collect();
 
-        // Verify fold-prefix pinned nodes by recomputing the accumulator.
+        // Verify fold-prefix pinned nodes by recomputing the accumulator (without the leaf
+        // count, which is hashed into the final root independently).
         if !bp.fold_prefix.is_empty() {
             if self.digests.is_empty() {
                 return false;
             }
-            let mut acc = hasher.digest(&self.leaves.to_be_bytes());
-            for pos in &bp.fold_prefix {
+            let Some(first) = pinned_map.remove(&bp.fold_prefix[0]) else {
+                return false;
+            };
+            let mut acc = first;
+            for pos in &bp.fold_prefix[1..] {
                 let Some(digest) = pinned_map.remove(pos) else {
                     return false;
                 };
@@ -462,12 +468,12 @@ impl<D: Digest> Proof<D> {
         let after_end = after_start + after_peaks.len();
         let siblings = &self.digests[after_end..];
 
-        // Start fold accumulator.
-        let mut acc = if has_prefix {
-            // The first digest is the pre-folded prefix accumulator.
-            self.digests[0]
+        // Fold all peaks into an accumulator (without the leaf count, which is hashed in at the
+        // end to prevent malleability via the `leaves` field).
+        let mut acc: Option<D> = if has_prefix {
+            Some(self.digests[0])
         } else {
-            hasher.digest(&self.leaves.to_be_bytes())
+            None
         };
 
         // Reconstruct each range peak and fold into acc.
@@ -490,7 +496,7 @@ impl<D: Digest> Proof<D> {
             if let Some(ref mut cd) = collected_digests {
                 cd.push((peak_pos, peak_digest));
             }
-            acc = hasher.fold(&acc, &peak_digest);
+            acc = Some(acc.map_or(peak_digest, |a| hasher.fold(&a, &peak_digest)));
         }
 
         // Fold after-peak digests.
@@ -499,7 +505,7 @@ impl<D: Digest> Proof<D> {
             if let Some(ref mut cd) = collected_digests {
                 cd.push((after_peak_pos, digest));
             }
-            acc = hasher.fold(&acc, &digest);
+            acc = Some(acc.map_or(digest, |a| hasher.fold(&a, &digest)));
         }
 
         // Verify all elements were consumed.
@@ -512,7 +518,12 @@ impl<D: Digest> Proof<D> {
             return Err(ReconstructionError::ExtraDigests);
         }
 
-        Ok(acc)
+        // Hash the leaf count into the final result.
+        Ok(if let Some(peaks_acc) = acc {
+            hasher.hash([self.leaves.to_be_bytes().as_slice(), peaks_acc.as_ref()])
+        } else {
+            hasher.digest(&self.leaves.to_be_bytes())
+        })
     }
 }
 
@@ -640,10 +651,11 @@ where
     let mut digests =
         Vec::with_capacity(if bp.fold_prefix.is_empty() { 0 } else { 1 } + bp.fetch_nodes.len());
 
-    // Fold prefix peaks into a single accumulator.
+    // Fold prefix peaks into a single accumulator (without the leaf count, which is always
+    // hashed into the final root independently).
     if !bp.fold_prefix.is_empty() {
-        let mut acc = hasher.digest(&leaves.to_be_bytes());
-        for &pos in &bp.fold_prefix {
+        let mut acc = get_node(bp.fold_prefix[0]).ok_or(Error::ElementPruned(bp.fold_prefix[0]))?;
+        for &pos in &bp.fold_prefix[1..] {
             let d = get_node(pos).ok_or(Error::ElementPruned(pos))?;
             acc = hasher.fold(&acc, &d);
         }
@@ -1822,6 +1834,43 @@ mod tests {
         );
     }
 
+    /// Regression test: mutating only the `leaves` field in a proof must invalidate it.
+    /// Before the fix, when a fold prefix existed, the leaf count was baked into the
+    /// pre-folded accumulator but not independently checked during verification, so a
+    /// different `leaves` value with a compatible peak structure would still verify.
+    #[test]
+    fn test_proof_leaves_malleability() {
+        let mut hasher: Standard<Sha256> = Standard::new();
+        let mut mmr = Mmr::new(&mut hasher);
+
+        // 252 leaves. Leaf 240 sits in a peak preceded by 4 prefix peaks.
+        let elements: Vec<Digest> = (0..252u16)
+            .map(|i| Sha256::hash(&i.to_be_bytes()))
+            .collect();
+        let changeset = {
+            let mut batch = mmr.new_batch();
+            for e in &elements {
+                batch.add(&mut hasher, e);
+            }
+            batch.merkleize(&mut hasher).finalize()
+        };
+        mmr.apply(changeset).unwrap();
+        let root = mmr.root();
+
+        let loc = Location::new(240);
+        let proof = mmr.proof(&mut hasher, loc).unwrap();
+        assert!(proof.verify_element_inclusion(&mut hasher, &elements[240], loc, root));
+
+        // Tamper with the leaves field (249 has the same peak layout for leaf 240).
+        let mut tampered = proof.clone();
+        tampered.leaves = Location::new(249);
+        assert_ne!(tampered, proof);
+        assert!(
+            !tampered.verify_element_inclusion(&mut hasher, &elements[240], loc, root),
+            "proof with tampered leaves field must not verify"
+        );
+    }
+
     #[cfg(feature = "arbitrary")]
     mod conformance {
         use super::*;