ObsidianDragon/third_party/silentdragonxlite/silentdragonxlite-cli/lib/tests/fulltx_fetch.rs

// Offline tests for the full-tx fetch engine (src/fulltx_fetch.rs). No network, no proving.
// Integration test on purpose: the in-crate #[cfg(test)] modules do not compile on BASE.
use std::collections::{BTreeSet, HashMap};
use std::sync::atomic::{AtomicUsize, Ordering};
use std::sync::{mpsc, Arc, Mutex, RwLock};
use std::time::{Duration, Instant};

use base58::FromBase58;
use rand::rngs::StdRng;
use rand::{Rng, SeedableRng};
use threadpool::ThreadPool;
use tokio::runtime::Runtime;
use zcash_primitives::legacy::{Script, TransparentAddress};
use zcash_primitives::transaction::components::{Amount, OutPoint, TxIn, TxOut};
use zcash_primitives::transaction::{Transaction, TransactionData, TxId};

use silentdragonxlitelib::fulltx_fetch::{
    plan_fetch_list, run_fetch_scan, wallet_scan_fn, FetchFn, FetchFuture, FetchPolicy, ScanFn,
};
use silentdragonxlitelib::lightclient::LightClientConfig;
use silentdragonxlitelib::lightwallet::LightWallet;

type CallLog = Arc<Mutex<Vec<(TxId, bool)>>>;
type ScanLog = Arc<Mutex<Vec<(TxId, i32)>>>;

fn tid(n: u32) -> TxId {
    let mut b = [0u8; 32];
    b[..4].copy_from_slice(&n.to_be_bytes()); // big-endian prefix => TxId order == n order
    TxId(b)
}

fn policy(workers: usize) -> FetchPolicy {
    FetchPolicy { workers, max_attempts: 3, backoff_ms: 1, attempt_timeout: Duration::from_secs(10), fail_breaker: 16 }
}

fn run(items: Vec<(TxId, i32)>, own: &[(TxId, i32)], p: FetchPolicy, fetch: FetchFn, scan: ScanFn) -> Result<(), String> {
    let own_set: BTreeSet<(TxId, i32)> = own.iter().cloned().collect();
    let mut rt = Runtime::new().unwrap();
    run_fetch_scan(&mut rt, items, own_set, p, ThreadPool::new(8), fetch, scan)
}

/// Fake GetTransaction: logs (txid, reconnect); returns the txid bytes, or Err when
/// `fail(txid, nth_call_for_this_txid)` says so.
fn fake_fetch<F>(log: CallLog, delay_ms: u64, fail: F) -> FetchFn
where
    F: Fn(TxId, usize) -> bool + Send + Sync + 'static,
{
    let calls: Arc<Mutex<HashMap<TxId, usize>>> = Arc::new(Mutex::new(HashMap::new()));
    let fail = Arc::new(fail);
    Arc::new(move |txid: TxId, reconnect: bool| {
        let log = log.clone();
        let calls = calls.clone();
        let fail = fail.clone();
        Box::pin(async move {
            log.lock().unwrap().push((txid, reconnect));
            let nth = {
                let mut c = calls.lock().unwrap();
                let e = c.entry(txid).or_insert(0);
                *e += 1;
                *e
            };
            if delay_ms > 0 {
                tokio::time::delay_for(Duration::from_millis(delay_ms)).await;
            }
            if (*fail)(txid, nth) { Err(format!("fake failure {}", txid)) } else { Ok(txid.0.to_vec()) }
        }) as FetchFuture
    })
}

fn logging_scan(log: ScanLog) -> ScanFn {
    Arc::new(move |bytes: Vec<u8>, height: i32| {
        let mut b = [0u8; 32];
        b.copy_from_slice(&bytes[..32]);
        log.lock().unwrap().push((TxId(b), height));
        Ok(())
    })
}

fn no_scan_log() -> ScanFn { logging_scan(Arc::new(Mutex::new(vec![]))) }

