ant_node/

node.rs

// Copyright 2024 MaidSafe.net limited.
//
// This SAFE Network Software is licensed to you under The General Public License (GPL), version 3.
// Unless required by applicable law or agreed to in writing, the SAFE Network Software distributed
// under the GPL Licence is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
// KIND, either express or implied. Please review the Licences for the specific language governing
// permissions and limitations relating to use of the SAFE Network Software.

use super::{
    Marker, NodeEvent, error::Result, event::NodeEventsChannel, quote::quotes_verification,
};
#[cfg(feature = "open-metrics")]
use crate::metrics::NodeMetricsRecorder;
#[cfg(feature = "open-metrics")]
use crate::networking::MetricsRegistries;
use crate::networking::{Addresses, Network, NetworkConfig, NetworkError, NetworkEvent, NodeIssue};
use crate::{PutValidationError, RunningNode};
use ant_bootstrap::bootstrap::Bootstrap;
use ant_evm::EvmNetwork;
use ant_evm::RewardsAddress;
use ant_evm::merkle_payments::MERKLE_PAYMENT_EXPIRATION;
use ant_protocol::{
    CLOSE_GROUP_SIZE, NetworkAddress, PrettyPrintRecordKey,
    error::Error as ProtocolError,
    messages::{ChunkProof, CmdResponse, Nonce, Query, QueryResponse, Request, Response},
    storage::{try_deserialize_record, Chunk, ValidationType},
};
use bytes::Bytes;
use futures::stream::{self, StreamExt};
use itertools::Itertools;
use libp2p::{
    Multiaddr, PeerId,
    identity::Keypair,
    kad::{KBucketDistance as Distance, Record, U256},
    request_response::OutboundFailure,
};
use num_traits::cast::ToPrimitive;
use rand::{
    Rng, SeedableRng,
    rngs::{OsRng, StdRng},
    seq::SliceRandom,
    thread_rng,
};
use std::{
    collections::{BTreeSet, HashMap, HashSet},
    net::SocketAddr,
    path::PathBuf,
    sync::{
        Arc,
        atomic::{AtomicUsize, Ordering},
    },
    time::{Duration, Instant},
};
use tokio::sync::{Mutex, watch};
use tokio::{
    sync::mpsc::Receiver,
    task::{JoinSet, spawn},
};

/// Interval to trigger replication of all records to all peers.
/// This is the max time it should take. Minimum interval at any node will be half this
pub const PERIODIC_REPLICATION_INTERVAL_MAX_S: u64 = 180;

/// Interval to trigger storage challenge.
/// This is the max time it should take. Minimum interval at any node will be half this
const STORE_CHALLENGE_INTERVAL_MAX_S: u64 = 7200;

/// Interval to update the nodes uptime metric
const UPTIME_METRICS_UPDATE_INTERVAL: Duration = Duration::from_secs(10);

/// Interval to clean up unrelevant records
/// This is the max time it should take. Minimum interval at any node will be half this
const UNRELEVANT_RECORDS_CLEANUP_INTERVAL_MAX_S: u64 = 7200;

/// Highest score to achieve from each metric sub-sector during StorageChallenge.
const HIGHEST_SCORE: usize = 100;

/// Any nodes bearing a score below this shall be considered as bad.
/// Max is to be 100 * 100
const MIN_ACCEPTABLE_HEALTHY_SCORE: usize = 3000;

/// in ms, expecting average StorageChallenge complete time to be around 250ms.
const TIME_STEP: usize = 20;

/// Delay before retrying a failed replicated record fetch.
const REPLICA_FETCH_RETRY_DELAY: Duration = Duration::from_secs(15 * 60);

/// Number of peers to query for a replicated record health check.
const REPLICA_FETCH_PEER_COUNT: usize = 5;

/// Minimum number of successful replicas required to consider the record healthy.
const MIN_HEALTHY_REPLICA_COUNT: usize = 3;

/// Index (0-based, including self) used to derive the distance threshold for nearby records.
const CLOSE_NEIGHBOUR_DISTANCE_INDEX: usize = 7;

/// Maximum number of closest peers to track for replication behaviour detection.
///
/// Set to 20 (roughly 4x CLOSE_GROUP_SIZE) to cover peers that are close enough
/// to affect data responsibility during churn. Tracking only the closest peers
/// keeps memory bounded while still detecting restart patterns in the critical
/// zone where replication matters most.
const CLOSE_GROUP_TRACKING_LIMIT: usize = 20;

/// Grace period before triggering replication after detecting a peer restart.
///
/// When a peer is detected as restarting (removed then quickly re-added),
/// replication is suppressed for this duration to allow the peer to stabilize.
/// Set to 90 seconds which is:
/// - Long enough to cover typical node restart times (30-60s)
/// - Short enough to ensure data availability isn't compromised
/// - Approximately half of PERIODIC_REPLICATION_INTERVAL_MAX_S, so the periodic
///   replication will catch any missed updates within a reasonable timeframe
const CLOSE_GROUP_RESTART_SUPPRESSION: Duration = Duration::from_secs(90);

/// Helper to build and run a Node
pub struct NodeBuilder {
    addr: SocketAddr,
    bootstrap: Bootstrap,
    evm_address: RewardsAddress,
    evm_network: EvmNetwork,
    identity_keypair: Keypair,
    local: bool,
    #[cfg(feature = "open-metrics")]
    /// Set to Some to enable the metrics server
    metrics_server_port: Option<u16>,
    no_upnp: bool,
    relay_client: bool,
    root_dir: PathBuf,
}

impl NodeBuilder {
    /// Instantiate the builder. The initial peers can either be supplied via the `initial_peers` method
    /// or fetched from the bootstrap cache set using `bootstrap_cache` method.
    pub fn new(
        identity_keypair: Keypair,
        bootstrap_flow: Bootstrap,
        evm_address: RewardsAddress,
        evm_network: EvmNetwork,
        addr: SocketAddr,
        root_dir: PathBuf,
    ) -> Self {
        Self {
            addr,
            bootstrap: bootstrap_flow,
            evm_address,
            evm_network,
            identity_keypair,
            local: false,
            #[cfg(feature = "open-metrics")]
            metrics_server_port: None,
            no_upnp: false,
            relay_client: false,
            root_dir,
        }
    }

    /// Set the flag to indicate if the node is running in local mode
    pub fn local(&mut self, local: bool) {
        self.local = local;
    }

    #[cfg(feature = "open-metrics")]
    /// Set the port for the OpenMetrics server. Defaults to a random port if not set
    pub fn metrics_server_port(&mut self, port: Option<u16>) {
        self.metrics_server_port = port;
    }

    /// Set the flag to make the node act as a relay client
    pub fn relay_client(&mut self, relay_client: bool) {
        self.relay_client = relay_client;
    }

    /// Set the flag to disable UPnP for the node
    pub fn no_upnp(&mut self, no_upnp: bool) {
        self.no_upnp = no_upnp;
    }

    /// Asynchronously runs a new node instance, setting up the swarm driver,
    /// creating a data storage, and handling network events. Returns the
    /// created `RunningNode` which contains a `NodeEventsChannel` for listening
    /// to node-related events.
    ///
    /// # Returns
    ///
    /// A `RunningNode` instance.
    ///
    /// # Errors
    ///
    /// Returns an error if there is a problem initializing the Network.
    pub fn build_and_run(self) -> Result<RunningNode> {
        // setup metrics
        #[cfg(feature = "open-metrics")]
        let (metrics_recorder, metrics_registries) = if self.metrics_server_port.is_some() {
            // metadata registry
            let mut metrics_registries = MetricsRegistries::default();
            let metrics_recorder = NodeMetricsRecorder::new(&mut metrics_registries);

            (Some(metrics_recorder), metrics_registries)
        } else {
            (None, MetricsRegistries::default())
        };

        // create a shutdown signal channel
        let (shutdown_tx, shutdown_rx) = watch::channel(false);

        // init network
        let network_config = NetworkConfig {
            keypair: self.identity_keypair,
            local: self.local,
            listen_addr: self.addr,
            root_dir: self.root_dir.clone(),
            shutdown_rx: shutdown_rx.clone(),
            bootstrap: self.bootstrap,
            no_upnp: self.no_upnp,
            relay_client: self.relay_client,
            custom_request_timeout: None,
            #[cfg(feature = "open-metrics")]
            metrics_registries,
            #[cfg(feature = "open-metrics")]
            metrics_server_port: self.metrics_server_port,
        };
        let (network, network_event_receiver) = Network::init(network_config)?;

        // init node
        let node_events_channel = NodeEventsChannel::default();
        let node = NodeInner {
            network: network.clone(),
            events_channel: node_events_channel.clone(),
            close_group_tracker: Mutex::new(CloseGroupTracker::new(network.peer_id())),
            reward_address: self.evm_address,
            #[cfg(feature = "open-metrics")]
            metrics_recorder,
            evm_network: self.evm_network,
        };
        let node = Node {
            inner: Arc::new(node),
        };

        // Run the node
        node.run(network_event_receiver, shutdown_rx);
        let running_node = RunningNode {
            shutdown_sender: shutdown_tx,
            network,
            node_events_channel,
            root_dir_path: self.root_dir,
            rewards_address: self.evm_address,
        };

        Ok(running_node)
    }
}

/// `Node` represents a single node in the distributed network. It handles
/// network events, processes incoming requests, interacts with the data
/// storage, and broadcasts node-related events.
#[derive(Clone)]
pub(crate) struct Node {
    inner: Arc<NodeInner>,
}

