From 971b38d8816ccee4321b9d2af5d02d6eb44b99fe Mon Sep 17 00:00:00 2001 From: ananas-block <58553958+ananas-block@users.noreply.github.com> Date: Wed, 4 Dec 2024 23:59:05 +0000 Subject: [PATCH] feat: rollover state trees (#1371) * feat: rollover batch state trees * test: batch state tree rollover --- programs/account-compression/src/errors.rs | 2 + .../intialize_batched_state_merkle_tree.rs | 11 +- .../src/instructions/mod.rs | 3 + .../rollover_address_merkle_tree_and_queue.rs | 2 +- .../rollover_batch_state_merkle_tree.rs | 971 ++++++++++++++++++ programs/account-compression/src/lib.rs | 8 + .../src/state/batched_merkle_tree.rs | 24 + .../src/state/batched_queue.rs | 17 + .../account-compression/src/state/rollover.rs | 4 + programs/account-compression/src/utils/mod.rs | 19 + .../initialize_batched_state_tree.rs | 4 +- .../initialize_tree_and_queue.rs | 5 + .../src/account_compression_cpi/mod.rs | 1 + .../rollover_batch_state_tree.rs | 71 ++ .../src/account_compression_cpi/sdk.rs | 47 +- programs/registry/src/lib.rs | 59 +- .../tests/batched_merkle_tree_test.rs | 495 +++++++-- test-programs/registry-test/tests/tests.rs | 208 +++- test-utils/src/test_batch_forester.rs | 201 +++- 19 files changed, 2051 insertions(+), 101 deletions(-) create mode 100644 programs/account-compression/src/instructions/rollover_batch_state_merkle_tree.rs create mode 100644 programs/registry/src/account_compression_cpi/rollover_batch_state_tree.rs diff --git a/programs/account-compression/src/errors.rs b/programs/account-compression/src/errors.rs index f70a8b4fae..45a2077470 100644 --- a/programs/account-compression/src/errors.rs +++ b/programs/account-compression/src/errors.rs @@ -76,4 +76,6 @@ pub enum AccountCompressionErrorCode { InvalidBatch, LeafIndexNotInBatch, UnsupportedParameters, + InvalidTreeType, + InvalidNetworkFee, } diff --git a/programs/account-compression/src/instructions/intialize_batched_state_merkle_tree.rs b/programs/account-compression/src/instructions/intialize_batched_state_merkle_tree.rs index f5c2d1661a..fd44e2af34 100644 --- a/programs/account-compression/src/instructions/intialize_batched_state_merkle_tree.rs +++ b/programs/account-compression/src/instructions/intialize_batched_state_merkle_tree.rs @@ -219,11 +219,13 @@ pub fn bytes_to_struct_checked BatchedQueueAccount { let rollover_fee = match rollover_threshold { Some(rollover_threshold) => compute_rollover_fee(rollover_threshold, height, rent) @@ -441,7 +441,7 @@ pub fn get_output_queue_account_default( rolledover_slot: u64::MAX, rollover_threshold: rollover_threshold.unwrap_or(u64::MAX), rollover_fee, - network_fee: 5000, + network_fee, additional_bytes, }, queue_type: QueueType::Output as u64, @@ -566,6 +566,7 @@ pub mod tests { mt_pubkey, params.height, params.output_queue_num_batches, + params.network_fee.unwrap_or_default(), ); assert_queue_zero_copy_inited( output_queue_account_data.as_mut_slice(), diff --git a/programs/account-compression/src/instructions/mod.rs b/programs/account-compression/src/instructions/mod.rs index 10534c9a2c..f195aeada8 100644 --- a/programs/account-compression/src/instructions/mod.rs +++ b/programs/account-compression/src/instructions/mod.rs @@ -42,3 +42,6 @@ pub use batch_nullify::*; pub mod batch_append; pub use batch_append::*; + +pub mod rollover_batch_state_merkle_tree; +pub use rollover_batch_state_merkle_tree::*; diff --git a/programs/account-compression/src/instructions/rollover_address_merkle_tree_and_queue.rs b/programs/account-compression/src/instructions/rollover_address_merkle_tree_and_queue.rs index 49b0a5fcb3..df98931687 100644 --- a/programs/account-compression/src/instructions/rollover_address_merkle_tree_and_queue.rs +++ b/programs/account-compression/src/instructions/rollover_address_merkle_tree_and_queue.rs @@ -48,7 +48,7 @@ impl<'info> GroupAccounts<'info> for RolloverAddressMerkleTreeAndQueue<'info> { /// /// Actions: /// 1. mark Merkle tree as rolled over in this slot -/// 2. initialize new Merkle tree and nullifier queue with the same parameters +/// 2. initialize new Merkle tree and queue with the same parameters pub fn process_rollover_address_merkle_tree_and_queue<'a, 'b, 'c: 'info, 'info>( ctx: Context<'a, 'b, 'c, 'info, RolloverAddressMerkleTreeAndQueue<'info>>, ) -> Result<()> { diff --git a/programs/account-compression/src/instructions/rollover_batch_state_merkle_tree.rs b/programs/account-compression/src/instructions/rollover_batch_state_merkle_tree.rs new file mode 100644 index 0000000000..8b53a94293 --- /dev/null +++ b/programs/account-compression/src/instructions/rollover_batch_state_merkle_tree.rs @@ -0,0 +1,971 @@ +use crate::{ + batched_merkle_tree::{BatchedMerkleTreeAccount, ZeroCopyBatchedMerkleTreeAccount}, + batched_queue::{ + assert_queue_zero_copy_inited, BatchedQueueAccount, ZeroCopyBatchedQueueAccount, + }, + init_batched_state_merkle_tree_accounts, + utils::{ + check_account::check_account_balance_is_rent_exempt, + check_signer_is_registered_or_authority::{ + check_signer_is_registered_or_authority, GroupAccounts, + }, + if_equals_none, + transfer_lamports::transfer_lamports, + }, + InitStateTreeAccountsInstructionData, RegisteredProgram, +}; +use anchor_lang::{prelude::*, solana_program::pubkey::Pubkey}; + +use super::assert_mt_zero_copy_inited; + +#[derive(Accounts)] +pub struct RolloverBatchStateMerkleTree<'info> { + #[account(mut)] + /// Signer used to receive rollover accounts rentexemption reimbursement. + pub fee_payer: Signer<'info>, + pub authority: Signer<'info>, + pub registered_program_pda: Option>, + /// CHECK: is initialized in this instruction. + #[account(zero)] + pub new_state_merkle_tree: AccountInfo<'info>, + /// CHECK: checked in manual deserialization. + #[account(mut)] + pub old_state_merkle_tree: AccountInfo<'info>, + /// CHECK: is initialized in this instruction. + #[account(zero)] + pub new_output_queue: AccountInfo<'info>, + /// CHECK: checked in manual deserialization. + #[account(mut)] + pub old_output_queue: AccountInfo<'info>, +} + +impl<'info> GroupAccounts<'info> for RolloverBatchStateMerkleTree<'info> { + fn get_authority(&self) -> &Signer<'info> { + &self.authority + } + fn get_registered_program_pda(&self) -> &Option> { + &self.registered_program_pda + } +} + +/// Checks: +/// 1. Merkle tree is ready to be rolled over +/// 2. Merkle tree is not already rolled over +/// 3. Rollover threshold is configured, if not tree cannot be rolled over +/// +/// Actions: +/// 1. mark Merkle tree as rolled over in this slot +/// 2. initialize new Merkle tree and output queue with the same parameters +pub fn process_rollover_batch_state_merkle_tree<'a, 'b, 'c: 'info, 'info>( + ctx: Context<'a, 'b, 'c, 'info, RolloverBatchStateMerkleTree<'info>>, + additional_bytes: u64, + network_fee: Option, +) -> Result<()> { + let old_merkle_tree_account = + &mut ZeroCopyBatchedMerkleTreeAccount::state_tree_from_account_info_mut( + &ctx.accounts.old_state_merkle_tree, + )?; + let old_output_queue = &mut ZeroCopyBatchedQueueAccount::output_queue_from_account_info_mut( + &ctx.accounts.old_output_queue, + )?; + check_signer_is_registered_or_authority::< + RolloverBatchStateMerkleTree, + ZeroCopyBatchedMerkleTreeAccount, + >(&ctx, old_merkle_tree_account)?; + + let merkle_tree_rent = check_account_balance_is_rent_exempt( + &ctx.accounts.new_state_merkle_tree.to_account_info(), + ctx.accounts + .old_state_merkle_tree + .to_account_info() + .data_len(), + )?; + let queue_rent = check_account_balance_is_rent_exempt( + &ctx.accounts.new_output_queue.to_account_info(), + ctx.accounts.old_output_queue.to_account_info().data_len(), + )?; + let additional_bytes_rent = Rent::get()?.minimum_balance(additional_bytes as usize); + let new_mt_data = &mut ctx.accounts.new_state_merkle_tree.try_borrow_mut_data()?; + rollover_batch_state_tree( + old_merkle_tree_account, + ctx.accounts.old_state_merkle_tree.key(), + new_mt_data, + merkle_tree_rent, + ctx.accounts.new_state_merkle_tree.key(), + old_output_queue, + ctx.accounts.old_output_queue.key(), + &mut ctx.accounts.new_output_queue.try_borrow_mut_data()?, + queue_rent, + ctx.accounts.new_output_queue.key(), + additional_bytes_rent, + additional_bytes, + network_fee, + )?; + + transfer_lamports( + &ctx.accounts.old_output_queue.to_account_info(), + &ctx.accounts.fee_payer.to_account_info(), + merkle_tree_rent + queue_rent + additional_bytes_rent, + )?; + + Ok(()) +} + +pub fn rollover_batch_state_tree( + old_merkle_tree: &mut ZeroCopyBatchedMerkleTreeAccount, + old_mt_pubkey: Pubkey, + new_mt_data: &mut [u8], + new_mt_rent: u64, + new_mt_pubkey: Pubkey, + old_output_queue: &mut ZeroCopyBatchedQueueAccount, + old_queue_pubkey: Pubkey, + new_output_queue_data: &mut [u8], + new_output_queue_rent: u64, + new_output_queue_pubkey: Pubkey, + additional_bytes_rent: u64, + additional_bytes: u64, + network_fee: Option, +) -> Result<()> { + old_merkle_tree + .get_account_mut() + .metadata + .rollover(old_queue_pubkey, new_mt_pubkey)?; + + old_output_queue + .get_account_mut() + .metadata + .rollover(old_mt_pubkey, new_output_queue_pubkey)?; + let old_merkle_tree_account = old_merkle_tree.get_account(); + + if old_merkle_tree_account.next_index + < ((1 << old_merkle_tree_account.height) + * old_merkle_tree_account + .metadata + .rollover_metadata + .rollover_threshold + / 100) + { + return err!(crate::errors::AccountCompressionErrorCode::NotReadyForRollover); + } + if old_merkle_tree_account + .metadata + .rollover_metadata + .network_fee + == 0 + && network_fee.is_some() + { + msg!("Network fee must be 0 for manually forested trees."); + return err!(crate::errors::AccountCompressionErrorCode::InvalidNetworkFee); + } + + let params = InitStateTreeAccountsInstructionData { + index: old_merkle_tree_account.metadata.rollover_metadata.index, + program_owner: if_equals_none( + old_merkle_tree_account + .metadata + .access_metadata + .program_owner, + Pubkey::default(), + ), + forester: if_equals_none( + old_merkle_tree_account.metadata.access_metadata.forester, + Pubkey::default(), + ), + height: old_merkle_tree_account.height, + input_queue_batch_size: old_merkle_tree_account.queue.batch_size, + input_queue_zkp_batch_size: old_merkle_tree_account.queue.zkp_batch_size, + bloom_filter_capacity: old_merkle_tree_account.queue.bloom_filter_capacity, + bloom_filter_num_iters: old_merkle_tree.batches[0].num_iters, + root_history_capacity: old_merkle_tree_account.root_history_capacity, + network_fee, + rollover_threshold: if_equals_none( + old_merkle_tree_account + .metadata + .rollover_metadata + .rollover_threshold, + u64::MAX, + ), + close_threshold: if_equals_none( + old_merkle_tree_account + .metadata + .rollover_metadata + .close_threshold, + u64::MAX, + ), + input_queue_num_batches: old_merkle_tree_account.queue.num_batches, + additional_bytes, + output_queue_batch_size: old_output_queue.get_account().queue.batch_size, + output_queue_zkp_batch_size: old_output_queue.get_account().queue.zkp_batch_size, + output_queue_num_batches: old_output_queue.batches.len() as u64, + }; + + init_batched_state_merkle_tree_accounts( + old_merkle_tree_account.metadata.access_metadata.owner, + params, + new_output_queue_data, + new_output_queue_pubkey, + new_output_queue_rent, + new_mt_data, + new_mt_pubkey, + new_mt_rent, + additional_bytes_rent, + ) +} + +#[cfg(test)] +mod batch_state_tree_rollover_tests { + use rand::{rngs::StdRng, Rng}; + use solana_sdk::pubkey::Pubkey; + + use crate::{ + assert_mt_zero_copy_inited, + batched_merkle_tree::{ + get_merkle_tree_account_size, get_merkle_tree_account_size_default, + BatchedMerkleTreeAccount, ZeroCopyBatchedMerkleTreeAccount, + }, + batched_queue::{ + assert_queue_zero_copy_inited, get_output_queue_account_size, + get_output_queue_account_size_default, ZeroCopyBatchedQueueAccount, + }, + get_output_queue_account_default, init_batched_state_merkle_tree_accounts, + rollover_batch_state_tree, + tests::get_output_queue, + InitStateTreeAccountsInstructionData, + }; + + use super::assert_state_mt_roll_over; + + /// Test rollover of state tree + /// 1. failing: not ready for rollover + /// 2. failing: already rolled over + /// 3. failing: mt size mismatch + /// 4. failing: queue size mismatch + /// 5. failing: invalid network fee + /// 5. functional: rollover address tree + /// 6. failing: rollover threshold not set + /// 7. failing: invalid network fee + /// 8. functional: rollover address tree with network fee 0 additional bytes 0 + #[test] + fn test_rollover() { + let owner = Pubkey::new_unique(); + + { + let mt_account_size = get_merkle_tree_account_size_default(); + let mut mt_account_data = vec![0; mt_account_size]; + let mt_pubkey = Pubkey::new_unique(); + + let queue_account_size = get_output_queue_account_size_default(); + let mut queue_account_data = vec![0; queue_account_size]; + let queue_pubkey = Pubkey::new_unique(); + let params = InitStateTreeAccountsInstructionData::test_default(); + let merkle_tree_rent = 1_000_000_000; + let queue_rent = 1_000_000_001; + let additional_bytes_rent = 1_000_000_002; + let additional_bytes = 1_000; + // create first merkle tree + + init_batched_state_merkle_tree_accounts( + owner, + params.clone(), + &mut queue_account_data, + queue_pubkey, + queue_rent, + &mut mt_account_data, + mt_pubkey, + merkle_tree_rent, + additional_bytes_rent, + ) + .unwrap(); + + let ref_mt_account = BatchedMerkleTreeAccount::get_state_tree_default( + owner, + None, + None, + params.rollover_threshold, + 0, + params.network_fee.unwrap_or_default(), + params.input_queue_batch_size, + params.input_queue_zkp_batch_size, + params.bloom_filter_capacity, + params.root_history_capacity, + queue_pubkey, + params.height, + params.input_queue_num_batches, + ); + assert_mt_zero_copy_inited( + &mut mt_account_data, + ref_mt_account, + params.bloom_filter_num_iters, + ); + + let ref_output_queue_account = get_output_queue_account_default( + owner, + None, + None, + params.rollover_threshold, + 0, + params.output_queue_batch_size, + params.output_queue_zkp_batch_size, + params.additional_bytes, + merkle_tree_rent + additional_bytes_rent + queue_rent, + mt_pubkey, + params.height, + params.output_queue_num_batches, + params.network_fee.unwrap_or_default(), + ); + assert_queue_zero_copy_inited( + queue_account_data.as_mut_slice(), + ref_output_queue_account, + 0, + ); + let mut new_mt_account_data = vec![0; mt_account_size]; + let new_mt_pubkey = Pubkey::new_unique(); + + let mut new_queue_account_data = vec![0; queue_account_size]; + let new_queue_pubkey = Pubkey::new_unique(); + + // 1. Failing: not ready for rollover + { + let mut mt_account_data = mt_account_data.clone(); + let mut queue_account_data = queue_account_data.clone(); + let result = rollover_batch_state_tree( + &mut ZeroCopyBatchedMerkleTreeAccount::state_tree_from_bytes_mut( + &mut mt_account_data, + ) + .unwrap(), + mt_pubkey, + &mut new_mt_account_data, + merkle_tree_rent, + new_mt_pubkey, + &mut ZeroCopyBatchedQueueAccount::from_bytes_mut(&mut queue_account_data) + .unwrap(), + queue_pubkey, + &mut new_queue_account_data, + queue_rent, + new_queue_pubkey, + additional_bytes_rent, + additional_bytes, + params.network_fee, + ); + assert_eq!( + result, + Err(crate::errors::AccountCompressionErrorCode::NotReadyForRollover.into()) + ); + } + // 2. Failing rollover threshold not set + { + let mut mt_account_data = mt_account_data.clone(); + let merkle_tree = &mut ZeroCopyBatchedMerkleTreeAccount::state_tree_from_bytes_mut( + &mut mt_account_data, + ) + .unwrap(); + merkle_tree + .get_account_mut() + .metadata + .rollover_metadata + .rollover_threshold = u64::MAX; + let mut queue_account_data = queue_account_data.clone(); + let result = rollover_batch_state_tree( + &mut ZeroCopyBatchedMerkleTreeAccount::state_tree_from_bytes_mut( + &mut mt_account_data, + ) + .unwrap(), + mt_pubkey, + &mut new_mt_account_data, + merkle_tree_rent, + new_mt_pubkey, + &mut ZeroCopyBatchedQueueAccount::from_bytes_mut(&mut queue_account_data) + .unwrap(), + queue_pubkey, + &mut new_queue_account_data, + queue_rent, + new_queue_pubkey, + additional_bytes_rent, + additional_bytes, + params.network_fee, + ); + assert_eq!( + result, + Err(crate::errors::AccountCompressionErrorCode::RolloverNotConfigured.into()) + ); + } + // 3. Failing: invalid mt size + { + let mut mt_account_data = mt_account_data.clone(); + let mut queue_account_data = queue_account_data.clone(); + let merkle_tree = &mut ZeroCopyBatchedMerkleTreeAccount::state_tree_from_bytes_mut( + &mut mt_account_data, + ) + .unwrap(); + merkle_tree.get_account_mut().next_index = 1 << merkle_tree.get_account().height; + let mut new_mt_account_data = vec![0; mt_account_size - 1]; + let mut new_queue_account_data = vec![0; queue_account_size]; + + let result = rollover_batch_state_tree( + merkle_tree, + mt_pubkey, + &mut new_mt_account_data, + merkle_tree_rent, + new_mt_pubkey, + &mut ZeroCopyBatchedQueueAccount::from_bytes_mut(&mut queue_account_data) + .unwrap(), + queue_pubkey, + &mut new_queue_account_data, + queue_rent, + new_queue_pubkey, + additional_bytes_rent, + additional_bytes, + params.network_fee, + ); + assert_eq!( + result, + Err(crate::errors::AccountCompressionErrorCode::SizeMismatch.into()) + ); + } + // 4. Failing: invalid queue size + { + let mut mt_account_data = mt_account_data.clone(); + let mut queue_account_data = queue_account_data.clone(); + let merkle_tree = &mut ZeroCopyBatchedMerkleTreeAccount::state_tree_from_bytes_mut( + &mut mt_account_data, + ) + .unwrap(); + merkle_tree.get_account_mut().next_index = 1 << merkle_tree.get_account().height; + let mut new_mt_account_data = vec![0; mt_account_size]; + let mut new_queue_account_data = vec![0; queue_account_size - 1]; + + let result = rollover_batch_state_tree( + merkle_tree, + mt_pubkey, + &mut new_mt_account_data, + merkle_tree_rent, + new_mt_pubkey, + &mut ZeroCopyBatchedQueueAccount::from_bytes_mut(&mut queue_account_data) + .unwrap(), + queue_pubkey, + &mut new_queue_account_data, + queue_rent, + new_queue_pubkey, + additional_bytes_rent, + additional_bytes, + params.network_fee, + ); + assert_eq!( + result, + Err(crate::errors::AccountCompressionErrorCode::SizeMismatch.into()) + ); + } + // 5. Functional: rollover address tree + { + let merkle_tree = &mut ZeroCopyBatchedMerkleTreeAccount::state_tree_from_bytes_mut( + &mut mt_account_data, + ) + .unwrap(); + merkle_tree.get_account_mut().next_index = 1 << merkle_tree.get_account().height; + println!("new_mt_pubkey {:?}", new_mt_pubkey); + println!("new_queue_pubkey {:?}", new_queue_pubkey); + rollover_batch_state_tree( + merkle_tree, + mt_pubkey, + &mut new_mt_account_data, + merkle_tree_rent, + new_mt_pubkey, + &mut ZeroCopyBatchedQueueAccount::from_bytes_mut(&mut queue_account_data) + .unwrap(), + queue_pubkey, + &mut new_queue_account_data, + queue_rent, + new_queue_pubkey, + additional_bytes_rent, + additional_bytes, + params.network_fee, + ) + .unwrap(); + + let mut ref_rolledover_mt = ref_mt_account.clone(); + ref_rolledover_mt.next_index = 1 << merkle_tree.get_account().height; + let mut new_ref_output_queue_account = ref_output_queue_account.clone(); + new_ref_output_queue_account + .metadata + .rollover_metadata + .additional_bytes = additional_bytes; + new_ref_output_queue_account.metadata.associated_merkle_tree = new_mt_pubkey; + let mut new_ref_merkle_tree_account = ref_mt_account.clone(); + new_ref_merkle_tree_account.metadata.associated_queue = new_queue_pubkey; + + assert_state_mt_roll_over( + mt_account_data.to_vec(), + new_ref_merkle_tree_account, + new_mt_account_data.to_vec(), + mt_pubkey, + new_mt_pubkey, + params.bloom_filter_num_iters, + ref_rolledover_mt, + queue_account_data.to_vec(), + new_ref_output_queue_account, + new_queue_account_data.to_vec(), + new_queue_pubkey, + ref_output_queue_account, + queue_pubkey, + 1, + ); + } + // 6. Failing: already rolled over + { + let mut mt_account_data = mt_account_data.clone(); + let mut queue_account_data = queue_account_data.clone(); + + let mut new_mt_account_data = vec![0; mt_account_size]; + let mut new_queue_account_data = vec![0; queue_account_size]; + + let result = rollover_batch_state_tree( + &mut ZeroCopyBatchedMerkleTreeAccount::state_tree_from_bytes_mut( + &mut mt_account_data, + ) + .unwrap(), + mt_pubkey, + &mut new_mt_account_data, + merkle_tree_rent, + new_mt_pubkey, + &mut ZeroCopyBatchedQueueAccount::from_bytes_mut(&mut queue_account_data) + .unwrap(), + queue_pubkey, + &mut new_queue_account_data, + queue_rent, + new_queue_pubkey, + additional_bytes_rent, + additional_bytes, + params.network_fee, + ); + assert_eq!( + result, + Err( + crate::errors::AccountCompressionErrorCode::MerkleTreeAlreadyRolledOver + .into() + ) + ); + } + } + { + let mt_account_size = get_merkle_tree_account_size_default(); + let mut mt_account_data = vec![0; mt_account_size]; + let mt_pubkey = Pubkey::new_unique(); + + let queue_account_size = get_output_queue_account_size_default(); + let mut queue_account_data = vec![0; queue_account_size]; + let queue_pubkey = Pubkey::new_unique(); + let forester = Pubkey::new_unique(); + let mut params = InitStateTreeAccountsInstructionData::test_default(); + params.network_fee = None; + params.forester = Some(forester); + let merkle_tree_rent = 1_000_000_000; + let queue_rent = 1_000_000_001; + let additional_bytes_rent = 0; + let additional_bytes = 0; + // create first merkle tree + + init_batched_state_merkle_tree_accounts( + owner, + params.clone(), + &mut queue_account_data, + queue_pubkey, + queue_rent, + &mut mt_account_data, + mt_pubkey, + merkle_tree_rent, + additional_bytes_rent, + ) + .unwrap(); + let new_mt_pubkey = Pubkey::new_unique(); + let new_queue_pubkey = Pubkey::new_unique(); + // 7. failing Invalid network fee + { + let mut mt_account_data = mt_account_data.clone(); + let mut queue_account_data = queue_account_data.clone(); + let merkle_tree = &mut ZeroCopyBatchedMerkleTreeAccount::state_tree_from_bytes_mut( + &mut mt_account_data, + ) + .unwrap(); + merkle_tree.get_account_mut().next_index = 1 << merkle_tree.get_account().height; + let mut new_mt_account_data = vec![0; mt_account_size]; + let mut new_queue_account_data = vec![0; queue_account_size]; + + let result = rollover_batch_state_tree( + merkle_tree, + mt_pubkey, + &mut new_mt_account_data, + merkle_tree_rent, + new_mt_pubkey, + &mut ZeroCopyBatchedQueueAccount::from_bytes_mut(&mut queue_account_data) + .unwrap(), + queue_pubkey, + &mut new_queue_account_data, + queue_rent, + new_queue_pubkey, + additional_bytes_rent, + additional_bytes, + Some(1), + ); + assert_eq!( + result, + Err(crate::errors::AccountCompressionErrorCode::InvalidNetworkFee.into()) + ); + } + let mut new_mt_account_data = vec![0; mt_account_size]; + let mut new_queue_account_data = vec![0; queue_account_size]; + let mut ref_mt_account = BatchedMerkleTreeAccount::get_state_tree_default( + owner, + None, + None, + params.rollover_threshold, + 0, + params.network_fee.unwrap_or_default(), + params.input_queue_batch_size, + params.input_queue_zkp_batch_size, + params.bloom_filter_capacity, + params.root_history_capacity, + queue_pubkey, + params.height, + params.input_queue_num_batches, + ); + ref_mt_account.metadata.access_metadata.forester = forester; + let mut ref_output_queue_account = get_output_queue_account_default( + owner, + None, + None, + params.rollover_threshold, + 0, + params.output_queue_batch_size, + params.output_queue_zkp_batch_size, + params.additional_bytes, + merkle_tree_rent + additional_bytes_rent + queue_rent, + mt_pubkey, + params.height, + params.output_queue_num_batches, + params.network_fee.unwrap_or_default(), + ); + ref_output_queue_account + .metadata + .rollover_metadata + .network_fee = 0; + ref_output_queue_account.metadata.access_metadata.forester = forester; + assert_queue_zero_copy_inited( + queue_account_data.as_mut_slice(), + ref_output_queue_account, + 0, + ); + // 8. Functional: rollover address tree with network fee 0 additional bytes 0 + { + let merkle_tree = &mut ZeroCopyBatchedMerkleTreeAccount::state_tree_from_bytes_mut( + &mut mt_account_data, + ) + .unwrap(); + merkle_tree.get_account_mut().next_index = 1 << merkle_tree.get_account().height; + println!("new_mt_pubkey {:?}", new_mt_pubkey); + println!("new_queue_pubkey {:?}", new_queue_pubkey); + rollover_batch_state_tree( + merkle_tree, + mt_pubkey, + &mut new_mt_account_data, + merkle_tree_rent, + new_mt_pubkey, + &mut ZeroCopyBatchedQueueAccount::from_bytes_mut(&mut queue_account_data) + .unwrap(), + queue_pubkey, + &mut new_queue_account_data, + queue_rent, + new_queue_pubkey, + additional_bytes_rent, + additional_bytes, + params.network_fee, + ) + .unwrap(); + + let mut ref_rolledover_mt = ref_mt_account.clone(); + ref_rolledover_mt.next_index = 1 << merkle_tree.get_account().height; + let mut new_ref_output_queue_account = ref_output_queue_account.clone(); + new_ref_output_queue_account + .metadata + .rollover_metadata + .additional_bytes = additional_bytes; + new_ref_output_queue_account.metadata.associated_merkle_tree = new_mt_pubkey; + let mut new_ref_merkle_tree_account = ref_mt_account.clone(); + new_ref_merkle_tree_account.metadata.associated_queue = new_queue_pubkey; + + assert_state_mt_roll_over( + mt_account_data.to_vec(), + new_ref_merkle_tree_account, + new_mt_account_data.to_vec(), + mt_pubkey, + new_mt_pubkey, + params.bloom_filter_num_iters, + ref_rolledover_mt, + queue_account_data.to_vec(), + new_ref_output_queue_account, + new_queue_account_data.to_vec(), + new_queue_pubkey, + ref_output_queue_account, + queue_pubkey, + 1, + ); + } + } + } + + #[test] + fn test_rnd_rollover() { + use rand::SeedableRng; + let mut rng = StdRng::seed_from_u64(0); + for _ in 0..1000 { + println!("next iter ------------------------------------"); + let owner = Pubkey::new_unique(); + + let program_owner = if rng.gen_bool(0.5) { + Some(Pubkey::new_unique()) + } else { + None + }; + let forester = if rng.gen_bool(0.5) { + Some(Pubkey::new_unique()) + } else { + None + }; + let input_queue_zkp_batch_size = rng.gen_range(1..1000); + let output_queue_zkp_batch_size = rng.gen_range(1..1000); + let network_fee = if rng.gen_bool(0.5) && forester.is_some() { + Some(rng.gen_range(0..1000)) + } else { + None + }; + + let params = InitStateTreeAccountsInstructionData { + index: rng.gen_range(0..1000), + program_owner, + forester, + additional_bytes: rng.gen_range(0..1000), + bloom_filter_num_iters: rng.gen_range(0..4), + input_queue_batch_size: rng.gen_range(1..1000) * input_queue_zkp_batch_size, + output_queue_batch_size: rng.gen_range(1..1000) * output_queue_zkp_batch_size, + input_queue_zkp_batch_size, + output_queue_zkp_batch_size, + // 8 bits per byte, divisible by 8 for aligned memory + bloom_filter_capacity: rng.gen_range(0..100) * 8 * 8, + network_fee, + rollover_threshold: Some(rng.gen_range(0..100)), + close_threshold: None, + root_history_capacity: rng.gen_range(1..1000), + input_queue_num_batches: rng.gen_range(1..4), + output_queue_num_batches: rng.gen_range(1..4), + height: rng.gen_range(1..32), + }; + + let queue_account_size = get_output_queue_account_size( + params.output_queue_batch_size, + params.output_queue_zkp_batch_size, + params.output_queue_num_batches, + ); + + let mut output_queue_account_data = vec![0; queue_account_size]; + let output_queue_pubkey = Pubkey::new_unique(); + + let mt_account_size = get_merkle_tree_account_size( + params.input_queue_batch_size, + params.bloom_filter_capacity, + params.input_queue_zkp_batch_size, + params.root_history_capacity, + params.height, + params.input_queue_num_batches, + ); + + let mut mt_account_data = vec![0; mt_account_size]; + let mt_pubkey = Pubkey::new_unique(); + + let merkle_tree_rent = rng.gen_range(0..10000000); + let queue_rent = rng.gen_range(0..10000000); + let additional_bytes_rent = rng.gen_range(0..10000000); + let additional_bytes = rng.gen_range(0..1000); + init_batched_state_merkle_tree_accounts( + owner, + params.clone(), + &mut output_queue_account_data, + output_queue_pubkey, + queue_rent, + &mut mt_account_data, + mt_pubkey, + merkle_tree_rent, + additional_bytes_rent, + ) + .unwrap(); + let ref_output_queue_account = get_output_queue( + owner, + program_owner, + forester, + params.rollover_threshold, + params.index, + params.output_queue_batch_size, + params.output_queue_zkp_batch_size, + params.additional_bytes, + merkle_tree_rent + additional_bytes_rent + queue_rent, + mt_pubkey, + params.network_fee.unwrap_or_default(), + params.output_queue_num_batches, + params.height, + ); + assert_queue_zero_copy_inited( + output_queue_account_data.as_mut_slice(), + ref_output_queue_account, + 0, + ); + let ref_mt_account = BatchedMerkleTreeAccount::get_state_tree_default( + owner, + program_owner, + forester, + params.rollover_threshold, + params.index, + params.network_fee.unwrap_or_default(), + params.input_queue_batch_size, + params.input_queue_zkp_batch_size, + params.bloom_filter_capacity, + params.root_history_capacity, + output_queue_pubkey, + params.height, + params.input_queue_num_batches, + ); + assert_mt_zero_copy_inited( + &mut mt_account_data, + ref_mt_account, + params.bloom_filter_num_iters, + ); + + let mut new_mt_account_data = vec![0; mt_account_size]; + let new_mt_pubkey = Pubkey::new_unique(); + + let mut new_queue_account_data = vec![0; queue_account_size]; + let new_queue_pubkey = Pubkey::new_unique(); + + let merkle_tree = &mut ZeroCopyBatchedMerkleTreeAccount::state_tree_from_bytes_mut( + &mut mt_account_data, + ) + .unwrap(); + merkle_tree.get_account_mut().next_index = 1 << merkle_tree.get_account().height; + rollover_batch_state_tree( + merkle_tree, + mt_pubkey, + &mut new_mt_account_data, + merkle_tree_rent, + new_mt_pubkey, + &mut ZeroCopyBatchedQueueAccount::from_bytes_mut(&mut output_queue_account_data) + .unwrap(), + output_queue_pubkey, + &mut new_queue_account_data, + queue_rent, + new_queue_pubkey, + additional_bytes_rent, + additional_bytes, + params.network_fee, + ) + .unwrap(); + + let mut ref_rolledover_mt = ref_mt_account.clone(); + ref_rolledover_mt.next_index = 1 << merkle_tree.get_account().height; + let mut new_ref_output_queue_account = ref_output_queue_account.clone(); + new_ref_output_queue_account + .metadata + .rollover_metadata + .additional_bytes = additional_bytes; + new_ref_output_queue_account.metadata.associated_merkle_tree = new_mt_pubkey; + let mut new_ref_merkle_tree_account = ref_mt_account.clone(); + new_ref_merkle_tree_account.metadata.associated_queue = new_queue_pubkey; + + assert_state_mt_roll_over( + mt_account_data.to_vec(), + new_ref_merkle_tree_account, + new_mt_account_data.to_vec(), + mt_pubkey, + new_mt_pubkey, + params.bloom_filter_num_iters, + ref_rolledover_mt, + output_queue_account_data.to_vec(), + new_ref_output_queue_account, + new_queue_account_data.to_vec(), + new_queue_pubkey, + ref_output_queue_account, + output_queue_pubkey, + 1, + ); + } + } +} + +pub fn assert_state_mt_roll_over( + mt_account_data: Vec, + ref_mt_account: BatchedMerkleTreeAccount, + new_mt_account_data: Vec, + old_mt_pubkey: Pubkey, + new_mt_pubkey: Pubkey, + bloom_filter_num_iters: u64, + ref_rolledover_mt: BatchedMerkleTreeAccount, + mut queue_account_data: Vec, + ref_queue_account: BatchedQueueAccount, + mut new_queue_account_data: Vec, + new_queue_pubkey: Pubkey, + mut ref_rolledover_queue: BatchedQueueAccount, + old_queue_pubkey: Pubkey, + slot: u64, +) { + ref_rolledover_queue + .metadata + .rollover(old_mt_pubkey, new_queue_pubkey) + .unwrap(); + ref_rolledover_queue + .metadata + .rollover_metadata + .rolledover_slot = slot; + + assert_queue_zero_copy_inited(&mut new_queue_account_data, ref_queue_account, 0); + // assert_queue_zero_copy_inited(&mut queue_account_data, ref_rolledover_queue, 0); + let zero_copy_queue = + ZeroCopyBatchedQueueAccount::from_bytes_mut(&mut queue_account_data).unwrap(); + assert_eq!( + zero_copy_queue.get_account().metadata, + ref_rolledover_queue.metadata + ); + assert_mt_roll_over( + mt_account_data, + ref_mt_account, + new_mt_account_data, + new_mt_pubkey, + bloom_filter_num_iters, + ref_rolledover_mt, + old_queue_pubkey, + slot, + ) +} + +pub fn assert_mt_roll_over( + mut mt_account_data: Vec, + ref_mt_account: BatchedMerkleTreeAccount, + mut new_mt_account_data: Vec, + new_mt_pubkey: Pubkey, + bloom_filter_num_iters: u64, + mut ref_rolledover_mt: BatchedMerkleTreeAccount, + old_queue_pubkey: Pubkey, + slot: u64, +) { + ref_rolledover_mt + .metadata + .rollover(old_queue_pubkey, new_mt_pubkey) + .unwrap(); + ref_rolledover_mt.metadata.rollover_metadata.rolledover_slot = slot; + + assert_mt_zero_copy_inited( + &mut mt_account_data, + ref_rolledover_mt, + bloom_filter_num_iters, + ); + assert_mt_zero_copy_inited( + &mut new_mt_account_data, + ref_mt_account, + bloom_filter_num_iters, + ); +} diff --git a/programs/account-compression/src/lib.rs b/programs/account-compression/src/lib.rs index 29e97f0d29..23978159f3 100644 --- a/programs/account-compression/src/lib.rs +++ b/programs/account-compression/src/lib.rs @@ -229,4 +229,12 @@ pub mod account_compression { process_batch_append_leaves(&ctx, instruction_data)?; Ok(()) } + + pub fn rollover_batch_state_merkle_tree<'a, 'b, 'c: 'info, 'info>( + ctx: Context<'a, 'b, 'c, 'info, RolloverBatchStateMerkleTree<'info>>, + additional_bytes: u64, + network_fee: Option, + ) -> Result<()> { + process_rollover_batch_state_merkle_tree(ctx, additional_bytes, network_fee) + } } diff --git a/programs/account-compression/src/state/batched_merkle_tree.rs b/programs/account-compression/src/state/batched_merkle_tree.rs index 1f788d9d6b..6b7efd8351 100644 --- a/programs/account-compression/src/state/batched_merkle_tree.rs +++ b/programs/account-compression/src/state/batched_merkle_tree.rs @@ -195,6 +195,30 @@ impl ZeroCopyBatchedMerkleTreeAccount { unsafe { self.account.as_mut() }.unwrap() } + pub fn state_tree_from_account_info_mut( + account_info: &AccountInfo<'_>, + ) -> Result { + if *account_info.owner != crate::ID { + return err!(ErrorCode::AccountOwnedByWrongProgram); + } + if !account_info.is_writable { + return err!(ErrorCode::AccountNotMutable); + } + let account_data = &mut account_info.try_borrow_mut_data()?; + let merkle_tree = ZeroCopyBatchedMerkleTreeAccount::from_bytes_mut(account_data)?; + Ok(merkle_tree) + } + + pub fn state_tree_from_bytes_mut( + account_data: &mut [u8], + ) -> Result { + let merkle_tree = ZeroCopyBatchedMerkleTreeAccount::from_bytes_mut(account_data)?; + if merkle_tree.get_account().tree_type != TreeType::BatchedState as u64 { + return err!(AccountCompressionErrorCode::InvalidTreeType); + } + Ok(merkle_tree) + } + pub fn from_bytes_mut(account_data: &mut [u8]) -> Result { unsafe { let account = bytes_to_struct_checked::(account_data)?; diff --git a/programs/account-compression/src/state/batched_queue.rs b/programs/account-compression/src/state/batched_queue.rs index 253722881e..f0d2ccabbb 100644 --- a/programs/account-compression/src/state/batched_queue.rs +++ b/programs/account-compression/src/state/batched_queue.rs @@ -177,6 +177,23 @@ impl ZeroCopyBatchedQueueAccount { unsafe { &mut *self.account } } + pub fn output_queue_from_account_info_mut( + account_info: &AccountInfo<'_>, + ) -> Result { + if *account_info.owner != crate::ID { + return err!(ErrorCode::AccountOwnedByWrongProgram); + } + if !account_info.is_writable { + return err!(ErrorCode::AccountNotMutable); + } + let account_data = &mut account_info.try_borrow_mut_data()?; + let queue = Self::from_bytes_mut(account_data)?; + if queue.get_account().metadata.queue_type != QueueType::Output as u64 { + return err!(AccountCompressionErrorCode::InvalidQueueType); + } + Ok(queue) + } + pub fn from_bytes_mut(account_data: &mut [u8]) -> Result { let account = bytes_to_struct_checked::(account_data)?; unsafe { diff --git a/programs/account-compression/src/state/rollover.rs b/programs/account-compression/src/state/rollover.rs index f4bc2e477e..91766527d9 100644 --- a/programs/account-compression/src/state/rollover.rs +++ b/programs/account-compression/src/state/rollover.rs @@ -55,6 +55,10 @@ impl RolloverMetadata { { self.rolledover_slot = Clock::get()?.slot; } + #[cfg(not(target_os = "solana"))] + { + self.rolledover_slot = 1; + } Ok(()) } } diff --git a/programs/account-compression/src/utils/mod.rs b/programs/account-compression/src/utils/mod.rs index fb6867afa5..e4a091540e 100644 --- a/programs/account-compression/src/utils/mod.rs +++ b/programs/account-compression/src/utils/mod.rs @@ -4,3 +4,22 @@ pub mod check_signer_is_registered_or_authority; pub mod constants; pub mod queue; pub mod transfer_lamports; + +pub fn if_equals_zero_u64(value: u64) -> Option { + if value == 0 { + None + } else { + Some(value) + } +} + +pub fn if_equals_none(value: T, default: T) -> Option +where + T: PartialEq, +{ + if value == default { + None + } else { + Some(value) + } +} diff --git a/programs/registry/src/account_compression_cpi/initialize_batched_state_tree.rs b/programs/registry/src/account_compression_cpi/initialize_batched_state_tree.rs index 1dcccd40fa..7c386d5e78 100644 --- a/programs/registry/src/account_compression_cpi/initialize_batched_state_tree.rs +++ b/programs/registry/src/account_compression_cpi/initialize_batched_state_tree.rs @@ -24,8 +24,8 @@ pub struct InitializeBatchedStateMerkleTreeAndQueue<'info> { pub account_compression_program: Program<'info, AccountCompression>, pub protocol_config_pda: Account<'info, ProtocolConfigPda>, /// CHECK: (system program) new cpi context account. - pub cpi_context_account: Option>, - pub light_system_program: Option>, + pub cpi_context_account: AccountInfo<'info>, + pub light_system_program: Program<'info, LightSystemProgram>, } pub fn process_initialize_batched_state_merkle_tree( diff --git a/programs/registry/src/account_compression_cpi/initialize_tree_and_queue.rs b/programs/registry/src/account_compression_cpi/initialize_tree_and_queue.rs index 743b1d0aa4..3ebe17aac1 100644 --- a/programs/registry/src/account_compression_cpi/initialize_tree_and_queue.rs +++ b/programs/registry/src/account_compression_cpi/initialize_tree_and_queue.rs @@ -128,6 +128,11 @@ pub fn check_cpi_context( ) -> Result { let config_cpi_context_account_len = protocol_config.cpi_context_size as usize; if account.data_len() != config_cpi_context_account_len { + msg!( + "CPI context account data len: {}, expected: {}", + account.data_len(), + config_cpi_context_account_len + ); return err!(RegistryError::CpiContextAccountInvalidDataLen); } let rent = Rent::get()?; diff --git a/programs/registry/src/account_compression_cpi/mod.rs b/programs/registry/src/account_compression_cpi/mod.rs index 066aa4bba1..758f56a0ca 100644 --- a/programs/registry/src/account_compression_cpi/mod.rs +++ b/programs/registry/src/account_compression_cpi/mod.rs @@ -4,6 +4,7 @@ pub mod initialize_batched_state_tree; pub mod initialize_tree_and_queue; pub mod nullify; pub mod register_program; +pub mod rollover_batch_state_tree; pub mod rollover_state_tree; pub mod sdk; pub mod update_address_tree; diff --git a/programs/registry/src/account_compression_cpi/rollover_batch_state_tree.rs b/programs/registry/src/account_compression_cpi/rollover_batch_state_tree.rs new file mode 100644 index 0000000000..9f0be78972 --- /dev/null +++ b/programs/registry/src/account_compression_cpi/rollover_batch_state_tree.rs @@ -0,0 +1,71 @@ +use crate::protocol_config::state::ProtocolConfigPda; +use crate::ForesterEpochPda; +use account_compression::utils::if_equals_zero_u64; +use account_compression::{program::AccountCompression, utils::constants::CPI_AUTHORITY_PDA_SEED}; +use anchor_lang::prelude::*; + +#[derive(Accounts)] +pub struct RolloverBatchStateMerkleTree<'info> { + /// CHECK: only eligible foresters can nullify leaves. Is checked in ix. + #[account(mut)] + pub registered_forester_pda: Option>, + #[account(mut)] + pub authority: Signer<'info>, + /// CHECK: initializated in account compression program. + #[account(zero)] + pub new_state_merkle_tree: AccountInfo<'info>, + /// CHECK: in account compression program. + #[account(mut)] + pub old_state_merkle_tree: AccountInfo<'info>, + /// CHECK: initializated in account compression program. + #[account(zero)] + pub new_output_queue: AccountInfo<'info>, + /// CHECK: in account compression program. + #[account(mut)] + pub old_output_queue: AccountInfo<'info>, + /// CHECK: (system program) new cpi context account. + pub cpi_context_account: AccountInfo<'info>, + /// CHECK: (account compression program) access control. + pub registered_program_pda: AccountInfo<'info>, + /// CHECK: (seed constraints) used to invoke account compression program via cpi. + #[account(mut, seeds = [CPI_AUTHORITY_PDA_SEED], bump)] + pub cpi_authority: AccountInfo<'info>, + pub account_compression_program: Program<'info, AccountCompression>, + pub protocol_config_pda: Account<'info, ProtocolConfigPda>, + pub light_system_program: Program<'info, light_system_program::program::LightSystemProgram>, +} + +pub fn process_rollover_batch_state_merkle_tree( + ctx: &Context, + bump: u8, +) -> Result<()> { + let bump = &[bump]; + let seeds = [CPI_AUTHORITY_PDA_SEED, bump]; + let signer_seeds = &[&seeds[..]]; + let accounts = account_compression::cpi::accounts::RolloverBatchStateMerkleTree { + fee_payer: ctx.accounts.authority.to_account_info(), + authority: ctx.accounts.cpi_authority.to_account_info(), + old_state_merkle_tree: ctx.accounts.old_state_merkle_tree.to_account_info(), + new_state_merkle_tree: ctx.accounts.new_state_merkle_tree.to_account_info(), + old_output_queue: ctx.accounts.old_output_queue.to_account_info(), + new_output_queue: ctx.accounts.new_output_queue.to_account_info(), + registered_program_pda: Some(ctx.accounts.registered_program_pda.clone()), + }; + + let cpi_ctx = CpiContext::new_with_signer( + ctx.accounts.account_compression_program.to_account_info(), + accounts, + signer_seeds, + ); + let network_fee = if ctx.accounts.registered_forester_pda.is_some() { + if_equals_zero_u64(ctx.accounts.protocol_config_pda.config.network_fee) + } else { + None + }; + + account_compression::cpi::rollover_batch_state_merkle_tree( + cpi_ctx, + ctx.accounts.protocol_config_pda.config.cpi_context_size, + network_fee, + ) +} diff --git a/programs/registry/src/account_compression_cpi/sdk.rs b/programs/registry/src/account_compression_cpi/sdk.rs index 415c3506a3..10e1bcce6d 100644 --- a/programs/registry/src/account_compression_cpi/sdk.rs +++ b/programs/registry/src/account_compression_cpi/sdk.rs @@ -303,8 +303,8 @@ pub fn create_initialize_batched_merkle_tree_instruction( cpi_authority, account_compression_program: account_compression::ID, protocol_config_pda, - light_system_program: Some(LightSystemProgram::id()), - cpi_context_account: Some(cpi_context_pubkey), + light_system_program: LightSystemProgram::id(), + cpi_context_account: cpi_context_pubkey, }; Instruction { program_id: crate::ID, @@ -370,3 +370,46 @@ pub fn create_batch_nullify_instruction( data: instruction_data.data(), } } + +pub fn create_rollover_batch_state_tree_instruction( + forester: Pubkey, + derivation_pubkey: Pubkey, + old_state_merkle_tree: Pubkey, + new_state_merkle_tree: Pubkey, + old_output_queue: Pubkey, + new_output_queue: Pubkey, + cpi_context_account: Pubkey, + epoch: u64, + light_forester: bool, +) -> Instruction { + let register_program_pda = get_registered_program_pda(&crate::ID); + let registered_forester_pda = + get_forester_epoch_pda_from_authority(&derivation_pubkey, epoch).0; + let (cpi_authority, bump) = get_cpi_authority_pda(); + let instruction_data = crate::instruction::RolloverBatchStateMerkleTree { bump }; + let registered_forester_pda = if !light_forester { + None + } else { + Some(registered_forester_pda) + }; + + let accounts = crate::accounts::RolloverBatchStateMerkleTree { + authority: forester, + registered_forester_pda, + registered_program_pda: register_program_pda, + old_state_merkle_tree, + new_state_merkle_tree, + old_output_queue, + new_output_queue, + cpi_context_account, + cpi_authority, + account_compression_program: account_compression::ID, + protocol_config_pda: get_protocol_config_pda_address().0, + light_system_program: LightSystemProgram::id(), + }; + Instruction { + program_id: crate::ID, + accounts: accounts.to_account_metas(Some(true)), + data: instruction_data.data(), + } +} diff --git a/programs/registry/src/lib.rs b/programs/registry/src/lib.rs index 0cc2ea84b9..8b5ce025c6 100644 --- a/programs/registry/src/lib.rs +++ b/programs/registry/src/lib.rs @@ -9,8 +9,8 @@ pub mod errors; pub use crate::epoch::{finalize_registration::*, register_epoch::*, report_work::*}; pub use account_compression_cpi::{ batch_append::*, batch_nullify::*, initialize_batched_state_tree::*, - initialize_tree_and_queue::*, nullify::*, register_program::*, rollover_state_tree::*, - update_address_tree::*, + initialize_tree_and_queue::*, nullify::*, register_program::*, rollover_batch_state_tree::*, + rollover_state_tree::*, update_address_tree::*, }; pub use protocol_config::{initialize::*, update::*}; @@ -44,6 +44,7 @@ declare_id!("Lighton6oQpVkeewmo2mcPTQQp7kYHr4fWpAgJyEmDX"); #[program] pub mod light_registry { + use account_compression::batched_merkle_tree::ZeroCopyBatchedMerkleTreeAccount; use constants::DEFAULT_WORK_V1; use super::*; @@ -453,16 +454,16 @@ pub mod light_registry { )?; check_cpi_context( - ctx.accounts.cpi_context_account.as_ref().to_account_info(), + ctx.accounts.cpi_context_account.to_account_info(), &ctx.accounts.protocol_config_pda.config, )?; process_rollover_state_merkle_tree_and_queue(&ctx, bump)?; process_initialize_cpi_context( bump, ctx.accounts.authority.to_account_info(), - ctx.accounts.cpi_context_account.as_ref().to_account_info(), + ctx.accounts.cpi_context_account.to_account_info(), ctx.accounts.new_merkle_tree.to_account_info(), - ctx.accounts.light_system_program.as_ref().to_account_info(), + ctx.accounts.light_system_program.to_account_info(), ) } @@ -485,11 +486,7 @@ pub mod light_registry { return err!(RegistryError::ForesterUndefined); } check_cpi_context( - ctx.accounts - .cpi_context_account - .as_ref() - .unwrap() - .to_account_info(), + ctx.accounts.cpi_context_account.to_account_info(), &ctx.accounts.protocol_config_pda.config, )?; @@ -498,17 +495,9 @@ pub mod light_registry { process_initialize_cpi_context( bump, ctx.accounts.authority.to_account_info(), - ctx.accounts - .cpi_context_account - .as_ref() - .unwrap() - .to_account_info(), + ctx.accounts.cpi_context_account.to_account_info(), ctx.accounts.merkle_tree.to_account_info(), - ctx.accounts - .light_system_program - .as_ref() - .unwrap() - .to_account_info(), + ctx.accounts.light_system_program.to_account_info(), ) } @@ -549,6 +538,36 @@ pub mod light_registry { } process_batch_append(&ctx, bump, data) } + + pub fn rollover_batch_state_merkle_tree<'info>( + ctx: Context<'_, '_, '_, 'info, RolloverBatchStateMerkleTree<'info>>, + bump: u8, + ) -> Result<()> { + let account = ZeroCopyBatchedMerkleTreeAccount::state_tree_from_account_info_mut( + &ctx.accounts.old_state_merkle_tree, + )?; + check_forester( + &account.get_account().metadata, + ctx.accounts.authority.key(), + ctx.accounts.old_state_merkle_tree.key(), + &mut ctx.accounts.registered_forester_pda, + DEFAULT_WORK_V1, + )?; + check_cpi_context( + ctx.accounts.cpi_context_account.to_account_info(), + &ctx.accounts.protocol_config_pda.config, + )?; + + process_rollover_batch_state_merkle_tree(&ctx, bump)?; + + process_initialize_cpi_context( + bump, + ctx.accounts.authority.to_account_info(), + ctx.accounts.cpi_context_account.to_account_info(), + ctx.accounts.new_state_merkle_tree.to_account_info(), + ctx.accounts.light_system_program.to_account_info(), + ) + } } /// if registered_forester_pda is not None check forester eligibility and network_fee is not 0 diff --git a/test-programs/account-compression-test/tests/batched_merkle_tree_test.rs b/test-programs/account-compression-test/tests/batched_merkle_tree_test.rs index d33e95b652..6be6dea0d1 100644 --- a/test-programs/account-compression-test/tests/batched_merkle_tree_test.rs +++ b/test-programs/account-compression-test/tests/batched_merkle_tree_test.rs @@ -14,13 +14,17 @@ use account_compression::{assert_mt_zero_copy_inited, get_output_queue_account_d use account_compression::{ batched_merkle_tree::BatchedMerkleTreeAccount, InitStateTreeAccountsInstructionData, ID, }; +use anchor_lang::error::ErrorCode; use anchor_lang::prelude::AccountMeta; use anchor_lang::{AnchorSerialize, InstructionData, ToAccountMetas}; use light_prover_client::gnark::helpers::{spawn_prover, ProofType, ProverConfig}; use light_prover_client::mock_batched_forester::{MockBatchedForester, MockTxEvent}; use light_system_program::invoke::verify_state_proof::create_tx_hash_offchain; +use light_test_utils::test_batch_forester::assert_perform_state_mt_roll_over; use light_test_utils::test_env::NOOP_PROGRAM_ID; -use light_test_utils::{assert_rpc_error, create_account_instruction, RpcConnection, RpcError}; +use light_test_utils::{ + airdrop_lamports, assert_rpc_error, create_account_instruction, RpcConnection, RpcError, +}; use light_test_utils::{rpc::ProgramTestRpcConnection, AccountZeroCopy}; use light_verifier::CompressedProof; use serial_test::serial; @@ -183,6 +187,7 @@ async fn test_batch_state_merkle_tree() { merkle_tree_pubkey, params.height, params.output_queue_num_batches, + params.network_fee.unwrap_or_default(), ); assert_queue_zero_copy_inited( &mut queue.account.data.as_mut_slice(), @@ -790,7 +795,6 @@ async fn test_init_batch_state_merkle_trees() { let context = program_test.start_with_context().await; let mut context = ProgramTestRpcConnection { context }; - let payer_pubkey = context.get_payer().pubkey(); let payer = context.get_payer().insecure_clone(); let params = InitStateTreeAccountsInstructionData::test_default(); let e2e_test_params = InitStateTreeAccountsInstructionData::e2e_test_default(); @@ -799,72 +803,18 @@ async fn test_init_batch_state_merkle_trees() { for params in param_vec.iter() { println!("Init new mt with params {:?}", params); let merkle_tree_keypair = Keypair::new(); - let merkle_tree_pubkey = merkle_tree_keypair.pubkey(); let nullifier_queue_keypair = Keypair::new(); + let merkle_tree_pubkey = merkle_tree_keypair.pubkey(); let output_queue_pubkey = nullifier_queue_keypair.pubkey(); - let queue_account_size = get_output_queue_account_size( - params.output_queue_batch_size, - params.output_queue_zkp_batch_size, - params.output_queue_num_batches, - ); - let mt_account_size = get_merkle_tree_account_size( - params.input_queue_batch_size, - params.bloom_filter_capacity, - params.input_queue_zkp_batch_size, - params.root_history_capacity, - params.height, - params.input_queue_num_batches, - ); - let queue_rent = context - .get_minimum_balance_for_rent_exemption(queue_account_size) - .await - .unwrap(); - let create_queue_account_ix = create_account_instruction( - &payer_pubkey, - queue_account_size, - queue_rent, - &ID, - Some(&nullifier_queue_keypair), - ); - let mt_rent = context - .get_minimum_balance_for_rent_exemption(mt_account_size) - .await - .unwrap(); - let additional_bytes_rent = context - .get_minimum_balance_for_rent_exemption(params.additional_bytes as usize) - .await - .unwrap(); - let total_rent = queue_rent + mt_rent + additional_bytes_rent; - let create_mt_account_ix = create_account_instruction( - &payer_pubkey, - mt_account_size, - mt_rent, - &ID, - Some(&merkle_tree_keypair), - ); - - let instruction = - account_compression::instruction::InitializeBatchedStateMerkleTree { params: *params }; - let accounts = account_compression::accounts::InitializeBatchedStateMerkleTreeAndQueue { - authority: context.get_payer().pubkey(), - merkle_tree: merkle_tree_pubkey, - queue: output_queue_pubkey, - registered_program_pda: None, - }; - - let instruction = Instruction { - program_id: ID, - accounts: accounts.to_account_metas(Some(true)), - data: instruction.data(), - }; - context - .create_and_send_transaction( - &[create_queue_account_ix, create_mt_account_ix, instruction], - &payer_pubkey, - &[&payer, &nullifier_queue_keypair, &merkle_tree_keypair], - ) - .await - .unwrap(); + let (_, total_rent) = perform_init_batch_state_merkle_tree( + &mut context, + &payer, + &merkle_tree_keypair, + &nullifier_queue_keypair, + params.clone(), + ) + .await + .unwrap(); let merkle_tree = AccountZeroCopy::::new(&mut context, merkle_tree_pubkey) .await; @@ -909,6 +859,7 @@ async fn test_init_batch_state_merkle_trees() { merkle_tree_pubkey, params.height, params.output_queue_num_batches, + params.network_fee.unwrap_or_default(), ); assert_queue_zero_copy_inited( &mut queue.account.data.as_mut_slice(), @@ -917,3 +868,415 @@ async fn test_init_batch_state_merkle_trees() { ); } } + +pub async fn perform_init_batch_state_merkle_tree( + context: &mut ProgramTestRpcConnection, + payer: &Keypair, + merkle_tree_keypair: &Keypair, + nullifier_queue_keypair: &Keypair, + params: InitStateTreeAccountsInstructionData, +) -> Result<(Signature, u64), RpcError> { + let payer_pubkey = payer.pubkey(); + let merkle_tree_pubkey = merkle_tree_keypair.pubkey(); + + let output_queue_pubkey = nullifier_queue_keypair.pubkey(); + let queue_account_size = get_output_queue_account_size( + params.output_queue_batch_size, + params.output_queue_zkp_batch_size, + params.output_queue_num_batches, + ); + let mt_account_size = get_merkle_tree_account_size( + params.input_queue_batch_size, + params.bloom_filter_capacity, + params.input_queue_zkp_batch_size, + params.root_history_capacity, + params.height, + params.input_queue_num_batches, + ); + let queue_rent = context + .get_minimum_balance_for_rent_exemption(queue_account_size) + .await + .unwrap(); + let create_queue_account_ix = create_account_instruction( + &payer_pubkey, + queue_account_size, + queue_rent, + &ID, + Some(&nullifier_queue_keypair), + ); + let mt_rent = context + .get_minimum_balance_for_rent_exemption(mt_account_size) + .await + .unwrap(); + let additional_bytes_rent = context + .get_minimum_balance_for_rent_exemption(params.additional_bytes as usize) + .await + .unwrap(); + let total_rent = queue_rent + mt_rent + additional_bytes_rent; + let create_mt_account_ix = create_account_instruction( + &payer_pubkey, + mt_account_size, + mt_rent, + &ID, + Some(&merkle_tree_keypair), + ); + + let instruction = account_compression::instruction::InitializeBatchedStateMerkleTree { params }; + let accounts = account_compression::accounts::InitializeBatchedStateMerkleTreeAndQueue { + authority: context.get_payer().pubkey(), + merkle_tree: merkle_tree_pubkey, + queue: output_queue_pubkey, + registered_program_pda: None, + }; + + let instruction = Instruction { + program_id: ID, + accounts: accounts.to_account_metas(Some(true)), + data: instruction.data(), + }; + Ok(( + context + .create_and_send_transaction( + &[create_queue_account_ix, create_mt_account_ix, instruction], + &payer_pubkey, + &[&payer, &nullifier_queue_keypair, &merkle_tree_keypair], + ) + .await?, + total_rent, + )) +} + +/// Tests: +/// 1. Failing - Invalid authority +/// 2. Failing - State tree invalid program owner +/// 3. Failing - Queue invalid program owner +/// 4. Failing - State tree invalid discriminator +/// 5. Failing - Queue invalid discriminator +/// 6. Failing - Merkle tree and queue not associated +/// 7. functional +#[serial] +#[tokio::test] +async fn test_rollover_batch_state_merkle_trees() { + let mut program_test = ProgramTest::default(); + program_test.add_program("account_compression", ID, None); + program_test.add_program( + "spl_noop", + Pubkey::new_from_array(account_compression::utils::constants::NOOP_PUBKEY), + None, + ); + program_test.set_compute_max_units(1_400_000u64); + let context = program_test.start_with_context().await; + let mut context = ProgramTestRpcConnection { context }; + let payer = context.get_payer().insecure_clone(); + let mut params = InitStateTreeAccountsInstructionData::test_default(); + params.rollover_threshold = Some(0); + println!("Init new mt with params {:?}", params); + let merkle_tree_keypair = Keypair::new(); + let nullifier_queue_keypair = Keypair::new(); + perform_init_batch_state_merkle_tree( + &mut context, + &payer, + &merkle_tree_keypair, + &nullifier_queue_keypair, + params.clone(), + ) + .await + .unwrap(); + let mut mock_indexer = MockBatchedForester::<26>::default(); + let output_queue_pubkey = nullifier_queue_keypair.pubkey(); + + perform_insert_into_output_queue( + &mut context, + &mut mock_indexer, + output_queue_pubkey, + &payer, + &mut 0, + 1, + ) + .await + .unwrap(); + let new_state_merkle_tree_keypair = Keypair::new(); + let new_output_queue_keypair = Keypair::new(); + // 1. failing - invalid authority + { + let invalid_authority = Keypair::new(); + airdrop_lamports(&mut context, &invalid_authority.pubkey(), 10_000_000_000) + .await + .unwrap(); + let result = perform_rollover_batch_state_merkle_tree( + &mut context, + &invalid_authority, + merkle_tree_keypair.pubkey(), + nullifier_queue_keypair.pubkey(), + &new_state_merkle_tree_keypair, + &new_output_queue_keypair, + params.additional_bytes, + params.network_fee, + BatchStateMerkleTreeRollOverTestMode::Functional, + ) + .await; + assert_rpc_error( + result, + 2, + AccountCompressionErrorCode::InvalidAuthority.into(), + ) + .unwrap(); + } + // 2. failing - state tree invalid program owner + { + let result = perform_rollover_batch_state_merkle_tree( + &mut context, + &payer, + merkle_tree_keypair.pubkey(), + nullifier_queue_keypair.pubkey(), + &new_state_merkle_tree_keypair, + &new_output_queue_keypair, + params.additional_bytes, + params.network_fee, + BatchStateMerkleTreeRollOverTestMode::InvalidProgramOwnerMerkleTree, + ) + .await; + assert_rpc_error(result, 2, ErrorCode::AccountOwnedByWrongProgram.into()).unwrap(); + } + // 3. failing - queue invalid program owner + { + let result = perform_rollover_batch_state_merkle_tree( + &mut context, + &payer, + merkle_tree_keypair.pubkey(), + nullifier_queue_keypair.pubkey(), + &new_state_merkle_tree_keypair, + &new_output_queue_keypair, + params.additional_bytes, + params.network_fee, + BatchStateMerkleTreeRollOverTestMode::InvalidProgramOwnerQueue, + ) + .await; + assert_rpc_error(result, 2, ErrorCode::AccountOwnedByWrongProgram.into()).unwrap(); + } + // 4. failing - state tree invalid discriminator + { + let result = perform_rollover_batch_state_merkle_tree( + &mut context, + &payer, + merkle_tree_keypair.pubkey(), + nullifier_queue_keypair.pubkey(), + &new_state_merkle_tree_keypair, + &new_output_queue_keypair, + params.additional_bytes, + params.network_fee, + BatchStateMerkleTreeRollOverTestMode::InvalidDiscriminatorMerkleTree, + ) + .await; + assert_rpc_error( + result, + 2, + AccountCompressionErrorCode::InvalidDiscriminator.into(), + ) + .unwrap(); + } + // 5. failing - queue invalid discriminator + { + let result = perform_rollover_batch_state_merkle_tree( + &mut context, + &payer, + merkle_tree_keypair.pubkey(), + nullifier_queue_keypair.pubkey(), + &new_state_merkle_tree_keypair, + &new_output_queue_keypair, + params.additional_bytes, + params.network_fee, + BatchStateMerkleTreeRollOverTestMode::InvalidDiscriminatorQueue, + ) + .await; + assert_rpc_error( + result, + 2, + AccountCompressionErrorCode::InvalidDiscriminator.into(), + ) + .unwrap(); + } + // 6. failing - merkle tree and queue not associated + { + let merkle_tree_keypair_1 = Keypair::new(); + let nullifier_queue_keypair_1 = Keypair::new(); + perform_init_batch_state_merkle_tree( + &mut context, + &payer, + &merkle_tree_keypair_1, + &nullifier_queue_keypair_1, + params.clone(), + ) + .await + .unwrap(); + let result = perform_rollover_batch_state_merkle_tree( + &mut context, + &payer, + merkle_tree_keypair_1.pubkey(), + nullifier_queue_keypair.pubkey(), + &new_state_merkle_tree_keypair, + &new_output_queue_keypair, + params.additional_bytes, + params.network_fee, + BatchStateMerkleTreeRollOverTestMode::Functional, + ) + .await; + assert_rpc_error( + result, + 2, + AccountCompressionErrorCode::MerkleTreeAndQueueNotAssociated.into(), + ) + .unwrap(); + } + // 7. functional + { + perform_rollover_batch_state_merkle_tree( + &mut context, + &payer, + merkle_tree_keypair.pubkey(), + nullifier_queue_keypair.pubkey(), + &new_state_merkle_tree_keypair, + &new_output_queue_keypair, + params.additional_bytes, + params.network_fee, + BatchStateMerkleTreeRollOverTestMode::Functional, + ) + .await + .unwrap(); + let additional_bytes_rent = context + .get_minimum_balance_for_rent_exemption(params.additional_bytes as usize) + .await + .unwrap(); + assert_perform_state_mt_roll_over( + &mut context, + payer.pubkey(), + merkle_tree_keypair.pubkey(), + new_state_merkle_tree_keypair.pubkey(), + nullifier_queue_keypair.pubkey(), + new_output_queue_keypair.pubkey(), + params, + additional_bytes_rent, + ) + .await; + } +} + +#[derive(Debug, PartialEq)] +pub enum BatchStateMerkleTreeRollOverTestMode { + Functional, + InvalidProgramOwnerMerkleTree, + InvalidProgramOwnerQueue, + InvalidDiscriminatorMerkleTree, + InvalidDiscriminatorQueue, +} + +pub async fn perform_rollover_batch_state_merkle_tree( + rpc: &mut R, + payer: &Keypair, + old_merkle_tree_pubkey: Pubkey, + old_output_queue_pubkey: Pubkey, + new_state_merkle_tree_keypair: &Keypair, + new_output_queue_keypair: &Keypair, + additional_bytes: u64, + network_fee: Option, + test_mode: BatchStateMerkleTreeRollOverTestMode, +) -> Result { + let payer_pubkey = payer.pubkey(); + let mut account = rpc.get_account(old_merkle_tree_pubkey).await?.unwrap(); + let old_merkle_tree = + ZeroCopyBatchedMerkleTreeAccount::state_tree_from_bytes_mut(account.data.as_mut_slice()) + .unwrap(); + let batch_zero = &old_merkle_tree.batches[0]; + let num_batches = old_merkle_tree.batches.len(); + let old_merkle_tree = old_merkle_tree.get_account(); + let mt_account_size = get_merkle_tree_account_size( + batch_zero.batch_size, + batch_zero.bloom_filter_capacity, + batch_zero.zkp_batch_size, + old_merkle_tree.root_history_capacity, + old_merkle_tree.height, + num_batches as u64, + ); + + let mt_rent = rpc + .get_minimum_balance_for_rent_exemption(mt_account_size) + .await + .unwrap(); + + let mut account = rpc.get_account(old_output_queue_pubkey).await?.unwrap(); + let old_queue_account = + ZeroCopyBatchedQueueAccount::from_bytes_mut(account.data.as_mut_slice()).unwrap(); + let batch_zero = &old_queue_account.batches[0]; + let queue_account_size = get_output_queue_account_size( + batch_zero.batch_size, + batch_zero.zkp_batch_size, + num_batches as u64, + ); + let queue_rent = rpc + .get_minimum_balance_for_rent_exemption(queue_account_size) + .await + .unwrap(); + + let create_mt_account_ix = create_account_instruction( + &payer_pubkey, + mt_account_size, + mt_rent, + &account_compression::ID, + Some(&new_state_merkle_tree_keypair), + ); + + let create_queue_account_ix = create_account_instruction( + &payer_pubkey, + queue_account_size, + queue_rent, + &account_compression::ID, + Some(new_output_queue_keypair), + ); + let old_state_merkle_tree = if test_mode + == BatchStateMerkleTreeRollOverTestMode::InvalidProgramOwnerMerkleTree + { + Pubkey::new_unique() + } else if test_mode == BatchStateMerkleTreeRollOverTestMode::InvalidDiscriminatorMerkleTree { + old_output_queue_pubkey + } else { + old_merkle_tree_pubkey + }; + let old_output_queue = + if test_mode == BatchStateMerkleTreeRollOverTestMode::InvalidProgramOwnerQueue { + Pubkey::new_unique() + } else if test_mode == BatchStateMerkleTreeRollOverTestMode::InvalidDiscriminatorQueue { + old_merkle_tree_pubkey + } else { + old_output_queue_pubkey + }; + let accounts = account_compression::accounts::RolloverBatchStateMerkleTree { + fee_payer: payer_pubkey, + authority: payer_pubkey, + old_state_merkle_tree, + new_state_merkle_tree: new_state_merkle_tree_keypair.pubkey(), + old_output_queue, + new_output_queue: new_output_queue_keypair.pubkey(), + registered_program_pda: None, + }; + let instruction_data = account_compression::instruction::RolloverBatchStateMerkleTree { + additional_bytes, + network_fee, + }; + let instruction = Instruction { + program_id: ID, + accounts: accounts.to_account_metas(Some(true)), + data: instruction_data.data(), + }; + + Ok(rpc + .create_and_send_transaction( + &[create_mt_account_ix, create_queue_account_ix, instruction], + &payer_pubkey, + &[ + &payer, + &new_state_merkle_tree_keypair, + &new_output_queue_keypair, + ], + ) + .await?) +} diff --git a/test-programs/registry-test/tests/tests.rs b/test-programs/registry-test/tests/tests.rs index dc7032008e..483db17a3b 100644 --- a/test-programs/registry-test/tests/tests.rs +++ b/test-programs/registry-test/tests/tests.rs @@ -1,4 +1,4 @@ -#![cfg(feature = "test-sbf")] +// #![cfg(feature = "test-sbf")] use account_compression::batched_merkle_tree::ZeroCopyBatchedMerkleTreeAccount; use account_compression::{ @@ -6,6 +6,7 @@ use account_compression::{ NullifierQueueConfig, StateMerkleTreeConfig, }; use anchor_lang::{InstructionData, ToAccountMetas}; +use forester_utils::airdrop_lamports; use forester_utils::forester_epoch::get_epoch_phases; use light_program_test::test_env::{ create_address_merkle_tree_and_queue_account, create_state_merkle_tree_and_queue_account, @@ -36,9 +37,12 @@ use light_test_utils::assert_epoch::{ assert_report_work, fetch_epoch_and_forester_pdas, }; use light_test_utils::e2e_test_env::{init_program_test_env, init_program_test_env_forester}; +use light_test_utils::indexer::TestIndexer; use light_test_utils::rpc::ProgramTestRpcConnection; use light_test_utils::test_batch_forester::{ - create_append_batch_ix_data, perform_batch_append, perform_batch_nullify, + assert_perform_state_mt_roll_over, create_append_batch_ix_data, + create_batched_state_merkle_tree, perform_batch_append, perform_batch_nullify, + perform_rollover_batch_state_merkle_tree, }; use light_test_utils::test_env::{ create_address_merkle_tree_and_queue_account, create_state_merkle_tree_and_queue_account, @@ -1318,3 +1322,203 @@ async fn init_accounts() { println!("forester pubkey: {:?}", keypairs.forester.pubkey()); setup_accounts(keypairs, SolanaRpcUrl::Localnet).await; } + +/// Test: +/// 1. Failing: rollover with invalid forester +/// 2. Functional: rollover with network fee +/// 3. Functional: rollover without network fee and custom forester +/// 4. failing: create with state tree with custom forester and invalid non-zero network fee +#[serial] +#[tokio::test] +async fn test_rollover_batch_state_tree() { + { + let mut params = InitStateTreeAccountsInstructionData::test_default(); + params.rollover_threshold = Some(0); + let is_light_forester = true; + + let (mut rpc, env_accounts) = + setup_test_programs_with_accounts_with_protocol_config_and_batched_tree_params( + None, + ProtocolConfig::default(), + true, + params, + ) + .await; + let payer = rpc.get_payer().insecure_clone(); + let mut test_indexer: TestIndexer = TestIndexer::init_from_env( + &env_accounts.forester.insecure_clone(), + &env_accounts, + None, + ) + .await; + light_test_utils::system_program::compress_sol_test( + &mut rpc, + &mut test_indexer, + &payer, + &[], + false, + 1_000_000, + &env_accounts.batched_output_queue, + None, + ) + .await + .unwrap(); + let new_merkle_tree_keypair = Keypair::new(); + let new_nullifier_queue_keypair = Keypair::new(); + let new_cpi_context = Keypair::new(); + + // 1. failing invalid forester + { + let unregistered_forester_keypair = Keypair::new(); + rpc.airdrop_lamports(&unregistered_forester_keypair.pubkey(), 1_000_000_000) + .await + .unwrap(); + + let result = perform_rollover_batch_state_merkle_tree( + &mut rpc, + &payer, + env_accounts.forester.pubkey(), + env_accounts.batched_state_merkle_tree, + env_accounts.batched_output_queue, + &new_merkle_tree_keypair, + &new_nullifier_queue_keypair, + &new_cpi_context, + 0, + is_light_forester, + ) + .await; + + assert_rpc_error(result, 3, RegistryError::InvalidForester.into()).unwrap(); + } + + // 2. functional with network fee + { + perform_rollover_batch_state_merkle_tree( + &mut rpc, + &env_accounts.forester, + env_accounts.forester.pubkey(), + env_accounts.batched_state_merkle_tree, + env_accounts.batched_output_queue, + &new_merkle_tree_keypair, + &new_nullifier_queue_keypair, + &new_cpi_context, + 0, + is_light_forester, + ) + .await + .unwrap(); + let new_cpi_ctx_account = rpc + .get_account(new_cpi_context.pubkey()) + .await + .unwrap() + .unwrap(); + assert_perform_state_mt_roll_over( + &mut rpc, + env_accounts.group_pda, + env_accounts.batched_state_merkle_tree, + new_merkle_tree_keypair.pubkey(), + env_accounts.batched_output_queue, + new_nullifier_queue_keypair.pubkey(), + params, + new_cpi_ctx_account.lamports, + ) + .await; + } + } + { + let custom_forester = Keypair::new(); + let mut params = InitStateTreeAccountsInstructionData::test_default(); + params.rollover_threshold = Some(0); + params.forester = Some(custom_forester.pubkey()); + params.network_fee = None; + let is_light_forester = false; + + let (mut rpc, env_accounts) = + setup_test_programs_with_accounts_with_protocol_config_and_batched_tree_params( + None, + ProtocolConfig::default(), + true, + params, + ) + .await; + airdrop_lamports(&mut rpc, &custom_forester.pubkey(), 10_000_000_000) + .await + .unwrap(); + let payer = rpc.get_payer().insecure_clone(); + let mut test_indexer: TestIndexer = TestIndexer::init_from_env( + &env_accounts.forester.insecure_clone(), + &env_accounts, + None, + ) + .await; + light_test_utils::system_program::compress_sol_test( + &mut rpc, + &mut test_indexer, + &payer, + &[], + false, + 1_000_000, + &env_accounts.batched_output_queue, + None, + ) + .await + .unwrap(); + let new_merkle_tree_keypair = Keypair::new(); + let new_nullifier_queue_keypair = Keypair::new(); + let new_cpi_context = Keypair::new(); + + // 3. functional without network fee and forester + + perform_rollover_batch_state_merkle_tree( + &mut rpc, + &custom_forester, + custom_forester.pubkey(), + env_accounts.batched_state_merkle_tree, + env_accounts.batched_output_queue, + &new_merkle_tree_keypair, + &new_nullifier_queue_keypair, + &new_cpi_context, + 0, + is_light_forester, + ) + .await + .unwrap(); + let new_cpi_ctx_account = rpc + .get_account(new_cpi_context.pubkey()) + .await + .unwrap() + .unwrap(); + assert_perform_state_mt_roll_over( + &mut rpc, + env_accounts.group_pda, + env_accounts.batched_state_merkle_tree, + new_merkle_tree_keypair.pubkey(), + env_accounts.batched_output_queue, + new_nullifier_queue_keypair.pubkey(), + params, + new_cpi_ctx_account.lamports, + ) + .await; + // 4. failing with custom forester and non-zero network fee + { + let mut params = InitStateTreeAccountsInstructionData::test_default(); + params.rollover_threshold = Some(0); + params.forester = Some(custom_forester.pubkey()); + params.network_fee = Some(1); + let new_merkle_tree_keypair = Keypair::new(); + let new_nullifier_queue_keypair = Keypair::new(); + let new_cpi_context = Keypair::new(); + let result = create_batched_state_merkle_tree( + &env_accounts.governance_authority, + true, + &mut rpc, + &new_merkle_tree_keypair, + &new_nullifier_queue_keypair, + &new_cpi_context, + params, + ) + .await; + assert_rpc_error(result, 3, RegistryError::InvalidNetworkFee.into()).unwrap(); + } + } +} diff --git a/test-utils/src/test_batch_forester.rs b/test-utils/src/test_batch_forester.rs index 85388ca160..6a1144917a 100644 --- a/test-utils/src/test_batch_forester.rs +++ b/test-utils/src/test_batch_forester.rs @@ -1,5 +1,5 @@ use account_compression::{ - assert_mt_zero_copy_inited, + assert_mt_zero_copy_inited, assert_state_mt_roll_over, batched_merkle_tree::{ get_merkle_tree_account_size, AppendBatchProofInputsIx, BatchAppendEvent, BatchNullifyEvent, BatchProofInputsIx, BatchedMerkleTreeAccount, @@ -12,7 +12,7 @@ use account_compression::{ }, get_output_queue_account_default, InitStateTreeAccountsInstructionData, }; -use anchor_lang::AnchorSerialize; +use anchor_lang::{AnchorDeserialize, AnchorSerialize}; use forester_utils::{create_account_instruction, indexer::StateMerkleTreeBundle, AccountZeroCopy}; use light_client::rpc::{RpcConnection, RpcError}; use light_hasher::Poseidon; @@ -32,7 +32,8 @@ use light_registry::{ create_batch_append_instruction, create_batch_nullify_instruction, create_initialize_batched_merkle_tree_instruction, }, - protocol_config::state::ProtocolConfig, + protocol_config::state::{ProtocolConfig, ProtocolConfigPda}, + utils::get_protocol_config_pda_address, }; use light_utils::bigint::bigint_to_be_bytes_array; use light_verifier::CompressedProof; @@ -543,6 +544,7 @@ pub async fn assert_registry_created_batched_state_merkle_tree merkle_tree_pubkey, params.height, params.output_queue_num_batches, + params.network_fee.unwrap_or_default(), ); assert_queue_zero_copy_inited( @@ -552,3 +554,196 @@ pub async fn assert_registry_created_batched_state_merkle_tree ); Ok(()) } +#[allow(clippy::too_many_arguments)] +pub async fn perform_rollover_batch_state_merkle_tree( + rpc: &mut R, + forester: &Keypair, + derivation_pubkey: Pubkey, + old_merkle_tree_pubkey: Pubkey, + old_output_queue_pubkey: Pubkey, + new_state_merkle_tree_keypair: &Keypair, + new_output_queue_keypair: &Keypair, + new_cpi_context_keypair: &Keypair, + epoch: u64, + light_forester: bool, +) -> Result { + let payer_pubkey = forester.pubkey(); + let mut account = rpc.get_account(old_merkle_tree_pubkey).await?.unwrap(); + let old_merkle_tree = + ZeroCopyBatchedMerkleTreeAccount::state_tree_from_bytes_mut(account.data.as_mut_slice()) + .unwrap(); + let batch_zero = &old_merkle_tree.batches[0]; + let num_batches = old_merkle_tree.batches.len(); + let old_merkle_tree = old_merkle_tree.get_account(); + let mt_account_size = get_merkle_tree_account_size( + batch_zero.batch_size, + batch_zero.bloom_filter_capacity, + batch_zero.zkp_batch_size, + old_merkle_tree.root_history_capacity, + old_merkle_tree.height, + num_batches as u64, + ); + + let mt_rent = rpc + .get_minimum_balance_for_rent_exemption(mt_account_size) + .await + .unwrap(); + + let mut account = rpc.get_account(old_output_queue_pubkey).await?.unwrap(); + let old_queue_account = + ZeroCopyBatchedQueueAccount::from_bytes_mut(account.data.as_mut_slice()).unwrap(); + let batch_zero = &old_queue_account.batches[0]; + let queue_account_size = get_output_queue_account_size( + batch_zero.batch_size, + batch_zero.zkp_batch_size, + num_batches as u64, + ); + let queue_rent = rpc + .get_minimum_balance_for_rent_exemption(queue_account_size) + .await + .unwrap(); + + let protocol_config_pubkey = get_protocol_config_pda_address().0; + let protocol_config_account = rpc.get_account(protocol_config_pubkey).await?.unwrap(); + let protocol_config = + ProtocolConfigPda::deserialize(&mut &protocol_config_account.data[8..]).unwrap(); + println!(" fetched protocol_config {:?}", protocol_config); + let create_mt_account_ix = create_account_instruction( + &payer_pubkey, + mt_account_size, + mt_rent, + &account_compression::ID, + Some(new_state_merkle_tree_keypair), + ); + + let create_queue_account_ix = create_account_instruction( + &payer_pubkey, + queue_account_size, + queue_rent, + &account_compression::ID, + Some(new_output_queue_keypair), + ); + let queue_rent = rpc + .get_minimum_balance_for_rent_exemption(protocol_config.config.cpi_context_size as usize) + .await + .unwrap(); + let create_cpi_context_account = create_account_instruction( + &payer_pubkey, + protocol_config.config.cpi_context_size as usize, + queue_rent, + &light_system_program::ID, + Some(new_cpi_context_keypair), + ); + + let instruction = + light_registry::account_compression_cpi::sdk::create_rollover_batch_state_tree_instruction( + forester.pubkey(), + derivation_pubkey, + old_merkle_tree_pubkey, + new_state_merkle_tree_keypair.pubkey(), + old_output_queue_pubkey, + new_output_queue_keypair.pubkey(), + new_cpi_context_keypair.pubkey(), + epoch, + light_forester, + ); + + rpc.create_and_send_transaction( + &[ + create_mt_account_ix, + create_queue_account_ix, + create_cpi_context_account, + instruction, + ], + &payer_pubkey, + &[ + forester, + new_state_merkle_tree_keypair, + new_output_queue_keypair, + new_cpi_context_keypair, + ], + ) + .await +} + +#[allow(clippy::too_many_arguments)] +pub async fn assert_perform_state_mt_roll_over( + rpc: &mut R, + owner: Pubkey, + old_state_merkle_tree_pubkey: Pubkey, + new_state_merkle_tree_pubkey: Pubkey, + old_queue_pubkey: Pubkey, + new_queue_pubkey: Pubkey, + params: InitStateTreeAccountsInstructionData, + additional_bytes_rent: u64, +) { + let old_state_merkle_tree = rpc + .get_account(old_state_merkle_tree_pubkey) + .await + .unwrap() + .unwrap(); + let new_state_merkle_tree = rpc + .get_account(new_state_merkle_tree_pubkey) + .await + .unwrap() + .unwrap(); + let ref_mt_account = BatchedMerkleTreeAccount::get_state_tree_default( + owner, + params.program_owner, + params.forester, + params.rollover_threshold, + params.index, + params.network_fee.unwrap_or_default(), + params.input_queue_batch_size, + params.input_queue_zkp_batch_size, + params.bloom_filter_capacity, + params.root_history_capacity, + old_queue_pubkey, + params.height, + params.input_queue_num_batches, + ); + let old_queue_account_data = rpc + .get_account(old_queue_pubkey) + .await + .unwrap() + .unwrap() + .data; + let new_queue_account = rpc.get_account(new_queue_pubkey).await.unwrap().unwrap(); + + let ref_queue_account = get_output_queue_account_default( + owner, + params.program_owner, + params.forester, + params.rollover_threshold, + params.index, + params.output_queue_batch_size, + params.output_queue_zkp_batch_size, + params.additional_bytes, + new_queue_account.lamports + new_state_merkle_tree.lamports + additional_bytes_rent, //new_cpi_ctx_account.lamports, + old_state_merkle_tree_pubkey, + params.height, + params.output_queue_num_batches, + params.network_fee.unwrap_or_default(), + ); + let mut new_ref_queue_account = ref_queue_account; + new_ref_queue_account.metadata.associated_merkle_tree = new_state_merkle_tree_pubkey; + let mut new_ref_mt_account = ref_mt_account; + new_ref_mt_account.metadata.associated_queue = new_queue_pubkey; + let slot = rpc.get_slot().await.unwrap(); + assert_state_mt_roll_over( + old_state_merkle_tree.data.to_vec(), + new_ref_mt_account, + new_state_merkle_tree.data.to_vec(), + old_state_merkle_tree_pubkey, + new_state_merkle_tree_pubkey, + params.bloom_filter_num_iters, + ref_mt_account, + old_queue_account_data.to_vec(), + new_ref_queue_account, + new_queue_account.data.to_vec(), + new_queue_pubkey, + ref_queue_account, + old_queue_pubkey, + slot, + ); +}