chore: add address tree height error (Lightprotocol#1590)

reilabs · Feb 15, 2025 · b9eebe4 · b9eebe4
1 parent c0b63e8
commit b9eebe4
Show file tree

Hide file tree

Showing 5 changed files with 17 additions and 8 deletions.
diff --git a/program-libs/batched-merkle-tree/src/merkle_tree.rs b/program-libs/batched-merkle-tree/src/merkle_tree.rs
@@ -287,6 +287,11 @@ impl<'a> BatchedMerkleTreeAccount<'a> {
             // Root for binary Merkle tree with all zero leaves.
             root_history.push(light_hasher::Poseidon::zero_bytes()[height as usize]);
         } else if tree_type == TreeType::BatchedAddress {
+            // Sanity check since init value is hardcoded.
+            #[cfg(not(test))]
+            if height != 40 {
+                return Err(MerkleTreeMetadataError::InvalidHeight.into());
+            }
             // Initialized indexed Merkle tree root.
             // See https://github.com/Lightprotocol/light-protocol/blob/c143c24f95c901e2eac96bc2bd498719958192cf/program-libs/indexed-merkle-tree/src/reference.rs#L69
             root_history.push(ADDRESS_TREE_INIT_ROOT_40);
@@ -1130,7 +1135,7 @@ mod test {
             root_history_len,
             batch_size,
             zkp_batch_size,
-            10,
+            40,
             num_iter,
             bloom_filter_capacity,
             TreeType::BatchedAddress,
@@ -1569,7 +1574,7 @@ mod test {
         let num_iter = 1;
         let mut current_slot = 1;
         let bloom_filter_capacity = 8000;
-        let height = 4;
+        let height = 40;
         let mut account = BatchedMerkleTreeAccount::init(
             &mut account_data,
             &Pubkey::new_unique(),
@@ -1703,7 +1708,7 @@ mod test {
         let root_history_len = 10;
         let num_iter = 1;
         let bloom_filter_capacity = 8000;
-        let height = 4;
+        let height = 40;
         let pubkey = Pubkey::new_unique();
         let mut account = BatchedMerkleTreeAccount::init(
             &mut account_data,
@@ -1741,7 +1746,7 @@ mod test {
         let root_history_len = 10;
         let num_iter = 1;
         let bloom_filter_capacity = 8000;
-        let height = 4;
+        let height = 40;
         let pubkey = Pubkey::new_unique();
         let associated_queue = Pubkey::new_unique();
         let account = BatchedMerkleTreeAccount::init(

diff --git a/program-libs/batched-merkle-tree/tests/initialize_address_tree.rs b/program-libs/batched-merkle-tree/tests/initialize_address_tree.rs
@@ -69,7 +69,7 @@ fn test_rnd_account_init() {
             rollover_threshold: Some(rng.gen_range(0..100)),
             close_threshold: None,
             root_history_capacity: rng.gen_range(1..1000),
-            height: rng.gen_range(1..32),
+            height: 40,
         };
 
         let mt_account_size = get_merkle_tree_account_size(

diff --git a/program-libs/batched-merkle-tree/tests/rollover_address_tree.rs b/program-libs/batched-merkle-tree/tests/rollover_address_tree.rs
@@ -187,7 +187,7 @@ fn test_rnd_rollover() {
             rollover_threshold: Some(rng.gen_range(0..100)),
             close_threshold: None,
             root_history_capacity: rng.gen_range(1..1000),
-            height: rng.gen_range(1..32),
+            height: 40,
         };
         if forester.is_some() {
             params.network_fee = None;
@@ -307,7 +307,7 @@ fn test_batched_tree_is_ready_for_rollover() {
         height,
         num_iter,
         bloom_filter_capacity,
-        TreeType::BatchedAddress,
+        TreeType::BatchedState,
     )
     .unwrap();
 
@@ -318,7 +318,7 @@ fn test_batched_tree_is_ready_for_rollover() {
     );
 
     let tree_capacity = 2u64.pow(height);
-    let start_index = 2;
+    let start_index = 0;
     let rollover_threshold =
         tree_capacity * metadata.rollover_metadata.rollover_threshold / 100 - start_index;
     // fill tree almost to the rollover threshold

diff --git a/program-libs/merkle-tree-metadata/src/errors.rs b/program-libs/merkle-tree-metadata/src/errors.rs
@@ -18,6 +18,8 @@ pub enum MerkleTreeMetadataError {
     InvalidTreeType,
     #[error("Invalid Rollover Threshold.")]
     InvalidRolloverThreshold,
+    #[error("Invalid Height.")]
+    InvalidHeight,
 }
 
 #[cfg(feature = "solana")]
@@ -32,6 +34,7 @@ impl From<MerkleTreeMetadataError> for u32 {
             MerkleTreeMetadataError::NotReadyForRollover => 14006,
             MerkleTreeMetadataError::InvalidTreeType => 14007,
             MerkleTreeMetadataError::InvalidRolloverThreshold => 14008,
+            MerkleTreeMetadataError::InvalidHeight => 14009,
         }
     }
 }

diff --git a/programs/system/src/processor/cpi.rs b/programs/system/src/processor/cpi.rs
@@ -39,6 +39,7 @@ pub fn create_cpi_data_and_context<
     ];
     let account_indices =
         Vec::<u8>::with_capacity((num_nullifiers + num_leaves + num_new_addresses) as usize);
+    // Min (remaining accounts or num values) for there cannot be more trees than accounts or values.
     let bytes_size = InsertIntoQueuesInstructionDataMut::required_size_for_capacity(
         num_leaves,
         num_nullifiers,