/// The actual implementation of the Node. The other is just a wrapper around this, so that we don't expose
/// the Arc from the interface.
struct NodeInner {
    events_channel: NodeEventsChannel,
    network: Network,
    close_group_tracker: Mutex<CloseGroupTracker>,
    #[cfg(feature = "open-metrics")]
    metrics_recorder: Option<NodeMetricsRecorder>,
    reward_address: RewardsAddress,
    evm_network: EvmNetwork,
}

impl Node {
    /// Returns the NodeEventsChannel
    pub(crate) fn events_channel(&self) -> &NodeEventsChannel {
        &self.inner.events_channel
    }

    /// Returns the instance of Network
    pub(crate) fn network(&self) -> &Network {
        &self.inner.network
    }

    fn close_group_tracker(&self) -> &Mutex<CloseGroupTracker> {
        &self.inner.close_group_tracker
    }

    #[cfg(feature = "open-metrics")]
    /// Returns a reference to the NodeMetricsRecorder if the `open-metrics` feature flag is enabled
    /// This is used to record various metrics for the node.
    pub(crate) fn metrics_recorder(&self) -> Option<&NodeMetricsRecorder> {
        self.inner.metrics_recorder.as_ref()
    }

    /// Returns the reward address of the node
    pub(crate) fn reward_address(&self) -> &RewardsAddress {
        &self.inner.reward_address
    }

    pub(crate) fn evm_network(&self) -> &EvmNetwork {
        &self.inner.evm_network
    }

    /// Spawns a task to process for `NetworkEvents`.
    /// Returns both tasks as JoinHandle<()>.
    fn run(
        self,
        mut network_event_receiver: Receiver<NetworkEvent>,
        mut shutdown_rx: watch::Receiver<bool>,
    ) {
        let mut rng = StdRng::from_entropy();

        let peers_connected = Arc::new(AtomicUsize::new(0));

        let _node_task = spawn(async move {
            // use a random activity timeout to ensure that the nodes do not sync when messages
            // are being transmitted.
            let replication_interval: u64 = rng.gen_range(
                PERIODIC_REPLICATION_INTERVAL_MAX_S / 2..PERIODIC_REPLICATION_INTERVAL_MAX_S,
            );
            let replication_interval_time = Duration::from_secs(replication_interval);
            debug!("Replication interval set to {replication_interval_time:?}");

            let mut replication_interval = tokio::time::interval(replication_interval_time);
            let _ = replication_interval.tick().await; // first tick completes immediately

            let mut uptime_metrics_update_interval =
                tokio::time::interval(UPTIME_METRICS_UPDATE_INTERVAL);
            let _ = uptime_metrics_update_interval.tick().await; // first tick completes immediately

            // use a random activity timeout to ensure that the nodes do not sync on work,
            // causing an overall CPU spike.
            let irrelevant_records_cleanup_interval: u64 = rng.gen_range(
                UNRELEVANT_RECORDS_CLEANUP_INTERVAL_MAX_S / 2
                    ..UNRELEVANT_RECORDS_CLEANUP_INTERVAL_MAX_S,
            );
            let irrelevant_records_cleanup_interval_time =
                Duration::from_secs(irrelevant_records_cleanup_interval);
            let mut irrelevant_records_cleanup_interval =
                tokio::time::interval(irrelevant_records_cleanup_interval_time);
            let _ = irrelevant_records_cleanup_interval.tick().await; // first tick completes immediately

            // use a random neighbour storage challenge ticker to ensure
            // neighbours do not carryout challenges at the same time
            let storage_challenge_interval: u64 =
                rng.gen_range(STORE_CHALLENGE_INTERVAL_MAX_S / 2..STORE_CHALLENGE_INTERVAL_MAX_S);
            let storage_challenge_interval_time = Duration::from_secs(storage_challenge_interval);
            debug!("Storage challenge interval set to {storage_challenge_interval_time:?}");

            let mut storage_challenge_interval =
                tokio::time::interval(storage_challenge_interval_time);
            let _ = storage_challenge_interval.tick().await; // first tick completes immediately

            loop {
                let peers_connected = &peers_connected;

                tokio::select! {
                    // Check for a shutdown command.
                    result = shutdown_rx.changed() => {
                        if result.is_ok() && *shutdown_rx.borrow() || result.is_err() {
                            info!("Shutdown signal received or sender dropped. Exiting network events loop.");
                            break;
                        }
                    },
                    net_event = network_event_receiver.recv() => {
                        match net_event {
                            Some(event) => {
                                let start = Instant::now();
                                let event_string = format!("{event:?}");

                                self.handle_network_event(event, peers_connected).await;
                                trace!("Handled non-blocking network event in {:?}: {:?}", start.elapsed(), event_string);

                            }
                            None => {
                                error!("The `NetworkEvent` channel is closed");
                                self.events_channel().broadcast(NodeEvent::ChannelClosed);
                                break;
                            }
                        }
                    }
                    // runs every replication_interval time
                    _ = replication_interval.tick() => {
                        let now = Instant::now();

                        {
                            let mut tracker = self.close_group_tracker().lock().await;
                            if tracker.handle_timer_expiry(now) {
                                trace!("Replication timer expired for tracked peers; periodic run will cover it");
                            }
                        }

                        let start = now;
                        let network = self.network().clone();
                        self.record_metrics(Marker::IntervalReplicationTriggered);

                        let _handle = spawn(async move {
                            Self::try_interval_replication(network);
                            trace!("Periodic replication took {:?}", start.elapsed());
                        });
                    }
                    _ = uptime_metrics_update_interval.tick() => {
                        #[cfg(feature = "open-metrics")]
                        if let Some(metrics_recorder) = self.metrics_recorder() {
                            let _ = metrics_recorder.uptime.set(metrics_recorder.started_instant.elapsed().as_secs() as i64);
                        }
                    }
                    _ = irrelevant_records_cleanup_interval.tick() => {
                        let network = self.network().clone();

                        let _handle = spawn(async move {
                            Self::trigger_irrelevant_record_cleanup(network);
                        });
                    }
                    // runs every storage_challenge_interval time
                    _ = storage_challenge_interval.tick() => {
                        let start = Instant::now();
                        debug!("Periodic storage challenge triggered");
                        let network = self.network().clone();

                        let _handle = spawn(async move {
                            Self::storage_challenge(network).await;
                            trace!("Periodic storage challenge took {:?}", start.elapsed());
                        });
                    }
                }
            }
        });
    }

    /// Calls Marker::log() to insert the marker into the log files.
    /// Also calls NodeMetrics::record() to record the metric if the `open-metrics` feature flag is enabled.
    pub(crate) fn record_metrics(&self, marker: Marker) {
        marker.log();
        #[cfg(feature = "open-metrics")]
        if let Some(metrics_recorder) = self.metrics_recorder() {
            metrics_recorder.record(marker)
        }
    }

    // **** Private helpers *****

