diff --git a/Cargo.toml b/Cargo.toml index 09b8906..f9404fe 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -61,5 +61,8 @@ features = ["ff_derive"] [build-dependencies] tower-grpc-build = { git = "https://github.com/tower-rs/tower-grpc", features = ["tower-hyper"] } +[dev-dependencies] +rand_core = "0.5.1" + [profile.release] debug = true \ No newline at end of file diff --git a/src/lightwallet/mod.rs b/src/lightwallet/mod.rs index d994f11..650b345 100644 --- a/src/lightwallet/mod.rs +++ b/src/lightwallet/mod.rs @@ -640,10 +640,10 @@ impl LightWallet { continue; } + // Update the WalletTx + info!("A sapling output was sent in {}", tx.txid()); { - // Update the WalletTx // Do it in a short scope because of the write lock. - info!("A sapling output was sent in {}", tx.txid()); let mut txs = self.txs.write().unwrap(); if txs.get(&tx.txid()).unwrap().outgoing_metadata.iter() @@ -1078,3 +1078,220 @@ impl LightWallet { Some(raw_tx.into_boxed_slice()) } } + + + +#[cfg(test)] +pub mod tests { + use std::convert::TryInto; + use ff::{Field, PrimeField, PrimeFieldRepr}; + use pairing::bls12_381::Bls12; + use rand_core::{RngCore, OsRng}; + use protobuf::Message; + use zcash_client_backend::{encoding::encode_payment_address, + proto::compact_formats::{ + CompactBlock, CompactOutput, CompactSpend, CompactTx, + } + }; + use zcash_primitives::{ + block::BlockHash, + jubjub::fs::Fs, + note_encryption::{Memo, SaplingNoteEncryption}, + primitives::{Note, PaymentAddress}, + transaction::components::Amount, + zip32::{ExtendedFullViewingKey, ExtendedSpendingKey}, + JUBJUB, + }; + + use crate::lightwallet::LightWallet; + use crate::LightClientConfig; + + /// Create a fake CompactBlock at the given height, containing a single output paying + /// the given address. Returns the CompactBlock and the nullifier for the new note. + pub(crate) fn fake_compact_block( + height: i32, + prev_hash: BlockHash, + extfvk: ExtendedFullViewingKey, + value: Amount, + ) -> (CompactBlock, Vec) { + let to = extfvk.default_address().unwrap().1; + + // Create a fake Note for the account + let mut rng = OsRng; + let note = Note { + g_d: to.diversifier.g_d::(&JUBJUB).unwrap(), + pk_d: to.pk_d.clone(), + value: value.into(), + r: Fs::random(&mut rng), + }; + let encryptor = SaplingNoteEncryption::new( + extfvk.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(); + + // Create a fake CompactBlock containing the note + let mut cout = CompactOutput::new(); + cout.set_cmu(cmu); + cout.set_epk(epk); + cout.set_ciphertext(enc_ciphertext[..52].to_vec()); + let mut ctx = CompactTx::new(); + let mut txid = vec![0; 32]; + rng.fill_bytes(&mut txid); + ctx.set_hash(txid); + ctx.outputs.push(cout); + let mut cb = CompactBlock::new(); + cb.set_height(height as u64); + cb.hash.resize(32, 0); + rng.fill_bytes(&mut cb.hash); + cb.prevHash.extend_from_slice(&prev_hash.0); + cb.vtx.push(ctx); + (cb, note.nf(&extfvk.fvk.vk, 0, &JUBJUB)) + } + + /// Create a fake CompactBlock at the given height, spending a single note from the + /// given address. + pub(crate) fn fake_compact_block_spending( + height: i32, + prev_hash: BlockHash, + (nf, in_value): (Vec, Amount), + extfvk: ExtendedFullViewingKey, + to: PaymentAddress, + value: Amount, + ) -> CompactBlock { + let mut rng = OsRng; + + // Create a fake CompactBlock containing the note + let mut cspend = CompactSpend::new(); + cspend.set_nf(nf); + let mut ctx = CompactTx::new(); + let mut txid = vec![0; 32]; + rng.fill_bytes(&mut txid); + ctx.set_hash(txid); + ctx.spends.push(cspend); + + // Create a fake Note for the payment + ctx.outputs.push({ + let note = Note { + g_d: to.diversifier.g_d::(&JUBJUB).unwrap(), + pk_d: to.pk_d.clone(), + value: value.into(), + r: Fs::random(&mut rng), + }; + let encryptor = SaplingNoteEncryption::new( + extfvk.fvk.ovk, + note.clone(), + to, + 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 + }); + + // Create a fake Note for the change + ctx.outputs.push({ + let change_addr = extfvk.default_address().unwrap().1; + let note = Note { + g_d: change_addr.diversifier.g_d::(&JUBJUB).unwrap(), + pk_d: change_addr.pk_d.clone(), + value: (in_value - value).into(), + r: Fs::random(&mut rng), + }; + let encryptor = SaplingNoteEncryption::new( + extfvk.fvk.ovk, + note.clone(), + change_addr, + 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 + }); + + let mut cb = CompactBlock::new(); + cb.set_height(height as u64); + cb.hash.resize(32, 0); + rng.fill_bytes(&mut cb.hash); + cb.prevHash.extend_from_slice(&prev_hash.0); + cb.vtx.push(ctx); + cb + } + + #[test] + fn z_balances() { + let wallet = LightWallet::new(None, &LightClientConfig { + server: "0.0.0.0:0".to_string(), + chain_name: "test".to_string(), + sapling_activation_height: 0 + }).unwrap(); + + const AMOUNT1:u64 = 5; + // Address is encoded in bech32 + let address = Some(encode_payment_address(wallet.config.hrp_sapling_address(), + &wallet.extfvks[0].default_address().unwrap().1)); + + let (cb1, _) = fake_compact_block( + 0, + BlockHash([0; 32]), + wallet.extfvks[0].clone(), + Amount::from_u64(AMOUNT1).unwrap(), + ); + + + // Make sure that the intial state is empty + assert_eq!(wallet.txs.read().unwrap().len(), 0); + assert_eq!(wallet.blocks.read().unwrap().len(), 0); + assert_eq!(wallet.zbalance(None), 0); + assert_eq!(wallet.zbalance(address.clone()), 0); + + wallet.scan_block(&cb1.write_to_bytes().unwrap()).unwrap(); + + assert_eq!(wallet.txs.read().unwrap().len(), 1); + assert_eq!(wallet.blocks.read().unwrap().len(), 1); + assert_eq!(wallet.zbalance(None), AMOUNT1); + assert_eq!(wallet.zbalance(address.clone()), AMOUNT1); + + const AMOUNT2:u64 = 10; + + // Add a second block + let (cb2, _) = fake_compact_block( + 1, // Block number 1 + BlockHash{0: cb1.hash[..].try_into().unwrap()}, + wallet.extfvks[0].clone(), + Amount::from_u64(AMOUNT2).unwrap(), + ); + + wallet.scan_block(&cb2.write_to_bytes().unwrap()).unwrap(); + + assert_eq!(wallet.txs.read().unwrap().len(), 2); + assert_eq!(wallet.blocks.read().unwrap().len(), 2); + assert_eq!(wallet.zbalance(None), AMOUNT1 + AMOUNT2); + assert_eq!(wallet.zbalance(address.clone()), AMOUNT1 + AMOUNT2); + } +} \ No newline at end of file