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 merkle_batch_sizes_with_cap(total, MAX_LEAVES)
378}
379
380#[must_use]
390pub fn merkle_batch_sizes_with_cap(total: usize, cap: usize) -> Vec<usize> {
391 if total < 2 {
392 return Vec::new();
393 }
394 let cap = cap.clamp(3, MAX_LEAVES);
395
396 let mut sizes = Vec::with_capacity(total.div_ceil(cap));
397 let mut remaining = total;
398 while remaining > cap {
399 let take = if remaining - cap == 1 { cap - 1 } else { cap };
402 sizes.push(take);
403 remaining -= take;
404 }
405 sizes.push(remaining);
406 sizes
407}
408
409#[must_use]
414pub fn merkle_batch_partitions(addresses: &[[u8; 32]]) -> Vec<&[[u8; 32]]> {
415 merkle_batch_partitions_with_cap(addresses, MAX_LEAVES)
416}
417
418#[must_use]
421pub fn merkle_batch_partitions_with_cap(addresses: &[[u8; 32]], cap: usize) -> Vec<&[[u8; 32]]> {
422 let mut partitions = Vec::new();
423 let mut rest = addresses;
424 for size in merkle_batch_sizes_with_cap(addresses.len(), cap) {
425 let (batch, tail) = rest.split_at(size);
426 partitions.push(batch);
427 rest = tail;
428 }
429 partitions
430}
431
432#[must_use]
441pub(crate) fn merge_merkle_batch_results(
442 results: Vec<MerkleBatchPaymentResult>,
443) -> MerkleBatchPaymentResult {
444 let mut merged = MerkleBatchPaymentResult {
445 proofs: HashMap::new(),
446 chunk_count: 0,
447 storage_cost_atto: "0".to_string(),
448 gas_cost_wei: 0,
449 merkle_payment_timestamp: 0,
450 };
451 let mut total_storage = Amount::ZERO;
452 for result in results {
453 merged.proofs.extend(result.proofs);
454 merged.chunk_count += result.chunk_count;
455 if let Ok(cost) = result.storage_cost_atto.parse::<Amount>() {
456 total_storage += cost;
457 }
458 merged.gas_cost_wei = merged.gas_cost_wei.saturating_add(result.gas_cost_wei);
459 if merged.merkle_payment_timestamp == 0
460 || (result.merkle_payment_timestamp > 0
461 && result.merkle_payment_timestamp < merged.merkle_payment_timestamp)
462 {
463 merged.merkle_payment_timestamp = result.merkle_payment_timestamp;
464 }
465 }
466 merged.storage_cost_atto = total_storage.to_string();
467 merged
468}
469
470fn padded_leaf_count(batch_size: usize) -> u64 {
475 let padded = batch_size
478 .max(2)
479 .checked_next_power_of_two()
480 .unwrap_or(usize::MAX);
481 u64::try_from(padded).unwrap_or(u64::MAX)
482}
483
484#[must_use]
493pub fn merkle_billable_leaves(chunk_count: u64) -> u64 {
494 let total = usize::try_from(chunk_count).unwrap_or(usize::MAX);
495 let batches = merkle_batch_sizes(total);
496 if batches.is_empty() {
497 return if total == 0 { 0 } else { 2 };
500 }
501
502 batches
503 .into_iter()
504 .map(padded_leaf_count)
505 .fold(0u64, u64::saturating_add)
506}
507
508fn ensure_single_merkle_tree_batch(address_count: usize) -> Result<()> {
517 if address_count > MAX_LEAVES {
518 return Err(Error::MerkleBatchTooLarge {
519 addresses: address_count,
520 max_leaves: MAX_LEAVES,
521 });
522 }
523 Ok(())
524}
525
526#[must_use]
529pub fn should_use_merkle(chunk_count: usize, mode: PaymentMode) -> bool {
530 match mode {
531 PaymentMode::Auto => chunk_count >= DEFAULT_MERKLE_THRESHOLD,
532 PaymentMode::Merkle => chunk_count >= 2,
533 PaymentMode::Single => false,
534 }
535}
536
537impl Client {
538 #[must_use]
540 pub fn should_use_merkle(&self, chunk_count: usize, mode: PaymentMode) -> bool {
541 should_use_merkle(chunk_count, mode)
542 }
543
544 pub async fn pay_for_merkle_batch(
559 &self,
560 addresses: &[[u8; 32]],
561 data_type: u32,
562 data_size: u64,
563 ) -> Result<MerkleBatchPaymentResult> {
564 let chunk_count = addresses.len();
565 if chunk_count < 2 {
566 return Err(Error::Payment(
567 "Merkle batch payment requires at least 2 chunks".to_string(),
568 ));
569 }
570
571 if chunk_count > MAX_LEAVES {
572 return self
573 .pay_for_merkle_multi_batch(addresses, data_type, data_size)
574 .await;
575 }
576
577 self.pay_for_merkle_single_batch(addresses, data_type, data_size)
578 .await
579 }
580
581 pub(crate) async fn plan_merkle_upload(
590 &self,
591 chunks: Vec<([u8; 32], u64)>,
592 data_type: u32,
593 progress: Option<&mpsc::Sender<UploadEvent>>,
594 ) -> Result<MerkleUploadPlan> {
595 let total_chunks = chunks.len();
596 if total_chunks == 0 {
597 return Ok(MerkleUploadPlan::default());
598 }
599
600 info!("Checking {total_chunks} merkle chunks for existing storage before payment");
601
602 let quote_limiter = self.controller().quote.clone();
603 let quote_concurrency = quote_limiter.current().min(total_chunks.max(1));
604 let mut check_stream = stream::iter(chunks.into_iter().enumerate())
605 .map(|(index, (address, data_size))| {
606 let limiter = quote_limiter.clone();
607 async move {
608 let result = observe_op(
609 &limiter,
610 || async move {
611 self.chunk_already_stored_for_merkle(&address, data_type, data_size)
612 .await
613 },
614 classify_error,
615 )
616 .await;
617 (index, address, data_size, result)
618 }
619 })
620 .buffer_unordered(quote_concurrency);
621
622 let mut already_stored: Vec<(usize, [u8; 32])> = Vec::new();
623 let mut to_upload: Vec<(usize, [u8; 32], u64)> = Vec::new();
624 let mut checked = 0usize;
625
626 while let Some((index, address, data_size, result)) = check_stream.next().await {
627 let is_already_stored = result?;
628 checked += 1;
629
630 if let Some(tx) = progress {
631 let _ = tx.try_send(UploadEvent::ChunkQuoted {
632 quoted: checked,
633 total: total_chunks,
634 });
635 }
636
637 if is_already_stored {
638 debug!(
639 "Merkle preflight {checked}/{total_chunks}: chunk {} already stored",
640 hex::encode(address)
641 );
642 already_stored.push((index, address));
643 if let Some(tx) = progress {
644 let _ = tx.try_send(UploadEvent::ChunkStored {
645 stored: already_stored.len(),
646 total: total_chunks,
647 });
648 }
649 } else {
650 debug!(
651 "Merkle preflight {checked}/{total_chunks}: chunk {} needs upload",
652 hex::encode(address)
653 );
654 to_upload.push((index, address, data_size));
655 }
656 }
657
658 already_stored.sort_by_key(|(index, _)| *index);
659 to_upload.sort_by_key(|(index, _, _)| *index);
660
661 let to_upload_total_bytes = to_upload.iter().fold(0u64, |acc, (_, _, data_size)| {
662 acc.saturating_add(*data_size)
663 });
664
665 let already_stored = already_stored
666 .into_iter()
667 .map(|(_, address)| address)
668 .collect::<Vec<_>>();
669 let to_upload = to_upload
670 .into_iter()
671 .map(|(_, address, _)| address)
672 .collect::<Vec<_>>();
673
674 info!(
675 "Merkle preflight complete: {} already stored, {} need upload",
676 already_stored.len(),
677 to_upload.len()
678 );
679
680 Ok(MerkleUploadPlan {
681 already_stored,
682 to_upload,
683 to_upload_total_bytes,
684 })
685 }
686
687 async fn chunk_already_stored_for_merkle(
688 &self,
689 address: &[u8; 32],
690 data_type: u32,
691 data_size: u64,
692 ) -> Result<bool> {
693 let result = self
694 .get_store_quotes_with_fault_tolerance(address, data_size, data_type)
695 .await;
696 if let Err(e) = &result {
697 if matches!(classify_error(e), Outcome::Timeout | Outcome::NetworkError) {
698 debug!(
699 "Merkle preflight: could not determine stored status for {} ({e}); \
700 treating as not stored and queuing for upload",
701 hex::encode(address)
702 );
703 }
704 }
705 preflight_stored_status(result)
706 }
707
708 pub async fn prepare_merkle_batches_external(
729 &self,
730 addresses: &[[u8; 32]],
731 data_type: u32,
732 data_size: u64,
733 cap: usize,
734 ) -> Result<Vec<PreparedMerkleBatch>> {
735 if addresses.len() < 2 {
736 return Err(Error::Payment(
737 "Merkle batch payment requires at least 2 chunks".to_string(),
738 ));
739 }
740 let partitions = merkle_batch_partitions_with_cap(addresses, cap);
741 let total = partitions.len();
742 let mut batches = Vec::with_capacity(total);
743 for (i, partition) in partitions.into_iter().enumerate() {
744 debug!(
745 "Preparing external merkle sub-batch {}/{total} ({} chunks)",
746 i + 1,
747 partition.len()
748 );
749 batches.push(
750 self.prepare_merkle_batch_external(partition, data_type, data_size)
751 .await?,
752 );
753 }
754 Ok(batches)
755 }
756
757 pub async fn prepare_merkle_batch_external(
773 &self,
774 addresses: &[[u8; 32]],
775 data_type: u32,
776 data_size: u64,
777 ) -> Result<PreparedMerkleBatch> {
778 ensure_single_merkle_tree_batch(addresses.len())?;
779
780 let chunk_count = addresses.len();
781 let xornames: Vec<XorName> = addresses.iter().map(|a| XorName(*a)).collect();
782
783 debug!("Building merkle tree for {chunk_count} chunks");
784
785 let tree = MerkleTree::from_xornames(xornames)
787 .map_err(|e| Error::Payment(format!("Failed to build merkle tree: {e}")))?;
788
789 let depth = tree.depth();
790 let merkle_payment_timestamp = std::time::SystemTime::now()
791 .duration_since(std::time::UNIX_EPOCH)
792 .map_err(|e| Error::Payment(format!("System time error: {e}")))?
793 .as_secs();
794
795 debug!("Merkle tree: depth={depth}, leaves={chunk_count}, ts={merkle_payment_timestamp}");
796
797 let midpoint_proofs = tree
799 .reward_candidates(merkle_payment_timestamp)
800 .map_err(|e| Error::Payment(format!("Failed to generate reward candidates: {e}")))?;
801
802 debug!(
803 "Collecting candidate pools from {} midpoints (concurrent)",
804 midpoint_proofs.len()
805 );
806
807 let candidate_pools = self
813 .build_candidate_pools(
814 &midpoint_proofs,
815 data_type,
816 data_size,
817 merkle_payment_timestamp,
818 )
819 .await?;
820
821 let pool_commitments: Vec<PoolCommitment> = candidate_pools
826 .iter()
827 .map(pool_commitment_with_payment_multiplier)
828 .collect::<Result<Vec<_>>>()?;
829
830 Ok(PreparedMerkleBatch {
831 depth,
832 pool_commitments,
833 merkle_payment_timestamp,
834 candidate_pools,
835 tree,
836 addresses: addresses.to_vec(),
837 })
838 }
839
840 async fn pay_for_merkle_single_batch(
842 &self,
843 addresses: &[[u8; 32]],
844 data_type: u32,
845 data_size: u64,
846 ) -> Result<MerkleBatchPaymentResult> {
847 let wallet = self.require_wallet()?;
848 let prepared = self
849 .prepare_merkle_batch_external(addresses, data_type, data_size)
850 .await?;
851
852 info!(
853 "Submitting merkle batch payment on-chain (depth={})",
854 prepared.depth
855 );
856 let (winner_pool_hash, amount, gas_info) = wallet
857 .pay_for_merkle_tree(
858 prepared.depth,
859 prepared.pool_commitments.clone(),
860 prepared.merkle_payment_timestamp,
861 )
862 .await
863 .map_err(|e| Error::Payment(format!("Merkle batch payment failed: {e}")))?;
864
865 info!(
866 "Merkle payment succeeded: winner pool {}",
867 hex::encode(winner_pool_hash)
868 );
869
870 let mut result = finalize_merkle_batch(prepared, winner_pool_hash)?;
871 result.storage_cost_atto = amount.to_string();
872 result.gas_cost_wei = gas_info.gas_cost_wei;
873 Ok(result)
874 }
875
876 async fn pay_for_merkle_multi_batch(
878 &self,
879 addresses: &[[u8; 32]],
880 data_type: u32,
881 data_size: u64,
882 ) -> Result<MerkleBatchPaymentResult> {
883 let sub_batches = merkle_batch_partitions(addresses);
888 let total_sub_batches = sub_batches.len();
889 let mut all_proofs = HashMap::with_capacity(addresses.len());
890 let mut total_storage = Amount::ZERO;
891 let mut total_gas: u128 = 0;
892 let mut oldest_ts: u64 = 0;
896
897 for (i, chunk) in sub_batches.into_iter().enumerate() {
898 match self
899 .pay_for_merkle_single_batch(chunk, data_type, data_size)
900 .await
901 {
902 Ok(sub_result) => {
903 if let Ok(cost) = sub_result.storage_cost_atto.parse::<Amount>() {
904 total_storage += cost;
905 }
906 total_gas = total_gas.saturating_add(sub_result.gas_cost_wei);
907 if oldest_ts == 0
908 || (sub_result.merkle_payment_timestamp > 0
909 && sub_result.merkle_payment_timestamp < oldest_ts)
910 {
911 oldest_ts = sub_result.merkle_payment_timestamp;
912 }
913 all_proofs.extend(sub_result.proofs);
914 }
915 Err(e) => {
916 if all_proofs.is_empty() {
917 return Err(e);
919 }
920 warn!(
922 "Merkle sub-batch {}/{total_sub_batches} failed: {e}. \
923 Returning {} proofs from prior sub-batches",
924 i + 1,
925 all_proofs.len()
926 );
927 return Ok(MerkleBatchPaymentResult {
928 chunk_count: all_proofs.len(),
929 proofs: all_proofs,
930 storage_cost_atto: total_storage.to_string(),
931 gas_cost_wei: total_gas,
932 merkle_payment_timestamp: oldest_ts,
933 });
934 }
935 }
936 }
937
938 Ok(MerkleBatchPaymentResult {
939 chunk_count: addresses.len(),
940 proofs: all_proofs,
941 storage_cost_atto: total_storage.to_string(),
942 gas_cost_wei: total_gas,
943 merkle_payment_timestamp: oldest_ts,
944 })
945 }
946
947 async fn build_candidate_pools(
949 &self,
950 midpoint_proofs: &[MidpointProof],
951 data_type: u32,
952 data_size: u64,
953 merkle_payment_timestamp: u64,
954 ) -> Result<Vec<MerklePaymentCandidatePool>> {
955 let mut pool_futures = FuturesUnordered::new();
956
957 for midpoint_proof in midpoint_proofs {
958 let pool_address = midpoint_proof.address();
959 let mp = midpoint_proof.clone();
960 pool_futures.push(async move {
961 let candidate_nodes = self
962 .get_merkle_candidate_pool(
963 &pool_address.0,
964 data_type,
965 data_size,
966 merkle_payment_timestamp,
967 )
968 .await?;
969 Ok::<_, Error>(MerklePaymentCandidatePool {
970 midpoint_proof: mp,
971 candidate_nodes,
972 })
973 });
974 }
975
976 let mut pools = Vec::with_capacity(midpoint_proofs.len());
977 while let Some(result) = pool_futures.next().await {
978 pools.push(result?);
979 }
980
981 Ok(pools)
982 }
983
984 #[allow(clippy::too_many_lines)]
986 async fn get_merkle_candidate_pool(
987 &self,
988 address: &[u8; 32],
989 data_type: u32,
990 data_size: u64,
991 merkle_payment_timestamp: u64,
992 ) -> Result<[MerklePaymentCandidateNode; CANDIDATES_PER_POOL]> {
993 let node = self.network().node();
994 let timeout = Duration::from_secs(self.config().quote_timeout_secs);
995
996 let query_count = CANDIDATES_PER_POOL * 2;
998 let mut remote_peers = self
999 .network()
1000 .find_closest_peers(address, query_count)
1001 .await?;
1002
1003 if remote_peers.len() < CANDIDATES_PER_POOL {
1007 let connected = self.network().connected_peers().await;
1008 for peer in connected {
1009 if !remote_peers.iter().any(|(id, _)| *id == peer) {
1010 remote_peers.push((peer, vec![]));
1011 }
1012 }
1013 }
1014
1015 if remote_peers.len() < CANDIDATES_PER_POOL {
1016 return Err(Error::InsufficientPeers(format!(
1017 "Found {} peers, need {CANDIDATES_PER_POOL} for merkle candidate pool. \
1018 Use --no-merkle or a larger network.",
1019 remote_peers.len()
1020 )));
1021 }
1022
1023 let mut candidate_futures = FuturesUnordered::new();
1024
1025 for (peer_id, peer_addrs) in &remote_peers {
1026 let request_id = self.next_request_id();
1027 let request = MerkleCandidateQuoteRequest {
1028 address: *address,
1029 data_type,
1030 data_size,
1031 merkle_payment_timestamp,
1032 };
1033 let message = ChunkMessage {
1034 request_id,
1035 body: ChunkMessageBody::MerkleCandidateQuoteRequest(request),
1036 };
1037
1038 let message_bytes = match message.encode() {
1039 Ok(bytes) => bytes,
1040 Err(e) => {
1041 warn!("Failed to encode merkle candidate request for {peer_id}: {e}");
1042 continue;
1043 }
1044 };
1045
1046 let peer_id_clone = *peer_id;
1047 let addrs_clone = peer_addrs.clone();
1048 let node_clone = node.clone();
1049
1050 let fut = async move {
1051 let result = send_and_await_chunk_response(
1052 &node_clone,
1053 &peer_id_clone,
1054 message_bytes,
1055 request_id,
1056 timeout,
1057 &addrs_clone,
1058 |body| match body {
1059 ChunkMessageBody::MerkleCandidateQuoteResponse(
1060 MerkleCandidateQuoteResponse::Success {
1061 candidate_node,
1062 commitment,
1063 },
1064 ) => {
1065 match rmp_serde::from_slice::<MerklePaymentCandidateNode>(
1066 &candidate_node,
1067 ) {
1068 Ok(node) => Some(Ok((node, commitment))),
1069 Err(e) => Some(Err(Error::Serialization(format!(
1070 "Failed to deserialize candidate node from {peer_id_clone}: {e}"
1071 )))),
1072 }
1073 }
1074 ChunkMessageBody::MerkleCandidateQuoteResponse(
1075 MerkleCandidateQuoteResponse::Error(e),
1076 ) => Some(Err(Error::Protocol(format!(
1077 "Merkle quote error from {peer_id_clone}: {e}"
1078 )))),
1079 _ => None,
1080 },
1081 |e| {
1082 Error::Network(format!(
1083 "Failed to send merkle candidate request to {peer_id_clone}: {e}"
1084 ))
1085 },
1086 || {
1087 Error::Timeout(format!(
1088 "Timeout waiting for merkle candidate from {peer_id_clone}"
1089 ))
1090 },
1091 )
1092 .await;
1093
1094 (peer_id_clone, result)
1095 };
1096
1097 candidate_futures.push(fut);
1098 }
1099
1100 self.collect_validated_candidates(&mut candidate_futures, address, merkle_payment_timestamp)
1101 .await
1102 }
1103
1104 async fn collect_validated_candidates(
1115 &self,
1116 futures: &mut FuturesUnordered<
1117 impl std::future::Future<
1118 Output = (
1119 PeerId,
1120 std::result::Result<(MerklePaymentCandidateNode, Option<Vec<u8>>), Error>,
1121 ),
1122 >,
1123 >,
1124 target_address: &[u8; 32],
1125 merkle_payment_timestamp: u64,
1126 ) -> Result<[MerklePaymentCandidateNode; CANDIDATES_PER_POOL]> {
1127 let mut valid: Vec<(PeerId, MerklePaymentCandidateNode)> = Vec::new();
1128 let mut failures: Vec<String> = Vec::new();
1129
1130 while let Some((peer_id, result)) = futures.next().await {
1131 match result {
1132 Ok((candidate, commitment)) => {
1133 if !verify_merkle_candidate_signature(&candidate) {
1134 warn!("Invalid ML-DSA-65 signature from merkle candidate {peer_id}");
1135 failures.push(format!("{peer_id}: invalid signature"));
1136 continue;
1137 }
1138 if candidate.merkle_payment_timestamp != merkle_payment_timestamp {
1139 warn!("Timestamp mismatch from merkle candidate {peer_id}");
1140 failures.push(format!("{peer_id}: timestamp mismatch"));
1141 continue;
1142 }
1143 let candidate_peer = PeerId::from_bytes(compute_address(&candidate.pub_key));
1148 if candidate_peer != peer_id {
1149 warn!(
1150 "Dropping merkle candidate {peer_id} — pub_key derives {candidate_peer}, \
1151 not the responding peer"
1152 );
1153 failures.push(format!("{peer_id}: candidate pub_key/peer mismatch"));
1154 continue;
1155 }
1156 if let Err(detail) =
1165 merkle_candidate_binding_is_valid(&candidate_peer, &candidate, &commitment)
1166 {
1167 warn!("Dropping merkle candidate {peer_id} — ADR-0004 binding invalid: {detail}");
1168 failures.push(format!("{peer_id}: bad commitment binding ({detail})"));
1169 continue;
1170 }
1171 valid.push((candidate_peer, candidate));
1172 }
1173 Err(e) => {
1174 debug!("Failed to get merkle candidate from {peer_id}: {e}");
1175 failures.push(format!("{peer_id}: {e}"));
1176 }
1177 }
1178 }
1179
1180 if valid.len() < CANDIDATES_PER_POOL {
1181 return Err(Error::InsufficientPeers(format!(
1182 "Got {} merkle candidates, need {CANDIDATES_PER_POOL}. Failures: [{}]",
1183 valid.len(),
1184 failures.join("; ")
1185 )));
1186 }
1187
1188 let target_peer = PeerId::from_bytes(*target_address);
1189 valid.sort_by_key(|(peer_id, _)| peer_id.xor_distance(&target_peer));
1190
1191 let candidates: Vec<MerklePaymentCandidateNode> = valid
1192 .into_iter()
1193 .take(CANDIDATES_PER_POOL)
1194 .map(|(_, candidate)| candidate)
1195 .collect();
1196
1197 let array: [MerklePaymentCandidateNode; CANDIDATES_PER_POOL] =
1198 candidates.try_into().map_err(|_| {
1199 Error::Payment("Failed to convert candidates to fixed array".to_string())
1200 })?;
1201 Ok(array)
1202 }
1203
1204 pub(crate) async fn merkle_upload_chunks(
1225 &self,
1226 chunk_contents: Vec<Bytes>,
1227 addresses: Vec<[u8; 32]>,
1228 batch_result: &MerkleBatchPaymentResult,
1229 progress: Option<&mpsc::Sender<UploadEvent>>,
1230 stored_offset: usize,
1231 total_chunks: usize,
1232 ) -> Result<MerkleStoreOutcome> {
1233 let store_limiter = self.controller().store.clone();
1234 let batch_size = chunk_contents.len();
1237 if batch_size != addresses.len() {
1238 return Err(Error::InvalidData(format!(
1239 "merkle upload has {batch_size} chunk contents but {} addresses",
1240 addresses.len()
1241 )));
1242 }
1243 let cap = || store_limiter.current().min(batch_size.max(1));
1247
1248 let bodies: std::collections::HashMap<[u8; 32], Bytes> =
1253 addresses.iter().copied().zip(chunk_contents).collect();
1254 let addrs = addresses;
1255
1256 let store_one = |addr: [u8; 32]| {
1261 let limiter = store_limiter.clone();
1262 let content = bodies.get(&addr).cloned();
1263 let proof_bytes = batch_result.proofs.get(&addr).cloned();
1264 async move {
1265 let started = std::time::Instant::now();
1266 let content = content.ok_or_else(|| {
1267 Error::InvalidData(format!("missing chunk body for {}", hex::encode(addr)))
1268 })?;
1269 let proof = proof_bytes.ok_or_else(|| {
1270 Error::Payment(format!(
1271 "Missing merkle proof for chunk {}",
1272 hex::encode(addr)
1273 ))
1274 })?;
1275 let peers = self.put_target_peers(&addr).await?;
1276 observe_op(
1277 &limiter,
1278 || async move { self.chunk_put_to_close_group(content, proof, &peers).await },
1279 classify_error,
1280 )
1281 .await
1282 .map(|_| started)
1283 }
1284 };
1285
1286 let outcome = merkle_store_with_retry(
1287 addrs,
1288 cap,
1289 MERKLE_STORE_MAX_ATTEMPTS,
1290 MERKLE_RETRY_BACKOFF,
1291 progress,
1292 stored_offset,
1293 total_chunks,
1294 store_one,
1295 )
1296 .await?;
1297
1298 if let Some(e) = outcome.fatal {
1306 return Err(e);
1307 }
1308 Ok(outcome)
1309 }
1310}
1311
1312pub(crate) const MERKLE_STORE_MAX_ATTEMPTS: usize = 4;
1324
1325pub(crate) const MERKLE_RETRY_BACKOFF: Duration = Duration::from_secs(30);
1331
1332const MERKLE_RETRY_JITTER: f64 = 0.1;
1335
1336#[derive(Debug, Default)]
1339pub(crate) struct MerkleStoreOutcome {
1340 pub stored: usize,
1343 pub stored_addresses: Vec<[u8; 32]>,
1350 pub failed: usize,
1352 pub failed_addresses: Vec<([u8; 32], String)>,
1357 pub fatal: Option<Error>,
1364 pub stats: crate::data::client::batch::WaveAggregateStats,
1366}
1367
1368#[allow(clippy::too_many_arguments)]
1394pub(crate) async fn merkle_store_with_retry<F, Fut, C>(
1395 addrs: Vec<[u8; 32]>,
1396 cap: C,
1397 max_attempts: usize,
1398 backoff: Duration,
1399 progress: Option<&mpsc::Sender<UploadEvent>>,
1400 stored_offset: usize,
1401 total: usize,
1402 store_one: F,
1403) -> Result<MerkleStoreOutcome>
1404where
1405 F: Fn([u8; 32]) -> Fut,
1406 Fut: std::future::Future<Output = Result<std::time::Instant>>,
1407 C: Fn() -> usize,
1408{
1409 let attempts = max_attempts.max(1);
1410 let mut outcome = MerkleStoreOutcome {
1411 stored: stored_offset,
1412 ..MerkleStoreOutcome::default()
1413 };
1414 let mut pending = addrs;
1415
1416 for attempt in 0..attempts {
1417 let mut next_failed: Vec<([u8; 32], String)> = Vec::new();
1423
1424 let mut pending_iter = pending.into_iter();
1429 let mut in_flight = FuturesUnordered::new();
1430 loop {
1431 let slots = cap().max(1);
1432 while in_flight.len() < slots {
1433 match pending_iter.next() {
1434 Some(addr) => {
1435 let fut = store_one(addr);
1436 in_flight.push(async move { (addr, fut.await) });
1437 }
1438 None => break,
1439 }
1440 }
1441 let Some((addr, result)) = in_flight.next().await else {
1442 break;
1443 };
1444 outcome.stats.chunk_attempts_total =
1445 outcome.stats.chunk_attempts_total.saturating_add(1);
1446 match result {
1447 Ok(started) => {
1448 let duration_ms =
1449 u64::try_from(started.elapsed().as_millis()).unwrap_or(u64::MAX);
1450 outcome.stats.store_durations_ms.push(duration_ms);
1451 let idx = attempt.min(outcome.stats.retries_histogram.len().saturating_sub(1));
1452 outcome.stats.retries_histogram[idx] =
1453 outcome.stats.retries_histogram[idx].saturating_add(1);
1454 outcome.stored += 1;
1455 outcome.stored_addresses.push(addr);
1456 if let Some(tx) = progress {
1457 let _ = tx.try_send(UploadEvent::ChunkStored {
1458 stored: outcome.stored,
1459 total,
1460 });
1461 }
1462 }
1463 Err(
1470 e @ (Error::InsufficientPeers(_)
1471 | Error::CloseGroupShortfall(_)
1472 | Error::RemotePut { .. }),
1473 ) => {
1474 next_failed.push((addr, e.to_string()));
1475 }
1476 Err(e) => {
1477 next_failed.push((addr, e.to_string()));
1485 outcome.fatal = Some(e);
1486 break;
1487 }
1488 }
1489 }
1490
1491 if outcome.fatal.is_some() {
1492 outcome.failed = next_failed.len();
1493 outcome.failed_addresses = next_failed;
1494 return Ok(outcome);
1495 }
1496
1497 if next_failed.is_empty() {
1498 break;
1499 }
1500
1501 if attempt + 1 < attempts {
1502 warn!(
1503 failed = next_failed.len(),
1504 attempt = attempt + 1,
1505 "merkle chunks short of quorum, retrying after backoff"
1506 );
1507 pending = next_failed.into_iter().map(|(addr, _msg)| addr).collect();
1508 if backoff > Duration::ZERO {
1509 let wait = {
1514 let mut rng = rand::thread_rng();
1515 let factor = 1.0 + rng.gen_range(-MERKLE_RETRY_JITTER..=MERKLE_RETRY_JITTER);
1516 backoff.mul_f64(factor)
1517 };
1518 tokio::time::sleep(wait).await;
1519 }
1520 } else {
1521 outcome.failed = next_failed.len();
1522 outcome.failed_addresses = next_failed;
1523 break;
1524 }
1525 }
1526
1527 Ok(outcome)
1528}
1529
1530pub(crate) const DEFERRED_ROUND_DELAYS_SECS: [u64; 3] = [0, 15, 45];
1539
1540pub(crate) fn deferred_round_histogram_slot(round: usize, hist_len: usize) -> usize {
1547 (round + 1).min(hist_len.saturating_sub(1))
1548}
1549
1550#[derive(Debug, Default)]
1552pub(crate) struct DeferredRetryOutcome {
1553 pub stored: usize,
1557 pub stored_addresses: Vec<[u8; 32]>,
1560 pub failed: usize,
1562 pub failed_addresses: Vec<([u8; 32], String)>,
1566 pub fatal: Option<String>,
1570 pub stats: crate::data::client::batch::WaveAggregateStats,
1573}
1574
1575#[allow(clippy::too_many_arguments)]
1593pub(crate) async fn merkle_deferred_retry<CF, SF, Fut>(
1594 deferred: Vec<([u8; 32], String)>,
1595 round_delays_secs: &[u64],
1596 concurrency_for: CF,
1597 progress: Option<&mpsc::Sender<UploadEvent>>,
1598 stored_offset: usize,
1599 total: usize,
1600 store_one: SF,
1601) -> Result<DeferredRetryOutcome>
1602where
1603 CF: Fn(usize) -> usize,
1604 SF: Fn([u8; 32]) -> Fut,
1605 Fut: std::future::Future<Output = Result<std::time::Instant>>,
1606{
1607 let mut outcome = DeferredRetryOutcome {
1608 stored: stored_offset,
1609 ..DeferredRetryOutcome::default()
1610 };
1611 let mut remaining = deferred;
1612 let rounds = round_delays_secs.len();
1613
1614 for (round, &delay_secs) in round_delays_secs.iter().enumerate() {
1615 if remaining.is_empty() {
1616 break;
1617 }
1618 if delay_secs > 0 {
1619 tokio::time::sleep(Duration::from_secs(delay_secs)).await;
1620 }
1621 info!(
1622 "Deferred merkle retry round {}/{}: {} chunk(s) short of quorum",
1623 round + 1,
1624 rounds,
1625 remaining.len(),
1626 );
1627
1628 let slot = deferred_round_histogram_slot(round, outcome.stats.retries_histogram.len());
1632 let round_addrs: Vec<[u8; 32]> = std::mem::take(&mut remaining)
1633 .into_iter()
1634 .map(|(addr, _msg)| addr)
1635 .collect();
1636 let round_len = round_addrs.len();
1637 let cap = || concurrency_for(round_len);
1640
1641 let round_outcome = merkle_store_with_retry(
1642 round_addrs,
1643 cap,
1644 1,
1645 Duration::ZERO,
1646 progress,
1647 outcome.stored,
1648 total,
1649 &store_one,
1650 )
1651 .await?;
1652
1653 outcome.stored = round_outcome.stored;
1654 outcome
1655 .stored_addresses
1656 .extend(round_outcome.stored_addresses);
1657
1658 outcome.stats.chunk_attempts_total = outcome
1660 .stats
1661 .chunk_attempts_total
1662 .saturating_add(round_outcome.stats.chunk_attempts_total);
1663 outcome
1664 .stats
1665 .store_durations_ms
1666 .extend(round_outcome.stats.store_durations_ms);
1667 let landed: usize = round_outcome.stats.retries_histogram.iter().sum();
1668 outcome.stats.retries_histogram[slot] =
1669 outcome.stats.retries_histogram[slot].saturating_add(landed);
1670
1671 if let Some(fatal) = round_outcome.fatal {
1672 outcome.fatal = Some(fatal.to_string());
1677 outcome.failed = round_outcome.failed_addresses.len();
1678 outcome.failed_addresses = round_outcome.failed_addresses;
1679 return Ok(outcome);
1680 }
1681
1682 remaining = round_outcome.failed_addresses;
1684 }
1685
1686 outcome.failed = remaining.len();
1687 outcome.failed_addresses = remaining;
1688 Ok(outcome)
1689}
1690
1691pub fn finalize_merkle_batch(
1696 prepared: PreparedMerkleBatch,
1697 winner_pool_hash: [u8; 32],
1698) -> Result<MerkleBatchPaymentResult> {
1699 let chunk_count = prepared.addresses.len();
1700 let xornames: Vec<XorName> = prepared.addresses.iter().map(|a| XorName(*a)).collect();
1701
1702 let winner_pool = prepared
1704 .candidate_pools
1705 .iter()
1706 .find(|pool| pool.hash() == winner_pool_hash)
1707 .ok_or_else(|| {
1708 Error::Payment(format!(
1709 "Winner pool {} not found in candidate pools",
1710 hex::encode(winner_pool_hash)
1711 ))
1712 })?;
1713
1714 info!("Generating merkle proofs for {chunk_count} chunks");
1725 let mut proofs = HashMap::with_capacity(chunk_count);
1726
1727 for (i, xorname) in xornames.iter().enumerate() {
1728 let address_proof = prepared
1729 .tree
1730 .generate_address_proof(i, *xorname)
1731 .map_err(|e| {
1732 Error::Payment(format!(
1733 "Failed to generate address proof for chunk {i}: {e}"
1734 ))
1735 })?;
1736
1737 let merkle_proof = MerklePaymentProof::new(*xorname, address_proof, winner_pool.clone());
1738
1739 let tagged_bytes = serialize_merkle_proof(&merkle_proof)
1740 .map_err(|e| Error::Serialization(format!("Failed to serialize merkle proof: {e}")))?;
1741
1742 proofs.insert(prepared.addresses[i], tagged_bytes);
1743 }
1744
1745 info!("Merkle batch payment complete: {chunk_count} proofs generated");
1746
1747 Ok(MerkleBatchPaymentResult {
1748 proofs,
1749 chunk_count,
1750 storage_cost_atto: "0".to_string(),
1751 gas_cost_wei: 0,
1752 merkle_payment_timestamp: prepared.merkle_payment_timestamp,
1753 })
1754}
1755
1756#[cfg(test)]
1758mod send_assertions {
1759 use super::*;
1760 use crate::data::client::Client;
1761
1762 fn _assert_send<T: Send>(_: &T) {}
1763
1764 #[allow(
1765 dead_code,
1766 unreachable_code,
1767 unused_variables,
1768 clippy::diverging_sub_expression
1769 )]
1770 async fn _merkle_upload_chunks_is_send(client: &Client) {
1771 let batch_result: MerkleBatchPaymentResult = todo!();
1772 let fut = client.merkle_upload_chunks(Vec::new(), Vec::new(), &batch_result, None, 0, 0);
1773 _assert_send(&fut);
1774 }
1775}
1776
1777#[cfg(test)]
1780#[allow(clippy::unwrap_used, clippy::expect_used)]
1781pub(crate) mod test_support {
1782 use super::*;
1783 use ant_protocol::evm::RewardsAddress;
1784
1785 pub(crate) fn make_test_addresses(count: usize) -> Vec<[u8; 32]> {
1786 (0..count)
1787 .map(|i| {
1788 let xn = XorName::from_content(&i.to_le_bytes());
1789 xn.0
1790 })
1791 .collect()
1792 }
1793
1794 pub(crate) fn make_dummy_candidate_nodes(
1795 timestamp: u64,
1796 ) -> [MerklePaymentCandidateNode; CANDIDATES_PER_POOL] {
1797 std::array::from_fn(|i| MerklePaymentCandidateNode {
1798 pub_key: vec![i as u8; 32],
1799 price: Amount::from(1024u64),
1800 reward_address: RewardsAddress::new([i as u8; 20]),
1801 merkle_payment_timestamp: timestamp,
1802 signature: vec![i as u8; 64],
1803 committed_key_count: 0,
1804 commitment_pin: None,
1805 })
1806 }
1807
1808 pub(crate) fn make_prepared_merkle_batch(count: usize) -> PreparedMerkleBatch {
1809 let addrs = make_test_addresses(count);
1810 let xornames: Vec<XorName> = addrs.iter().map(|a| XorName(*a)).collect();
1811 let tree = MerkleTree::from_xornames(xornames).unwrap();
1812
1813 let timestamp = std::time::SystemTime::now()
1814 .duration_since(std::time::UNIX_EPOCH)
1815 .unwrap()
1816 .as_secs();
1817
1818 let midpoints = tree.reward_candidates(timestamp).unwrap();
1819
1820 let candidate_pools: Vec<MerklePaymentCandidatePool> = midpoints
1821 .into_iter()
1822 .map(|mp| MerklePaymentCandidatePool {
1823 midpoint_proof: mp,
1824 candidate_nodes: make_dummy_candidate_nodes(timestamp),
1825 })
1826 .collect();
1827
1828 let pool_commitments = candidate_pools
1829 .iter()
1830 .map(pool_commitment_with_payment_multiplier)
1831 .collect::<Result<Vec<_>>>()
1832 .unwrap();
1833
1834 PreparedMerkleBatch {
1835 depth: tree.depth(),
1836 pool_commitments,
1837 merkle_payment_timestamp: timestamp,
1838 candidate_pools,
1839 tree,
1840 addresses: addrs,
1841 }
1842 }
1843
1844 pub(crate) fn winner_hash_for(batch: &PreparedMerkleBatch) -> [u8; 32] {
1849 batch.candidate_pools[0].hash()
1850 }
1851}
1852
1853#[cfg(test)]
1854#[allow(clippy::unwrap_used, clippy::expect_used, clippy::panic)]
1855mod tests {
1856 use super::test_support::*;
1857 use super::*;
1858 use ant_protocol::evm::{Amount, MerkleTree, RewardsAddress, CANDIDATES_PER_POOL};
1859
1860 #[test]
1865 fn test_auto_below_threshold() {
1866 assert!(!should_use_merkle(1, PaymentMode::Auto));
1867 assert!(!should_use_merkle(10, PaymentMode::Auto));
1868 assert!(!should_use_merkle(63, PaymentMode::Auto));
1869 }
1870
1871 #[test]
1872 fn test_auto_at_and_above_threshold() {
1873 assert!(should_use_merkle(64, PaymentMode::Auto));
1874 assert!(should_use_merkle(65, PaymentMode::Auto));
1875 assert!(should_use_merkle(1000, PaymentMode::Auto));
1876 }
1877
1878 #[test]
1879 fn test_merkle_mode_forces_at_2() {
1880 assert!(!should_use_merkle(1, PaymentMode::Merkle));
1881 assert!(should_use_merkle(2, PaymentMode::Merkle));
1882 assert!(should_use_merkle(3, PaymentMode::Merkle));
1883 }
1884
1885 #[test]
1886 fn test_single_mode_always_false() {
1887 assert!(!should_use_merkle(0, PaymentMode::Single));
1888 assert!(!should_use_merkle(64, PaymentMode::Single));
1889 assert!(!should_use_merkle(1000, PaymentMode::Single));
1890 }
1891
1892 #[test]
1893 fn test_default_mode_is_auto() {
1894 assert_eq!(PaymentMode::default(), PaymentMode::Auto);
1895 }
1896
1897 #[test]
1898 fn test_threshold_value() {
1899 assert_eq!(DEFAULT_MERKLE_THRESHOLD, 64);
1900 }
1901
1902 #[test]
1907 fn test_preflight_quotes_gathered_means_not_stored() {
1908 assert!(matches!(preflight_stored_status(Ok(())), Ok(false)));
1909 }
1910
1911 #[test]
1912 fn test_preflight_already_stored_is_stored() {
1913 let r: Result<()> = Err(Error::AlreadyStored);
1914 assert!(matches!(preflight_stored_status(r), Ok(true)));
1915 }
1916
1917 #[test]
1921 fn test_preflight_transient_quote_failure_does_not_abort() {
1922 let insufficient: Result<()> =
1924 Err(Error::InsufficientPeers("Got 5 quotes, need 7".to_string()));
1925 assert!(
1926 matches!(preflight_stored_status(insufficient), Ok(false)),
1927 "insufficient-peers during preflight must degrade to not-stored, not error"
1928 );
1929
1930 let timeout: Result<()> = Err(Error::Timeout("Timeout waiting for quote".to_string()));
1931 assert!(matches!(preflight_stored_status(timeout), Ok(false)));
1932
1933 let network: Result<()> = Err(Error::Network("connection reset".to_string()));
1934 assert!(matches!(preflight_stored_status(network), Ok(false)));
1935 }
1936
1937 #[test]
1940 fn test_preflight_application_error_propagates() {
1941 let payment: Result<()> = Err(Error::Payment("bad payment".to_string()));
1942 assert!(matches!(
1943 preflight_stored_status(payment),
1944 Err(Error::Payment(_))
1945 ));
1946 }
1947
1948 #[test]
1949 fn chunk_contents_for_upload_addresses_preserves_requested_order() {
1950 let first = Bytes::from_static(b"first");
1951 let second = Bytes::from_static(b"second");
1952 let first_addr = compute_address(&first);
1953 let second_addr = compute_address(&second);
1954
1955 let selected = chunk_contents_for_upload_addresses(
1956 vec![first.clone(), second.clone()],
1957 &[second_addr, first_addr],
1958 )
1959 .unwrap();
1960
1961 assert_eq!(selected, vec![second, first]);
1962 }
1963
1964 #[test]
1965 fn chunk_contents_for_upload_addresses_preserves_duplicate_requests() {
1966 let repeated = Bytes::from_static(b"same-content");
1967 let other = Bytes::from_static(b"other-content");
1968 let repeated_addr = compute_address(&repeated);
1969
1970 let selected = chunk_contents_for_upload_addresses(
1971 vec![repeated.clone(), other, repeated.clone()],
1972 &[repeated_addr, repeated_addr],
1973 )
1974 .unwrap();
1975
1976 assert_eq!(selected, vec![repeated.clone(), repeated]);
1977 }
1978
1979 #[test]
1980 fn chunk_contents_for_upload_addresses_ignores_unrequested_duplicates() {
1981 let requested = Bytes::from_static(b"requested-content");
1982 let unrequested = Bytes::from_static(b"unrequested-content");
1983 let requested_addr = compute_address(&requested);
1984
1985 let selected = chunk_contents_for_upload_addresses(
1986 vec![
1987 unrequested.clone(),
1988 requested.clone(),
1989 unrequested.clone(),
1990 unrequested,
1991 ],
1992 &[requested_addr],
1993 )
1994 .unwrap();
1995
1996 assert_eq!(selected, vec![requested]);
1997 }
1998
1999 #[test]
2000 fn chunk_contents_for_upload_addresses_errors_for_missing_content() {
2001 let present = Bytes::from_static(b"present-content");
2002 let missing = Bytes::from_static(b"missing-content");
2003 let missing_addr = compute_address(&missing);
2004
2005 let result = chunk_contents_for_upload_addresses(vec![present], &[missing_addr]);
2006
2007 assert!(matches!(result, Err(Error::InvalidData(_))));
2008 }
2009
2010 #[test]
2015 fn test_tree_depth_for_known_sizes() {
2016 let cases = [(2, 1), (4, 2), (16, 4), (100, 7), (256, 8)];
2017 for (count, expected_depth) in cases {
2018 let addrs = make_test_addresses(count);
2019 let xornames: Vec<XorName> = addrs.iter().map(|a| XorName(*a)).collect();
2020 let tree = MerkleTree::from_xornames(xornames).unwrap();
2021 assert_eq!(
2022 tree.depth(),
2023 expected_depth,
2024 "depth mismatch for {count} leaves"
2025 );
2026 }
2027 }
2028
2029 #[test]
2030 fn test_proof_generation_and_verification_for_all_leaves() {
2031 let addrs = make_test_addresses(16);
2032 let xornames: Vec<XorName> = addrs.iter().map(|a| XorName(*a)).collect();
2033 let tree = MerkleTree::from_xornames(xornames.clone()).unwrap();
2034
2035 for (i, xn) in xornames.iter().enumerate() {
2036 let proof = tree.generate_address_proof(i, *xn).unwrap();
2037 assert!(proof.verify(), "proof for leaf {i} should verify");
2038 assert_eq!(proof.depth(), tree.depth() as usize);
2039 }
2040 }
2041
2042 #[test]
2043 fn test_proof_fails_for_wrong_address() {
2044 let addrs = make_test_addresses(8);
2045 let xornames: Vec<XorName> = addrs.iter().map(|a| XorName(*a)).collect();
2046 let tree = MerkleTree::from_xornames(xornames).unwrap();
2047
2048 let wrong = XorName::from_content(b"wrong");
2049 let proof = tree.generate_address_proof(0, wrong).unwrap();
2050 assert!(!proof.verify(), "proof with wrong address should fail");
2051 }
2052
2053 #[test]
2054 fn test_tree_too_few_leaves() {
2055 let xornames = vec![XorName::from_content(b"only_one")];
2056 let result = MerkleTree::from_xornames(xornames);
2057 assert!(result.is_err());
2058 }
2059
2060 #[test]
2061 fn test_tree_at_max_leaves() {
2062 let addrs = make_test_addresses(MAX_LEAVES);
2063 let xornames: Vec<XorName> = addrs.iter().map(|a| XorName(*a)).collect();
2064 let tree = MerkleTree::from_xornames(xornames).unwrap();
2065 assert_eq!(tree.leaf_count(), MAX_LEAVES);
2066 }
2067
2068 #[test]
2073 fn test_merkle_proof_serialize_deserialize_roundtrip() {
2074 use ant_protocol::evm::{Amount, MerklePaymentCandidateNode, RewardsAddress};
2075 use ant_protocol::payment::{deserialize_merkle_proof, serialize_merkle_proof};
2076
2077 let addrs = make_test_addresses(4);
2078 let xornames: Vec<XorName> = addrs.iter().map(|a| XorName(*a)).collect();
2079 let tree = MerkleTree::from_xornames(xornames.clone()).unwrap();
2080
2081 let timestamp = std::time::SystemTime::now()
2082 .duration_since(std::time::UNIX_EPOCH)
2083 .unwrap()
2084 .as_secs();
2085
2086 let candidates = tree.reward_candidates(timestamp).unwrap();
2087 let midpoint = candidates.first().unwrap().clone();
2088
2089 #[allow(clippy::cast_possible_truncation)]
2091 let candidate_nodes: [MerklePaymentCandidateNode; CANDIDATES_PER_POOL] =
2092 std::array::from_fn(|i| MerklePaymentCandidateNode {
2093 pub_key: vec![i as u8; 32],
2094 price: Amount::from(1024u64),
2095 reward_address: RewardsAddress::new([i as u8; 20]),
2096 merkle_payment_timestamp: timestamp,
2097 signature: vec![i as u8; 64],
2098 committed_key_count: 0,
2099 commitment_pin: None,
2100 });
2101
2102 let pool = MerklePaymentCandidatePool {
2103 midpoint_proof: midpoint,
2104 candidate_nodes,
2105 };
2106
2107 let address_proof = tree.generate_address_proof(0, xornames[0]).unwrap();
2108 let merkle_proof = MerklePaymentProof::new(xornames[0], address_proof, pool);
2109
2110 let tagged = serialize_merkle_proof(&merkle_proof).unwrap();
2111 assert_eq!(
2112 tagged.first().copied(),
2113 Some(0x02),
2114 "tag should be PROOF_TAG_MERKLE"
2115 );
2116
2117 let deserialized = deserialize_merkle_proof(&tagged).unwrap();
2118 assert_eq!(deserialized.address, merkle_proof.address);
2119 assert_eq!(
2120 deserialized.winner_pool.candidate_nodes.len(),
2121 CANDIDATES_PER_POOL
2122 );
2123 }
2124
2125 #[test]
2130 fn test_candidate_wrong_timestamp_rejected() {
2131 let candidate = MerklePaymentCandidateNode {
2133 pub_key: vec![0u8; 32],
2134 price: ant_protocol::evm::Amount::ZERO,
2135 reward_address: ant_protocol::evm::RewardsAddress::new([0u8; 20]),
2136 merkle_payment_timestamp: 1000,
2137 signature: vec![0u8; 64],
2138 committed_key_count: 0,
2139 commitment_pin: None,
2140 };
2141
2142 assert_ne!(candidate.merkle_payment_timestamp, 2000);
2144 }
2145
2146 fn pool_with_varied_prices(timestamp: u64) -> MerklePaymentCandidatePool {
2153 let addrs = make_test_addresses(4);
2154 let xornames: Vec<XorName> = addrs.iter().map(|a| XorName(*a)).collect();
2155 let tree = MerkleTree::from_xornames(xornames).unwrap();
2156 let midpoint = tree
2157 .reward_candidates(timestamp)
2158 .unwrap()
2159 .into_iter()
2160 .next()
2161 .unwrap();
2162
2163 let candidate_nodes = std::array::from_fn(|i| MerklePaymentCandidateNode {
2164 pub_key: vec![i as u8; 32],
2165 price: Amount::from((i as u64 + 1) * 100),
2167 reward_address: RewardsAddress::new([i as u8; 20]),
2168 merkle_payment_timestamp: timestamp,
2169 signature: vec![i as u8; 64],
2170 committed_key_count: 0,
2171 commitment_pin: None,
2172 });
2173
2174 MerklePaymentCandidatePool {
2175 midpoint_proof: midpoint,
2176 candidate_nodes,
2177 }
2178 }
2179
2180 fn median16(mut amounts: Vec<Amount>) -> Amount {
2182 amounts.sort_unstable();
2183 *amounts.get(amounts.len() / 2).unwrap()
2184 }
2185
2186 #[test]
2187 fn pool_commitment_applies_payment_multiplier_to_every_candidate() {
2188 let pool = pool_with_varied_prices(1_700_000_000);
2189 let commitment = pool_commitment_with_payment_multiplier(&pool).unwrap();
2190
2191 for (candidate, signed) in commitment
2192 .candidates
2193 .iter()
2194 .zip(pool.candidate_nodes.iter())
2195 {
2196 assert_eq!(
2197 candidate.price,
2198 signed.price * Amount::from(MERKLE_PAYMENT_MULTIPLIER),
2199 "on-chain payable amount must be {MERKLE_PAYMENT_MULTIPLIER}x the quoted price"
2200 );
2201 }
2202 }
2203
2204 #[test]
2205 fn pool_commitment_multiplier_leaves_signed_prices_and_pool_hash_untouched() {
2206 let pool = pool_with_varied_prices(1_700_000_000);
2207 let before: Vec<Amount> = pool.candidate_nodes.iter().map(|c| c.price).collect();
2208
2209 let commitment = pool_commitment_with_payment_multiplier(&pool).unwrap();
2210
2211 let after: Vec<Amount> = pool.candidate_nodes.iter().map(|c| c.price).collect();
2212 assert_eq!(before, after, "signed candidate prices must not change");
2213 assert_eq!(
2214 commitment.pool_hash,
2215 pool.hash(),
2216 "pool hash is the storer's on-chain lookup key and must be \
2217 computed over the signed 1x prices"
2218 );
2219 }
2220
2221 #[test]
2230 fn merkle_settlement_per_padded_leaf_is_the_multiplied_pool_median() {
2231 let pool = pool_with_varied_prices(1_700_000_000);
2232 let commitment = pool_commitment_with_payment_multiplier(&pool).unwrap();
2233
2234 let quoted_median = median16(pool.candidate_nodes.iter().map(|c| c.price).collect());
2235 let per_chunk = median16(commitment.candidates.iter().map(|c| c.price).collect());
2236
2237 assert_eq!(quoted_median, Amount::from(900u64));
2238 assert_eq!(
2239 per_chunk,
2240 quoted_median * Amount::from(MERKLE_PAYMENT_MULTIPLIER),
2241 "merkle per-chunk settlement must equal the single-node \
2242 {MERKLE_PAYMENT_MULTIPLIER}x median, not the bare quoted price"
2243 );
2244 }
2245
2246 #[test]
2247 fn test_finalize_merkle_batch_with_valid_winner() {
2248 let prepared = make_prepared_merkle_batch(4);
2249 let winner_hash = prepared.candidate_pools[0].hash();
2250
2251 let result = finalize_merkle_batch(prepared, winner_hash);
2252 assert!(
2253 result.is_ok(),
2254 "should succeed with valid winner: {result:?}"
2255 );
2256
2257 let batch = result.unwrap();
2258 assert_eq!(batch.chunk_count, 4);
2259 assert_eq!(batch.proofs.len(), 4);
2260
2261 for proof_bytes in batch.proofs.values() {
2263 assert!(!proof_bytes.is_empty());
2264 }
2265 }
2266
2267 #[test]
2275 fn test_finalize_merkle_batch_ships_no_commitment_sidecars() {
2276 use ant_protocol::payment::deserialize_merkle_proof;
2277
2278 let mut prepared = make_prepared_merkle_batch(4);
2279 for pool in &mut prepared.candidate_pools {
2282 for candidate in &mut pool.candidate_nodes {
2283 candidate.committed_key_count = 9_000;
2284 candidate.commitment_pin = Some([7u8; 32]);
2285 }
2286 }
2287 let winner_hash = prepared.candidate_pools[0].hash();
2288
2289 let batch = finalize_merkle_batch(prepared, winner_hash).unwrap();
2290 assert_eq!(batch.proofs.len(), 4);
2291 for proof_bytes in batch.proofs.values() {
2292 let proof = deserialize_merkle_proof(proof_bytes).unwrap();
2293 assert!(
2294 proof.commitment_sidecars.is_empty(),
2295 "per-chunk merkle proofs must not ship commitment sidecars"
2296 );
2297 }
2298 }
2299
2300 #[test]
2301 fn test_finalize_merkle_batch_with_invalid_winner() {
2302 let prepared = make_prepared_merkle_batch(4);
2303 let bad_hash = [0xFF; 32];
2304
2305 let result = finalize_merkle_batch(prepared, bad_hash);
2306 assert!(result.is_err());
2307 let err = result.unwrap_err().to_string();
2308 assert!(err.contains("not found in candidate pools"), "got: {err}");
2309 }
2310
2311 #[test]
2312 fn test_finalize_merkle_batch_proofs_are_deserializable() {
2313 use ant_protocol::payment::deserialize_merkle_proof;
2314
2315 let prepared = make_prepared_merkle_batch(8);
2316 let winner_hash = prepared.candidate_pools[0].hash();
2317
2318 let batch = finalize_merkle_batch(prepared, winner_hash).unwrap();
2319
2320 for (addr, proof_bytes) in &batch.proofs {
2321 let proof = deserialize_merkle_proof(proof_bytes);
2322 assert!(
2323 proof.is_ok(),
2324 "proof for {} should deserialize: {:?}",
2325 hex::encode(addr),
2326 proof.err()
2327 );
2328 }
2329 }
2330
2331 const PARTITION_CASES: [(usize, &[usize]); 10] = [
2340 (2, &[2]),
2341 (64, &[64]),
2342 (65, &[65]),
2343 (255, &[255]),
2344 (256, &[256]),
2345 (257, &[255, 2]),
2346 (300, &[256, 44]),
2347 (512, &[256, 256]),
2348 (513, &[256, 255, 2]),
2349 (769, &[256, 256, 255, 2]),
2350 ];
2351
2352 #[test]
2353 fn merkle_batch_sizes_rebalance_singleton_remainders() {
2354 for (total, expected) in PARTITION_CASES {
2355 assert_eq!(
2356 merkle_batch_sizes(total),
2357 expected,
2358 "{total} addresses must partition as {expected:?}"
2359 );
2360 }
2361 }
2362
2363 #[test]
2367 fn merkle_batch_sizes_with_cap_partitions_and_clamps() {
2368 assert_eq!(merkle_batch_sizes_with_cap(6, 3), vec![3, 3]);
2370 assert_eq!(merkle_batch_sizes_with_cap(7, 3), vec![3, 2, 2]);
2371 assert_eq!(merkle_batch_sizes_with_cap(4, 3), vec![2, 2]);
2372 assert_eq!(merkle_batch_sizes_with_cap(5, 2), vec![3, 2]);
2374 assert_eq!(
2376 merkle_batch_sizes_with_cap(MAX_LEAVES + 1, MAX_LEAVES * 4),
2377 vec![MAX_LEAVES - 1, 2]
2378 );
2379 for total in 2..200usize {
2382 let sizes = merkle_batch_sizes_with_cap(total, 3);
2383 assert_eq!(sizes.iter().sum::<usize>(), total, "cover for {total}");
2384 assert!(
2385 sizes.iter().all(|&s| (2..=3).contains(&s)),
2386 "unpayable part for {total}: {sizes:?}"
2387 );
2388 }
2389 }
2390
2391 #[test]
2395 fn merge_merkle_batch_results_unions_proofs_and_keeps_oldest_timestamp() {
2396 let a = MerkleBatchPaymentResult {
2397 proofs: [([1u8; 32], vec![1u8])].into_iter().collect(),
2398 chunk_count: 1,
2399 storage_cost_atto: "100".into(),
2400 gas_cost_wei: 7,
2401 merkle_payment_timestamp: 2_000,
2402 };
2403 let b = MerkleBatchPaymentResult {
2404 proofs: [([2u8; 32], vec![2u8]), ([3u8; 32], vec![3u8])]
2405 .into_iter()
2406 .collect(),
2407 chunk_count: 2,
2408 storage_cost_atto: "50".into(),
2409 gas_cost_wei: 5,
2410 merkle_payment_timestamp: 1_500,
2411 };
2412 let merged = merge_merkle_batch_results(vec![a, b]);
2413 assert_eq!(merged.proofs.len(), 3);
2414 assert_eq!(merged.chunk_count, 3);
2415 assert_eq!(merged.storage_cost_atto, "150");
2416 assert_eq!(merged.gas_cost_wei, 12);
2417 assert_eq!(merged.merkle_payment_timestamp, 1_500);
2418 }
2419
2420 #[test]
2424 fn merkle_batch_sizes_are_always_buildable_trees() {
2425 for total in 2..=(4 * MAX_LEAVES + 3) {
2426 let sizes = merkle_batch_sizes(total);
2427 assert!(!sizes.is_empty(), "{total} addresses must produce batches");
2428 assert_eq!(
2429 sizes.iter().sum::<usize>(),
2430 total,
2431 "{total} addresses: partition must cover every address"
2432 );
2433 for size in sizes {
2434 assert!(
2435 (2..=MAX_LEAVES).contains(&size),
2436 "{total} addresses produced a batch of {size}, outside 2..={MAX_LEAVES}"
2437 );
2438 }
2439 }
2440 }
2441
2442 #[test]
2443 fn merkle_batch_sizes_below_two_have_no_payable_partition() {
2444 assert!(merkle_batch_sizes(0).is_empty());
2445 assert!(merkle_batch_sizes(1).is_empty());
2446 }
2447
2448 #[test]
2449 fn merkle_batch_partitions_preserve_order_and_use_each_address_once() {
2450 for (total, _) in PARTITION_CASES {
2451 let addrs = make_test_addresses(total);
2452 let partitions = merkle_batch_partitions(&addrs);
2453
2454 let flattened: Vec<[u8; 32]> = partitions.concat();
2455 assert_eq!(
2456 flattened, addrs,
2457 "{total} addresses: partitions must concatenate back to the input in order"
2458 );
2459
2460 let unique: std::collections::HashSet<[u8; 32]> = flattened.iter().copied().collect();
2461 assert_eq!(
2462 unique.len(),
2463 total,
2464 "{total} addresses: no address may be duplicated or synthesised"
2465 );
2466 }
2467 }
2468
2469 #[test]
2473 fn post_preflight_plan_of_257_partitions_into_payable_batches() {
2474 let plan = MerkleUploadPlan {
2475 already_stored: make_test_addresses(3),
2476 to_upload: make_test_addresses(257),
2477 to_upload_total_bytes: 257 * 1024,
2478 };
2479 assert_eq!(plan.to_upload.len(), 257);
2480
2481 let partitions = merkle_batch_partitions(&plan.to_upload);
2482 let sizes: Vec<usize> = partitions.iter().map(|batch| batch.len()).collect();
2483 assert_eq!(sizes, vec![255, 2]);
2484 for batch in partitions {
2485 let xornames: Vec<XorName> = batch.iter().map(|a| XorName(*a)).collect();
2486 assert!(
2487 MerkleTree::from_xornames(xornames).is_ok(),
2488 "every partition of a 257-chunk plan must build a tree"
2489 );
2490 }
2491 }
2492
2493 #[test]
2497 fn no_partition_pays_before_a_singleton_tree_failure() {
2498 for total in [257usize, 513, 769] {
2499 let addrs = make_test_addresses(total);
2500 for batch in merkle_batch_partitions(&addrs) {
2501 let xornames: Vec<XorName> = batch.iter().map(|a| XorName(*a)).collect();
2502 assert!(
2503 MerkleTree::from_xornames(xornames).is_ok(),
2504 "{total} addresses: batch of {} is unpayable",
2505 batch.len()
2506 );
2507 }
2508 }
2509 }
2510
2511 #[test]
2512 fn merkle_billable_leaves_sum_the_padded_partitions() {
2513 for (total, expected) in PARTITION_CASES {
2514 let padded: u64 = expected
2515 .iter()
2516 .map(|size| size.next_power_of_two() as u64)
2517 .sum();
2518 assert_eq!(
2519 merkle_billable_leaves(total as u64),
2520 padded,
2521 "{total} chunks must bill for the padded partition {expected:?}"
2522 );
2523 }
2524
2525 assert_eq!(merkle_billable_leaves(65), 128);
2527 assert_eq!(merkle_billable_leaves(257), 256 + 2);
2528 assert_eq!(merkle_billable_leaves(300), 256 + 64);
2529 assert_eq!(merkle_billable_leaves(0), 0);
2532 assert_eq!(merkle_billable_leaves(1), 2);
2533 }
2534
2535 #[test]
2536 fn merkle_billable_leaves_never_under_quote() {
2537 for chunks in 1..2000u64 {
2538 assert!(
2539 merkle_billable_leaves(chunks) >= chunks,
2540 "{chunks} chunks must never be billed as fewer leaves"
2541 );
2542 }
2543 }
2544
2545 #[test]
2549 fn external_preparation_refuses_more_than_one_tree_of_addresses() {
2550 assert!(ensure_single_merkle_tree_batch(2).is_ok());
2551 assert!(ensure_single_merkle_tree_batch(MAX_LEAVES).is_ok());
2552
2553 for oversized in [MAX_LEAVES + 1, 300, 513] {
2554 match ensure_single_merkle_tree_batch(oversized) {
2555 Err(Error::MerkleBatchTooLarge {
2556 addresses,
2557 max_leaves,
2558 }) => {
2559 assert_eq!(addresses, oversized);
2560 assert_eq!(max_leaves, MAX_LEAVES);
2561 }
2562 other => panic!("{oversized} addresses should be refused, got {other:?}"),
2563 }
2564 }
2565 }
2566
2567 use std::sync::{Arc, Mutex};
2572
2573 fn make_addrs(count: usize) -> Vec<[u8; 32]> {
2576 make_test_addresses(count)
2577 }
2578
2579 #[tokio::test]
2583 async fn store_with_retry_collects_failures_instead_of_aborting() {
2584 let chunks = make_addrs(6);
2585 let failing: std::collections::HashSet<[u8; 32]> = chunks.iter().take(2).copied().collect();
2586 let failing_for_closure = failing.clone();
2587
2588 let store_one = move |addr: [u8; 32]| {
2589 let fail = failing_for_closure.contains(&addr);
2590 async move {
2591 if fail {
2592 Err(Error::InsufficientPeers("test shortfall".into()))
2593 } else {
2594 Ok(std::time::Instant::now())
2595 }
2596 }
2597 };
2598
2599 let outcome =
2600 merkle_store_with_retry(chunks, || 8, 1, Duration::ZERO, None, 0, 6, store_one)
2601 .await
2602 .expect("quorum shortfalls must not abort the batch");
2603
2604 assert_eq!(outcome.stored, 4);
2605 assert_eq!(outcome.failed, 2);
2606 assert_eq!(outcome.stats.retries_histogram[0], 4);
2608 assert_eq!(outcome.stats.chunk_attempts_total, 6);
2609 }
2610
2611 #[tokio::test]
2619 async fn quorum_shortfall_survives_deferred_retries_with_exact_accounting() {
2620 let chunks = make_addrs(5);
2621 let short: std::collections::HashSet<[u8; 32]> = chunks.iter().take(2).copied().collect();
2622 let short_for_closure = short.clone();
2623 let store_one = move |addr: [u8; 32]| {
2624 let fail = short_for_closure.contains(&addr);
2625 async move {
2626 if fail {
2627 Err(Error::InsufficientPeers("still short of quorum".into()))
2628 } else {
2629 Ok(std::time::Instant::now())
2630 }
2631 }
2632 };
2633
2634 let pass = merkle_store_with_retry(
2637 chunks.clone(),
2638 || 8,
2639 1,
2640 Duration::ZERO,
2641 None,
2642 0,
2643 5,
2644 &store_one,
2645 )
2646 .await
2647 .expect("quorum shortfalls must not abort the pass");
2648 assert!(pass.fatal.is_none());
2649 assert_eq!(pass.stored, 3);
2650 assert_eq!(pass.failed, 2);
2651
2652 let dr = merkle_deferred_retry(
2655 pass.failed_addresses.clone(),
2656 &[0, 0, 0],
2657 |n: usize| n.max(1),
2658 None,
2659 pass.stored,
2660 5,
2661 &store_one,
2662 )
2663 .await
2664 .expect("deferred shortfalls must not abort");
2665
2666 assert!(dr.fatal.is_none());
2667 assert_eq!(
2668 dr.stored + dr.failed,
2669 5,
2670 "stored + failed must account for every chunk"
2671 );
2672 assert_eq!(dr.stored, 3, "paid-and-stored chunks must stay counted");
2673 assert_eq!(dr.failed, 2);
2674 let failed_set: std::collections::HashSet<[u8; 32]> =
2675 dr.failed_addresses.iter().map(|(a, _)| *a).collect();
2676 assert_eq!(
2677 failed_set, short,
2678 "failed set must be exactly the shortfall chunks"
2679 );
2680 }
2681
2682 #[tokio::test]
2688 async fn store_with_retry_rereads_cap_per_slot() {
2689 let count = 6;
2690 let chunks = make_addrs(count);
2691 let cap_calls = Arc::new(Mutex::new(0usize));
2692 let cap_calls_for_closure = cap_calls.clone();
2693 let cap = move || {
2694 *cap_calls_for_closure.lock().expect("cap counter poisoned") += 1;
2695 2
2696 };
2697 let store_one = move |_addr: [u8; 32]| async move { Ok(std::time::Instant::now()) };
2698
2699 let outcome =
2700 merkle_store_with_retry(chunks, cap, 1, Duration::ZERO, None, 0, count, store_one)
2701 .await
2702 .expect("all stores succeed");
2703
2704 assert_eq!(outcome.stored, count);
2705 let calls = *cap_calls.lock().expect("cap counter poisoned");
2706 assert!(
2707 calls >= count,
2708 "cap must be re-read per drained slot (rolling), not snapshotted once — \
2709 expected >= {count} invocations, got {calls}",
2710 );
2711 }
2712
2713 #[tokio::test]
2720 async fn store_pass_has_no_barrier() {
2721 use std::sync::atomic::{AtomicUsize, Ordering};
2722 let count = 8;
2723 let addrs = make_addrs(count);
2724 let slow = addrs[0];
2725 let fast_completed = Arc::new(AtomicUsize::new(0));
2726 let release_slow = Arc::new(tokio::sync::Notify::new());
2727
2728 let store_one = move |addr: [u8; 32]| {
2729 let fast_completed = fast_completed.clone();
2730 let release_slow = release_slow.clone();
2731 async move {
2732 if addr == slow {
2733 release_slow.notified().await;
2737 } else if fast_completed.fetch_add(1, Ordering::SeqCst) + 1 == count - 1 {
2738 release_slow.notify_one();
2739 }
2740 Ok(std::time::Instant::now())
2741 }
2742 };
2743
2744 let outcome = tokio::time::timeout(
2745 Duration::from_secs(5),
2746 merkle_store_with_retry(addrs, || 8, 1, Duration::ZERO, None, 0, count, store_one),
2747 )
2748 .await
2749 .expect("store pass must not deadlock — a slow chunk must not block the others")
2750 .expect("all stores succeed");
2751
2752 assert_eq!(outcome.stored, count);
2753 }
2754
2755 #[tokio::test]
2760 async fn store_pass_keeps_at_most_cap_in_flight() {
2761 use std::sync::atomic::{AtomicUsize, Ordering};
2762 let count = 40;
2763 let cap = 4;
2764 let addrs = make_addrs(count);
2765 let in_flight = Arc::new(AtomicUsize::new(0));
2766 let max_in_flight = Arc::new(AtomicUsize::new(0));
2767 let max_in_flight_for_closure = max_in_flight.clone();
2768
2769 let store_one = move |_addr: [u8; 32]| {
2770 let in_flight = in_flight.clone();
2771 let max_in_flight = max_in_flight_for_closure.clone();
2772 async move {
2773 let now = in_flight.fetch_add(1, Ordering::SeqCst) + 1;
2774 max_in_flight.fetch_max(now, Ordering::SeqCst);
2775 tokio::task::yield_now().await;
2778 in_flight.fetch_sub(1, Ordering::SeqCst);
2779 Ok(std::time::Instant::now())
2780 }
2781 };
2782
2783 let outcome = merkle_store_with_retry(
2784 addrs,
2785 move || cap,
2786 1,
2787 Duration::ZERO,
2788 None,
2789 0,
2790 count,
2791 store_one,
2792 )
2793 .await
2794 .expect("all stores succeed");
2795
2796 assert_eq!(outcome.stored, count);
2797 let peak = max_in_flight.load(Ordering::SeqCst);
2798 assert!(
2799 peak <= cap,
2800 "at most `cap` bodies may be in flight (memory bound), got peak {peak} > cap {cap}",
2801 );
2802 assert!(
2803 peak > 1,
2804 "the pass must actually run concurrently, not serialize (peak {peak})",
2805 );
2806 }
2807
2808 #[tokio::test]
2813 async fn store_with_retry_treats_remote_put_as_recoverable() {
2814 let chunks = make_addrs(6);
2815 let failing: std::collections::HashSet<[u8; 32]> = chunks.iter().take(2).copied().collect();
2816 let failing_for_closure = failing.clone();
2817
2818 let store_one = move |addr: [u8; 32]| {
2819 let fail = failing_for_closure.contains(&addr);
2820 async move {
2821 if fail {
2822 Err(Error::RemotePut {
2823 address: hex::encode(addr),
2824 source: ant_protocol::ProtocolError::StorageFailed(
2825 "insufficient disk space".into(),
2826 ),
2827 })
2828 } else {
2829 Ok(std::time::Instant::now())
2830 }
2831 }
2832 };
2833
2834 let outcome =
2835 merkle_store_with_retry(chunks, || 8, 1, Duration::ZERO, None, 0, 6, store_one)
2836 .await
2837 .expect("remote app-rejections must not abort the batch");
2838
2839 assert_eq!(outcome.stored, 4);
2840 assert_eq!(outcome.failed, 2);
2841 }
2842
2843 #[tokio::test]
2847 async fn store_with_retry_reports_non_quorum_errors_as_fatal() {
2848 let chunks = make_addrs(3);
2849 let store_one = |_addr: [u8; 32]| async move {
2850 Err::<std::time::Instant, _>(Error::Payment("missing proof".into()))
2851 };
2852
2853 let outcome =
2854 merkle_store_with_retry(chunks, || 8, 3, Duration::ZERO, None, 0, 3, store_one)
2855 .await
2856 .expect("fatal is carried in the outcome, not returned as Err");
2857 assert!(matches!(outcome.fatal, Some(Error::Payment(_))));
2858 }
2859
2860 #[tokio::test]
2865 async fn store_with_retry_fatal_preserves_same_pass_successes() {
2866 let chunks = make_addrs(6);
2867 let bad = chunks[5];
2868 let store_one = move |addr: [u8; 32]| async move {
2869 if addr == bad {
2870 Err(Error::Payment("fatal".into()))
2871 } else {
2872 Ok(std::time::Instant::now())
2873 }
2874 };
2875
2876 let outcome =
2877 merkle_store_with_retry(chunks, || 1, 1, Duration::ZERO, None, 0, 6, store_one)
2878 .await
2879 .expect("fatal carried in outcome, not returned as Err");
2880 assert!(matches!(outcome.fatal, Some(Error::Payment(_))));
2881 assert_eq!(outcome.stored, 5);
2883 assert_eq!(outcome.stored_addresses.len(), 5);
2884 assert!(!outcome.stored_addresses.contains(&bad));
2885 assert!(outcome.failed_addresses.iter().any(|(a, _)| *a == bad));
2887 }
2888
2889 #[tokio::test]
2891 async fn store_with_retry_retries_only_the_failed_set() {
2892 let chunks = make_addrs(5);
2893 let total = chunks.len();
2894 let failing: std::collections::HashSet<[u8; 32]> = chunks.iter().take(2).copied().collect();
2895 let failing_for_closure = failing.clone();
2896
2897 let calls = Arc::new(Mutex::new(Vec::<[u8; 32]>::new()));
2899 let calls_for_closure = calls.clone();
2900
2901 let store_one = move |addr: [u8; 32]| {
2902 let calls = calls_for_closure.clone();
2903 let already_seen = calls.lock().unwrap().iter().filter(|&&a| a == addr).count();
2905 let fail = failing_for_closure.contains(&addr) && already_seen == 0;
2906 calls.lock().unwrap().push(addr);
2907 async move {
2908 if fail {
2909 Err(Error::InsufficientPeers("round-1 shortfall".into()))
2910 } else {
2911 Ok(std::time::Instant::now())
2912 }
2913 }
2914 };
2915
2916 let outcome =
2917 merkle_store_with_retry(chunks, || 8, 3, Duration::ZERO, None, 0, total, store_one)
2918 .await
2919 .expect("should converge after retry");
2920
2921 assert_eq!(outcome.stored, total);
2922 assert_eq!(outcome.failed, 0);
2923
2924 let calls = calls.lock().unwrap();
2928 assert_eq!(calls.len(), total + failing.len());
2929 let round_two: std::collections::HashSet<[u8; 32]> =
2930 calls[total..].iter().copied().collect();
2931 assert_eq!(round_two, failing);
2932 }
2933
2934 #[tokio::test]
2937 async fn store_with_retry_counts_retry_success_once_in_histogram() {
2938 let chunks = make_addrs(4);
2939 let total = chunks.len();
2940 let flaky_addr = chunks[0];
2941
2942 let attempts = Arc::new(Mutex::new(HashMap::<[u8; 32], usize>::new()));
2943 let attempts_for_closure = attempts.clone();
2944
2945 let store_one = move |addr: [u8; 32]| {
2946 let attempts = attempts_for_closure.clone();
2947 let n = {
2948 let mut m = attempts.lock().unwrap();
2949 let entry = m.entry(addr).or_insert(0);
2950 *entry += 1;
2951 *entry
2952 };
2953 let fail = addr == flaky_addr && n == 1;
2954 async move {
2955 if fail {
2956 Err(Error::InsufficientPeers("transient".into()))
2957 } else {
2958 Ok(std::time::Instant::now())
2959 }
2960 }
2961 };
2962
2963 let outcome =
2964 merkle_store_with_retry(chunks, || 8, 3, Duration::ZERO, None, 0, total, store_one)
2965 .await
2966 .expect("flaky chunk should recover on retry");
2967
2968 assert_eq!(outcome.stored, total);
2969 assert_eq!(outcome.failed, 0);
2970 assert_eq!(outcome.stats.retries_histogram[0], total - 1);
2972 assert_eq!(outcome.stats.retries_histogram[1], 1);
2973 assert_eq!(outcome.stats.chunk_attempts_total, total + 1);
2975 }
2976
2977 #[tokio::test]
2982 async fn store_with_retry_reports_all_failed_when_retries_exhausted() {
2983 let chunks = make_addrs(3);
2984 let total = chunks.len();
2985
2986 let store_one = |_addr: [u8; 32]| async move {
2987 Err::<std::time::Instant, _>(Error::InsufficientPeers("never converges".into()))
2988 };
2989
2990 let outcome = merkle_store_with_retry(
2991 chunks,
2992 || 8,
2993 MERKLE_STORE_MAX_ATTEMPTS,
2994 Duration::ZERO,
2995 None,
2996 0,
2997 total,
2998 store_one,
2999 )
3000 .await
3001 .expect("an exhausted retry budget is reported, not propagated as Err");
3002
3003 assert_eq!(outcome.stored, 0);
3004 assert_eq!(outcome.failed, total);
3005 assert_eq!(
3007 outcome.stats.chunk_attempts_total,
3008 total * MERKLE_STORE_MAX_ATTEMPTS
3009 );
3010 assert_eq!(outcome.stats.retries_histogram, [0; 4]);
3012 }
3013
3014 #[tokio::test]
3019 async fn store_with_retry_records_failed_addresses_when_exhausted() {
3020 let chunks = make_addrs(6);
3021 let failing: std::collections::HashSet<[u8; 32]> = chunks.iter().take(2).copied().collect();
3022 let failing_for_closure = failing.clone();
3023
3024 let store_one = move |addr: [u8; 32]| {
3025 let fail = failing_for_closure.contains(&addr);
3026 async move {
3027 if fail {
3028 Err(Error::InsufficientPeers("permanent shortfall".into()))
3029 } else {
3030 Ok(std::time::Instant::now())
3031 }
3032 }
3033 };
3034
3035 let outcome = merkle_store_with_retry(
3036 chunks,
3037 || 8,
3038 MERKLE_STORE_MAX_ATTEMPTS,
3039 Duration::ZERO,
3040 None,
3041 0,
3042 6,
3043 store_one,
3044 )
3045 .await
3046 .expect("quorum shortfalls must not abort the batch");
3047
3048 assert_eq!(outcome.stored, 4);
3049 assert_eq!(outcome.failed, 2);
3050 assert_eq!(outcome.failed_addresses.len(), 2);
3052 let reported: std::collections::HashSet<[u8; 32]> =
3053 outcome.failed_addresses.iter().map(|(a, _)| *a).collect();
3054 assert_eq!(reported, failing);
3055 for (_, msg) in &outcome.failed_addresses {
3057 assert!(msg.contains("permanent shortfall"));
3058 }
3059 }
3060
3061 #[tokio::test]
3064 async fn store_with_retry_failed_addresses_empty_on_full_success() {
3065 let chunks = make_addrs(4);
3066 let total = chunks.len();
3067 let store_one = |_addr: [u8; 32]| async move { Ok(std::time::Instant::now()) };
3068
3069 let outcome = merkle_store_with_retry(
3070 chunks,
3071 || 8,
3072 MERKLE_STORE_MAX_ATTEMPTS,
3073 Duration::ZERO,
3074 None,
3075 0,
3076 total,
3077 store_one,
3078 )
3079 .await
3080 .expect("all chunks store");
3081
3082 assert_eq!(outcome.stored, total);
3083 assert_eq!(outcome.failed, 0);
3084 assert!(outcome.failed_addresses.is_empty());
3085 }
3086
3087 #[test]
3094 fn deferred_round_histogram_slot_maps_and_clamps() {
3095 assert_eq!(deferred_round_histogram_slot(0, 4), 1);
3096 assert_eq!(deferred_round_histogram_slot(1, 4), 2);
3097 assert_eq!(deferred_round_histogram_slot(2, 4), 3);
3098 assert_eq!(deferred_round_histogram_slot(3, 4), 3);
3100 assert_eq!(deferred_round_histogram_slot(9, 4), 3);
3101 }
3102
3103 fn deferred_set(count: usize) -> Vec<([u8; 32], String)> {
3104 make_test_addresses(count)
3105 .into_iter()
3106 .map(|addr| (addr, "short of quorum".to_string()))
3107 .collect()
3108 }
3109
3110 #[tokio::test]
3114 async fn deferred_retry_succeeds_on_a_later_round() {
3115 let deferred = deferred_set(3);
3116 let attempts = Arc::new(Mutex::new(HashMap::<[u8; 32], usize>::new()));
3119 let attempts_for_closure = attempts.clone();
3120 let store_one = move |addr: [u8; 32]| {
3121 let attempts = attempts_for_closure.clone();
3122 async move {
3123 let n = {
3124 let mut map = attempts.lock().unwrap();
3125 let e = map.entry(addr).or_insert(0);
3126 *e += 1;
3127 *e
3128 };
3129 if n < 2 {
3130 Err(Error::InsufficientPeers("still short".into()))
3131 } else {
3132 Ok(std::time::Instant::now())
3133 }
3134 }
3135 };
3136
3137 let outcome = merkle_deferred_retry(
3138 deferred,
3139 &[0, 0, 0],
3140 |n: usize| n.max(1),
3141 None,
3142 0,
3143 3,
3144 store_one,
3145 )
3146 .await
3147 .expect("deferred retry must not abort on quorum shortfalls");
3148
3149 assert_eq!(outcome.stored, 3, "all three land by round 1");
3150 assert_eq!(outcome.stored_addresses.len(), 3);
3151 assert_eq!(outcome.failed, 0);
3152 assert!(outcome.failed_addresses.is_empty());
3153 assert!(outcome.fatal.is_none());
3154 assert_eq!(outcome.stats.retries_histogram[1], 0);
3156 assert_eq!(outcome.stats.retries_histogram[2], 3);
3157 assert_eq!(outcome.stats.chunk_attempts_total, 6);
3159 }
3160
3161 #[tokio::test]
3164 async fn deferred_retry_leftovers_become_failed() {
3165 let deferred = deferred_set(2);
3166 let store_one = |_addr: [u8; 32]| async move {
3167 Err::<std::time::Instant, _>(Error::InsufficientPeers("always short".into()))
3168 };
3169
3170 let outcome = merkle_deferred_retry(
3171 deferred,
3172 &[0, 0, 0],
3173 |n: usize| n.max(1),
3174 None,
3175 0,
3176 2,
3177 store_one,
3178 )
3179 .await
3180 .expect("exhausted retries report failures, not an error");
3181
3182 assert_eq!(outcome.stored, 0);
3183 assert!(outcome.stored_addresses.is_empty());
3184 assert_eq!(outcome.failed, 2);
3185 assert_eq!(outcome.failed_addresses.len(), 2);
3186 assert!(outcome.fatal.is_none());
3187 assert_eq!(outcome.stats.chunk_attempts_total, 6);
3189 }
3190
3191 #[tokio::test]
3196 async fn deferred_retry_fatal_error_preserves_prior_progress() {
3197 let addrs = make_test_addresses(2);
3198 let good = addrs[0];
3199 let bad = addrs[1];
3200 let deferred = vec![(good, "short".to_string()), (bad, "short".to_string())];
3201
3202 let attempts = Arc::new(Mutex::new(HashMap::<[u8; 32], usize>::new()));
3205 let attempts_for_closure = attempts.clone();
3206 let store_one = move |addr: [u8; 32]| {
3207 let attempts = attempts_for_closure.clone();
3208 async move {
3209 let n = {
3210 let mut map = attempts.lock().unwrap();
3211 let e = map.entry(addr).or_insert(0);
3212 *e += 1;
3213 *e
3214 };
3215 if addr == good {
3216 Ok(std::time::Instant::now())
3217 } else if n == 1 {
3218 Err(Error::InsufficientPeers("short".into()))
3219 } else {
3220 Err(Error::Payment("fatal on retry".into()))
3221 }
3222 }
3223 };
3224
3225 let outcome = merkle_deferred_retry(
3226 deferred,
3227 &[0, 0, 0],
3228 |n: usize| n.max(1),
3229 None,
3230 0,
3231 2,
3232 store_one,
3233 )
3234 .await
3235 .expect("a fatal round error is reported via `fatal`, not as Err");
3236
3237 assert!(outcome.fatal.is_some(), "fatal error must be captured");
3238 assert_eq!(outcome.stored, 1, "round-0 success preserved");
3239 assert_eq!(outcome.stored_addresses, vec![good]);
3240 assert_eq!(outcome.failed, 1);
3241 assert_eq!(outcome.failed_addresses.len(), 1);
3242 assert_eq!(outcome.failed_addresses[0].0, bad);
3243 }
3244
3245 #[tokio::test]
3247 async fn deferred_retry_empty_set_is_a_noop() {
3248 let store_one = |_addr: [u8; 32]| async move {
3249 Err::<std::time::Instant, _>(Error::InsufficientPeers("unused".into()))
3250 };
3251
3252 let outcome = merkle_deferred_retry(
3253 Vec::new(),
3254 &DEFERRED_ROUND_DELAYS_SECS,
3255 |n: usize| n.max(1),
3256 None,
3257 7,
3258 7,
3259 store_one,
3260 )
3261 .await
3262 .expect("empty deferred set is a no-op");
3263
3264 assert_eq!(outcome.stored, 7, "stored_offset carried through unchanged");
3265 assert_eq!(outcome.failed, 0);
3266 assert!(outcome.stored_addresses.is_empty());
3267 assert!(outcome.failed_addresses.is_empty());
3268 assert!(outcome.fatal.is_none());
3269 }
3270}