    /// Handle a network event.
    /// Spawns a thread for any likely long running tasks
    async fn handle_network_event(&self, event: NetworkEvent, peers_connected: &Arc<AtomicUsize>) {
        let start = Instant::now();
        let event_string = format!("{event:?}");
        let event_header;

        // Reducing non-mandatory logging
        if let NetworkEvent::QueryRequestReceived {
            query: Query::GetVersion { .. },
            ..
        } = event
        {
            trace!("Handling NetworkEvent {event_string}");
        } else {
            debug!("Handling NetworkEvent {event_string}");
        }

        match event {
            NetworkEvent::PeerAdded(peer_id, connected_peers) => {
                event_header = "PeerAdded";
                // increment peers_connected and send ConnectedToNetwork event if have connected to K_VALUE peers
                let _ = peers_connected.fetch_add(1, Ordering::SeqCst);
                if peers_connected.load(Ordering::SeqCst) == CLOSE_GROUP_SIZE {
                    self.events_channel()
                        .broadcast(NodeEvent::ConnectedToNetwork);
                }

                self.record_metrics(Marker::PeersInRoutingTable(connected_peers));
                self.record_metrics(Marker::PeerAddedToRoutingTable(&peer_id));

                // try query peer version
                let network = self.network().clone();
                let _handle = spawn(async move {
                    Self::try_query_peer_version(network, peer_id, Default::default()).await;
                });

                let replication_decision = {
                    let mut tracker = self.close_group_tracker().lock().await;
                    tracker.record_peer_added(peer_id)
                };

                if replication_decision.should_trigger() {
                    let network = self.network().clone();
                    self.record_metrics(Marker::IntervalReplicationTriggered);
                    let _handle = spawn(async move {
                        Self::try_interval_replication(network);
                    });
                } else if matches!(replication_decision, ReplicationDirective::Skip) {
                    trace!(
                        "Replication skipped for {peer_id:?} addition due to restart behaviour tracking"
                    );
                }
            }
            NetworkEvent::PeerRemoved(peer_id, connected_peers) => {
                event_header = "PeerRemoved";
                self.record_metrics(Marker::PeersInRoutingTable(connected_peers));
                self.record_metrics(Marker::PeerRemovedFromRoutingTable(&peer_id));

                let self_id = self.network().peer_id();
                let distance =
                    NetworkAddress::from(self_id).distance(&NetworkAddress::from(peer_id));
                info!(
                    "Node {self_id:?} removed peer from routing table: {peer_id:?}. It has a {:?} distance to us.",
                    distance.ilog2()
                );

                let replication_decision = {
                    let mut tracker = self.close_group_tracker().lock().await;
                    tracker.record_peer_removed(peer_id)
                };

                if replication_decision.should_trigger() {
                    let network = self.network().clone();
                    self.record_metrics(Marker::IntervalReplicationTriggered);
                    let _handle = spawn(async move {
                        Self::try_interval_replication(network);
                    });
                } else if matches!(replication_decision, ReplicationDirective::Skip) {
                    trace!(
                        "Replication skipped for {peer_id:?} removal due to restart behaviour tracking"
                    );
                }
            }
            NetworkEvent::PeerWithUnsupportedProtocol { .. } => {
                event_header = "PeerWithUnsupportedProtocol";
            }
            NetworkEvent::NewListenAddr(_) => {
                event_header = "NewListenAddr";
            }
            NetworkEvent::ResponseReceived { res } => {
                event_header = "ResponseReceived";
                if let Err(err) = self.handle_response(res) {
                    error!("Error while handling NetworkEvent::ResponseReceived {err:?}");
                }
            }
            NetworkEvent::KeysToFetchForReplication(keys) => {
                event_header = "KeysToFetchForReplication";
                self.record_metrics(Marker::fetching_keys_for_replication(&keys));

                if let Err(err) = self.fetch_replication_keys_without_wait(keys) {
                    error!("Error while trying to fetch replicated data {err:?}");
                }
            }
            NetworkEvent::QueryRequestReceived { query, channel } => {
                event_header = "QueryRequestReceived";
                let node = self.clone();
                let payment_address = *self.reward_address();

                let _handle = spawn(async move {
                    let network = node.network().clone();
                    let res = Self::handle_query(node, query, payment_address).await;

                    // Reducing non-mandatory logging
                    if let Response::Query(QueryResponse::GetVersion { .. }) = res {
                        trace!("Sending response {res:?}");
                    } else {
                        debug!("Sending response {res:?}");
                    }

                    network.send_response(res, channel);
                });
            }
            NetworkEvent::UnverifiedRecord(record) => {
                event_header = "UnverifiedRecord";
                // queries can be long running and require validation, so we spawn a task to handle them
                let self_clone = self.clone();
                let _handle = spawn(async move {
                    let key = PrettyPrintRecordKey::from(&record.key).into_owned();
                    match self_clone.validate_and_store_record(record).await {
                        Ok(()) => debug!("UnverifiedRecord {key} has been stored"),
                        Err(err) => {
                            self_clone.record_metrics(Marker::RecordRejected(&key, &err));
                        }
                    }
                });
            }
            NetworkEvent::TerminateNode { reason } => {
                event_header = "TerminateNode";
                error!("Received termination from swarm_driver due to {reason:?}");
                self.events_channel()
                    .broadcast(NodeEvent::TerminateNode(format!("{reason}")));
            }
            NetworkEvent::FailedToFetchHolders(bad_nodes) => {
                event_header = "FailedToFetchHolders";
                let network = self.network().clone();
                let pretty_log: Vec<_> = bad_nodes
                    .iter()
                    .map(|(peer_id, record_key)| {
                        let pretty_key = PrettyPrintRecordKey::from(record_key);
                        (peer_id, pretty_key)
                    })
                    .collect();
                // Note: this log will be checked in CI, and expecting `not appear`.
                //       any change to the keyword `failed to fetch` shall incur
                //       correspondent CI script change as well.
                debug!(
                    "Received notification from replication_fetcher, notifying {pretty_log:?} failed to fetch replication copies from."
                );
                let _handle = spawn(async move {
                    for (peer_id, record_key) in bad_nodes {
                        // Obsoleted fetch request (due to flooded in fresh replicates) could result
                        // in peer to be claimed as bad, as local copy blocks the entry to be cleared.
                        if let Ok(false) = network.is_record_key_present_locally(&record_key).await
                        {
                            error!(
                                "From peer {peer_id:?}, failed to fetch record {:?}",
                                PrettyPrintRecordKey::from(&record_key)
                            );
                            network.record_node_issues(peer_id, NodeIssue::ReplicationFailure);
                        }
                    }
                });
            }
            NetworkEvent::QuoteVerification { quotes } => {
                event_header = "QuoteVerification";
                let network = self.network().clone();

                let _handle = spawn(async move {
                    quotes_verification(&network, quotes).await;
                });
            }
            NetworkEvent::FreshReplicateToFetch { holder, keys } => {
                event_header = "FreshReplicateToFetch";
                self.fresh_replicate_to_fetch(holder, keys);
            }
            NetworkEvent::PeersForVersionQuery(peers) => {
                event_header = "PeersForVersionQuery";
                let network = self.network().clone();
                let _handle = spawn(async move {
                    Self::query_peers_version(network, peers).await;
                });
            }
            NetworkEvent::NetworkWideReplication { keys } => {
                event_header = "NetworkWideReplication";
                self.perform_network_wide_replication(keys);
            }
        }

        trace!(
            "Network handling statistics, Event {event_header:?} handled in {:?} : {event_string:?}",
            start.elapsed()
        );
    }

    // Handle the response that was not awaited at the call site
    fn handle_response(&self, response: Response) -> Result<()> {
        match response {
            Response::Cmd(CmdResponse::Replicate(Ok(()))) => {
                // This should actually have been short-circuted when received
                warn!("Mishandled replicate response, should be handled earlier");
            }
            Response::Query(QueryResponse::GetReplicatedRecord(resp)) => {
                error!(
                    "Response to replication shall be handled by called not by common handler, {resp:?}"
                );
            }
            Response::Cmd(CmdResponse::FreshReplicate(Ok(()))) => {
                // No need to handle
            }
            other => {
                warn!("handle_response not implemented for {other:?}");
            }
        };

        Ok(())
    }

    async fn handle_query(node: Self, query: Query, payment_address: RewardsAddress) -> Response {
        let network = node.network();
        let resp: QueryResponse = match query {
            Query::GetStoreQuote {
                key,
                data_type,
                data_size,
                nonce,
                difficulty,
            } => {
                let record_key = key.to_record_key();
                let self_id = network.peer_id();

                let maybe_quoting_metrics = network
                    .get_local_quoting_metrics(record_key.clone(), data_type, data_size)
                    .await;

                let storage_proofs = if let Some(nonce) = nonce {
                    Self::respond_x_closest_record_proof(
                        network,
                        key.clone(),
                        nonce,
                        difficulty,
                        false,
                    )
                    .await
                } else {
                    vec![]
                };

                match maybe_quoting_metrics {
                    Ok((quoting_metrics, is_already_stored)) => {
                        if is_already_stored {
                            QueryResponse::GetStoreQuote {
                                quote: Err(ProtocolError::RecordExists(
                                    PrettyPrintRecordKey::from(&record_key).into_owned(),
                                )),
                                peer_address: NetworkAddress::from(self_id),
                                storage_proofs,
                            }
                        } else {
                            QueryResponse::GetStoreQuote {
                                quote: Self::create_quote_for_storecost(
                                    network,
                                    &key,
                                    &quoting_metrics,
                                    &payment_address,
                                ),
                                peer_address: NetworkAddress::from(self_id),
                                storage_proofs,
                            }
                        }
                    }
                    Err(err) => {
                        warn!("GetStoreQuote failed for {key:?}: {err}");
                        QueryResponse::GetStoreQuote {
                            quote: Err(ProtocolError::GetStoreQuoteFailed),
                            peer_address: NetworkAddress::from(self_id),
                            storage_proofs,
                        }
                    }
                }
            }
            Query::GetReplicatedRecord { requester: _, key } => {
                let our_address = NetworkAddress::from(network.peer_id());
                let record_key = key.to_record_key();

                let result = match network.get_local_record(&record_key).await {
                    Ok(Some(record)) => Ok((our_address, Bytes::from(record.value))),
                    Ok(None) => Err(ProtocolError::ReplicatedRecordNotFound {
                        holder: Box::new(our_address),
                        key: Box::new(key.clone()),
                    }),
                    // Use `PutRecordFailed` as place holder
                    Err(err) => Err(ProtocolError::PutRecordFailed(format!(
                        "Error to fetch local record for GetReplicatedRecord {err:?}"
                    ))),
                };

                QueryResponse::GetReplicatedRecord(result)
            }
            Query::GetChunkExistenceProof {
                key,
                nonce,
                difficulty,
            } => QueryResponse::GetChunkExistenceProof(
                Self::respond_x_closest_record_proof(network, key, nonce, difficulty, true).await,
            ),
            Query::CheckNodeInProblem(target_address) => {
                debug!("Got CheckNodeInProblem for peer {target_address:?}");

                let is_in_trouble =
                    if let Ok(result) = network.is_peer_shunned(target_address.clone()).await {
                        result
                    } else {
                        debug!("Could not get status of {target_address:?}.");
                        false
                    };

                QueryResponse::CheckNodeInProblem {
                    reporter_address: NetworkAddress::from(network.peer_id()),
                    target_address,
                    is_in_trouble,
                }
            }
            Query::GetClosestPeers {
                key,
                num_of_peers,
                range,
                sign_result,
            } => {
                debug!(
                    "Got GetClosestPeers targeting {key:?} with {num_of_peers:?} peers or {range:?} range, signature {sign_result} required."
                );
                Self::respond_get_closest_peers(network, key, num_of_peers, range, sign_result)
                    .await
            }
            Query::GetVersion(_) => QueryResponse::GetVersion {
                peer: NetworkAddress::from(network.peer_id()),
                version: ant_build_info::package_version(),
            },
            Query::PutRecord {
                holder,
                address,
                serialized_record,
            } => {
                let record = Record {
                    key: address.to_record_key(),
                    value: serialized_record,
                    publisher: None,
                    expires: None,
                };

                let key = PrettyPrintRecordKey::from(&record.key).into_owned();
                let result = match node.validate_and_store_record(record).await {
                    Ok(()) => Ok(()),
                    Err(PutValidationError::OutdatedRecordCounter { counter, expected }) => {
                        node.record_metrics(Marker::RecordRejected(
                            &key,
                            &PutValidationError::OutdatedRecordCounter { counter, expected },
                        ));
                        Err(ProtocolError::OutdatedRecordCounter { counter, expected })
                    }
                    Err(PutValidationError::TopologyVerificationFailed {
                        target_address,
                        valid_count,
                        total_paid,
                        closest_count,
                        node_peers,
                        paid_peers,
                    }) => {
                        node.record_metrics(Marker::RecordRejected(
                            &key,
                            &PutValidationError::TopologyVerificationFailed {
                                target_address: target_address.clone(),
                                valid_count,
                                total_paid,
                                closest_count,
                                node_peers: node_peers.clone(),
                                paid_peers: paid_peers.clone(),
                            },
                        ));
                        Err(ProtocolError::TopologyVerificationFailed {
                            target_address: Box::new(target_address),
                            valid_count,
                            total_paid,
                            closest_count,
                            node_peers,
                            paid_peers,
                        })
                    }
                    Err(err) => {
                        node.record_metrics(Marker::RecordRejected(&key, &err));
                        Err(ProtocolError::PutRecordFailed(format!("{err:?}")))
                    }
                };
                QueryResponse::PutRecord {
                    result,
                    peer_address: holder,
                    record_addr: address,
                }
            }
            Query::GetMerkleCandidateQuote {
                key,
                data_type,
                data_size,
                merkle_payment_timestamp,
            } => {
                Self::respond_merkle_candidate_quote(
                    network,
                    key,
                    data_type,
                    data_size,
                    merkle_payment_timestamp,
                    payment_address,
                )
                .await
            }
        };
        Response::Query(resp)
    }

