diff --git a/antd/src/grpc/service.rs b/antd/src/grpc/service.rs index 0e72ae0..4ad4fd5 100644 --- a/antd/src/grpc/service.rs +++ b/antd/src/grpc/service.rs @@ -6,7 +6,7 @@ use tonic::{Request, Response, Status}; use crate::error::AntdError; use crate::state::AppState; -use crate::types::format_payment_mode; +use crate::types::{adjust_for_public_upload, format_payment_mode}; // Generated protobuf modules #[allow(dead_code)] @@ -536,10 +536,16 @@ impl pb::file_service_server::FileService for FileServiceImpl { .map_err(AntdError::from_core) .map_err(tonic::Status::from)?; + let (chunk_count, atto_tokens) = if req.is_public { + adjust_for_public_upload(estimate.chunk_count, &estimate.storage_cost_atto) + } else { + (estimate.chunk_count, estimate.storage_cost_atto) + }; + Ok(Response::new(pb::Cost { - atto_tokens: estimate.storage_cost_atto, + atto_tokens, file_size: estimate.file_size, - chunk_count: estimate.chunk_count as u32, + chunk_count: chunk_count as u32, estimated_gas_cost_wei: estimate.estimated_gas_cost_wei, payment_mode: format_payment_mode(estimate.payment_mode), })) diff --git a/antd/src/rest/files.rs b/antd/src/rest/files.rs index 73ca3a9..245aeb9 100644 --- a/antd/src/rest/files.rs +++ b/antd/src/rest/files.rs @@ -218,10 +218,16 @@ pub async fn file_cost( .map_err(|e| AntdError::Internal(format!("task failed: {e}")))? .map_err(AntdError::from_core)?; + let (chunk_count, cost) = if req.is_public { + adjust_for_public_upload(estimate.chunk_count, &estimate.storage_cost_atto) + } else { + (estimate.chunk_count, estimate.storage_cost_atto) + }; + Ok(Json(CostResponse { - cost: estimate.storage_cost_atto, + cost, file_size: estimate.file_size, - chunk_count: estimate.chunk_count, + chunk_count, estimated_gas_cost_wei: estimate.estimated_gas_cost_wei, payment_mode: format_payment_mode(estimate.payment_mode), })) diff --git a/antd/src/types.rs b/antd/src/types.rs index e363bc8..ef21d3a 100644 --- a/antd/src/types.rs +++ b/antd/src/types.rs @@ -225,7 +225,6 @@ pub struct CostResponse { #[derive(Deserialize)] pub struct FileCostRequest { pub path: String, - #[allow(dead_code)] #[serde(default = "default_true")] pub is_public: bool, } @@ -234,6 +233,27 @@ fn default_true() -> bool { true } +/// Adjust a chunk count + storage cost estimate to reflect a public upload, +/// which bundles one additional DataMap chunk into the same payment batch. +/// +/// Returns `(adjusted_chunk_count, adjusted_storage_cost_atto)`. If the input +/// cost cannot be parsed, only the chunk count is bumped and the cost is +/// returned unchanged. +pub fn adjust_for_public_upload(chunk_count: usize, storage_cost_atto: &str) -> (usize, String) { + let new_chunk_count = chunk_count.saturating_add(1); + if chunk_count == 0 { + return (new_chunk_count, storage_cost_atto.to_string()); + } + let total: ant_core::data::U256 = match storage_cost_atto.parse() { + Ok(v) => v, + Err(_) => return (new_chunk_count, storage_cost_atto.to_string()), + }; + let divisor = ant_core::data::U256::from(chunk_count as u64); + let per_chunk = total / divisor; + let new_total = total + per_chunk; + (new_chunk_count, new_total.to_string()) +} + /// Parse a payment mode string into ant-core's PaymentMode. pub fn parse_payment_mode(mode: Option<&str>) -> Result { match mode { @@ -451,4 +471,59 @@ mod tests { assert_eq!(json["payment_token_address"], ""); assert_eq!(json["payment_vault_address"], ""); } + + #[test] + fn adjust_for_public_upload_bumps_count_and_scales_cost() { + // 5 chunks, 1000 atto total → per-chunk = 200 atto. + // Public upload adds 1 chunk → 6 chunks, 1200 atto total. + let (chunks, cost) = adjust_for_public_upload(5, "1000"); + assert_eq!(chunks, 6); + assert_eq!(cost, "1200"); + } + + #[test] + fn adjust_for_public_upload_handles_uneven_division() { + // 3 chunks, 100 atto total → per-chunk = 33 atto (integer divide). + // Result: 4 chunks, 133 atto. Slight rounding is acceptable for + // an estimate; exact pricing is the on-chain quote. + let (chunks, cost) = adjust_for_public_upload(3, "100"); + assert_eq!(chunks, 4); + assert_eq!(cost, "133"); + } + + #[test] + fn adjust_for_public_upload_handles_large_atto_values() { + // Real-world atto costs frequently exceed u64. Verify U256 handles it. + // 10 chunks at 1e22 atto total = 10K ANT — well above u64::MAX. + let total = "10000000000000000000000"; // 1e22 + let (chunks, cost) = adjust_for_public_upload(10, total); + assert_eq!(chunks, 11); + // 1e22 + 1e21 = 1.1e22 + assert_eq!(cost, "11000000000000000000000"); + } + + #[test] + fn adjust_for_public_upload_zero_chunks_only_bumps_count() { + // Defensive: shouldn't happen in practice, but division-by-zero would. + let (chunks, cost) = adjust_for_public_upload(0, "0"); + assert_eq!(chunks, 1); + assert_eq!(cost, "0"); + } + + #[test] + fn adjust_for_public_upload_unparseable_cost_only_bumps_count() { + // ant-core returns "0" for already-stored chunks; a non-numeric or + // negative value would be a contract bug, but stay graceful. + let (chunks, cost) = adjust_for_public_upload(5, "not-a-number"); + assert_eq!(chunks, 6); + assert_eq!(cost, "not-a-number"); + } + + #[test] + fn adjust_for_public_upload_zero_cost_round_trips() { + // Already-stored case: cost stays 0, count still bumps. + let (chunks, cost) = adjust_for_public_upload(5, "0"); + assert_eq!(chunks, 6); + assert_eq!(cost, "0"); + } }