#[test] // T1
fn plan_is_exactly_the_old_list() {
    let mut rng = StdRng::seed_from_u64(7);
    for round in 0..300 {
        let n_own = rng.gen_range(0, 12);
        let own: Vec<(TxId, i32)> = (0..n_own).map(|_| (tid(rng.gen_range(0, 40)), rng.gen_range(1, 4))).collect();
        let n_dec = rng.gen_range(0, 60);
        let decoys: Vec<(TxId, i32)> = (0..n_dec).map(|_| (tid(rng.gen_range(0, 40)), rng.gen_range(1, 4))).collect();
        let mut old = own.clone();
        old.extend_from_slice(&decoys[..]);
        old.sort();
        old.dedup();
        let (items, own_set) = plan_fetch_list(own.clone(), decoys);
        assert_eq!(items, old, "round {}", round);
        assert_eq!(own_set, own.iter().cloned().collect::<BTreeSet<_>>());
    }
}

#[test] // T2
fn every_item_fetched_and_scanned_exactly_once() {
    for &n in &[0u32, 1, 7, 200] {
        let items: Vec<(TxId, i32)> = (0..n).map(|i| (tid(i), (i % 5) as i32 + 1)).collect();
        let own: Vec<(TxId, i32)> = items.iter().cloned().step_by(3).collect();
        let calls: CallLog = Arc::new(Mutex::new(vec![]));
        let scans: ScanLog = Arc::new(Mutex::new(vec![]));
        let r = run(items.clone(), &own, policy(8), fake_fetch(calls.clone(), 1, |_, _| false), logging_scan(scans.clone()));
        assert_eq!(r, Ok(()));
        let mut fetched: Vec<TxId> = calls.lock().unwrap().iter().map(|(t, re)| { assert!(!*re); *t }).collect();
        fetched.sort();
        let mut scanned = scans.lock().unwrap().clone();
        scanned.sort();
        assert_eq!(fetched, items.iter().map(|(t, _)| *t).collect::<Vec<TxId>>());
        assert_eq!(scanned, items);
    }
}

#[test] // T3
fn concurrency_is_bounded_by_workers() {
    for &workers in &[1usize, 2, 8] {
        let inflight = Arc::new(AtomicUsize::new(0));
        let peak = Arc::new(AtomicUsize::new(0));
        let fetch: FetchFn = {
            let (inflight, peak) = (inflight.clone(), peak.clone());
            Arc::new(move |txid: TxId, _r: bool| {
                let (inflight, peak) = (inflight.clone(), peak.clone());
                Box::pin(async move {
                    let now = inflight.fetch_add(1, Ordering::SeqCst) + 1;
                    peak.fetch_max(now, Ordering::SeqCst);
                    tokio::time::delay_for(Duration::from_millis(5)).await;
                    inflight.fetch_sub(1, Ordering::SeqCst);
                    Ok::<Vec<u8>, String>(txid.0.to_vec())
                }) as FetchFuture
            })
        };
        let items: Vec<(TxId, i32)> = (0..200).map(|i| (tid(i), 1)).collect();
        assert_eq!(run(items, &[], policy(workers), fetch, no_scan_log()), Ok(()));
        let p = peak.load(Ordering::SeqCst);
        assert!(p <= workers, "peak {} > workers {}", p, workers);
        if workers > 1 { assert!(p >= 2, "fetches never overlapped"); }
    }
}

#[test] // T4: privacy guard
fn request_sequence_does_not_depend_on_ownership() {
    let items: Vec<(TxId, i32)> = (0..40).map(|i| (tid(i), 7)).collect();
    let evens: Vec<(TxId, i32)> = items.iter().cloned().filter(|(t, _)| t.0[3] % 2 == 0).collect();
    let first10: Vec<(TxId, i32)> = items[..10].to_vec();
    let mut logs = vec![];
    for own in &[evens, first10, vec![]] {
        let calls: CallLog = Arc::new(Mutex::new(vec![]));
        let fetch = fake_fetch(calls.clone(), 0, |t, nth| t == tid(5) || t == tid(17) || (t == tid(9) && nth == 1));
        let _ = run(items.clone(), own, policy(1), fetch, no_scan_log());
        logs.push(calls.lock().unwrap().clone());
    }
    assert_eq!(logs[0], logs[1]);
    assert_eq!(logs[0], logs[2]);
    let mut seen: Vec<TxId> = vec![];
    for (t, _) in &logs[0] { if !seen.contains(t) { seen.push(*t); } }
    assert_eq!(seen, items.iter().map(|(t, _)| *t).collect::<Vec<TxId>>()); // FIFO, plan order
}