    async fn respond_get_closest_peers(
        network: &Network,
        target: NetworkAddress,
        num_of_peers: Option<usize>,
        range: Option<[u8; 32]>,
        sign_result: bool,
    ) -> QueryResponse {
        let local_peers = network.get_local_peers_with_multiaddr().await;
        let peers: Vec<(NetworkAddress, Vec<Multiaddr>)> = if let Ok(local_peers) = local_peers {
            Self::calculate_get_closest_peers(local_peers, target.clone(), num_of_peers, range)
        } else {
            vec![]
        };

        let signature = if sign_result {
            let mut bytes = rmp_serde::to_vec(&target).unwrap_or_default();
            bytes.extend_from_slice(&rmp_serde::to_vec(&peers).unwrap_or_default());
            network.sign(&bytes).ok()
        } else {
            None
        };

        QueryResponse::GetClosestPeers {
            target,
            peers,
            signature,
        }
    }

    fn calculate_get_closest_peers(
        peer_addrs: Vec<(PeerId, Vec<Multiaddr>)>,
        target: NetworkAddress,
        num_of_peers: Option<usize>,
        range: Option<[u8; 32]>,
    ) -> Vec<(NetworkAddress, Vec<Multiaddr>)> {
        match (num_of_peers, range) {
            (_, Some(value)) => {
                let distance = U256::from_big_endian(&value);
                peer_addrs
                    .iter()
                    .filter_map(|(peer_id, multi_addrs)| {
                        let addr = NetworkAddress::from(*peer_id);
                        if target.distance(&addr).0 <= distance {
                            Some((addr, multi_addrs.clone()))
                        } else {
                            None
                        }
                    })
                    .collect()
            }
            (Some(num_of_peers), _) => {
                let mut result: Vec<(NetworkAddress, Vec<Multiaddr>)> = peer_addrs
                    .iter()
                    .map(|(peer_id, multi_addrs)| {
                        let addr = NetworkAddress::from(*peer_id);
                        (addr, multi_addrs.clone())
                    })
                    .collect();
                result.sort_by_key(|(addr, _multi_addrs)| target.distance(addr));
                result.into_iter().take(num_of_peers).collect()
            }
            (None, None) => vec![],
        }
    }

    /// Handle GetMerkleCandidateQuote query
    /// Returns a signed MerklePaymentCandidateNode containing the node's current quoting metrics,
    /// reward address, and timestamp commitment
    async fn respond_merkle_candidate_quote(
        network: &Network,
        key: NetworkAddress,
        data_type: u32,
        data_size: usize,
        merkle_payment_timestamp: u64,
        payment_address: RewardsAddress,
    ) -> QueryResponse {
        debug!(
            "merkle payment: GetMerkleCandidateQuote for target {key:?}, timestamp: {merkle_payment_timestamp}, data_type: {data_type}, data_size: {data_size}"
        );

        // Validate timestamp before signing to prevent committing to invalid times.
        // Nodes will reject proofs with expired/future timestamps during payment verification,
        // so signing such timestamps would create useless quotes that can't be used.
        //
        // Allow ±24 hours tolerance to handle timezone differences and clock skew between
        // clients and nodes. This prevents valid payments from being rejected due to minor
        // time differences while still catching truly expired/invalid timestamps.
        const TIMESTAMP_TOLERANCE: u64 = 24 * 60 * 60; // 24 hours

        let now = std::time::SystemTime::now()
            .duration_since(std::time::UNIX_EPOCH)
            .unwrap_or_default()
            .as_secs();

        // Reject future timestamps (with 24h tolerance for clock skew)
        let future_threshold = now + TIMESTAMP_TOLERANCE;
        if merkle_payment_timestamp > future_threshold {
            let error_msg = format!(
                "Rejected future timestamp {merkle_payment_timestamp} (current time: {now}, threshold: {future_threshold})"
            );
            warn!("{error_msg} for {key:?}");
            return QueryResponse::GetMerkleCandidateQuote(Err(
                ProtocolError::GetMerkleCandidateQuoteFailed(error_msg),
            ));
        }

        // Reject expired timestamps (with 24h tolerance for clock skew)
        let expiration_threshold = MERKLE_PAYMENT_EXPIRATION + TIMESTAMP_TOLERANCE;
        let age = now.saturating_sub(merkle_payment_timestamp);
        if age > expiration_threshold {
            let error_msg = format!(
                "Rejected expired timestamp {merkle_payment_timestamp} (age: {age}s, max: {expiration_threshold}s)",
            );
            warn!("{error_msg} for {key:?}");
            return QueryResponse::GetMerkleCandidateQuote(Err(
                ProtocolError::GetMerkleCandidateQuoteFailed(error_msg),
            ));
        }

        // Get node's current quoting metrics
        let record_key = key.to_record_key();
        let (quoting_metrics, _is_already_stored) = match network
            .get_local_quoting_metrics(record_key, data_type, data_size)
            .await
        {
            Ok(metrics) => metrics,
            Err(err) => {
                let error_msg = format!("Failed to get quoting metrics for {key:?}: {err}");
                warn!("{error_msg}");
                return QueryResponse::GetMerkleCandidateQuote(Err(
                    ProtocolError::GetMerkleCandidateQuoteFailed(error_msg),
                ));
            }
        };

        // Create the MerklePaymentCandidateNode with node's signed commitment
        let pub_key = network.get_pub_key();
        let reward_address = payment_address;
        let bytes = ant_evm::merkle_payments::MerklePaymentCandidateNode::bytes_to_sign(
            &quoting_metrics,
            &reward_address,
            merkle_payment_timestamp,
        );
        let signature = match network.sign(&bytes) {
            Ok(sig) => sig,
            Err(e) => {
                let error_msg = format!("Failed to sign candidate node for {key:?}: {e}");
                error!("{error_msg}");
                return QueryResponse::GetMerkleCandidateQuote(Err(
                    ProtocolError::FailedToSignMerkleCandidate(error_msg),
                ));
            }
        };

        let candidate = ant_evm::merkle_payments::MerklePaymentCandidateNode {
            quoting_metrics,
            reward_address,
            merkle_payment_timestamp,
            pub_key,
            signature,
        };
        QueryResponse::GetMerkleCandidateQuote(Ok(candidate))
    }

    // Nodes only check ChunkProof each other, to avoid `multi-version` issue
    // Client check proof against all records, as have to fetch from network anyway.
    async fn respond_x_closest_record_proof(
        network: &Network,
        key: NetworkAddress,
        nonce: Nonce,
        difficulty: usize,
        chunk_only: bool,
    ) -> Vec<(NetworkAddress, Result<ChunkProof, ProtocolError>)> {
        let start = Instant::now();
        let mut results = vec![];
        if difficulty == 1 {
            // Client checking existence of published chunk.
            let mut result = Err(ProtocolError::ChunkDoesNotExist(key.clone()));
            if let Ok(Some(record)) = network.get_local_record(&key.to_record_key()).await {
                let proof = ChunkProof::new(&record.value, nonce);
                debug!("Chunk proof for {key:?} is {proof:?}");
                result = Ok(proof)
            } else {
                debug!("Could not get ChunkProof for {key:?} as we don't have the record locally.");
            }

            results.push((key.clone(), result));
        } else {
            let all_local_records = network.get_all_local_record_addresses().await;

            if let Ok(all_local_records) = all_local_records {
                let mut all_chunk_addrs: Vec<_> = if chunk_only {
                    all_local_records
                        .iter()
                        .filter_map(|(addr, record_type)| {
                            if *record_type == ValidationType::Chunk {
                                Some(addr.clone())
                            } else {
                                None
                            }
                        })
                        .collect()
                } else {
                    all_local_records.keys().cloned().collect()
                };

                // Sort by distance and only take first X closest entries
                all_chunk_addrs.sort_by_key(|addr| key.distance(addr));

                // TODO: this shall be deduced from resource usage dynamically
                let workload_factor = std::cmp::min(difficulty, CLOSE_GROUP_SIZE);

                for addr in all_chunk_addrs.iter().take(workload_factor) {
                    if let Ok(Some(record)) = network.get_local_record(&addr.to_record_key()).await
                    {
                        let proof = ChunkProof::new(&record.value, nonce);
                        debug!("Chunk proof for {key:?} is {proof:?}");
                        results.push((addr.clone(), Ok(proof)));
                    }
                }
            }

            info!(
                "Respond with {} answers to the StorageChallenge targeting {key:?} with {difficulty} difficulty, in {:?}",
                results.len(),
                start.elapsed()
            );
        }

        results
    }

