zcash_client_backend/data_api/testing/
sapling.rs1use 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
34pub struct SaplingPoolTester;
36impl ShieldedPoolTester for SaplingPoolTester {
37 const SHIELDED_PROTOCOL: ShieldedPool = ShieldedPool::Sapling;
38 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 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 spend.proof_generation_key().is_none().then_some(index)
211 })
212 .collect::<Vec<_>>();
213
214 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 for index in 0.. {
237 match signer.sign_sapling(UnwrapErr(SysRng), index, extsk.expsk().ask()) {
238 Err(pczt::roles::signer::Error::InvalidIndex) => break,
240 Ok(())
242 | Err(pczt::roles::signer::Error::SaplingSign(
243 sapling::pczt::SignerError::WrongSpendAuthorizingKey,
244 )) => Ok(()),
245 Err(e) => Err(e),
247 }?;
248 }
249
250 Ok(())
251 }
252}