use crate::data::client::adaptive::{observe_op, rebucketed_ordered};
use crate::data::client::batch::{PaymentIntent, PreparedChunk};
use crate::data::client::classify_error;
use crate::data::client::file::{ExternalPaymentInfo, PreparedUpload, Visibility};
use crate::data::client::merkle::{chunk_contents_for_upload_addresses, PaymentMode};
use crate::data::client::Client;
use crate::data::error::{Error, Result};
use ant_protocol::{compute_address, DATA_TYPE_CHUNK};
use bytes::Bytes;
use futures::stream::StreamExt;
use self_encryption::{decrypt, encrypt, get_root_data_map, DataMap, EncryptedChunk};
use std::num::NonZeroUsize;
use tokio::runtime::{Handle, RuntimeFlavor};
use tracing::{debug, info};
use xor_name::XorName;
#[derive(Debug, Clone)]
pub struct DataUploadResult {
pub data_map: DataMap,
pub chunks_stored: usize,
pub payment_mode_used: PaymentMode,
}
impl Client {
pub async fn data_upload(&self, content: Bytes) -> Result<DataUploadResult> {
let content_len = content.len();
debug!("Encrypting data ({content_len} bytes)");
let (data_map, encrypted_chunks) = encrypt(content)
.map_err(|e| Error::Encryption(format!("Failed to encrypt data: {e}")))?;
info!("Data encrypted into {} chunks", encrypted_chunks.len());
let chunk_contents: Vec<Bytes> = encrypted_chunks
.into_iter()
.map(|chunk| chunk.content)
.collect();
let (addresses, _storage_cost, _gas_cost) =
self.batch_upload_chunks(chunk_contents).await?;
let chunks_stored = addresses.len();
info!("Data uploaded: {chunks_stored} chunks stored ({content_len} bytes original)");
Ok(DataUploadResult {
data_map,
chunks_stored,
payment_mode_used: PaymentMode::Single,
})
}
pub async fn data_upload_with_mode(
&self,
content: Bytes,
mode: PaymentMode,
) -> Result<DataUploadResult> {
let content_len = content.len();
debug!("Encrypting data ({content_len} bytes) with mode {mode:?}");
let (data_map, encrypted_chunks) = encrypt(content)
.map_err(|e| Error::Encryption(format!("Failed to encrypt data: {e}")))?;
let chunk_count = encrypted_chunks.len();
info!("Data encrypted into {chunk_count} chunks");
let chunk_contents: Vec<Bytes> = encrypted_chunks
.into_iter()
.map(|chunk| chunk.content)
.collect();
if self.should_use_merkle(chunk_count, mode) {
info!("Using merkle batch payment for {chunk_count} chunks");
let chunk_entries: Vec<([u8; 32], u64)> = chunk_contents
.iter()
.map(|chunk| {
let size = u64::try_from(chunk.len())
.map_err(|e| Error::InvalidData(format!("chunk size too large: {e}")))?;
Ok((compute_address(chunk), size))
})
.collect::<Result<Vec<_>>>()?;
let merkle_plan = match self
.plan_merkle_upload(chunk_entries, DATA_TYPE_CHUNK, None)
.await
{
Ok(plan) => plan,
Err(Error::InsufficientPeers(ref msg)) if mode == PaymentMode::Auto => {
info!("Merkle preflight needs more peers ({msg}), falling back to wave-batch");
let (addresses, _sc, _gc) = self.batch_upload_chunks(chunk_contents).await?;
return Ok(DataUploadResult {
data_map,
chunks_stored: addresses.len(),
payment_mode_used: PaymentMode::Single,
});
}
Err(e) => return Err(e),
};
if merkle_plan.to_upload.is_empty() {
info!("All {chunk_count} chunks already stored; skipping merkle payment");
return Ok(DataUploadResult {
data_map,
chunks_stored: chunk_count,
payment_mode_used: PaymentMode::Merkle,
});
}
let chunk_contents =
chunk_contents_for_upload_addresses(chunk_contents, &merkle_plan.to_upload)?;
let remaining_chunks = merkle_plan.to_upload.len();
if !self.should_use_merkle(remaining_chunks, mode) {
info!(
"{remaining_chunks} chunks need upload after merkle preflight; \
using single-node payment"
);
let (addresses, _sc, _gc) = self.batch_upload_chunks(chunk_contents).await?;
return Ok(DataUploadResult {
data_map,
chunks_stored: merkle_plan.already_stored.len() + addresses.len(),
payment_mode_used: PaymentMode::Single,
});
}
let batch_result = match self
.pay_for_merkle_batch(
&merkle_plan.to_upload,
DATA_TYPE_CHUNK,
merkle_plan.to_upload_avg_size(),
)
.await
{
Ok(result) => result,
Err(Error::InsufficientPeers(ref msg)) if mode == PaymentMode::Auto => {
info!("Merkle needs more peers ({msg}), falling back to wave-batch");
let (addresses, _sc, _gc) = self.batch_upload_chunks(chunk_contents).await?;
return Ok(DataUploadResult {
data_map,
chunks_stored: merkle_plan.already_stored.len() + addresses.len(),
payment_mode_used: PaymentMode::Single,
});
}
Err(e) => return Err(e),
};
let outcome = self
.merkle_upload_chunks(
chunk_contents,
merkle_plan.to_upload,
&batch_result,
None,
merkle_plan.already_stored.len(),
chunk_count,
)
.await?;
if outcome.failed > 0 {
return Err(Error::InsufficientPeers(format!(
"Data merkle upload incomplete: {} of {} chunk(s) short of quorum after retries",
outcome.failed, chunk_count
)));
}
info!(
"Data uploaded via merkle: {} chunks stored ({content_len} bytes)",
outcome.stored
);
Ok(DataUploadResult {
data_map,
chunks_stored: outcome.stored,
payment_mode_used: PaymentMode::Merkle,
})
} else {
let (addresses, _sc, _gc) = self.batch_upload_chunks(chunk_contents).await?;
info!(
"Data uploaded: {} chunks stored ({content_len} bytes original)",
addresses.len()
);
Ok(DataUploadResult {
data_map,
chunks_stored: addresses.len(),
payment_mode_used: PaymentMode::Single,
})
}
}
pub async fn data_prepare_upload(&self, content: Bytes) -> Result<PreparedUpload> {
self.data_prepare_upload_with_visibility(content, Visibility::Private)
.await
}
pub async fn data_prepare_upload_with_visibility(
&self,
content: Bytes,
visibility: Visibility,
) -> Result<PreparedUpload> {
let content_len = content.len();
debug!("Preparing data upload for external signing (visibility={visibility:?}, {content_len} bytes)");
let (data_map, encrypted_chunks) = encrypt(content)
.map_err(|e| Error::Encryption(format!("Failed to encrypt data: {e}")))?;
let mut chunk_contents: Vec<Bytes> = encrypted_chunks
.into_iter()
.map(|chunk| chunk.content)
.collect();
info!("Data encrypted into {} chunks", chunk_contents.len());
let data_map_address = match visibility {
Visibility::Private => None,
Visibility::Public => {
let serialized = rmp_serde::to_vec(&data_map).map_err(|e| {
Error::Serialization(format!("Failed to serialize DataMap: {e}"))
})?;
let bytes = Bytes::from(serialized);
let address = compute_address(&bytes);
info!(
"Public upload: bundling DataMap chunk ({} bytes) at address {}",
bytes.len(),
hex::encode(address)
);
chunk_contents.push(bytes);
Some(address)
}
};
let chunk_count = chunk_contents.len();
let chunks_with_addr: Vec<(Bytes, [u8; 32])> = chunk_contents
.into_iter()
.map(|content| {
let address = compute_address(&content);
(content, address)
})
.collect();
let quote_limiter = self.controller().quote.clone();
let quote_concurrency = quote_limiter.current().min(chunk_count.max(1));
let results: Vec<([u8; 32], Result<Option<PreparedChunk>>)> =
futures::stream::iter(chunks_with_addr)
.map(|(content, address)| {
let limiter = quote_limiter.clone();
async move {
let result = observe_op(
&limiter,
|| async move { self.prepare_chunk_payment(content).await },
classify_error,
)
.await;
(address, result)
}
})
.buffer_unordered(quote_concurrency)
.collect()
.await;
let mut prepared_chunks = Vec::with_capacity(results.len());
let mut already_stored_addresses = Vec::new();
for (address, result) in results {
match result? {
Some(prepared) => prepared_chunks.push(prepared),
None => already_stored_addresses.push(address),
}
}
if let Some(addr) = data_map_address {
if already_stored_addresses.contains(&addr) {
info!(
"Public upload: DataMap chunk {} was already stored \
on the network — address is retrievable without a \
new payment",
hex::encode(addr)
);
}
}
let payment_intent = PaymentIntent::from_prepared_chunks(&prepared_chunks);
info!(
"Data prepared for external signing: {} chunks, {} already stored, total {} atto ({content_len} bytes)",
prepared_chunks.len(),
already_stored_addresses.len(),
payment_intent.total_amount,
);
Ok(PreparedUpload {
data_map,
payment_info: ExternalPaymentInfo::WaveBatch {
prepared_chunks,
payment_intent,
},
data_map_address,
already_stored_addresses,
total_chunks: chunk_count,
})
}
pub async fn data_map_store(&self, data_map: &DataMap) -> Result<[u8; 32]> {
let serialized = rmp_serde::to_vec(data_map)
.map_err(|e| Error::Serialization(format!("Failed to serialize DataMap: {e}")))?;
info!(
"Storing DataMap as public chunk ({} bytes serialized)",
serialized.len()
);
self.chunk_put(Bytes::from(serialized)).await
}
pub async fn data_map_fetch(&self, address: &[u8; 32]) -> Result<DataMap> {
let chunk = self.chunk_get(address).await?.ok_or_else(|| {
Error::NotFound(format!(
"DataMap chunk not found at {}",
hex::encode(address)
))
})?;
decode_data_map_chunk(&chunk.content)
}
pub async fn data_map_fetch_from_closest_peers(
&self,
address: &[u8; 32],
peer_count: NonZeroUsize,
) -> Result<DataMap> {
let chunk = self
.chunk_get_from_closest_peers(address, peer_count.get())
.await?
.ok_or_else(|| {
Error::NotFound(format!(
"DataMap chunk not found at {}",
hex::encode(address)
))
})?;
decode_data_map_chunk(&chunk.content)
}
pub async fn data_download(&self, data_map: &DataMap) -> Result<Bytes> {
let root_data_map = self.resolve_root_data_map(data_map).await?;
let chunk_infos = root_data_map.infos();
debug!("Downloading data ({} chunks)", chunk_infos.len());
let addresses: Vec<[u8; 32]> = chunk_infos.iter().map(|info| info.dst_hash.0).collect();
let fetch_limiter = self.controller().fetch.clone();
let encrypted_chunks: Vec<EncryptedChunk> = rebucketed_ordered(
&fetch_limiter,
addresses.into_iter().enumerate(),
|(idx, address)| {
async move {
let chunk = self.chunk_get_observed(&address).await?.ok_or_else(|| {
Error::NotFound(format!(
"Missing chunk {} required for data reconstruction",
hex::encode(address)
))
})?;
Ok::<_, Error>((
idx,
EncryptedChunk {
content: chunk.content,
},
))
}
},
)
.await?;
debug!(
"All {} chunks retrieved, decrypting",
encrypted_chunks.len()
);
let content = decrypt(&root_data_map, &encrypted_chunks)
.map_err(|e| Error::Encryption(format!("Failed to decrypt data: {e}")))?;
info!("Data downloaded and decrypted ({} bytes)", content.len());
Ok(content)
}
async fn resolve_root_data_map(&self, data_map: &DataMap) -> Result<DataMap> {
if !data_map.is_child() {
return Ok(data_map.clone());
}
debug!("DataMap is shrunk (child); resolving root data map");
let handle = Handle::current();
ensure_shrunk_resolution_runtime(handle.runtime_flavor())?;
let mut fetch_error: Option<Error> = None;
let resolve_result = tokio::task::block_in_place(|| {
let mut get_chunk =
|name: XorName| -> std::result::Result<Bytes, self_encryption::Error> {
let address = name.0;
handle.block_on(async {
match self.chunk_get_observed(&address).await {
Ok(Some(chunk)) => Ok(chunk.content),
Ok(None) => Err(record_wrapper_fetch_error(
&mut fetch_error,
Error::NotFound(format!(
"Missing wrapper chunk {} required to resolve root DataMap",
hex::encode(address)
)),
)),
Err(e) => Err(record_wrapper_fetch_error(&mut fetch_error, e)),
}
})
};
get_root_data_map(data_map.clone(), &mut get_chunk)
});
resolve_result.map_err(|e| {
fetch_error.take().unwrap_or_else(|| {
Error::Encryption(format!("Failed to resolve root data map: {e}"))
})
})
}
}
fn record_wrapper_fetch_error(slot: &mut Option<Error>, error: Error) -> self_encryption::Error {
let message = error.to_string();
*slot = Some(error);
self_encryption::Error::Generic(message)
}
fn ensure_shrunk_resolution_runtime(flavor: RuntimeFlavor) -> Result<()> {
if flavor == RuntimeFlavor::CurrentThread {
return Err(Error::Config(
"resolving a shrunk DataMap requires a multi-threaded Tokio runtime, \
but data_download was called on a current-thread runtime"
.to_string(),
));
}
Ok(())
}
fn decode_data_map_chunk(content: &[u8]) -> Result<DataMap> {
rmp_serde::from_slice(content)
.map_err(|e| Error::Serialization(format!("Failed to deserialize DataMap: {e}")))
}
#[cfg(test)]
mod send_assertions {
use super::*;
fn _assert_send<T: Send>(_: &T) {}
#[allow(
dead_code,
unreachable_code,
unused_variables,
clippy::diverging_sub_expression
)]
async fn _data_download_is_send(client: &Client) {
let dm: DataMap = todo!();
let fut = client.data_download(&dm);
_assert_send(&fut);
}
#[allow(dead_code, unreachable_code, clippy::diverging_sub_expression)]
async fn _data_upload_is_send(client: &Client) {
let fut = client.data_upload(Bytes::new());
_assert_send(&fut);
}
#[allow(dead_code, unreachable_code, clippy::diverging_sub_expression)]
async fn _data_upload_with_mode_is_send(client: &Client) {
let fut = client.data_upload_with_mode(Bytes::new(), PaymentMode::Auto);
_assert_send(&fut);
}
#[allow(dead_code, unreachable_code, clippy::diverging_sub_expression)]
async fn _data_prepare_upload_is_send(client: &Client) {
let fut = client.data_prepare_upload(Bytes::new());
_assert_send(&fut);
}
#[allow(dead_code, unreachable_code, clippy::diverging_sub_expression)]
async fn _data_prepare_upload_with_visibility_is_send(client: &Client) {
let fut = client.data_prepare_upload_with_visibility(Bytes::new(), Visibility::Public);
_assert_send(&fut);
}
}
#[cfg(test)]
mod runtime_guard_tests {
use super::*;
#[test]
fn shrunk_resolution_rejects_current_thread_runtime() {
assert!(matches!(
ensure_shrunk_resolution_runtime(RuntimeFlavor::CurrentThread),
Err(Error::Config(_))
));
}
#[test]
fn shrunk_resolution_accepts_multi_thread_runtime() {
assert!(ensure_shrunk_resolution_runtime(RuntimeFlavor::MultiThread).is_ok());
}
}