#[test] // T5
fn decoy_failures_are_ignored_and_retries_redial() {
    let items: Vec<(TxId, i32)> = (0..10).map(|i| (tid(i), 3)).collect();
    let calls: CallLog = Arc::new(Mutex::new(vec![]));
    let scans: ScanLog = Arc::new(Mutex::new(vec![]));
    let fetch = fake_fetch(calls.clone(), 0, |t, nth| t == tid(5) || t == tid(7) || (t == tid(3) && nth == 1));
    assert_eq!(run(items, &[(tid(3), 3)], policy(4), fetch, logging_scan(scans.clone())), Ok(()));
    let calls = calls.lock().unwrap().clone();
    let flags = |id: TxId| calls.iter().filter(|(t, _)| *t == id).map(|(_, r)| *r).collect::<Vec<bool>>();
    assert_eq!(flags(tid(3)), vec![false, true]);
    assert_eq!(flags(tid(5)), vec![false, true, true]);
    assert_eq!(flags(tid(0)), vec![false]);
    assert!(scans.lock().unwrap().contains(&(tid(3), 3)));
}

#[test] // T6
fn own_failure_is_an_error_but_never_stops_the_batch() {
    let items: Vec<(TxId, i32)> = (0..20).map(|i| (tid(i), 1)).collect();
    let calls: CallLog = Arc::new(Mutex::new(vec![]));
    let scans: ScanLog = Arc::new(Mutex::new(vec![]));
    let r = run(items, &[(tid(0), 1), (tid(10), 1)], policy(1), fake_fetch(calls.clone(), 0, |t, _| t == tid(0)), logging_scan(scans.clone()));
    assert!(r.unwrap_err().contains("fake failure"));
    let fetched: BTreeSet<TxId> = calls.lock().unwrap().iter().map(|(t, _)| *t).collect();
    assert_eq!(fetched.len(), 20);
    assert!(scans.lock().unwrap().contains(&(tid(10), 1)));
}

#[test] // T7
fn dead_network_trips_the_breaker() {
    let p = policy(8);
    for &with_own in &[true, false] {
        let items: Vec<(TxId, i32)> = (0..1000).map(|i| (tid(i), 1)).collect();
        let own = if with_own { vec![(tid(500), 1), (tid(999), 1)] } else { vec![] };
        let calls: CallLog = Arc::new(Mutex::new(vec![]));
        let start = Instant::now();
        let r = run(items, &own, p, fake_fetch(calls.clone(), 0, |_, _| true), no_scan_log());
        assert_eq!(r.is_err(), with_own);
        let n = calls.lock().unwrap().len();
        assert!(n <= (p.fail_breaker + 2 * p.workers) * p.max_attempts as usize, "{} calls", n);
        assert!(start.elapsed() < Duration::from_secs(20));
    }
}

#[test] // T8
fn stalled_request_times_out_and_redials() {
    let calls: CallLog = Arc::new(Mutex::new(vec![]));
    let fetch: FetchFn = {
        let calls = calls.clone();
        Arc::new(move |txid: TxId, reconnect: bool| {
            let calls = calls.clone();
            Box::pin(async move {
                calls.lock().unwrap().push((txid, reconnect));
                if !reconnect { tokio::time::delay_for(Duration::from_secs(30)).await; }
                Ok::<Vec<u8>, String>(txid.0.to_vec())
            }) as FetchFuture
        })
    };
    let mut p = policy(2);
    p.attempt_timeout = Duration::from_millis(100);
    let start = Instant::now();
    assert_eq!(run(vec![(tid(1), 1), (tid(2), 1)], &[(tid(1), 1)], p, fetch, no_scan_log()), Ok(()));
    assert!(start.elapsed() < Duration::from_secs(5));
    assert_eq!(calls.lock().unwrap().iter().filter(|(t, _)| *t == tid(1)).count(), 2);
}

