From 703192f45bad6ba9347258e283b1a6469357477c Mon Sep 17 00:00:00 2001
From: Aditya Kulkarni <adityapk@gmail.com>
Date: Wed, 29 Apr 2020 12:21:07 -0700
Subject: [PATCH] Add guards around updating memo, save wallet

---
 lib/src/lightclient.rs | 14 ++++++++++++--
 lib/src/lightwallet.rs | 21 ++++++++++++++-------
 2 files changed, 26 insertions(+), 9 deletions(-)

diff --git a/lib/src/lightclient.rs b/lib/src/lightclient.rs
index 5fa9ee8..3925c79 100644
--- a/lib/src/lightclient.rs
+++ b/lib/src/lightclient.rs
@@ -342,6 +342,9 @@ impl LightClient {
         info!("Created new wallet with a new seed!");
         info!("Created LightClient to {}", &config.server);
 
+        // Save
+        l.do_save().map_err(|s| io::Error::new(ErrorKind::PermissionDenied, s))?;
+
         Ok(l)
     }
 
@@ -367,6 +370,9 @@ impl LightClient {
         info!("Created new wallet!");
         info!("Created LightClient to {}", &config.server);
 
+        // Save
+        l.do_save().map_err(|s| io::Error::new(ErrorKind::PermissionDenied, s))?;
+
         Ok(l)
     }
 
@@ -576,14 +582,18 @@ impl LightClient {
             1_000_000, // 1 MB write buffer
             File::create(self.config.get_wallet_path()).unwrap());
         
-        match self.wallet.write().unwrap().write(&mut file_buffer) {
+        let r = match self.wallet.write().unwrap().write(&mut file_buffer) {
             Ok(_) => Ok(()),
             Err(e) => {
                 let err = format!("ERR: {}", e);
                 error!("{}", err);
                 Err(e.to_string())
             }
-        }
+        };
+
+        file_buffer.flush().map_err(|e| format!("{}", e))?;
+
+        r
     }
 
     pub fn get_server_uri(&self) -> http::Uri {
diff --git a/lib/src/lightwallet.rs b/lib/src/lightwallet.rs
index d2b9b0b..1ab90aa 100644
--- a/lib/src/lightwallet.rs
+++ b/lib/src/lightwallet.rs
@@ -1062,14 +1062,21 @@ impl LightWallet {
                 };
 
                 {
-                    info!("A sapling note was spent in {}", tx.txid());
-                    // Update the WalletTx 
-                    // Do it in a short scope because of the write lock.
+                    info!("A sapling note was sent in {}, getting memo", tx.txid());
+                    
+                    // Do it in a short scope because of the write lock.   
                     let mut txs = self.txs.write().unwrap();
-                    txs.get_mut(&tx.txid()).unwrap()
-                        .notes.iter_mut()
-                        .find(|nd| nd.note == note).unwrap()
-                        .memo = Some(memo);
+
+                    // Update memo if we have this Tx. 
+                    match txs.get_mut(&tx.txid())
+                        .and_then(|t| {
+                            t.notes.iter_mut().find(|nd| nd.note == note)
+                        }) {
+                            None => (),
+                            Some(nd) => {
+                                nd.memo = Some(memo)
+                            }
+                        }
                 }
             }