miden-client 0.18.0

use alloc::boxed::Box;
use alloc::collections::BTreeMap;
use alloc::string::ToString;
use alloc::sync::Arc;
use alloc::vec::Vec;
use core::sync::atomic::{AtomicUsize, Ordering};

use miden_protocol::block::BlockNumber;
use miden_protocol::note::{NoteHeader, NoteId, NoteInclusionProof, NoteTag};
use miden_tx::utils::serde::{
    ByteReader,
    ByteWriter,
    Deserializable,
    DeserializationError,
    Serializable,
};
use miden_tx::utils::sync::RwLock;

use crate::note_transport::{
    NoteInfo,
    NoteTransportClient,
    NoteTransportCursor,
    NoteTransportError,
    NoteTransportPage,
    TransportNote,
};

/// Mock Note Transport Node
///
/// Simulates the functionality of the note transport node.
#[derive(Clone)]
pub struct MockNoteTransportNode {
    notes: BTreeMap<NoteTag, Vec<(NoteInfo, NoteTransportCursor)>>,
    nonce: u64,
    next_sequence: u64,
    /// Optional per-response batch cap; if `Some(n)`, `get_notes` returns at most `n` entries
    /// (total, across all tags) in one call. Used to exercise client-side pagination drain loops.
    /// `None` = unbounded (legacy behavior).
    max_batch: Option<usize>,
    /// Notes stored through the with-proof path, with the block their proof named.
    proven_notes: BTreeMap<NoteId, BlockNumber>,
}

impl MockNoteTransportNode {
    pub fn new() -> Self {
        Self {
            notes: BTreeMap::default(),
            nonce: 1,
            next_sequence: 1,
            max_batch: None,
            proven_notes: BTreeMap::default(),
        }
    }

    /// Build a mock that caps each `get_notes` response at `max_batch` entries.
    pub fn with_max_batch(max_batch: usize) -> Self {
        Self {
            notes: BTreeMap::default(),
            nonce: 1,
            next_sequence: 1,
            max_batch: Some(max_batch),
            proven_notes: BTreeMap::default(),
        }
    }

    /// Seed a note relayed with its inclusion proof. The real service verifies the proof against
    /// its node; the mock only records the proof's block and serves it as the commitment block.
    ///
    /// The real service stores a note only once. The mock also ignores a note with an id that it
    /// already stores.
    pub fn add_note_with_proof(
        &mut self,
        header: NoteHeader,
        details_bytes: Vec<u8>,
        inclusion_proof: &NoteInclusionProof,
    ) {
        if self.contains_note(&header.id()) {
            return;
        }
        let block_num = inclusion_proof.location().block_num();
        self.proven_notes.insert(header.id(), block_num);
        self.add_note_after(header, details_bytes, Some(block_num));
    }

    /// Returns whether the mock stores a note with `note_id`.
    fn contains_note(&self, note_id: &NoteId) -> bool {
        self.notes.values().flatten().any(|(info, _)| info.header.id() == *note_id)
    }

    /// Returns the block named by the proof a note was stored with, or `None` when the note was not
    /// stored through the with-proof path.
    pub fn proven_block(&self, note_id: &NoteId) -> Option<BlockNumber> {
        self.proven_notes.get(note_id).copied()
    }

    pub fn add_note(&mut self, header: NoteHeader, details_bytes: Vec<u8>) {
        self.add_note_after(header, details_bytes, None);
    }

    /// Seed a note that carries a commitment block floor, as a transport that stored the note
    /// without verifying a proof would serve it.
    pub fn add_note_after(
        &mut self,
        header: NoteHeader,
        details_bytes: Vec<u8>,
        block_hint: Option<BlockNumber>,
    ) {
        let tag = header.metadata().tag();
        let info = NoteInfo { header, details_bytes, block_hint };
        let cursor = NoteTransportCursor {
            nonce: self.nonce,
            sequence: self.next_sequence,
        };
        self.next_sequence += 1;
        self.notes.entry(tag).or_default().push((info, cursor));
    }

    /// Seed a note under an arbitrary transport tag key, regardless of the note's own tag.
    pub fn add_note_with_tag_key(
        &mut self,
        tag: NoteTag,
        header: NoteHeader,
        details_bytes: Vec<u8>,
    ) {
        let info = NoteInfo { header, details_bytes, block_hint: None };
        let cursor = NoteTransportCursor {
            nonce: self.nonce,
            sequence: self.next_sequence,
        };
        self.next_sequence += 1;
        self.notes.entry(tag).or_default().push((info, cursor));
    }

    pub fn get_notes(
        &self,
        tags: &[NoteTag],
        cursor: Option<NoteTransportCursor>,
    ) -> (Vec<NoteInfo>, NoteTransportCursor) {
        self.notes_up_to(tags, cursor, self.head())
    }

