DragonX/ObsidianDragon
1
1
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases 10 Wiki Activity

perf(lite-backend): encode note addresses from pk_d instead of deriving

Browse Source 
note_address() re-derived a note's payment address from its viewing key
and diversifier (a group hash, an ivk hash and a 252-bit scalar
multiplication) for every note in `notes` and `list` and for each new
note in the block scan. The z-balance sums did the same derivation in
their address filter, so `balance`, which sums four balances for each of
Z z-addresses, cost Z x (N + 2U) derivations for N notes, U of them
unspent. do_list_notes also encoded up to Z + Z^2 addresses to find the
spendable ones and cloned that set once per transaction.

Now:
- note_address() encodes the note's stored pk_d with its diversifier.
  Both decryption and wallet load build that pk_d from the same key and
  diversifier, so the string is byte-identical;
- the three z-balance sums share one address filter. The per-address
  sums behind `balance` compare (diversifier, pk_d) with each z-address
  directly; any other address string keeps the old exact string compare;
- those sums work from a copy of the z-addresses whose lock is released
  before any note is read, so the txs-then-zkeys lock order that sync
  uses is unchanged;
- do_list_notes builds the spendable-address set with Z encodes and
  borrows it.

No output changes: balances and the notes/list/balance JSON are
byte-identical. The address filter no longer contains an unwrap (it
could not fail).

Tests compare against verbatim copies of the old functions over
randomized chains (note addresses, also after a reload; every filtered
balance for own, diversified, foreign and junk address strings; the
per-address balances; the spendable set), and a concurrency test fails
with a deadlock if the zkeys lock is held across the sums. A
differential run against the unmodified backend compared per-address
balances and the JSON after every operation, through reorgs and reloads,
and the old build re-saves the new files byte for byte.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
This commit is contained in:
DanS
2026-10-10 09:50:03 -05:00
co-authored by Claude Opus 5.5
parent d28cc0da31
commit bd4c06dfd0
3 changed files with 739 additions and 56 deletions
Download Patch File Download Diff File Expand all files Collapse all files
third_party/silentdragonxlite/silentdragonxlite-cli/lib/src/lightclient.rs
Vendored
+10 -11
View File
@@ -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::<Vec<JsonValue>>();
@@ -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<String> = wallet.get_all_zaddresses().iter()
.filter(|address| wallet.have_spending_key_for_zaddress(address))
.map(|address| address.clone())
.collect();
let spendable_address_set: HashSet<String> = wallet.spendable_zaddress_set();
let spendable_address: &HashSet<String> = &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<String> = spendable_address;
wtx.notes.iter().filter_map(move |nd|
if !all_notes && nd.spent.is_some() {
None
third_party/silentdragonxlite/silentdragonxlite-cli/lib/src/lightwallet.rs
Vendored
+100 -45
View File
@@ -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<Bls12>),
Str(String),
}
impl ZAddrFilter {
fn from_opt(addr: &Option<String>) -> 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<String> {
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<String>) -> 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::<u64>()
})
.sum::<u64>();
confirmed_balance + unconfirmed_balance
.sum::<u64>()
}
// Get all (unspent) utxos. Unconfirmed spent utxos are included
@@ -1119,10 +1141,15 @@ impl LightWallet {
}
pub fn verified_zbalance(&self, addr: Option<String>) -> 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::<u64>()
} else {
@@ -1153,7 +1171,14 @@ impl LightWallet {
}
pub fn spendable_zbalance(&self, addr: Option<String>) -> 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::<u64>()
} 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<Bls12>)> = 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<String> {
let hrp = self.config.hrp_sapling_address();
let mut first: HashMap<String, bool> = 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) {
third_party/silentdragonxlite/silentdragonxlite-cli/lib/src/lightwallet/tests.rs
Vendored
+629
View File
@@ -124,6 +124,12 @@ impl FakeCompactBlock {
fn add_tx_paying(&mut self, extfvk: ExtendedFullViewingKey, value: u64)
-> (Vec<u8>, 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<Bls12>, value: u64)
-> (Vec<u8>, 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<String> {
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<String>) -> 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::<u64>()
})
.sum::<u64>();
confirmed_balance + unconfirmed_balance
}
fn old_verified_zbalance(w: &LightWallet, addr: Option<String>) -> 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::<u64>()
} else {
0
}
})
.sum::<u64>()
}
fn old_spendable_zbalance(w: &LightWallet, addr: Option<String>) -> 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::<u64>()
} else {
0
}
})
.sum::<u64>() as u64
}
// The old do_list_notes construction of the spendable-address set.
fn old_spendable_address_set(w: &LightWallet) -> HashSet<String> {
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<Bls12> {
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<u8> = 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<u8> = 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<ExtendedFullViewingKey>, // the 8 keys notes were paid to (6 HD, imported vk, imported sk)
diversified: Vec<String>, // 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<ExtendedFullViewingKey> = 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<String> = 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<u8>, 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::<u64>(), 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<Option<String>> {
let hrp = g.wallet.config.hrp_sapling_address().to_string();
let mut inputs: Vec<Option<String>> = 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<String>], 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<String> = 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::<Vec<String>>(), all);
assert_eq!(all.iter().cloned().collect::<HashSet<String>>().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::<u64>());
assert_eq!(wallet.z_balances_by_address().len(), nkeys);
}