Move Budget's verify_plan() into tests
This functionality is supposed to be the the interpreter
This commit is contained in:
		| @@ -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<SystemInstruction> { | ||||
| #[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<SystemInstruction> { | ||||
|         deserialize(&tx.data(index)).ok() | ||||
|     } | ||||
|  | ||||
|     pub fn instruction(tx: &Transaction, index: usize) -> Option<BudgetInstruction> { | ||||
|     fn deserialize_budget_instruction(tx: &Transaction, index: usize) -> Option<BudgetInstruction> { | ||||
|         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)); | ||||
|     } | ||||
| } | ||||
|   | ||||
		Reference in New Issue
	
	Block a user