1use crate::data::client::adaptive::{observe_op, Outcome};
8use crate::data::client::classify_error;
9use crate::data::client::file::UploadEvent;
10use crate::data::client::Client;
11use crate::data::error::{Error, Result};
12use ant_protocol::evm::{
13 Amount, MerklePaymentCandidateNode, MerklePaymentCandidatePool, MerklePaymentProof, MerkleTree,
14 MidpointProof, PoolCommitment, CANDIDATES_PER_POOL, MAX_LEAVES,
15};
16use ant_protocol::payment::commitment::{
17 commitment_hash, verify_commitment_signature, StorageCommitment, MAX_COMMITMENT_KEY_COUNT,
18 MAX_COMMITMENT_SIDECAR_BYTES,
19};
20use ant_protocol::payment::{
21 calculate_price, serialize_merkle_proof, verify_merkle_candidate_signature,
22};
23use ant_protocol::transport::PeerId;
24use ant_protocol::{
25 compute_address, send_and_await_chunk_response, ChunkMessage, ChunkMessageBody,
26 MerkleCandidateQuoteRequest, MerkleCandidateQuoteResponse,
27};
28use bytes::Bytes;
29use futures::stream::{self, FuturesUnordered, StreamExt};
30use rand::Rng;
31use std::collections::{HashMap, VecDeque};
32use std::time::Duration;
33use tokio::sync::mpsc;
34use tracing::{debug, info, warn};
35use xor_name::XorName;
36
37pub const DEFAULT_MERKLE_THRESHOLD: usize = 64;
39
40use crate::data::client::payment::SINGLE_NODE_PAYMENT_MULTIPLIER as MERKLE_PAYMENT_MULTIPLIER;
62
63fn merkle_candidate_binding_is_valid(
73 peer_id: &PeerId,
74 candidate: &MerklePaymentCandidateNode,
75 commitment: &Option<Vec<u8>>,
76) -> std::result::Result<(), String> {
77 let count = candidate.committed_key_count;
78 let pin = candidate.commitment_pin;
79 match (count, pin.is_some()) {
80 (0, false) | (1.., true) => {}
81 (1.., false) => {
82 return Err(format!(
83 "committed_key_count={count} > 0 but commitment_pin is None (unauditable count)"
84 ));
85 }
86 (0, true) => {
87 return Err("committed_key_count=0 with a commitment_pin (incoherent baseline)".into());
88 }
89 }
90 if count > MAX_COMMITMENT_KEY_COUNT {
91 return Err(format!(
92 "committed_key_count={count} exceeds MAX_COMMITMENT_KEY_COUNT={MAX_COMMITMENT_KEY_COUNT}"
93 ));
94 }
95 let expected = calculate_price(count as usize);
96 if candidate.price != expected {
97 return Err(format!(
98 "price {} does not equal calculate_price(committed_key_count={count}) = {expected}",
99 candidate.price
100 ));
101 }
102
103 let Some(pin) = pin else {
104 return Ok(()); };
106 let Some(blob) = commitment else {
107 return Err("bound candidate did not ship its commitment; pin is unresolvable".into());
108 };
109 if blob.len() > MAX_COMMITMENT_SIDECAR_BYTES {
110 return Err(format!(
111 "shipped commitment is {} bytes, exceeds MAX_COMMITMENT_SIDECAR_BYTES={MAX_COMMITMENT_SIDECAR_BYTES}",
112 blob.len()
113 ));
114 }
115 let commitment: StorageCommitment = rmp_serde::from_slice(blob)
116 .map_err(|e| format!("shipped commitment did not deserialize: {e}"))?;
117 if compute_address(&commitment.sender_public_key) != *peer_id.as_bytes()
118 || commitment.sender_peer_id != *peer_id.as_bytes()
119 {
120 return Err("shipped commitment is not bound to the candidate peer".into());
121 }
122 if !verify_commitment_signature(&commitment) {
123 return Err("shipped commitment has an invalid signature".into());
124 }
125 if commitment_hash(&commitment) != Some(pin) {
126 return Err("shipped commitment does not hash to the candidate's pin".into());
127 }
128 if commitment.key_count != count {
129 return Err(format!(
130 "shipped commitment attests key_count={} but the candidate claims {count}",
131 commitment.key_count
132 ));
133 }
134 Ok(())
135}
136
137fn pool_commitment_with_payment_multiplier(
150 pool: &MerklePaymentCandidatePool,
151) -> Result<PoolCommitment> {
152 let mut commitment = pool.to_commitment();
153 let multiplier = Amount::from(MERKLE_PAYMENT_MULTIPLIER);
154 for candidate in &mut commitment.candidates {
155 candidate.price = candidate.price.checked_mul(multiplier).ok_or_else(|| {
156 Error::Payment(format!(
157 "Merkle candidate amount overflow applying {MERKLE_PAYMENT_MULTIPLIER}x to price {}",
158 candidate.price
159 ))
160 })?;
161 }
162 Ok(commitment)
163}
164
165#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
167#[serde(rename_all = "snake_case")]
168pub enum PaymentMode {
169 #[default]
171 Auto,
172 Merkle,
174 Single,
176}
177
178#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
183pub struct MerkleBatchPaymentResult {
184 pub proofs: HashMap<[u8; 32], Vec<u8>>,
186 pub chunk_count: usize,
188 pub storage_cost_atto: String,
190 pub gas_cost_wei: u128,
192 #[serde(default)]
197 pub merkle_payment_timestamp: u64,
198}
199
200pub struct PreparedMerkleBatch {
205 pub depth: u8,
207 pub pool_commitments: Vec<PoolCommitment>,
209 pub merkle_payment_timestamp: u64,
211 candidate_pools: Vec<MerklePaymentCandidatePool>,
213 tree: MerkleTree,
215 addresses: Vec<[u8; 32]>,
217}
218
219#[derive(Debug, Clone, Default)]
221pub(crate) struct MerkleUploadPlan {
222 pub already_stored: Vec<[u8; 32]>,
224 pub to_upload: Vec<[u8; 32]>,
226 to_upload_total_bytes: u64,
228}
229
230impl MerkleUploadPlan {
231 #[must_use]
233 pub fn to_upload_avg_size(&self) -> u64 {
234 if self.to_upload.is_empty() {
235 return 0;
236 }
237
238 self.to_upload_total_bytes / self.to_upload.len() as u64
239 }
240}
241
242impl std::fmt::Debug for PreparedMerkleBatch {
243 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
244 f.debug_struct("PreparedMerkleBatch")
245 .field("depth", &self.depth)
246 .field("pool_commitments", &self.pool_commitments.len())
247 .field("merkle_payment_timestamp", &self.merkle_payment_timestamp)
248 .field("candidate_pools", &self.candidate_pools.len())
249 .field("addresses", &self.addresses.len())
250 .finish()
251 }
252}
253
254pub(crate) fn chunk_contents_for_upload_addresses(
259 chunk_contents: Vec<Bytes>,
260 addresses: &[[u8; 32]],
261) -> Result<Vec<Bytes>> {
262 if addresses.is_empty() {
263 return Ok(Vec::new());
264 }
265
266 let mut needed_by_address: HashMap<[u8; 32], usize> = HashMap::new();
267 for address in addresses {
268 *needed_by_address.entry(*address).or_default() += 1;
269 }
270
271 let mut chunks_by_address: HashMap<[u8; 32], VecDeque<Bytes>> =
272 HashMap::with_capacity(needed_by_address.len());
273 let mut remaining = addresses.len();
274 for chunk in chunk_contents {
275 let address = compute_address(&chunk);
276 if let Some(needed) = needed_by_address.get_mut(&address) {
277 if *needed > 0 {
278 chunks_by_address
279 .entry(address)
280 .or_default()
281 .push_back(chunk);
282 *needed -= 1;
283 remaining -= 1;
284 if remaining == 0 {
285 break;
286 }
287 }
288 }
289 }
290
291 for (address, needed) in &needed_by_address {
292 if *needed == 0 {
293 continue;
294 }
295
296 if chunks_by_address.contains_key(address) {
297 return Err(Error::InvalidData(format!(
298 "missing duplicate chunk content for merkle address {}",
299 hex::encode(address)
300 )));
301 }
302
303 return Err(Error::InvalidData(format!(
304 "missing chunk content for merkle address {}",
305 hex::encode(address)
306 )));
307 }
308
309 let mut selected = Vec::with_capacity(addresses.len());
310 for address in addresses {
311 let chunks = chunks_by_address.get_mut(address).ok_or_else(|| {
312 Error::InvalidData(format!(
313 "missing chunk content for merkle address {}",
314 hex::encode(address)
315 ))
316 })?;
317 let chunk = chunks.pop_front().ok_or_else(|| {
318 Error::InvalidData(format!(
319 "missing duplicate chunk content for merkle address {}",
320 hex::encode(address)
321 ))
322 })?;
323 selected.push(chunk);
324 }
325
326 Ok(selected)
327}
328
329fn preflight_stored_status<T>(result: Result<T>) -> Result<bool> {
348 match result {
349 Ok(_) => Ok(false),
350 Err(Error::AlreadyStored) => Ok(true),
351 Err(e) if matches!(classify_error(&e), Outcome::Timeout | Outcome::NetworkError) => {
352 Ok(false)
353 }
354 Err(e) => Err(e),
355 }
356}
357
358#[must_use]
376pub fn merkle_batch_sizes(total: usize) -> Vec<usize> {
377 if total < 2 {
378 return Vec::new();
379 }
380
381 let mut sizes = Vec::with_capacity(total.div_ceil(MAX_LEAVES));
382 let mut remaining = total;
383 while remaining > MAX_LEAVES {
384 let take = if remaining - MAX_LEAVES == 1 {
387 MAX_LEAVES - 1
388 } else {
389 MAX_LEAVES
390 };
391 sizes.push(take);
392 remaining -= take;
393 }
394 sizes.push(remaining);
395 sizes
396}
397
398#[must_use]
403pub fn merkle_batch_partitions(addresses: &[[u8; 32]]) -> Vec<&[[u8; 32]]> {
404 let mut partitions = Vec::new();
405 let mut rest = addresses;
406 for size in merkle_batch_sizes(addresses.len()) {
407 let (batch, tail) = rest.split_at(size);
408 partitions.push(batch);
409 rest = tail;
410 }
411 partitions
412}
413
414fn padded_leaf_count(batch_size: usize) -> u64 {
419 let padded = batch_size
422 .max(2)
423 .checked_next_power_of_two()
424 .unwrap_or(usize::MAX);
425 u64::try_from(padded).unwrap_or(u64::MAX)
426}
427
428#[must_use]
437pub fn merkle_billable_leaves(chunk_count: u64) -> u64 {
438 let total = usize::try_from(chunk_count).unwrap_or(usize::MAX);
439 let batches = merkle_batch_sizes(total);
440 if batches.is_empty() {
441 return if total == 0 { 0 } else { 2 };
444 }
445
446 batches
447 .into_iter()
448 .map(padded_leaf_count)
449 .fold(0u64, u64::saturating_add)
450}
451
452fn ensure_single_merkle_tree_batch(address_count: usize) -> Result<()> {
461 if address_count > MAX_LEAVES {
462 return Err(Error::MerkleBatchTooLarge {
463 addresses: address_count,
464 max_leaves: MAX_LEAVES,
465 });
466 }
467 Ok(())
468}
469
470#[must_use]
473pub fn should_use_merkle(chunk_count: usize, mode: PaymentMode) -> bool {
474 match mode {
475 PaymentMode::Auto => chunk_count >= DEFAULT_MERKLE_THRESHOLD,
476 PaymentMode::Merkle => chunk_count >= 2,
477 PaymentMode::Single => false,
478 }
479}
480
481impl Client {
482 #[must_use]
484 pub fn should_use_merkle(&self, chunk_count: usize, mode: PaymentMode) -> bool {
485 should_use_merkle(chunk_count, mode)
486 }
487
488 pub async fn pay_for_merkle_batch(
503 &self,
504 addresses: &[[u8; 32]],
505 data_type: u32,
506 data_size: u64,
507 ) -> Result<MerkleBatchPaymentResult> {
508 let chunk_count = addresses.len();
509 if chunk_count < 2 {
510 return Err(Error::Payment(
511 "Merkle batch payment requires at least 2 chunks".to_string(),
512 ));
513 }
514
515 if chunk_count > MAX_LEAVES {
516 return self
517 .pay_for_merkle_multi_batch(addresses, data_type, data_size)
518 .await;
519 }
520
521 self.pay_for_merkle_single_batch(addresses, data_type, data_size)
522 .await
523 }
524
525 pub(crate) async fn plan_merkle_upload(
534 &self,
535 chunks: Vec<([u8; 32], u64)>,
536 data_type: u32,
537 progress: Option<&mpsc::Sender<UploadEvent>>,
538 ) -> Result<MerkleUploadPlan> {
539 let total_chunks = chunks.len();
540 if total_chunks == 0 {
541 return Ok(MerkleUploadPlan::default());
542 }
543
544 info!("Checking {total_chunks} merkle chunks for existing storage before payment");
545
546 let quote_limiter = self.controller().quote.clone();
547 let quote_concurrency = quote_limiter.current().min(total_chunks.max(1));
548 let mut check_stream = stream::iter(chunks.into_iter().enumerate())
549 .map(|(index, (address, data_size))| {
550 let limiter = quote_limiter.clone();
551 async move {
552 let result = observe_op(
553 &limiter,
554 || async move {
555 self.chunk_already_stored_for_merkle(&address, data_type, data_size)
556 .await
557 },
558 classify_error,
559 )
560 .await;
561 (index, address, data_size, result)
562 }
563 })
564 .buffer_unordered(quote_concurrency);
565
566 let mut already_stored: Vec<(usize, [u8; 32])> = Vec::new();
567 let mut to_upload: Vec<(usize, [u8; 32], u64)> = Vec::new();
568 let mut checked = 0usize;
569
570 while let Some((index, address, data_size, result)) = check_stream.next().await {
571 let is_already_stored = result?;
572 checked += 1;
573
574 if let Some(tx) = progress {
575 let _ = tx.try_send(UploadEvent::ChunkQuoted {
576 quoted: checked,
577 total: total_chunks,
578 });
579 }
580
581 if is_already_stored {
582 debug!(
583 "Merkle preflight {checked}/{total_chunks}: chunk {} already stored",
584 hex::encode(address)
585 );
586 already_stored.push((index, address));
587 if let Some(tx) = progress {
588 let _ = tx.try_send(UploadEvent::ChunkStored {
589 stored: already_stored.len(),
590 total: total_chunks,
591 });
592 }
593 } else {
594 debug!(
595 "Merkle preflight {checked}/{total_chunks}: chunk {} needs upload",
596 hex::encode(address)
597 );
598 to_upload.push((index, address, data_size));
599 }
600 }
601
602 already_stored.sort_by_key(|(index, _)| *index);
603 to_upload.sort_by_key(|(index, _, _)| *index);
604
605 let to_upload_total_bytes = to_upload.iter().fold(0u64, |acc, (_, _, data_size)| {
606 acc.saturating_add(*data_size)
607 });
608
609 let already_stored = already_stored
610 .into_iter()
611 .map(|(_, address)| address)
612 .collect::<Vec<_>>();
613 let to_upload = to_upload
614 .into_iter()
615 .map(|(_, address, _)| address)
616 .collect::<Vec<_>>();
617
618 info!(
619 "Merkle preflight complete: {} already stored, {} need upload",
620 already_stored.len(),
621 to_upload.len()
622 );
623
624 Ok(MerkleUploadPlan {
625 already_stored,
626 to_upload,
627 to_upload_total_bytes,
628 })
629 }
630
631 async fn chunk_already_stored_for_merkle(
632 &self,
633 address: &[u8; 32],
634 data_type: u32,
635 data_size: u64,
636 ) -> Result<bool> {
637 let result = self
638 .get_store_quotes_with_fault_tolerance(address, data_size, data_type)
639 .await;
640 if let Err(e) = &result {
641 if matches!(classify_error(e), Outcome::Timeout | Outcome::NetworkError) {
642 debug!(
643 "Merkle preflight: could not determine stored status for {} ({e}); \
644 treating as not stored and queuing for upload",
645 hex::encode(address)
646 );
647 }
648 }
649 preflight_stored_status(result)
650 }
651
652 pub async fn prepare_merkle_batch_external(
666 &self,
667 addresses: &[[u8; 32]],
668 data_type: u32,
669 data_size: u64,
670 ) -> Result<PreparedMerkleBatch> {
671 ensure_single_merkle_tree_batch(addresses.len())?;
672
673 let chunk_count = addresses.len();
674 let xornames: Vec<XorName> = addresses.iter().map(|a| XorName(*a)).collect();
675
676 debug!("Building merkle tree for {chunk_count} chunks");
677
678 let tree = MerkleTree::from_xornames(xornames)
680 .map_err(|e| Error::Payment(format!("Failed to build merkle tree: {e}")))?;
681
682 let depth = tree.depth();
683 let merkle_payment_timestamp = std::time::SystemTime::now()
684 .duration_since(std::time::UNIX_EPOCH)
685 .map_err(|e| Error::Payment(format!("System time error: {e}")))?
686 .as_secs();
687
688 debug!("Merkle tree: depth={depth}, leaves={chunk_count}, ts={merkle_payment_timestamp}");
689
690 let midpoint_proofs = tree
692 .reward_candidates(merkle_payment_timestamp)
693 .map_err(|e| Error::Payment(format!("Failed to generate reward candidates: {e}")))?;
694
695 debug!(
696 "Collecting candidate pools from {} midpoints (concurrent)",
697 midpoint_proofs.len()
698 );
699
700 let candidate_pools = self
706 .build_candidate_pools(
707 &midpoint_proofs,
708 data_type,
709 data_size,
710 merkle_payment_timestamp,
711 )
712 .await?;
713
714 let pool_commitments: Vec<PoolCommitment> = candidate_pools
719 .iter()
720 .map(pool_commitment_with_payment_multiplier)
721 .collect::<Result<Vec<_>>>()?;
722
723 Ok(PreparedMerkleBatch {
724 depth,
725 pool_commitments,
726 merkle_payment_timestamp,
727 candidate_pools,
728 tree,
729 addresses: addresses.to_vec(),
730 })
731 }
732
733 async fn pay_for_merkle_single_batch(
735 &self,
736 addresses: &[[u8; 32]],
737 data_type: u32,
738 data_size: u64,
739 ) -> Result<MerkleBatchPaymentResult> {
740 let wallet = self.require_wallet()?;
741 let prepared = self
742 .prepare_merkle_batch_external(addresses, data_type, data_size)
743 .await?;
744
745 info!(
746 "Submitting merkle batch payment on-chain (depth={})",
747 prepared.depth
748 );
749 let (winner_pool_hash, amount, gas_info) = wallet
750 .pay_for_merkle_tree(
751 prepared.depth,
752 prepared.pool_commitments.clone(),
753 prepared.merkle_payment_timestamp,
754 )
755 .await
756 .map_err(|e| Error::Payment(format!("Merkle batch payment failed: {e}")))?;
757
758 info!(
759 "Merkle payment succeeded: winner pool {}",
760 hex::encode(winner_pool_hash)
761 );
762
763 let mut result = finalize_merkle_batch(prepared, winner_pool_hash)?;
764 result.storage_cost_atto = amount.to_string();
765 result.gas_cost_wei = gas_info.gas_cost_wei;
766 Ok(result)
767 }
768
769 async fn pay_for_merkle_multi_batch(
771 &self,
772 addresses: &[[u8; 32]],
773 data_type: u32,
774 data_size: u64,
775 ) -> Result<MerkleBatchPaymentResult> {
776 let sub_batches = merkle_batch_partitions(addresses);
781 let total_sub_batches = sub_batches.len();
782 let mut all_proofs = HashMap::with_capacity(addresses.len());
783 let mut total_storage = Amount::ZERO;
784 let mut total_gas: u128 = 0;
785 let mut oldest_ts: u64 = 0;
789
790 for (i, chunk) in sub_batches.into_iter().enumerate() {
791 match self
792 .pay_for_merkle_single_batch(chunk, data_type, data_size)
793 .await
794 {
795 Ok(sub_result) => {
796 if let Ok(cost) = sub_result.storage_cost_atto.parse::<Amount>() {
797 total_storage += cost;
798 }
799 total_gas = total_gas.saturating_add(sub_result.gas_cost_wei);
800 if oldest_ts == 0
801 || (sub_result.merkle_payment_timestamp > 0
802 && sub_result.merkle_payment_timestamp < oldest_ts)
803 {
804 oldest_ts = sub_result.merkle_payment_timestamp;
805 }
806 all_proofs.extend(sub_result.proofs);
807 }
808 Err(e) => {
809 if all_proofs.is_empty() {
810 return Err(e);
812 }
813 warn!(
815 "Merkle sub-batch {}/{total_sub_batches} failed: {e}. \
816 Returning {} proofs from prior sub-batches",
817 i + 1,
818 all_proofs.len()
819 );
820 return Ok(MerkleBatchPaymentResult {
821 chunk_count: all_proofs.len(),
822 proofs: all_proofs,
823 storage_cost_atto: total_storage.to_string(),
824 gas_cost_wei: total_gas,
825 merkle_payment_timestamp: oldest_ts,
826 });
827 }
828 }
829 }
830
831 Ok(MerkleBatchPaymentResult {
832 chunk_count: addresses.len(),
833 proofs: all_proofs,
834 storage_cost_atto: total_storage.to_string(),
835 gas_cost_wei: total_gas,
836 merkle_payment_timestamp: oldest_ts,
837 })
838 }
839
840 async fn build_candidate_pools(
842 &self,
843 midpoint_proofs: &[MidpointProof],
844 data_type: u32,
845 data_size: u64,
846 merkle_payment_timestamp: u64,
847 ) -> Result<Vec<MerklePaymentCandidatePool>> {
848 let mut pool_futures = FuturesUnordered::new();
849
850 for midpoint_proof in midpoint_proofs {
851 let pool_address = midpoint_proof.address();
852 let mp = midpoint_proof.clone();
853 pool_futures.push(async move {
854 let candidate_nodes = self
855 .get_merkle_candidate_pool(
856 &pool_address.0,
857 data_type,
858 data_size,
859 merkle_payment_timestamp,
860 )
861 .await?;
862 Ok::<_, Error>(MerklePaymentCandidatePool {
863 midpoint_proof: mp,
864 candidate_nodes,
865 })
866 });
867 }
868
869 let mut pools = Vec::with_capacity(midpoint_proofs.len());
870 while let Some(result) = pool_futures.next().await {
871 pools.push(result?);
872 }
873
874 Ok(pools)
875 }
876
877 #[allow(clippy::too_many_lines)]
879 async fn get_merkle_candidate_pool(
880 &self,
881 address: &[u8; 32],
882 data_type: u32,
883 data_size: u64,
884 merkle_payment_timestamp: u64,
885 ) -> Result<[MerklePaymentCandidateNode; CANDIDATES_PER_POOL]> {
886 let node = self.network().node();
887 let timeout = Duration::from_secs(self.config().quote_timeout_secs);
888
889 let query_count = CANDIDATES_PER_POOL * 2;
891 let mut remote_peers = self
892 .network()
893 .find_closest_peers(address, query_count)
894 .await?;
895
896 if remote_peers.len() < CANDIDATES_PER_POOL {
900 let connected = self.network().connected_peers().await;
901 for peer in connected {
902 if !remote_peers.iter().any(|(id, _)| *id == peer) {
903 remote_peers.push((peer, vec![]));
904 }
905 }
906 }
907
908 if remote_peers.len() < CANDIDATES_PER_POOL {
909 return Err(Error::InsufficientPeers(format!(
910 "Found {} peers, need {CANDIDATES_PER_POOL} for merkle candidate pool. \
911 Use --no-merkle or a larger network.",
912 remote_peers.len()
913 )));
914 }
915
916 let mut candidate_futures = FuturesUnordered::new();
917
918 for (peer_id, peer_addrs) in &remote_peers {
919 let request_id = self.next_request_id();
920 let request = MerkleCandidateQuoteRequest {
921 address: *address,
922 data_type,
923 data_size,
924 merkle_payment_timestamp,
925 };
926 let message = ChunkMessage {
927 request_id,
928 body: ChunkMessageBody::MerkleCandidateQuoteRequest(request),
929 };
930
931 let message_bytes = match message.encode() {
932 Ok(bytes) => bytes,
933 Err(e) => {
934 warn!("Failed to encode merkle candidate request for {peer_id}: {e}");
935 continue;
936 }
937 };
938
939 let peer_id_clone = *peer_id;
940 let addrs_clone = peer_addrs.clone();
941 let node_clone = node.clone();
942
943 let fut = async move {
944 let result = send_and_await_chunk_response(
945 &node_clone,
946 &peer_id_clone,
947 message_bytes,
948 request_id,
949 timeout,
950 &addrs_clone,
951 |body| match body {
952 ChunkMessageBody::MerkleCandidateQuoteResponse(
953 MerkleCandidateQuoteResponse::Success {
954 candidate_node,
955 commitment,
956 },
957 ) => {
958 match rmp_serde::from_slice::<MerklePaymentCandidateNode>(
959 &candidate_node,
960 ) {
961 Ok(node) => Some(Ok((node, commitment))),
962 Err(e) => Some(Err(Error::Serialization(format!(
963 "Failed to deserialize candidate node from {peer_id_clone}: {e}"
964 )))),
965 }
966 }
967 ChunkMessageBody::MerkleCandidateQuoteResponse(
968 MerkleCandidateQuoteResponse::Error(e),
969 ) => Some(Err(Error::Protocol(format!(
970 "Merkle quote error from {peer_id_clone}: {e}"
971 )))),
972 _ => None,
973 },
974 |e| {
975 Error::Network(format!(
976 "Failed to send merkle candidate request to {peer_id_clone}: {e}"
977 ))
978 },
979 || {
980 Error::Timeout(format!(
981 "Timeout waiting for merkle candidate from {peer_id_clone}"
982 ))
983 },
984 )
985 .await;
986
987 (peer_id_clone, result)
988 };
989
990 candidate_futures.push(fut);
991 }
992
993 self.collect_validated_candidates(&mut candidate_futures, address, merkle_payment_timestamp)
994 .await
995 }
996
997 async fn collect_validated_candidates(
1008 &self,
1009 futures: &mut FuturesUnordered<
1010 impl std::future::Future<
1011 Output = (
1012 PeerId,
1013 std::result::Result<(MerklePaymentCandidateNode, Option<Vec<u8>>), Error>,
1014 ),
1015 >,
1016 >,
1017 target_address: &[u8; 32],
1018 merkle_payment_timestamp: u64,
1019 ) -> Result<[MerklePaymentCandidateNode; CANDIDATES_PER_POOL]> {
1020 let mut valid: Vec<(PeerId, MerklePaymentCandidateNode)> = Vec::new();
1021 let mut failures: Vec<String> = Vec::new();
1022
1023 while let Some((peer_id, result)) = futures.next().await {
1024 match result {
1025 Ok((candidate, commitment)) => {
1026 if !verify_merkle_candidate_signature(&candidate) {
1027 warn!("Invalid ML-DSA-65 signature from merkle candidate {peer_id}");
1028 failures.push(format!("{peer_id}: invalid signature"));
1029 continue;
1030 }
1031 if candidate.merkle_payment_timestamp != merkle_payment_timestamp {
1032 warn!("Timestamp mismatch from merkle candidate {peer_id}");
1033 failures.push(format!("{peer_id}: timestamp mismatch"));
1034 continue;
1035 }
1036 let candidate_peer = PeerId::from_bytes(compute_address(&candidate.pub_key));
1041 if candidate_peer != peer_id {
1042 warn!(
1043 "Dropping merkle candidate {peer_id} — pub_key derives {candidate_peer}, \
1044 not the responding peer"
1045 );
1046 failures.push(format!("{peer_id}: candidate pub_key/peer mismatch"));
1047 continue;
1048 }
1049 if let Err(detail) =
1058 merkle_candidate_binding_is_valid(&candidate_peer, &candidate, &commitment)
1059 {
1060 warn!("Dropping merkle candidate {peer_id} — ADR-0004 binding invalid: {detail}");
1061 failures.push(format!("{peer_id}: bad commitment binding ({detail})"));
1062 continue;
1063 }
1064 valid.push((candidate_peer, candidate));
1065 }
1066 Err(e) => {
1067 debug!("Failed to get merkle candidate from {peer_id}: {e}");
1068 failures.push(format!("{peer_id}: {e}"));
1069 }
1070 }
1071 }
1072
1073 if valid.len() < CANDIDATES_PER_POOL {
1074 return Err(Error::InsufficientPeers(format!(
1075 "Got {} merkle candidates, need {CANDIDATES_PER_POOL}. Failures: [{}]",
1076 valid.len(),
1077 failures.join("; ")
1078 )));
1079 }
1080
1081 let target_peer = PeerId::from_bytes(*target_address);
1082 valid.sort_by_key(|(peer_id, _)| peer_id.xor_distance(&target_peer));
1083
1084 let candidates: Vec<MerklePaymentCandidateNode> = valid
1085 .into_iter()
1086 .take(CANDIDATES_PER_POOL)
1087 .map(|(_, candidate)| candidate)
1088 .collect();
1089
1090 let array: [MerklePaymentCandidateNode; CANDIDATES_PER_POOL] =
1091 candidates.try_into().map_err(|_| {
1092 Error::Payment("Failed to convert candidates to fixed array".to_string())
1093 })?;
1094 Ok(array)
1095 }
1096
1097 pub(crate) async fn merkle_upload_chunks(
1118 &self,
1119 chunk_contents: Vec<Bytes>,
1120 addresses: Vec<[u8; 32]>,
1121 batch_result: &MerkleBatchPaymentResult,
1122 progress: Option<&mpsc::Sender<UploadEvent>>,
1123 stored_offset: usize,
1124 total_chunks: usize,
1125 ) -> Result<MerkleStoreOutcome> {
1126 let store_limiter = self.controller().store.clone();
1127 let batch_size = chunk_contents.len();
1130 if batch_size != addresses.len() {
1131 return Err(Error::InvalidData(format!(
1132 "merkle upload has {batch_size} chunk contents but {} addresses",
1133 addresses.len()
1134 )));
1135 }
1136 let cap = || store_limiter.current().min(batch_size.max(1));
1140
1141 let bodies: std::collections::HashMap<[u8; 32], Bytes> =
1146 addresses.iter().copied().zip(chunk_contents).collect();
1147 let addrs = addresses;
1148
1149 let store_one = |addr: [u8; 32]| {
1154 let limiter = store_limiter.clone();
1155 let content = bodies.get(&addr).cloned();
1156 let proof_bytes = batch_result.proofs.get(&addr).cloned();
1157 async move {
1158 let started = std::time::Instant::now();
1159 let content = content.ok_or_else(|| {
1160 Error::InvalidData(format!("missing chunk body for {}", hex::encode(addr)))
1161 })?;
1162 let proof = proof_bytes.ok_or_else(|| {
1163 Error::Payment(format!(
1164 "Missing merkle proof for chunk {}",
1165 hex::encode(addr)
1166 ))
1167 })?;
1168 let peers = self.put_target_peers(&addr).await?;
1169 observe_op(
1170 &limiter,
1171 || async move { self.chunk_put_to_close_group(content, proof, &peers).await },
1172 classify_error,
1173 )
1174 .await
1175 .map(|_| started)
1176 }
1177 };
1178
1179 let outcome = merkle_store_with_retry(
1180 addrs,
1181 cap,
1182 MERKLE_STORE_MAX_ATTEMPTS,
1183 MERKLE_RETRY_BACKOFF,
1184 progress,
1185 stored_offset,
1186 total_chunks,
1187 store_one,
1188 )
1189 .await?;
1190
1191 if let Some(e) = outcome.fatal {
1197 return Err(e);
1198 }
1199 Ok(outcome)
1200 }
1201}
1202
1203pub(crate) const MERKLE_STORE_MAX_ATTEMPTS: usize = 4;
1215
1216pub(crate) const MERKLE_RETRY_BACKOFF: Duration = Duration::from_secs(30);
1222
1223const MERKLE_RETRY_JITTER: f64 = 0.1;
1226
1227#[derive(Debug, Default)]
1230pub(crate) struct MerkleStoreOutcome {
1231 pub stored: usize,
1234 pub stored_addresses: Vec<[u8; 32]>,
1241 pub failed: usize,
1243 pub failed_addresses: Vec<([u8; 32], String)>,
1248 pub fatal: Option<Error>,
1255 pub stats: crate::data::client::batch::WaveAggregateStats,
1257}
1258
1259#[allow(clippy::too_many_arguments)]
1285pub(crate) async fn merkle_store_with_retry<F, Fut, C>(
1286 addrs: Vec<[u8; 32]>,
1287 cap: C,
1288 max_attempts: usize,
1289 backoff: Duration,
1290 progress: Option<&mpsc::Sender<UploadEvent>>,
1291 stored_offset: usize,
1292 total: usize,
1293 store_one: F,
1294) -> Result<MerkleStoreOutcome>
1295where
1296 F: Fn([u8; 32]) -> Fut,
1297 Fut: std::future::Future<Output = Result<std::time::Instant>>,
1298 C: Fn() -> usize,
1299{
1300 let attempts = max_attempts.max(1);
1301 let mut outcome = MerkleStoreOutcome {
1302 stored: stored_offset,
1303 ..MerkleStoreOutcome::default()
1304 };
1305 let mut pending = addrs;
1306
1307 for attempt in 0..attempts {
1308 let mut next_failed: Vec<([u8; 32], String)> = Vec::new();
1314
1315 let mut pending_iter = pending.into_iter();
1320 let mut in_flight = FuturesUnordered::new();
1321 loop {
1322 let slots = cap().max(1);
1323 while in_flight.len() < slots {
1324 match pending_iter.next() {
1325 Some(addr) => {
1326 let fut = store_one(addr);
1327 in_flight.push(async move { (addr, fut.await) });
1328 }
1329 None => break,
1330 }
1331 }
1332 let Some((addr, result)) = in_flight.next().await else {
1333 break;
1334 };
1335 outcome.stats.chunk_attempts_total =
1336 outcome.stats.chunk_attempts_total.saturating_add(1);
1337 match result {
1338 Ok(started) => {
1339 let duration_ms =
1340 u64::try_from(started.elapsed().as_millis()).unwrap_or(u64::MAX);
1341 outcome.stats.store_durations_ms.push(duration_ms);
1342 let idx = attempt.min(outcome.stats.retries_histogram.len().saturating_sub(1));
1343 outcome.stats.retries_histogram[idx] =
1344 outcome.stats.retries_histogram[idx].saturating_add(1);
1345 outcome.stored += 1;
1346 outcome.stored_addresses.push(addr);
1347 if let Some(tx) = progress {
1348 let _ = tx.try_send(UploadEvent::ChunkStored {
1349 stored: outcome.stored,
1350 total,
1351 });
1352 }
1353 }
1354 Err(
1361 e @ (Error::InsufficientPeers(_)
1362 | Error::CloseGroupShortfall(_)
1363 | Error::RemotePut { .. }),
1364 ) => {
1365 next_failed.push((addr, e.to_string()));
1366 }
1367 Err(e) => {
1368 next_failed.push((addr, e.to_string()));
1376 outcome.fatal = Some(e);
1377 break;
1378 }
1379 }
1380 }
1381
1382 if outcome.fatal.is_some() {
1383 outcome.failed = next_failed.len();
1384 outcome.failed_addresses = next_failed;
1385 return Ok(outcome);
1386 }
1387
1388 if next_failed.is_empty() {
1389 break;
1390 }
1391
1392 if attempt + 1 < attempts {
1393 warn!(
1394 failed = next_failed.len(),
1395 attempt = attempt + 1,
1396 "merkle chunks short of quorum, retrying after backoff"
1397 );
1398 pending = next_failed.into_iter().map(|(addr, _msg)| addr).collect();
1399 if backoff > Duration::ZERO {
1400 let wait = {
1405 let mut rng = rand::thread_rng();
1406 let factor = 1.0 + rng.gen_range(-MERKLE_RETRY_JITTER..=MERKLE_RETRY_JITTER);
1407 backoff.mul_f64(factor)
1408 };
1409 tokio::time::sleep(wait).await;
1410 }
1411 } else {
1412 outcome.failed = next_failed.len();
1413 outcome.failed_addresses = next_failed;
1414 break;
1415 }
1416 }
1417
1418 Ok(outcome)
1419}
1420
1421pub(crate) const DEFERRED_ROUND_DELAYS_SECS: [u64; 3] = [0, 15, 45];
1430
1431pub(crate) fn deferred_round_histogram_slot(round: usize, hist_len: usize) -> usize {
1438 (round + 1).min(hist_len.saturating_sub(1))
1439}
1440
1441#[derive(Debug, Default)]
1443pub(crate) struct DeferredRetryOutcome {
1444 pub stored: usize,
1448 pub stored_addresses: Vec<[u8; 32]>,
1451 pub failed: usize,
1453 pub failed_addresses: Vec<([u8; 32], String)>,
1457 pub fatal: Option<String>,
1461 pub stats: crate::data::client::batch::WaveAggregateStats,
1464}
1465
1466#[allow(clippy::too_many_arguments)]
1484pub(crate) async fn merkle_deferred_retry<CF, SF, Fut>(
1485 deferred: Vec<([u8; 32], String)>,
1486 round_delays_secs: &[u64],
1487 concurrency_for: CF,
1488 progress: Option<&mpsc::Sender<UploadEvent>>,
1489 stored_offset: usize,
1490 total: usize,
1491 store_one: SF,
1492) -> Result<DeferredRetryOutcome>
1493where
1494 CF: Fn(usize) -> usize,
1495 SF: Fn([u8; 32]) -> Fut,
1496 Fut: std::future::Future<Output = Result<std::time::Instant>>,
1497{
1498 let mut outcome = DeferredRetryOutcome {
1499 stored: stored_offset,
1500 ..DeferredRetryOutcome::default()
1501 };
1502 let mut remaining = deferred;
1503 let rounds = round_delays_secs.len();
1504
1505 for (round, &delay_secs) in round_delays_secs.iter().enumerate() {
1506 if remaining.is_empty() {
1507 break;
1508 }
1509 if delay_secs > 0 {
1510 tokio::time::sleep(Duration::from_secs(delay_secs)).await;
1511 }
1512 info!(
1513 "Deferred merkle retry round {}/{}: {} chunk(s) short of quorum",
1514 round + 1,
1515 rounds,
1516 remaining.len(),
1517 );
1518
1519 let slot = deferred_round_histogram_slot(round, outcome.stats.retries_histogram.len());
1523 let round_addrs: Vec<[u8; 32]> = std::mem::take(&mut remaining)
1524 .into_iter()
1525 .map(|(addr, _msg)| addr)
1526 .collect();
1527 let round_len = round_addrs.len();
1528 let cap = || concurrency_for(round_len);
1531
1532 let round_outcome = merkle_store_with_retry(
1533 round_addrs,
1534 cap,
1535 1,
1536 Duration::ZERO,
1537 progress,
1538 outcome.stored,
1539 total,
1540 &store_one,
1541 )
1542 .await?;
1543
1544 outcome.stored = round_outcome.stored;
1545 outcome
1546 .stored_addresses
1547 .extend(round_outcome.stored_addresses);
1548
1549 outcome.stats.chunk_attempts_total = outcome
1551 .stats
1552 .chunk_attempts_total
1553 .saturating_add(round_outcome.stats.chunk_attempts_total);
1554 outcome
1555 .stats
1556 .store_durations_ms
1557 .extend(round_outcome.stats.store_durations_ms);
1558 let landed: usize = round_outcome.stats.retries_histogram.iter().sum();
1559 outcome.stats.retries_histogram[slot] =
1560 outcome.stats.retries_histogram[slot].saturating_add(landed);
1561
1562 if let Some(fatal) = round_outcome.fatal {
1563 outcome.fatal = Some(fatal.to_string());
1568 outcome.failed = round_outcome.failed_addresses.len();
1569 outcome.failed_addresses = round_outcome.failed_addresses;
1570 return Ok(outcome);
1571 }
1572
1573 remaining = round_outcome.failed_addresses;
1575 }
1576
1577 outcome.failed = remaining.len();
1578 outcome.failed_addresses = remaining;
1579 Ok(outcome)
1580}
1581
1582pub fn finalize_merkle_batch(
1587 prepared: PreparedMerkleBatch,
1588 winner_pool_hash: [u8; 32],
1589) -> Result<MerkleBatchPaymentResult> {
1590 let chunk_count = prepared.addresses.len();
1591 let xornames: Vec<XorName> = prepared.addresses.iter().map(|a| XorName(*a)).collect();
1592
1593 let winner_pool = prepared
1595 .candidate_pools
1596 .iter()
1597 .find(|pool| pool.hash() == winner_pool_hash)
1598 .ok_or_else(|| {
1599 Error::Payment(format!(
1600 "Winner pool {} not found in candidate pools",
1601 hex::encode(winner_pool_hash)
1602 ))
1603 })?;
1604
1605 info!("Generating merkle proofs for {chunk_count} chunks");
1616 let mut proofs = HashMap::with_capacity(chunk_count);
1617
1618 for (i, xorname) in xornames.iter().enumerate() {
1619 let address_proof = prepared
1620 .tree
1621 .generate_address_proof(i, *xorname)
1622 .map_err(|e| {
1623 Error::Payment(format!(
1624 "Failed to generate address proof for chunk {i}: {e}"
1625 ))
1626 })?;
1627
1628 let merkle_proof = MerklePaymentProof::new(*xorname, address_proof, winner_pool.clone());
1629
1630 let tagged_bytes = serialize_merkle_proof(&merkle_proof)
1631 .map_err(|e| Error::Serialization(format!("Failed to serialize merkle proof: {e}")))?;
1632
1633 proofs.insert(prepared.addresses[i], tagged_bytes);
1634 }
1635
1636 info!("Merkle batch payment complete: {chunk_count} proofs generated");
1637
1638 Ok(MerkleBatchPaymentResult {
1639 proofs,
1640 chunk_count,
1641 storage_cost_atto: "0".to_string(),
1642 gas_cost_wei: 0,
1643 merkle_payment_timestamp: prepared.merkle_payment_timestamp,
1644 })
1645}
1646
1647#[cfg(test)]
1649mod send_assertions {
1650 use super::*;
1651 use crate::data::client::Client;
1652
1653 fn _assert_send<T: Send>(_: &T) {}
1654
1655 #[allow(
1656 dead_code,
1657 unreachable_code,
1658 unused_variables,
1659 clippy::diverging_sub_expression
1660 )]
1661 async fn _merkle_upload_chunks_is_send(client: &Client) {
1662 let batch_result: MerkleBatchPaymentResult = todo!();
1663 let fut = client.merkle_upload_chunks(Vec::new(), Vec::new(), &batch_result, None, 0, 0);
1664 _assert_send(&fut);
1665 }
1666}
1667
1668#[cfg(test)]
1669#[allow(clippy::unwrap_used, clippy::expect_used, clippy::panic)]
1670mod tests {
1671 use super::*;
1672 use ant_protocol::evm::{Amount, MerkleTree, RewardsAddress, CANDIDATES_PER_POOL};
1673
1674 #[test]
1679 fn test_auto_below_threshold() {
1680 assert!(!should_use_merkle(1, PaymentMode::Auto));
1681 assert!(!should_use_merkle(10, PaymentMode::Auto));
1682 assert!(!should_use_merkle(63, PaymentMode::Auto));
1683 }
1684
1685 #[test]
1686 fn test_auto_at_and_above_threshold() {
1687 assert!(should_use_merkle(64, PaymentMode::Auto));
1688 assert!(should_use_merkle(65, PaymentMode::Auto));
1689 assert!(should_use_merkle(1000, PaymentMode::Auto));
1690 }
1691
1692 #[test]
1693 fn test_merkle_mode_forces_at_2() {
1694 assert!(!should_use_merkle(1, PaymentMode::Merkle));
1695 assert!(should_use_merkle(2, PaymentMode::Merkle));
1696 assert!(should_use_merkle(3, PaymentMode::Merkle));
1697 }
1698
1699 #[test]
1700 fn test_single_mode_always_false() {
1701 assert!(!should_use_merkle(0, PaymentMode::Single));
1702 assert!(!should_use_merkle(64, PaymentMode::Single));
1703 assert!(!should_use_merkle(1000, PaymentMode::Single));
1704 }
1705
1706 #[test]
1707 fn test_default_mode_is_auto() {
1708 assert_eq!(PaymentMode::default(), PaymentMode::Auto);
1709 }
1710
1711 #[test]
1712 fn test_threshold_value() {
1713 assert_eq!(DEFAULT_MERKLE_THRESHOLD, 64);
1714 }
1715
1716 #[test]
1721 fn test_preflight_quotes_gathered_means_not_stored() {
1722 assert!(matches!(preflight_stored_status(Ok(())), Ok(false)));
1723 }
1724
1725 #[test]
1726 fn test_preflight_already_stored_is_stored() {
1727 let r: Result<()> = Err(Error::AlreadyStored);
1728 assert!(matches!(preflight_stored_status(r), Ok(true)));
1729 }
1730
1731 #[test]
1735 fn test_preflight_transient_quote_failure_does_not_abort() {
1736 let insufficient: Result<()> =
1738 Err(Error::InsufficientPeers("Got 5 quotes, need 7".to_string()));
1739 assert!(
1740 matches!(preflight_stored_status(insufficient), Ok(false)),
1741 "insufficient-peers during preflight must degrade to not-stored, not error"
1742 );
1743
1744 let timeout: Result<()> = Err(Error::Timeout("Timeout waiting for quote".to_string()));
1745 assert!(matches!(preflight_stored_status(timeout), Ok(false)));
1746
1747 let network: Result<()> = Err(Error::Network("connection reset".to_string()));
1748 assert!(matches!(preflight_stored_status(network), Ok(false)));
1749 }
1750
1751 #[test]
1754 fn test_preflight_application_error_propagates() {
1755 let payment: Result<()> = Err(Error::Payment("bad payment".to_string()));
1756 assert!(matches!(
1757 preflight_stored_status(payment),
1758 Err(Error::Payment(_))
1759 ));
1760 }
1761
1762 #[test]
1763 fn chunk_contents_for_upload_addresses_preserves_requested_order() {
1764 let first = Bytes::from_static(b"first");
1765 let second = Bytes::from_static(b"second");
1766 let first_addr = compute_address(&first);
1767 let second_addr = compute_address(&second);
1768
1769 let selected = chunk_contents_for_upload_addresses(
1770 vec![first.clone(), second.clone()],
1771 &[second_addr, first_addr],
1772 )
1773 .unwrap();
1774
1775 assert_eq!(selected, vec![second, first]);
1776 }
1777
1778 #[test]
1779 fn chunk_contents_for_upload_addresses_preserves_duplicate_requests() {
1780 let repeated = Bytes::from_static(b"same-content");
1781 let other = Bytes::from_static(b"other-content");
1782 let repeated_addr = compute_address(&repeated);
1783
1784 let selected = chunk_contents_for_upload_addresses(
1785 vec![repeated.clone(), other, repeated.clone()],
1786 &[repeated_addr, repeated_addr],
1787 )
1788 .unwrap();
1789
1790 assert_eq!(selected, vec![repeated.clone(), repeated]);
1791 }
1792
1793 #[test]
1794 fn chunk_contents_for_upload_addresses_ignores_unrequested_duplicates() {
1795 let requested = Bytes::from_static(b"requested-content");
1796 let unrequested = Bytes::from_static(b"unrequested-content");
1797 let requested_addr = compute_address(&requested);
1798
1799 let selected = chunk_contents_for_upload_addresses(
1800 vec![
1801 unrequested.clone(),
1802 requested.clone(),
1803 unrequested.clone(),
1804 unrequested,
1805 ],
1806 &[requested_addr],
1807 )
1808 .unwrap();
1809
1810 assert_eq!(selected, vec![requested]);
1811 }
1812
1813 #[test]
1814 fn chunk_contents_for_upload_addresses_errors_for_missing_content() {
1815 let present = Bytes::from_static(b"present-content");
1816 let missing = Bytes::from_static(b"missing-content");
1817 let missing_addr = compute_address(&missing);
1818
1819 let result = chunk_contents_for_upload_addresses(vec![present], &[missing_addr]);
1820
1821 assert!(matches!(result, Err(Error::InvalidData(_))));
1822 }
1823
1824 fn make_test_addresses(count: usize) -> Vec<[u8; 32]> {
1829 (0..count)
1830 .map(|i| {
1831 let xn = XorName::from_content(&i.to_le_bytes());
1832 xn.0
1833 })
1834 .collect()
1835 }
1836
1837 #[test]
1838 fn test_tree_depth_for_known_sizes() {
1839 let cases = [(2, 1), (4, 2), (16, 4), (100, 7), (256, 8)];
1840 for (count, expected_depth) in cases {
1841 let addrs = make_test_addresses(count);
1842 let xornames: Vec<XorName> = addrs.iter().map(|a| XorName(*a)).collect();
1843 let tree = MerkleTree::from_xornames(xornames).unwrap();
1844 assert_eq!(
1845 tree.depth(),
1846 expected_depth,
1847 "depth mismatch for {count} leaves"
1848 );
1849 }
1850 }
1851
1852 #[test]
1853 fn test_proof_generation_and_verification_for_all_leaves() {
1854 let addrs = make_test_addresses(16);
1855 let xornames: Vec<XorName> = addrs.iter().map(|a| XorName(*a)).collect();
1856 let tree = MerkleTree::from_xornames(xornames.clone()).unwrap();
1857
1858 for (i, xn) in xornames.iter().enumerate() {
1859 let proof = tree.generate_address_proof(i, *xn).unwrap();
1860 assert!(proof.verify(), "proof for leaf {i} should verify");
1861 assert_eq!(proof.depth(), tree.depth() as usize);
1862 }
1863 }
1864
1865 #[test]
1866 fn test_proof_fails_for_wrong_address() {
1867 let addrs = make_test_addresses(8);
1868 let xornames: Vec<XorName> = addrs.iter().map(|a| XorName(*a)).collect();
1869 let tree = MerkleTree::from_xornames(xornames).unwrap();
1870
1871 let wrong = XorName::from_content(b"wrong");
1872 let proof = tree.generate_address_proof(0, wrong).unwrap();
1873 assert!(!proof.verify(), "proof with wrong address should fail");
1874 }
1875
1876 #[test]
1877 fn test_tree_too_few_leaves() {
1878 let xornames = vec![XorName::from_content(b"only_one")];
1879 let result = MerkleTree::from_xornames(xornames);
1880 assert!(result.is_err());
1881 }
1882
1883 #[test]
1884 fn test_tree_at_max_leaves() {
1885 let addrs = make_test_addresses(MAX_LEAVES);
1886 let xornames: Vec<XorName> = addrs.iter().map(|a| XorName(*a)).collect();
1887 let tree = MerkleTree::from_xornames(xornames).unwrap();
1888 assert_eq!(tree.leaf_count(), MAX_LEAVES);
1889 }
1890
1891 #[test]
1896 fn test_merkle_proof_serialize_deserialize_roundtrip() {
1897 use ant_protocol::evm::{Amount, MerklePaymentCandidateNode, RewardsAddress};
1898 use ant_protocol::payment::{deserialize_merkle_proof, serialize_merkle_proof};
1899
1900 let addrs = make_test_addresses(4);
1901 let xornames: Vec<XorName> = addrs.iter().map(|a| XorName(*a)).collect();
1902 let tree = MerkleTree::from_xornames(xornames.clone()).unwrap();
1903
1904 let timestamp = std::time::SystemTime::now()
1905 .duration_since(std::time::UNIX_EPOCH)
1906 .unwrap()
1907 .as_secs();
1908
1909 let candidates = tree.reward_candidates(timestamp).unwrap();
1910 let midpoint = candidates.first().unwrap().clone();
1911
1912 #[allow(clippy::cast_possible_truncation)]
1914 let candidate_nodes: [MerklePaymentCandidateNode; CANDIDATES_PER_POOL] =
1915 std::array::from_fn(|i| MerklePaymentCandidateNode {
1916 pub_key: vec![i as u8; 32],
1917 price: Amount::from(1024u64),
1918 reward_address: RewardsAddress::new([i as u8; 20]),
1919 merkle_payment_timestamp: timestamp,
1920 signature: vec![i as u8; 64],
1921 committed_key_count: 0,
1922 commitment_pin: None,
1923 });
1924
1925 let pool = MerklePaymentCandidatePool {
1926 midpoint_proof: midpoint,
1927 candidate_nodes,
1928 };
1929
1930 let address_proof = tree.generate_address_proof(0, xornames[0]).unwrap();
1931 let merkle_proof = MerklePaymentProof::new(xornames[0], address_proof, pool);
1932
1933 let tagged = serialize_merkle_proof(&merkle_proof).unwrap();
1934 assert_eq!(
1935 tagged.first().copied(),
1936 Some(0x02),
1937 "tag should be PROOF_TAG_MERKLE"
1938 );
1939
1940 let deserialized = deserialize_merkle_proof(&tagged).unwrap();
1941 assert_eq!(deserialized.address, merkle_proof.address);
1942 assert_eq!(
1943 deserialized.winner_pool.candidate_nodes.len(),
1944 CANDIDATES_PER_POOL
1945 );
1946 }
1947
1948 #[test]
1953 fn test_candidate_wrong_timestamp_rejected() {
1954 let candidate = MerklePaymentCandidateNode {
1956 pub_key: vec![0u8; 32],
1957 price: ant_protocol::evm::Amount::ZERO,
1958 reward_address: ant_protocol::evm::RewardsAddress::new([0u8; 20]),
1959 merkle_payment_timestamp: 1000,
1960 signature: vec![0u8; 64],
1961 committed_key_count: 0,
1962 commitment_pin: None,
1963 };
1964
1965 assert_ne!(candidate.merkle_payment_timestamp, 2000);
1967 }
1968
1969 fn make_dummy_candidate_nodes(
1974 timestamp: u64,
1975 ) -> [MerklePaymentCandidateNode; CANDIDATES_PER_POOL] {
1976 std::array::from_fn(|i| MerklePaymentCandidateNode {
1977 pub_key: vec![i as u8; 32],
1978 price: Amount::from(1024u64),
1979 reward_address: RewardsAddress::new([i as u8; 20]),
1980 merkle_payment_timestamp: timestamp,
1981 signature: vec![i as u8; 64],
1982 committed_key_count: 0,
1983 commitment_pin: None,
1984 })
1985 }
1986
1987 fn make_prepared_merkle_batch(count: usize) -> PreparedMerkleBatch {
1988 let addrs = make_test_addresses(count);
1989 let xornames: Vec<XorName> = addrs.iter().map(|a| XorName(*a)).collect();
1990 let tree = MerkleTree::from_xornames(xornames).unwrap();
1991
1992 let timestamp = std::time::SystemTime::now()
1993 .duration_since(std::time::UNIX_EPOCH)
1994 .unwrap()
1995 .as_secs();
1996
1997 let midpoints = tree.reward_candidates(timestamp).unwrap();
1998
1999 let candidate_pools: Vec<MerklePaymentCandidatePool> = midpoints
2000 .into_iter()
2001 .map(|mp| MerklePaymentCandidatePool {
2002 midpoint_proof: mp,
2003 candidate_nodes: make_dummy_candidate_nodes(timestamp),
2004 })
2005 .collect();
2006
2007 let pool_commitments = candidate_pools
2008 .iter()
2009 .map(pool_commitment_with_payment_multiplier)
2010 .collect::<Result<Vec<_>>>()
2011 .unwrap();
2012
2013 PreparedMerkleBatch {
2014 depth: tree.depth(),
2015 pool_commitments,
2016 merkle_payment_timestamp: timestamp,
2017 candidate_pools,
2018 tree,
2019 addresses: addrs,
2020 }
2021 }
2022
2023 fn pool_with_varied_prices(timestamp: u64) -> MerklePaymentCandidatePool {
2026 let addrs = make_test_addresses(4);
2027 let xornames: Vec<XorName> = addrs.iter().map(|a| XorName(*a)).collect();
2028 let tree = MerkleTree::from_xornames(xornames).unwrap();
2029 let midpoint = tree
2030 .reward_candidates(timestamp)
2031 .unwrap()
2032 .into_iter()
2033 .next()
2034 .unwrap();
2035
2036 let candidate_nodes = std::array::from_fn(|i| MerklePaymentCandidateNode {
2037 pub_key: vec![i as u8; 32],
2038 price: Amount::from((i as u64 + 1) * 100),
2040 reward_address: RewardsAddress::new([i as u8; 20]),
2041 merkle_payment_timestamp: timestamp,
2042 signature: vec![i as u8; 64],
2043 committed_key_count: 0,
2044 commitment_pin: None,
2045 });
2046
2047 MerklePaymentCandidatePool {
2048 midpoint_proof: midpoint,
2049 candidate_nodes,
2050 }
2051 }
2052
2053 fn median16(mut amounts: Vec<Amount>) -> Amount {
2055 amounts.sort_unstable();
2056 *amounts.get(amounts.len() / 2).unwrap()
2057 }
2058
2059 #[test]
2060 fn pool_commitment_applies_payment_multiplier_to_every_candidate() {
2061 let pool = pool_with_varied_prices(1_700_000_000);
2062 let commitment = pool_commitment_with_payment_multiplier(&pool).unwrap();
2063
2064 for (candidate, signed) in commitment
2065 .candidates
2066 .iter()
2067 .zip(pool.candidate_nodes.iter())
2068 {
2069 assert_eq!(
2070 candidate.price,
2071 signed.price * Amount::from(MERKLE_PAYMENT_MULTIPLIER),
2072 "on-chain payable amount must be {MERKLE_PAYMENT_MULTIPLIER}x the quoted price"
2073 );
2074 }
2075 }
2076
2077 #[test]
2078 fn pool_commitment_multiplier_leaves_signed_prices_and_pool_hash_untouched() {
2079 let pool = pool_with_varied_prices(1_700_000_000);
2080 let before: Vec<Amount> = pool.candidate_nodes.iter().map(|c| c.price).collect();
2081
2082 let commitment = pool_commitment_with_payment_multiplier(&pool).unwrap();
2083
2084 let after: Vec<Amount> = pool.candidate_nodes.iter().map(|c| c.price).collect();
2085 assert_eq!(before, after, "signed candidate prices must not change");
2086 assert_eq!(
2087 commitment.pool_hash,
2088 pool.hash(),
2089 "pool hash is the storer's on-chain lookup key and must be \
2090 computed over the signed 1x prices"
2091 );
2092 }
2093
2094 #[test]
2103 fn merkle_settlement_per_padded_leaf_is_the_multiplied_pool_median() {
2104 let pool = pool_with_varied_prices(1_700_000_000);
2105 let commitment = pool_commitment_with_payment_multiplier(&pool).unwrap();
2106
2107 let quoted_median = median16(pool.candidate_nodes.iter().map(|c| c.price).collect());
2108 let per_chunk = median16(commitment.candidates.iter().map(|c| c.price).collect());
2109
2110 assert_eq!(quoted_median, Amount::from(900u64));
2111 assert_eq!(
2112 per_chunk,
2113 quoted_median * Amount::from(MERKLE_PAYMENT_MULTIPLIER),
2114 "merkle per-chunk settlement must equal the single-node \
2115 {MERKLE_PAYMENT_MULTIPLIER}x median, not the bare quoted price"
2116 );
2117 }
2118
2119 #[test]
2120 fn test_finalize_merkle_batch_with_valid_winner() {
2121 let prepared = make_prepared_merkle_batch(4);
2122 let winner_hash = prepared.candidate_pools[0].hash();
2123
2124 let result = finalize_merkle_batch(prepared, winner_hash);
2125 assert!(
2126 result.is_ok(),
2127 "should succeed with valid winner: {result:?}"
2128 );
2129
2130 let batch = result.unwrap();
2131 assert_eq!(batch.chunk_count, 4);
2132 assert_eq!(batch.proofs.len(), 4);
2133
2134 for proof_bytes in batch.proofs.values() {
2136 assert!(!proof_bytes.is_empty());
2137 }
2138 }
2139
2140 #[test]
2148 fn test_finalize_merkle_batch_ships_no_commitment_sidecars() {
2149 use ant_protocol::payment::deserialize_merkle_proof;
2150
2151 let mut prepared = make_prepared_merkle_batch(4);
2152 for pool in &mut prepared.candidate_pools {
2155 for candidate in &mut pool.candidate_nodes {
2156 candidate.committed_key_count = 9_000;
2157 candidate.commitment_pin = Some([7u8; 32]);
2158 }
2159 }
2160 let winner_hash = prepared.candidate_pools[0].hash();
2161
2162 let batch = finalize_merkle_batch(prepared, winner_hash).unwrap();
2163 assert_eq!(batch.proofs.len(), 4);
2164 for proof_bytes in batch.proofs.values() {
2165 let proof = deserialize_merkle_proof(proof_bytes).unwrap();
2166 assert!(
2167 proof.commitment_sidecars.is_empty(),
2168 "per-chunk merkle proofs must not ship commitment sidecars"
2169 );
2170 }
2171 }
2172
2173 #[test]
2174 fn test_finalize_merkle_batch_with_invalid_winner() {
2175 let prepared = make_prepared_merkle_batch(4);
2176 let bad_hash = [0xFF; 32];
2177
2178 let result = finalize_merkle_batch(prepared, bad_hash);
2179 assert!(result.is_err());
2180 let err = result.unwrap_err().to_string();
2181 assert!(err.contains("not found in candidate pools"), "got: {err}");
2182 }
2183
2184 #[test]
2185 fn test_finalize_merkle_batch_proofs_are_deserializable() {
2186 use ant_protocol::payment::deserialize_merkle_proof;
2187
2188 let prepared = make_prepared_merkle_batch(8);
2189 let winner_hash = prepared.candidate_pools[0].hash();
2190
2191 let batch = finalize_merkle_batch(prepared, winner_hash).unwrap();
2192
2193 for (addr, proof_bytes) in &batch.proofs {
2194 let proof = deserialize_merkle_proof(proof_bytes);
2195 assert!(
2196 proof.is_ok(),
2197 "proof for {} should deserialize: {:?}",
2198 hex::encode(addr),
2199 proof.err()
2200 );
2201 }
2202 }
2203
2204 const PARTITION_CASES: [(usize, &[usize]); 10] = [
2213 (2, &[2]),
2214 (64, &[64]),
2215 (65, &[65]),
2216 (255, &[255]),
2217 (256, &[256]),
2218 (257, &[255, 2]),
2219 (300, &[256, 44]),
2220 (512, &[256, 256]),
2221 (513, &[256, 255, 2]),
2222 (769, &[256, 256, 255, 2]),
2223 ];
2224
2225 #[test]
2226 fn merkle_batch_sizes_rebalance_singleton_remainders() {
2227 for (total, expected) in PARTITION_CASES {
2228 assert_eq!(
2229 merkle_batch_sizes(total),
2230 expected,
2231 "{total} addresses must partition as {expected:?}"
2232 );
2233 }
2234 }
2235
2236 #[test]
2240 fn merkle_batch_sizes_are_always_buildable_trees() {
2241 for total in 2..=(4 * MAX_LEAVES + 3) {
2242 let sizes = merkle_batch_sizes(total);
2243 assert!(!sizes.is_empty(), "{total} addresses must produce batches");
2244 assert_eq!(
2245 sizes.iter().sum::<usize>(),
2246 total,
2247 "{total} addresses: partition must cover every address"
2248 );
2249 for size in sizes {
2250 assert!(
2251 (2..=MAX_LEAVES).contains(&size),
2252 "{total} addresses produced a batch of {size}, outside 2..={MAX_LEAVES}"
2253 );
2254 }
2255 }
2256 }
2257
2258 #[test]
2259 fn merkle_batch_sizes_below_two_have_no_payable_partition() {
2260 assert!(merkle_batch_sizes(0).is_empty());
2261 assert!(merkle_batch_sizes(1).is_empty());
2262 }
2263
2264 #[test]
2265 fn merkle_batch_partitions_preserve_order_and_use_each_address_once() {
2266 for (total, _) in PARTITION_CASES {
2267 let addrs = make_test_addresses(total);
2268 let partitions = merkle_batch_partitions(&addrs);
2269
2270 let flattened: Vec<[u8; 32]> = partitions.concat();
2271 assert_eq!(
2272 flattened, addrs,
2273 "{total} addresses: partitions must concatenate back to the input in order"
2274 );
2275
2276 let unique: std::collections::HashSet<[u8; 32]> = flattened.iter().copied().collect();
2277 assert_eq!(
2278 unique.len(),
2279 total,
2280 "{total} addresses: no address may be duplicated or synthesised"
2281 );
2282 }
2283 }
2284
2285 #[test]
2289 fn post_preflight_plan_of_257_partitions_into_payable_batches() {
2290 let plan = MerkleUploadPlan {
2291 already_stored: make_test_addresses(3),
2292 to_upload: make_test_addresses(257),
2293 to_upload_total_bytes: 257 * 1024,
2294 };
2295 assert_eq!(plan.to_upload.len(), 257);
2296
2297 let partitions = merkle_batch_partitions(&plan.to_upload);
2298 let sizes: Vec<usize> = partitions.iter().map(|batch| batch.len()).collect();
2299 assert_eq!(sizes, vec![255, 2]);
2300 for batch in partitions {
2301 let xornames: Vec<XorName> = batch.iter().map(|a| XorName(*a)).collect();
2302 assert!(
2303 MerkleTree::from_xornames(xornames).is_ok(),
2304 "every partition of a 257-chunk plan must build a tree"
2305 );
2306 }
2307 }
2308
2309 #[test]
2313 fn no_partition_pays_before_a_singleton_tree_failure() {
2314 for total in [257usize, 513, 769] {
2315 let addrs = make_test_addresses(total);
2316 for batch in merkle_batch_partitions(&addrs) {
2317 let xornames: Vec<XorName> = batch.iter().map(|a| XorName(*a)).collect();
2318 assert!(
2319 MerkleTree::from_xornames(xornames).is_ok(),
2320 "{total} addresses: batch of {} is unpayable",
2321 batch.len()
2322 );
2323 }
2324 }
2325 }
2326
2327 #[test]
2328 fn merkle_billable_leaves_sum_the_padded_partitions() {
2329 for (total, expected) in PARTITION_CASES {
2330 let padded: u64 = expected
2331 .iter()
2332 .map(|size| size.next_power_of_two() as u64)
2333 .sum();
2334 assert_eq!(
2335 merkle_billable_leaves(total as u64),
2336 padded,
2337 "{total} chunks must bill for the padded partition {expected:?}"
2338 );
2339 }
2340
2341 assert_eq!(merkle_billable_leaves(65), 128);
2343 assert_eq!(merkle_billable_leaves(257), 256 + 2);
2344 assert_eq!(merkle_billable_leaves(300), 256 + 64);
2345 assert_eq!(merkle_billable_leaves(0), 0);
2348 assert_eq!(merkle_billable_leaves(1), 2);
2349 }
2350
2351 #[test]
2352 fn merkle_billable_leaves_never_under_quote() {
2353 for chunks in 1..2000u64 {
2354 assert!(
2355 merkle_billable_leaves(chunks) >= chunks,
2356 "{chunks} chunks must never be billed as fewer leaves"
2357 );
2358 }
2359 }
2360
2361 #[test]
2365 fn external_preparation_refuses_more_than_one_tree_of_addresses() {
2366 assert!(ensure_single_merkle_tree_batch(2).is_ok());
2367 assert!(ensure_single_merkle_tree_batch(MAX_LEAVES).is_ok());
2368
2369 for oversized in [MAX_LEAVES + 1, 300, 513] {
2370 match ensure_single_merkle_tree_batch(oversized) {
2371 Err(Error::MerkleBatchTooLarge {
2372 addresses,
2373 max_leaves,
2374 }) => {
2375 assert_eq!(addresses, oversized);
2376 assert_eq!(max_leaves, MAX_LEAVES);
2377 }
2378 other => panic!("{oversized} addresses should be refused, got {other:?}"),
2379 }
2380 }
2381 }
2382
2383 use std::sync::{Arc, Mutex};
2388
2389 fn make_addrs(count: usize) -> Vec<[u8; 32]> {
2392 make_test_addresses(count)
2393 }
2394
2395 #[tokio::test]
2399 async fn store_with_retry_collects_failures_instead_of_aborting() {
2400 let chunks = make_addrs(6);
2401 let failing: std::collections::HashSet<[u8; 32]> = chunks.iter().take(2).copied().collect();
2402 let failing_for_closure = failing.clone();
2403
2404 let store_one = move |addr: [u8; 32]| {
2405 let fail = failing_for_closure.contains(&addr);
2406 async move {
2407 if fail {
2408 Err(Error::InsufficientPeers("test shortfall".into()))
2409 } else {
2410 Ok(std::time::Instant::now())
2411 }
2412 }
2413 };
2414
2415 let outcome =
2416 merkle_store_with_retry(chunks, || 8, 1, Duration::ZERO, None, 0, 6, store_one)
2417 .await
2418 .expect("quorum shortfalls must not abort the batch");
2419
2420 assert_eq!(outcome.stored, 4);
2421 assert_eq!(outcome.failed, 2);
2422 assert_eq!(outcome.stats.retries_histogram[0], 4);
2424 assert_eq!(outcome.stats.chunk_attempts_total, 6);
2425 }
2426
2427 #[tokio::test]
2433 async fn store_with_retry_rereads_cap_per_slot() {
2434 let count = 6;
2435 let chunks = make_addrs(count);
2436 let cap_calls = Arc::new(Mutex::new(0usize));
2437 let cap_calls_for_closure = cap_calls.clone();
2438 let cap = move || {
2439 *cap_calls_for_closure.lock().expect("cap counter poisoned") += 1;
2440 2
2441 };
2442 let store_one = move |_addr: [u8; 32]| async move { Ok(std::time::Instant::now()) };
2443
2444 let outcome =
2445 merkle_store_with_retry(chunks, cap, 1, Duration::ZERO, None, 0, count, store_one)
2446 .await
2447 .expect("all stores succeed");
2448
2449 assert_eq!(outcome.stored, count);
2450 let calls = *cap_calls.lock().expect("cap counter poisoned");
2451 assert!(
2452 calls >= count,
2453 "cap must be re-read per drained slot (rolling), not snapshotted once — \
2454 expected >= {count} invocations, got {calls}",
2455 );
2456 }
2457
2458 #[tokio::test]
2465 async fn store_pass_has_no_barrier() {
2466 use std::sync::atomic::{AtomicUsize, Ordering};
2467 let count = 8;
2468 let addrs = make_addrs(count);
2469 let slow = addrs[0];
2470 let fast_completed = Arc::new(AtomicUsize::new(0));
2471 let release_slow = Arc::new(tokio::sync::Notify::new());
2472
2473 let store_one = move |addr: [u8; 32]| {
2474 let fast_completed = fast_completed.clone();
2475 let release_slow = release_slow.clone();
2476 async move {
2477 if addr == slow {
2478 release_slow.notified().await;
2482 } else if fast_completed.fetch_add(1, Ordering::SeqCst) + 1 == count - 1 {
2483 release_slow.notify_one();
2484 }
2485 Ok(std::time::Instant::now())
2486 }
2487 };
2488
2489 let outcome = tokio::time::timeout(
2490 Duration::from_secs(5),
2491 merkle_store_with_retry(addrs, || 8, 1, Duration::ZERO, None, 0, count, store_one),
2492 )
2493 .await
2494 .expect("store pass must not deadlock — a slow chunk must not block the others")
2495 .expect("all stores succeed");
2496
2497 assert_eq!(outcome.stored, count);
2498 }
2499
2500 #[tokio::test]
2505 async fn store_pass_keeps_at_most_cap_in_flight() {
2506 use std::sync::atomic::{AtomicUsize, Ordering};
2507 let count = 40;
2508 let cap = 4;
2509 let addrs = make_addrs(count);
2510 let in_flight = Arc::new(AtomicUsize::new(0));
2511 let max_in_flight = Arc::new(AtomicUsize::new(0));
2512 let max_in_flight_for_closure = max_in_flight.clone();
2513
2514 let store_one = move |_addr: [u8; 32]| {
2515 let in_flight = in_flight.clone();
2516 let max_in_flight = max_in_flight_for_closure.clone();
2517 async move {
2518 let now = in_flight.fetch_add(1, Ordering::SeqCst) + 1;
2519 max_in_flight.fetch_max(now, Ordering::SeqCst);
2520 tokio::task::yield_now().await;
2523 in_flight.fetch_sub(1, Ordering::SeqCst);
2524 Ok(std::time::Instant::now())
2525 }
2526 };
2527
2528 let outcome = merkle_store_with_retry(
2529 addrs,
2530 move || cap,
2531 1,
2532 Duration::ZERO,
2533 None,
2534 0,
2535 count,
2536 store_one,
2537 )
2538 .await
2539 .expect("all stores succeed");
2540
2541 assert_eq!(outcome.stored, count);
2542 let peak = max_in_flight.load(Ordering::SeqCst);
2543 assert!(
2544 peak <= cap,
2545 "at most `cap` bodies may be in flight (memory bound), got peak {peak} > cap {cap}",
2546 );
2547 assert!(
2548 peak > 1,
2549 "the pass must actually run concurrently, not serialize (peak {peak})",
2550 );
2551 }
2552
2553 #[tokio::test]
2558 async fn store_with_retry_treats_remote_put_as_recoverable() {
2559 let chunks = make_addrs(6);
2560 let failing: std::collections::HashSet<[u8; 32]> = chunks.iter().take(2).copied().collect();
2561 let failing_for_closure = failing.clone();
2562
2563 let store_one = move |addr: [u8; 32]| {
2564 let fail = failing_for_closure.contains(&addr);
2565 async move {
2566 if fail {
2567 Err(Error::RemotePut {
2568 address: hex::encode(addr),
2569 source: ant_protocol::ProtocolError::StorageFailed(
2570 "insufficient disk space".into(),
2571 ),
2572 })
2573 } else {
2574 Ok(std::time::Instant::now())
2575 }
2576 }
2577 };
2578
2579 let outcome =
2580 merkle_store_with_retry(chunks, || 8, 1, Duration::ZERO, None, 0, 6, store_one)
2581 .await
2582 .expect("remote app-rejections must not abort the batch");
2583
2584 assert_eq!(outcome.stored, 4);
2585 assert_eq!(outcome.failed, 2);
2586 }
2587
2588 #[tokio::test]
2592 async fn store_with_retry_reports_non_quorum_errors_as_fatal() {
2593 let chunks = make_addrs(3);
2594 let store_one = |_addr: [u8; 32]| async move {
2595 Err::<std::time::Instant, _>(Error::Payment("missing proof".into()))
2596 };
2597
2598 let outcome =
2599 merkle_store_with_retry(chunks, || 8, 3, Duration::ZERO, None, 0, 3, store_one)
2600 .await
2601 .expect("fatal is carried in the outcome, not returned as Err");
2602 assert!(matches!(outcome.fatal, Some(Error::Payment(_))));
2603 }
2604
2605 #[tokio::test]
2610 async fn store_with_retry_fatal_preserves_same_pass_successes() {
2611 let chunks = make_addrs(6);
2612 let bad = chunks[5];
2613 let store_one = move |addr: [u8; 32]| async move {
2614 if addr == bad {
2615 Err(Error::Payment("fatal".into()))
2616 } else {
2617 Ok(std::time::Instant::now())
2618 }
2619 };
2620
2621 let outcome =
2622 merkle_store_with_retry(chunks, || 1, 1, Duration::ZERO, None, 0, 6, store_one)
2623 .await
2624 .expect("fatal carried in outcome, not returned as Err");
2625 assert!(matches!(outcome.fatal, Some(Error::Payment(_))));
2626 assert_eq!(outcome.stored, 5);
2628 assert_eq!(outcome.stored_addresses.len(), 5);
2629 assert!(!outcome.stored_addresses.contains(&bad));
2630 assert!(outcome.failed_addresses.iter().any(|(a, _)| *a == bad));
2632 }
2633
2634 #[tokio::test]
2636 async fn store_with_retry_retries_only_the_failed_set() {
2637 let chunks = make_addrs(5);
2638 let total = chunks.len();
2639 let failing: std::collections::HashSet<[u8; 32]> = chunks.iter().take(2).copied().collect();
2640 let failing_for_closure = failing.clone();
2641
2642 let calls = Arc::new(Mutex::new(Vec::<[u8; 32]>::new()));
2644 let calls_for_closure = calls.clone();
2645
2646 let store_one = move |addr: [u8; 32]| {
2647 let calls = calls_for_closure.clone();
2648 let already_seen = calls.lock().unwrap().iter().filter(|&&a| a == addr).count();
2650 let fail = failing_for_closure.contains(&addr) && already_seen == 0;
2651 calls.lock().unwrap().push(addr);
2652 async move {
2653 if fail {
2654 Err(Error::InsufficientPeers("round-1 shortfall".into()))
2655 } else {
2656 Ok(std::time::Instant::now())
2657 }
2658 }
2659 };
2660
2661 let outcome =
2662 merkle_store_with_retry(chunks, || 8, 3, Duration::ZERO, None, 0, total, store_one)
2663 .await
2664 .expect("should converge after retry");
2665
2666 assert_eq!(outcome.stored, total);
2667 assert_eq!(outcome.failed, 0);
2668
2669 let calls = calls.lock().unwrap();
2673 assert_eq!(calls.len(), total + failing.len());
2674 let round_two: std::collections::HashSet<[u8; 32]> =
2675 calls[total..].iter().copied().collect();
2676 assert_eq!(round_two, failing);
2677 }
2678
2679 #[tokio::test]
2682 async fn store_with_retry_counts_retry_success_once_in_histogram() {
2683 let chunks = make_addrs(4);
2684 let total = chunks.len();
2685 let flaky_addr = chunks[0];
2686
2687 let attempts = Arc::new(Mutex::new(HashMap::<[u8; 32], usize>::new()));
2688 let attempts_for_closure = attempts.clone();
2689
2690 let store_one = move |addr: [u8; 32]| {
2691 let attempts = attempts_for_closure.clone();
2692 let n = {
2693 let mut m = attempts.lock().unwrap();
2694 let entry = m.entry(addr).or_insert(0);
2695 *entry += 1;
2696 *entry
2697 };
2698 let fail = addr == flaky_addr && n == 1;
2699 async move {
2700 if fail {
2701 Err(Error::InsufficientPeers("transient".into()))
2702 } else {
2703 Ok(std::time::Instant::now())
2704 }
2705 }
2706 };
2707
2708 let outcome =
2709 merkle_store_with_retry(chunks, || 8, 3, Duration::ZERO, None, 0, total, store_one)
2710 .await
2711 .expect("flaky chunk should recover on retry");
2712
2713 assert_eq!(outcome.stored, total);
2714 assert_eq!(outcome.failed, 0);
2715 assert_eq!(outcome.stats.retries_histogram[0], total - 1);
2717 assert_eq!(outcome.stats.retries_histogram[1], 1);
2718 assert_eq!(outcome.stats.chunk_attempts_total, total + 1);
2720 }
2721
2722 #[tokio::test]
2727 async fn store_with_retry_reports_all_failed_when_retries_exhausted() {
2728 let chunks = make_addrs(3);
2729 let total = chunks.len();
2730
2731 let store_one = |_addr: [u8; 32]| async move {
2732 Err::<std::time::Instant, _>(Error::InsufficientPeers("never converges".into()))
2733 };
2734
2735 let outcome = merkle_store_with_retry(
2736 chunks,
2737 || 8,
2738 MERKLE_STORE_MAX_ATTEMPTS,
2739 Duration::ZERO,
2740 None,
2741 0,
2742 total,
2743 store_one,
2744 )
2745 .await
2746 .expect("an exhausted retry budget is reported, not propagated as Err");
2747
2748 assert_eq!(outcome.stored, 0);
2749 assert_eq!(outcome.failed, total);
2750 assert_eq!(
2752 outcome.stats.chunk_attempts_total,
2753 total * MERKLE_STORE_MAX_ATTEMPTS
2754 );
2755 assert_eq!(outcome.stats.retries_histogram, [0; 4]);
2757 }
2758
2759 #[tokio::test]
2764 async fn store_with_retry_records_failed_addresses_when_exhausted() {
2765 let chunks = make_addrs(6);
2766 let failing: std::collections::HashSet<[u8; 32]> = chunks.iter().take(2).copied().collect();
2767 let failing_for_closure = failing.clone();
2768
2769 let store_one = move |addr: [u8; 32]| {
2770 let fail = failing_for_closure.contains(&addr);
2771 async move {
2772 if fail {
2773 Err(Error::InsufficientPeers("permanent shortfall".into()))
2774 } else {
2775 Ok(std::time::Instant::now())
2776 }
2777 }
2778 };
2779
2780 let outcome = merkle_store_with_retry(
2781 chunks,
2782 || 8,
2783 MERKLE_STORE_MAX_ATTEMPTS,
2784 Duration::ZERO,
2785 None,
2786 0,
2787 6,
2788 store_one,
2789 )
2790 .await
2791 .expect("quorum shortfalls must not abort the batch");
2792
2793 assert_eq!(outcome.stored, 4);
2794 assert_eq!(outcome.failed, 2);
2795 assert_eq!(outcome.failed_addresses.len(), 2);
2797 let reported: std::collections::HashSet<[u8; 32]> =
2798 outcome.failed_addresses.iter().map(|(a, _)| *a).collect();
2799 assert_eq!(reported, failing);
2800 for (_, msg) in &outcome.failed_addresses {
2802 assert!(msg.contains("permanent shortfall"));
2803 }
2804 }
2805
2806 #[tokio::test]
2809 async fn store_with_retry_failed_addresses_empty_on_full_success() {
2810 let chunks = make_addrs(4);
2811 let total = chunks.len();
2812 let store_one = |_addr: [u8; 32]| async move { Ok(std::time::Instant::now()) };
2813
2814 let outcome = merkle_store_with_retry(
2815 chunks,
2816 || 8,
2817 MERKLE_STORE_MAX_ATTEMPTS,
2818 Duration::ZERO,
2819 None,
2820 0,
2821 total,
2822 store_one,
2823 )
2824 .await
2825 .expect("all chunks store");
2826
2827 assert_eq!(outcome.stored, total);
2828 assert_eq!(outcome.failed, 0);
2829 assert!(outcome.failed_addresses.is_empty());
2830 }
2831
2832 #[test]
2839 fn deferred_round_histogram_slot_maps_and_clamps() {
2840 assert_eq!(deferred_round_histogram_slot(0, 4), 1);
2841 assert_eq!(deferred_round_histogram_slot(1, 4), 2);
2842 assert_eq!(deferred_round_histogram_slot(2, 4), 3);
2843 assert_eq!(deferred_round_histogram_slot(3, 4), 3);
2845 assert_eq!(deferred_round_histogram_slot(9, 4), 3);
2846 }
2847
2848 fn deferred_set(count: usize) -> Vec<([u8; 32], String)> {
2849 make_test_addresses(count)
2850 .into_iter()
2851 .map(|addr| (addr, "short of quorum".to_string()))
2852 .collect()
2853 }
2854
2855 #[tokio::test]
2859 async fn deferred_retry_succeeds_on_a_later_round() {
2860 let deferred = deferred_set(3);
2861 let attempts = Arc::new(Mutex::new(HashMap::<[u8; 32], usize>::new()));
2864 let attempts_for_closure = attempts.clone();
2865 let store_one = move |addr: [u8; 32]| {
2866 let attempts = attempts_for_closure.clone();
2867 async move {
2868 let n = {
2869 let mut map = attempts.lock().unwrap();
2870 let e = map.entry(addr).or_insert(0);
2871 *e += 1;
2872 *e
2873 };
2874 if n < 2 {
2875 Err(Error::InsufficientPeers("still short".into()))
2876 } else {
2877 Ok(std::time::Instant::now())
2878 }
2879 }
2880 };
2881
2882 let outcome = merkle_deferred_retry(
2883 deferred,
2884 &[0, 0, 0],
2885 |n: usize| n.max(1),
2886 None,
2887 0,
2888 3,
2889 store_one,
2890 )
2891 .await
2892 .expect("deferred retry must not abort on quorum shortfalls");
2893
2894 assert_eq!(outcome.stored, 3, "all three land by round 1");
2895 assert_eq!(outcome.stored_addresses.len(), 3);
2896 assert_eq!(outcome.failed, 0);
2897 assert!(outcome.failed_addresses.is_empty());
2898 assert!(outcome.fatal.is_none());
2899 assert_eq!(outcome.stats.retries_histogram[1], 0);
2901 assert_eq!(outcome.stats.retries_histogram[2], 3);
2902 assert_eq!(outcome.stats.chunk_attempts_total, 6);
2904 }
2905
2906 #[tokio::test]
2909 async fn deferred_retry_leftovers_become_failed() {
2910 let deferred = deferred_set(2);
2911 let store_one = |_addr: [u8; 32]| async move {
2912 Err::<std::time::Instant, _>(Error::InsufficientPeers("always short".into()))
2913 };
2914
2915 let outcome = merkle_deferred_retry(
2916 deferred,
2917 &[0, 0, 0],
2918 |n: usize| n.max(1),
2919 None,
2920 0,
2921 2,
2922 store_one,
2923 )
2924 .await
2925 .expect("exhausted retries report failures, not an error");
2926
2927 assert_eq!(outcome.stored, 0);
2928 assert!(outcome.stored_addresses.is_empty());
2929 assert_eq!(outcome.failed, 2);
2930 assert_eq!(outcome.failed_addresses.len(), 2);
2931 assert!(outcome.fatal.is_none());
2932 assert_eq!(outcome.stats.chunk_attempts_total, 6);
2934 }
2935
2936 #[tokio::test]
2941 async fn deferred_retry_fatal_error_preserves_prior_progress() {
2942 let addrs = make_test_addresses(2);
2943 let good = addrs[0];
2944 let bad = addrs[1];
2945 let deferred = vec![(good, "short".to_string()), (bad, "short".to_string())];
2946
2947 let attempts = Arc::new(Mutex::new(HashMap::<[u8; 32], usize>::new()));
2950 let attempts_for_closure = attempts.clone();
2951 let store_one = move |addr: [u8; 32]| {
2952 let attempts = attempts_for_closure.clone();
2953 async move {
2954 let n = {
2955 let mut map = attempts.lock().unwrap();
2956 let e = map.entry(addr).or_insert(0);
2957 *e += 1;
2958 *e
2959 };
2960 if addr == good {
2961 Ok(std::time::Instant::now())
2962 } else if n == 1 {
2963 Err(Error::InsufficientPeers("short".into()))
2964 } else {
2965 Err(Error::Payment("fatal on retry".into()))
2966 }
2967 }
2968 };
2969
2970 let outcome = merkle_deferred_retry(
2971 deferred,
2972 &[0, 0, 0],
2973 |n: usize| n.max(1),
2974 None,
2975 0,
2976 2,
2977 store_one,
2978 )
2979 .await
2980 .expect("a fatal round error is reported via `fatal`, not as Err");
2981
2982 assert!(outcome.fatal.is_some(), "fatal error must be captured");
2983 assert_eq!(outcome.stored, 1, "round-0 success preserved");
2984 assert_eq!(outcome.stored_addresses, vec![good]);
2985 assert_eq!(outcome.failed, 1);
2986 assert_eq!(outcome.failed_addresses.len(), 1);
2987 assert_eq!(outcome.failed_addresses[0].0, bad);
2988 }
2989
2990 #[tokio::test]
2992 async fn deferred_retry_empty_set_is_a_noop() {
2993 let store_one = |_addr: [u8; 32]| async move {
2994 Err::<std::time::Instant, _>(Error::InsufficientPeers("unused".into()))
2995 };
2996
2997 let outcome = merkle_deferred_retry(
2998 Vec::new(),
2999 &DEFERRED_ROUND_DELAYS_SECS,
3000 |n: usize| n.max(1),
3001 None,
3002 7,
3003 7,
3004 store_one,
3005 )
3006 .await
3007 .expect("empty deferred set is a no-op");
3008
3009 assert_eq!(outcome.stored, 7, "stored_offset carried through unchanged");
3010 assert_eq!(outcome.failed, 0);
3011 assert!(outcome.stored_addresses.is_empty());
3012 assert!(outcome.failed_addresses.is_empty());
3013 assert!(outcome.fatal.is_none());
3014 }
3015}