    /// Returns the position of the last stored note.
    fn head(&self) -> NoteTransportCursor {
        NoteTransportCursor {
            nonce: self.nonce,
            sequence: self.next_sequence - 1,
        }
    }

    /// Returns the notes after `cursor` and up to `sync_to`.
    fn notes_up_to(
        &self,
        tags: &[NoteTag],
        cursor: Option<NoteTransportCursor>,
        sync_to: NoteTransportCursor,
    ) -> (Vec<NoteInfo>, NoteTransportCursor) {
        // A cursor from another database starts from the first retained note.
        let after = cursor.filter(|cursor| cursor.nonce == self.nonce).map_or(0, |c| c.sequence);
        let mut collected: Vec<(NoteInfo, NoteTransportCursor)> = tags
            .iter()
            .filter_map(|tag| self.notes.get(tag))
            .flatten()
            .filter(|(_, tcursor)| tcursor.sequence > after && tcursor.sequence <= sync_to.sequence)
            .cloned()
            .collect();

        // Deterministic ordering across tags: sort by cursor ascending so the client sees notes in
        // per-cursor order regardless of tag iteration order, matching the real server's `ORDER BY
        // seq ASC`.
        collected.sort_by_key(|(_, c)| c.sequence);

        // Apply the batch cap, if configured.
        if let Some(max) = self.max_batch {
            collected.truncate(max);
        }

        // Matches the real server's contract (`rcursor = max(cursor, max_seq_returned)`), so an
        // empty batch returns the caller's own position.
        let sequence = collected.last().map_or(after, |(_, c)| c.sequence);
        let notes = collected.into_iter().map(|(n, _)| n).collect();
        (notes, NoteTransportCursor { nonce: self.nonce, sequence })
    }
}

impl Default for MockNoteTransportNode {
    fn default() -> Self {
        Self::new()
    }
}

/// Mock Note Transport API
///
/// Simulates communications with the note transport node.
#[derive(Clone, Default)]
pub struct MockNoteTransportApi {
    pub mock_node: Arc<RwLock<MockNoteTransportNode>>,
    fetch_tag_counts: Arc<RwLock<Vec<usize>>>,
    fetch_bounds: Arc<RwLock<Vec<Option<NoteTransportCursor>>>>,
}

impl MockNoteTransportApi {
    pub fn new(mock_node: Arc<RwLock<MockNoteTransportNode>>) -> Self {
        Self {
            mock_node,
            fetch_tag_counts: Arc::default(),
            fetch_bounds: Arc::default(),
        }
    }

    /// Returns the number of tags in each fetch request.
    pub fn fetch_tag_counts(&self) -> Vec<usize> {
        self.fetch_tag_counts.read().clone()
    }

    /// Returns the `sync_to` bound of each fetch request.
    pub fn fetch_bounds(&self) -> Vec<Option<NoteTransportCursor>> {
        self.fetch_bounds.read().clone()
    }
}

impl MockNoteTransportApi {
    pub fn send_note_with_proof(&self, note: TransportNote, inclusion_proof: &NoteInclusionProof) {
        let (header, details) = note.into_parts();
        let details_bytes = details.to_bytes();
        self.mock_node
            .write()
            .add_note_with_proof(header, details_bytes, inclusion_proof);
    }
}

#[cfg_attr(not(target_arch = "wasm32"), async_trait::async_trait)]
#[cfg_attr(target_arch = "wasm32", async_trait::async_trait(?Send))]
impl NoteTransportClient for MockNoteTransportApi {
    async fn send_note_with_proof(
        &self,
        note: TransportNote,
        inclusion_proof: NoteInclusionProof,
    ) -> Result<(), NoteTransportError> {
        self.send_note_with_proof(note, &inclusion_proof);
        Ok(())
    }

    async fn fetch_notes_page(
        &self,
        tags: &[NoteTag],
        cursor: Option<NoteTransportCursor>,
        sync_to: Option<NoteTransportCursor>,
    ) -> Result<NoteTransportPage, NoteTransportError> {
        self.fetch_tag_counts.write().push(tags.len());
        self.fetch_bounds.write().push(sync_to);
        let node = self.mock_node.read();
        let sync_to = sync_to.unwrap_or_else(|| node.head());
        let (notes, cursor) = node.notes_up_to(tags, cursor, sync_to);
        let has_more = !node.notes_up_to(tags, Some(cursor), sync_to).0.is_empty();
        Ok(NoteTransportPage { notes, cursor, has_more, sync_to })
    }
}

// FAULTY NOTE TRANSPORT API
// ================================================================================================

