From 6ff2572ebe44b4c7f73c5f98ee7cc376fef93f39 Mon Sep 17 00:00:00 2001 From: Tyera Eulberg Date: Tue, 5 Mar 2019 12:11:05 -0700 Subject: [PATCH] Refactor system entrypoint to use helper fns; add unit tests --- programs/system/src/lib.rs | 258 +++++++++++++++++++++++++++++-------- 1 file changed, 201 insertions(+), 57 deletions(-) diff --git a/programs/system/src/lib.rs b/programs/system/src/lib.rs index d130c1c9a..cf2902e2b 100644 --- a/programs/system/src/lib.rs +++ b/programs/system/src/lib.rs @@ -6,6 +6,74 @@ use solana_sdk::solana_entrypoint; use solana_sdk::system_instruction::SystemInstruction; use solana_sdk::system_program; +const FROM_ACCOUNT_INDEX: usize = 0; +const TO_ACCOUNT_INDEX: usize = 1; + +#[derive(Debug, Clone, PartialEq)] +enum SystemError { + AccountAlreadyInUse, + ResultWithNegativeTokens, + SourceNotSystemAccount, +} + +fn create_system_account( + keyed_accounts: &mut [KeyedAccount], + tokens: u64, + space: u64, + program_id: Pubkey, +) -> Result<(), SystemError> { + if !system_program::check_id(&keyed_accounts[FROM_ACCOUNT_INDEX].account.owner) { + info!("CreateAccount: invalid account[from] owner"); + Err(SystemError::SourceNotSystemAccount)?; + } + + if !keyed_accounts[TO_ACCOUNT_INDEX].account.userdata.is_empty() + || !system_program::check_id(&keyed_accounts[TO_ACCOUNT_INDEX].account.owner) + { + info!( + "CreateAccount: invalid argument; account {} already in use", + keyed_accounts[TO_ACCOUNT_INDEX].unsigned_key() + ); + Err(SystemError::AccountAlreadyInUse)?; + } + if tokens > keyed_accounts[FROM_ACCOUNT_INDEX].account.tokens { + info!( + "CreateAccount: insufficient tokens ({}, need {})", + keyed_accounts[FROM_ACCOUNT_INDEX].account.tokens, tokens + ); + Err(SystemError::ResultWithNegativeTokens)?; + } + keyed_accounts[FROM_ACCOUNT_INDEX].account.tokens -= tokens; + keyed_accounts[TO_ACCOUNT_INDEX].account.tokens += tokens; + keyed_accounts[TO_ACCOUNT_INDEX].account.owner = program_id; + keyed_accounts[TO_ACCOUNT_INDEX].account.userdata = vec![0; space as usize]; + keyed_accounts[TO_ACCOUNT_INDEX].account.executable = false; + Ok(()) +} + +fn assign_account_to_program( + keyed_accounts: &mut [KeyedAccount], + program_id: Pubkey, +) -> Result<(), ProgramError> { + if !system_program::check_id(&keyed_accounts[FROM_ACCOUNT_INDEX].account.owner) { + Err(ProgramError::AssignOfUnownedAccount)?; + } + keyed_accounts[FROM_ACCOUNT_INDEX].account.owner = program_id; + Ok(()) +} +fn move_tokens(keyed_accounts: &mut [KeyedAccount], tokens: u64) -> Result<(), ProgramError> { + if tokens > keyed_accounts[FROM_ACCOUNT_INDEX].account.tokens { + info!( + "Move: insufficient tokens ({}, need {})", + keyed_accounts[FROM_ACCOUNT_INDEX].account.tokens, tokens + ); + Err(ProgramError::ResultWithNegativeTokens)?; + } + keyed_accounts[FROM_ACCOUNT_INDEX].account.tokens -= tokens; + keyed_accounts[TO_ACCOUNT_INDEX].account.tokens += tokens; + Ok(()) +} + solana_entrypoint!(entrypoint); pub fn entrypoint( _program_id: &Pubkey, @@ -16,10 +84,9 @@ pub fn entrypoint( if let Ok(syscall) = bincode::deserialize(data) { trace!("process_instruction: {:?}", syscall); trace!("keyed_accounts: {:?}", keyed_accounts); - let from = 0; // all system instructions require that accounts_keys[0] be a signer - if keyed_accounts[from].signer_key().is_none() { + if keyed_accounts[FROM_ACCOUNT_INDEX].signer_key().is_none() { info!("account[from] is unsigned"); Err(ProgramError::InvalidArgument)?; } @@ -30,66 +97,143 @@ pub fn entrypoint( space, program_id, } => { - let to = 1; - if !system_program::check_id(&keyed_accounts[from].account.owner) { - info!("CreateAccount: invalid account[from] owner"); - Err(ProgramError::InvalidArgument)?; - } - - if !keyed_accounts[to].account.userdata.is_empty() - || !system_program::check_id(&keyed_accounts[to].account.owner) - { - if space > 0 { - info!( - "CreateAccount: invalid argument space: {} accounts.userdata.len(): {}", - space, - keyed_accounts[to].account.userdata.len(), - ); - } else { - info!( - "CreateAccount: invalid argument; account {} already in use", - keyed_accounts[to].unsigned_key() - ); - } - Err(ProgramError::InvalidArgument)?; - } - if tokens > keyed_accounts[from].account.tokens { - info!( - "CreateAccount: insufficient tokens ({}, need {})", - keyed_accounts[from].account.tokens, tokens - ); - Err(ProgramError::ResultWithNegativeTokens)?; - } - keyed_accounts[from].account.tokens -= tokens; - keyed_accounts[to].account.tokens += tokens; - keyed_accounts[to].account.owner = program_id; - keyed_accounts[to].account.userdata = vec![0; space as usize]; - keyed_accounts[to].account.executable = false; + create_system_account(keyed_accounts, tokens, space, program_id).map_err( + |e| match e { + SystemError::ResultWithNegativeTokens => { + ProgramError::ResultWithNegativeTokens + } + _ => ProgramError::InvalidArgument, + }, + ) } SystemInstruction::Assign { program_id } => { - if !system_program::check_id(&keyed_accounts[from].account.owner) { - Err(ProgramError::AssignOfUnownedAccount)?; - } - keyed_accounts[from].account.owner = program_id; - } - SystemInstruction::Move { tokens } => { - let to = 1; - - // bank should be verifying correctness - if tokens > keyed_accounts[from].account.tokens { - info!( - "Move: insufficient tokens ({}, need {})", - keyed_accounts[from].account.tokens, tokens - ); - Err(ProgramError::ResultWithNegativeTokens)?; - } - keyed_accounts[from].account.tokens -= tokens; - keyed_accounts[to].account.tokens += tokens; + assign_account_to_program(keyed_accounts, program_id) } + SystemInstruction::Move { tokens } => move_tokens(keyed_accounts, tokens), } - Ok(()) } else { info!("Invalid transaction instruction userdata: {:?}", data); - Err(ProgramError::InvalidArgument) + Err(ProgramError::InvalidUserdata) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use solana_sdk::account::Account; + use solana_sdk::signature::{Keypair, KeypairUtil}; + + #[test] + fn test_create_system_account() { + let new_program_owner = Pubkey::new(&[9; 32]); + + let from = Keypair::new().pubkey(); + let mut from_account = Account::new(100, 0, system_program::id()); + let to = Keypair::new().pubkey(); + let mut to_account = Account::new(0, 0, Pubkey::default()); + let mut keyed_accounts = [ + KeyedAccount::new(&from, true, &mut from_account), + KeyedAccount::new(&to, false, &mut to_account), + ]; + create_system_account(&mut keyed_accounts, 50, 2, new_program_owner).unwrap(); + let from_tokens = from_account.tokens; + let to_tokens = to_account.tokens; + let to_owner = to_account.owner; + let to_userdata = to_account.userdata.clone(); + assert_eq!(from_tokens, 50); + assert_eq!(to_tokens, 50); + assert_eq!(to_owner, new_program_owner); + assert_eq!(to_userdata, [0, 0]); + + // Attempt to move more tokens than remaining in from_account + let unchanged_account = to_account.clone(); + keyed_accounts = [ + KeyedAccount::new(&from, true, &mut from_account), + KeyedAccount::new(&to, false, &mut to_account), + ]; + let result = move_tokens(&mut keyed_accounts, 100); + assert!(result.is_err()); + let from_tokens = from_account.tokens; + assert_eq!(from_tokens, 50); + assert_eq!(to_account, unchanged_account); + + // Attempt to create system account in account already owned by another program + let original_program_owner = Pubkey::new(&[5; 32]); + let owned_key = Keypair::new().pubkey(); + let mut owned_account = Account::new(0, 0, original_program_owner); + let unchanged_account = owned_account.clone(); + keyed_accounts = [ + KeyedAccount::new(&from, true, &mut from_account), + KeyedAccount::new(&owned_key, false, &mut owned_account), + ]; + let result = create_system_account(&mut keyed_accounts, 50, 2, new_program_owner); + assert!(result.is_err()); + let from_tokens = from_account.tokens; + assert_eq!(from_tokens, 50); + assert_eq!(owned_account, unchanged_account); + + // Attempt to create system account in account with populated userdata + let populated_key = Keypair::new().pubkey(); + let mut populated_account = Account { + tokens: 0, + userdata: vec![0, 1, 2, 3], + owner: Pubkey::default(), + executable: false, + }; + let unchanged_account = populated_account.clone(); + keyed_accounts = [ + KeyedAccount::new(&from, true, &mut from_account), + KeyedAccount::new(&populated_key, false, &mut populated_account), + ]; + let result = create_system_account(&mut keyed_accounts, 50, 2, new_program_owner); + assert!(result.is_err()); + assert_eq!(from_account.tokens, 50); + assert_eq!(populated_account, unchanged_account); + } + + #[test] + fn test_assign_account_to_program() { + let new_program_owner = Pubkey::new(&[9; 32]); + + let from = Keypair::new().pubkey(); + let mut from_account = Account::new(100, 0, system_program::id()); + let mut keyed_accounts = [KeyedAccount::new(&from, true, &mut from_account)]; + assign_account_to_program(&mut keyed_accounts, new_program_owner).unwrap(); + let from_owner = from_account.owner; + assert_eq!(from_owner, new_program_owner); + + // Attempt to assign account not owned by system program + let another_program_owner = Pubkey::new(&[8; 32]); + keyed_accounts = [KeyedAccount::new(&from, true, &mut from_account)]; + let result = assign_account_to_program(&mut keyed_accounts, another_program_owner); + assert!(result.is_err()); + assert_eq!(from_account.owner, new_program_owner); + } + + #[test] + fn test_move_tokens() { + let from = Keypair::new().pubkey(); + let mut from_account = Account::new(100, 0, Pubkey::new(&[2; 32])); // account owner should not matter + let to = Keypair::new().pubkey(); + let mut to_account = Account::new(1, 0, Pubkey::new(&[3; 32])); // account owner should not matter + let mut keyed_accounts = [ + KeyedAccount::new(&from, true, &mut from_account), + KeyedAccount::new(&to, false, &mut to_account), + ]; + move_tokens(&mut keyed_accounts, 50).unwrap(); + let from_tokens = from_account.tokens; + let to_tokens = to_account.tokens; + assert_eq!(from_tokens, 50); + assert_eq!(to_tokens, 51); + + // Attempt to move more tokens than remaining in from_account + keyed_accounts = [ + KeyedAccount::new(&from, true, &mut from_account), + KeyedAccount::new(&to, false, &mut to_account), + ]; + let result = move_tokens(&mut keyed_accounts, 100); + assert!(result.is_err()); + assert_eq!(from_account.tokens, 50); + assert_eq!(to_account.tokens, 51); } }