Fix canopy size check (#101)

* fix: canopy size check

* chore: change tests

programs/bubblegum/program/src/processor/create_tree.rs

@@ -1,4 +1,4 @@
-use std::mem::size_of;
+use bytemuck::cast_slice;
 
 use anchor_lang::{prelude::*, system_program::System};
 use spl_account_compression::{
@@ -92,14 +92,46 @@ fn check_canopy_size(
 
     let (_tree_bytes, canopy_bytes) = rest.split_at(merkle_tree_size);
 
-    let required_canopy = max_depth.saturating_sub(MAX_ACC_PROOFS_SIZE);
+    let canopy = cast_slice::<u8, Node>(canopy_bytes);
+
+    let cached_path_len = get_cached_path_length(canopy, max_depth)?;
 
-    let actual_canopy_size = canopy_bytes.len() / size_of::<Node>();
+    let required_canopy = max_depth.saturating_sub(MAX_ACC_PROOFS_SIZE);
 
     require!(
-        (actual_canopy_size as u32) >= required_canopy,
+        (cached_path_len as u32) >= required_canopy,
         BubblegumError::InvalidCanopySize
     );
 
     Ok(())
 }
+
+// Method is taken from account-compression Solana program
+#[inline(always)]
+fn get_cached_path_length(canopy: &[Node], max_depth: u32) -> Result<u32> {
+    // The offset of 2 is applied because the canopy is a full binary tree without the root node
+    // Size: (2^n - 2) -> Size + 2 must be a power of 2
+    let closest_power_of_2 = (canopy.len() + 2) as u32;
+    // This expression will return true if `closest_power_of_2` is actually a power of 2
+    if closest_power_of_2 & (closest_power_of_2 - 1) == 0 {
+        // (1 << max_depth) returns the number of leaves in the full merkle tree
+        // (1 << (max_depth + 1)) - 1 returns the number of nodes in the full tree
+        // The canopy size cannot exceed the size of the tree
+        if closest_power_of_2 > (1 << (max_depth + 1)) {
+            msg!(
+                "Canopy size is too large. Size: {}. Max size: {}",
+                closest_power_of_2 - 2,
+                (1 << (max_depth + 1)) - 2
+            );
+            return err!(BubblegumError::InvalidCanopySize);
+        }
+    } else {
+        msg!(
+            "Canopy length {} is not 2 less than a power of 2",
+            canopy.len()
+        );
+        return err!(BubblegumError::InvalidCanopySize);
+    }
+    // 1 is subtracted from the trailing zeros because the root is not stored in the canopy
+    Ok(closest_power_of_2.trailing_zeros() - 1)
+}

programs/bubblegum/program/tests/simple.rs

@@ -250,6 +250,31 @@ async fn test_create_public_tree_with_canopy() {
     assert_eq!(cfg.tree_creator, payer.pubkey());
     assert_eq!(cfg.tree_delegate, payer.pubkey());
     assert!(cfg.is_public);
+}
+
+#[tokio::test]
+async fn test_cannot_create_tree_needing_too_many_proofs_with_too_small_canopy() {
+    let context = BubblegumTestContext::new().await.unwrap();
+
+    let tree_create_result = context.create_tree_with_canopy::<19, 64>(1, true).await;
+
+    if let Err(err) = tree_create_result {
+        if let BanksClient(BanksClientError::TransactionError(e)) = *err {
+            assert_eq!(
+                e,
+                TransactionError::InstructionError(0, InstructionError::Custom(6041),)
+            );
+        } else {
+            panic!("Wrong variant");
+        }
+    } else {
+        panic!("Should have failed");
+    }
+}
+
+#[tokio::test]
+async fn test_cannot_create_tree_needing_too_many_proofs_with_no_canopy() {
+    let context = BubblegumTestContext::new().await.unwrap();
 
     let tree_create_result = context.create_tree_with_canopy::<18, 64>(0, true).await;