    /// Check among all chunk type records that we have,
    /// and randomly pick one as the verification candidate.
    /// This will challenge all closest peers at once.
    async fn storage_challenge(network: Network) {
        let start = Instant::now();
        let closest_peers: Vec<(PeerId, Addresses)> = if let Ok(closest_peers) =
            network.get_k_closest_local_peers_to_the_target(None).await
        {
            closest_peers
                .into_iter()
                .take(CLOSE_GROUP_SIZE)
                .collect_vec()
        } else {
            error!("Cannot get local neighbours");
            return;
        };
        if closest_peers.len() < CLOSE_GROUP_SIZE {
            debug!(
                "Not enough neighbours ({}/{}) to carry out storage challenge.",
                closest_peers.len(),
                CLOSE_GROUP_SIZE
            );
            return;
        }

        let mut verify_candidates: Vec<NetworkAddress> =
            if let Ok(all_keys) = network.get_all_local_record_addresses().await {
                all_keys
                    .iter()
                    .filter_map(|(addr, record_type)| {
                        if ValidationType::Chunk == *record_type {
                            Some(addr.clone())
                        } else {
                            None
                        }
                    })
                    .collect()
            } else {
                error!("Failed to get local record addresses.");
                return;
            };
        let num_of_targets = verify_candidates.len();
        if num_of_targets < 50 {
            debug!("Not enough candidates({num_of_targets}/50) to be checked against neighbours.");
            return;
        }

        // To ensure the neighbours sharing same knowledge as to us,
        // The target is choosen to be not far from us.
        let self_addr = NetworkAddress::from(network.peer_id());
        verify_candidates.sort_by_key(|addr| self_addr.distance(addr));
        let index: usize = OsRng.gen_range(0..num_of_targets / 2);
        let target = verify_candidates[index].clone();
        // TODO: workload shall be dynamically deduced from resource usage
        let difficulty = CLOSE_GROUP_SIZE;
        verify_candidates.sort_by_key(|addr| target.distance(addr));
        let expected_targets = verify_candidates.into_iter().take(difficulty);
        let nonce: Nonce = thread_rng().r#gen::<u64>();
        let mut expected_proofs = HashMap::new();
        for addr in expected_targets {
            if let Ok(Some(record)) = network.get_local_record(&addr.to_record_key()).await {
                let expected_proof = ChunkProof::new(&record.value, nonce);
                let _ = expected_proofs.insert(addr, expected_proof);
            } else {
                error!("Local record {addr:?} cann't be loaded from disk.");
            }
        }
        let request = Request::Query(Query::GetChunkExistenceProof {
            key: target.clone(),
            nonce,
            difficulty,
        });

        let mut tasks = JoinSet::new();
        for (peer_id, addresses) in closest_peers {
            if peer_id == network.peer_id() {
                continue;
            }
            let network_clone = network.clone();
            let request_clone = request.clone();
            let expected_proofs_clone = expected_proofs.clone();
            let _ = tasks.spawn(async move {
                let res = scoring_peer(
                    network_clone,
                    (peer_id, addresses),
                    request_clone,
                    expected_proofs_clone,
                )
                .await;
                (peer_id, res)
            });
        }

        let mut peer_scores = vec![];
        while let Some(res) = tasks.join_next().await {
            match res {
                Ok((peer_id, score)) => {
                    let is_healthy = score > MIN_ACCEPTABLE_HEALTHY_SCORE;
                    if !is_healthy {
                        info!(
                            "Peer {peer_id:?} failed storage challenge with low score {score}/{MIN_ACCEPTABLE_HEALTHY_SCORE}."
                        );
                        // TODO: shall the challenge failure immediately triggers the node to be removed?
                        network.record_node_issues(peer_id, NodeIssue::FailedChunkProofCheck);
                    }
                    peer_scores.push((peer_id, is_healthy));
                }
                Err(e) => {
                    info!("StorageChallenge task completed with error {e:?}");
                }
            }
        }
        if !peer_scores.is_empty() {
            network.notify_peer_scores(peer_scores);
        }

        Self::verify_local_replication_health(network.clone()).await;

        info!(
            "Completed node StorageChallenge against neighbours in {:?}!",
            start.elapsed()
        );
    }

    /// Perform a direct replicated record fetch from nearby peers to verify replication health.
    async fn verify_local_replication_health(network: Network) {
        let closest_peers = match network
            .get_k_closest_local_peers_to_the_target(None)
            .await
        {
            Ok(peers) => peers,
            Err(err) => {
                warn!("Cannot fetch closest peers for replica verification: {err:?}");
                return;
            }
        };

        if closest_peers.len() <= CLOSE_NEIGHBOUR_DISTANCE_INDEX {
            debug!(
                "Skipping replica verification as we only know {} neighbours (need > {}).",
                closest_peers.len(),
                CLOSE_NEIGHBOUR_DISTANCE_INDEX
            );
            return;
        }

        let self_address = NetworkAddress::from(network.peer_id());
        let Some((threshold_peer, _)) = closest_peers.get(CLOSE_NEIGHBOUR_DISTANCE_INDEX) else {
            debug!("Unable to determine distance threshold for replica verification.");
            return;
        };

        let threshold_distance =
            self_address.distance(&NetworkAddress::from(*threshold_peer));
        let Some(threshold_ilog2) = threshold_distance.ilog2() else {
            debug!("Threshold distance lacks ilog2; cannot proceed with replica verification.");
            return;
        };

        let local_records = match network.get_all_local_record_addresses().await {
            Ok(records) => records,
            Err(err) => {
                warn!("Failed to list local records for replica verification: {err:?}");
                return;
            }
        };

        let mut nearby_records: Vec<NetworkAddress> = local_records
            .into_iter()
            .filter_map(|(address, record_type)| {
                if record_type != ValidationType::Chunk {
                    return None;
                }
                let distance = self_address.distance(&address);
                distance
                    .ilog2()
                    .and_then(|record_ilog2| {
                        if record_ilog2 <= threshold_ilog2 {
                            Some(address)
                        } else {
                            None
                        }
                    })
            })
            .collect();

        nearby_records.shuffle(&mut thread_rng());
        let target_record = if let Some(entry) = nearby_records.first().cloned() {
            entry
        } else {
            debug!("No nearby chunk records available for replica verification.");
            return;
        };

        let pretty_key =
            PrettyPrintRecordKey::from(&target_record.to_record_key()).into_owned();

        let candidate_peers = match network
            .get_k_closest_local_peers_to_the_target(Some(target_record.clone()))
            .await
        {
            Ok(peers) => peers
                .into_iter()
                .filter(|(peer_id, _)| peer_id != &network.peer_id())
                .take(REPLICA_FETCH_PEER_COUNT)
                .collect::<Vec<_>>(),
            Err(err) => {
                warn!(
                    "Cannot fetch record-specific closest peers for replica verification: {err:?}"
                );
                return;
            }
        };

        if candidate_peers.len() < REPLICA_FETCH_PEER_COUNT {
            debug!(
                "Only {} peers available for replica verification (need at least {}).",
                candidate_peers.len(),
                REPLICA_FETCH_PEER_COUNT
            );
            return;
        }

        debug!(
            "Verifying replicated record {pretty_key:?} against {} closest peers.",
            candidate_peers.len()
        );

        let (successful_peers, failed_peers) = Self::fetch_record_from_peers_with_addresses(
            network.clone(),
            target_record.clone(),
            candidate_peers,
        )
        .await;

        if failed_peers.is_empty() {
            debug!(
                "All peers returned record {pretty_key:?} during replica verification."
            );
            return;
        }

        if successful_peers.len() < MIN_HEALTHY_REPLICA_COUNT {
            warn!(
                "Replica verification fetched only {} copies of {pretty_key:?}. Record is unhealthy; skipping peer classification.",
                successful_peers.len()
            );
            return;
        }

        if !failed_peers.is_empty() {
            info!(
                "Scheduling replica verification retry for {} peers on record {pretty_key:?}.",
                failed_peers.len()
            );
            Self::schedule_record_fetch_retry(network.clone(), target_record.clone(), failed_peers).await;
        }
    }

    /// Fetch a replicated record from the provided peers and return the success/failure split.
    async fn fetch_record_from_peers_with_addresses(
        network: Network,
        record_address: NetworkAddress,
        peers: Vec<(PeerId, Addresses)>,
    ) -> (Vec<PeerId>, Vec<PeerId>) {
        let request = Request::Query(Query::GetReplicatedRecord {
            requester: NetworkAddress::from(network.peer_id()),
            key: record_address.clone(),
        });
        let expected_key = record_address.to_record_key();
        let pretty_key = PrettyPrintRecordKey::from(&expected_key).into_owned();

        let mut successes = Vec::new();
        let mut failures = Vec::new();
        let concurrency = peers.len();
        let results = stream::iter(
            peers
                .into_iter()
                .map(|(peer_id, addrs)| {
                    let network_clone = network.clone();
                    let request_clone = request.clone();
                    async move {
                        let result = network_clone
                            .send_request(request_clone, peer_id, addrs.clone())
                            .await;
                        (peer_id, addrs, result)
                    }
                }),
        )
        .buffer_unordered(concurrency)
        .collect::<Vec<_>>()
        .await;

        for res in results {
            match res {
                (peer_id, _addrs, Ok((Response::Query(QueryResponse::GetReplicatedRecord(result)), _))) => {
                    match result {
                        Ok((_holder, value)) => {
                            // Try to deserialize the record and ensure the Chunk content matches
                            let record = Record::new(record_address.to_record_key(), value.to_vec());
                            if let Ok(chunk) = try_deserialize_record::<Chunk>(&record) 
                            && chunk.network_address().to_record_key() == expected_key {
                                successes.push(peer_id);
                            } else {
                                warn!(
                                    "Peer {peer_id:?} responded with an incorrect chunk copy of {pretty_key:?}."
                                );
                                failures.push(peer_id);
                            }
                        }
                        Err(err) => {
                            info!(
                                "Peer {peer_id:?} responded with error {err:?} for replicated record {pretty_key:?}."
                            );
                            failures.push(peer_id);
                        }
                    }
                }
                (peer_id, _addrs, Ok((other_response, _))) => {
                    warn!(
                        "Peer {peer_id:?} responded with unexpected message {other_response:?} for {pretty_key:?}."
                    );
                    failures.push(peer_id);
                }
                (peer_id, _addrs, Err(err)) => {
                    info!(
                        "Failed to reach peer {peer_id:?} for replicated record {pretty_key:?}: {err:?}"
                    );
                    failures.push(peer_id);
                }
            }
        }

        (successes, failures)
    }

