From 5691bf557c82945728b06ad8b98bf3586237e34c Mon Sep 17 00:00:00 2001 From: Michael Vines Date: Thu, 20 Sep 2018 09:38:37 -0700 Subject: [PATCH] Handle bad account userdata better --- src/budget_contract.rs | 114 +++++++++++++++++++++++++++++++---------- 1 file changed, 87 insertions(+), 27 deletions(-) diff --git a/src/budget_contract.rs b/src/budget_contract.rs index 86776398c1..153f1f3e4b 100644 --- a/src/budget_contract.rs +++ b/src/budget_contract.rs @@ -19,6 +19,7 @@ pub enum BudgetError { NegativeTokens, DestinationMissing(Pubkey), FailedWitness, + UserdataTooSmall, } #[derive(Serialize, Deserialize, Debug, Clone, Default, PartialEq)] @@ -148,39 +149,42 @@ impl BudgetContract { state.pending_budget = Some(budget); accounts[1].tokens += contract.tokens; state.initialized = true; - state.serialize(&mut accounts[1].userdata); - Ok(()) + state.serialize(&mut accounts[1].userdata) } } } Instruction::ApplyTimestamp(dt) => { - let mut state = Self::deserialize(&accounts[1].userdata).unwrap(); - if !state.is_pending() { - Err(BudgetError::ContractNotPending(tx.keys[1])) - } else if !state.initialized { - trace!("contract is uninitialized"); - Err(BudgetError::UninitializedContract(tx.keys[1])) + if let Ok(mut state) = Self::deserialize(&accounts[1].userdata) { + if !state.is_pending() { + Err(BudgetError::ContractNotPending(tx.keys[1])) + } else if !state.initialized { + trace!("contract is uninitialized"); + Err(BudgetError::UninitializedContract(tx.keys[1])) + } else { + trace!("apply timestamp"); + state.apply_timestamp(&tx.keys, accounts, *dt)?; + trace!("apply timestamp committed"); + state.serialize(&mut accounts[1].userdata) + } } else { - trace!("apply timestamp"); - state.apply_timestamp(&tx.keys, accounts, *dt)?; - trace!("apply timestamp committed"); - state.serialize(&mut accounts[1].userdata); - Ok(()) + Err(BudgetError::UninitializedContract(tx.keys[1])) } } Instruction::ApplySignature => { - let mut state = Self::deserialize(&accounts[1].userdata).unwrap(); - if !state.is_pending() { - Err(BudgetError::ContractNotPending(tx.keys[1])) - } else if !state.initialized { - trace!("contract is uninitialized"); - Err(BudgetError::UninitializedContract(tx.keys[1])) + if let Ok(mut state) = Self::deserialize(&accounts[1].userdata) { + if !state.is_pending() { + Err(BudgetError::ContractNotPending(tx.keys[1])) + } else if !state.initialized { + trace!("contract is uninitialized"); + Err(BudgetError::UninitializedContract(tx.keys[1])) + } else { + trace!("apply signature"); + state.apply_signature(&tx.keys, accounts)?; + trace!("apply signature committed"); + state.serialize(&mut accounts[1].userdata) + } } else { - trace!("apply signature"); - state.apply_signature(&tx.keys, accounts)?; - trace!("apply signature committed"); - state.serialize(&mut accounts[1].userdata); - Ok(()) + Err(BudgetError::UninitializedContract(tx.keys[1])) } } Instruction::NewVote(_vote) => { @@ -190,16 +194,26 @@ impl BudgetContract { } } } - fn serialize(&self, output: &mut [u8]) { + fn serialize(&self, output: &mut [u8]) -> Result<(), BudgetError> { let len = serialized_size(self).unwrap() as u64; + if output.len() < len as usize { + warn!( + "{} bytes required to serialize, only have {} bytes", + len, + output.len() + ); + return Err(BudgetError::UserdataTooSmall); + } { let writer = io::BufWriter::new(&mut output[..8]); serialize_into(writer, &len).unwrap(); } + { let writer = io::BufWriter::new(&mut output[8..8 + len as usize]); serialize_into(writer, self).unwrap(); } + Ok(()) } pub fn deserialize(input: &[u8]) -> bincode::Result { @@ -220,7 +234,7 @@ impl BudgetContract { if let Ok(mut state) = BudgetContract::deserialize(&accounts[1].userdata) { trace!("saved error {:?}", e); state.last_error = Some(e); - state.serialize(&mut accounts[1].userdata); + state.serialize(&mut accounts[1].userdata).unwrap(); } else { trace!("error in uninitialized contract {:?}", e,); } @@ -233,6 +247,7 @@ impl BudgetContract { /// 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| { @@ -268,13 +283,23 @@ mod test { fn test_serializer() { let mut a = Account::new(0, 512, BudgetContract::id()); let b = BudgetContract::default(); - b.serialize(&mut a.userdata); + b.serialize(&mut a.userdata).unwrap(); let buf = serialize(&b).unwrap(); assert_eq!(a.userdata[8..8 + buf.len()], buf[0..]); let c = BudgetContract::deserialize(&a.userdata).unwrap(); assert_eq!(b, c); } + #[test] + fn test_serializer_userdata_too_small() { + let mut a = Account::new(0, 1, BudgetContract::id()); + let b = BudgetContract::default(); + assert_eq!( + b.serialize(&mut a.userdata), + Err(BudgetError::UserdataTooSmall) + ); + } + #[test] fn test_transfer_on_date() { let mut accounts = vec![ @@ -425,6 +450,41 @@ mod test { ); } + #[test] + fn test_userdata_too_small() { + let mut accounts = vec![ + Account::new(1, 0, BudgetContract::id()), + Account::new(1, 0, BudgetContract::id()), // <== userdata is 0, which is not enough + Account::new(1, 0, BudgetContract::id()), + ]; + let from = Keypair::new(); + let contract = Keypair::new(); + let to = Keypair::new(); + let tx = Transaction::budget_new_on_date( + &from, + to.pubkey(), + contract.pubkey(), + Utc::now(), + 1, + Hash::default(), + ); + + BudgetContract::process_transaction(&tx, &mut accounts); + assert!(BudgetContract::deserialize(&accounts[1].userdata).is_err()); + + let tx = Transaction::budget_new_timestamp( + &from, + contract.pubkey(), + to.pubkey(), + Utc::now(), + Hash::default(), + ); + BudgetContract::process_transaction(&tx, &mut accounts); + assert!(BudgetContract::deserialize(&accounts[1].userdata).is_err()); + + // Success if there was no panic... + } + /// Detect binary changes in the serialized contract userdata, which could have a downstream /// affect on SDKs and DApps #[test]