/// Test-only [`NoteTransportClient`] decorator that injects controlled failures into
/// `send_note_with_proof` calls.
///
/// Reproduces the failure mode where the NTL is reachable but rejects a send. Tests use it to check
/// how [`Client::send_private_note_with_proof`](crate::Client::send_private_note_with_proof)
/// reports the failure and how a later send by the caller delivers the note.
///
/// The decorator counts attempts (`send_attempts`) and lets a test specify how many of the next
/// `send_note_with_proof` calls should fail (`fail_next`); successful calls delegate to an inner
/// [`MockNoteTransportApi`]. `fetch_notes_page` failures can be injected separately via
/// [`FaultyNoteTransportApi::fail_next_n_fetches`].
pub struct FaultyNoteTransportApi {
    inner: MockNoteTransportApi,
    fail_next: AtomicUsize,
    send_attempts: AtomicUsize,
    fail_next_fetches: AtomicUsize,
    fail_on_fetch: AtomicUsize,
    fetch_attempts: AtomicUsize,
}

impl FaultyNoteTransportApi {
    /// Create a faulty transport that fails the next `fail_next` `send_note_with_proof` calls
    /// before delegating to the inner mock.
    pub fn new(mock_node: Arc<RwLock<MockNoteTransportNode>>, fail_next: usize) -> Self {
        Self {
            inner: MockNoteTransportApi::new(mock_node),
            fail_next: AtomicUsize::new(fail_next),
            send_attempts: AtomicUsize::new(0),
            fail_next_fetches: AtomicUsize::new(0),
            fail_on_fetch: AtomicUsize::new(0),
            fetch_attempts: AtomicUsize::new(0),
        }
    }

    /// Total `send_note_with_proof` calls observed (success + failure).
    pub fn send_attempts(&self) -> usize {
        self.send_attempts.load(Ordering::SeqCst)
    }

    /// Fail the next `n` `fetch_notes_page` calls before delegating to the inner mock again.
    pub fn fail_next_n_fetches(&self, n: usize) {
        self.fail_next_fetches.store(n, Ordering::SeqCst);
    }

    /// Fails one fetch attempt. Attempt numbers start at one.
    pub fn fail_on_fetch_attempt(&self, attempt: usize) {
        self.fail_on_fetch.store(attempt, Ordering::SeqCst);
    }

    /// Total `fetch_notes_page` calls observed (success + failure).
    pub fn fetch_attempts(&self) -> usize {
        self.fetch_attempts.load(Ordering::SeqCst)
    }

    /// Records a send attempt and returns whether it must fail.
    fn take_send_failure(&self) -> Option<NoteTransportError> {
        self.send_attempts.fetch_add(1, Ordering::SeqCst);
        self.fail_next
            .fetch_update(Ordering::SeqCst, Ordering::SeqCst, |n| n.checked_sub(1))
            .is_ok()
            .then(|| {
                NoteTransportError::Network(
                    "FaultyNoteTransportApi: simulated send_note_with_proof failure".to_string(),
                )
            })
    }
}

#[cfg_attr(not(target_arch = "wasm32"), async_trait::async_trait)]
#[cfg_attr(target_arch = "wasm32", async_trait::async_trait(?Send))]
impl NoteTransportClient for FaultyNoteTransportApi {
    async fn send_note_with_proof(
        &self,
        note: TransportNote,
        inclusion_proof: NoteInclusionProof,
    ) -> Result<(), NoteTransportError> {
        if let Some(error) = self.take_send_failure() {
            return Err(error);
        }
        self.inner.send_note_with_proof(note, &inclusion_proof);
        Ok(())
    }

    async fn fetch_notes_page(
        &self,
        tags: &[NoteTag],
        cursor: Option<NoteTransportCursor>,
        sync_to: Option<NoteTransportCursor>,
    ) -> Result<NoteTransportPage, NoteTransportError> {
        let attempt = self.fetch_attempts.fetch_add(1, Ordering::SeqCst) + 1;
        let should_fail = self
            .fail_next_fetches
            .fetch_update(Ordering::SeqCst, Ordering::SeqCst, |n| n.checked_sub(1))
            .is_ok();
        if should_fail || attempt == self.fail_on_fetch.load(Ordering::SeqCst) {
            return Err(NoteTransportError::Network(
                "FaultyNoteTransportApi: simulated fetch_notes_page failure".to_string(),
            ));
        }
        self.inner.fetch_notes_page(tags, cursor, sync_to).await
    }
}

// SERIALIZATION
// ================================================================================================

impl Serializable for MockNoteTransportNode {
    fn write_into<W: ByteWriter>(&self, target: &mut W) {
        self.notes.write_into(target);
        self.nonce.write_into(target);
        self.next_sequence.write_into(target);
    }
}

impl Deserializable for MockNoteTransportNode {
    fn read_from<R: ByteReader>(source: &mut R) -> Result<Self, DeserializationError> {
        let notes = BTreeMap::<NoteTag, Vec<(NoteInfo, NoteTransportCursor)>>::read_from(source)?;
        let nonce = u64::read_from(source)?;
        let next_sequence = u64::read_from(source)?;

        Ok(Self {
            notes,
            nonce,
            next_sequence,
            max_batch: None,
            proven_notes: BTreeMap::default(),
        })
    }
}