    /// Schedule a future retry for peers that failed to return the replicated record.
    async fn schedule_record_fetch_retry(
        network: Network,
        record_address: NetworkAddress,
        failed_peers: Vec<PeerId>,
    ) {
        let retry_peers: HashSet<_> = failed_peers.into_iter().collect();
        if retry_peers.is_empty() {
            return;
        }

        let record_clone = record_address.clone();
        let network_clone = network.clone();

        tokio::time::sleep(REPLICA_FETCH_RETRY_DELAY).await;
        let pretty_key =
            PrettyPrintRecordKey::from(&record_clone.to_record_key()).into_owned();

        let refreshed_candidates = Self::refresh_retry_candidate_addresses(
            &network_clone,
            &record_clone,
            &retry_peers,
        )
        .await;

        if refreshed_candidates.is_empty() {
            info!(
                "Skipping replica retry for {pretty_key:?}; no tracked peers remain close to the record."
            );
            return;
        }

        let (_, still_failed) = Self::fetch_record_from_peers_with_addresses(
            network_clone.clone(),
            record_clone.clone(),
            refreshed_candidates,
        )
        .await;

        if still_failed.is_empty() {
            info!(
                "All peers successfully returned {pretty_key:?} during replica retry."
            );
            return;
        }

        warn!(
            "{} peers still failed to provide {pretty_key:?}; evicting and blacklisting.",
            still_failed.len()
        );
        for peer_id in still_failed {
            Self::evict_and_blacklist_peer(network_clone.clone(), peer_id);
        }
    }

    /// Get the latest address information for tracked peers via a DHT query.
    async fn refresh_retry_candidate_addresses(
        network: &Network,
        record_address: &NetworkAddress,
        retry_peers: &HashSet<PeerId>,
    ) -> Vec<(PeerId, Addresses)> {
        let pretty_key =
            PrettyPrintRecordKey::from(&record_address.to_record_key()).into_owned();
        let Ok(closest_peers) = network.get_closest_peers(record_address).await else {
            warn!(
                "Failed to refresh peer addresses for {pretty_key:?}; unable to retry replica fetch."
            );
            return Vec::new();
        };

        let closest_map: HashMap<PeerId, Addresses> = closest_peers
            .into_iter()
            .collect();

        retry_peers
            .iter()
            .filter_map(|peer_id| closest_map.get(peer_id).cloned().map(|addrs| (*peer_id, addrs)))
            .collect()
    }

    /// Remove the peer from the routing table and add it to the blacklist.
    fn evict_and_blacklist_peer(network: Network, peer_id: PeerId) {
        network.remove_peer(peer_id);
        network.add_peer_to_blocklist(peer_id);
    }

    /// Query peers' versions and update local knowledge.
    async fn query_peers_version(network: Network, peers: Vec<(PeerId, Addresses)>) {
        // To avoid choking, carry out the queries one by one
        for (peer_id, addrs) in peers {
            Self::try_query_peer_version(network.clone(), peer_id, addrs).await;
        }
    }

    /// Query peer's version and update local knowledge.
    async fn try_query_peer_version(network: Network, peer: PeerId, addrs: Addresses) {
        let request = Request::Query(Query::GetVersion(NetworkAddress::from(peer)));
        // We can skip passing `addrs` here as the new peer should be part of the kad::RT and swarm can get the addr.
        let version = match network.send_request(request, peer, addrs).await {
            Ok((Response::Query(QueryResponse::GetVersion { version, .. }), _conn_info)) => {
                trace!("Fetched peer version {peer:?} as {version:?}");
                version
            }
            Ok(other) => {
                info!("Not a fetched peer version {peer:?}, {other:?}");
                "none".to_string()
            }
            Err(err) => {
                info!("Failed to fetch peer version {peer:?} with error {err:?}");
                // Failed version fetch (which contains dial then re-attempt by itself)
                // with error of `DialFailure` indicates the peer could be dead with high chance.
                // In that case, the peer shall be removed from the routing table.
                if let NetworkError::OutboundError(OutboundFailure::DialFailure) = err {
                    network.remove_peer(peer);
                    return;
                }
                "old".to_string()
            }
        };
        network.notify_node_version(peer, version);
    }
}

#[derive(Debug)]
struct CloseGroupTracker {
    self_address: NetworkAddress,
    /// Tracks the closest N peers by (distance, peer_id), sorted by distance to self.
    /// Using a tuple key handles the theoretical case where two peers have identical
    /// XOR distance to self.
    close_up_peers: BTreeSet<(Distance, PeerId)>,
    tracked_entries: HashMap<PeerId, BehaviourEntry>,
}

impl CloseGroupTracker {
    fn new(self_peer_id: PeerId) -> Self {
        Self {
            self_address: NetworkAddress::from(self_peer_id),
            close_up_peers: BTreeSet::new(),
            tracked_entries: HashMap::new(),
        }
    }

    fn record_peer_added(&mut self, peer_id: PeerId) -> ReplicationDirective {
        let distance = self.distance_to_peer(peer_id);

        // Check if this peer was already present
        let was_present = self.close_up_peers.contains(&(distance, peer_id));

        // Determine if the peer is close enough to track
        let is_close_enough = self.is_distance_within_close_range(distance);

        if is_close_enough {
            // Add to close_up_peers
            let _ = self.close_up_peers.insert((distance, peer_id));

            // If we exceed the limit, remove the farthest peer
            if self.close_up_peers.len() > CLOSE_GROUP_TRACKING_LIMIT
                && let Some(&(farthest_distance, farthest_peer)) =
                    self.close_up_peers.iter().next_back()
            {
                let _ = self.close_up_peers.remove(&(farthest_distance, farthest_peer));
            }
        } else {
            // Too far, don't track and clean up any existing entry
            let _ = self.tracked_entries.remove(&peer_id);
            return ReplicationDirective::Ignore;
        }

        use std::collections::hash_map::Entry;
        match self.tracked_entries.entry(peer_id) {
            Entry::Vacant(vacant_entry) => {
                let _ = vacant_entry.insert(BehaviourEntry::default());
                if !was_present {
                    ReplicationDirective::Trigger
                } else {
                    ReplicationDirective::Ignore
                }
            }
            Entry::Occupied(mut occupied_entry) => {
                let entry = occupied_entry.get_mut();
                if entry.awaiting_rejoin || entry.timer_deadline.is_some() {
                    entry.restart_detected = true;
                    entry.awaiting_rejoin = false;
                    entry.timer_deadline = None;
                    ReplicationDirective::Skip
                } else if !was_present {
                    entry.restart_detected = false;
                    ReplicationDirective::Trigger
                } else {
                    ReplicationDirective::Ignore
                }
            }
        }
    }

    fn record_peer_removed(&mut self, peer_id: PeerId) -> ReplicationDirective {
        let distance = self.distance_to_peer(peer_id);

        // Check if this peer is in our close_up_peers
        let was_tracked = self.close_up_peers.contains(&(distance, peer_id));

        if !was_tracked {
            // Not being tracked, i.e. being out-of-range, hence Ignore
            return ReplicationDirective::Ignore;
        }

        // Remove from close_up_peers
        let _ = self.close_up_peers.remove(&(distance, peer_id));

        let entry = self.tracked_entries.entry(peer_id).or_default();

        if entry.restart_detected {
            entry.timer_deadline = Some(Instant::now() + CLOSE_GROUP_RESTART_SUPPRESSION);
            entry.awaiting_rejoin = false;
            ReplicationDirective::Skip
        } else {
            entry.awaiting_rejoin = true;
            entry.timer_deadline = None;
            ReplicationDirective::Trigger
        }
    }

    fn handle_timer_expiry(&mut self, now: Instant) -> bool {
        let mut expired = false;
        for entry in self.tracked_entries.values_mut() {
            if let Some(deadline) = entry.timer_deadline
                && now >= deadline
            {
                entry.timer_deadline = None;
                entry.restart_detected = false;
                entry.awaiting_rejoin = false;
                expired = true;
            }
        }

        // Keep only entries that are still relevant
        self.tracked_entries.retain(|peer_id, entry| {
            let is_currently_present = self
                .close_up_peers
                .iter()
                .any(|(_, present_peer)| present_peer == peer_id);

            entry.awaiting_rejoin
                || entry.restart_detected
                || entry.timer_deadline.is_some()
                || is_currently_present
        });

        expired
    }

    fn is_distance_within_close_range(&self, distance: Distance) -> bool {
        // If we haven't reached the limit yet, any peer is close enough
        if self.close_up_peers.len() < CLOSE_GROUP_TRACKING_LIMIT {
            return true;
        }

        // Otherwise, check if this distance is closer than the farthest tracked peer
        if let Some(&(farthest_distance, _)) = self.close_up_peers.iter().next_back() {
            distance < farthest_distance
        } else {
            true
        }
    }

