From 22ec436123317fbaa747b1b02639135a9d03a9a0 Mon Sep 17 00:00:00 2001 From: Zelimir Fedoran Date: Thu, 30 Jan 2025 14:55:15 -0500 Subject: [PATCH 1/3] initial airdrop opcode --- api/src/cvm/messages/airdrop.rs | 60 ++++++++++++ api/src/cvm/messages/mod.rs | 2 + api/src/opcode.rs | 37 ++++++++ program/src/instruction/exec.rs | 2 + program/src/opcode/airdrop.rs | 162 ++++++++++++++++++++++++++++++++ program/src/opcode/mod.rs | 2 + program/tests/utils/context.rs | 12 ++- program/tests/utils/state.rs | 12 +++ program/tests/utils/svm.rs | 7 +- program/tests/vm_exec.rs | 77 ++++++++++++--- 10 files changed, 357 insertions(+), 16 deletions(-) create mode 100644 api/src/cvm/messages/airdrop.rs create mode 100644 program/src/opcode/airdrop.rs diff --git a/api/src/cvm/messages/airdrop.rs b/api/src/cvm/messages/airdrop.rs new file mode 100644 index 0000000..656aad5 --- /dev/null +++ b/api/src/cvm/messages/airdrop.rs @@ -0,0 +1,60 @@ +use steel::*; + +use crate::utils; +use crate::types::Hash; +use crate::cvm::{ + CodeVmAccount, + VirtualDurableNonce, + VirtualTimelockAccount +}; + +pub fn compact_airdrop_message( + src_timelock_address: &Pubkey, + dst_timelock_addresses: &[Pubkey], + amount: u64, + vdn: &VirtualDurableNonce, +) -> Hash { + let mut msg = Vec::new(); + + msg.push(b"airdrop" as &[u8]); + msg.push(src_timelock_address.as_ref()); + msg.push(vdn.address.as_ref()); + msg.push(vdn.value.as_ref()); + + // Store the little-endian bytes in a local variable so it won't go out of scope + let amount_bytes = amount.to_le_bytes(); + msg.push(&amount_bytes); + + // Push each destination pubkey + for dst_pubkey in dst_timelock_addresses { + msg.push(dst_pubkey.as_ref()); + } + + utils::hashv(&msg) +} + +pub fn create_airdrop_message( + vm: &CodeVmAccount, + src_vta: &VirtualTimelockAccount, + destinations: &[Pubkey], + amount: u64, + vdn: &VirtualDurableNonce, +) -> Hash { + + let src_timelock_address = src_vta.get_timelock_address( + &vm.get_mint(), + &vm.get_authority(), + vm.get_lock_duration(), + ); + + let src_token_address = src_vta.get_token_address( + &src_timelock_address, + ); + + compact_airdrop_message( + &src_token_address, + destinations, + amount, + vdn, + ) +} \ No newline at end of file diff --git a/api/src/cvm/messages/mod.rs b/api/src/cvm/messages/mod.rs index ae4da06..c462ca7 100644 --- a/api/src/cvm/messages/mod.rs +++ b/api/src/cvm/messages/mod.rs @@ -1,5 +1,7 @@ +mod airdrop; mod transfer; mod withdraw; +pub use airdrop::*; pub use transfer::*; pub use withdraw::*; \ No newline at end of file diff --git a/api/src/opcode.rs b/api/src/opcode.rs index 3c8ee42..c28fb15 100644 --- a/api/src/opcode.rs +++ b/api/src/opcode.rs @@ -15,6 +15,8 @@ pub enum Opcode { ExternalRelayOp = 20, ConditionalTransferOp = 12, + + AirdropOp = 30, } instruction!(Opcode, TransferOp); @@ -24,6 +26,7 @@ instruction!(Opcode, ExternalTransferOp); instruction!(Opcode, ExternalWithdrawOp); instruction!(Opcode, ExternalRelayOp); instruction!(Opcode, ConditionalTransferOp); +instruction!(Opcode, AirdropOp); #[repr(C)] #[derive(Clone, Copy, Debug, Pod, Zeroable)] @@ -210,3 +213,37 @@ pub struct ParsedConditionalTransferOp { pub signature: [u8; 64], pub amount: u64, } + +#[repr(C)] +#[derive(Clone, Copy, Debug, Pod, Zeroable)] +pub struct AirdropOp { + pub signature: [u8; 64], + pub amount: [u8; 8], // Pack u64 as [u8; 8] + pub count: u8, // Up to 255 airdrops in a single tx (but CU limit will be hit first) +} + +impl AirdropOp { + /// Converts the byte array `amount` to `u64`. + pub fn to_struct(&self) -> Result { + Ok(ParsedAirdropOp { + signature: self.signature, + amount: u64::from_le_bytes(self.amount), + count: self.count, + }) + } + + /// Creates `AirdropOp` from the parsed struct by converting `u64` back to byte array. + pub fn from_struct(parsed: ParsedAirdropOp) -> Self { + AirdropOp { + signature: parsed.signature, + amount: parsed.amount.to_le_bytes(), + count: parsed.count, + } + } +} + +pub struct ParsedAirdropOp { + pub signature: [u8; 64], + pub amount: u64, + pub count: u8, +} diff --git a/program/src/instruction/exec.rs b/program/src/instruction/exec.rs index bbcded6..043c30d 100644 --- a/program/src/instruction/exec.rs +++ b/program/src/instruction/exec.rs @@ -69,6 +69,8 @@ pub fn process_exec(accounts: &[AccountInfo<'_>], data: &[u8]) -> ProgramResult Opcode::ConditionalTransferOp => process_conditional_transfer(&ctx, &args), + Opcode::AirdropOp => process_airdrop(&ctx, &args), + _ => Err(ProgramError::InvalidInstructionData), }?; diff --git a/program/src/opcode/airdrop.rs b/program/src/opcode/airdrop.rs new file mode 100644 index 0000000..bc8fe16 --- /dev/null +++ b/program/src/opcode/airdrop.rs @@ -0,0 +1,162 @@ +use code_vm_api::prelude::*; +use steel::*; + +use crate::ExecContext; + +/* + This instruction is used to transfer tokens from *one* virtual account to a + number of virtual accounts. The signature of the source account is required + to authorize the transfer. + + Extra accounts required by this instruction: + + | # | R/W | Type | Req | PDA | Name | Description | + |---|-----|------------- |-----|-----|--------|--------------| + |...| The same as the vm_exec instruction. | + |---|-----|------------- |-----|-----|--------|--------------| + | 6 | | | | | | | + | 7 | | | | | | | + | 8 | | | | | | | + | 9 | | | | | | | + |10 | | | | | | | + + + Instruction data: + + 0. signature: [u8;64] - The opcode to execute. + 1. amount: [u64] - The account_indicies of the virtual accounts to use. + 2. count: [u8] - The number of destinations. +*/ +pub fn process_airdrop( + ctx: &ExecContext, + data: &ExecIxData, +) -> ProgramResult { + + let vm = load_vm(ctx.vm_info)?; + let args = AirdropOp::try_from_bytes(&data.data)?.to_struct()?; + + let mem_indicies = &data.mem_indicies; + let mem_banks = &data.mem_banks; + let num_accounts = 2 + (args.count as usize); + + check_condition( + mem_indicies.len() == num_accounts, + "invalid number of memory indicies", + )?; + + check_condition( + mem_banks.len() == num_accounts, + "invalid number of memory banks", + )?; + + let nonce_index = mem_indicies[0]; + let nonce_mem = mem_banks[0]; + + let src_index = mem_indicies[1]; + let src_mem = mem_banks[1]; + + let vm_mem = ctx.get_banks(); + + check_condition( + vm_mem[nonce_mem as usize].is_some(), + "the nonce memory account must be provided", + )?; + + check_condition( + vm_mem[src_mem as usize].is_some(), + "the source memory account must be provided", + )?; + + let nonce_mem_info = vm_mem[nonce_mem as usize].unwrap(); + let src_mem_info = vm_mem[src_mem as usize].unwrap(); + + let va = try_read(&nonce_mem_info, nonce_index)?; + let mut vdn = va.into_inner_nonce().unwrap(); + + let va = try_read(&src_mem_info, src_index)?; + let mut src_vta = va.into_inner_timelock().unwrap(); + + let total_amount = args.amount + .checked_mul(args.count as u64) + .ok_or(ProgramError::ArithmeticOverflow)?; + + if src_vta.balance < total_amount { + return Err(ProgramError::InsufficientFunds); + } + + src_vta.balance = src_vta.balance + .checked_sub(total_amount) + .ok_or(ProgramError::ArithmeticOverflow)?; + + let mut dst_pubkeys = Vec::new(); + for i in 0..args.count as usize { + let dst_index = mem_indicies[2 + i]; + let dst_mem = mem_banks[2 + i]; + + check_condition( + vm_mem[dst_mem as usize].is_some(), + "a destination memory account must be provided", + )?; + + let dst_mem_info = vm_mem[dst_mem as usize].unwrap(); + + let va = try_read(&dst_mem_info, dst_index)?; + let mut dst_vta = va.into_inner_timelock().unwrap(); + + // Check if this destination is actually the source. + let is_same_account = (src_mem == dst_mem) && (src_index == dst_index); + if is_same_account { + // If the source is also in the destinations list, it receives the airdrop as well. + src_vta.balance = src_vta.balance + .checked_add(args.amount) + .ok_or(ProgramError::ArithmeticOverflow)?; + + } else { + // Normal destination: add the airdrop to its balance + dst_vta.balance = dst_vta.balance + .checked_add(args.amount) + .ok_or(ProgramError::ArithmeticOverflow)?; + + // Write the updated destination back + try_write( + dst_mem_info, + dst_index, + &VirtualAccount::Timelock(dst_vta) + )?; + } + + dst_pubkeys.push(dst_vta.owner); + } + + let hash = create_airdrop_message( + &vm, + &src_vta, + &dst_pubkeys, + args.amount, + &vdn, + ); + + sig_verify( + src_vta.owner.as_ref(), + args.signature.as_ref(), + hash.as_ref(), + )?; + + vdn.value = vm.get_current_poh(); + + // Finally, write back the updated source (which now includes + // any airdrop shares if the source was also in the destination list). + try_write( + src_mem_info, + src_index, + &VirtualAccount::Timelock(src_vta) + )?; + + try_write( + nonce_mem_info, + nonce_index, + &VirtualAccount::Nonce(vdn) + )?; + + Ok(()) +} \ No newline at end of file diff --git a/program/src/opcode/mod.rs b/program/src/opcode/mod.rs index 7d4110c..976471a 100644 --- a/program/src/opcode/mod.rs +++ b/program/src/opcode/mod.rs @@ -1,3 +1,4 @@ +mod airdrop; mod conditional_transfer; mod external_relay; mod external_transfer; @@ -6,6 +7,7 @@ mod relay; mod transfer; mod withdraw; +pub use airdrop::*; pub use conditional_transfer::*; pub use external_relay::*; pub use external_transfer::*; diff --git a/program/tests/utils/context.rs b/program/tests/utils/context.rs index cf62a3c..478ed0a 100644 --- a/program/tests/utils/context.rs +++ b/program/tests/utils/context.rs @@ -16,9 +16,9 @@ pub struct TestContext { } impl TestContext { - pub fn new(param: u8) -> Self { + pub fn new(lock_duration: u8) -> Self { let (svm, payer, mint_owner, mint_pk, vm_address) = - setup_svm_with_payer_and_vm(param); + setup_svm_with_payer_and_vm(lock_duration); let vm = get_vm_account(&svm, vm_address); @@ -231,6 +231,14 @@ impl TestContext { pub fn get_virtual_timelock(&self, mem: Pubkey, index: u16) -> VirtualTimelockAccount { get_virtual_timelock(&self.svm, mem, index) } + + pub fn has_virtual_account(&self, mem: Pubkey, index: u16) -> bool { + has_virtual_account(&self.svm, mem, index) + } + + pub fn get_ata_balance(&self, ata: Pubkey) -> u64 { + get_ata_balance(&self.svm, &ata) + } } pub struct RelayContext { diff --git a/program/tests/utils/state.rs b/program/tests/utils/state.rs index c847fa4..42e3af7 100644 --- a/program/tests/utils/state.rs +++ b/program/tests/utils/state.rs @@ -46,6 +46,18 @@ pub fn get_unlock_state(svm: &LiteSVM, unlock_address: Pubkey) -> UnlockStateAcc UnlockStateAccount::unpack(&account.data) } +pub fn has_virtual_account(svm: &LiteSVM, vm_memory: Pubkey, account_index: u16) -> bool { + let info = svm.get_account(&vm_memory).unwrap(); + let mem_account = MemoryAccount::unpack(&info.data); + let capacity = mem_account.get_capacity(); + let max_item_size = mem_account.get_account_size(); + + let offset = MemoryAccount::get_size(); + let data = info.data.split_at(offset).1; + let mem = SliceAllocator::try_from_slice(data, capacity, max_item_size).unwrap(); + mem.has_item(account_index) +} + pub fn get_virtual_account_data(svm: &LiteSVM, vm_memory: Pubkey, account_index: u16) -> Option> { let info = svm.get_account(&vm_memory).unwrap(); let mem_account = MemoryAccount::unpack(&info.data); diff --git a/program/tests/utils/svm.rs b/program/tests/utils/svm.rs index 76e3c60..bbb0e6c 100644 --- a/program/tests/utils/svm.rs +++ b/program/tests/utils/svm.rs @@ -3,7 +3,7 @@ use std::path::PathBuf; use code_vm_api::prelude::CodeInstruction; use solana_sdk::{pubkey::Pubkey, signature::Keypair, signer::Signer, transaction::Transaction}; use litesvm::{types::{FailedTransactionMetadata, TransactionMetadata, TransactionResult}, LiteSVM}; -use litesvm_token::{CreateAssociatedTokenAccount, CreateMint, MintTo}; +use litesvm_token::{CreateAssociatedTokenAccount, CreateMint, MintTo, spl_token::{state::Account}, get_spl_account}; use pretty_hex::*; pub fn program_bytes() -> Vec { @@ -60,6 +60,11 @@ pub fn create_ata(svm: &mut LiteSVM, payer_kp: &Keypair, mint_pk: &Pubkey, owner .unwrap() } +pub fn get_ata_balance(svm: &LiteSVM, ata: &Pubkey) -> u64 { + let info:Account = get_spl_account(svm, &ata).unwrap(); + info.amount +} + pub fn mint_to(svm: &mut LiteSVM, payer: &Keypair, mint: &Pubkey, diff --git a/program/tests/vm_exec.rs b/program/tests/vm_exec.rs index 5d131f8..373bb2b 100644 --- a/program/tests/vm_exec.rs +++ b/program/tests/vm_exec.rs @@ -21,8 +21,15 @@ fn run_transfer() { // Create durable nonce account let vdn_ctx = ctx.create_durable_nonce_account(mem_a, 0); + // -- 1) Deposit tokens into `vta_a_ctx` so we have something to send + let deposit_amount = 100; + ctx.deposit_tokens_to_timelock(mem_b, &vta_a_ctx, deposit_amount) + .unwrap(); + + // -- 2) Set a non-zero transfer amount + let amount = 42; + // Create the transfer message and signature - let amount = 0; // Sending 0 tokens to keep the test simple let hash = create_transfer_message( &ctx.vm, &vta_a_ctx.account, @@ -40,11 +47,9 @@ fn run_transfer() { // Prepare the opcode data let mem_indices = vec![vdn_ctx.index, vta_a_ctx.index, vta_b_ctx.index]; let mem_banks = vec![0, 1, 1]; - let data = TransferOp::from_struct( - ParsedTransferOp { amount, signature } - ).to_bytes(); + let data = TransferOp::from_struct(ParsedTransferOp { amount, signature }).to_bytes(); - // Execute the opcode + // -- 3) Execute the transfer opcode ctx.exec_opcode( [Some(mem_a), Some(mem_b), None, None], None, // vm_omnibus @@ -57,6 +62,12 @@ fn run_transfer() { mem_banks, ) .unwrap(); + + // -- 4) Verify final balances + let src_vta = ctx.get_virtual_timelock(mem_b, vta_a_ctx.index); + let dst_vta = ctx.get_virtual_timelock(mem_b, vta_b_ctx.index); + assert_eq!(src_vta.balance, deposit_amount - amount); + assert_eq!(dst_vta.balance, amount); } #[test] @@ -75,10 +86,17 @@ fn run_transfer_to_external() { let vdn_ctx = ctx.create_durable_nonce_account(mem_a, 0); // Prepare the destination pubkey - let dst_pubkey = ctx.vm.omnibus.vault; // Self send (to keep things simple) + let dst_pubkey = ctx.vm.omnibus.vault; // e.g. the VM's omnibus vault + + // -- 1) Deposit tokens into `vta_a_ctx` + let deposit_amount = 50; + ctx.deposit_tokens_to_timelock(mem_b, &vta_a_ctx, deposit_amount) + .unwrap(); + + // -- 2) Set a non-zero transfer amount + let amount = 10; // Create the transfer message and signature - let amount = 0; // Sending 0 tokens to keep the test simple let hash = create_transfer_message_to_external( &ctx.vm, &vta_a_ctx.account, @@ -96,11 +114,11 @@ fn run_transfer_to_external() { // Prepare the opcode data let mem_indices = vec![vdn_ctx.index, vta_a_ctx.index]; let mem_banks = vec![0, 1]; - let data = ExternalTransferOp::from_struct( + let data = ExternalTransferOp::from_struct( ParsedExternalTransferOp { amount, signature } ).to_bytes(); - // Execute the opcode + // -- 3) Execute the transfer opcode ctx.exec_opcode( [Some(mem_a), Some(mem_b), None, None], Some(ctx.vm.omnibus.vault), // vm_omnibus @@ -113,6 +131,13 @@ fn run_transfer_to_external() { mem_banks, ) .unwrap(); + + // -- 4) Verify final balance in the timelock + let src_vta = ctx.get_virtual_timelock(mem_b, vta_a_ctx.index); + assert_eq!(src_vta.balance, deposit_amount - amount); + + // Optionally, if you want to verify the omnibus vault gained tokens, + // you'd look up the vault’s balance in the test context (implementation dependent). } #[test] @@ -131,7 +156,12 @@ fn run_withdraw() { // Create durable nonce account let vdn_ctx = ctx.create_durable_nonce_account(mem_a, 0); - // Create the withdraw message and signature + // -- 1) Deposit tokens into vta_a + let deposit_amount = 100; + ctx.deposit_tokens_to_timelock(mem_b, &vta_a_ctx, deposit_amount) + .unwrap(); + + // Create the transfer message and signature let hash = create_withdraw_message( &ctx.vm, &vta_a_ctx.account, @@ -150,7 +180,7 @@ fn run_withdraw() { let mem_banks = vec![0, 1, 1]; let data = WithdrawOp { signature }.to_bytes(); - // Execute the opcode + // -- 2) Execute the withdraw opcode ctx.exec_opcode( [Some(mem_a), Some(mem_b), None, None], None, // vm_omnibus @@ -163,6 +193,14 @@ fn run_withdraw() { mem_banks, ) .unwrap(); + + // -- 3) Verify final balances + let dst_vta = ctx.get_virtual_timelock(mem_b, vta_b_ctx.index); + assert_eq!(dst_vta.balance, deposit_amount); + + // We expect the source account to be deleted after the withdraw + let src_exists = ctx.has_virtual_account(mem_b, vta_a_ctx.index); + assert_eq!(src_exists, false); } #[test] @@ -181,7 +219,12 @@ fn run_withdraw_to_external() { let vdn_ctx = ctx.create_durable_nonce_account(mem_a, 0); // Prepare the destination pubkey - let dst_pubkey = ctx.vm.omnibus.vault; // Self send + let dst_pubkey = ctx.vm.omnibus.vault; // e.g. the VM's omnibus vault + + // -- 1) Deposit tokens into `vta_a_ctx` + let deposit_amount = 100; + ctx.deposit_tokens_to_timelock(mem_b, &vta_a_ctx, deposit_amount) + .unwrap(); // Create the withdraw message and signature let hash = create_withdraw_message_to_external( @@ -202,7 +245,7 @@ fn run_withdraw_to_external() { let mem_banks = vec![0, 1]; let data = ExternalWithdrawOp { signature }.to_bytes(); - // Execute the opcode + // -- 2) Execute the withdraw-to-external opcode ctx.exec_opcode( [Some(mem_a), Some(mem_b), None, None], Some(ctx.vm.omnibus.vault), // vm_omnibus @@ -215,4 +258,12 @@ fn run_withdraw_to_external() { mem_banks, ) .unwrap(); + + // -- 3) Verify final balances + + let src_exists = ctx.has_virtual_account(mem_b, vta_a_ctx.index); + assert_eq!(src_exists, false); + + let dst_balance = ctx.get_ata_balance(dst_pubkey); + assert_eq!(dst_balance, deposit_amount); } From 78a936e2d17dad253875bd7c1bd64bfc988fb8ca Mon Sep 17 00:00:00 2001 From: Zelimir Fedoran Date: Thu, 30 Jan 2025 16:57:19 -0500 Subject: [PATCH 2/3] added airdrop tests --- program/tests/vm_exec_airdrop.rs | 292 +++++++++++++++++++++++++++++++ 1 file changed, 292 insertions(+) create mode 100644 program/tests/vm_exec_airdrop.rs diff --git a/program/tests/vm_exec_airdrop.rs b/program/tests/vm_exec_airdrop.rs new file mode 100644 index 0000000..8c682a2 --- /dev/null +++ b/program/tests/vm_exec_airdrop.rs @@ -0,0 +1,292 @@ +#![cfg(test)] +pub mod utils; +use utils::*; + +use solana_sdk::signature::Signer; +use code_vm_api::prelude::*; + + +#[test] +fn run_airdrop_10() { + run_airdrop_test(10, 100); +} + +#[test] +fn run_airdrop_25() { + run_airdrop_test(25, 100); +} + +#[test] +fn run_airdrop_50() { + run_airdrop_test(50, 100); +} + +#[test] +fn run_airdrop_0() { + run_airdrop_test(0, 0); + run_airdrop_test(0, 100); + run_airdrop_test(10, 0); +} + +#[test] +fn run_airdrop_include_self() { + run_airdrop_with_self(0, 100); + run_airdrop_with_self(1, 100); + run_airdrop_with_self(10, 100); +} + +#[test] +fn run_airdrop_only_to_self() { + run_self_edgecase(0, 100); + run_self_edgecase(10, 100); +} + +/// Runs an airdrop test with the specified number of destination accounts. +/// Each destination receives 100 tokens from a single source timelock. +fn run_airdrop_test(count: usize, amount_each: u64) { + let mut ctx = TestContext::new(21); + + let mem_a = ctx.create_memory(10, VirtualDurableNonce::LEN + 1, "mem_nonce_0"); + let mem_b = ctx.create_memory(count+2, VirtualTimelockAccount::LEN + 1, "mem_timelock_0"); + + let vta_source = ctx.create_timelock_account(mem_b, 0); // move occurs because vta_source has type TimelockAccountContext, which does not implement the Copy + + let mut destinations = Vec::with_capacity(count); + for i in 1..=count { + let dst_vta = ctx.create_timelock_account(mem_b, i as u16); + destinations.push(dst_vta); + } + + let vdn_ctx = ctx.create_durable_nonce_account(mem_a, 0); + + let total_outflow = amount_each.checked_mul(count as u64).unwrap(); + let deposit_amount = total_outflow; + ctx.deposit_tokens_to_timelock(mem_b, &vta_source, deposit_amount) + .unwrap(); + + let dst_pubkeys: Vec<_> = destinations + .iter() + .map(|dst_vta| dst_vta.account.owner) + .collect(); + + let hash = create_airdrop_message( + &ctx.vm, + &vta_source.account, + &dst_pubkeys, + amount_each, + &vdn_ctx.account, + ); + + let sig = vta_source + .key + .sign_message(hash.as_ref()) + .as_ref() + .try_into() + .unwrap(); + + let data = AirdropOp::from_struct(ParsedAirdropOp { + signature: sig, + amount: amount_each, + count: count as u8 + }).to_bytes(); + + let mut mem_indices = vec![vdn_ctx.index, vta_source.index]; + mem_indices.extend(destinations.iter().map(|d| d.index)); + + let mut mem_banks = vec![0, 1]; // 0 for mem_a (nonce), 1 for mem_b (source/dest) + mem_banks.extend(std::iter::repeat(1).take(count)); + + ctx.exec_opcode( + [Some(mem_a), Some(mem_b), None, None], + None, // vm_omnibus + None, // relay + None, // relay_vault + None, // external_address + None, // token_program + data, + mem_indices, + mem_banks, + ) + .unwrap(); + + let src_after = ctx.get_virtual_timelock(mem_b, vta_source.index); + assert_eq!(src_after.balance, deposit_amount - total_outflow); + + for (i, dst) in destinations.iter().enumerate() { + let dst_balance = ctx.get_virtual_timelock(mem_b, dst.index).balance; + assert_eq!( + dst_balance, + amount_each, + "Destination #{} did not receive 100 tokens", + i + ); + } +} + +// Same as run_airdrop_test, but includes the source account in the destination list. +fn run_airdrop_with_self(count: usize, amount_each: u64) { + let mut ctx = TestContext::new(21); + + let final_count = count + 1; + let mem_a = ctx.create_memory(10, VirtualDurableNonce::LEN + 1, "mem_nonce_0"); + let mem_b = ctx.create_memory(count + 2, VirtualTimelockAccount::LEN + 1, "mem_timelock_0"); + + let vta_source = ctx.create_timelock_account(mem_b, 0); + + let mut destinations = Vec::with_capacity(count); + for i in 1..=count { + let dst_vta = ctx.create_timelock_account(mem_b, i as u16); + destinations.push(dst_vta); + } + + let vdn_ctx = ctx.create_durable_nonce_account(mem_a, 0); + let total_outflow = amount_each.checked_mul(final_count as u64).unwrap(); + + let deposit_amount = total_outflow; + ctx.deposit_tokens_to_timelock(mem_b, &vta_source, deposit_amount) + .unwrap(); + + let mut dst_pubkeys: Vec<_> = destinations + .iter() + .map(|d| d.account.owner) + .collect(); + + dst_pubkeys.push(vta_source.account.owner); + + let hash = create_airdrop_message( + &ctx.vm, + &vta_source.account, + &dst_pubkeys, + amount_each, + &vdn_ctx.account, + ); + + let sig = vta_source + .key + .sign_message(hash.as_ref()) + .as_ref() + .try_into() + .unwrap(); + + destinations.push(vta_source); + + let data = AirdropOp::from_struct(ParsedAirdropOp { + signature: sig, + amount: amount_each, + count: final_count as u8, + }) + .to_bytes(); + + let mut mem_indices = vec![vdn_ctx.index, 0]; // 0 for source + let mut mem_banks = vec![0, 1]; // 0 for mem_a, 1 for mem_b + + // For all destinations, push their memory index and bank=1 + for dst in &destinations { + mem_indices.push(dst.index); + mem_banks.push(1); + } + + ctx.exec_opcode( + [Some(mem_a), Some(mem_b), None, None], + None, // vm_omnibus + None, // relay + None, // relay_vault + None, // external_address + None, // token_program + data, + mem_indices, + mem_banks, + ) + .unwrap(); + + let src_after = ctx.get_virtual_timelock(mem_b, 0); + assert_eq!(src_after.balance, deposit_amount - total_outflow + amount_each); + + for (i, dst) in destinations.iter().enumerate() { + let dst_vta = ctx.get_virtual_timelock(mem_b, dst.index); + let dst_balance = dst_vta.balance; + + if i == count { + assert_eq!(dst_balance, deposit_amount - total_outflow + amount_each); + } else { + assert_eq!( + dst_balance, + amount_each, + "Destination #{} did not receive {} tokens", + i, + amount_each + ); + } + } +} + +/// Airdrops exclusively to the same source `count` times, each worth `amount_each`. +/// (This is an edge case, we're testing that the source account doesn't get more +/// or less tokens than what it started with) +fn run_self_edgecase(count: usize, amount_each: u64) { + let mut ctx = TestContext::new(21); + + let mem_a = ctx.create_memory(1, VirtualDurableNonce::LEN + 1, "mem_nonce_0"); + let mem_b = ctx.create_memory(1, VirtualTimelockAccount::LEN + 1, "mem_timelock_0"); + + let vta_source = ctx.create_timelock_account(mem_b, 0); + + let vdn_ctx = ctx.create_durable_nonce_account(mem_a, 0); + + let total_outflow = amount_each + .checked_mul(count as u64) + .expect("overflow computing total_outflow"); + + let deposit_amount = total_outflow; + ctx.deposit_tokens_to_timelock(mem_b, &vta_source, deposit_amount) + .unwrap(); + + let dst_pubkeys = vec![vta_source.account.owner; count]; + + let hash = create_airdrop_message( + &ctx.vm, + &vta_source.account, + &dst_pubkeys, + amount_each, + &vdn_ctx.account, + ); + let signature = vta_source + .key + .sign_message(hash.as_ref()) + .as_ref() + .try_into() + .unwrap(); + + let data = AirdropOp::from_struct(ParsedAirdropOp { + signature, + amount: amount_each, + count: count as u8, + }) + .to_bytes(); + + let mut mem_indices = vec![vdn_ctx.index, vta_source.index]; + mem_indices.extend(std::iter::repeat(vta_source.index).take(count)); + + let mut mem_banks = vec![0, 1]; + mem_banks.extend(std::iter::repeat(1).take(count)); + + ctx.exec_opcode( + [Some(mem_a), Some(mem_b), None, None], + None, // vm_omnibus + None, // relay + None, // relay_vault + None, // external_address + None, // token_program + data, + mem_indices, + mem_banks, + ) + .unwrap(); + + let src_after = ctx.get_virtual_timelock(mem_b, vta_source.index); + assert_eq!( + src_after.balance, + deposit_amount, + "Source final balance mismatch after repeated self-airdrop" + ); +} \ No newline at end of file From 7c58accead9090271c1e4e6e6f85e1a504b1240d Mon Sep 17 00:00:00 2001 From: Zelimir Fedoran Date: Fri, 31 Jan 2025 13:07:51 -0500 Subject: [PATCH 3/3] added chunked test --- program/tests/utils/context.rs | 52 +++++++++++- program/tests/vm_exec_airdrop.rs | 136 +++++++++++++++++++++++++++++++ 2 files changed, 187 insertions(+), 1 deletion(-) diff --git a/program/tests/utils/context.rs b/program/tests/utils/context.rs index 478ed0a..a5f2ca7 100644 --- a/program/tests/utils/context.rs +++ b/program/tests/utils/context.rs @@ -3,7 +3,7 @@ use super::*; use steel::*; use litesvm::{types::TransactionResult, LiteSVM}; -use solana_sdk::signature::{Keypair, Signer}; +use solana_sdk::{signature::{Keypair, Signer }, transaction::Transaction}; use code_vm_api::prelude::*; pub struct TestContext { @@ -228,6 +228,56 @@ impl TestContext { ) } + pub fn ix_send( + &mut self, + ix: &[Instruction], + ) -> TransactionResult { + let payer_pk = self.payer.pubkey(); + let blockhash = self.svm.latest_blockhash(); + let tx = Transaction::new_signed_with_payer( + ix, + Some(&payer_pk), + &[&self.payer], + blockhash + ); + + send_tx(&mut self.svm, tx) + } + + pub fn get_exec_ix( + &mut self, + mems: [Option; 4], + vm_omnibus: Option, + relay: Option, + relay_vault: Option, + external_address: Option, + token_program: Option, + data: Vec, + mem_indices: Vec, + mem_banks: Vec, + ) -> Instruction { + let opcode = data[0]; + let data = data[1..].to_vec(); + + vm_exec( + self.payer.pubkey(), + self.vm_address, + mems[0], + mems[1], + mems[2], + mems[3], + vm_omnibus, + relay, + relay_vault, + external_address, + token_program, + opcode, + mem_indices, + mem_banks, + data, + ) + } + pub fn get_virtual_timelock(&self, mem: Pubkey, index: u16) -> VirtualTimelockAccount { get_virtual_timelock(&self.svm, mem, index) } diff --git a/program/tests/vm_exec_airdrop.rs b/program/tests/vm_exec_airdrop.rs index 8c682a2..92cc7a7 100644 --- a/program/tests/vm_exec_airdrop.rs +++ b/program/tests/vm_exec_airdrop.rs @@ -1,4 +1,7 @@ #![cfg(test)] + +use steel::*; + pub mod utils; use utils::*; @@ -21,6 +24,11 @@ fn run_airdrop_50() { run_airdrop_test(50, 100); } +#[test] +fn run_airdrop_chunked() { + run_airdrop_test_chunked(200, 100); +} + #[test] fn run_airdrop_0() { run_airdrop_test(0, 0); @@ -289,4 +297,132 @@ fn run_self_edgecase(count: usize, amount_each: u64) { deposit_amount, "Source final balance mismatch after repeated self-airdrop" ); +} + +/// Runs an airdrop test with the specified number of destination accounts, +/// grouping them in chunks of up to 50. Each destination receives `amount_each` +/// tokens from a single source account. +fn run_airdrop_test_chunked(count: usize, amount_each: u64) { + let mut ctx = TestContext::new(21); + + let mem_a = ctx.create_memory(10, VirtualDurableNonce::LEN + 1, "mem_nonce_0"); + let mem_b = ctx.create_memory(count+2, VirtualTimelockAccount::LEN + 1, "mem_timelock_0"); + + let vta_source = ctx.create_timelock_account(mem_b, 0); + + let mut destinations = Vec::with_capacity(count); + for i in 1..=count { + let dst_vta = ctx.create_timelock_account(mem_b, i as u16); + destinations.push(dst_vta); + } + + let total_outflow = amount_each.checked_mul(count as u64).unwrap(); + ctx.deposit_tokens_to_timelock(mem_b, &vta_source, total_outflow) + .unwrap(); + + let chunk_size = 50; + let mut instructions = Vec::new(); + + let num_chunks = (count + chunk_size - 1) / chunk_size; + + for i in 0..num_chunks { + let start = i * chunk_size; + let end = std::cmp::min(start + chunk_size, count); + let chunk = &destinations[start..end]; + + // Create a fresh durable nonce account for this chunk (in prod, these + // would be pre-created and reused) + let vdn_ctx = ctx.create_durable_nonce_account(mem_a, i as u16); + + // Build the instruction for this chunk + let ix = create_airdrop_ix( + &mut ctx, + mem_a, // The memory account for the nonce + mem_b, // The memory account for source/dest + &vdn_ctx, // This chunk's nonce + &vta_source, // Source Timelock + chunk, // slice of up to 50 + amount_each, + ); + instructions.push(ix); + } + + ctx.ix_send(&instructions).unwrap(); + + let src_after = ctx.get_virtual_timelock(mem_b, vta_source.index); + assert_eq!( + src_after.balance, + total_outflow - total_outflow, + "Source did not properly deduct the outflow" + ); + + for (i, dst) in destinations.iter().enumerate() { + let dst_balance = ctx.get_virtual_timelock(mem_b, dst.index).balance; + assert_eq!( + dst_balance, + amount_each, + "Destination #{} did not receive the expected tokens", + i + ); + } +} + +/// Creates a single `Instruction` for a batch (chunk) of destinations. +/// This instructs the VM to execute the bulk airdrop for all `destinations`. +fn create_airdrop_ix( + ctx: &mut TestContext, + mem_a_key: Pubkey, // The memory account for the DurableNonce + mem_b_key: Pubkey, // The memory account for Source & Dest accounts + vdn_ctx: &DurableNonceContext, // Contains .index and .account (nonce data) + vta_source: &TimelockAccountContext, // Source Timelock (has .index, .account, .key) + destinations: &[TimelockAccountContext], // A slice of up to ~50 destinations + amount_each: u64, +) -> Instruction +{ + let dst_pubkeys: Vec<_> = destinations + .iter() + .map(|dst| dst.account.owner) // or whatever pubkey you want to reference + .collect(); + + let hash = create_airdrop_message( + &ctx.vm, + &vta_source.account, + &dst_pubkeys, + amount_each, + &vdn_ctx.account, + ); + + let sig = vta_source + .key + .sign_message(hash.as_ref()) + .as_ref() + .try_into() + .unwrap(); + + let data = AirdropOp::from_struct(ParsedAirdropOp { + signature: sig, + amount: amount_each, + count: destinations.len() as u8, + }) + .to_bytes(); + + let mut mem_indices = vec![vdn_ctx.index, vta_source.index]; + for d in destinations { + mem_indices.push(d.index); + } + + let mut mem_banks = vec![0u8, 1u8]; // 0 => mem_a, 1 => mem_b + mem_banks.extend(std::iter::repeat(1u8).take(destinations.len())); + + ctx.get_exec_ix( + [Some(mem_a_key), Some(mem_b_key), None, None], // up to 4 memory accounts + None, // vm_omnibus + None, // relay + None, // relay_vault + None, // external_address + None, // token_program + data, + mem_indices, + mem_banks, + ) } \ No newline at end of file