From 11ea9e7c4bdf90bf386b040175f4079cc568513b Mon Sep 17 00:00:00 2001
From: Tyera Eulberg <tyera@solana.com>
Date: Fri, 21 Sep 2018 19:51:42 -0600
Subject: [PATCH] Add cancelable handling

---
 src/bin/wallet.rs     |  7 +++----
 src/budget_program.rs |  5 +++++
 src/transaction.rs    | 13 +++++++++----
 src/wallet.rs         | 10 +++++++++-
 4 files changed, 26 insertions(+), 9 deletions(-)

diff --git a/src/bin/wallet.rs b/src/bin/wallet.rs
index a7b617a992..3670338043 100644
--- a/src/bin/wallet.rs
+++ b/src/bin/wallet.rs
@@ -189,10 +189,9 @@ fn main() -> Result<(), Box<error::Error>> {
                         .use_delimiter(true)
                         .help("Any third party signatures required to unlock the tokens"),
                 ).arg(
-                    Arg::with_name("cancellable")
-                        .long("cancellable")
-                        .takes_value(false)
-                        .requires("witness"),
+                    Arg::with_name("cancelable")
+                        .long("cancelable")
+                        .takes_value(false),
                 ),
         ).subcommand(
             SubCommand::with_name("send-signature")
diff --git a/src/budget_program.rs b/src/budget_program.rs
index ebb24d1da8..1206a0c6ed 100644
--- a/src/budget_program.rs
+++ b/src/budget_program.rs
@@ -333,6 +333,7 @@ mod test {
             contract.pubkey(),
             dt,
             from.pubkey(),
+            true,
             1,
             Hash::default(),
         );
@@ -407,6 +408,7 @@ mod test {
             contract.pubkey(),
             dt,
             from.pubkey(),
+            true,
             1,
             Hash::default(),
         );
@@ -471,6 +473,8 @@ mod test {
             to.pubkey(),
             contract.pubkey(),
             Utc::now(),
+            from.pubkey(),
+            true,
             1,
             Hash::default(),
         );
@@ -524,6 +528,7 @@ mod test {
             contract,
             date,
             keypair.pubkey(),
+            true,
             192,
             Hash::default(),
         );
diff --git a/src/transaction.rs b/src/transaction.rs
index 1634c1fd68..62f512574b 100644
--- a/src/transaction.rs
+++ b/src/transaction.rs
@@ -153,14 +153,19 @@ impl Transaction {
         contract: Pubkey,
         dt: DateTime<Utc>,
         dt_pubkey: Pubkey,
+        cancelable: bool,
         tokens: i64,
         last_id: Hash,
     ) -> Self {
         let from = from_keypair.pubkey();
-        let budget = Budget::Or(
-            (Condition::Timestamp(dt, dt_pubkey), Payment { tokens, to }),
-            (Condition::Signature(from), Payment { tokens, to: from }),
-        );
+        let budget = if cancelable {
+            Budget::Or(
+                (Condition::Timestamp(dt, dt_pubkey), Payment { tokens, to }),
+                (Condition::Signature(from), Payment { tokens, to: from }),
+            )
+        } else {
+            Budget::After(Condition::Timestamp(dt, dt_pubkey), Payment { tokens, to })
+        };
         let instruction = Instruction::NewContract(Contract { budget, tokens });
         let userdata = serialize(&instruction).expect("serialize instruction");
         Self::new_with_userdata(
diff --git a/src/wallet.rs b/src/wallet.rs
index cd2d8151d8..7e469f9a86 100644
--- a/src/wallet.rs
+++ b/src/wallet.rs
@@ -374,6 +374,7 @@ pub fn process_command(config: &WalletConfig) -> Result<String, Box<error::Error
                 .into_vec()
                 .map_err(|_| WalletError::RpcRequestError("Received bad last_id".to_string()))?;
             let last_id = Hash::new(&last_id_vec);
+            let cancelable_bool = cancelable.is_some();
 
             if timestamp == None && *witnesses == None {
                 let tx = Transaction::new(&config.id, to, tokens, last_id);
@@ -400,7 +401,14 @@ pub fn process_command(config: &WalletConfig) -> Result<String, Box<error::Error
                 };
                 let process_id = Keypair::new().pubkey();
                 let tx = Transaction::budget_new_on_date(
-                    &config.id, to, process_id, dt, dt_pubkey, tokens, last_id,
+                    &config.id,
+                    to,
+                    process_id,
+                    dt,
+                    dt_pubkey,
+                    cancelable_bool,
+                    tokens,
+                    last_id,
                 );
                 let serialized = serialize(&tx).unwrap();
                 let params = json!(serialized);