    fn distance_to_peer(&self, peer_id: PeerId) -> Distance {
        self.self_address.distance(&NetworkAddress::from(peer_id))
    }
}

#[derive(Debug, Default)]
struct BehaviourEntry {
    awaiting_rejoin: bool,
    restart_detected: bool,
    timer_deadline: Option<Instant>,
}

#[derive(Debug, PartialEq, Eq)]
enum ReplicationDirective {
    Trigger,
    Skip,
    Ignore,
}

impl ReplicationDirective {
    fn should_trigger(&self) -> bool {
        matches!(self, Self::Trigger)
    }
}

#[cfg(test)]
mod close_group_tracker_tests {
    use super::*;
    use std::time::Duration;

    fn random_peer() -> PeerId {
        let keypair = libp2p::identity::Keypair::generate_ed25519();
        PeerId::from(keypair.public())
    }

    #[test]
    fn new_peer_triggers_replication() {
        let mut tracker = CloseGroupTracker::new(random_peer());
        let peer = random_peer();

        assert_eq!(
            tracker.record_peer_added(peer),
            ReplicationDirective::Trigger
        );

        assert_eq!(
            tracker.record_peer_removed(peer),
            ReplicationDirective::Trigger
        );
    }

    #[test]
    fn restart_detection_skips_replication() {
        let mut tracker = CloseGroupTracker::new(random_peer());
        let peer = random_peer();

        assert_eq!(
            tracker.record_peer_added(peer),
            ReplicationDirective::Trigger
        );

        assert_eq!(
            tracker.record_peer_removed(peer),
            ReplicationDirective::Trigger
        );

        assert_eq!(tracker.record_peer_added(peer), ReplicationDirective::Skip);

        assert_eq!(
            tracker.record_peer_removed(peer),
            ReplicationDirective::Skip
        );

        assert!(tracker.handle_timer_expiry(
            Instant::now() + CLOSE_GROUP_RESTART_SUPPRESSION + Duration::from_secs(1)
        ));

        assert_eq!(
            tracker.record_peer_added(peer),
            ReplicationDirective::Trigger
        );
    }

    #[test]
    fn peer_outside_close_group_is_ignored() {
        let mut tracker = CloseGroupTracker::new(random_peer());

        // Fill the tracker to capacity with random peers
        let mut added_peers = Vec::new();
        for _ in 0..CLOSE_GROUP_TRACKING_LIMIT {
            let peer = random_peer();
            let result = tracker.record_peer_added(peer);
            // All should trigger since we're under capacity
            assert_eq!(result, ReplicationDirective::Trigger);
            added_peers.push(peer);
        }

        assert_eq!(tracker.close_up_peers.len(), CLOSE_GROUP_TRACKING_LIMIT);

        // Now add more peers - some will be ignored if they're farther than existing ones
        let mut ignored_count = 0;
        let mut trigger_count = 0;
        for _ in 0..50 {
            let peer = random_peer();
            match tracker.record_peer_added(peer) {
                ReplicationDirective::Ignore => ignored_count += 1,
                ReplicationDirective::Trigger => trigger_count += 1,
                _ => {}
            }
        }

        // At capacity, new peers are only accepted if closer than the farthest
        // Some should be ignored (too far), some should trigger (closer than farthest)
        assert!(
            ignored_count > 0 || trigger_count > 0,
            "Expected some peers to be processed"
        );

        // Tracker should still only have CLOSE_GROUP_TRACKING_LIMIT peers
        assert_eq!(tracker.close_up_peers.len(), CLOSE_GROUP_TRACKING_LIMIT);
    }

    #[test]
    fn removal_of_untracked_peer_is_ignored() {
        let mut tracker = CloseGroupTracker::new(random_peer());

        // Try to remove a peer that was never added
        let unknown_peer = random_peer();
        assert_eq!(
            tracker.record_peer_removed(unknown_peer),
            ReplicationDirective::Ignore
        );
    }

    #[test]
    fn timer_expiry_resets_restart_state() {
        let mut tracker = CloseGroupTracker::new(random_peer());
        let peer = random_peer();

        // Add and remove peer (normal behavior)
        assert_eq!(
            tracker.record_peer_added(peer),
            ReplicationDirective::Trigger
        );
        assert_eq!(
            tracker.record_peer_removed(peer),
            ReplicationDirective::Trigger
        );

        // Peer rejoins quickly - detected as restart, skipped
        assert_eq!(tracker.record_peer_added(peer), ReplicationDirective::Skip);

        // Remove again - still in restart mode, skipped with timer set
        assert_eq!(
            tracker.record_peer_removed(peer),
            ReplicationDirective::Skip
        );

        // Timer hasn't expired yet - no change
        assert!(!tracker.handle_timer_expiry(Instant::now()));

        // Timer expires
        let after_expiry = Instant::now() + CLOSE_GROUP_RESTART_SUPPRESSION + Duration::from_secs(1);
        assert!(tracker.handle_timer_expiry(after_expiry));

        // Now the peer should trigger replication again (state reset)
        assert_eq!(
            tracker.record_peer_added(peer),
            ReplicationDirective::Trigger
        );
    }

    #[test]
    fn eviction_maintains_closest_peers() {
        let self_peer = random_peer();
        let mut tracker = CloseGroupTracker::new(self_peer);

        // Add exactly CLOSE_GROUP_TRACKING_LIMIT peers
        let mut peers_with_distances: Vec<(PeerId, Distance)> = Vec::new();
        for _ in 0..CLOSE_GROUP_TRACKING_LIMIT {
            let peer = random_peer();
            let distance = tracker.distance_to_peer(peer);
            let _ = tracker.record_peer_added(peer);
            peers_with_distances.push((peer, distance));
        }

        assert_eq!(tracker.close_up_peers.len(), CLOSE_GROUP_TRACKING_LIMIT);

        // Sort by distance to know which is farthest
        peers_with_distances.sort_by_key(|(_, d)| *d);
        let farthest_distance = peers_with_distances.last().map(|(_, d)| *d).unwrap();

        // Add a new peer that's closer than the farthest
        // Keep trying until we find one (random, so may take a few tries)
        let mut found_closer = false;
        for _ in 0..100 {
            let new_peer = random_peer();
            let new_distance = tracker.distance_to_peer(new_peer);
            if new_distance < farthest_distance {
                let result = tracker.record_peer_added(new_peer);
                assert_eq!(result, ReplicationDirective::Trigger);
                // Should still have exactly CLOSE_GROUP_TRACKING_LIMIT peers
                assert_eq!(tracker.close_up_peers.len(), CLOSE_GROUP_TRACKING_LIMIT);
                found_closer = true;
                break;
            }
        }

        // Note: This assertion might occasionally fail due to randomness,
        // but with 100 attempts it's extremely unlikely
        assert!(
            found_closer,
            "Could not find a peer closer than the farthest in 100 attempts"
        );
    }

    #[test]
    fn tracked_entries_cleaned_up_on_timer_expiry() {
        let mut tracker = CloseGroupTracker::new(random_peer());
        let peer = random_peer();

        // Add, remove, re-add (restart detected), remove again (timer set)
        let _ = tracker.record_peer_added(peer);
        let _ = tracker.record_peer_removed(peer);
        let _ = tracker.record_peer_added(peer);
        let _ = tracker.record_peer_removed(peer);

        // Peer should have a tracked entry with timer
        assert!(tracker.tracked_entries.contains_key(&peer));

        // After timer expiry, entry should be cleaned up since peer is not present
        let after_expiry = Instant::now() + CLOSE_GROUP_RESTART_SUPPRESSION + Duration::from_secs(1);
        let _ = tracker.handle_timer_expiry(after_expiry);

        // Entry should be removed since peer is not in close_up_peers
        // and no longer has active flags
        assert!(!tracker.tracked_entries.contains_key(&peer));
    }
}

async fn scoring_peer(
    network: Network,
    peer: (PeerId, Addresses),
    request: Request,
    expected_proofs: HashMap<NetworkAddress, ChunkProof>,
) -> usize {
    let peer_id = peer.0;
    let start = Instant::now();
    let responses = network
        .send_and_get_responses(&[peer], &request, true)
        .await;

    if let Some(Ok((Response::Query(QueryResponse::GetChunkExistenceProof(answers)), _conn_info))) =
        responses.get(&peer_id)
    {
        if answers.is_empty() {
            info!("Peer {peer_id:?} didn't answer the ChunkProofChallenge.");
            return 0;
        }
        let elapsed = start.elapsed();

        let mut received_proofs = vec![];
        for (addr, proof) in answers {
            if let Ok(proof) = proof {
                received_proofs.push((addr.clone(), proof.clone()));
            }
        }

        let score = mark_peer(elapsed, received_proofs, &expected_proofs);
        info!(
            "Received {} answers from peer {peer_id:?} after {elapsed:?}, score it as {score}.",
            answers.len()
        );
        score
    } else {
        info!("Peer {peer_id:?} doesn't reply the ChunkProofChallenge, or replied with error.");
        0
    }
}

// Based on following metrics:
//   * the duration
//   * is there false answer
//   * percentage of correct answers among the expected closest
// The higher the score, the better confidence on the peer
fn mark_peer(
    duration: Duration,
    answers: Vec<(NetworkAddress, ChunkProof)>,
    expected_proofs: &HashMap<NetworkAddress, ChunkProof>,
) -> usize {
    let duration_score = duration_score_scheme(duration);
    let challenge_score = challenge_score_scheme(answers, expected_proofs);

    duration_score * challenge_score
}

