diff --git a/third_party/silentdragonxlite/silentdragonxlite-cli/lib/src/lightclient.rs b/third_party/silentdragonxlite/silentdragonxlite-cli/lib/src/lightclient.rs index 8b76b29..ebd2e71 100644 --- a/third_party/silentdragonxlite/silentdragonxlite-cli/lib/src/lightclient.rs +++ b/third_party/silentdragonxlite/silentdragonxlite-cli/lib/src/lightclient.rs @@ -659,13 +659,14 @@ impl LightClient { let wallet = self.wallet.read().unwrap(); // Collect z addresses - let z_addresses = wallet.get_all_zaddresses().iter().map(|zaddress| { + let z_addresses = wallet.z_balances_by_address().into_iter() + .map(|(address, zbalance, unconfirmed, verified_zbalance, spendable_zbalance)| { object!{ - "address" => zaddress.clone(), - "zbalance" => wallet.zbalance(Some(zaddress.clone())), - "unconfirmed" => wallet.unconfirmed_zbalance(Some(zaddress.clone())), - "verified_zbalance" => wallet.verified_zbalance(Some(zaddress.clone())), - "spendable_zbalance" => wallet.spendable_zbalance(Some(zaddress.clone())) + "address" => address, + "zbalance" => zbalance, + "unconfirmed" => unconfirmed, + "verified_zbalance" => verified_zbalance, + "spendable_zbalance" => spendable_zbalance } }).collect::>(); @@ -847,15 +848,13 @@ impl LightClient { let wallet = self.wallet.read().unwrap(); // First, collect all extfvk's that are spendable (i.e., we have the private key) - let spendable_address: HashSet = wallet.get_all_zaddresses().iter() - .filter(|address| wallet.have_spending_key_for_zaddress(address)) - .map(|address| address.clone()) - .collect(); + let spendable_address_set: HashSet = wallet.spendable_zaddress_set(); + let spendable_address: &HashSet = &spendable_address_set; // Collect Sapling notes wallet.txs.read().unwrap().iter() .flat_map( |(txid, wtx)| { - let spendable_address = spendable_address.clone(); + let spendable_address: &HashSet = spendable_address; wtx.notes.iter().filter_map(move |nd| if !all_notes && nd.spent.is_some() { None diff --git a/third_party/silentdragonxlite/silentdragonxlite-cli/lib/src/lightwallet.rs b/third_party/silentdragonxlite/silentdragonxlite-cli/lib/src/lightwallet.rs index d605c9e..f538111 100644 --- a/third_party/silentdragonxlite/silentdragonxlite-cli/lib/src/lightwallet.rs +++ b/third_party/silentdragonxlite/silentdragonxlite-cli/lib/src/lightwallet.rs @@ -105,6 +105,33 @@ impl ToBase58Check for [u8] { } } +/// Which notes a z-balance sum counts. `Addr` compares (diversifier, pk_d) directly. Because +/// encode_payment_address is injective on (diversifier, pk_d), it selects exactly the notes whose +/// old derived string == encode(pa). `Str` keeps the old exact string compare (on the cheap +/// note_address) for arbitrary caller strings, so no decoding and no new panic paths. +enum ZAddrFilter { + All, + Addr(PaymentAddress), + Str(String), +} + +impl ZAddrFilter { + fn from_opt(addr: &Option) -> Self { + match addr { + Some(a) => ZAddrFilter::Str(a.clone()), + None => ZAddrFilter::All, + } + } + + fn matches(&self, hrp: &str, nd: &SaplingNoteData) -> bool { + match self { + ZAddrFilter::All => true, + ZAddrFilter::Addr(pa) => nd.diversifier == pa.diversifier && nd.note.pk_d == pa.pk_d, + ZAddrFilter::Str(a) => LightWallet::note_address(hrp, nd).as_ref() == Some(a), + } + } +} + pub struct LightWallet { // Is the wallet encrypted? If it is, then when writing to disk, the seed is always encrypted // and the individual spending keys are not written @@ -490,10 +517,12 @@ impl LightWallet { } pub fn note_address(hrp: &str, note: &SaplingNoteData) -> Option { - match note.extfvk.fvk.vk.into_payment_address(note.diversifier, &JUBJUB) { - Some(pa) => Some(encode_payment_address(hrp, &pa)), - None => None - } + // note.pk_d == g_d(diversifier) * ivk(note.extfvk) by construction: decryption + // (parse_note_plaintext_without_memo) and SaplingNoteData::read (data.rs:158) both build + // it that way. So this encodes exactly what into_payment_address() would, without the + // group hash, ivk hash and 252-bit scalar multiplication. + let pa = PaymentAddress { pk_d: note.note.pk_d.clone(), diversifier: note.diversifier }; + Some(encode_payment_address(hrp, &pa)) } pub fn get_birthday(&self) -> u64 { @@ -1065,33 +1094,26 @@ impl LightWallet { pub fn zbalance(&self, addr: Option) -> u64 { let unconfirmed_balance = self.unconfirmed_zbalance(addr.clone()); + let confirmed_balance = self.confirmed_zbalance_f(&ZAddrFilter::from_opt(&addr)); - let confirmed_balance = self.txs.read().unwrap() + confirmed_balance + unconfirmed_balance + } + + // Unspent (and not pending-spent) note value matching `f`. The spent checks run first so spent + // notes never reach the address filter; both are pure, so the sum equals the old + // filter(addr).map(if unspent {v} else {0}).sum(). + fn confirmed_zbalance_f(&self, f: &ZAddrFilter) -> u64 { + let hrp = self.config.hrp_sapling_address(); + self.txs.read().unwrap() .values() .map(|tx| { tx.notes.iter() - .filter(|nd| { - match addr.as_ref() { - Some(a) => *a == encode_payment_address( - self.config.hrp_sapling_address(), - &nd.extfvk.fvk.vk - .into_payment_address(nd.diversifier, &JUBJUB).unwrap() - ), - None => true - } - }) - .map(|nd| { - if nd.spent.is_none() && nd.unconfirmed_spent.is_none() { - nd.note.value - } else { - 0 - } - }) + .filter(|nd| nd.spent.is_none() && nd.unconfirmed_spent.is_none()) + .filter(|nd| f.matches(hrp, nd)) + .map(|nd| nd.note.value) .sum::() }) - .sum::(); - - confirmed_balance + unconfirmed_balance + .sum::() } // Get all (unspent) utxos. Unconfirmed spent utxos are included @@ -1119,10 +1141,15 @@ impl LightWallet { } pub fn verified_zbalance(&self, addr: Option) -> u64 { + self.verified_zbalance_f(&ZAddrFilter::from_opt(&addr)) + } + + fn verified_zbalance_f(&self, f: &ZAddrFilter) -> u64 { let anchor_height = match self.get_target_height_and_anchor_offset() { Some((height, anchor_offset)) => height - anchor_offset as u32 , None => return 0, }; + let hrp = self.config.hrp_sapling_address(); self.txs .read() @@ -1133,16 +1160,7 @@ impl LightWallet { tx.notes .iter() .filter(|nd| nd.spent.is_none() && nd.unconfirmed_spent.is_none()) - .filter(|nd| { // TODO, this whole section is shared with verified_balance. Refactor it. - match addr.as_ref() { - Some(a) => *a == encode_payment_address( - self.config.hrp_sapling_address(), - &nd.extfvk.fvk.vk - .into_payment_address(nd.diversifier, &JUBJUB).unwrap() - ), - None => true - } - }) + .filter(|nd| f.matches(hrp, nd)) .map(|nd| nd.note.value) .sum::() } else { @@ -1153,7 +1171,14 @@ impl LightWallet { } pub fn spendable_zbalance(&self, addr: Option) -> u64 { + self.spendable_zbalance_f(&ZAddrFilter::from_opt(&addr)) + } + + // Lock order: txs.read, then zkeys.read inside have_spendingkey_for_extfvk (same order as + // scan_block). Never call this while holding a zkeys guard. + fn spendable_zbalance_f(&self, f: &ZAddrFilter) -> u64 { let anchor_height = self.get_anchor_height(); + let hrp = self.config.hrp_sapling_address(); self.txs .read() @@ -1164,20 +1189,11 @@ impl LightWallet { tx.notes .iter() .filter(|nd| nd.spent.is_none() && nd.unconfirmed_spent.is_none()) + .filter(|nd| f.matches(hrp, nd)) .filter(|nd| { // Check to see if we have this note's spending key. self.have_spendingkey_for_extfvk(&nd.extfvk) }) - .filter(|nd| { // TODO, this whole section is shared with verified_balance. Refactor it. - match addr.as_ref() { - Some(a) => *a == encode_payment_address( - self.config.hrp_sapling_address(), - &nd.extfvk.fvk.vk - .into_payment_address(nd.diversifier, &JUBJUB).unwrap() - ), - None => true - } - }) .map(|nd| nd.note.value) .sum::() } else { @@ -1227,6 +1243,45 @@ impl LightWallet { } } + /// Per-zkey z balances for do_balance, in zkeys order: + /// (address, zbalance, unconfirmed, verified_zbalance, spendable_zbalance). + /// Equal to calling zbalance/unconfirmed_zbalance/verified_zbalance/spendable_zbalance with + /// Some(encoded zaddress), in that order, without re-deriving any note's address. + pub fn z_balances_by_address(&self) -> Vec<(String, u64, u64, u64, u64)> { + let hrp = self.config.hrp_sapling_address(); + // Snapshot, then RELEASE the zkeys guard before any txs lock: scan_block holds txs.write + // while taking zkeys.read/write (ensure_hd_zaddresses -> add_zaddr), and + // spendable_zbalance_f re-reads zkeys. Holding zkeys here would invert the lock order + // and re-enter a writer-preferring RwLock (deadlock). + let addrs: Vec<(String, PaymentAddress)> = self.zkeys.read().unwrap().iter() + .map(|zk| (encode_payment_address(hrp, &zk.zaddress), zk.zaddress.clone())) + .collect(); + + addrs.into_iter().map(|(address, pa)| { + let f = ZAddrFilter::Addr(pa); + let zbalance = { + let unconfirmed_balance = self.unconfirmed_zbalance(Some(address.clone())); + self.confirmed_zbalance_f(&f) + unconfirmed_balance + }; + let unconfirmed = self.unconfirmed_zbalance(Some(address.clone())); + let verified = self.verified_zbalance_f(&f); + let spendable = self.spendable_zbalance_f(&f); + (address, zbalance, unconfirmed, verified, spendable) + }).collect() + } + + /// Encoded z-addresses whose FIRST zkey with that address has a spending key. Same set as + /// get_all_zaddresses() filtered by have_spending_key_for_zaddress (Z encodes, not Z + Z^2). + pub fn spendable_zaddress_set(&self) -> HashSet { + let hrp = self.config.hrp_sapling_address(); + let mut first: HashMap = HashMap::new(); + for zk in self.zkeys.read().unwrap().iter() { + first.entry(encode_payment_address(hrp, &zk.zaddress)) + .or_insert_with(|| zk.have_spending_key()); + } + first.into_iter().filter(|(_, spendable)| *spendable).map(|(a, _)| a).collect() + } + fn add_toutput_to_wtx(&self, height: i32, timestamp: u64, txid: &TxId, vout: &TxOut, n: u64) { diff --git a/third_party/silentdragonxlite/silentdragonxlite-cli/lib/src/lightwallet/tests.rs b/third_party/silentdragonxlite/silentdragonxlite-cli/lib/src/lightwallet/tests.rs index 65339fc..4e51c79 100644 --- a/third_party/silentdragonxlite/silentdragonxlite-cli/lib/src/lightwallet/tests.rs +++ b/third_party/silentdragonxlite/silentdragonxlite-cli/lib/src/lightwallet/tests.rs @@ -124,6 +124,12 @@ impl FakeCompactBlock { fn add_tx_paying(&mut self, extfvk: ExtendedFullViewingKey, value: u64) -> (Vec, TxId) { let to = extfvk.default_address().unwrap().1; + self.add_tx_paying_to(extfvk, to, value) + } + + // add_tx_paying with an explicit recipient (e.g. a diversified address of `extfvk`). + fn add_tx_paying_to(&mut self, extfvk: ExtendedFullViewingKey, to: PaymentAddress, value: u64) + -> (Vec, TxId) { let value = Amount::from_u64(value).unwrap(); // Create a fake Note for the account @@ -2346,3 +2352,626 @@ fn test_encrypted_treceive() { assert_eq!(utxo2.unconfirmed_spent, None); } } + +// --------------------------------------------------------------------------------------------- +// g1 (C249): note addresses come from the stored pk_d, the z-balance address filters compare +// (diversifier, pk_d), and do_balance / do_list_notes use z_balances_by_address / +// spendable_zaddress_set. All of it must be exactly equivalent to the pre-change code. The old_* +// functions below are verbatim copies of the replaced bodies, taking the wallet as a parameter. +// --------------------------------------------------------------------------------------------- + +use std::collections::HashSet; +use super::{SaplingNoteData, WalletTx, IncomingTxMetadata}; + +fn old_note_address(hrp: &str, note: &SaplingNoteData) -> Option { + match note.extfvk.fvk.vk.into_payment_address(note.diversifier, &JUBJUB) { + Some(pa) => Some(encode_payment_address(hrp, &pa)), + None => None + } +} + +fn old_zbalance(w: &LightWallet, addr: Option) -> u64 { + let unconfirmed_balance = w.unconfirmed_zbalance(addr.clone()); + + let confirmed_balance = w.txs.read().unwrap() + .values() + .map(|tx| { + tx.notes.iter() + .filter(|nd| { + match addr.as_ref() { + Some(a) => *a == encode_payment_address( + w.config.hrp_sapling_address(), + &nd.extfvk.fvk.vk + .into_payment_address(nd.diversifier, &JUBJUB).unwrap() + ), + None => true + } + }) + .map(|nd| { + if nd.spent.is_none() && nd.unconfirmed_spent.is_none() { + nd.note.value + } else { + 0 + } + }) + .sum::() + }) + .sum::(); + + confirmed_balance + unconfirmed_balance +} + +fn old_verified_zbalance(w: &LightWallet, addr: Option) -> u64 { + let anchor_height = match w.get_target_height_and_anchor_offset() { + Some((height, anchor_offset)) => height - anchor_offset as u32 , + None => return 0, + }; + + w.txs + .read() + .unwrap() + .values() + .map(|tx| { + if tx.block as u32 <= anchor_height { + tx.notes + .iter() + .filter(|nd| nd.spent.is_none() && nd.unconfirmed_spent.is_none()) + .filter(|nd| { + match addr.as_ref() { + Some(a) => *a == encode_payment_address( + w.config.hrp_sapling_address(), + &nd.extfvk.fvk.vk + .into_payment_address(nd.diversifier, &JUBJUB).unwrap() + ), + None => true + } + }) + .map(|nd| nd.note.value) + .sum::() + } else { + 0 + } + }) + .sum::() +} + +fn old_spendable_zbalance(w: &LightWallet, addr: Option) -> u64 { + let anchor_height = w.get_anchor_height(); + + w.txs + .read() + .unwrap() + .values() + .map(|tx| { + if tx.block as u32 <= anchor_height { + tx.notes + .iter() + .filter(|nd| nd.spent.is_none() && nd.unconfirmed_spent.is_none()) + .filter(|nd| { + // Check to see if we have this note's spending key. + w.have_spendingkey_for_extfvk(&nd.extfvk) + }) + .filter(|nd| { + match addr.as_ref() { + Some(a) => *a == encode_payment_address( + w.config.hrp_sapling_address(), + &nd.extfvk.fvk.vk + .into_payment_address(nd.diversifier, &JUBJUB).unwrap() + ), + None => true + } + }) + .map(|nd| nd.note.value) + .sum::() + } else { + 0 + } + }) + .sum::() as u64 +} + +// The old do_list_notes construction of the spendable-address set. +fn old_spendable_address_set(w: &LightWallet) -> HashSet { + w.get_all_zaddresses().iter() + .filter(|address| w.have_spending_key_for_zaddress(address)) + .map(|address| address.clone()) + .collect() +} + +/// Deterministic xorshift64, so the shape of the g1 chains is reproducible (note randomness and +/// txids still come from OsRng). +struct G1Rng(u64); + +impl G1Rng { + fn next(&mut self) -> u64 { + let mut x = self.0; + x ^= x << 13; + x ^= x >> 7; + x ^= x << 17; + self.0 = x; + x + } + + fn below(&mut self, n: u64) -> u64 { + self.next() % n + } +} + +/// A diversified (non-default) address of `extfvk`. +fn g1_diversified_address(extfvk: &ExtendedFullViewingKey) -> PaymentAddress { + let (mut j, _) = extfvk.default_address().unwrap(); + j.increment().unwrap(); + extfvk.address(j).unwrap().1 +} + +/// Minimal BIP-173 bech32 encoder, only used to build syntactically valid junk addresses. +fn g1_bech32(hrp: &str, data: &[u8]) -> String { + const CHARSET: &[u8] = b"qpzry9x8gf2tvdw0s3jn54khce6mua7l"; + const GEN: [u32; 5] = [0x3b6a57b2, 0x26508e6d, 0x1ea119fa, 0x3d4233dd, 0x2a1462b3]; + + let mut five: Vec = vec![]; + let mut acc: u32 = 0; + let mut bits: u32 = 0; + for b in data { + acc = (acc << 8) | (*b as u32); + bits += 8; + while bits >= 5 { + bits -= 5; + five.push(((acc >> bits) & 31) as u8); + } + acc &= (1 << bits) - 1; + } + if bits > 0 { + five.push(((acc << (5 - bits)) & 31) as u8); + } + + let mut values: Vec = hrp.bytes().map(|c| c >> 5).collect(); + values.push(0); + values.extend(hrp.bytes().map(|c| c & 31)); + values.extend(five.iter().cloned()); + values.extend([0u8; 6].iter().cloned()); + let mut chk: u32 = 1; + for v in values { + let top = chk >> 25; + chk = ((chk & 0x1ffffff) << 5) ^ (v as u32); + for (i, g) in GEN.iter().enumerate() { + if (top >> i) & 1 == 1 { + chk ^= *g; + } + } + } + let polymod = chk ^ 1; + + let mut s = String::from(hrp); + s.push('1'); + for d in five { + s.push(CHARSET[d as usize] as char); + } + for i in 0..6 { + s.push(CHARSET[((polymod >> (5 * (5 - i))) & 31) as usize] as char); + } + s +} + +struct G1Wallet { + wallet: LightWallet, + keys: Vec, // the 8 keys notes were paid to (6 HD, imported vk, imported sk) + diversified: Vec, // diversified addresses of our keys that received notes + foreign: String, // default address of a key that is not in the wallet + vk_address: String, // the imported view-only key's address + sk_address: String, // the imported spending key's address +} + +/// 6 HD keys, another wallet's key imported view-only and an imported spending key. Then `blocks` +/// blocks of payments chosen by `seed`: random key and value (zero-value and dust included), +/// foreign outputs, compact spends with change, and 3 notes paid to diversified addresses. The +/// view-only key's notes are never spent (so verified != spendable), and the tip block always pays +/// us (so anchor offsets and the spendable "- 1" matter). +fn g1_build(seed: u64, blocks: i32) -> G1Wallet { + use zcash_client_backend::encoding::{encode_extended_full_viewing_key, encode_extended_spending_key}; + + let config = get_test_config(); + let hrp = config.hrp_sapling_address().to_string(); + let mut rng = G1Rng(seed | 1); + + let mut wallet = LightWallet::new(None, &config, 0, 0).unwrap(); + for _ in 0..5 { + wallet.add_zaddr(); + } + + let other = LightWallet::new(None, &config, 0, 0).unwrap(); + let vk_extfvk = other.zkeys.read().unwrap()[0].extfvk.clone(); + let vk_address = wallet.add_imported_vk( + encode_extended_full_viewing_key(config.hrp_sapling_viewing_key(), &vk_extfvk), 0); + assert!(vk_address.starts_with(&hrp), "{}", vk_address); + + let mut sk_seed = [0x5au8; 32]; + sk_seed[..8].copy_from_slice(&seed.to_le_bytes()); + let extsk = ExtendedSpendingKey::master(&sk_seed); + let sk_address = wallet.add_imported_sk( + encode_extended_spending_key(config.hrp_sapling_private_key(), &extsk), 0); + assert!(sk_address.starts_with(&hrp), "{}", sk_address); + + let foreign_extfvk = ExtendedFullViewingKey::from(&ExtendedSpendingKey::master(&[7u8; 32])); + let foreign_pa = foreign_extfvk.default_address().unwrap().1; + let foreign = encode_payment_address(&hrp, &foreign_pa); + + let keys: Vec = wallet.zkeys.read().unwrap().iter() + .map(|zk| zk.extfvk.clone()).collect(); + assert_eq!(keys.len(), 8); + assert!(keys[6] == vk_extfvk); + + let diversified_heights = [blocks / 6, blocks / 2, blocks - 3]; + let mut diversified: Vec = vec![]; + let mut prev = BlockHash([0; 32]); + for h in 0..blocks { + let mut cb = FakeCompactBlock::new(h, prev); + + // Sometimes a compact spend (with change) of an unspent, non-zero, spendable-key note. + if h > 2 && rng.below(4) == 0 { + let candidate = { + let txs = wallet.txs.read().unwrap(); + let mut c: Vec<(i32, u64, Vec, ExtendedFullViewingKey)> = txs.values() + .flat_map(|tx| tx.notes.iter().map(move |nd| (tx.block, nd))) + .filter(|(_, nd)| nd.spent.is_none() && nd.note.value > 0 && nd.extfvk != vk_extfvk + && !diversified.contains(&old_note_address(&hrp, nd).unwrap())) + .map(|(b, nd)| (b, nd.note.value, nd.nullifier.to_vec(), nd.extfvk.clone())) + .collect(); + // HashMap order is per instance; sort so the pick depends only on the chain. + c.sort_by(|x, y| (x.0, x.1, &x.2).cmp(&(y.0, y.1, &y.2))); + if c.is_empty() { + None + } else { + let i = rng.below(c.len() as u64) as usize; + Some(c.swap_remove(i)) + } + }; + if let Some((_, value, nf, extfvk)) = candidate { + let sent = rng.below(value + 1); + cb.add_tx_spending((nf, value), extfvk, foreign_pa.clone(), sent); + } + } + + // 0-2 incoming payments to random keys (always at least one in the last 4 blocks). + let n = if h >= blocks - 4 { 1 + rng.below(2) } else { rng.below(3) }; + for _ in 0..n { + let k = rng.below(keys.len() as u64) as usize; + let value = match rng.below(10) { + 0 => 0, // zero-value (HushChat-style) note + 1 | 2 => 1 + rng.below(1000), // dust + _ => 1 + rng.below(1_000_000_000), + }; + cb.add_tx_paying(keys[k].clone(), value); + } + if h == 0 { + cb.add_tx_paying(vk_extfvk.clone(), 777_777); + } + if h == blocks - 1 { + cb.add_tx_paying(keys[0].clone(), 333_333); + } + + // A foreign output now and then. + if rng.below(3) == 0 { + cb.add_tx_paying(foreign_extfvk.clone(), 1 + rng.below(1000)); + } + + // 3 notes paid to diversified addresses of our keys. + if diversified_heights.contains(&h) { + let k = rng.below(keys.len() as u64) as usize; + let to = g1_diversified_address(&keys[k]); + assert!(to != keys[k].default_address().unwrap().1); + diversified.push(encode_payment_address(&hrp, &to)); + cb.add_tx_paying_to(keys[k].clone(), to, 10_000 + rng.below(1_000_000)); + } + + wallet.scan_block(&cb.as_bytes()).unwrap(); + prev = cb.hash(); + } + + G1Wallet { wallet, keys, diversified, foreign, vk_address, sk_address } +} + +/// Marks about a third of the unspent notes spent or pending-spent (deterministic per seed), except +/// the view-only key's notes and the tip block's notes. +fn g1_mark_spent(g: &G1Wallet, seed: u64) { + let mut rng = G1Rng(seed | 1); + let tip = g.wallet.last_scanned_height(); + let mut txs = g.wallet.txs.write().unwrap(); + let mut ids: Vec<(i32, u64, TxId)> = txs.values() + .map(|tx| (tx.block, tx.notes.iter().map(|nd| nd.note.value).sum::(), tx.txid)) + .collect(); + ids.sort_by(|x, y| (x.0, x.1, x.2 .0).cmp(&(y.0, y.1, y.2 .0))); + for (block, _, txid) in ids { + if block == tip { + continue; + } + for nd in txs.get_mut(&txid).unwrap().notes.iter_mut() { + if nd.spent.is_some() || nd.extfvk == g.keys[6] { + continue; + } + match rng.below(6) { + 0 => nd.spent = Some(TxId([0xab; 32])), + 1 => nd.unconfirmed_spent = Some(TxId([0xcd; 32])), + _ => {} + } + } + } +} + +/// Adds an incoming mempool tx paying two of our addresses and a diversified one, so the +/// unconfirmed term of zbalance is non-zero. +fn g1_add_incoming_mempool(g: &G1Wallet) { + let hrp = g.wallet.config.hrp_sapling_address().to_string(); + let txid = TxId([0x11; 32]); + let mut wtx = WalletTx::new(g.wallet.last_scanned_height() + 1, 0, &txid); + for (address, value) in vec![ + (encode_payment_address(&hrp, &g.keys[1].default_address().unwrap().1), 4242u64), + (g.sk_address.clone(), 17), + (g.diversified[0].clone(), 99), + ] { + wtx.incoming_metadata.push(IncomingTxMetadata { + address, + value, + memo: Memo::default(), + incoming_mempool: true, + }); + } + g.wallet.incoming_mempool_txs.write().unwrap().insert(txid, vec![wtx]); +} + +/// Every address-filter input of tests 2 and 3: None, the wallet's addresses, the diversified +/// addresses, and junk that must keep the old exact string compare (never decoded). +fn g1_inputs(g: &G1Wallet) -> Vec> { + let hrp = g.wallet.config.hrp_sapling_address().to_string(); + let mut inputs: Vec> = vec![None]; + let all = g.wallet.get_all_zaddresses(); + inputs.extend(all.iter().cloned().map(Some)); + inputs.extend(g.diversified.iter().cloned().map(Some)); + + let valid = all[1].clone(); + let mut last_changed = valid.clone(); + let last = last_changed.pop().unwrap(); + last_changed.push(if last == 'q' { 'p' } else { 'q' }); + + // The bech32 helper is exact: it reproduces encode_payment_address on a real address. + let pa = g.keys[2].default_address().unwrap().1; + let mut payload = pa.diversifier.0.to_vec(); + pa.pk_d.write(&mut payload).unwrap(); + assert_eq!(g1_bech32(&hrp, &payload), encode_payment_address(&hrp, &pa)); + + inputs.extend(vec![ + String::new(), + valid.to_uppercase(), + last_changed, + g.foreign.clone(), + "z".repeat(300), + g1_bech32(&hrp, &[1, 2, 3, 4, 5]), // valid bech32, short payload + g1_bech32(&hrp, &payload[..42]), // valid bech32, one byte short + g1_bech32("ztestsapling", &payload), // our key under the testnet hrp + format!("{} ", valid), + ].into_iter().map(Some)); + inputs +} + +fn g1_assert_filters_match_old(w: &LightWallet, inputs: &[Option], what: &str) { + for a in inputs { + assert_eq!(w.zbalance(a.clone()), old_zbalance(w, a.clone()), "{}: zbalance({:?})", what, a); + assert_eq!(w.verified_zbalance(a.clone()), old_verified_zbalance(w, a.clone()), + "{}: verified_zbalance({:?})", what, a); + assert_eq!(w.spendable_zbalance(a.clone()), old_spendable_zbalance(w, a.clone()), + "{}: spendable_zbalance({:?})", what, a); + } +} + +#[test] +fn test_note_address_pk_d_matches_derivation() { + for seed in &[0x9e3779b97f4a7c15u64, 0x0123456789abcdef] { + let g = g1_build(*seed, 40); + let config = g.wallet.config.clone(); + let hrp = config.hrp_sapling_address().to_string(); + + let check = |w: &LightWallet, what: &str| -> usize { + let txs = w.txs.read().unwrap(); + let mut n = 0; + let mut seen: HashSet = HashSet::new(); + for tx in txs.values() { + for nd in tx.notes.iter() { + let new = LightWallet::note_address(&hrp, nd); + assert!(new.is_some(), "{}: no address", what); + assert_eq!(new, old_note_address(&hrp, nd), "{}: note address differs", what); + seen.insert(new.unwrap()); + n += 1; + } + } + for d in g.diversified.iter() { + assert!(seen.contains(d), "{}: diversified address {} not derived", what, d); + } + n + }; + + let n = check(&g.wallet, "scanned"); + assert!(n >= 30, "only {} notes", n); + let all = g.wallet.get_all_zaddresses(); + for d in g.diversified.iter() { + assert!(!all.contains(d)); + } + + // Notes rebuilt by SaplingNoteData::read (data.rs:158) give the same addresses, and the + // file format is unchanged: write -> read -> write is byte-identical. + let mut first = vec![]; + g.wallet.write(&mut first).unwrap(); + let reread = LightWallet::read(&first[..], &config).unwrap(); + assert_eq!(check(&reread, "reread"), n); + let mut second = vec![]; + reread.write(&mut second).unwrap(); + assert!(first == second, "write -> read -> write changed the wallet bytes"); + } +} + +#[test] +fn test_zbalance_filters_match_old() { + for seed in &[0x2545f4914f6cdd1du64, 0x00c0ffee12345678] { + let g = g1_build(*seed, 40); + g1_mark_spent(&g, *seed ^ 0xff); + g1_add_incoming_mempool(&g); + let inputs = g1_inputs(&g); + + // Non-vacuous: pending and confirmed spends exist, view-only notes make verified != + // spendable, and the junk inputs match nothing. + { + let txs = g.wallet.txs.read().unwrap(); + let notes = txs.values().flat_map(|tx| tx.notes.iter()); + assert!(notes.clone().any(|nd| nd.spent.is_some())); + assert!(notes.clone().any(|nd| nd.unconfirmed_spent.is_some())); + } + assert!(g.wallet.verified_zbalance(None) > g.wallet.spendable_zbalance(None)); + assert!(g.wallet.verified_zbalance(Some(g.vk_address.clone())) >= 777_777); + assert_eq!(g.wallet.spendable_zbalance(Some(g.vk_address.clone())), 0); + assert_eq!(g.wallet.zbalance(Some(g.foreign.clone())), 0); + g1_assert_filters_match_old(&g.wallet, &inputs, "offset 0"); + + // Same wallet read back with anchor_offset 3. read() drops the pending spends, so mark + // again. + let mut bytes = vec![]; + g.wallet.write(&mut bytes).unwrap(); + let mut config3 = g.wallet.config.clone(); + config3.anchor_offset = 3; + let w3 = LightWallet::read(&bytes[..], &config3).unwrap(); + g1_assert_filters_match_old(&w3, &inputs, "offset 3, reread"); + // The tip block's 333_333 note is above the offset-3 anchor. + assert!(w3.verified_zbalance(None) + 333_333 <= w3.zbalance(None) - w3.unconfirmed_zbalance(None)); + let g3 = G1Wallet { wallet: w3, keys: g.keys.clone(), diversified: g.diversified.clone(), + foreign: g.foreign.clone(), vk_address: g.vk_address.clone(), + sk_address: g.sk_address.clone() }; + g1_mark_spent(&g3, *seed ^ 0x5555); + g1_add_incoming_mempool(&g3); + g1_assert_filters_match_old(&g3.wallet, &inputs, "offset 3, re-marked"); + } +} + +#[test] +fn test_z_balances_by_address_matches_per_address_calls() { + for seed in &[0x6a09e667f3bcc908u64, 0x3c6ef372fe94f82b] { + let g = g1_build(*seed, 40); + g1_add_incoming_mempool(&g); + let w = &g.wallet; + + let by_address = w.z_balances_by_address(); + let all = w.get_all_zaddresses(); + assert_eq!(by_address.iter().map(|t| t.0.clone()).collect::>(), all); + assert_eq!(all.iter().cloned().collect::>().len(), all.len()); + + let mut confirmed_listed: u64 = 0; + for (a, zbalance, unconfirmed, verified, spendable) in by_address.iter() { + let s = Some(a.clone()); + assert_eq!(*zbalance, w.zbalance(s.clone()), "zbalance {}", a); + assert_eq!(*unconfirmed, w.unconfirmed_zbalance(s.clone()), "unconfirmed {}", a); + assert_eq!(*verified, w.verified_zbalance(s.clone()), "verified {}", a); + assert_eq!(*spendable, w.spendable_zbalance(s.clone()), "spendable {}", a); + assert_eq!(*zbalance, old_zbalance(w, s.clone()), "old zbalance {}", a); + assert_eq!(*verified, old_verified_zbalance(w, s.clone()), "old verified {}", a); + assert_eq!(*spendable, old_spendable_zbalance(w, s.clone()), "old spendable {}", a); + confirmed_listed += zbalance - unconfirmed; + } + assert!(by_address.iter().any(|t| t.2 > 0), "no unconfirmed term exercised"); + + // Notes paid to diversified addresses belong to no listed address, as before. + let hrp = w.config.hrp_sapling_address().to_string(); + let diversified_unspent: u64 = w.txs.read().unwrap().values() + .flat_map(|tx| tx.notes.iter()) + .filter(|nd| nd.spent.is_none() && nd.unconfirmed_spent.is_none() + && g.diversified.contains(&old_note_address(&hrp, nd).unwrap())) + .map(|nd| nd.note.value) + .sum(); + assert!(diversified_unspent > 0); + let confirmed_all = w.zbalance(None) - w.unconfirmed_zbalance(None); + assert_eq!(confirmed_listed + diversified_unspent, confirmed_all); + } +} + +#[test] +fn test_spendable_zaddress_set_matches_old() { + let g = g1_build(0x510e527fade682d1, 12); + let vk_address = g.vk_address.clone(); + let sk_address = g.sk_address.clone(); + let mut w = g.wallet; + + let check = |w: &LightWallet, what: &str| { + let set = w.spendable_zaddress_set(); + assert_eq!(set, old_spendable_address_set(w), "{}", what); + assert!(!set.contains(&vk_address), "{}: view-only key counted as spendable", what); + assert!(set.contains(&sk_address), "{}: imported spending key missing", what); + assert_eq!(set.len(), w.get_all_zaddresses().len() - 1, "{}", what); + }; + + check(&w, "plain"); + w.encrypt("g1 password".to_string()).unwrap(); + check(&w, "encrypted + locked"); + w.unlock("g1 password".to_string()).unwrap(); + check(&w, "unlocked"); + w.remove_encryption("g1 password".to_string()).unwrap(); + check(&w, "decrypted"); +} + +#[test] +fn test_balances_no_deadlock_with_zkeys_writer() { + use std::sync::{Arc, mpsc, atomic::{AtomicBool, Ordering}}; + use std::time::Duration; + + let wallet = Arc::new(LightWallet::new(None, &get_test_config(), 0, 0).unwrap()); + for _ in 0..5 { + wallet.add_zaddr(); + } + let writer_done = Arc::new(AtomicBool::new(false)); + + // Reader: the do_balance and send-path balance calls, back to back. + let (rtx, rrx) = mpsc::channel(); + let reader = { + let w = wallet.clone(); + let done = writer_done.clone(); + std::thread::spawn(move || { + let mut rounds = 0u64; + let mut seen = 0u64; + while rounds < 300 || (!done.load(Ordering::SeqCst) && rounds < 100_000) { + seen += w.z_balances_by_address().len() as u64; + seen += w.spendable_zbalance(None) % 2; + rounds += 1; + } + rtx.send((rounds, seen)).unwrap(); + }) + }; + + // Writer: add_zaddr (zkeys.write) interleaved with scan_block of blocks paying the newest HD + // address, so ensure_hd_zaddresses takes zkeys.write while scan_block holds txs.write. + let (wtx, wrx) = mpsc::channel(); + let writer = { + let w = wallet.clone(); + let done = writer_done.clone(); + std::thread::spawn(move || { + let mut prev = BlockHash([0; 32]); + for i in 0..100 { + w.add_zaddr(); + let newest = w.zkeys.read().unwrap().last().unwrap().extfvk.clone(); + let mut cb = FakeCompactBlock::new(i, prev); + cb.add_tx_paying(newest, 1000 + i as u64); + w.scan_block(&cb.as_bytes()).unwrap(); + prev = cb.hash(); + } + done.store(true, Ordering::SeqCst); + wtx.send(w.zkeys.read().unwrap().len()).unwrap(); + }) + }; + + let nkeys = wrx.recv_timeout(Duration::from_secs(120)) + .expect("add_zaddr/scan_block did not finish within 120 s (lock-order deadlock?)"); + let (rounds, _) = rrx.recv_timeout(Duration::from_secs(120)) + .expect("z_balances_by_address/spendable_zbalance did not finish within 120 s (deadlock?)"); + writer.join().unwrap(); + reader.join().unwrap(); + + // 6 + 100 add_zaddr + 5 per block from the gap rule. + assert_eq!(nkeys, 6 + 100 + 5 * 100); + assert!(rounds >= 300); + assert_eq!(wallet.zbalance(None), (0..100u64).map(|i| 1000 + i).sum::()); + assert_eq!(wallet.z_balances_by_address().len(), nkeys); +}