add spent height

lib/src/lightclient.rs

@@ -790,6 +790,8 @@ impl LightClient {
                                 "address"            => address,
                                 "spendable"          => spendable,
                                 "spent"              => nd.spent.map(|spent_txid| format!("{}", spent_txid)),
+                                "spent_at_height"    => nd.spent_at_height.map(|h| format!("{}", h)),
+                                "witness_size"       => nd.witnesses.len(),
                                 "unconfirmed_spent"  => nd.unconfirmed_spent.map(|spent_txid| format!("{}", spent_txid)),
                             })
                         }

lib/src/lightwallet.rs

@@ -198,7 +198,7 @@ impl LightWallet {
     let zdustaddress = zdustextfvk.default_address().unwrap().1;
 
     (zdustaddress)
-}
+    }
 
     pub fn is_shielded_address(addr: &String, config: &LightClientConfig) -> bool {
         match address::RecipientAddress::from_str(addr,
@@ -414,7 +414,8 @@ impl LightWallet {
 
         let birthday = reader.read_u64::<LittleEndian>()?;
 
-        Ok(LightWallet{
+        
+            let lw = LightWallet{
             encrypted:   encrypted,
             unlocked:    !encrypted, // When reading from disk, if wallet is encrypted, it starts off locked. 
             enc_seed:    enc_seed,
@@ -428,7 +429,14 @@ impl LightWallet {
             config:      config.clone(),
             birthday,
             total_scan_duration: Arc::new(RwLock::new(vec![Duration::new(0, 0)])),
-        })
+        };
+
+        // Do a one-time fix of the spent_at_height for older wallets
+        if version <= 7 {
+            lw.fix_spent_at_height();
+        }
+
+        Ok(lw)
     }
 
     pub fn write<W: Write>(&self, mut writer: W) -> io::Result<()> {
@@ -1835,6 +1843,16 @@ impl LightWallet {
             // Create a write lock 
             let mut txs = self.txs.write().unwrap();
 
+             // Trim the older witnesses
+             txs.values_mut().for_each(|wtx| {
+                wtx.notes
+                    .iter_mut()
+                    .filter(|nd| nd.spent.is_some() && nd.spent_at_height.is_some() && nd.spent_at_height.unwrap() < height - (MAX_REORG as i32) - 1)
+                    .for_each(|nd| {
+                        nd.witnesses.clear()
+                    })
+            });
+
             // Create a Vec containing all unspent nullifiers.
             // Include only the confirmed spent nullifiers, since unconfirmed ones still need to be included
             // during scan_block below.
@@ -1872,11 +1890,22 @@ impl LightWallet {
                 let nf_refs = nfs.iter().map(|(nf, account, _)| (nf.to_vec(), *account)).collect::<Vec<_>>();
                 let extfvks: Vec<ExtendedFullViewingKey> = self.zkeys.read().unwrap().iter().map(|zk| zk.extfvk.clone()).collect();
 
-                // Create a single mutable slice of all the newly-added witnesses.
+                 // Create a single mutable slice of all the wallet's note's witnesses.
                 let mut witness_refs: Vec<_> = txs
                     .values_mut()
-                    .map(|tx| tx.notes.iter_mut().filter_map(
-                        |nd| if nd.spent.is_none() && nd.unconfirmed_spent.is_none() { nd.witnesses.last_mut() } else { None }))
+                    .map(|tx| 
+                        tx.notes.iter_mut()
+                            .filter_map(|nd| 
+                                // Note was not spent 
+                                if nd.spent.is_none() && nd.unconfirmed_spent.is_none() {
+                                    nd.witnesses.last_mut() 
+                                } else if nd.spent.is_some() && nd.spent_at_height.is_some() && nd.spent_at_height.unwrap() < height - (MAX_REORG as i32) - 1 {
+                                   // Note was spent in the last 100 blocks
+                                    nd.witnesses.last_mut() 
+                                } else { 
+                                    // If note was old (spent NOT in the last 100 blocks)
+                                    None 
+                                }))
                     .flatten()
                     .collect();
 
@@ -1930,6 +1959,7 @@ impl LightWallet {
                 // Mark the note as spent, and remove the unconfirmed part of it
                 info!("Marked a note as spent");
                 spent_note.spent = Some(tx.txid);
+                spent_note.spent_at_height = Some(height);
                 spent_note.unconfirmed_spent = None::<TxId>;
 
                 total_shielded_value_spent += spent_note.note.value;
@@ -1992,6 +2022,22 @@ impl LightWallet {
         Ok(all_txs)
     }
 
+    // Add the spent_at_height for each sapling note that has been spent. This field was added in wallet version 8, 
+    // so for older wallets, it will need to be added
+    pub fn fix_spent_at_height(&self) {
+        // First, build an index of all the txids and the heights at which they were spent. 
+        let spent_txid_map: HashMap<_, _> = self.txs.read().unwrap().iter().map(|(txid, wtx)| (txid.clone(), wtx.block)).collect();
+
+        // Go over all the sapling notes that might need updating
+        self.txs.write().unwrap().values_mut().for_each(|wtx| {
+            wtx.notes.iter_mut()
+                .filter(|nd| nd.spent.is_some() && nd.spent_at_height.is_none())
+                .for_each(|nd| {
+                    nd.spent_at_height = spent_txid_map.get(&nd.spent.unwrap()).map(|b| *b);
+                })
+        });
+    }
+
     pub fn send_to_address<F> (
         &self,
         consensus_branch_id: u32,

lib/src/lightwallet/data.rs

@@ -66,9 +66,10 @@ pub struct SaplingNoteData {
     pub(super) extfvk: ExtendedFullViewingKey, // Technically, this should be recoverable from the account number, but we're going to refactor this in the future, so I'll write it again here.
     pub diversifier: Diversifier,
     pub note: Note<Bls12>,
-    pub(super) witnesses: Vec<IncrementalWitness<Node>>,
+    pub witnesses: Vec<IncrementalWitness<Node>>,
     pub(super) nullifier: [u8; 32],
     pub spent: Option<TxId>,             // If this note was confirmed spent
+    pub spent_at_height: Option<i32>,        // The height at which this note was spent
     pub unconfirmed_spent: Option<TxId>, // If this note was spent in a send, but has not yet been confirmed.
     pub memo:  Option<Memo>,
     pub is_change: bool,
@@ -107,7 +108,7 @@ pub fn read_note<R: Read>(mut reader: R) -> io::Result<(u64, Fs)> {
 
 impl SaplingNoteData {
     fn serialized_version() -> u64 {
-        1
+        2
     }
 
     pub fn new(
@@ -133,6 +134,7 @@ impl SaplingNoteData {
             witnesses: vec![witness],
             nullifier: nf,
             spent: None,
+            spent_at_height: None,
             unconfirmed_spent: None,
             memo: None,
             is_change: output.is_change,
@@ -141,7 +143,7 @@ impl SaplingNoteData {
 
     // Reading a note also needs the corresponding address to read from.
     pub fn read<R: Read>(mut reader: R) -> io::Result<Self> {
-        let _version = reader.read_u64::<LittleEndian>()?;
+        let version = reader.read_u64::<LittleEndian>()?;
 
         let account = reader.read_u64::<LittleEndian>()? as usize;
 
@@ -176,6 +178,12 @@ impl SaplingNoteData {
             Ok(TxId{0: txid_bytes})
         })?;
 
+        let spent_at_height = if version >=2 {
+            Optional::read(&mut reader, |r| r.read_i32::<LittleEndian>())?
+        } else {
+            None
+        };
+
         let memo = Optional::read(&mut reader, |r| {
             let mut memo_bytes = [0u8; 512];
             r.read_exact(&mut memo_bytes)?;
@@ -195,6 +203,7 @@ impl SaplingNoteData {
             witnesses,
             nullifier,
             spent,
+            spent_at_height,
             unconfirmed_spent: None,
             memo,
             is_change,
@@ -224,6 +233,8 @@ impl SaplingNoteData {
         writer.write_all(&self.nullifier)?;
         Optional::write(&mut writer, &self.spent, |w, t| w.write_all(&t.0))?;
 
+        Optional::write(&mut writer, &self.spent_at_height, |w, h| w.write_i32::<LittleEndian>(*h))?;
+
         Optional::write(&mut writer, &self.memo, |w, m| w.write_all(m.as_bytes()))?;
 
         writer.write_u8(if self.is_change {1} else {0})?;