// Less duration shall get higher score
fn duration_score_scheme(duration: Duration) -> usize {
    // So far just a simple stepped scheme, capped by HIGHEST_SCORE
    let in_ms = if let Some(value) = duration.as_millis().to_usize() {
        value
    } else {
        info!("Cannot get milli seconds from {duration:?}, using a default value of 1000ms.");
        1000
    };

    let step = std::cmp::min(HIGHEST_SCORE, in_ms / TIME_STEP);
    HIGHEST_SCORE - step
}

// Any false answer shall result in 0 score immediately
fn challenge_score_scheme(
    answers: Vec<(NetworkAddress, ChunkProof)>,
    expected_proofs: &HashMap<NetworkAddress, ChunkProof>,
) -> usize {
    let mut correct_answers = 0;
    for (addr, chunk_proof) in answers {
        if let Some(expected_proof) = expected_proofs.get(&addr) {
            if expected_proof.verify(&chunk_proof) {
                correct_answers += 1;
            } else {
                info!("Spot a false answer to the challenge regarding {addr:?}");
                // Any false answer shall result in 0 score immediately
                return 0;
            }
        }
    }
    // TODO: For those answers not among the expected_proofs,
    //       it could be due to having different knowledge of records to us.
    //       shall we:
    //         * set the target being close to us, so that neighbours sharing same knowledge in higher chance
    //         * fetch from local to testify
    //         * fetch from network to testify
    std::cmp::min(
        HIGHEST_SCORE,
        HIGHEST_SCORE * correct_answers / expected_proofs.len(),
    )
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::str::FromStr;

    #[test]
    fn test_no_local_peers() {
        let local_peers: Vec<(PeerId, Vec<Multiaddr>)> = vec![];
        let target = NetworkAddress::from(PeerId::random());
        let num_of_peers = Some(5);
        let range = None;
        let result = Node::calculate_get_closest_peers(local_peers, target, num_of_peers, range);

        assert_eq!(result, vec![]);
    }

    #[test]
    fn test_fewer_local_peers_than_num_of_peers() {
        let local_peers: Vec<(PeerId, Vec<Multiaddr>)> = vec![
            (
                PeerId::random(),
                vec![Multiaddr::from_str("/ip4/192.168.1.1/tcp/8080").unwrap()],
            ),
            (
                PeerId::random(),
                vec![Multiaddr::from_str("/ip4/192.168.1.2/tcp/8080").unwrap()],
            ),
            (
                PeerId::random(),
                vec![Multiaddr::from_str("/ip4/192.168.1.2/tcp/8080").unwrap()],
            ),
        ];
        let target = NetworkAddress::from(PeerId::random());
        let num_of_peers = Some(2);
        let range = None;
        let result = Node::calculate_get_closest_peers(
            local_peers.clone(),
            target.clone(),
            num_of_peers,
            range,
        );

        // Result shall be sorted and truncated
        let mut expected_result: Vec<(NetworkAddress, Vec<Multiaddr>)> = local_peers
            .iter()
            .map(|(peer_id, multi_addrs)| {
                let addr = NetworkAddress::from(*peer_id);
                (addr, multi_addrs.clone())
            })
            .collect();
        expected_result.sort_by_key(|(addr, _multi_addrs)| target.distance(addr));
        let expected_result: Vec<_> = expected_result.into_iter().take(2).collect();

        assert_eq!(expected_result, result);
    }

    #[test]
    fn test_with_range_and_num_of_peers() {
        let local_peers: Vec<(PeerId, Vec<Multiaddr>)> = vec![
            (
                PeerId::random(),
                vec![Multiaddr::from_str("/ip4/192.168.1.1/tcp/8080").unwrap()],
            ),
            (
                PeerId::random(),
                vec![Multiaddr::from_str("/ip4/192.168.1.2/tcp/8080").unwrap()],
            ),
            (
                PeerId::random(),
                vec![Multiaddr::from_str("/ip4/192.168.1.2/tcp/8080").unwrap()],
            ),
        ];
        let target = NetworkAddress::from(PeerId::random());
        let num_of_peers = Some(0);
        let range_value = [128; 32];
        let range = Some(range_value);
        let result = Node::calculate_get_closest_peers(
            local_peers.clone(),
            target.clone(),
            num_of_peers,
            range,
        );

        // Range shall be preferred, i.e. the result peers shall all within the range
        let distance = U256::from_big_endian(&range_value);
        let expected_result: Vec<(NetworkAddress, Vec<Multiaddr>)> = local_peers
            .into_iter()
            .filter_map(|(peer_id, multi_addrs)| {
                let addr = NetworkAddress::from(peer_id);
                if target.distance(&addr).0 <= distance {
                    Some((addr, multi_addrs.clone()))
                } else {
                    None
                }
            })
            .collect();

        assert_eq!(expected_result, result);
    }

    mod merkle_payment_tests {
        use super::*;

        /// Test that timestamp validation accepts valid timestamps (within the acceptable window)
        #[test]
        fn test_timestamp_validation_accepts_valid_timestamp() {
            let now = std::time::SystemTime::now()
                .duration_since(std::time::UNIX_EPOCH)
                .unwrap()
                .as_secs();

            // Valid timestamp: 1 hour ago
            let valid_timestamp = now - 3600;

            // Validate timestamp
            let age = now.saturating_sub(valid_timestamp);

            assert!(
                valid_timestamp <= now,
                "Valid timestamp should not be in the future"
            );
            assert!(
                age <= MERKLE_PAYMENT_EXPIRATION,
                "Valid timestamp should not be expired"
            );
        }

        /// Test that timestamp validation rejects future timestamps
        #[test]
        fn test_timestamp_validation_rejects_future_timestamp() {
            let now = std::time::SystemTime::now()
                .duration_since(std::time::UNIX_EPOCH)
                .unwrap()
                .as_secs();

            // Future timestamp: 1 hour in the future
            let future_timestamp = now + 3600;

            // Timestamp should be rejected
            assert!(
                future_timestamp > now,
                "Future timestamp should be rejected"
            );
        }

        /// Test that timestamp validation rejects expired timestamps
        #[test]
        fn test_timestamp_validation_rejects_expired_timestamp() {
            let now = std::time::SystemTime::now()
                .duration_since(std::time::UNIX_EPOCH)
                .unwrap()
                .as_secs();

            // Expired timestamp: 8 days ago (> 7 day expiration)
            let expired_timestamp = now - (MERKLE_PAYMENT_EXPIRATION + 86400);

            // Calculate age
            let age = now.saturating_sub(expired_timestamp);

            // Timestamp should be rejected
            assert!(
                age > MERKLE_PAYMENT_EXPIRATION,
                "Expired timestamp should be rejected"
            );
        }

        /// Test timestamp at the exact expiration boundary (should be rejected)
        #[test]
        fn test_timestamp_validation_at_expiration_boundary() {
            let now = std::time::SystemTime::now()
                .duration_since(std::time::UNIX_EPOCH)
                .unwrap()
                .as_secs();

            // Timestamp exactly at expiration boundary
            let boundary_timestamp = now - MERKLE_PAYMENT_EXPIRATION;

            let age = now.saturating_sub(boundary_timestamp);

            // At the boundary, age == MERKLE_PAYMENT_EXPIRATION
            assert_eq!(age, MERKLE_PAYMENT_EXPIRATION);
            // The validation uses >, so this should pass
            assert!(
                age <= MERKLE_PAYMENT_EXPIRATION,
                "Timestamp exactly at boundary should not be rejected"
            );
        }

        /// Test timestamp just beyond expiration boundary (should be rejected)
        #[test]
        fn test_timestamp_validation_beyond_expiration_boundary() {
            let now = std::time::SystemTime::now()
                .duration_since(std::time::UNIX_EPOCH)
                .unwrap()
                .as_secs();

            // Timestamp just beyond expiration boundary (1 second past)
            let beyond_boundary_timestamp = now - (MERKLE_PAYMENT_EXPIRATION + 1);

            let age = now.saturating_sub(beyond_boundary_timestamp);

            assert!(
                age > MERKLE_PAYMENT_EXPIRATION,
                "Timestamp beyond boundary should be rejected"
            );
        }

        /// Test timestamp at current time (should be accepted)
        #[test]
        fn test_timestamp_validation_at_current_time() {
            let now = std::time::SystemTime::now()
                .duration_since(std::time::UNIX_EPOCH)
                .unwrap()
                .as_secs();

            // Timestamp at current time
            let current_timestamp = now;

            let age = now.saturating_sub(current_timestamp);

            assert!(
                current_timestamp <= now,
                "Current timestamp should not be in future"
            );
            assert!(
                age <= MERKLE_PAYMENT_EXPIRATION,
                "Current timestamp should not be expired"
            );
            assert_eq!(age, 0, "Age should be 0 for current timestamp");
        }

        /// Test timestamp near future boundary (1 second in future)
        #[test]
        fn test_timestamp_validation_near_future_boundary() {
            let now = std::time::SystemTime::now()
                .duration_since(std::time::UNIX_EPOCH)
                .unwrap()
                .as_secs();

            // Timestamp 1 second in the future
            let near_future_timestamp = now + 1;

            assert!(
                near_future_timestamp > now,
                "Near-future timestamp should be rejected"
            );
        }

        /// Test expiration constant is set correctly (7 days = 604800 seconds)
        #[test]
        fn test_merkle_payment_expiration_constant() {
            const SEVEN_DAYS_IN_SECONDS: u64 = 7 * 24 * 60 * 60;
            assert_eq!(
                MERKLE_PAYMENT_EXPIRATION, SEVEN_DAYS_IN_SECONDS,
                "MERKLE_PAYMENT_EXPIRATION should be 7 days"
            );
        }
    }
}