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