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    pub fn add_note_with_proof(
69        &mut self,
70        header: NoteHeader,
71        details_bytes: Vec<u8>,
72        inclusion_proof: &NoteInclusionProof,
73    ) {
74        let block_num = inclusion_proof.location().block_num();
75        self.proven_notes.insert(header.id(), block_num);
76        self.add_note_after(header, details_bytes, Some(block_num));
77    }
78
79    /// Returns the block named by the proof a note was stored with, or `None` when the note was not
80    /// stored through the with-proof path.
81    pub fn proven_block(&self, note_id: &NoteId) -> Option<BlockNumber> {
82        self.proven_notes.get(note_id).copied()
83    }
84
85    pub fn add_note(&mut self, header: NoteHeader, details_bytes: Vec<u8>) {
86        self.add_note_after(header, details_bytes, None);
87    }
88
89    /// Seed a note that carries a commitment block floor, as a transport that stored the note
90    /// without verifying a proof would serve it.
91    pub fn add_note_after(
92        &mut self,
93        header: NoteHeader,
94        details_bytes: Vec<u8>,
95        block_hint: Option<BlockNumber>,
96    ) {
97        let tag = header.metadata().tag();
98        let info = NoteInfo { header, details_bytes, block_hint };
99        let cursor = NoteTransportCursor::from_parts(self.nonce, self.next_sequence);
100        self.next_sequence += 1;
101        self.notes.entry(tag).or_default().push((info, cursor));
102    }
103
104    /// Seed a note under an arbitrary transport tag key, regardless of the note's own tag.
105    pub fn add_note_with_tag_key(
106        &mut self,
107        tag: NoteTag,
108        header: NoteHeader,
109        details_bytes: Vec<u8>,
110    ) {
111        let info = NoteInfo { header, details_bytes, block_hint: None };
112        let cursor = NoteTransportCursor::from_parts(self.nonce, self.next_sequence);
113        self.next_sequence += 1;
114        self.notes.entry(tag).or_default().push((info, cursor));
115    }
116
117    pub fn get_notes(
118        &self,
119        tags: &[NoteTag],
120        cursor: NoteTransportCursor,
121    ) -> (Vec<NoteInfo>, NoteTransportCursor) {
122        let cursor = if cursor.parts().is_some_and(|(nonce, _)| nonce != self.nonce) {
123            NoteTransportCursor::init()
124        } else {
125            cursor
126        };
127        // Start `rcursor` at the input — matches the real server's contract (`rcursor = max(cursor,
128        // max_seq_returned)`), so an empty batch returns the caller's own cursor rather than
129        // `init()`.
130        let mut collected: Vec<(NoteInfo, NoteTransportCursor)> = vec![];
131        for tag in tags {
132            // Assumes stored notes are ordered by cursor
133            let tnotes = self
134                .notes
135                .get(tag)
136                .map(|pg_notes| {
137                    // Find first element after cursor
138                    if let Some(pos) = pg_notes.iter().position(|(_, tcursor)| *tcursor > cursor) {
139                        &pg_notes[pos..]
140                    } else {
141                        &[]
142                    }
143                })
144                .map(Vec::from)
145                .unwrap_or_default();
146            collected.extend(tnotes);
147        }
148
149        // Deterministic ordering across tags: sort by cursor ascending so the client sees notes in
150        // per-cursor order regardless of tag iteration order, matching the real server's `ORDER BY
151        // seq ASC`.
152        collected.sort_by_key(|(_, c)| *c);
153
154        // Apply the batch cap, if configured.
155        if let Some(max) = self.max_batch {
156            collected.truncate(max);
157        }
158
159        let initial_cursor = cursor
160            .parts()
161            .map_or(NoteTransportCursor::from_parts(self.nonce, 0), |_| cursor);
162        let rcursor = collected.iter().map(|(_, c)| *c).max().unwrap_or(initial_cursor);
163        let notes = collected.into_iter().map(|(n, _)| n).collect();
164        (notes, rcursor)
165    }
166}
167
168impl Default for MockNoteTransportNode {
169    fn default() -> Self {
170        Self::new()
171    }
172}
173
174/// Mock Note Transport API
175///
176/// Simulates communications with the note transport node.
177#[derive(Clone, Default)]
178pub struct MockNoteTransportApi {
179    pub mock_node: Arc<RwLock<MockNoteTransportNode>>,
180    fetch_tag_counts: Arc<RwLock<Vec<usize>>>,
181}
182
183impl MockNoteTransportApi {
184    pub fn new(mock_node: Arc<RwLock<MockNoteTransportNode>>) -> Self {
185        Self {
186            mock_node,
187            fetch_tag_counts: Arc::default(),
188        }
189    }
190
191    /// Returns the number of tags in each fetch request.
192    pub fn fetch_tag_counts(&self) -> Vec<usize> {
193        self.fetch_tag_counts.read().clone()
194    }
195}
196
197impl MockNoteTransportApi {
198    pub fn send_note_with_proof(&self, note: TransportNote, inclusion_proof: &NoteInclusionProof) {
199        let (header, details) = note.into_parts();
200        let details_bytes = details.to_bytes();
201        self.mock_node
202            .write()
203            .add_note_with_proof(header, details_bytes, inclusion_proof);
204    }
205
206    pub fn fetch_notes(
207        &self,
208        tags: &[NoteTag],
209        cursor: NoteTransportCursor,
210    ) -> (Vec<NoteInfo>, NoteTransportCursor) {
211        self.mock_node.read().get_notes(tags, cursor)
212    }
213}
214
215#[cfg_attr(not(target_arch = "wasm32"), async_trait::async_trait)]
216#[cfg_attr(target_arch = "wasm32", async_trait::async_trait(?Send))]
217impl NoteTransportClient for MockNoteTransportApi {
218    async fn send_note_with_proof(
219        &self,
220        note: TransportNote,
221        inclusion_proof: NoteInclusionProof,
222    ) -> Result<(), NoteTransportError> {
223        self.send_note_with_proof(note, &inclusion_proof);
224        Ok(())
225    }
226
227    async fn fetch_notes(
228        &self,
229        tags: &[NoteTag],
230        cursor: NoteTransportCursor,
231    ) -> Result<(Vec<NoteInfo>, NoteTransportCursor), NoteTransportError> {
232        let page = self.fetch_notes_page(tags, cursor).await?;
233        Ok((page.notes, page.cursor))
234    }
235
236    async fn fetch_notes_page(
237        &self,
238        tags: &[NoteTag],
239        cursor: NoteTransportCursor,
240    ) -> Result<NoteTransportPage, NoteTransportError> {
241        self.fetch_tag_counts.write().push(tags.len());
242        let node = self.mock_node.read();
243        let (notes, cursor) = node.get_notes(tags, cursor);
244        let has_more = !node.get_notes(tags, cursor).0.is_empty();
245        Ok(NoteTransportPage { notes, cursor, has_more })
246    }
247}
248
249// FAULTY NOTE TRANSPORT API
250// ================================================================================================
251
252/// Test-only [`NoteTransportClient`] decorator that injects controlled failures into
253/// `send_note_with_proof` calls.
254///
255/// Reproduces the failure mode where the NTL is reachable but rejects (or silently drops) a relay
256/// attempt, exercising the durable outbox in
257/// [`Client::send_private_note_with_proof`](crate::Client::send_private_note_with_proof): without
258/// retry/persistence a failed relay would leave the recipient unable to discover the note.
259///
260/// The decorator counts attempts (`send_attempts`) and lets a test specify how many of the next
261/// `send_note_with_proof` calls should fail (`fail_next`); successful calls delegate to an inner
262/// [`MockNoteTransportApi`]. `fetch_notes` failures can be injected separately via
263/// [`FaultyNoteTransportApi::fail_next_n_fetches`].
264pub struct FaultyNoteTransportApi {
265    inner: MockNoteTransportApi,
266    fail_next: AtomicUsize,
267    send_attempts: AtomicUsize,
268    fail_next_fetches: AtomicUsize,
269    fail_on_fetch: AtomicUsize,
270    fetch_attempts: AtomicUsize,
271}
272
273impl FaultyNoteTransportApi {
274    /// Create a faulty transport that fails the next `fail_next` `send_note_with_proof` calls
275    /// before delegating to the inner mock.
276    pub fn new(mock_node: Arc<RwLock<MockNoteTransportNode>>, fail_next: usize) -> Self {
277        Self {
278            inner: MockNoteTransportApi::new(mock_node),
279            fail_next: AtomicUsize::new(fail_next),
280            send_attempts: AtomicUsize::new(0),
281            fail_next_fetches: AtomicUsize::new(0),
282            fail_on_fetch: AtomicUsize::new(0),
283            fetch_attempts: AtomicUsize::new(0),
284        }
285    }
286
287    /// Total `send_note_with_proof` calls observed (success + failure).
288    pub fn send_attempts(&self) -> usize {
289        self.send_attempts.load(Ordering::SeqCst)
290    }
291
292    /// Fail the next `n` `fetch_notes` calls before delegating to the inner mock again.
293    pub fn fail_next_n_fetches(&self, n: usize) {
294        self.fail_next_fetches.store(n, Ordering::SeqCst);
295    }
296
297    /// Fails one fetch attempt. Attempt numbers start at one.
298    pub fn fail_on_fetch_attempt(&self, attempt: usize) {
299        self.fail_on_fetch.store(attempt, Ordering::SeqCst);
300    }
301
302    /// Total `fetch_notes` calls observed (success + failure).
303    pub fn fetch_attempts(&self) -> usize {
304        self.fetch_attempts.load(Ordering::SeqCst)
305    }
306
307    /// Records a send attempt and returns whether it must fail.
308    fn take_send_failure(&self) -> Option<NoteTransportError> {
309        self.send_attempts.fetch_add(1, Ordering::SeqCst);
310        self.fail_next
311            .fetch_update(Ordering::SeqCst, Ordering::SeqCst, |n| n.checked_sub(1))
312            .is_ok()
313            .then(|| {
314                NoteTransportError::Network(
315                    "FaultyNoteTransportApi: simulated send_note_with_proof failure".to_string(),
316                )
317            })
318    }
319}
320
321#[cfg_attr(not(target_arch = "wasm32"), async_trait::async_trait)]
322#[cfg_attr(target_arch = "wasm32", async_trait::async_trait(?Send))]
323impl NoteTransportClient for FaultyNoteTransportApi {
324    async fn send_note_with_proof(
325        &self,
326        note: TransportNote,
327        inclusion_proof: NoteInclusionProof,
328    ) -> Result<(), NoteTransportError> {
329        if let Some(error) = self.take_send_failure() {
330            return Err(error);
331        }
332        self.inner.send_note_with_proof(note, &inclusion_proof);
333        Ok(())
334    }
335
336    async fn fetch_notes(
337        &self,
338        tags: &[NoteTag],
339        cursor: NoteTransportCursor,
340    ) -> Result<(Vec<NoteInfo>, NoteTransportCursor), NoteTransportError> {
341        let attempt = self.fetch_attempts.fetch_add(1, Ordering::SeqCst) + 1;
342        let should_fail = self
343            .fail_next_fetches
344            .fetch_update(Ordering::SeqCst, Ordering::SeqCst, |n| n.checked_sub(1))
345            .is_ok();
346        if should_fail || attempt == self.fail_on_fetch.load(Ordering::SeqCst) {
347            return Err(NoteTransportError::Network(
348                "FaultyNoteTransportApi: simulated fetch_notes failure".to_string(),
349            ));
350        }
351        Ok(self.inner.fetch_notes(tags, cursor))
352    }
353}
354
355// SERIALIZATION
356// ================================================================================================
357
358impl Serializable for MockNoteTransportNode {
359    fn write_into<W: ByteWriter>(&self, target: &mut W) {
360        self.notes.write_into(target);
361        self.nonce.write_into(target);
362        self.next_sequence.write_into(target);
363    }
364}
365
366impl Deserializable for MockNoteTransportNode {
367    fn read_from<R: ByteReader>(source: &mut R) -> Result<Self, DeserializationError> {
368        let notes = BTreeMap::<NoteTag, Vec<(NoteInfo, NoteTransportCursor)>>::read_from(source)?;
369        let nonce = u64::read_from(source)?;
370        let next_sequence = u64::read_from(source)?;
371
372        Ok(Self {
373            notes,
374            nonce,
375            next_sequence,
376            max_batch: None,
377            proven_notes: BTreeMap::default(),
378        })
379    }
380}