#[test] // T9
fn panicking_scan_never_hangs() {
    for &(own_panics, expect_ok) in &[(false, true), (true, false)] {
        let (tx, rx) = mpsc::channel();
        std::thread::spawn(move || {
            let items: Vec<(TxId, i32)> = (0..30).map(|i| (tid(i), 1)).collect();
            let bad = if own_panics { tid(1) } else { tid(2) };
            let scan: ScanFn = Arc::new(move |bytes: Vec<u8>, _h: i32| {
                if bytes[..] == bad.0[..] { panic!("scan boom"); }
                Ok(())
            });
            let r = run(items, &[(tid(1), 1)], policy(4), fake_fetch(Arc::new(Mutex::new(vec![])), 0, |_, _| false), scan);
            let _ = tx.send(r);
        });
        let r = rx.recv_timeout(Duration::from_secs(30)).expect("engine hung after a scan panic");
        assert_eq!(r.is_ok(), expect_ok);
    }
}

fn main_config() -> LightClientConfig {
    // chain "main": LightClientConfig::base58_pubkey_address() panics for any other chain name.
    LightClientConfig {
        server: "0.0.0.0:0".parse().unwrap(),
        chain_name: "main".to_string(),
        sapling_activation_height: 0,
        consensus_branch_id: "000000".to_string(),
        anchor_offset: 0,
        no_cert_verification: false,
        data_dir: None,
    }
}

#[test] // T10
fn unparsable_tx_bytes_are_an_error_not_a_panic_or_hang() {
    let wallet = Arc::new(RwLock::new(LightWallet::new(None, &main_config(), 0, 0).unwrap()));
    let scan = wallet_scan_fn(wallet);
    assert!(scan(vec![], 1).is_err());
    for &(own_bad, expect_ok) in &[(true, false), (false, true)] {
        let fetch: FetchFn = Arc::new(move |_txid: TxId, _r: bool| {
            Box::pin(async move { Ok::<Vec<u8>, String>(vec![0xffu8; 64]) }) as FetchFuture
        });
        let own = if own_bad { vec![(tid(1), 1)] } else { vec![] };
        let scan = scan.clone();
        let (tx, rx) = mpsc::channel();
        std::thread::spawn(move || {
            let _ = tx.send(run(vec![(tid(1), 1), (tid(2), 1)], &own, policy(2), fetch, scan));
        });
        let r = rx.recv_timeout(Duration::from_secs(30)).expect("engine hung on garbage bytes");
        assert_eq!(r.is_ok(), expect_ok);
    }
}

fn my_taddr_hash(wallet: &LightWallet) -> [u8; 20] {
    let raw = wallet.get_all_taddresses()[0].from_base58().unwrap();
    assert_eq!(raw.len(), 25, "prefix(1) + hash160(20) + checksum(4)");
    let mut h = [0u8; 20];
    h.copy_from_slice(&raw[1..21]);
    h
}

fn fake_tx(txid: TxId, spends: &[(TxId, u32)], pays: &[([u8; 20], u64)]) -> Transaction {
    let mut tx = Transaction { txid, data: TransactionData::new() };
    for &(prev, n) in spends {
        tx.data.vin.push(TxIn { prevout: OutPoint { hash: prev.0, n }, script_sig: Script(vec![]), sequence: 0 });
    }
    for &(hash, value) in pays {
        tx.data.vout.push(TxOut { value: Amount::from_u64(value).unwrap(), script_pubkey: TransparentAddress::PublicKey(hash).script() });
    }
    tx
}

fn wallet_bytes(w: &Arc<RwLock<LightWallet>>) -> Vec<u8> {
    let mut v = vec![];
    w.read().unwrap().write(&mut v).unwrap();
    v
}

