Merge pull request #18 from HerodotusDev/audit

Audit

src/data_structures/mmr/mmr.cairo

@@ -1,6 +1,8 @@
 use cairo_lib::data_structures::mmr::peaks::{Peaks, PeaksTrait};
 use cairo_lib::data_structures::mmr::proof::{Proof, ProofTrait};
-use cairo_lib::data_structures::mmr::utils::{compute_root, get_height};
+use cairo_lib::data_structures::mmr::utils::{
+    compute_root, get_height, mmr_size_to_leaf_count, leaf_count_to_peaks_count, get_peak_info
+};
 use cairo_lib::hashing::poseidon::PoseidonHasher;
 
 // @notice Merkle Mountatin Range struct
@@ -34,6 +36,10 @@ impl MMRImpl of MMRTrait {
     // @param peaks The peaks of the MMR
     // @return Result with the new root and new peaks of the MMR
     fn append(ref self: MMR, hash: felt252, peaks: Peaks) -> Result<(felt252, Peaks), felt252> {
+        let leaf_count = mmr_size_to_leaf_count(self.last_pos.into());
+        if leaf_count_to_peaks_count(leaf_count) != peaks.len().into() {
+            return Result::Err('Invalid peaks count');
+        }
         if !peaks.valid(self.last_pos, self.root) {
             return Result::Err('Invalid peaks');
         }
@@ -99,11 +105,22 @@ impl MMRImpl of MMRTrait {
     fn verify_proof(
         self: @MMR, index: usize, hash: felt252, peaks: Peaks, proof: Proof
     ) -> Result<bool, felt252> {
+        let leaf_count = mmr_size_to_leaf_count((*self.last_pos).into());
+        if leaf_count_to_peaks_count(leaf_count) != peaks.len().into() {
+            return Result::Err('Invalid peaks count');
+        }
         if !peaks.valid(*self.last_pos, *self.root) {
             return Result::Err('Invalid peaks');
         }
 
+        let (peak_index, peak_height) = get_peak_info(*self.last_pos, index);
+
+        if proof.len() != peak_height {
+            return Result::Ok(false);
+        }
+
         let peak = proof.compute_peak(index, hash);
-        Result::Ok(peaks.contains_peak(peak))
+
+        Result::Ok(*peaks.at(peak_index) == peak)
     }
 }

src/data_structures/mmr/tests/test_mmr.cairo

@@ -226,3 +226,26 @@ fn test_verify_proof_invalid_peaks() {
 
     assert(mmr.verify_proof(2, *elems.at(1), peaks, proof).is_err(), 'Proof wrong peaks')
 }
+
+#[test]
+#[available_gas(99999999)]
+fn test_attack_forge_peaks() {
+    let elems = helper_test_get_elements();
+    let mut mmr_real: MMR = MMRTrait::new(
+        0x21aea73dea77022a4882e1f656b76c9195161ed1cff2b065a74d7246b02d5d6, 0x8
+    );
+    let mut mmr_fake: MMR = MMRTrait::new(
+        0x21aea73dea77022a4882e1f656b76c9195161ed1cff2b065a74d7246b02d5d6, 0x8
+    );
+
+    // add the next element normally to mmr_real and get the root;
+    let peaks_real = array![*elems.at(6), *elems.at(7)].span();
+    mmr_real.append(9, peaks_real);
+
+    // add the next element abnormally to mmr_real and get the root;
+    let forged_peak = PoseidonHasher::hash_double(*elems.at(6), *elems.at(7));
+    let peaks_fake = array![forged_peak].span();
+    let res = mmr_fake.append(9, peaks_fake);
+
+    assert(res.is_err(), 'attack success: forged peak');
+}

src/data_structures/mmr/tests/test_utils.cairo

@@ -1,5 +1,6 @@
 use cairo_lib::data_structures::mmr::utils::{
-    get_height, compute_root, count_ones, leaf_index_to_mmr_index
+    get_height, compute_root, count_ones, leaf_index_to_mmr_index, get_peak_info,
+    mmr_size_to_leaf_count,
 };
 use cairo_lib::hashing::poseidon::PoseidonHasher;
 use cairo_lib::data_structures::mmr::peaks::PeaksTrait;
@@ -72,3 +73,27 @@ fn test_leaf_index_to_mmr_index() {
     assert(leaf_index_to_mmr_index(10) == 17, 'leaf_..._index(10) != 17');
     assert(leaf_index_to_mmr_index(11) == 19, 'leaf_..._index(11) != 19');
 }
+
+#[test]
+#[available_gas(999999999)]
+fn test_mmr_size_to_leaf_count() {
+    assert(mmr_size_to_leaf_count(1) == 1, 'mmr_size_to_leaf_count(1) != 1');
+    assert(mmr_size_to_leaf_count(3) == 2, 'mmr_size_to_leaf_count(3) != 2');
+    assert(mmr_size_to_leaf_count(4) == 3, 'mmr_size_to_leaf_count(4) != 3');
+    assert(mmr_size_to_leaf_count(7) == 4, 'mmr_size_to_leaf_count(7) != 4');
+    assert(mmr_size_to_leaf_count(8) == 5, 'mmr_size_to_leaf_count(8) != 5');
+    assert(mmr_size_to_leaf_count(10) == 6, 'mmr_size_to_leaf_count(10) != 6');
+    assert(mmr_size_to_leaf_count(11) == 7, 'mmr_size_to_leaf_count(11) != 7');
+    assert(mmr_size_to_leaf_count(15) == 8, 'mmr_size_to_leaf_count(15) != 8');
+}
+
+#[test]
+#[available_gas(999999999)]
+fn test_get_peak_info() {
+    assert(get_peak_info(11, 11) == (2, 0), 'get_peak_info(11, 11) != (2, 0)');
+    assert(get_peak_info(15, 11) == (0, 3), 'get_peak_info(15, 11) != (0, 3)');
+    assert(get_peak_info(18, 16) == (1, 1), 'get_peak_info(18, 16) != (1, 1)');
+    assert(get_peak_info(26, 16) == (1, 2), 'get_peak_info(26, 16) != (1, 2)');
+    assert(get_peak_info(26, 16) == (1, 2), 'get_peak_info(26, 16) != (1, 2)');
+    assert(get_peak_info(31, 16) == (0, 4), 'get_peak_info(31, 16) != (0, 4)');
+}

src/data_structures/mmr/utils.cairo

@@ -28,10 +28,10 @@ fn compute_root(last_pos: felt252, peaks: Peaks) -> felt252 {
 }
 
 // @notice Count the number of bits set to 1 in a 256-bit unsigned integer
-// @param n The 256-bit unsigned integer
+// @param arg The 256-bit unsigned integer
 // @return The number of bits set to 1 in n
-fn count_ones(n: u256) -> u256 {
-    let mut n = n;
+fn count_ones(arg: u256) -> u256 {
+    let mut n = arg;
     let mut count = 0;
     loop {
         if n == 0 {
@@ -42,9 +42,62 @@ fn count_ones(n: u256) -> u256 {
     }
 }
 
-// @notice Convert a leaf index to an Merkle Mountain Range tree leaf index
+// @notice Convert a leaf index to an Merkle Mountain Range tree index
 // @param n The leaf index
 // @return The MMR index
 fn leaf_index_to_mmr_index(n: u256) -> u256 {
     2 * n - 1 - count_ones(n - 1)
 }
+
+// @notice Convert a Merkle Mountain Range tree size to number of leaves
+// @param n MMR size
+// @result Number of leaves
+fn mmr_size_to_leaf_count(n: u256) -> u256 {
+    let mut mmr_size = n;
+    let bits = bit_length(mmr_size + 1);
+    let mut mountain_leaf_count = pow(2, bits - 1);
+    let mut leaf_count = 0;
+    loop {
+        if mountain_leaf_count == 0 {
+            break leaf_count;
+        }
+        let mountain_size = 2 * mountain_leaf_count - 1;
+        if mountain_size <= mmr_size {
+            mmr_size -= mountain_size;
+            leaf_count += mountain_leaf_count;
+        }
+        mountain_leaf_count /= 2;
+    }
+}
+
+// @notice Convert a number of leaves to number of peaks
+// @param leaf_count Number of leaves
+// @return Number of peaks
+fn leaf_count_to_peaks_count(leaf_count: u256) -> u256 {
+    count_ones(leaf_count)
+}
+
+// @notice Get peak size and index of the peak the element is in
+// @param elements_count The size of the MMR (number of elements in the MMR)
+// @param element_index The index of the element in the MMR
+// @return (peak index, peak height)
+fn get_peak_info(elements_count: u32, element_index: u32) -> (u32, u32) {
+    let mut elements_count = elements_count;
+    let mut element_index = element_index;
+
+    let mut mountain_height = bit_length(elements_count);
+    let mut mountain_elements_count = pow(2, mountain_height) - 1;
+    let mut mountain_index = 0;
+    loop {
+        if mountain_elements_count <= elements_count {
+            if element_index <= mountain_elements_count {
+                break (mountain_index, mountain_height - 1);
+            }
+            elements_count -= mountain_elements_count;
+            element_index -= mountain_elements_count;
+            mountain_index += 1;
+        }
+        mountain_height -= 1;
+        mountain_elements_count /= 2;
+    }
+}