From ea01ff2aabfec1e9f80b87308e1e9ffa55391424 Mon Sep 17 00:00:00 2001
From: Greg Fitzgerald <greg@solana.com>
Date: Mon, 18 Mar 2019 06:12:21 -0600
Subject: [PATCH] Add pubkey to BudgetExpr::new_cancelable_future_payment for
 wallet

---
 programs/budget_api/src/budget_expr.rs   | 22 ++++++++++++----------
 programs/budget_api/src/budget_script.rs | 24 ++++++------------------
 2 files changed, 18 insertions(+), 28 deletions(-)

diff --git a/programs/budget_api/src/budget_expr.rs b/programs/budget_api/src/budget_expr.rs
index 124faa1e91..0bfc690aac 100644
--- a/programs/budget_api/src/budget_expr.rs
+++ b/programs/budget_api/src/budget_expr.rs
@@ -98,35 +98,37 @@ impl BudgetExpr {
         )
     }
 
-    /// Create a budget that pays `lamports` to `to` after the given DateTime.
+    /// Create a budget that pays `lamports` to `to` after the given DateTime signed
+    /// by `dt_pubkey`.
     pub fn new_future_payment(
         dt: DateTime<Utc>,
-        from: &Pubkey,
+        dt_pubkey: &Pubkey,
         lamports: u64,
         to: &Pubkey,
     ) -> Self {
         BudgetExpr::After(
-            Condition::Timestamp(dt, *from),
+            Condition::Timestamp(dt, *dt_pubkey),
             Box::new(Self::new_payment(lamports, to)),
         )
     }
 
     /// Create a budget that pays `lamports` to `to` after the given DateTime
-    /// unless cancelled by `from`.
+    /// signed by `dt_pubkey` unless canceled by `from`.
     pub fn new_cancelable_future_payment(
         dt: DateTime<Utc>,
-        from: &Pubkey,
+        dt_pubkey: &Pubkey,
         lamports: u64,
         to: &Pubkey,
+        from: &Pubkey,
     ) -> Self {
         BudgetExpr::Or(
             (
-                Condition::Timestamp(dt, *from),
+                Condition::Timestamp(dt, *dt_pubkey),
                 Box::new(Self::new_payment(lamports, to)),
             ),
             (
                 Condition::Signature(*from),
-                Box::new(Self::new_payment(lamports, to)),
+                Box::new(Self::new_payment(lamports, from)),
             ),
         )
     }
@@ -211,7 +213,7 @@ mod tests {
         assert!(BudgetExpr::new_payment(42, &to).verify(42));
         assert!(BudgetExpr::new_authorized_payment(&from, 42, &to).verify(42));
         assert!(BudgetExpr::new_future_payment(dt, &from, 42, &to).verify(42));
-        assert!(BudgetExpr::new_cancelable_future_payment(dt, &from, 42, &to).verify(42));
+        assert!(BudgetExpr::new_cancelable_future_payment(dt, &from, 42, &to, &from).verify(42));
     }
 
     #[test]
@@ -255,11 +257,11 @@ mod tests {
         let from = Pubkey::default();
         let to = Pubkey::default();
 
-        let mut expr = BudgetExpr::new_cancelable_future_payment(dt, &from, 42, &to);
+        let mut expr = BudgetExpr::new_cancelable_future_payment(dt, &from, 42, &to, &from);
         expr.apply_witness(&Witness::Timestamp(dt), &from);
         assert_eq!(expr, BudgetExpr::new_payment(42, &to));
 
-        let mut expr = BudgetExpr::new_cancelable_future_payment(dt, &from, 42, &to);
+        let mut expr = BudgetExpr::new_cancelable_future_payment(dt, &from, 42, &to, &from);
         expr.apply_witness(&Witness::Signature, &from);
         assert_eq!(expr, BudgetExpr::new_payment(42, &from));
     }
diff --git a/programs/budget_api/src/budget_script.rs b/programs/budget_api/src/budget_script.rs
index 8c1865014b..07b9b44142 100644
--- a/programs/budget_api/src/budget_script.rs
+++ b/programs/budget_api/src/budget_script.rs
@@ -1,4 +1,4 @@
-use crate::budget_expr::{BudgetExpr, Condition};
+use crate::budget_expr::BudgetExpr;
 use crate::budget_instruction::BudgetInstruction;
 use crate::budget_state::BudgetState;
 use crate::id;
@@ -43,22 +43,10 @@ impl BudgetScript {
         cancelable: Option<Pubkey>,
         lamports: u64,
     ) -> Script {
-        let expr = if let Some(from) = cancelable {
-            BudgetExpr::Or(
-                (
-                    Condition::Timestamp(dt, *dt_pubkey),
-                    Box::new(BudgetExpr::new_payment(lamports, to)),
-                ),
-                (
-                    Condition::Signature(from),
-                    Box::new(BudgetExpr::new_payment(lamports, &from)),
-                ),
-            )
+        let expr = if let Some(from) = &cancelable {
+            BudgetExpr::new_cancelable_future_payment(dt, dt_pubkey, lamports, to, from)
         } else {
-            BudgetExpr::After(
-                Condition::Timestamp(dt, *dt_pubkey),
-                Box::new(BudgetExpr::new_payment(lamports, to)),
-            )
+            BudgetExpr::new_future_payment(dt, dt_pubkey, lamports, to)
         };
 
         Self::new_account(from, contract, lamports, expr)
@@ -73,8 +61,8 @@ impl BudgetScript {
         cancelable: Option<Pubkey>,
         lamports: u64,
     ) -> Script {
-        let expr = if let Some(from) = cancelable {
-            BudgetExpr::new_cancelable_authorized_payment(witness, lamports, to, &from)
+        let expr = if let Some(from) = &cancelable {
+            BudgetExpr::new_cancelable_authorized_payment(witness, lamports, to, from)
         } else {
             BudgetExpr::new_authorized_payment(witness, lamports, to)
         };