diff --git a/third_party/silentdragonxlite/silentdragonxlite-cli/lib/src/lightwallet.rs b/third_party/silentdragonxlite/silentdragonxlite-cli/lib/src/lightwallet.rs index f538111..1606e0e 100644 --- a/third_party/silentdragonxlite/silentdragonxlite-cli/lib/src/lightwallet.rs +++ b/third_party/silentdragonxlite/silentdragonxlite-cli/lib/src/lightwallet.rs @@ -442,7 +442,7 @@ impl LightWallet { let birthday = reader.read_u64::()?; - Ok(LightWallet{ + let wallet = LightWallet{ encrypted: encrypted, unlocked: !encrypted, // When reading from disk, if wallet is encrypted, it starts off locked. enc_seed: enc_seed, @@ -457,7 +457,16 @@ impl LightWallet { config: config.clone(), birthday, total_scan_duration: Arc::new(RwLock::new(vec![Duration::new(0, 0)])), - }) + }; + + // Existing wallets carry up to 100 witnesses per long-spent note; drop them now so the next + // save is small. Same predicate as the per-block call in scan_block_with_pool. + let cleared = wallet.clear_old_witnesses(); + if cleared > 0 { + info!("Dropped witness history of {} notes whose spend is deeper than {} blocks", cleared, MAX_REORG); + } + + Ok(wallet) } pub fn write(&self, mut writer: W) -> io::Result<()> { @@ -1753,6 +1762,64 @@ pub fn scan_full_mempool_tx(&self, tx: &Transaction, height: i32, _datetime: u64 } } + /// The height below which a confirmed spend can no longer be rolled back, or None while the block + /// window is not full. When the window holds > MAX_REORG blocks, every retained block is at or above + /// the returned height. An invalidate_block() that removes a spender whose WalletTx.block is below it + /// therefore pops the whole window, and that already empties every note's witness list + /// (each list holds <= MAX_REORG entries, and the window holds > MAX_REORG blocks). + fn witness_clear_height(&self) -> Option { + let blks = match self.blocks.read() { + Ok(b) => b, + Err(p) => p.into_inner(), + }; + if blks.len() > MAX_REORG { + blks.first().map(|b| b.height) + } else { + None + } + } + + /// Drop the witness history of notes whose confirmed spender (keyed by WalletTx.block, the same + /// field invalidate_block uses to decide removal) is below `clear_below`. Returns the number cleared. + fn clear_spent_witnesses_below(txs: &mut HashMap, clear_below: i32) -> usize { + let mut to_clear: Vec<(TxId, usize)> = Vec::new(); + for (txid, tx) in txs.iter() { + for (i, nd) in tx.notes.iter().enumerate() { + if nd.witnesses.is_empty() { + continue; + } + if let Some(spend_txid) = nd.spent { + if let Some(spend_tx) = txs.get(&spend_txid) { + if spend_tx.block < clear_below { + to_clear.push((*txid, i)); + } + } + } + } + } + for (txid, i) in &to_clear { + if let Some(tx) = txs.get_mut(txid) { + // Vec::new(), not clear(): release the buffer. + tx.notes[*i].witnesses = Vec::new(); + } + } + to_clear.len() + } + + /// Clear the history of long-spent notes (used right after loading a wallet). + pub fn clear_old_witnesses(&self) -> usize { + match self.witness_clear_height() { + None => 0, + Some(h) => { + let mut txs = match self.txs.write() { + Ok(t) => t, + Err(p) => p.into_inner(), + }; + LightWallet::clear_spent_witnesses_below(&mut *txs, h) + } + } + } + // Invalidate all blocks including and after "at_height". // Returns the number of blocks invalidated pub fn invalidate_block(&self, at_height: i32) -> u64 { @@ -2137,6 +2204,9 @@ pub fn scan_full_mempool_tx(&self, tx: &Transaction, height: i32, _datetime: u64 // These are filled in inside the block let new_txs; let nfs: Vec<_>; + + // Read before taking the txs lock (no lock nesting). The window ends at height-1 here. + let witness_clear_below = self.witness_clear_height(); { // Create a write lock let mut txs = self.txs.write().unwrap(); @@ -2159,6 +2229,11 @@ pub fn scan_full_mempool_tx(&self, tx: &Transaction, height: i32, _datetime: u64 .flatten() .collect(); + // Free the witness history of notes whose spend can no longer be rolled back. + if let Some(h) = witness_clear_below { + LightWallet::clear_spent_witnesses_below(&mut *txs, h); + } + // Prepare the note witnesses for updating for tx in txs.values_mut() { for nd in tx.notes.iter_mut() { @@ -2656,3 +2731,6 @@ pub mod tests; #[cfg(test)] mod oracle_tests; + +#[cfg(test)] +mod witness_tests; diff --git a/third_party/silentdragonxlite/silentdragonxlite-cli/lib/src/lightwallet/oracle_tests.rs b/third_party/silentdragonxlite/silentdragonxlite-cli/lib/src/lightwallet/oracle_tests.rs index cf3062d..d089b71 100644 --- a/third_party/silentdragonxlite/silentdragonxlite-cli/lib/src/lightwallet/oracle_tests.rs +++ b/third_party/silentdragonxlite/silentdragonxlite-cli/lib/src/lightwallet/oracle_tests.rs @@ -411,6 +411,10 @@ struct Sim { step: usize, op: &'static str, cleared_seen: Cell, + /// The wallet drops the history of a note whose spender is below a full block window, at each + /// block scan and at each read. A later partial reorg leaves it dropped, so the highest such + /// bound seen is kept (the main wallet's scans and reloads; a copy's read is the current window). + clear_bound: i32, stats: Stats, } @@ -487,6 +491,7 @@ impl Sim { step: 0, op: "init", cleared_seen: Cell::new(0), + clear_bound: i32::MIN, stats: Stats::default(), }; sim.check(true); @@ -866,6 +871,10 @@ impl Sim { while self.blocks.front().unwrap().height < h - KEEP { self.blocks.pop_front(); } + // The wallet scanned this block with the window as it was before it (see clear_bound). + if self.win.len() > MAX_REORG { + self.clear_bound = cmp::max(self.clear_bound, *self.win.front().unwrap()); + } self.win.push_back(h); while self.win.len() > MAX_REORG + 1 { self.win.pop_front(); @@ -1021,6 +1030,9 @@ impl Sim { w3.write(&mut b3).unwrap(); assert!(b2 == b3, "{} write(read(x)) is not idempotent", self.ctx()); self.w = w3; + if self.win.len() > MAX_REORG { + self.clear_bound = cmp::max(self.clear_bound, *self.win.front().unwrap()); + } } /// Build a send. Probe: on a throwaway copy with the given anchor offset. Real: on the @@ -1225,6 +1237,7 @@ impl Sim { let win_first = *self.win.front().unwrap(); let window_full = self.win.len() > MAX_REORG; + let clear_bound = if window_full { cmp::max(self.clear_bound, win_first) } else { self.clear_bound }; let mut cleared = 0; for (idx, n) in self.notes.iter().enumerate() { let (txid, nd) = by_nf.get(&n.nf) @@ -1244,7 +1257,7 @@ impl Sim { if let Some((_, spender_block)) = n.spent { // A spent note keeps its legacy list, or (once the spend can no longer be // rolled back) an empty one. - let cleared_ok = l == 0 && window_full && spender_block < win_first; + let cleared_ok = l == 0 && spender_block < clear_bound; assert!(l == n.len || cleared_ok, "{} spent note has {} witnesses, expected {}", nctx, l, n.len); if l != n.len { cleared += 1; diff --git a/third_party/silentdragonxlite/silentdragonxlite-cli/lib/src/lightwallet/witness_tests.rs b/third_party/silentdragonxlite/silentdragonxlite-cli/lib/src/lightwallet/witness_tests.rs new file mode 100644 index 0000000..c68f2de --- /dev/null +++ b/third_party/silentdragonxlite/silentdragonxlite-cli/lib/src/lightwallet/witness_tests.rs @@ -0,0 +1,698 @@ +// Copyright The Hush Developers 2019-2022 +// Released under the GPLv3 + +//! Witness-history tests for the Tier-6 g2 changes (audit C232/C233). +//! +//! `ChainModel` is the OLD (pre-g2) witness algorithm, written independently of the wallet: +//! per block, every note's newest witness is duplicated and its list trimmed to 100 entries; +//! then each leaf of the block is appended to the newest witness of every unspent note, to the +//! witnesses of our notes created earlier in the same block, and to the tree. A k-block +//! invalidate pops k entries from every list. The wallet must equal the model entry for entry, +//! with one declared exception: a note whose confirmed spender is below a full +//! (MAX_REORG + 1)-block window may lose its history, and then must have lost it. +//! +//! `check_anchor_alignment` (T0) is the spend-validity oracle: the k-th newest witness of +//! every live note must have the root of the k-th newest retained block, which is exactly the +//! anchor a spend at anchor offset k proves against. + +use std::cmp; +use std::convert::TryInto; + +use ff::{Field, PrimeField, PrimeFieldRepr}; +use pairing::bls12_381::{Bls12, Fr}; +use protobuf::Message; +use rand::{rngs::OsRng, RngCore}; + +use zcash_client_backend::proto::compact_formats::{CompactBlock, CompactOutput, CompactSpend, CompactTx}; +use zcash_primitives::{ + block::BlockHash, + jubjub::fs::Fs, + merkle_tree::{CommitmentTree, IncrementalWitness}, + note_encryption::{Memo, SaplingNoteEncryption}, + primitives::Note, + sapling::Node, + transaction::TxId, + zip32::{ExtendedFullViewingKey, ExtendedSpendingKey}, + JUBJUB, +}; + +use super::data::SaplingNoteData; +use super::{LightWallet, MAX_REORG}; +use crate::lightclient::LightClientConfig; + +fn cfg() -> LightClientConfig { + LightClientConfig { + server: "0.0.0.0:0".parse().unwrap(), + chain_name: "main".to_string(), // "test" panics in base58_pubkey_address + sapling_activation_height: 0, + consensus_branch_id: "000000".to_string(), + anchor_offset: 0, + no_cert_verification: false, + data_dir: None, // a LightWallet never touches the disk itself + } +} + +fn wbytes(w: &IncrementalWitness) -> Vec { + let mut v = vec![]; + w.write(&mut v).unwrap(); + v +} + +fn tbytes(t: &CommitmentTree) -> Vec { + let mut v = vec![]; + t.write(&mut v).unwrap(); + v +} + +fn random_hash() -> [u8; 32] { + let mut h = [0u8; 32]; + OsRng.fill_bytes(&mut h); + h +} + +fn foreign() -> ExtendedFullViewingKey { + ExtendedFullViewingKey::from(&ExtendedSpendingKey::master(&[7u8; 32])) +} + +/// An output paying `value` to the default address of `fvk` (the body of tests.rs add_tx_paying_to). +fn out_to(fvk: &ExtendedFullViewingKey, value: u64) -> CompactOutput { + let mut rng = OsRng; + let to = fvk.default_address().unwrap().1; + let note = Note { + g_d: to.diversifier.g_d::(&JUBJUB).unwrap(), + pk_d: to.pk_d.clone(), + value, + r: Fs::random(&mut rng), + }; + let encryptor = SaplingNoteEncryption::new(fvk.fvk.ovk, note.clone(), to.clone(), Memo::default(), &mut rng); + let mut cmu = vec![]; + note.cm(&JUBJUB).into_repr().write_le(&mut cmu).unwrap(); + let mut epk = vec![]; + encryptor.epk().write(&mut epk).unwrap(); + let enc_ciphertext = encryptor.encrypt_note_plaintext(); + + let mut cout = CompactOutput::new(); + cout.set_cmu(cmu); + cout.set_epk(epk); + cout.set_ciphertext(enc_ciphertext[..52].to_vec()); + cout +} + +fn random_cmu_bytes() -> Vec { + let mut v = vec![]; + Fr::random(&mut OsRng).into_repr().write_le(&mut v).unwrap(); + v +} + +/// A valid cmu with a 31-byte epk: the leaf is committed, decryption is skipped. +#[allow(dead_code)] +fn bad_epk() -> CompactOutput { + let mut o = CompactOutput::new(); + o.set_cmu(random_cmu_bytes()); + o.set_epk(vec![1u8; 31]); + o.set_ciphertext(vec![0u8; 52]); + o +} + +/// A 31-byte cmu: nothing is appended anywhere. +#[allow(dead_code)] +fn bad_cmu() -> CompactOutput { + let mut o = CompactOutput::new(); + o.set_cmu(vec![1u8; 31]); + o.set_epk(vec![0u8; 32]); + o.set_ciphertext(vec![0u8; 52]); + o +} + +fn tx_with(outputs: Vec, spend_nfs: Vec>) -> (CompactTx, TxId) { + let hash = random_hash(); + let mut ctx = CompactTx::new(); + ctx.set_hash(hash.to_vec()); + for nf in spend_nfs { + let mut s = CompactSpend::new(); + s.set_nf(nf); + ctx.spends.push(s); + } + for o in outputs { + ctx.outputs.push(o); + } + (ctx, TxId(hash)) +} + +/// A deterministic schedule (the note contents themselves use OsRng). +struct Lcg(u64); + +impl Lcg { + fn below(&mut self, n: u64) -> u64 { + self.0 = self.0.wrapping_mul(6364136223846793005).wrapping_add(1442695040888963407); + (self.0 >> 33) % n + } +} + +/// A tx for `Chain::mine`: its compact form, which of its outputs pays the wallet, and the +/// creating txids of the model notes it spends. +struct GenTx { + ctx: CompactTx, + txid: TxId, + ours: Option, + spends: Vec, +} + +struct ModelNote { + /// Creating txid (every test tx pays the wallet at most one note). + txid: TxId, + recv: i32, + /// Block of the spending tx. + spent_at: Option, + hist: Vec>, +} + +/// The OLD witness algorithm (see the module comment). +struct ChainModel { + tree: CommitmentTree, + /// Retained window: (height, tree at the end of that block), at most MAX_REORG + 1 like the wallet. + win: Vec<(i32, CommitmentTree)>, + notes: Vec, + /// A spent note whose spender block is below this may (and must) have lost its history: + /// the max first height of a full window seen by a scan (before the block) or by a read. + clear_bound: i32, +} + +impl ChainModel { + fn new() -> Self { + ChainModel { tree: CommitmentTree::new(), win: vec![], notes: vec![], clear_bound: i32::MIN } + } + + fn note(&self, txid: &TxId) -> &ModelNote { + self.notes.iter().find(|n| n.txid == *txid).expect("model note") + } + + fn note_full_window(&mut self) { + if self.win.len() > MAX_REORG { + self.clear_bound = cmp::max(self.clear_bound, self.win[0].0); + } + } + + fn must_be_cleared(&self, n: &ModelNote) -> bool { + n.spent_at.map_or(false, |s| s < self.clear_bound) + } + + /// nodes: the block's leaves in order (outputs whose cmu parses); ours: (index into nodes, + /// creating txid) of the outputs paying the wallet; spends: creating txids of spent notes. + fn apply_block(&mut self, height: i32, nodes: &[Node], ours: &[(usize, TxId)], spends: &[TxId]) { + for n in self.notes.iter_mut() { + if let Some(w) = n.hist.last() { + let c = w.clone(); + n.hist.push(c); + } + n.hist = n.hist.split_off(n.hist.len().saturating_sub(100)); + } + let active: Vec = (0..self.notes.len()).filter(|i| self.notes[*i].spent_at.is_none()).collect(); + let mut fresh: Vec = vec![]; + for (k, node) in nodes.iter().enumerate() { + for i in &active { + if let Some(w) = self.notes[*i].hist.last_mut() { + w.append(*node).unwrap(); + } + } + for f in fresh.iter_mut() { + f.hist[0].append(*node).unwrap(); + } + self.tree.append(*node).unwrap(); + if let Some((_, txid)) = ours.iter().find(|(idx, _)| *idx == k) { + fresh.push(ModelNote { + txid: *txid, + recv: height, + spent_at: None, + hist: vec![IncrementalWitness::from_tree(&self.tree)], + }); + } + } + for t in spends { + for n in self.notes.iter_mut() { + if n.txid == *t { + n.spent_at = Some(height); + } + } + } + self.notes.extend(fresh); + self.win.push((height, self.tree.clone())); + if self.win.len() > MAX_REORG + 1 { + let d = self.win.len() - (MAX_REORG + 1); + self.win.drain(..d); + } + } + + /// OLD invalidate_block: pop the retained blocks at or above `at`, drop the notes created + /// there, un-spend the notes spent there, and pop the same count from every list. + fn invalidate(&mut self, at: i32) -> u64 { + let mut k = 0usize; + while self.win.last().map_or(false, |(h, _)| *h >= at) { + self.win.pop(); + k += 1; + } + self.tree = match self.win.last() { + Some((_, t)) => t.clone(), + None => CommitmentTree::new(), + }; + self.notes.retain(|n| n.recv < at); + for n in self.notes.iter_mut() { + if n.spent_at.map_or(false, |s| s >= at) { + n.spent_at = None; + } + let keep = n.hist.len().saturating_sub(k); + n.hist.truncate(keep); + } + k as u64 + } +} + +struct Chain { + wallet: LightWallet, + model: ChainModel, + prev: BlockHash, + next: i32, +} + +impl Chain { + fn new() -> Self { + Chain { + wallet: LightWallet::new(None, &cfg(), 0, 0).unwrap(), + model: ChainModel::new(), + prev: BlockHash([0u8; 32]), + next: 0, + } + } + + fn ours(&self) -> ExtendedFullViewingKey { + self.wallet.zkeys.read().unwrap()[0].extfvk.clone() + } + + /// The stored nullifier of the wallet note created by `txid` (its true tree position). + fn nf(&self, txid: &TxId) -> Vec { + self.wallet.txs.read().unwrap()[txid].notes[0].nullifier.to_vec() + } + + fn pay_us(&self, value: u64) -> GenTx { + let (ctx, txid) = tx_with(vec![out_to(&self.ours(), value)], vec![]); + GenTx { ctx, txid, ours: Some(0), spends: vec![] } + } + + fn pay_foreign(&self, n: usize) -> GenTx { + let f = foreign(); + let (ctx, txid) = tx_with((0..n).map(|i| out_to(&f, 5_000 + i as u64)).collect(), vec![]); + GenTx { ctx, txid, ours: None, spends: vec![] } + } + + /// Spend the note created by `note`: [foreign payment, change to us]. + fn spend(&self, note: TxId, change: u64) -> GenTx { + let (ctx, txid) = tx_with(vec![out_to(&foreign(), 7_000), out_to(&self.ours(), change)], vec![self.nf(¬e)]); + GenTx { ctx, txid, ours: Some(1), spends: vec![note] } + } + + fn mine(&mut self, gtxs: Vec) -> CompactBlock { + let height = self.next; + let mut cb = CompactBlock::new(); + cb.set_height(height as u64); + cb.hash = random_hash().to_vec(); + cb.prevHash = self.prev.0.to_vec(); + let mut nodes: Vec = vec![]; + let mut ours: Vec<(usize, TxId)> = vec![]; + let mut spends: Vec = vec![]; + for (i, g) in gtxs.into_iter().enumerate() { + for (j, o) in g.ctx.outputs.iter().enumerate() { + if let Ok(cmu) = o.cmu() { + if g.ours == Some(j) { + ours.push((nodes.len(), g.txid)); + } + nodes.push(Node::new(cmu.into_repr())); + } + } + spends.extend(g.spends.iter().cloned()); + let mut ctx = g.ctx; + ctx.set_index(i as u64); + cb.vtx.push(ctx); + } + self.model.note_full_window(); + let r = self.wallet.scan_block(&cb.write_to_bytes().unwrap()); + assert!(r.is_ok(), "scan_block({}) failed: {:?}", height, r); + self.model.apply_block(height, &nodes, &ours, &spends); + self.prev = BlockHash(cb.hash[..].try_into().unwrap()); + self.next += 1; + cb + } + + fn rewind(&mut self) { + self.next = self.wallet.last_scanned_height() + 1; + self.prev = self.wallet.blocks.read().unwrap().last().map(|b| b.hash).unwrap_or(BlockHash([0u8; 32])); + } + + fn invalidate(&mut self, at: i32) -> u64 { + let n = self.wallet.invalidate_block(at); + let m = self.model.invalidate(at); + assert_eq!(n, m, "invalidate_block({}) count vs the model", at); + self.rewind(); + n + } + + /// Replace the wallet by `read(write(wallet))`. Returns (first write, write of the re-read wallet). + fn reload(&mut self) -> (Vec, Vec) { + let mut b1 = vec![]; + self.wallet.write(&mut b1).unwrap(); + self.wallet = LightWallet::read(&b1[..], &cfg()).unwrap(); + self.model.note_full_window(); + let mut b2 = vec![]; + self.wallet.write(&mut b2).unwrap(); + (b1, b2) + } + + fn list_len(&self, txid: &TxId) -> usize { + self.wallet.txs.read().unwrap()[txid].notes[0].witnesses.len() + } + + /// Compare the wallet with the model: the block window and trees, every note's spent state, + /// and its witness list (`full`: every entry's bytes; otherwise the count and the newest entry). + /// Returns how many notes have the history the model requires to be dropped. + fn compare(&self, full: bool, ctx: &str) -> usize { + { + let blocks = self.wallet.blocks.read().unwrap(); + let hs: Vec = blocks.iter().map(|b| b.height).collect(); + let mh: Vec = self.model.win.iter().map(|(h, _)| *h).collect(); + assert_eq!(hs, mh, "{}: retained block window", ctx); + if let Some(b) = blocks.last() { + assert!(tbytes(&b.tree) == tbytes(&self.model.tree), "{}: tip tree differs from the model", ctx); + } + if full { + for (b, (h, t)) in blocks.iter().zip(self.model.win.iter()) { + assert!(tbytes(&b.tree) == tbytes(t), "{}: tree of block {} differs from the model", ctx, h); + } + } + } + let txs = self.wallet.txs.read().unwrap(); + let total: usize = txs.values().map(|t| t.notes.len()).sum(); + assert_eq!(total, self.model.notes.len(), "{}: number of wallet notes", ctx); + let mut cleared = 0; + for m in &self.model.notes { + let wtx = txs.get(&m.txid).unwrap_or_else(|| panic!("{}: no wallet tx {}", ctx, m.txid)); + assert_eq!(wtx.notes.len(), 1, "{}: notes of tx {}", ctx, m.txid); + assert_eq!(wtx.block, m.recv, "{}: block of tx {}", ctx, m.txid); + let nd = &wtx.notes[0]; + match (nd.spent, m.spent_at) { + (None, None) => {} + (Some(s), Some(h)) => { + assert_eq!(txs.get(&s).map(|t| t.block), Some(h), "{}: spender block of note {}", ctx, m.txid) + } + (a, b) => panic!("{}: note {} spent {:?}, model spent at {:?}", ctx, m.txid, a, b), + } + assert!(nd.unconfirmed_spent.is_none(), "{}: note {} pending", ctx, m.txid); + let l = nd.witnesses.len(); + assert!(l <= MAX_REORG, "{}: note {} has {} witnesses", ctx, m.txid, l); + if self.model.must_be_cleared(m) { + assert!(l == 0, "{}: note {} (spent at {:?}, clear bound {}) still has {} witnesses", + ctx, m.txid, m.spent_at, self.model.clear_bound, l); + cleared += 1; + continue; + } + assert_eq!(l, m.hist.len(), "{}: witness count of note {} (received {}, spent at {:?})", + ctx, m.txid, m.recv, m.spent_at); + let from = if full { 0 } else { l.saturating_sub(1) }; + for i in from..l { + assert!(wbytes(&nd.witnesses[i]) == wbytes(&m.hist[i]), + "{}: witness {} of {} of note {} (received {}, spent at {:?}) differs from the model", + ctx, i, l, m.txid, m.recv, m.spent_at); + } + } + cleared + } +} + +/// T0: for every live note (L entries, B retained blocks) assert L <= 100, L <= B and, for k in +/// {0, 1, L-1}, witnesses[L-1-k].root() == blocks[B-1-k].tree.root(): the anchor at offset k. +fn check_anchor_alignment(w: &LightWallet, ctx: &str) -> usize { + let blocks = w.blocks.read().unwrap(); + let b = blocks.len(); + let txs = w.txs.read().unwrap(); + let mut checked = 0; + for tx in txs.values() { + for nd in tx.notes.iter() { + if nd.spent.is_some() || nd.unconfirmed_spent.is_some() { + continue; + } + let l = nd.witnesses.len(); + assert!(l <= MAX_REORG && l <= b, "{}: live note of tx {} has {} witnesses, {} blocks", ctx, tx.txid, l, b); + let mut ks = vec![0usize, 1, l.saturating_sub(1)]; + ks.dedup(); + for k in ks { + if k < l { + assert!(nd.witnesses[l - 1 - k].root() == blocks[b - 1 - k].tree.root(), + "{}: live note of tx {}: witness at offset {} is not the root of block {}", + ctx, tx.txid, k, blocks[b - 1 - k].height); + checked += 1; + } + } + } + } + checked +} + +/// Test 1's chain, heights 0..=259: block 0 pays note A; about 25% of blocks pay us 1-3 notes and +/// about 20% carry foreign outputs; at 40, 120 and 200 a tx spends the oldest live note (A first) +/// with [foreign, change to us]. With `check`, the wallet is compared with the model after every +/// block (all entries plus T0 and the roots at the checkpoints). Returns the (height, note) spends. +fn t1_chain(check: bool) -> (Chain, Vec<(i32, TxId)>) { + let mut c = Chain::new(); + let mut rng = Lcg(0x5eed_0001); + let mut spent: Vec<(i32, TxId)> = vec![]; + let checkpoints = [41, 100, 101, 141, 142, 143, 221, 222, 259]; + for h in 0..=259 { + let mut txs = vec![]; + if h == 0 { + txs.push(c.pay_us(100_000)); + } + if h == 40 || h == 120 || h == 200 { + let note = c.model.notes.iter() + .find(|n| n.spent_at.is_none()) + .map(|n| n.txid) + .expect("a live note to spend"); + txs.push(c.spend(note, 1_000)); + spent.push((h, note)); + } + if h > 0 && rng.below(100) < 25 { + for _ in 0..1 + rng.below(3) { + let v = 1_000 + rng.below(1_000); + txs.push(c.pay_us(v)); + } + } + if rng.below(100) < 20 { + let n = 1 + rng.below(2) as usize; + txs.push(c.pay_foreign(n)); + } + c.mine(txs); + if !check { + continue; + } + let ctx = format!("t1 h={}", h); + let full = checkpoints.contains(&h); + c.compare(full, &ctx); + // The clearing rule, independently of the model: a note spent at s is empty after + // block h iff h >= 101 && s <= h - 102. + for (s, note) in &spent { + let expect_empty = h >= 101 && *s <= h - 102; + assert_eq!(c.list_len(note) == 0, expect_empty, "{}: note spent at {} cleared?", ctx, s); + } + if full { + check_anchor_alignment(&c.wallet, &ctx); + let root = c.wallet.blocks.read().unwrap().last().unwrap().tree.root(); + assert!(root == c.model.tree.root(), "{}: tip root differs from the model", ctx); + } + } + assert_eq!(spent[0].1, c.model.notes[0].txid, "the first spend must spend note A"); + (c, spent) +} + +#[test] +fn witness_history_matches_old_algorithm() { + let (c, spent) = t1_chain(true); + assert_eq!(spent.iter().map(|(s, _)| *s).collect::>(), vec![40, 120, 200]); + assert_eq!(c.compare(true, "t1 end"), 2, "notes spent at 40 and 120 are cleared at 259, 200 is not"); + assert!(c.model.notes.len() > 40, "only {} notes", c.model.notes.len()); + let live_at_cap = c.model.notes.iter().filter(|n| n.spent_at.is_none() && n.hist.len() == MAX_REORG).count(); + assert!(live_at_cap > 10, "only {} live notes at the 100-entry cap", live_at_cap); +} + +/// Test 2's chain: block 0 pays A, every 4th block pays us, some blocks carry foreign outputs, +/// block 10 spends A. Mined up to `tip`, compared with the model after every block; also pins +/// witness_clear_height() (None while <= 100 blocks are retained, then the window's first height) +/// and the block at which A is cleared (112). +fn t2_chain(tip: i32) -> (Chain, TxId) { + let mut c = Chain::new(); + let mut a = None; + for h in 0..=tip { + let mut txs = vec![]; + if h == 0 { + let g = c.pay_us(50_000); + a = Some(g.txid); + txs.push(g); + } + if h == 10 { + txs.push(c.spend(a.unwrap(), 500)); + } + if h % 4 == 1 { + txs.push(c.pay_us(1_000 + h as u64)); + } + if h % 7 == 3 { + txs.push(c.pay_foreign(1)); + } + c.mine(txs); + let ctx = format!("t2 h={}", h); + c.compare(h % 25 == 0 || h == tip, &ctx); + let expect = if h >= MAX_REORG as i32 { Some(h - MAX_REORG as i32) } else { None }; + assert_eq!(c.wallet.witness_clear_height(), expect, "{}: witness_clear_height", ctx); + if h >= 10 { + assert_eq!(c.list_len(&a.unwrap()) == 0, h >= 112, "{}: note A cleared?", ctx); + } + } + (c, a.unwrap()) +} + +#[test] +fn clear_then_reorg_matches_old_algorithm() { + // (a) Tip 60: the window (61 blocks) is not full, nothing is cleared. + let (mut c, a) = t2_chain(60); + assert_eq!(c.compare(true, "2a tip 60"), 0); + assert_eq!(c.list_len(&a), 61); + assert_eq!(c.wallet.clear_old_witnesses(), 0); + assert_eq!(c.invalidate(10), 51); + assert!(c.wallet.txs.read().unwrap()[&a].notes[0].spent.is_none(), "A is unspent again"); + assert_eq!(c.list_len(&a), 10); + c.compare(true, "2a after invalidate(10)"); + assert!(check_anchor_alignment(&c.wallet, "2a after invalidate(10)") > 0); + // A new branch: the revived note advances like every other live note. + for h in 0..6 { + let txs = if h == 2 { vec![c.pay_us(4_242)] } else { vec![c.pay_foreign(1)] }; + c.mine(txs); + c.compare(true, "2a new branch"); + check_anchor_alignment(&c.wallet, "2a new branch"); + } + assert_eq!(c.list_len(&a), 16); + + // (b) Tip 115: A was cleared at 112. invalidate(10) pops the whole window; A is unspent and + // empty, exactly like the old algorithm's trimmed list. + let (mut c, a) = t2_chain(115); + assert_eq!(c.compare(true, "2b tip 115"), 1); + assert_eq!(c.invalidate(10), 101); + assert!(c.wallet.txs.read().unwrap()[&a].notes[0].spent.is_none(), "A is unspent again"); + assert_eq!(c.list_len(&a), 0); + assert!(c.model.note(&a).hist.is_empty(), "the old algorithm also holds an empty list"); + assert_eq!(c.compare(true, "2b after invalidate(10)"), 0); + check_anchor_alignment(&c.wallet, "2b after invalidate(10)"); + + // (c) Tip 115, a partial reorg: A stays spent and empty, live notes equal the model trimmed + // by 36; then the rest of the window (65 blocks) is popped. + let (mut c, a) = t2_chain(115); + assert_eq!(c.invalidate(80), 36); + assert!(c.wallet.txs.read().unwrap()[&a].notes[0].spent.is_some(), "A is still spent"); + assert_eq!(c.list_len(&a), 0); + assert_eq!(c.compare(true, "2c after invalidate(80)"), 1); + assert!(check_anchor_alignment(&c.wallet, "2c after invalidate(80)") > 0); + assert_eq!(c.invalidate(10), 65); + assert_eq!(c.list_len(&a), 0); + assert!(c.model.note(&a).hist.is_empty()); + assert_eq!(c.compare(true, "2c after invalidate(10)"), 0); +} + +#[test] +fn witness_serialization_compat() { + // (a) write -> read -> write is byte-identical (clearing on read is idempotent). + let (mut c, spent) = t1_chain(false); + assert_eq!(c.compare(true, "3a built"), 2); + let (b1, b2) = c.reload(); + assert!(b1 == b2, "write(read(write(w))) differs from write(w)"); + assert_eq!(c.wallet.clear_old_witnesses(), 0); + assert_eq!(c.compare(true, "3a reloaded"), 2); + check_anchor_alignment(&c.wallet, "3a reloaded"); + + // (b) A legacy file: the cleared notes carry the old algorithm's full lists. Reading it + // drops them again; everything else is unchanged, so the re-write equals the new file. + let cleared: Vec = spent.iter().map(|(_, t)| *t).filter(|t| c.list_len(t) == 0).collect(); + assert_eq!(cleared.len(), 2); + for t in &cleared { + let legacy = { + let txs = c.wallet.txs.read().unwrap(); + let mut b = vec![]; + txs[t].notes[0].write(&mut b).unwrap(); + // version, account, extfvk (169), diversifier (11), value, rcm (32), then the witness Vector + let at = 8 + 8 + 169 + 11 + 8 + 32; + assert_eq!(b[at], 0, "an empty witness Vector (CompactSize 0) at offset {}", at); + let hist = &c.model.note(t).hist; + assert_eq!(hist.len(), MAX_REORG); + let mut v = b[..at].to_vec(); + v.push(hist.len() as u8); // CompactSize below 253 is one byte + for w in hist.iter() { + v.extend(wbytes(w)); + } + v.extend_from_slice(&b[at + 1..]); + v + }; + let nd = SaplingNoteData::read(&legacy[..]).unwrap(); + assert_eq!(nd.witnesses.len(), MAX_REORG); + let mut again = vec![]; + nd.write(&mut again).unwrap(); + assert!(again == legacy, "the legacy note record does not round-trip"); + c.wallet.txs.write().unwrap().get_mut(t).unwrap().notes[0] = nd; + assert_eq!(c.list_len(t), MAX_REORG); + } + let (legacy_file, rewritten) = c.reload(); + for t in &cleared { + assert_eq!(c.list_len(t), 0, "reading a legacy file clears the long-spent note again"); + } + assert!(rewritten == b1, "the legacy file re-writes to exactly the new file"); + assert!(rewritten.len() < legacy_file.len(), "{} vs {}", rewritten.len(), legacy_file.len()); + assert_eq!(c.compare(true, "3b reloaded"), 2); + + // (c) A 50-block chain (window not full): the spent note keeps its list through a reload. + let (mut c50, a) = { + let mut c = Chain::new(); + let mut a = None; + for h in 0..50 { + let mut txs = vec![]; + if h == 0 { + let g = c.pay_us(30_000); + a = Some(g.txid); + txs.push(g); + } + if h == 10 { + txs.push(c.spend(a.unwrap(), 300)); + } + if h % 3 == 2 { + txs.push(c.pay_us(2_000 + h as u64)); + } + c.mine(txs); + } + (c, a.unwrap()) + }; + assert_eq!(c50.list_len(&a), 50); + let (x1, x2) = c50.reload(); + assert!(x1 == x2); + assert_eq!(c50.list_len(&a), 50); + assert_eq!(c50.compare(true, "3c reloaded"), 0); + + // (d) A pending note and a note spent by a tx the wallet does not know are never cleared. + let (x, y) = { + let live: Vec = c.model.notes.iter().filter(|n| n.spent_at.is_none()).map(|n| n.txid).take(2).collect(); + (live[0], live[1]) + }; + let before = |c: &Chain, t: &TxId| -> Vec> { + c.wallet.txs.read().unwrap()[t].notes[0].witnesses.iter().map(|w| wbytes(w)).collect() + }; + let (bx, by) = (before(&c, &x), before(&c, &y)); + assert!(!bx.is_empty() && !by.is_empty()); + { + let mut txs = c.wallet.txs.write().unwrap(); + txs.get_mut(&x).unwrap().notes[0].unconfirmed_spent = Some(TxId(random_hash())); + txs.get_mut(&y).unwrap().notes[0].spent = Some(TxId(random_hash())); + } + assert!(c.wallet.witness_clear_height().is_some(), "the window is full"); + assert_eq!(c.wallet.clear_old_witnesses(), 0); + assert!(before(&c, &x) == bx && before(&c, &y) == by, "pending / unknown-spender notes were touched"); +}