feat: add canopy size check

programs/bubblegum/program/src/lib.rs

@@ -104,10 +104,9 @@ pub mod bubblegum {
         ctx: Context<'_, '_, '_, 'info, PrepareTree<'info>>,
         max_depth: u32,
         max_buffer_size: u32,
-        num_minted: u64,
         public: Option<bool>,
     ) -> Result<()> {
-        processor::prepare_tree(ctx, max_depth, max_buffer_size, num_minted, public)
+        processor::prepare_tree(ctx, max_depth, max_buffer_size, public)
     }
 
     pub fn add_canopy<'info>(
@@ -119,8 +118,8 @@ pub mod bubblegum {
     }
 
     /// Creates a new tree.
-    pub fn create_tree(
-        ctx: Context<CreateTree>,
+    pub fn create_tree<'info>(
+        ctx: Context<'_, '_, '_, 'info, CreateTree<'info>>,
         max_depth: u32,
         max_buffer_size: u32,
         public: Option<bool>,

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

@@ -1,21 +1,13 @@
-use std::mem::size_of;
-
 use anchor_lang::{prelude::*, system_program::System};
 use spl_account_compression::{
     program::SplAccountCompression,
-    state::{
-        merkle_tree_get_size, ConcurrentMerkleTreeHeader, CONCURRENT_MERKLE_TREE_HEADER_SIZE_V1,
-    },
-    Node, Noop,
+    Noop,
 };
 
 use crate::{
-    error::BubblegumError,
-    state::{DecompressibleState, TreeConfig, TREE_AUTHORITY_SIZE},
+    state::{DecompressibleState, TreeConfig, TREE_AUTHORITY_SIZE}, utils::check_canopy_size,
 };
 
-pub const MAX_ACC_PROOFS_SIZE: u32 = 17;
-
 #[derive(Accounts)]
 pub struct CreateTree<'info> {
     #[account(
@@ -37,15 +29,15 @@ pub struct CreateTree<'info> {
     pub system_program: Program<'info, System>,
 }
 
-pub(crate) fn create_tree(
-    ctx: Context<CreateTree>,
+pub(crate) fn create_tree<'info>(
+    ctx: Context<'_, '_, '_, 'info, CreateTree<'info>>,
     max_depth: u32,
     max_buffer_size: u32,
     public: Option<bool>,
 ) -> Result<()> {
     let merkle_tree = ctx.accounts.merkle_tree.to_account_info();
 
-    check_canopy_size(&ctx, max_depth, max_buffer_size)?;
+    check_canopy_size(ctx.accounts.merkle_tree.to_account_info(), ctx.accounts.tree_authority.to_account_info(), max_depth, max_buffer_size)?;
 
     let seed = merkle_tree.key();
     let seeds = &[seed.as_ref(), &[ctx.bumps.tree_authority]];
@@ -70,36 +62,3 @@ pub(crate) fn create_tree(
     );
     spl_account_compression::cpi::init_empty_merkle_tree(cpi_ctx, max_depth, max_buffer_size)
 }
-
-fn check_canopy_size(
-    ctx: &Context<CreateTree>,
-    max_depth: u32,
-    max_buffer_size: u32,
-) -> Result<()> {
-    let merkle_tree_bytes = ctx.accounts.merkle_tree.data.borrow();
-
-    let (header_bytes, rest) = merkle_tree_bytes.split_at(CONCURRENT_MERKLE_TREE_HEADER_SIZE_V1);
-
-    let mut header = ConcurrentMerkleTreeHeader::try_from_slice(header_bytes)?;
-    header.initialize(
-        max_depth,
-        max_buffer_size,
-        &ctx.accounts.tree_authority.key(),
-        Clock::get()?.slot,
-    );
-
-    let merkle_tree_size = merkle_tree_get_size(&header)?;
-
-    let (_tree_bytes, canopy_bytes) = rest.split_at(merkle_tree_size);
-
-    let required_canopy = max_depth.saturating_sub(MAX_ACC_PROOFS_SIZE);
-
-    let actual_canopy_size = canopy_bytes.len() / size_of::<Node>();
-
-    require!(
-        (actual_canopy_size as u32) >= required_canopy,
-        BubblegumError::InvalidCanopySize
-    );
-
-    Ok(())
-}

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

@@ -9,7 +9,7 @@ use crate::{
     state::{
         DecompressibleState, TreeConfig, MINIMUM_STAKE, REALM, REALM_GOVERNING_MINT,
         TREE_AUTHORITY_SIZE,
-    },
+    }, utils::check_canopy_size,
 };
 
 #[derive(Accounts)]
@@ -42,7 +42,6 @@ pub(crate) fn prepare_tree<'info>(
     ctx: Context<'_, '_, '_, 'info, PrepareTree<'info>>,
     max_depth: u32,
     max_buffer_size: u32,
