diff --git a/programs/budget_api/src/budget_transaction.rs b/programs/budget_api/src/budget_transaction.rs index 3162ef700c..1c53247191 100644 --- a/programs/budget_api/src/budget_transaction.rs +++ b/programs/budget_api/src/budget_transaction.rs @@ -2,13 +2,11 @@ use crate::budget_instruction::BudgetInstruction; use crate::budget_script::BudgetScript; -use bincode::deserialize; use chrono::prelude::*; use solana_sdk::hash::Hash; use solana_sdk::pubkey::Pubkey; use solana_sdk::script::Script; use solana_sdk::signature::{Keypair, KeypairUtil}; -use solana_sdk::system_instruction::SystemInstruction; use solana_sdk::transaction::Transaction; pub struct BudgetTransaction {} @@ -110,22 +108,30 @@ impl BudgetTransaction { ); Self::new_signed(from_keypair, script, recent_blockhash, 0) } +} - pub fn system_instruction(tx: &Transaction, index: usize) -> Option { +#[cfg(test)] +mod tests { + use super::*; + use crate::budget_expr::BudgetExpr; + use bincode::{deserialize, serialize}; + use solana_sdk::system_instruction::SystemInstruction; + + fn deserialize_system_instruction(tx: &Transaction, index: usize) -> Option { deserialize(&tx.data(index)).ok() } - pub fn instruction(tx: &Transaction, index: usize) -> Option { + fn deserialize_budget_instruction(tx: &Transaction, index: usize) -> Option { deserialize(&tx.data(index)).ok() } /// Verify only the payment plan. - pub fn verify_plan(tx: &Transaction) -> bool { + fn verify_plan(tx: &Transaction) -> bool { if let Some(SystemInstruction::CreateAccount { lamports, .. }) = - Self::system_instruction(tx, 0) + deserialize_system_instruction(tx, 0) { if let Some(BudgetInstruction::InitializeAccount(expr)) = - BudgetTransaction::instruction(&tx, 1) + deserialize_budget_instruction(&tx, 1) { if !expr.verify(lamports) { return false; @@ -134,20 +140,13 @@ impl BudgetTransaction { } true } -} - -#[cfg(test)] -mod tests { - use super::*; - use crate::budget_expr::BudgetExpr; - use bincode::{deserialize, serialize}; #[test] fn test_claim() { let keypair = Keypair::new(); let zero = Hash::default(); let tx0 = BudgetTransaction::new_payment(&keypair, &keypair.pubkey(), 42, zero, 0); - assert!(BudgetTransaction::verify_plan(&tx0)); + assert!(verify_plan(&tx0)); } #[test] @@ -157,7 +156,7 @@ mod tests { let keypair1 = Keypair::new(); let pubkey1 = keypair1.pubkey(); let tx0 = BudgetTransaction::new_payment(&keypair0, &pubkey1, 42, zero, 0); - assert!(BudgetTransaction::verify_plan(&tx0)); + assert!(verify_plan(&tx0)); } #[test] @@ -166,7 +165,7 @@ mod tests { let keypair0 = Keypair::new(); let pubkey1 = Keypair::new().pubkey(); let tx0 = BudgetTransaction::new_payment(&keypair0, &pubkey1, 1, zero, 1); - assert!(BudgetTransaction::verify_plan(&tx0)); + assert!(verify_plan(&tx0)); } #[test] @@ -186,13 +185,13 @@ mod tests { let keypair = Keypair::new(); let pubkey = keypair.pubkey(); let mut tx = BudgetTransaction::new_payment(&keypair, &pubkey, 42, zero, 0); - let mut system_instruction = BudgetTransaction::system_instruction(&tx, 0).unwrap(); + let mut system_instruction = deserialize_system_instruction(&tx, 0).unwrap(); if let SystemInstruction::CreateAccount { ref mut lamports, .. } = system_instruction { *lamports = 1_000_000; // <-- attack, part 1! - let mut instruction = BudgetTransaction::instruction(&tx, 1).unwrap(); + let mut instruction = deserialize_budget_instruction(&tx, 1).unwrap(); if let BudgetInstruction::InitializeAccount(ref mut expr) = instruction { if let BudgetExpr::Pay(ref mut payment) = expr { payment.lamports = *lamports; // <-- attack, part 2! @@ -201,7 +200,7 @@ mod tests { tx.instructions[1].data = serialize(&instruction).unwrap(); } tx.instructions[0].data = serialize(&system_instruction).unwrap(); - assert!(BudgetTransaction::verify_plan(&tx)); + assert!(verify_plan(&tx)); assert!(!tx.verify_signature()); } @@ -213,14 +212,14 @@ mod tests { let pubkey1 = keypair1.pubkey(); let zero = Hash::default(); let mut tx = BudgetTransaction::new_payment(&keypair0, &pubkey1, 42, zero, 0); - let mut instruction = BudgetTransaction::instruction(&tx, 1); + let mut instruction = deserialize_budget_instruction(&tx, 1); if let Some(BudgetInstruction::InitializeAccount(ref mut expr)) = instruction { if let BudgetExpr::Pay(ref mut payment) = expr { payment.to = thief_keypair.pubkey(); // <-- attack! } } tx.instructions[1].data = serialize(&instruction).unwrap(); - assert!(BudgetTransaction::verify_plan(&tx)); + assert!(verify_plan(&tx)); assert!(!tx.verify_signature()); } @@ -230,23 +229,23 @@ mod tests { let keypair1 = Keypair::new(); let zero = Hash::default(); let mut tx = BudgetTransaction::new_payment(&keypair0, &keypair1.pubkey(), 1, zero, 0); - let mut instruction = BudgetTransaction::instruction(&tx, 1).unwrap(); + let mut instruction = deserialize_budget_instruction(&tx, 1).unwrap(); if let BudgetInstruction::InitializeAccount(ref mut expr) = instruction { if let BudgetExpr::Pay(ref mut payment) = expr { payment.lamports = 2; // <-- attack! } } tx.instructions[1].data = serialize(&instruction).unwrap(); - assert!(!BudgetTransaction::verify_plan(&tx)); + assert!(!verify_plan(&tx)); // Also, ensure all branchs of the plan spend all lamports - let mut instruction = BudgetTransaction::instruction(&tx, 1).unwrap(); + let mut instruction = deserialize_budget_instruction(&tx, 1).unwrap(); if let BudgetInstruction::InitializeAccount(ref mut expr) = instruction { if let BudgetExpr::Pay(ref mut payment) = expr { payment.lamports = 0; // <-- whoops! } } tx.instructions[1].data = serialize(&instruction).unwrap(); - assert!(!BudgetTransaction::verify_plan(&tx)); + assert!(!verify_plan(&tx)); } }