Skip to main content

miden_client/test_utils/
note_transport.rs

1use alloc::boxed::Box;
2use alloc::collections::BTreeMap;
3use alloc::string::ToString;
4use alloc::sync::Arc;
5use alloc::vec::Vec;
6use core::sync::atomic::{AtomicUsize, Ordering};
7
8use miden_protocol::block::BlockNumber;
9use miden_protocol::note::{NoteHeader, NoteId, NoteInclusionProof, NoteTag};
10use miden_tx::utils::serde::{
11    ByteReader,
12    ByteWriter,
13    Deserializable,
14    DeserializationError,
15    Serializable,
16};
17use miden_tx::utils::sync::RwLock;
18
19use crate::note_transport::{
20    NoteInfo,
21    NoteTransportClient,
22    NoteTransportCursor,
23    NoteTransportError,
24    NoteTransportPage,
25    TransportNote,
26};
27
28/// Mock Note Transport Node
29///
30/// Simulates the functionality of the note transport node.
31#[derive(Clone)]
32pub struct MockNoteTransportNode {
33    notes: BTreeMap<NoteTag, Vec<(NoteInfo, NoteTransportCursor)>>,
34    nonce: u64,
35    next_sequence: u64,
36    /// Optional per-response batch cap; if `Some(n)`, `get_notes` returns at most `n` entries
37    /// (total, across all tags) in one call. Used to exercise client-side pagination drain loops.
38    /// `None` = unbounded (legacy behavior).
39    max_batch: Option<usize>,
40    /// Notes stored through the with-proof path, with the block their proof named.
41    proven_notes: BTreeMap<NoteId, BlockNumber>,
42}
43
44impl MockNoteTransportNode {
45    pub fn new() -> Self {
46        Self {
47            notes: BTreeMap::default(),
48            nonce: 1,
49            next_sequence: 1,
50            max_batch: None,
51            proven_notes: BTreeMap::default(),
52        }
53    }
54
55    /// Build a mock that caps each `get_notes` response at `max_batch` entries.
56    pub fn with_max_batch(max_batch: usize) -> Self {
57        Self {
58            notes: BTreeMap::default(),
59            nonce: 1,
60            next_sequence: 1,
61            max_batch: Some(max_batch),
62            proven_notes: BTreeMap::default(),
63        }
64    }
65
66    /// Seed a note relayed with its inclusion proof. The real service verifies the proof against
67    /// its node; the mock only records the proof's block and serves it as the commitment block.
68    ///
69    /// The real service stores a note only once. The mock also ignores a note with an id that it
70    /// already stores.
71    pub fn add_note_with_proof(
72        &mut self,
73        header: NoteHeader,
74        details_bytes: Vec<u8>,
75        inclusion_proof: &NoteInclusionProof,
76    ) {
77        if self.contains_note(&header.id()) {
78            return;
79        }
80        let block_num = inclusion_proof.location().block_num();
81        self.proven_notes.insert(header.id(), block_num);
82        self.add_note_after(header, details_bytes, Some(block_num));
83    }
84
85    /// Returns whether the mock stores a note with `note_id`.
86    fn contains_note(&self, note_id: &NoteId) -> bool {
87        self.notes.values().flatten().any(|(info, _)| info.header.id() == *note_id)
88    }
89
90    /// Returns the block named by the proof a note was stored with, or `None` when the note was not
91    /// stored through the with-proof path.
92    pub fn proven_block(&self, note_id: &NoteId) -> Option<BlockNumber> {
93        self.proven_notes.get(note_id).copied()
94    }
95
96    pub fn add_note(&mut self, header: NoteHeader, details_bytes: Vec<u8>) {
97        self.add_note_after(header, details_bytes, None);
98    }
99
100    /// Seed a note that carries a commitment block floor, as a transport that stored the note
101    /// without verifying a proof would serve it.
102    pub fn add_note_after(
103        &mut self,
104        header: NoteHeader,
105        details_bytes: Vec<u8>,
106        block_hint: Option<BlockNumber>,
107    ) {
108        let tag = header.metadata().tag();
109        let info = NoteInfo { header, details_bytes, block_hint };
110        let cursor = NoteTransportCursor {
111            nonce: self.nonce,
112            sequence: self.next_sequence,
113        };
114        self.next_sequence += 1;
115        self.notes.entry(tag).or_default().push((info, cursor));
116    }
117
118    /// Seed a note under an arbitrary transport tag key, regardless of the note's own tag.
119    pub fn add_note_with_tag_key(
120        &mut self,
121        tag: NoteTag,
122        header: NoteHeader,
123        details_bytes: Vec<u8>,
124    ) {
125        let info = NoteInfo { header, details_bytes, block_hint: None };
126        let cursor = NoteTransportCursor {
127            nonce: self.nonce,
128            sequence: self.next_sequence,
129        };
130        self.next_sequence += 1;
131        self.notes.entry(tag).or_default().push((info, cursor));
132    }
133
134    pub fn get_notes(
135        &self,
136        tags: &[NoteTag],
137        cursor: Option<NoteTransportCursor>,
138    ) -> (Vec<NoteInfo>, NoteTransportCursor) {
139        self.notes_up_to(tags, cursor, self.head())
140    }
141
142    /// Returns the position of the last stored note.
143    fn head(&self) -> NoteTransportCursor {
144        NoteTransportCursor {
145            nonce: self.nonce,
146            sequence: self.next_sequence - 1,
147        }
148    }
149
150    /// Returns the notes after `cursor` and up to `sync_to`.
151    fn notes_up_to(
152        &self,
153        tags: &[NoteTag],
154        cursor: Option<NoteTransportCursor>,
155        sync_to: NoteTransportCursor,
156    ) -> (Vec<NoteInfo>, NoteTransportCursor) {
157        // A cursor from another database starts from the first retained note.
158        let after = cursor.filter(|cursor| cursor.nonce == self.nonce).map_or(0, |c| c.sequence);
159        let mut collected: Vec<(NoteInfo, NoteTransportCursor)> = tags
160            .iter()
161            .filter_map(|tag| self.notes.get(tag))
162            .flatten()
163            .filter(|(_, tcursor)| tcursor.sequence > after && tcursor.sequence <= sync_to.sequence)
164            .cloned()
165            .collect();
166
167        // Deterministic ordering across tags: sort by cursor ascending so the client sees notes in
168        // per-cursor order regardless of tag iteration order, matching the real server's `ORDER BY
169        // seq ASC`.
170        collected.sort_by_key(|(_, c)| c.sequence);
171
172        // Apply the batch cap, if configured.
173        if let Some(max) = self.max_batch {
174            collected.truncate(max);
175        }
176
177        // Matches the real server's contract (`rcursor = max(cursor, max_seq_returned)`), so an
178        // empty batch returns the caller's own position.
179        let sequence = collected.last().map_or(after, |(_, c)| c.sequence);
180        let notes = collected.into_iter().map(|(n, _)| n).collect();
181        (notes, NoteTransportCursor { nonce: self.nonce, sequence })
182    }
183}
184
185impl Default for MockNoteTransportNode {
186    fn default() -> Self {
187        Self::new()
188    }
189}
190
191/// Mock Note Transport API
192///
193/// Simulates communications with the note transport node.
194#[derive(Clone, Default)]
195pub struct MockNoteTransportApi {
196    pub mock_node: Arc<RwLock<MockNoteTransportNode>>,
197    fetch_tag_counts: Arc<RwLock<Vec<usize>>>,
198    fetch_bounds: Arc<RwLock<Vec<Option<NoteTransportCursor>>>>,
199}
200
201impl MockNoteTransportApi {
202    pub fn new(mock_node: Arc<RwLock<MockNoteTransportNode>>) -> Self {
203        Self {
204            mock_node,
205            fetch_tag_counts: Arc::default(),
206            fetch_bounds: Arc::default(),
207        }
208    }
209
210    /// Returns the number of tags in each fetch request.
211    pub fn fetch_tag_counts(&self) -> Vec<usize> {
212        self.fetch_tag_counts.read().clone()
213    }
214
215    /// Returns the `sync_to` bound of each fetch request.
216    pub fn fetch_bounds(&self) -> Vec<Option<NoteTransportCursor>> {
217        self.fetch_bounds.read().clone()
218    }
219}
220
221impl MockNoteTransportApi {
222    pub fn send_note_with_proof(&self, note: TransportNote, inclusion_proof: &NoteInclusionProof) {
223        let (header, details) = note.into_parts();
224        let details_bytes = details.to_bytes();
225        self.mock_node
226            .write()
227            .add_note_with_proof(header, details_bytes, inclusion_proof);
228    }
229}
230
231#[cfg_attr(not(target_arch = "wasm32"), async_trait::async_trait)]
232#[cfg_attr(target_arch = "wasm32", async_trait::async_trait(?Send))]
233impl NoteTransportClient for MockNoteTransportApi {
234    async fn send_note_with_proof(
235        &self,
236        note: TransportNote,
237        inclusion_proof: NoteInclusionProof,
238    ) -> Result<(), NoteTransportError> {
239        self.send_note_with_proof(note, &inclusion_proof);
240        Ok(())
241    }
242
243    async fn fetch_notes_page(
244        &self,
245        tags: &[NoteTag],
246        cursor: Option<NoteTransportCursor>,
247        sync_to: Option<NoteTransportCursor>,
248    ) -> Result<NoteTransportPage, NoteTransportError> {
249        self.fetch_tag_counts.write().push(tags.len());
250        self.fetch_bounds.write().push(sync_to);
251        let node = self.mock_node.read();
252        let sync_to = sync_to.unwrap_or_else(|| node.head());
253        let (notes, cursor) = node.notes_up_to(tags, cursor, sync_to);
254        let has_more = !node.notes_up_to(tags, Some(cursor), sync_to).0.is_empty();
255        Ok(NoteTransportPage { notes, cursor, has_more, sync_to })
256    }
257}
258
259// FAULTY NOTE TRANSPORT API
260// ================================================================================================
261
262/// Test-only [`NoteTransportClient`] decorator that injects controlled failures into
263/// `send_note_with_proof` calls.
264///
265/// Reproduces the failure mode where the NTL is reachable but rejects a send. Tests use it to check
266/// how [`Client::send_private_note_with_proof`](crate::Client::send_private_note_with_proof)
267/// reports the failure and how a later send by the caller delivers the note.
268///
269/// The decorator counts attempts (`send_attempts`) and lets a test specify how many of the next
270/// `send_note_with_proof` calls should fail (`fail_next`); successful calls delegate to an inner
271/// [`MockNoteTransportApi`]. `fetch_notes_page` failures can be injected separately via
272/// [`FaultyNoteTransportApi::fail_next_n_fetches`].
273pub struct FaultyNoteTransportApi {
274    inner: MockNoteTransportApi,
275    fail_next: AtomicUsize,
276    send_attempts: AtomicUsize,
277    fail_next_fetches: AtomicUsize,
278    fail_on_fetch: AtomicUsize,
279    fetch_attempts: AtomicUsize,
280}
281
282impl FaultyNoteTransportApi {
283    /// Create a faulty transport that fails the next `fail_next` `send_note_with_proof` calls
284    /// before delegating to the inner mock.
285    pub fn new(mock_node: Arc<RwLock<MockNoteTransportNode>>, fail_next: usize) -> Self {
286        Self {
287            inner: MockNoteTransportApi::new(mock_node),
288            fail_next: AtomicUsize::new(fail_next),
289            send_attempts: AtomicUsize::new(0),
290            fail_next_fetches: AtomicUsize::new(0),
291            fail_on_fetch: AtomicUsize::new(0),
292            fetch_attempts: AtomicUsize::new(0),
293        }
294    }
295
296    /// Total `send_note_with_proof` calls observed (success + failure).
297    pub fn send_attempts(&self) -> usize {
298        self.send_attempts.load(Ordering::SeqCst)
299    }
300
301    /// Fail the next `n` `fetch_notes_page` calls before delegating to the inner mock again.
302    pub fn fail_next_n_fetches(&self, n: usize) {
303        self.fail_next_fetches.store(n, Ordering::SeqCst);
304    }
305
306    /// Fails one fetch attempt. Attempt numbers start at one.
307    pub fn fail_on_fetch_attempt(&self, attempt: usize) {
308        self.fail_on_fetch.store(attempt, Ordering::SeqCst);
309    }
310
311    /// Total `fetch_notes_page` calls observed (success + failure).
312    pub fn fetch_attempts(&self) -> usize {
313        self.fetch_attempts.load(Ordering::SeqCst)
314    }
315
316    /// Records a send attempt and returns whether it must fail.
317    fn take_send_failure(&self) -> Option<NoteTransportError> {
318        self.send_attempts.fetch_add(1, Ordering::SeqCst);
319        self.fail_next
320            .fetch_update(Ordering::SeqCst, Ordering::SeqCst, |n| n.checked_sub(1))
321            .is_ok()
322            .then(|| {
323                NoteTransportError::Network(
324                    "FaultyNoteTransportApi: simulated send_note_with_proof failure".to_string(),
325                )
326            })
327    }
328}
329
330#[cfg_attr(not(target_arch = "wasm32"), async_trait::async_trait)]
331#[cfg_attr(target_arch = "wasm32", async_trait::async_trait(?Send))]
332impl NoteTransportClient for FaultyNoteTransportApi {
333    async fn send_note_with_proof(
334        &self,
335        note: TransportNote,
336        inclusion_proof: NoteInclusionProof,
337    ) -> Result<(), NoteTransportError> {
338        if let Some(error) = self.take_send_failure() {
339            return Err(error);
340        }
341        self.inner.send_note_with_proof(note, &inclusion_proof);
342        Ok(())
343    }
344
345    async fn fetch_notes_page(
346        &self,
347        tags: &[NoteTag],
348        cursor: Option<NoteTransportCursor>,
349        sync_to: Option<NoteTransportCursor>,
350    ) -> Result<NoteTransportPage, NoteTransportError> {
351        let attempt = self.fetch_attempts.fetch_add(1, Ordering::SeqCst) + 1;
352        let should_fail = self
353            .fail_next_fetches
354            .fetch_update(Ordering::SeqCst, Ordering::SeqCst, |n| n.checked_sub(1))
355            .is_ok();
356        if should_fail || attempt == self.fail_on_fetch.load(Ordering::SeqCst) {
357            return Err(NoteTransportError::Network(
358                "FaultyNoteTransportApi: simulated fetch_notes_page failure".to_string(),
359            ));
360        }
361        self.inner.fetch_notes_page(tags, cursor, sync_to).await
362    }
363}
364
365// SERIALIZATION
366// ================================================================================================
367
368impl Serializable for MockNoteTransportNode {
369    fn write_into<W: ByteWriter>(&self, target: &mut W) {
370        self.notes.write_into(target);
371        self.nonce.write_into(target);
372        self.next_sequence.write_into(target);
373    }
374}
375
376impl Deserializable for MockNoteTransportNode {
377    fn read_from<R: ByteReader>(source: &mut R) -> Result<Self, DeserializationError> {
378        let notes = BTreeMap::<NoteTag, Vec<(NoteInfo, NoteTransportCursor)>>::read_from(source)?;
379        let nonce = u64::read_from(source)?;
380        let next_sequence = u64::read_from(source)?;
381
382        Ok(Self {
383            notes,
384            nonce,
385            next_sequence,
386            max_batch: None,
387            proven_notes: BTreeMap::default(),
388        })
389    }
390}