-    num_minted: u64,
     public: Option<bool>,
 ) -> Result<()> {
     assert_eq!(ctx.accounts.registrar.owner, &mplx_staking_states::ID);
@@ -108,7 +107,7 @@ pub(crate) fn prepare_tree<'info>(
         return Err(BubblegumError::NotEnoughStakeForOperation.into());
     }
 
-    // TODO: check canopy size
+    check_canopy_size(ctx.accounts.merkle_tree.to_account_info(), ctx.accounts.tree_authority.to_account_info(), max_depth, max_buffer_size)?;
 
     let merkle_tree = ctx.accounts.merkle_tree.to_account_info();
     let seed = merkle_tree.key();

programs/bubblegum/program/src/state/mod.rs

@@ -12,6 +12,8 @@ pub const VOUCHER_PREFIX: &str = "voucher";
 pub const ASSET_PREFIX: &str = "asset";
 pub const COLLECTION_CPI_PREFIX: &str = "collection_cpi";
 
+pub const MAX_ACC_PROOFS_SIZE: u32 = 17;
+
 // TODO: set real keys before mainnet deploy
 pub const REALM: Pubkey = solana_program::pubkey!("EzsKaQq68FLZwRaiUx7t17LWVVzsE8wRkhBghFrZGGwG");
 pub const REALM_GOVERNING_MINT: Pubkey =

programs/bubblegum/program/src/utils.rs

@@ -1,13 +1,15 @@
-use crate::state::{
+use std::mem::size_of;
+
+use crate::{error::BubblegumError, state::{
     metaplex_adapter::{Creator, MetadataArgs},
-    ASSET_PREFIX,
-};
+    ASSET_PREFIX, MAX_ACC_PROOFS_SIZE,
+}};
 use anchor_lang::{
     prelude::*,
     solana_program::{program_memory::sol_memcmp, pubkey::PUBKEY_BYTES},
 };
 use solana_program::keccak;
-use spl_account_compression::Node;
+use spl_account_compression::{state::{merkle_tree_get_size, ConcurrentMerkleTreeHeader, CONCURRENT_MERKLE_TREE_HEADER_SIZE_V1}, Node};
 
 pub fn hash_creators(creators: &[Creator]) -> Result<[u8; 32]> {
     // Convert creator Vec to bytes Vec.
@@ -106,3 +108,37 @@ pub fn get_asset_id(tree_id: &Pubkey, nonce: u64) -> Pubkey {
     )
     .0
 }
+
+pub(crate) fn check_canopy_size<'info>(
+    merkle_tree: AccountInfo<'info>,
+    tree_authority: AccountInfo<'info>,
+    max_depth: u32,
+    max_buffer_size: u32,
+) -> Result<()> {
+    let merkle_tree_bytes = merkle_tree.data.borrow();
+
+    let (header_bytes, rest) = merkle_tree_bytes.split_at(CONCURRENT_MERKLE_TREE_HEADER_SIZE_V1);
+
+    let mut header = ConcurrentMerkleTreeHeader::try_from_slice(header_bytes)?;
+    header.initialize(
+        max_depth,
+        max_buffer_size,
+        &tree_authority.key(),
+        Clock::get()?.slot,
+    );
+
+    let merkle_tree_size = merkle_tree_get_size(&header)?;
+
+    let (_tree_bytes, canopy_bytes) = rest.split_at(merkle_tree_size);
+
+    let required_canopy = max_depth.saturating_sub(MAX_ACC_PROOFS_SIZE);
+
+    let actual_canopy_size = canopy_bytes.len() / size_of::<Node>();
+
+    require!(
+        (actual_canopy_size as u32) >= required_canopy,
+        BubblegumError::InvalidCanopySize
+    );
+
+    Ok(())
+}

programs/bubblegum/program/tests/rollup.rs

@@ -178,7 +178,6 @@ async fn test_prepare_tree_without_canopy() {
         false,
         MAX_DEPTH as u32,
         MAX_BUF_SIZE as u32,
-        1000,
         registrar_key,
         voter_key,
     );
@@ -362,7 +361,6 @@ async fn test_prepare_tree_with_canopy() {
         false,
         MAX_DEPTH as u32,
         MAX_BUF_SIZE as u32,
-        1000,
         registrar_key,
         voter_key,
     );
@@ -563,7 +561,6 @@ async fn test_put_wrong_canopy() {
         false,
         MAX_DEPTH as u32,
         MAX_BUF_SIZE as u32,
-        1000,
         registrar_key,
         voter_key,
     );
@@ -622,3 +619,167 @@ async fn test_put_wrong_canopy() {
         panic!("Should have failed");
     }
 }
