diff --git a/ant-core/src/data/client/file.rs b/ant-core/src/data/client/file.rs index 9c10dca..479d6f5 100644 --- a/ant-core/src/data/client/file.rs +++ b/ant-core/src/data/client/file.rs @@ -6,7 +6,9 @@ //! For in-memory data uploads, see the `data` module. use crate::data::client::batch::{finalize_batch_payment, PaymentIntent, PreparedChunk}; -use crate::data::client::merkle::PaymentMode; +use crate::data::client::merkle::{ + finalize_merkle_batch, should_use_merkle, PaymentMode, PreparedMerkleBatch, +}; use crate::data::client::Client; use crate::data::error::{Error, Result}; use ant_evm::QuoteHash; @@ -33,23 +35,42 @@ pub struct FileUploadResult { pub payment_mode_used: PaymentMode, } +/// Payment information for external signing — either wave-batch or merkle. +#[derive(Debug)] +pub enum ExternalPaymentInfo { + /// Wave-batch: individual (quote_hash, rewards_address, amount) tuples. + WaveBatch { + /// Chunks ready for payment (needed for finalize). + prepared_chunks: Vec, + /// Payment intent for external signing. + payment_intent: PaymentIntent, + }, + /// Merkle: single on-chain call with depth, pool commitments, timestamp. + Merkle { + /// The prepared merkle batch (public fields sent to frontend, private fields stay in Rust). + prepared_batch: PreparedMerkleBatch, + /// Raw chunk contents (needed for upload after payment). + chunk_contents: Vec, + /// Chunk addresses in order (needed for upload after payment). + chunk_addresses: Vec<[u8; 32]>, + }, +} + /// Prepared upload ready for external payment. /// /// Contains everything needed to construct the on-chain payment transaction /// externally (e.g. via WalletConnect in a desktop app) and then finalize /// the upload without a Rust-side wallet. /// -/// Note: This struct stays in Rust memory — only `payment_intent` is sent -/// to the frontend. `PreparedChunk` contains non-serializable network types, -/// so the full struct cannot derive `Serialize`. +/// Note: This struct stays in Rust memory — only the public fields of +/// `payment_info` are sent to the frontend. `PreparedChunk` contains +/// non-serializable network types, so the full struct cannot derive `Serialize`. #[derive(Debug)] pub struct PreparedUpload { /// The data map for later retrieval. pub data_map: DataMap, - /// Chunks ready for payment. - pub prepared_chunks: Vec, - /// Payment intent for external signing. - pub payment_intent: PaymentIntent, + /// Payment information — either wave-batch or merkle depending on chunk count. + pub payment_info: ExternalPaymentInfo, } /// Return type for [`spawn_file_encryption`]: chunk receiver, `DataMap` oneshot, join handle. @@ -212,54 +233,146 @@ impl Client { .await .map_err(|_| Error::Encryption("no DataMap from encryption thread".to_string()))?; - // Prepare each chunk (collect quotes, fetch contract prices). - let mut prepared_chunks = Vec::with_capacity(chunk_contents.len()); - for content in chunk_contents { - if let Some(prepared) = self.prepare_chunk_payment(content).await? { - prepared_chunks.push(prepared); + let chunk_count = chunk_contents.len(); + + let payment_info = if should_use_merkle(chunk_count, PaymentMode::Auto) { + // Merkle path: build tree, collect candidate pools, return for external payment. + info!("Using merkle batch preparation for {chunk_count} file chunks"); + + let addresses: Vec<[u8; 32]> = + chunk_contents.iter().map(|c| compute_address(c)).collect(); + + let avg_size = + chunk_contents.iter().map(bytes::Bytes::len).sum::() / chunk_count.max(1); + let avg_size_u64 = u64::try_from(avg_size).unwrap_or(0); + + let prepared_batch = self + .prepare_merkle_batch_external(&addresses, DATA_TYPE_CHUNK, avg_size_u64) + .await?; + + info!( + "File prepared for external merkle signing: {} chunks, depth={} ({})", + chunk_count, + prepared_batch.depth, + path.display() + ); + + ExternalPaymentInfo::Merkle { + prepared_batch, + chunk_contents, + chunk_addresses: addresses, + } + } else { + // Wave-batch path: collect quotes per chunk. + let mut prepared_chunks = Vec::with_capacity(chunk_contents.len()); + for content in chunk_contents { + if let Some(prepared) = self.prepare_chunk_payment(content).await? { + prepared_chunks.push(prepared); + } } - } - let payment_intent = PaymentIntent::from_prepared_chunks(&prepared_chunks); + let payment_intent = PaymentIntent::from_prepared_chunks(&prepared_chunks); - info!( - "File prepared for external signing: {} chunks, total {} atto ({})", - prepared_chunks.len(), - payment_intent.total_amount, - path.display() - ); + info!( + "File prepared for external signing: {} chunks, total {} atto ({})", + prepared_chunks.len(), + payment_intent.total_amount, + path.display() + ); + + ExternalPaymentInfo::WaveBatch { + prepared_chunks, + payment_intent, + } + }; Ok(PreparedUpload { data_map, - prepared_chunks, - payment_intent, + payment_info, }) } - /// Phase 2 of external-signer upload: finalize with externally-signed tx hashes. + /// Phase 2 of external-signer upload (wave-batch): finalize with externally-signed tx hashes. /// - /// Takes a [`PreparedUpload`] from [`Client::file_prepare_upload`] and a map + /// Takes a [`PreparedUpload`] that used wave-batch payment and a map /// of `quote_hash -> tx_hash` provided by the external signer after on-chain /// payment. Builds payment proofs and stores chunks on the network. /// /// # Errors /// - /// Returns an error if proof construction fails or any chunk cannot be stored. + /// Returns an error if the prepared upload used merkle payment (use + /// [`Client::finalize_upload_merkle`] instead), proof construction fails, + /// or any chunk cannot be stored. pub async fn finalize_upload( &self, prepared: PreparedUpload, tx_hash_map: &HashMap, ) -> Result { - let paid_chunks = finalize_batch_payment(prepared.prepared_chunks, tx_hash_map)?; - let chunks_stored = self.store_paid_chunks(paid_chunks).await?.len(); - - info!("External-signer upload finalized: {chunks_stored} chunks stored"); + match prepared.payment_info { + ExternalPaymentInfo::WaveBatch { + prepared_chunks, + payment_intent: _, + } => { + let paid_chunks = finalize_batch_payment(prepared_chunks, tx_hash_map)?; + let chunks_stored = self.store_paid_chunks(paid_chunks).await?.len(); + + info!("External-signer upload finalized: {chunks_stored} chunks stored"); + + Ok(FileUploadResult { + data_map: prepared.data_map, + chunks_stored, + payment_mode_used: PaymentMode::Single, + }) + } + ExternalPaymentInfo::Merkle { .. } => Err(Error::Payment( + "Cannot finalize merkle upload with wave-batch tx hashes. \ + Use finalize_upload_merkle() instead." + .to_string(), + )), + } + } - Ok(FileUploadResult { - data_map: prepared.data_map, - chunks_stored, - payment_mode_used: PaymentMode::Single, - }) + /// Phase 2 of external-signer upload (merkle): finalize with winner pool hash. + /// + /// Takes a [`PreparedUpload`] that used merkle payment and the `winner_pool_hash` + /// returned by the on-chain merkle payment transaction. Generates proofs and + /// stores chunks on the network. + /// + /// # Errors + /// + /// Returns an error if the prepared upload used wave-batch payment (use + /// [`Client::finalize_upload`] instead), proof generation fails, + /// or any chunk cannot be stored. + pub async fn finalize_upload_merkle( + &self, + prepared: PreparedUpload, + winner_pool_hash: [u8; 32], + ) -> Result { + match prepared.payment_info { + ExternalPaymentInfo::Merkle { + prepared_batch, + chunk_contents, + chunk_addresses, + } => { + let batch_result = finalize_merkle_batch(prepared_batch, winner_pool_hash)?; + let chunks_stored = self + .merkle_upload_chunks(chunk_contents, chunk_addresses, &batch_result) + .await?; + + info!("External-signer merkle upload finalized: {chunks_stored} chunks stored"); + + Ok(FileUploadResult { + data_map: prepared.data_map, + chunks_stored, + payment_mode_used: PaymentMode::Merkle, + }) + } + ExternalPaymentInfo::WaveBatch { .. } => Err(Error::Payment( + "Cannot finalize wave-batch upload with merkle winner hash. \ + Use finalize_upload() instead." + .to_string(), + )), + } } /// Upload a file with a specific payment mode. diff --git a/ant-core/src/data/client/merkle.rs b/ant-core/src/data/client/merkle.rs index d94db9b..02b037b 100644 --- a/ant-core/src/data/client/merkle.rs +++ b/ant-core/src/data/client/merkle.rs @@ -48,6 +48,37 @@ pub struct MerkleBatchPaymentResult { pub chunk_count: usize, } +/// Prepared merkle batch ready for external payment. +/// +/// Contains everything needed to submit the on-chain merkle payment +/// and then finalize proof generation without a wallet. +pub struct PreparedMerkleBatch { + /// Merkle tree depth (needed for the on-chain call). + pub depth: u8, + /// Pool commitments for the on-chain call. + pub pool_commitments: Vec, + /// Timestamp used for the merkle payment. + pub merkle_payment_timestamp: u64, + /// Internal: candidate pools (needed for proof generation after payment). + candidate_pools: Vec, + /// Internal: the merkle tree (needed for proof generation). + tree: MerkleTree, + /// Internal: chunk addresses in order. + addresses: Vec<[u8; 32]>, +} + +impl std::fmt::Debug for PreparedMerkleBatch { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("PreparedMerkleBatch") + .field("depth", &self.depth) + .field("pool_commitments", &self.pool_commitments.len()) + .field("merkle_payment_timestamp", &self.merkle_payment_timestamp) + .field("candidate_pools", &self.candidate_pools.len()) + .field("addresses", &self.addresses.len()) + .finish() + } +} + /// Determine whether to use merkle payments for a given batch size. /// Free function — no Client needed. #[must_use] @@ -101,21 +132,24 @@ impl Client { .await } - /// Pay for a single batch (up to `MAX_LEAVES` chunks). - async fn pay_for_merkle_single_batch( + /// Phase 1 of external-signer merkle payment: prepare batch without paying. + /// + /// Builds the merkle tree, collects candidate pools from the network, + /// and returns the data needed for the on-chain payment call. + /// Requires `EvmNetwork` but NOT a wallet. + pub async fn prepare_merkle_batch_external( &self, addresses: &[[u8; 32]], data_type: u32, data_size: u64, - ) -> Result { + ) -> Result { let chunk_count = addresses.len(); - let wallet = self.require_wallet()?; let xornames: Vec = addresses.iter().map(|a| XorName(*a)).collect(); info!("Building merkle tree for {chunk_count} chunks"); // 1. Build merkle tree - let tree = MerkleTree::from_xornames(xornames.clone()) + let tree = MerkleTree::from_xornames(xornames) .map_err(|e| Error::Payment(format!("Failed to build merkle tree: {e}")))?; let depth = tree.depth(); @@ -152,10 +186,38 @@ impl Client { .map(MerklePaymentCandidatePool::to_commitment) .collect(); - // 5. Pay on-chain (single transaction) - info!("Submitting merkle batch payment on-chain (depth={depth})"); + Ok(PreparedMerkleBatch { + depth, + pool_commitments, + merkle_payment_timestamp, + candidate_pools, + tree, + addresses: addresses.to_vec(), + }) + } + + /// Pay for a single batch (up to `MAX_LEAVES` chunks). + async fn pay_for_merkle_single_batch( + &self, + addresses: &[[u8; 32]], + data_type: u32, + data_size: u64, + ) -> Result { + let wallet = self.require_wallet()?; + let prepared = self + .prepare_merkle_batch_external(addresses, data_type, data_size) + .await?; + + info!( + "Submitting merkle batch payment on-chain (depth={})", + prepared.depth + ); let (winner_pool_hash, _amount, _gas_info) = wallet - .pay_for_merkle_tree(depth, pool_commitments, merkle_payment_timestamp) + .pay_for_merkle_tree( + prepared.depth, + prepared.pool_commitments.clone(), + prepared.merkle_payment_timestamp, + ) .await .map_err(|e| Error::Payment(format!("Merkle batch payment failed: {e}")))?; @@ -164,44 +226,7 @@ impl Client { hex::encode(winner_pool_hash) ); - // 6. Find the winner pool - let winner_pool = candidate_pools - .iter() - .find(|pool| pool.hash() == winner_pool_hash) - .ok_or_else(|| { - Error::Payment(format!( - "Winner pool {} not found in candidate pools", - hex::encode(winner_pool_hash) - )) - })?; - - // 7. Generate proofs for each chunk - info!("Generating merkle proofs for {chunk_count} chunks"); - let mut proofs = HashMap::with_capacity(chunk_count); - - for (i, xorname) in xornames.iter().enumerate() { - let address_proof = tree.generate_address_proof(i, *xorname).map_err(|e| { - Error::Payment(format!( - "Failed to generate address proof for chunk {i}: {e}" - )) - })?; - - let merkle_proof = - MerklePaymentProof::new(*xorname, address_proof, winner_pool.clone()); - - let tagged_bytes = serialize_merkle_proof(&merkle_proof).map_err(|e| { - Error::Serialization(format!("Failed to serialize merkle proof: {e}")) - })?; - - proofs.insert(addresses[i], tagged_bytes); - } - - info!("Merkle batch payment complete: {chunk_count} proofs generated"); - - Ok(MerkleBatchPaymentResult { - proofs, - chunk_count, - }) + finalize_merkle_batch(prepared, winner_pool_hash) } /// Handle batches larger than `MAX_LEAVES` by splitting into sub-batches. @@ -490,6 +515,59 @@ impl Client { } } +/// Phase 2 of external-signer merkle payment: generate proofs from winner. +/// +/// Takes the prepared batch and the winner pool hash returned by the +/// on-chain payment transaction. Generates per-chunk merkle proofs. +pub fn finalize_merkle_batch( + prepared: PreparedMerkleBatch, + winner_pool_hash: [u8; 32], +) -> Result { + let chunk_count = prepared.addresses.len(); + let xornames: Vec = prepared.addresses.iter().map(|a| XorName(*a)).collect(); + + // Find the winner pool + let winner_pool = prepared + .candidate_pools + .iter() + .find(|pool| pool.hash() == winner_pool_hash) + .ok_or_else(|| { + Error::Payment(format!( + "Winner pool {} not found in candidate pools", + hex::encode(winner_pool_hash) + )) + })?; + + // Generate proofs for each chunk + info!("Generating merkle proofs for {chunk_count} chunks"); + let mut proofs = HashMap::with_capacity(chunk_count); + + for (i, xorname) in xornames.iter().enumerate() { + let address_proof = prepared + .tree + .generate_address_proof(i, *xorname) + .map_err(|e| { + Error::Payment(format!( + "Failed to generate address proof for chunk {i}: {e}" + )) + })?; + + let merkle_proof = MerklePaymentProof::new(*xorname, address_proof, winner_pool.clone()); + + let tagged_bytes = serialize_merkle_proof(&merkle_proof) + .map_err(|e| Error::Serialization(format!("Failed to serialize merkle proof: {e}")))?; + + proofs.insert(prepared.addresses[i], tagged_bytes); + } + + info!("Merkle batch payment complete: {chunk_count} proofs generated"); + + Ok(MerkleBatchPaymentResult { + proofs, + chunk_count, + }) +} + /// Compile-time assertions that merkle method futures are Send. #[cfg(test)] mod send_assertions { @@ -746,6 +824,119 @@ mod tests { assert_ne!(candidate.quoting_metrics.data_type, 0); } + // ========================================================================= + // finalize_merkle_batch (external signer) + // ========================================================================= + + fn make_dummy_candidate_nodes( + timestamp: u64, + ) -> [MerklePaymentCandidateNode; CANDIDATES_PER_POOL] { + std::array::from_fn(|i| MerklePaymentCandidateNode { + pub_key: vec![i as u8; 32], + quoting_metrics: ant_evm::QuotingMetrics { + data_size: 1024, + data_type: 0, + close_records_stored: 0, + records_per_type: vec![], + max_records: 100, + received_payment_count: 0, + live_time: 0, + network_density: None, + network_size: None, + }, + reward_address: ant_evm::RewardsAddress::new([i as u8; 20]), + merkle_payment_timestamp: timestamp, + signature: vec![i as u8; 64], + }) + } + + fn make_prepared_merkle_batch(count: usize) -> PreparedMerkleBatch { + let addrs = make_test_addresses(count); + let xornames: Vec = addrs.iter().map(|a| XorName(*a)).collect(); + let tree = MerkleTree::from_xornames(xornames).unwrap(); + + let timestamp = std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .unwrap() + .as_secs(); + + let midpoints = tree.reward_candidates(timestamp).unwrap(); + + let candidate_pools: Vec = midpoints + .into_iter() + .map(|mp| MerklePaymentCandidatePool { + midpoint_proof: mp, + candidate_nodes: make_dummy_candidate_nodes(timestamp), + }) + .collect(); + + let pool_commitments = candidate_pools + .iter() + .map(MerklePaymentCandidatePool::to_commitment) + .collect(); + + PreparedMerkleBatch { + depth: tree.depth(), + pool_commitments, + merkle_payment_timestamp: timestamp, + candidate_pools, + tree, + addresses: addrs, + } + } + + #[test] + fn test_finalize_merkle_batch_with_valid_winner() { + let prepared = make_prepared_merkle_batch(4); + let winner_hash = prepared.candidate_pools[0].hash(); + + let result = finalize_merkle_batch(prepared, winner_hash); + assert!( + result.is_ok(), + "should succeed with valid winner: {result:?}" + ); + + let batch = result.unwrap(); + assert_eq!(batch.chunk_count, 4); + assert_eq!(batch.proofs.len(), 4); + + // Every proof should be non-empty + for proof_bytes in batch.proofs.values() { + assert!(!proof_bytes.is_empty()); + } + } + + #[test] + fn test_finalize_merkle_batch_with_invalid_winner() { + let prepared = make_prepared_merkle_batch(4); + let bad_hash = [0xFF; 32]; + + let result = finalize_merkle_batch(prepared, bad_hash); + assert!(result.is_err()); + let err = result.unwrap_err().to_string(); + assert!(err.contains("not found in candidate pools"), "got: {err}"); + } + + #[test] + fn test_finalize_merkle_batch_proofs_are_deserializable() { + use ant_node::payment::deserialize_merkle_proof; + + let prepared = make_prepared_merkle_batch(8); + let winner_hash = prepared.candidate_pools[0].hash(); + + let batch = finalize_merkle_batch(prepared, winner_hash).unwrap(); + + for (addr, proof_bytes) in &batch.proofs { + let proof = deserialize_merkle_proof(proof_bytes); + assert!( + proof.is_ok(), + "proof for {} should deserialize: {:?}", + hex::encode(addr), + proof.err() + ); + } + } + // ========================================================================= // Batch splitting edge cases // ========================================================================= diff --git a/ant-core/src/data/mod.rs b/ant-core/src/data/mod.rs index 7b28e5d..1102796 100644 --- a/ant-core/src/data/mod.rs +++ b/ant-core/src/data/mod.rs @@ -21,8 +21,11 @@ pub use ant_node::client::{compute_address, DataChunk, XorName}; // Re-export client data types pub use client::batch::{finalize_batch_payment, PaidChunk, PaymentIntent, PreparedChunk}; pub use client::data::DataUploadResult; -pub use client::file::{FileUploadResult, PreparedUpload}; -pub use client::merkle::{MerkleBatchPaymentResult, PaymentMode, DEFAULT_MERKLE_THRESHOLD}; +pub use client::file::{ExternalPaymentInfo, FileUploadResult, PreparedUpload}; +pub use client::merkle::{ + finalize_merkle_batch, MerkleBatchPaymentResult, PaymentMode, PreparedMerkleBatch, + DEFAULT_MERKLE_THRESHOLD, +}; // Re-export self-encryption types pub use self_encryption::DataMap;