From a54854abc7ce5d99f0e7865ada98f71bfcf1c28d Mon Sep 17 00:00:00 2001
From: Greg Fitzgerald <greg@solana.com>
Date: Mon, 18 Mar 2019 06:50:17 -0600
Subject: [PATCH] Do Budget verification in BudgetScript

---
 programs/budget_api/src/budget_script.rs      | 36 ++++++++++
 programs/budget_api/src/budget_transaction.rs | 68 -------------------
 2 files changed, 36 insertions(+), 68 deletions(-)

diff --git a/programs/budget_api/src/budget_script.rs b/programs/budget_api/src/budget_script.rs
index 07b9b44142..f4a0220ccb 100644
--- a/programs/budget_api/src/budget_script.rs
+++ b/programs/budget_api/src/budget_script.rs
@@ -18,6 +18,9 @@ impl BudgetScript {
         lamports: u64,
         expr: BudgetExpr,
     ) -> Script {
+        if !expr.verify(lamports) {
+            panic!("invalid budget expression");
+        }
         let space = serialized_size(&BudgetState::new(expr.clone())).unwrap();
         let instructions = vec![
             SystemInstruction::new_program_account(&from, contract, lamports, space, &id()),
@@ -70,3 +73,36 @@ impl BudgetScript {
         Self::new_account(from, contract, lamports, expr)
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::budget_expr::BudgetExpr;
+
+    #[test]
+    fn test_budget_script_verify() {
+        let alice_pubkey = Keypair::new().pubkey();
+        let bob_pubkey = Keypair::new().pubkey();
+        BudgetScript::pay(&alice_pubkey, &bob_pubkey, 1); // No panic! indicates success.
+    }
+
+    #[test]
+    #[should_panic]
+    fn test_budget_script_overspend() {
+        let alice_pubkey = Keypair::new().pubkey();
+        let bob_pubkey = Keypair::new().pubkey();
+        let budget_pubkey = Keypair::new().pubkey();
+        let expr = BudgetExpr::new_payment(2, &bob_pubkey);
+        BudgetScript::new_account(&alice_pubkey, &budget_pubkey, 1, expr);
+    }
+
+    #[test]
+    #[should_panic]
+    fn test_budget_script_underspend() {
+        let alice_pubkey = Keypair::new().pubkey();
+        let bob_pubkey = Keypair::new().pubkey();
+        let budget_pubkey = Keypair::new().pubkey();
+        let expr = BudgetExpr::new_payment(1, &bob_pubkey);
+        BudgetScript::new_account(&alice_pubkey, &budget_pubkey, 2, expr);
+    }
+}
diff --git a/programs/budget_api/src/budget_transaction.rs b/programs/budget_api/src/budget_transaction.rs
index c7854c775e..c08a7ff494 100644
--- a/programs/budget_api/src/budget_transaction.rs
+++ b/programs/budget_api/src/budget_transaction.rs
@@ -109,71 +109,3 @@ impl BudgetTransaction {
         Self::new_signed(from_keypair, script, recent_blockhash, 0)
     }
 }
-
-#[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()
-    }
-
-    fn deserialize_budget_instruction(tx: &Transaction, index: usize) -> Option<BudgetInstruction> {
-        deserialize(&tx.data(index)).ok()
-    }
-
-    /// Verify only the payment plan.
-    fn verify_plan(tx: &Transaction) -> bool {
-        if let Some(SystemInstruction::CreateAccount { lamports, .. }) =
-            deserialize_system_instruction(tx, 0)
-        {
-            if let Some(BudgetInstruction::InitializeAccount(expr)) =
-                deserialize_budget_instruction(&tx, 1)
-            {
-                if !expr.verify(lamports) {
-                    return false;
-                }
-            }
-        }
-        true
-    }
-
-    #[test]
-    fn test_payment() {
-        let zero = Hash::default();
-        let keypair0 = Keypair::new();
-        let keypair1 = Keypair::new();
-        let pubkey1 = keypair1.pubkey();
-        let tx0 = BudgetTransaction::new_payment(&keypair0, &pubkey1, 42, zero, 0);
-        assert!(verify_plan(&tx0));
-    }
-
-    #[test]
-    fn test_overspend_attack() {
-        let keypair0 = Keypair::new();
-        let keypair1 = Keypair::new();
-        let zero = Hash::default();
-        let mut tx = BudgetTransaction::new_payment(&keypair0, &keypair1.pubkey(), 1, zero, 0);
-        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!(!verify_plan(&tx));
-
-        // Also, ensure all branchs of the plan spend all lamports
-        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!(!verify_plan(&tx));
-    }
-}