diff --git a/src/bank.rs b/src/bank.rs index 320fbc4fe00652..d9f3182895d2d3 100644 --- a/src/bank.rs +++ b/src/bank.rs @@ -388,7 +388,7 @@ impl Bank { .expect("write() in apply_signature") .entry(tx_sig) { - e.get_mut().apply_witness(&Witness::Signature(from)); + e.get_mut().apply_witness(&Witness::Signature, &from); if let Some(payment) = e.get().final_payment() { self.apply_payment(&payment, &mut self.balances.write().unwrap()); e.remove_entry(); @@ -400,7 +400,7 @@ impl Bank { /// Process a Witness Timestamp. Any payment plans waiting on this timestamp /// will progress one step. - fn apply_timestamp(&self, _from: PublicKey, dt: DateTime) -> Result<()> { + fn apply_timestamp(&self, from: PublicKey, dt: DateTime) -> Result<()> { // Check to see if any timelocked transactions can be completed. let mut completed = vec![]; @@ -410,7 +410,7 @@ impl Bank { .write() .expect("'pending' write lock in apply_timestamp"); for (key, plan) in pending.iter_mut() { - plan.apply_witness(&Witness::Timestamp(dt)); + plan.apply_witness(&Witness::Timestamp(dt), &from); if let Some(payment) = plan.final_payment() { self.apply_payment(&payment, &mut self.balances.write().unwrap()); completed.push(key.clone()); diff --git a/src/budget.rs b/src/budget.rs index f2e9af55051f45..41f4f1033229bc 100644 --- a/src/budget.rs +++ b/src/budget.rs @@ -12,7 +12,7 @@ use std::mem; #[derive(Serialize, Deserialize, Debug, PartialEq, Eq, Clone)] pub enum Condition { /// Wait for a `Timestamp` `Witness` at or after the given `DateTime`. - Timestamp(DateTime), + Timestamp(DateTime, PublicKey), /// Wait for a `Signature` `Witness` from `PublicKey`. Signature(PublicKey), @@ -20,10 +20,12 @@ pub enum Condition { impl Condition { /// Return true if the given Witness satisfies this Condition. - pub fn is_satisfied(&self, witness: &Witness) -> bool { + pub fn is_satisfied(&self, witness: &Witness, from: &PublicKey) -> bool { match (self, witness) { - (Condition::Signature(pubkey), Witness::Signature(from)) => pubkey == from, - (Condition::Timestamp(dt), Witness::Timestamp(last_time)) => dt <= last_time, + (Condition::Signature(pubkey), Witness::Signature) => pubkey == from, + (Condition::Timestamp(dt, pubkey), Witness::Timestamp(last_time)) => { + pubkey == from && dt <= last_time + } _ => false, } } @@ -56,8 +58,13 @@ impl Budget { } /// Create a budget that pays `tokens` to `to` after the given DateTime. - pub fn new_future_payment(dt: DateTime, tokens: i64, to: PublicKey) -> Self { - Budget::After(Condition::Timestamp(dt), Payment { tokens, to }) + pub fn new_future_payment( + dt: DateTime, + from: PublicKey, + tokens: i64, + to: PublicKey, + ) -> Self { + Budget::After(Condition::Timestamp(dt, from), Payment { tokens, to }) } /// Create a budget that pays `tokens` to `to` after the given DateTime @@ -69,7 +76,7 @@ impl Budget { to: PublicKey, ) -> Self { Budget::Or( - (Condition::Timestamp(dt), Payment { tokens, to }), + (Condition::Timestamp(dt, from), Payment { tokens, to }), (Condition::Signature(from), Payment { tokens, to: from }), ) } @@ -94,11 +101,11 @@ impl PaymentPlan for Budget { /// Apply a witness to the budget to see if the budget can be reduced. /// If so, modify the budget in-place. - fn apply_witness(&mut self, witness: &Witness) { + fn apply_witness(&mut self, witness: &Witness, from: &PublicKey) { let new_payment = match self { - Budget::After(cond, payment) if cond.is_satisfied(witness) => Some(payment), - Budget::Or((cond, payment), _) if cond.is_satisfied(witness) => Some(payment), - Budget::Or(_, (cond, payment)) if cond.is_satisfied(witness) => Some(payment), + Budget::After(cond, payment) if cond.is_satisfied(witness, from) => Some(payment), + Budget::Or((cond, payment), _) if cond.is_satisfied(witness, from) => Some(payment), + Budget::Or(_, (cond, payment)) if cond.is_satisfied(witness, from) => Some(payment), _ => None, }.cloned(); @@ -111,20 +118,22 @@ impl PaymentPlan for Budget { #[cfg(test)] mod tests { use super::*; + use signature::{KeyPair, KeyPairUtil}; #[test] fn test_signature_satisfied() { - let sig = PublicKey::default(); - assert!(Condition::Signature(sig).is_satisfied(&Witness::Signature(sig))); + let from = PublicKey::default(); + assert!(Condition::Signature(from).is_satisfied(&Witness::Signature, &from)); } #[test] fn test_timestamp_satisfied() { let dt1 = Utc.ymd(2014, 11, 14).and_hms(8, 9, 10); let dt2 = Utc.ymd(2014, 11, 14).and_hms(10, 9, 8); - assert!(Condition::Timestamp(dt1).is_satisfied(&Witness::Timestamp(dt1))); - assert!(Condition::Timestamp(dt1).is_satisfied(&Witness::Timestamp(dt2))); - assert!(!Condition::Timestamp(dt2).is_satisfied(&Witness::Timestamp(dt1))); + let from = PublicKey::default(); + assert!(Condition::Timestamp(dt1, from).is_satisfied(&Witness::Timestamp(dt1), &from)); + assert!(Condition::Timestamp(dt1, from).is_satisfied(&Witness::Timestamp(dt2), &from)); + assert!(!Condition::Timestamp(dt2, from).is_satisfied(&Witness::Timestamp(dt1), &from)); } #[test] @@ -134,7 +143,7 @@ mod tests { let to = PublicKey::default(); assert!(Budget::new_payment(42, to).verify(42)); assert!(Budget::new_authorized_payment(from, 42, to).verify(42)); - assert!(Budget::new_future_payment(dt, 42, to).verify(42)); + assert!(Budget::new_future_payment(dt, from, 42, to).verify(42)); assert!(Budget::new_cancelable_future_payment(dt, from, 42, to).verify(42)); } @@ -144,20 +153,35 @@ mod tests { let to = PublicKey::default(); let mut budget = Budget::new_authorized_payment(from, 42, to); - budget.apply_witness(&Witness::Signature(from)); + budget.apply_witness(&Witness::Signature, &from); assert_eq!(budget, Budget::new_payment(42, to)); } #[test] fn test_future_payment() { let dt = Utc.ymd(2014, 11, 14).and_hms(8, 9, 10); - let to = PublicKey::default(); + let from = KeyPair::new().pubkey(); + let to = KeyPair::new().pubkey(); - let mut budget = Budget::new_future_payment(dt, 42, to); - budget.apply_witness(&Witness::Timestamp(dt)); + let mut budget = Budget::new_future_payment(dt, from, 42, to); + budget.apply_witness(&Witness::Timestamp(dt), &from); assert_eq!(budget, Budget::new_payment(42, to)); } + #[test] + fn test_unauthorized_future_payment() { + // Ensure timestamp will only be acknowledged if it came from the + // whitelisted public key. + let dt = Utc.ymd(2014, 11, 14).and_hms(8, 9, 10); + let from = KeyPair::new().pubkey(); + let to = KeyPair::new().pubkey(); + + let mut budget = Budget::new_future_payment(dt, from, 42, to); + let orig_budget = budget.clone(); + budget.apply_witness(&Witness::Timestamp(dt), &to); // <-- Attack! + assert_eq!(budget, orig_budget); + } + #[test] fn test_cancelable_future_payment() { let dt = Utc.ymd(2014, 11, 14).and_hms(8, 9, 10); @@ -165,11 +189,11 @@ mod tests { let to = PublicKey::default(); let mut budget = Budget::new_cancelable_future_payment(dt, from, 42, to); - budget.apply_witness(&Witness::Timestamp(dt)); + budget.apply_witness(&Witness::Timestamp(dt), &from); assert_eq!(budget, Budget::new_payment(42, to)); let mut budget = Budget::new_cancelable_future_payment(dt, from, 42, to); - budget.apply_witness(&Witness::Signature(from)); + budget.apply_witness(&Witness::Signature, &from); assert_eq!(budget, Budget::new_payment(42, from)); } } diff --git a/src/payment_plan.rs b/src/payment_plan.rs index 059c63df60459d..29b82a0056a93e 100644 --- a/src/payment_plan.rs +++ b/src/payment_plan.rs @@ -13,7 +13,7 @@ pub enum Witness { Timestamp(DateTime), /// A siganture from PublicKey. - Signature(PublicKey), + Signature, } /// Some amount of tokens that should be sent to the `to` `PublicKey`. @@ -36,5 +36,5 @@ pub trait PaymentPlan { /// Apply a witness to the payment plan to see if the plan can be reduced. /// If so, modify the plan in-place. - fn apply_witness(&mut self, witness: &Witness); + fn apply_witness(&mut self, witness: &Witness, from: &PublicKey); } diff --git a/src/transaction.rs b/src/transaction.rs index 2a8288c123c88f..f8df220673c136 100644 --- a/src/transaction.rs +++ b/src/transaction.rs @@ -32,9 +32,9 @@ impl PaymentPlan for Plan { } } - fn apply_witness(&mut self, witness: &Witness) { + fn apply_witness(&mut self, witness: &Witness, from: &PublicKey) { match self { - Plan::Budget(budget) => budget.apply_witness(witness), + Plan::Budget(budget) => budget.apply_witness(witness, from), } } } @@ -145,7 +145,7 @@ impl Transaction { ) -> Self { let from = from_keypair.pubkey(); let budget = Budget::Or( - (Condition::Timestamp(dt), Payment { tokens, to }), + (Condition::Timestamp(dt, from), Payment { tokens, to }), (Condition::Signature(from), Payment { tokens, to: from }), ); let plan = Plan::Budget(budget);