Add test for zcash#899.

Signed-off-by: Daira Emma Hopwood <[email protected]>

zcash_client_sqlite/src/testing.rs

@@ -34,7 +34,7 @@ use zcash_client_backend::{
     wallet::OvkPolicy,
     zip321,
 };
-use zcash_note_encryption::Domain;
+use zcash_note_encryption::{Domain, COMPACT_NOTE_SIZE};
 use zcash_primitives::{
     block::BlockHash,
     consensus::{self, BlockHeight, Network, NetworkUpgrade, Parameters},
@@ -49,7 +49,7 @@ use zcash_primitives::{
     transaction::{
         components::{amount::BalanceError, Amount},
         fees::FeeRule,
-        TxId,
+        Transaction, TxId,
     },
     zip32::{sapling::DiversifiableFullViewingKey, DiversifierIndex},
 };
@@ -263,6 +263,34 @@ where
         (height, res)
     }
 
+    /// Creates a fake block at the expected next height containing only the given
+    /// transaction, and inserts it into the cache.
+    /// This assumes that the transaction only has Sapling spends and outputs.
+    ///
+    /// This generated block will be treated as the latest block, and subsequent calls to
+    /// [`Self::generate_next_block`] will build on it.
+    pub(crate) fn generate_next_block_from_tx(
+        &mut self,
+        tx: &Transaction,
+    ) -> (BlockHeight, Cache::InsertResult) {
+        let (height, prev_hash, initial_sapling_tree_size) = self
+            .latest_cached_block
+            .map(|(prev_height, prev_hash, end_size)| (prev_height + 1, prev_hash, end_size))
+            .unwrap_or_else(|| (self.sapling_activation_height(), BlockHash([0; 32]), 0));
+
+        let cb = fake_compact_block_from_tx(height, prev_hash, tx, initial_sapling_tree_size);
+        let res = self.cache.insert(&cb);
+
+        self.latest_cached_block = Some((
+            height,
+            cb.hash(),
+            initial_sapling_tree_size
+                + cb.vtx.iter().map(|tx| tx.outputs.len() as u32).sum::<u32>(),
+        ));
+
+        (height, res)
+    }
+
     /// Invokes [`scan_cached_blocks`] with the given arguments, expecting success.
     pub(crate) fn scan_cached_blocks(&mut self, from_height: BlockHeight, limit: usize) {
         assert_matches!(self.try_scan_cached_blocks(from_height, limit), Ok(_));
@@ -288,6 +316,18 @@ where
             limit,
         )
     }
+
+    /// Reset the wallet using a new wallet database, but with the same cache of blocks.
+    /// This does not recreate accounts, nor does it rescan the cached blocks.
+    /// The resulting wallet has no test account.
+    pub(crate) fn reset(&mut self) {
+        let network = self.network();
+        self.latest_cached_block = None;
+        self._data_file = NamedTempFile::new().unwrap();
+        self.db_data = WalletDb::for_path(self._data_file.path(), network).unwrap();
+        self.test_account = None;
+        init_wallet_db(&mut self.db_data, None).unwrap();
+    }
 }
 
 impl<Cache> TestState<Cache> {
@@ -632,6 +672,36 @@ pub(crate) fn fake_compact_block<P: consensus::Parameters>(
     (cb, note.nf(&dfvk.fvk().vk.nk, 0))
 }
 
+/// Create a fake CompactBlock at the given height containing only the given transaction.
+/// This assumes that the transaction only has Sapling spends and outputs.
+pub(crate) fn fake_compact_block_from_tx(
+    height: BlockHeight,
+    prev_hash: BlockHash,
+    tx: &Transaction,
+    initial_sapling_tree_size: u32,
+) -> CompactBlock {
+    // Create a fake CompactTx
+    let mut ctx = CompactTx::default();
+    ctx.hash = tx.txid().as_ref().to_vec();
+
+    if let Some(bundle) = tx.sapling_bundle() {
+        for spend in bundle.shielded_spends() {
+            ctx.spends.push(CompactSaplingSpend {
+                nf: spend.nullifier().to_vec(),
+            });
+        }
+        for output in bundle.shielded_outputs() {
+            ctx.outputs.push(CompactSaplingOutput {
+                cmu: output.cmu().to_bytes().to_vec(),
+                ephemeral_key: output.ephemeral_key().0.to_vec(),
+                ciphertext: output.enc_ciphertext()[..COMPACT_NOTE_SIZE].to_vec(),
+            });
+        }
+    }
+
+    fake_compact_block_from_compact_tx(ctx, height, prev_hash, initial_sapling_tree_size)
+}
+
 /// Create a fake CompactBlock at the given height, spending a single note from the
 /// given address.
 #[allow(clippy::too_many_arguments)]
@@ -706,6 +776,16 @@ pub(crate) fn fake_compact_block_spending<P: consensus::Parameters>(
         }
     });
 
