diff --git a/lib/src/lightclient.rs b/lib/src/lightclient.rs index fa181a1..ca87b52 100644 --- a/lib/src/lightclient.rs +++ b/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)), }) } diff --git a/lib/src/lightwallet.rs b/lib/src/lightwallet.rs index 9c66639..8062a3c 100644 --- a/lib/src/lightwallet.rs +++ b/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::()?; - 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(&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::>(); let extfvks: Vec = 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::; 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 ( &self, consensus_branch_id: u32, diff --git a/lib/src/lightwallet/data.rs b/lib/src/lightwallet/data.rs index 99d4d97..14aded6 100644 --- a/lib/src/lightwallet/data.rs +++ b/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, - pub(super) witnesses: Vec>, + pub witnesses: Vec>, pub(super) nullifier: [u8; 32], pub spent: Option, // If this note was confirmed spent + pub spent_at_height: Option, // The height at which this note was spent pub unconfirmed_spent: Option, // If this note was spent in a send, but has not yet been confirmed. pub memo: Option, pub is_change: bool, @@ -107,7 +108,7 @@ pub fn read_note(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(mut reader: R) -> io::Result { - let _version = reader.read_u64::()?; + let version = reader.read_u64::()?; let account = reader.read_u64::()? 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::())? + } 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::(*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})?;