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#[derive(Clone)]
32pub struct MockNoteTransportNode {
33 notes: BTreeMap<NoteTag, Vec<(NoteInfo, NoteTransportCursor)>>,
34 nonce: u64,
35 next_sequence: u64,
36 max_batch: Option<usize>,
40 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 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 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 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 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 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 let mut collected: Vec<(NoteInfo, NoteTransportCursor)> = vec![];
131 for tag in tags {
132 let tnotes = self
134 .notes
135 .get(tag)
136 .map(|pg_notes| {
137 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 collected.sort_by_key(|(_, c)| *c);
153
154 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#[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 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
249pub 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 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 pub fn send_attempts(&self) -> usize {
289 self.send_attempts.load(Ordering::SeqCst)
290 }
291
292 pub fn fail_next_n_fetches(&self, n: usize) {
294 self.fail_next_fetches.store(n, Ordering::SeqCst);
295 }
296
297 pub fn fail_on_fetch_attempt(&self, attempt: usize) {
299 self.fail_on_fetch.store(attempt, Ordering::SeqCst);
300 }
301
302 pub fn fetch_attempts(&self) -> usize {
304 self.fetch_attempts.load(Ordering::SeqCst)
305 }
306
307 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
355impl 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}