From a6c15684c9f66360c198ba0b2a5672f3382a5b1e Mon Sep 17 00:00:00 2001
From: Michael Vines <mvines@gmail.com>
Date: Thu, 20 Sep 2018 13:54:42 -0700
Subject: [PATCH] Avoid panicking invalid instructions

---
 src/budget_contract.rs | 40 +++++++++++++++++++++------
 src/system_contract.rs | 61 ++++++++++++++++++++++--------------------
 src/transaction.rs     |  2 +-
 3 files changed, 65 insertions(+), 38 deletions(-)

diff --git a/src/budget_contract.rs b/src/budget_contract.rs
index 153f1f3e4b..f173f4f1a6 100644
--- a/src/budget_contract.rs
+++ b/src/budget_contract.rs
@@ -246,14 +246,17 @@ impl BudgetContract {
     /// * accounts[1] - The contract context.  Once the contract has been completed, the tokens can
     /// be spent from this account .
     pub fn process_transaction(tx: &Transaction, accounts: &mut [Account]) -> () {
-        let instruction = deserialize(&tx.userdata).unwrap();
-        trace!("process_transaction: {:?}", instruction);
-        let _ = Self::apply_debits_to_budget_state(tx, accounts, &instruction)
-            .and_then(|_| Self::apply_credits_to_budget_state(tx, accounts, &instruction))
-            .map_err(|e| {
-                trace!("saving error {:?}", e);
-                Self::save_error_to_budget_state(e, accounts);
-            });
+        if let Ok(instruction) = deserialize(&tx.userdata) {
+            trace!("process_transaction: {:?}", instruction);
+            let _ = Self::apply_debits_to_budget_state(tx, accounts, &instruction)
+                .and_then(|_| Self::apply_credits_to_budget_state(tx, accounts, &instruction))
+                .map_err(|e| {
+                    trace!("saving error {:?}", e);
+                    Self::save_error_to_budget_state(e, accounts);
+                });
+        } else {
+            info!("Invalid transaction userdata: {:?}", tx.userdata);
+        }
     }
 
     //TODO the contract needs to provide a "get_balance" introspection call of the userdata
@@ -299,6 +302,27 @@ mod test {
             Err(BudgetError::UserdataTooSmall)
         );
     }
+    #[test]
+    fn test_invalid_instruction() {
+        let mut accounts = vec![
+            Account::new(1, 0, BudgetContract::id()),
+            Account::new(0, 512, BudgetContract::id()),
+        ];
+        let from = Keypair::new();
+        let contract = Keypair::new();
+
+        let tx = Transaction::new_with_userdata(
+            &from,
+            &[contract.pubkey()],
+            BudgetContract::id(),
+            vec![1, 2, 3], // <== garbage instruction
+            Hash::default(),
+            0,
+        );
+        BudgetContract::process_transaction(&tx, &mut accounts);
+
+        // Success if there was no panic...
+    }
 
     #[test]
     fn test_transfer_on_date() {
diff --git a/src/system_contract.rs b/src/system_contract.rs
index 4b0d94585e..f9591c7381 100644
--- a/src/system_contract.rs
+++ b/src/system_contract.rs
@@ -41,39 +41,42 @@ impl SystemContract {
         account.tokens
     }
     pub fn process_transaction(tx: &Transaction, accounts: &mut [Account]) {
-        let syscall: SystemContract = deserialize(&tx.userdata).unwrap();
-        trace!("process_transaction: {:?}", syscall);
-        match syscall {
-            SystemContract::CreateAccount {
-                tokens,
-                space,
-                contract_id,
-            } => {
-                if !Self::check_id(&accounts[0].contract_id) {
-                    return;
+        if let Ok(syscall) = deserialize(&tx.userdata) {
+            trace!("process_transaction: {:?}", syscall);
+            match syscall {
+                SystemContract::CreateAccount {
+                    tokens,
+                    space,
+                    contract_id,
+                } => {
+                    if !Self::check_id(&accounts[0].contract_id) {
+                        return;
+                    }
+                    if space > 0
+                        && (!accounts[1].userdata.is_empty()
+                            || !Self::check_id(&accounts[1].contract_id))
+                    {
+                        return;
+                    }
+                    accounts[0].tokens -= tokens;
+                    accounts[1].tokens += tokens;
+                    accounts[1].contract_id = contract_id;
+                    accounts[1].userdata = vec![0; space as usize];
                 }
-                if space > 0
-                    && (!accounts[1].userdata.is_empty()
-                        || !Self::check_id(&accounts[1].contract_id))
-                {
-                    return;
+                SystemContract::Assign { contract_id } => {
+                    if !Self::check_id(&accounts[0].contract_id) {
+                        return;
+                    }
+                    accounts[0].contract_id = contract_id;
                 }
-                accounts[0].tokens -= tokens;
-                accounts[1].tokens += tokens;
-                accounts[1].contract_id = contract_id;
-                accounts[1].userdata = vec![0; space as usize];
-            }
-            SystemContract::Assign { contract_id } => {
-                if !Self::check_id(&accounts[0].contract_id) {
-                    return;
+                SystemContract::Move { tokens } => {
+                    //bank should be verifying correctness
+                    accounts[0].tokens -= tokens;
+                    accounts[1].tokens += tokens;
                 }
-                accounts[0].contract_id = contract_id;
-            }
-            SystemContract::Move { tokens } => {
-                //bank should be verifying correctness
-                accounts[0].tokens -= tokens;
-                accounts[1].tokens += tokens;
             }
+        } else {
+            info!("Invalid transaction userdata: {:?}", tx.userdata);
         }
     }
 }
diff --git a/src/transaction.rs b/src/transaction.rs
index e1991d5e8f..43a685745c 100644
--- a/src/transaction.rs
+++ b/src/transaction.rs
@@ -49,7 +49,7 @@ impl Transaction {
     /// * `userdata` - The input data that the contract will execute with
     /// * `last_id` - The PoH hash.
     /// * `fee` - The transaction fee.
-    fn new_with_userdata(
+    pub fn new_with_userdata(
         from_keypair: &Keypair,
         transaction_keys: &[Pubkey],
         contract_id: Pubkey,