diff --git a/src/budget.rs b/src/budget.rs index 7de861bc5..344038c33 100644 --- a/src/budget.rs +++ b/src/budget.rs @@ -44,6 +44,9 @@ pub enum Budget { /// Either make a payment after one condition or a different payment after another /// condition, which ever condition is satisfied first. Or((Condition, Payment), (Condition, Payment)), + + /// Make a payment after both of two conditions are satisfied + And(Condition, Condition, Payment), } impl Budget { @@ -57,6 +60,15 @@ impl Budget { Budget::After(Condition::Signature(from), Payment { tokens, to }) } + /// Create a budget that pays tokens` to `to` after being witnessed by 2x `from`s + pub fn new_2_2_multisig_payment(from0: Pubkey, from1: Pubkey, tokens: i64, to: Pubkey) -> Self { + Budget::And( + Condition::Signature(from0), + Condition::Signature(from1), + Payment { tokens, to }, + ) + } + /// Create a budget that pays `tokens` to `to` after the given DateTime. pub fn new_future_payment(dt: DateTime, from: Pubkey, tokens: i64, to: Pubkey) -> Self { Budget::After(Condition::Timestamp(dt, from), Payment { tokens, to }) @@ -87,7 +99,9 @@ impl Budget { /// Return true if the budget spends exactly `spendable_tokens`. pub fn verify(&self, spendable_tokens: i64) -> bool { match self { - Budget::Pay(payment) | Budget::After(_, payment) => payment.tokens == spendable_tokens, + Budget::Pay(payment) | Budget::After(_, payment) | Budget::And(_, _, payment) => { + payment.tokens == spendable_tokens + } Budget::Or(a, b) => a.1.tokens == spendable_tokens && b.1.tokens == spendable_tokens, } } @@ -95,15 +109,29 @@ impl Budget { /// Apply a witness to the budget to see if the budget can be reduced. /// If so, modify the budget in-place. pub fn apply_witness(&mut self, witness: &Witness, from: &Pubkey) { - let new_payment = match self { - 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), + let new_budget = match self { + Budget::After(cond, payment) if cond.is_satisfied(witness, from) => { + Some(Budget::Pay(payment.clone())) + } + Budget::Or((cond, payment), _) if cond.is_satisfied(witness, from) => { + Some(Budget::Pay(payment.clone())) + } + Budget::Or(_, (cond, payment)) if cond.is_satisfied(witness, from) => { + Some(Budget::Pay(payment.clone())) + } + Budget::And(cond0, cond1, payment) => { + if cond0.is_satisfied(witness, from) { + Some(Budget::After(cond1.clone(), payment.clone())) + } else if cond1.is_satisfied(witness, from) { + Some(Budget::After(cond0.clone(), payment.clone())) + } else { + None + } + } _ => None, - }.cloned(); - - if let Some(payment) = new_payment { - mem::replace(self, Budget::Pay(payment)); + }; + if let Some(budget) = new_budget { + mem::replace(self, budget); } } } @@ -189,4 +217,14 @@ mod tests { budget.apply_witness(&Witness::Signature, &from); assert_eq!(budget, Budget::new_payment(42, from)); } + #[test] + fn test_2_2_multisig_payment() { + let from0 = Keypair::new().pubkey(); + let from1 = Keypair::new().pubkey(); + let to = Pubkey::default(); + + let mut budget = Budget::new_2_2_multisig_payment(from0, from1, 42, to); + budget.apply_witness(&Witness::Signature, &from0); + assert_eq!(budget, Budget::new_authorized_payment(from1, 42, to)); + } }