+
+#[tokio::test]
+async fn test_prepare_with_small_canopy() {
+    let mut program_context = BubblegumTestContext::new().await.unwrap();
+
+    let merkle_tree = MerkleTree::new(vec![Node::default(); 1 << 20].as_slice());
+
+    let governance_program_id =
+        Pubkey::from_str("CuyWCRdHT8pZLG793UR5R9z31AC49d47ZW9ggN6P7qZ4").unwrap();
+    let realm_authority = Pubkey::from_str("Euec5oQGN3Y9kqVrz6PQRfTpYSn6jK3k1JonDiMTzAtA").unwrap();
+    let voter_authority = program_context.test_context().payer.pubkey();
+
+    let mplx_mint_key = Pubkey::new_unique();
+    let grant_authority = Pubkey::new_unique();
+
+    let registrar_key = Pubkey::find_program_address(
+        &[
+            REALM.to_bytes().as_ref(),
+            b"registrar".as_ref(),
+            REALM_GOVERNING_MINT.to_bytes().as_ref(),
+        ],
+        &mplx_staking_states::ID,
+    )
+    .0;
+
+    let (voter_key, voter_bump) = Pubkey::find_program_address(
+        &[
+            registrar_key.to_bytes().as_ref(),
+            b"voter".as_ref(),
+            voter_authority.to_bytes().as_ref(),
+        ],
+        &mplx_staking_states::ID,
+    );
+
+    // // init structs for Registrar and Voter and fill it in with data
+    let voting_mint_config = VotingMintConfig {
+        mint: mplx_mint_key,
+        grant_authority,
+        baseline_vote_weight_scaled_factor: 0,
+        max_extra_lockup_vote_weight_scaled_factor: 0,
+        lockup_saturation_secs: 0,
+        digit_shift: 0,
+        padding: [0, 0, 0, 0, 0, 0, 0],
+    };
+
+    let registrar = Registrar {
+        governance_program_id,
+        realm: REALM,
+        realm_governing_token_mint: REALM_GOVERNING_MINT,
+        realm_authority,
+        voting_mints: [
+            voting_mint_config,
+            voting_mint_config,
+            voting_mint_config,
+            voting_mint_config,
+        ],
+        time_offset: 0,
+        bump: 0,
+    };
+
+    let current_time = SystemTime::now()
+        .duration_since(UNIX_EPOCH)
+        .unwrap()
+        .as_millis() as u64;
+
+    let lockup = Lockup {
+        start_ts: 0,
+        end_ts: current_time + 100,
+        cooldown_ends_at: 0,
+        cooldown_requested: false,
+        kind: LockupKind::Constant,
+        period: LockupPeriod::ThreeMonths,
+        _reserved1: [0; 5],
+    };
+
+    let deposit_entry = DepositEntry {
+        lockup: lockup.clone(),
+        amount_deposited_native: 100000000,
+        voting_mint_config_idx: 0,
+        is_used: true,
+        _reserved1: [0; 6],
+    };
+
+    let deposit_entries = [deposit_entry; 32];
+
+    let voter = Voter {
+        deposits: deposit_entries,
+        voter_authority,
+        registrar: registrar_key,
+        voter_bump,
+        voter_weight_record_bump: 0,
+        _reserved1: [0; 14],
+    };
+
+    let registrar_acc_data = [
+        REGISTRAR_DISCRIMINATOR.as_ref(),
+        bytemuck::bytes_of(&registrar),
+    ]
+    .concat();
+    let voter_acc_data = [VOTER_DISCRIMINATOR.as_ref(), bytemuck::bytes_of(&voter)].concat();
+
+    // for next two accounts set arbitrary balance because it doesn't meter for test
+    let mut registrar_account = AccountSharedData::new(
+        10000000000000000,
+        registrar_acc_data.len(),
+        &mplx_staking_states::ID,
+    );
+    registrar_account.set_data_from_slice(registrar_acc_data.as_ref());
+
+    let mut voter_account = AccountSharedData::new(
+        10000000000000000,
+        voter_acc_data.len(),
+        &mplx_staking_states::ID,
+    );
+    voter_account.set_data_from_slice(voter_acc_data.as_ref());
+
+    program_context
+        .mut_test_context()
+        .set_account(&registrar_key, &registrar_account);
+    program_context
+        .mut_test_context()
+        .set_account(&voter_key, &voter_account);
+
+    let tree_creator = Keypair::from_bytes(TREE_CREATOR.as_ref()).unwrap();
+
+    let tree_key = Keypair::from_bytes(TREE_KEY.as_ref()).unwrap();
+
+    let mut tree = Tree::<20, 64>::with_preinitialised_tree(
+        &tree_creator,
+        &tree_key,
+        program_context.client(),
+        merkle_tree,
+        1000,
+        0,
+    );
+
+    tree.alloc(&program_context.test_context().payer)
+        .await
+        .unwrap();
+
+    let mut tree_tx_builder = tree.prepare_tree_tx(
+        &program_context.test_context().payer,
+        false,
+        20,
+        64,
+        registrar_key,
+        voter_key,
+    );
+
+    let res = tree_tx_builder.execute_without_root_check().await;
+
+    if let Err(err) = res {
+        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");
+    }
+}