Add spent_at_height for notes

lib/src/lightclient.rs

@@ -835,6 +835,8 @@ impl LightClient {
                                 "is_change"          => nd.is_change,
                                 "address"            => LightWallet::note_address(self.config.hrp_sapling_address(), nd),
                                 "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

@@ -143,7 +143,7 @@ pub struct LightWallet {
 
 impl LightWallet {
     pub fn serialized_version() -> u64 {
-        return 7;
+        return 8;
     }
 
     fn get_taddr_from_bip39seed(config: &LightClientConfig, bip39_seed: &[u8], pos: u32) -> SecretKey {
@@ -379,7 +379,7 @@ 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,
@@ -394,7 +394,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<()> {
@@ -1686,6 +1693,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.
@@ -1723,11 +1740,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();
 
@@ -1780,6 +1808,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;
@@ -1842,6 +1871,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(
         &self,
         consensus_branch_id: u32,

lib/src/lightwallet/data.rs

@@ -65,9 +65,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,
@@ -106,7 +107,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(
@@ -132,6 +133,7 @@ impl SaplingNoteData {
             witnesses: vec![witness],
             nullifier: nf,
             spent: None,
+            spent_at_height: None,
             unconfirmed_spent: None,
             memo: None,
             is_change: output.is_change,
@@ -140,7 +142,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;
         
@@ -175,6 +177,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)?;
@@ -194,6 +202,7 @@ impl SaplingNoteData {
             witnesses,
             nullifier,
             spent,
+            spent_at_height,
             unconfirmed_spent: None,
             memo,
             is_change,
@@ -221,6 +230,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})?;