#[test] // T11: the engine leaves the same wallet file as the old serial fill
fn engine_leaves_same_wallet_bytes_as_serial_fill() {
    let cfg = main_config();
    let w0 = LightWallet::new(None, &cfg, 0, 0).unwrap();
    let me = my_taddr_hash(&w0);
    let other = [7u8; 20];
    let mut txs: HashMap<TxId, Transaction> = HashMap::new();
    // 12 own txs paying us. The first one expands the HD gap (serially, before the snapshot), so the
    // batch below never races ensure_hd_taddresses (a pre-existing race that exists today too).
    let mut own: Vec<(TxId, i32)> = vec![];
    for i in 0..12u32 {
        let tx = fake_tx(tid(1000 + i), &[], &[(me, 1000 + i as u64), (other, 5)]);
        w0.scan_full_tx(&tx, 10 + i as i32, 0);
        own.push((tx.txid, 10 + i as i32));
        txs.insert(tx.txid, tx);
    }
    assert!(w0.tbalance(None) > 0, "t-address hash did not match the wallet");
    for (id, _) in &own { w0.txs.write().unwrap().get_mut(id).unwrap().full_tx_scanned = false; }
    let mut snap = vec![];
    w0.write(&mut snap).unwrap();
    // Decoys with no intra-batch dependencies: spenders of pre-batch utxos, payers to us, unrelated.
    let mut decoys: Vec<(TxId, i32)> = own.clone();
    for i in 0..6u32 { let t = fake_tx(tid(2000 + i), &[(tid(1000 + i), 0)], &[(other, 900)]); decoys.push((t.txid, 200 + i as i32)); txs.insert(t.txid, t); }
    for i in 0..8u32 { let t = fake_tx(tid(3000 + i), &[], &[(me, 77 + i as u64)]); decoys.push((t.txid, 300 + i as i32)); txs.insert(t.txid, t); }
    for i in 0..30u32 { let t = fake_tx(tid(4000 + i), &[], &[(other, 1)]); decoys.push((t.txid, 400 + i as i32)); txs.insert(t.txid, t); }
    decoys.push((tid(4000), 999)); // same unrelated txid at a second height
    let (items, own_set) = plan_fetch_list(own.clone(), decoys);
    let txs = Arc::new(txs);

    let wa = Arc::new(RwLock::new(LightWallet::read(&snap[..], &cfg).unwrap()));
    for (id, h) in &items { wa.read().unwrap().scan_full_tx(&txs[id], *h, 0); } // old fill, serial
    let expected = wallet_bytes(&wa);

    for seed in 0..5u8 {
        let wb = Arc::new(RwLock::new(LightWallet::read(&snap[..], &cfg).unwrap()));
        let scan: ScanFn = {
            let (wb, txs) = (wb.clone(), txs.clone());
            Arc::new(move |bytes: Vec<u8>, h: i32| {
                let mut b = [0u8; 32];
                b.copy_from_slice(&bytes[..32]);
                match txs.get(&TxId(b)) {
                    Some(tx) => { wb.read().unwrap().scan_full_tx(tx, h, 0); Ok(()) }
                    None => Err("unknown txid".to_string()),
                }
            })
        };
        let fetch: FetchFn = Arc::new(move |txid: TxId, _r: bool| {
            Box::pin(async move {
                tokio::time::delay_for(Duration::from_millis(((txid.0[3] ^ seed) % 4) as u64)).await;
                Ok::<Vec<u8>, String>(txid.0.to_vec())
            }) as FetchFuture
        });
        let mut rt = Runtime::new().unwrap();
        assert_eq!(run_fetch_scan(&mut rt, items.clone(), own_set.clone(), policy(8), ThreadPool::new(8), fetch, scan), Ok(()));
        assert!(wallet_bytes(&wb) == expected, "wallet bytes differ from the serial fill (seed {})", seed);
        for (id, _) in &own { assert!(wb.read().unwrap().txs.read().unwrap()[id].full_tx_scanned); }
    }
}