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(
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 fn contains_note(&self, note_id: &NoteId) -> bool {
87 self.notes.values().flatten().any(|(info, _)| info.header.id() == *note_id)
88 }
89
90 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 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 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 fn head(&self) -> NoteTransportCursor {
144 NoteTransportCursor {
145 nonce: self.nonce,
146 sequence: self.next_sequence - 1,
147 }
148 }
149
150 fn notes_up_to(
152 &self,
153 tags: &[NoteTag],
154 cursor: Option<NoteTransportCursor>,
155 sync_to: NoteTransportCursor,
156 ) -> (Vec<NoteInfo>, NoteTransportCursor) {
157 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 collected.sort_by_key(|(_, c)| c.sequence);
171
172 if let Some(max) = self.max_batch {
174 collected.truncate(max);
175 }
176
177 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#[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 pub fn fetch_tag_counts(&self) -> Vec<usize> {
212 self.fetch_tag_counts.read().clone()
213 }
214
215 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
259pub 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 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 pub fn send_attempts(&self) -> usize {
298 self.send_attempts.load(Ordering::SeqCst)
299 }
300
301 pub fn fail_next_n_fetches(&self, n: usize) {
303 self.fail_next_fetches.store(n, Ordering::SeqCst);
304 }
305
306 pub fn fail_on_fetch_attempt(&self, attempt: usize) {
308 self.fail_on_fetch.store(attempt, Ordering::SeqCst);
309 }
310
311 pub fn fetch_attempts(&self) -> usize {
313 self.fetch_attempts.load(Ordering::SeqCst)
314 }
315
316 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
365impl 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}