From ad15d633ad7f8d0bba7acc59b726065dfafd1154 Mon Sep 17 00:00:00 2001
From: Aditya Kulkarni <zcash@adityapk.com>
Date: Sat, 21 Sep 2019 20:39:20 -0700
Subject: [PATCH] trim stored blocks in the wallet

---
 src/lightclient.rs     |  8 ++++----
 src/lightwallet/mod.rs | 19 +++++++++++++++----
 2 files changed, 19 insertions(+), 8 deletions(-)

diff --git a/src/lightclient.rs b/src/lightclient.rs
index 23197f3..15f2fb1 100644
--- a/src/lightclient.rs
+++ b/src/lightclient.rs
@@ -545,7 +545,7 @@ impl LightClient {
         // Count how many bytes we've downloaded
         let bytes_downloaded = Arc::new(AtomicUsize::new(0));
 
-        let mut total_reorg = 0u64;
+        let mut total_reorg = 0;
 
         // Fetch CompactBlocks in increments
         loop {
@@ -592,9 +592,9 @@ impl LightClient {
             }
 
             // Make sure we're not re-orging too much!
-            if total_reorg > 99 {
-                error!("Reorg has now exceeded 100 blocks!");
-                return "Reorg has exceeded 100 blocks. Aborting.".to_string();
+            if total_reorg > (crate::lightwallet::MAX_REORG - 1) as u64 {
+                error!("Reorg has now exceeded {} blocks!", crate::lightwallet::MAX_REORG);
+                return format!("Reorg has exceeded {} blocks. Aborting.", crate::lightwallet::MAX_REORG);
             } 
             
             if invalid_height > 0 {
diff --git a/src/lightwallet/mod.rs b/src/lightwallet/mod.rs
index b3461bf..d994f11 100644
--- a/src/lightwallet/mod.rs
+++ b/src/lightwallet/mod.rs
@@ -42,14 +42,13 @@ use crate::LightClientConfig;
 
 use sha2::{Sha256, Digest};
 
-
-
 pub mod data;
 pub mod extended_key;
 
 use extended_key::{KeyIndex, ExtendedPrivKey};
 
 const ANCHOR_OFFSET: u32 = 1;
+pub const MAX_REORG: usize = 100;
 
 fn now() -> f64 {
     SystemTime::now().duration_since(SystemTime::UNIX_EPOCH).unwrap().as_secs() as f64
@@ -765,6 +764,8 @@ impl LightWallet {
                 .map(|block| block.tree.clone())
                 .unwrap_or(CommitmentTree::new()),
         };
+
+        // Create a write lock that will last for the rest of the function.
         let mut txs = self.txs.write().unwrap();
 
         // Create a Vec containing all unspent nullifiers.
@@ -870,9 +871,19 @@ impl LightWallet {
             }
         }
 
-        // Store scanned data for this block.
-        self.blocks.write().unwrap().push(block_data);
+        {
+            let mut blks = self.blocks.write().unwrap();
+            
+            // Store scanned data for this block.
+            blks.push(block_data);
 
+            // Trim the old blocks, keeping only as many as needed for a worst-case reorg (i.e. 101 blocks)
+            let len = blks.len();
+            if len > MAX_REORG + 1 {
+                blks.drain(..(len-MAX_REORG+1));
+            }
+        }
+        
         // Print info about the block every 10,000 blocks
         if height % 10_000 == 0 {
             match self.get_sapling_tree() {