Skip to main content

zcash_client_backend/data_api/testing/
sapling.rs

1use std::hash::Hash;
2
3use incrementalmerkletree::{Address as TreeAddress, Hashable, Level, Position};
4use sapling::{
5    note_encryption::try_sapling_output_recovery,
6    zip32::{DiversifiableFullViewingKey, ExtendedSpendingKey},
7};
8use shardtree::error::ShardTreeError;
9use zcash_keys::{address::Address, keys::UnifiedSpendingKey};
10use zcash_primitives::transaction::{Transaction, components::sapling::zip212_enforcement};
11use zcash_protocol::{
12    ShieldedPool,
13    consensus::{self, BlockHeight},
14    memo::MemoBytes,
15    value::Zatoshis,
16};
17use zip32::Scope;
18
19use crate::{
20    data_api::{
21        DecryptedTransaction, InputSource, TargetValue, WalletCommitmentTrees, WalletSummary,
22        WalletTest,
23        chain::{CommitmentTreeRoot, ScanSummary},
24        wallet::{
25            ConfirmationsPolicy, TargetHeight,
26            input_selection::{LockFilter, LockedInputPolicy},
27        },
28    },
29    wallet::{Note, ReceivedNote},
30};
31
32use super::{TestState, pool::ShieldedPoolTester};
33
34/// Type for running pool-agnostic tests on the Sapling pool.
35pub struct SaplingPoolTester;
36impl ShieldedPoolTester for SaplingPoolTester {
37    const SHIELDED_PROTOCOL: ShieldedPool = ShieldedPool::Sapling;
38    // const MERKLE_TREE_DEPTH: u8 = sapling::NOTE_COMMITMENT_TREE_DEPTH;
39
40    type Sk = ExtendedSpendingKey;
41    type Fvk = DiversifiableFullViewingKey;
42    type MerkleTreeHash = sapling::Node;
43    type Note = sapling::Note;
44
45    fn test_account_fvk<Cache, DbT: WalletTest, P: consensus::Parameters>(
46        st: &TestState<Cache, DbT, P>,
47    ) -> Self::Fvk {
48        st.test_account_sapling().unwrap().clone()
49    }
50
51    fn usk_to_sk(usk: &UnifiedSpendingKey) -> &Self::Sk {
52        usk.sapling()
53    }
54
55    fn sk(seed: &[u8]) -> Self::Sk {
56        ExtendedSpendingKey::master(seed).expect("the derivation path yields a valid key")
57    }
58
59    fn sk_to_fvk(sk: &Self::Sk) -> Self::Fvk {
60        sk.to_diversifiable_full_viewing_key()
61    }
62
63    fn sk_default_address(sk: &Self::Sk) -> Address {
64        sk.default_address().1.into()
65    }
66
67    fn fvk_default_address(fvk: &Self::Fvk) -> Address {
68        fvk.default_address().1.into()
69    }
70
71    fn fvks_equal(a: &Self::Fvk, b: &Self::Fvk) -> bool {
72        a.to_bytes() == b.to_bytes()
73    }
74
75    fn empty_tree_leaf() -> Self::MerkleTreeHash {
76        ::sapling::Node::empty_leaf()
77    }
78
79    fn empty_tree_root(level: Level) -> Self::MerkleTreeHash {
80        ::sapling::Node::empty_root(level)
81    }
82
83    fn put_subtree_roots<Cache, DbT: WalletTest + WalletCommitmentTrees, P>(
84        st: &mut TestState<Cache, DbT, P>,
85        start_index: u64,
86        roots: &[CommitmentTreeRoot<Self::MerkleTreeHash>],
87    ) -> Result<(), ShardTreeError<<DbT as WalletCommitmentTrees>::Error>> {
88        st.wallet_mut()
89            .put_sapling_subtree_roots(start_index, roots)
90    }
91
92    fn shard_root<Cache, DbT: WalletTest + WalletCommitmentTrees, P>(
93        st: &mut TestState<Cache, DbT, P>,
94        shard_index: u64,
95    ) -> Result<Self::MerkleTreeHash, ShardTreeError<<DbT as WalletCommitmentTrees>::Error>> {
96        let shard_height = crate::data_api::SAPLING_SHARD_HEIGHT;
97        let addr = TreeAddress::from_parts(Level::from(shard_height), shard_index);
98        let end_position = Position::from((shard_index + 1) << shard_height);
99        st.wallet_mut()
100            .with_sapling_tree_mut(|tree| tree.root(addr, end_position))
101    }
102
103    fn next_subtree_index<A: Hash + Eq>(s: &WalletSummary<A>) -> u64 {
104        s.next_sapling_subtree_index()
105    }
106
107    fn note_value(note: &Self::Note) -> Zatoshis {
108        Zatoshis::const_from_u64(note.value().inner())
109    }
110
111    fn select_spendable_notes<Cache, DbT: InputSource + WalletTest, P>(
112        st: &TestState<Cache, DbT, P>,
113        account: <DbT as InputSource>::AccountId,
114        target_value: TargetValue,
115        target_height: TargetHeight,
116        confirmations_policy: ConfirmationsPolicy,
117        exclude: &[DbT::NoteRef],
118    ) -> Result<Vec<ReceivedNote<DbT::NoteRef, Self::Note>>, <DbT as InputSource>::Error> {
119        st.wallet()
120            .select_spendable_notes(
121                account,
122                target_value,
123                &[ShieldedPool::Sapling],
124                target_height,
125                confirmations_policy,
126                exclude,
127                LockFilter::Policy(&LockedInputPolicy::Exclude),
128            )
129            .map(|n| n.take_sapling())
130    }
131
132    fn select_unspent_notes<Cache, DbT: InputSource + WalletTest, P>(
133        st: &TestState<Cache, DbT, P>,
134        account: <DbT as InputSource>::AccountId,
135        target_height: TargetHeight,
136        exclude: &[DbT::NoteRef],
137    ) -> Result<Vec<ReceivedNote<DbT::NoteRef, Self::Note>>, <DbT as InputSource>::Error> {
138        st.wallet()
139            .select_unspent_notes(
140                account,
141                &[ShieldedPool::Sapling],
142                target_height,
143                exclude,
144                LockFilter::Policy(&LockedInputPolicy::Exclude),
145            )
146            .map(|n| n.take_sapling())
147    }
148
149    fn decrypted_pool_outputs_count<A>(d_tx: &DecryptedTransaction<Transaction, A>) -> usize {
150        d_tx.sapling_outputs().len()
151    }
152
153    fn with_decrypted_pool_memos<A>(
154        d_tx: &DecryptedTransaction<Transaction, A>,
155        mut f: impl FnMut(&MemoBytes),
156    ) {
157        for output in d_tx.sapling_outputs() {
158            f(output.memo());
159        }
160    }
161
162    fn try_output_recovery<P: consensus::Parameters>(
163        params: &P,
164        height: BlockHeight,
165        tx: &Transaction,
166        fvk: &Self::Fvk,
167    ) -> Option<(Note, Address, MemoBytes)> {
168        for output in tx.sapling_bundle().unwrap().shielded_outputs() {
169            // Find the output that decrypts with the external OVK
170            let result = try_sapling_output_recovery(
171                &fvk.to_ovk(Scope::External),
172                output,
173                zip212_enforcement(params, height),
174            );
175
176            if result.is_some() {
177                return result.map(|(note, addr, memo)| {
178                    (
179                        Note::Sapling(note),
180                        addr.into(),
181                        MemoBytes::from_bytes(&memo).expect("correct length"),
182                    )
183                });
184            }
185        }
186
187        None
188    }
189
190    fn received_note_count(summary: &ScanSummary) -> usize {
191        summary.received_sapling_note_count()
192    }
193
194    #[cfg(feature = "pczt")]
195    fn add_proof_generation_keys(
196        pczt: pczt::Pczt,
197        usk: &UnifiedSpendingKey,
198    ) -> Result<pczt::Pczt, pczt::roles::updater::SaplingError> {
199        let extsk = Self::usk_to_sk(usk);
200
201        Ok(pczt::roles::updater::Updater::new(pczt)
202            .update_sapling_with(|mut updater| {
203                let non_dummy_spends = updater
204                    .bundle()
205                    .spends()
206                    .iter()
207                    .enumerate()
208                    .filter_map(|(index, spend)| {
209                        // Dummy spends will already have a proof generation key.
210                        spend.proof_generation_key().is_none().then_some(index)
211                    })
212                    .collect::<Vec<_>>();
213
214                // Assume all non-dummy spent notes are from the same account.
215                for index in non_dummy_spends {
216                    updater.update_spend_with(index, |mut spend_updater| {
217                        spend_updater.set_proof_generation_key(extsk.expsk().proof_generation_key())
218                    })?;
219                }
220
221                Ok(())
222            })?
223            .finish())
224    }
225
226    #[cfg(feature = "pczt")]
227    fn apply_signatures_to_pczt(
228        signer: &mut pczt::roles::signer::Signer,
229        usk: &UnifiedSpendingKey,
230    ) -> Result<(), pczt::roles::signer::Error> {
231        use rand::{rand_core::UnwrapErr, rngs::SysRng};
232
233        let extsk = Self::usk_to_sk(usk);
234
235        // Figuring out which one is for us is hard. Let's just try signing all of them!
236        for index in 0.. {
237            match signer.sign_sapling(UnwrapErr(SysRng), index, extsk.expsk().ask()) {
238                // Loop termination.
239                Err(pczt::roles::signer::Error::InvalidIndex) => break,
240                // Ignore any errors due to using the wrong key.
241                Ok(())
242                | Err(pczt::roles::signer::Error::SaplingSign(
243                    sapling::pczt::SignerError::WrongSpendAuthorizingKey,
244                )) => Ok(()),
245                // Raise any unexpected errors.
246                Err(e) => Err(e),
247            }?;
248        }
249
250        Ok(())
251    }
252}