penumbra_sdk_mock_client/
lib.rs

1use anyhow::Error;
2use cnidarium::StateRead;
3use penumbra_sdk_compact_block::{component::StateReadExt as _, CompactBlock, StatePayload};
4use penumbra_sdk_dex::{swap::SwapPlaintext, swap_claim::SwapClaimPlan};
5use penumbra_sdk_keys::{keys::SpendKey, FullViewingKey};
6use penumbra_sdk_sct::{
7    component::{clock::EpochRead, tree::SctRead},
8    Nullifier,
9};
10use penumbra_sdk_shielded_pool::{note, Note, SpendPlan};
11use penumbra_sdk_tct as tct;
12use penumbra_sdk_transaction::{AuthorizationData, Transaction, TransactionPlan, WitnessData};
13use rand_core::OsRng;
14use std::collections::BTreeMap;
15
16/// A bare-bones mock client for use exercising the state machine.
17pub struct MockClient {
18    latest_height: u64,
19    sk: SpendKey,
20    pub fvk: FullViewingKey,
21    /// All notes, whether spent or not.
22    pub notes: BTreeMap<note::StateCommitment, Note>,
23    pub nullifiers: BTreeMap<note::StateCommitment, Nullifier>,
24    /// Whether a note was spent or not.
25    pub spent_notes: BTreeMap<note::StateCommitment, ()>,
26    swaps: BTreeMap<tct::StateCommitment, SwapPlaintext>,
27    pub sct: penumbra_sdk_tct::Tree,
28}
29
30impl MockClient {
31    pub fn new(sk: SpendKey) -> MockClient {
32        Self {
33            latest_height: u64::MAX,
34            fvk: sk.full_viewing_key().clone(),
35            sk,
36            notes: Default::default(),
37            spent_notes: Default::default(),
38            nullifiers: Default::default(),
39            sct: Default::default(),
40            swaps: Default::default(),
41        }
42    }
43
44    pub async fn with_sync_to_storage(
45        mut self,
46        storage: impl AsRef<cnidarium::Storage>,
47    ) -> anyhow::Result<Self> {
48        let latest = storage.as_ref().latest_snapshot();
49        self.sync_to_latest(latest).await?;
50
51        Ok(self)
52    }
53
54    pub async fn with_sync_to_inner_storage(
55        mut self,
56        storage: cnidarium::Storage,
57    ) -> anyhow::Result<Self> {
58        let latest = storage.latest_snapshot();
59        self.sync_to_latest(latest).await?;
60
61        Ok(self)
62    }
63
64    pub async fn sync_to_latest<R: StateRead>(&mut self, state: R) -> anyhow::Result<()> {
65        let height = state.get_block_height().await?;
66        self.sync_to(height, state).await?;
67        Ok(())
68    }
69
70    pub async fn sync_to<R: StateRead>(
71        &mut self,
72        target_height: u64,
73        state: R,
74    ) -> anyhow::Result<()> {
75        let start_height = self.latest_height.wrapping_add(1);
76        for height in start_height..=target_height {
77            let compact_block = state
78                .compact_block(height)
79                .await?
80                .ok_or_else(|| anyhow::anyhow!("missing compact block for height {}", height))?;
81            self.scan_block(compact_block.try_into()?)?;
82            let (latest_height, root) = self.latest_height_and_sct_root();
83            anyhow::ensure!(latest_height == height, "latest height should be updated");
84            let expected_root = state
85                .get_anchor_by_height(height)
86                .await?
87                .ok_or_else(|| anyhow::anyhow!("missing sct anchor for height {}", height))?;
88            anyhow::ensure!(
89                root == expected_root,
90                format!(
91                    "client sct root should match chain state: {:?} != {:?}",
92                    root, expected_root
93                )
94            );
95        }
96        Ok(())
97    }
98
99    pub fn scan_block(&mut self, block: CompactBlock) -> anyhow::Result<()> {
100        use penumbra_sdk_tct::Witness::*;
101
102        if self.latest_height.wrapping_add(1) != block.height {
103            anyhow::bail!(
104                "wrong block height {} for latest height {}",
105                block.height,
106                self.latest_height
107            );
108        }
109
110        for payload in block.state_payloads {
111            match payload {
112                StatePayload::Note { note: payload, .. } => {
113                    match payload.trial_decrypt(&self.fvk) {
114                        Some(note) => {
115                            self.sct.insert(Keep, payload.note_commitment)?;
116                            let nullifier = self
117                                .nullifier(payload.note_commitment)
118                                .expect("newly inserted note should be present in sct");
119                            self.notes.insert(payload.note_commitment, note.clone());
120                            self.nullifiers.insert(payload.note_commitment, nullifier);
121                        }
122                        None => {
123                            self.sct.insert(Forget, payload.note_commitment)?;
124                        }
125                    }
126                }
127                StatePayload::Swap { swap: payload, .. } => {
128                    match payload.trial_decrypt(&self.fvk) {
129                        Some(swap) => {
130                            self.sct.insert(Keep, payload.commitment)?;
131                            // At this point, we need to retain the swap plaintext,
132                            // and also derive the expected output notes so we can
133                            // notice them while scanning later blocks.
134                            self.swaps.insert(payload.commitment, swap.clone());
135
136                            let batch_data =
137                                block.swap_outputs.get(&swap.trading_pair).ok_or_else(|| {
138                                    anyhow::anyhow!("server gave invalid compact block")
139                                })?;
140
141                            let (output_1, output_2) = swap.output_notes(batch_data);
142                            // Pre-insert the output notes into our notes table, so that
143                            // we can notice them when we scan the block where they are claimed.
144                            // TODO: We should handle tracking the nullifiers for these notes,
145                            // however they aren't inserted into the SCT at this point.
146                            // let nullifier_1 = self
147                            //     .nullifier(output_1.commit())
148                            //     .expect("newly inserted swap should be present in sct");
149                            // let nullifier_2 = self
150                            //     .nullifier(output_2.commit())
151                            //     .expect("newly inserted swap should be present in sct");
152                            self.notes.insert(output_1.commit(), output_1.clone());
153                            // self.nullifiers.insert(output_1.commit(), nullifier_1);
154                            self.notes.insert(output_2.commit(), output_2.clone());
155                            // self.nullifiers.insert(output_2.commit(), nullifier_2);
156                        }
157                        None => {
158                            self.sct.insert(Forget, payload.commitment)?;
159                        }
160                    }
161                }
162                StatePayload::RolledUp { commitment, .. } => {
163                    if self.notes.contains_key(&commitment) {
164                        // This is a note we anticipated, so retain its auth path.
165                        self.sct.insert(Keep, commitment)?;
166                    } else {
167                        // This is someone else's note.
168                        self.sct.insert(Forget, commitment)?;
169                    }
170                }
171            }
172        }
173
174        // Mark spent nullifiers
175        for nullifier in block.nullifiers {
176            // skip if we don't know about this nullifier
177            if !self.nullifiers.values().any(move |n| *n == nullifier) {
178                continue;
179            }
180
181            self.spent_notes.insert(
182                *self
183                    .nullifiers
184                    .iter()
185                    .find_map(|(k, v)| if *v == nullifier { Some(k) } else { None })
186                    .unwrap(),
187                (),
188            );
189        }
190
191        self.sct.end_block()?;
192        if block.epoch_root.is_some() {
193            self.sct.end_epoch()?;
194        }
195
196        self.latest_height = block.height;
197
198        Ok(())
199    }
200
201    pub fn latest_height_and_sct_root(&self) -> (u64, penumbra_sdk_tct::Root) {
202        (self.latest_height, self.sct.root())
203    }
204
205    pub fn note_by_commitment(&self, commitment: &note::StateCommitment) -> Option<Note> {
206        self.notes.get(commitment).cloned()
207    }
208
209    pub fn swap_by_commitment(&self, commitment: &note::StateCommitment) -> Option<SwapPlaintext> {
210        self.swaps.get(commitment).cloned()
211    }
212
213    pub fn position(
214        &self,
215        commitment: note::StateCommitment,
216    ) -> Option<penumbra_sdk_tct::Position> {
217        self.sct.witness(commitment).map(|proof| proof.position())
218    }
219
220    pub fn nullifier(&self, commitment: note::StateCommitment) -> Option<Nullifier> {
221        let position = self.position(commitment);
222
223        if position.is_none() {
224            return None;
225        }
226        let nk = self.fvk.nullifier_key();
227
228        Some(Nullifier::derive(&nk, position.unwrap(), &commitment))
229    }
230
231    pub fn witness_commitment(
232        &self,
233        commitment: note::StateCommitment,
234    ) -> Option<penumbra_sdk_tct::Proof> {
235        self.sct.witness(commitment)
236    }
237
238    pub fn witness_plan(&self, plan: &TransactionPlan) -> Result<WitnessData, Error> {
239        let spend_commitment = |spend: &SpendPlan| spend.note.commit();
240        let spends = plan.spend_plans().map(spend_commitment);
241
242        let swap_claim_commitment = |swap: &SwapClaimPlan| swap.swap_plaintext.swap_commitment();
243        let swap_claims = plan.swap_claim_plans().map(swap_claim_commitment);
244
245        let witness = |commitment| {
246            self.sct
247                .witness(commitment)
248                .ok_or_else(|| anyhow::anyhow!("note commitment {commitment:?} unknown to client"))
249                .map(|proof| (commitment, proof))
250        };
251
252        Ok(WitnessData {
253            anchor: self.sct.root(),
254            // TODO: this will only witness spends and swap claims, but not other proofs
255            state_commitment_proofs: spends
256                .chain(swap_claims)
257                .map(witness)
258                .collect::<Result<_, Error>>()?,
259        })
260    }
261
262    pub fn authorize_plan(&self, plan: &TransactionPlan) -> Result<AuthorizationData, Error> {
263        plan.authorize(OsRng, &self.sk)
264    }
265
266    pub async fn witness_auth_build(&self, plan: &TransactionPlan) -> Result<Transaction, Error> {
267        let witness_data = self.witness_plan(plan)?;
268        let auth_data = self.authorize_plan(plan)?;
269        plan.clone()
270            .build_concurrent(&self.fvk, &witness_data, &auth_data)
271            .await
272    }
273
274    pub fn notes_by_asset(
275        &self,
276        asset_id: penumbra_sdk_asset::asset::Id,
277    ) -> impl Iterator<Item = &Note> + '_ {
278        self.notes
279            .values()
280            .filter(move |n| n.asset_id() == asset_id)
281    }
282
283    pub fn spent_note(&self, commitment: &note::StateCommitment) -> bool {
284        self.spent_notes.contains_key(commitment)
285    }
286
287    pub fn spendable_notes_by_asset(
288        &self,
289        asset_id: penumbra_sdk_asset::asset::Id,
290    ) -> impl Iterator<Item = &Note> + '_ {
291        self.notes
292            .values()
293            .filter(move |n| n.asset_id() == asset_id && !self.spent_note(&n.commit()))
294    }
295}
penumbra_sdk_mock_client/lib.rs

penumbra_sdk_mock_client/
lib.rs