+    fake_compact_block_from_compact_tx(ctx, height, prev_hash, initial_sapling_tree_size)
+}
+
+pub(crate) fn fake_compact_block_from_compact_tx(
+    ctx: CompactTx,
+    height: BlockHeight,
+    prev_hash: BlockHash,
+    initial_sapling_tree_size: u32,
+) -> CompactBlock {
+    let mut rng = OsRng;
     let mut cb = CompactBlock {
         hash: {
             let mut hash = vec![0; 32];

zcash_client_sqlite/src/wallet/sapling.rs

@@ -408,6 +408,7 @@ pub(crate) mod tests {
     use zcash_proofs::prover::LocalTxProver;
 
     use zcash_primitives::{
+        block::BlockHash,
         consensus::BranchId,
         legacy::TransparentAddress,
         memo::{Memo, MemoBytes},
@@ -427,7 +428,10 @@ pub(crate) mod tests {
         data_api::{
             self,
             error::Error,
-            wallet::input_selection::{GreedyInputSelector, GreedyInputSelectorError},
+            wallet::{
+                decrypt_and_store_transaction,
+                input_selection::{GreedyInputSelector, GreedyInputSelectorError},
+            },
             ShieldedProtocol, WalletRead, WalletWrite,
         },
         decrypt_transaction,
@@ -1044,6 +1048,113 @@ pub(crate) mod tests {
         );
     }
 
+    #[test]
+    fn external_address_change_spends_detected_in_restore_from_seed() {
+        let mut st = TestBuilder::new().with_block_cache().build();
+
+        // Add an account to the wallet
+        let seed = Secret::new([0u8; 32].to_vec());
+        let (_, usk) = st.wallet_mut().create_account(&seed).unwrap();
+        let dfvk = usk.sapling().to_diversifiable_full_viewing_key();
+
+        let (_, usk2) = st.wallet_mut().create_account(&seed).unwrap();
+        let dfvk2 = usk2.sapling().to_diversifiable_full_viewing_key();
+
+        // Add funds to the wallet in a single note
+        let value = Amount::from_u64(100000).unwrap();
+        let (h, _, _) = st.generate_next_block(&dfvk, AddressType::DefaultExternal, value);
+        st.scan_cached_blocks(h, 1);
+
+        // Verified balance matches total balance
+        let (_, anchor_height) = st
+            .wallet()
+            .get_target_and_anchor_heights(NonZeroU32::new(1).unwrap())
+            .unwrap()
+            .unwrap();
+        assert_eq!(st.get_balance(AccountId::from(0)).unwrap(), value);
+        assert_eq!(
+            st.get_balance_at(AccountId::from(0), anchor_height)
+                .unwrap(),
+            value
+        );
+        assert_eq!(st.get_balance(AccountId::from(1)).unwrap(), Amount::zero());
+
+        let amount_sent = Amount::from_u64(20000).unwrap();
+        let addr = dfvk.default_address().1;
+        let addr2 = dfvk2.default_address().1;
+        let req = TransactionRequest::new(vec![
+            // payment to an external recipient
+            Payment {
+                recipient_address: RecipientAddress::Shielded(addr2),
+                amount: amount_sent,
+                memo: None,
+                label: None,
+                message: None,
+                other_params: vec![],
+            },
+            // payment back to the originating wallet, simulating legacy change
+            Payment {
+                recipient_address: RecipientAddress::Shielded(addr),
+                amount: Amount::from_u64(30000).unwrap(),
+                memo: None,
+                label: None,
+                message: None,
+                other_params: vec![],
+            },
+        ])
+        .unwrap();
+
+        let fee_rule = FixedFeeRule::standard();
+        let input_selector = GreedyInputSelector::new(
+            fixed::SingleOutputChangeStrategy::new(fee_rule),
+            DustOutputPolicy::default(),
+        );
+
+        let txid = st
+            .spend(
+                &input_selector,
+                &usk,
+                req,
+                OvkPolicy::Sender,
+                NonZeroU32::new(1).unwrap(),
+            )
+            .unwrap();
+        let tx = &st.wallet().get_transaction(txid).unwrap();
+
+        let (h, _) = st.generate_next_block_from_tx(tx);
+        st.scan_cached_blocks(h, 1);
+
+        let amount_left = (value - (amount_sent + fee_rule.fixed_fee()).unwrap()).unwrap();
+        assert_eq!(st.get_balance(AccountId::from(1)).unwrap(), amount_sent);
+        assert_eq!(st.get_balance(AccountId::from(0)).unwrap(), amount_left);
+
+        st.reset();
+
+        // Account creation and DFVK derivation should be deterministic.
+        let (_, restored_usk) = st.wallet_mut().create_account(&seed).unwrap();
+        assert_eq!(
+            restored_usk
+                .sapling()
+                .to_diversifiable_full_viewing_key()
+                .to_bytes(),
+            dfvk.to_bytes()
+        );
+
+        let (_, restored_usk2) = st.wallet_mut().create_account(&seed).unwrap();
+        assert_eq!(
+            restored_usk2
+                .sapling()
+                .to_diversifiable_full_viewing_key()
+                .to_bytes(),
+            dfvk2.to_bytes()
+        );
+
+        st.scan_cached_blocks(st.sapling_activation_height(), 2);
+
+        assert_eq!(st.get_balance(AccountId::from(1)).unwrap(), amount_sent);
+        assert_eq!(st.get_balance(AccountId::from(0)).unwrap(), amount_left);
+    }
+
     #[test]
     fn zip317_spend() {
         let mut st = TestBuilder::new().with_block_cache().build();