Skip to main content

ant_protocol/
chunk.rs

1//! Chunk message types for the ANT protocol.
2//!
3//! Chunks are immutable, content-addressed data blocks where the address
4//! is the BLAKE3 hash of the content. Maximum size is 4MB.
5//!
6//! This module defines the wire protocol messages for chunk operations
7//! using postcard serialization for compact, fast encoding.
8
9use serde::{Deserialize, Serialize};
10
11/// Protocol identifier for chunk operations.
12pub const CHUNK_PROTOCOL_ID: &str = "autonomi.ant.chunk.v1";
13
14/// Current protocol version.
15pub const PROTOCOL_VERSION: u16 = 1;
16
17/// Maximum chunk size in bytes (4MB).
18pub const MAX_CHUNK_SIZE: usize = 4 * 1024 * 1024;
19
20/// Maximum wire message size in bytes (5MB).
21///
22/// Limits the input buffer accepted by [`ChunkMessage::decode`] to prevent
23/// unbounded allocation from malicious or corrupted payloads. Set slightly
24/// above [`MAX_CHUNK_SIZE`] to accommodate message envelope overhead.
25pub const MAX_WIRE_MESSAGE_SIZE: usize = 5 * 1024 * 1024;
26
27/// Data type identifier for chunks.
28pub const DATA_TYPE_CHUNK: u32 = 0;
29
30/// Number of nodes in a Kademlia close group.
31///
32/// Clients fetch quotes from the `CLOSE_GROUP_SIZE` closest nodes to a target
33/// address and select the median-priced quote for payment.
34pub const CLOSE_GROUP_SIZE: usize = 7;
35
36/// Minimum number of close group members that must agree for a decision to be valid.
37///
38/// This is a simple majority: `(CLOSE_GROUP_SIZE / 2) + 1`.
39pub const CLOSE_GROUP_MAJORITY: usize = (CLOSE_GROUP_SIZE / 2) + 1;
40
41/// Content-addressed identifier (32 bytes).
42pub type XorName = [u8; 32];
43
44/// Byte length of an [`XorName`].
45pub const XORNAME_LEN: usize = std::mem::size_of::<XorName>();
46
47/// Enum of all chunk protocol message types.
48///
49/// Uses a single-byte discriminant for efficient wire encoding.
50///
51/// Marked `#[non_exhaustive]` so new message variants can be added
52/// in a minor release without breaking downstream `match` expressions.
53#[derive(Debug, Clone, Serialize, Deserialize)]
54#[non_exhaustive]
55pub enum ChunkMessageBody {
56    /// Request to store a chunk.
57    PutRequest(ChunkPutRequest),
58    /// Response to a PUT request.
59    PutResponse(ChunkPutResponse),
60    /// Request to retrieve a chunk.
61    GetRequest(ChunkGetRequest),
62    /// Response to a GET request.
63    GetResponse(ChunkGetResponse),
64    /// Request a storage quote.
65    QuoteRequest(ChunkQuoteRequest),
66    /// Response with a storage quote.
67    QuoteResponse(ChunkQuoteResponse),
68    /// Request a merkle candidate quote for batch payments.
69    MerkleCandidateQuoteRequest(MerkleCandidateQuoteRequest),
70    /// Response with a merkle candidate quote.
71    MerkleCandidateQuoteResponse(MerkleCandidateQuoteResponse),
72}
73
74/// Wire-format wrapper that pairs a sender-assigned `request_id` with
75/// a [`ChunkMessageBody`].
76///
77/// The sender picks a unique `request_id`; the handler echoes it back
78/// in the response so callers can correlate replies by ID rather than
79/// by source peer.
80#[derive(Debug, Clone, Serialize, Deserialize)]
81pub struct ChunkMessage {
82    /// Sender-assigned identifier, echoed back in the response.
83    pub request_id: u64,
84    /// The protocol message body.
85    pub body: ChunkMessageBody,
86}
87
88impl ChunkMessage {
89    /// Encode the message to bytes using postcard.
90    ///
91    /// # Errors
92    ///
93    /// Returns an error if serialization fails.
94    pub fn encode(&self) -> Result<Vec<u8>, ProtocolError> {
95        postcard::to_stdvec(self).map_err(|e| ProtocolError::SerializationFailed(e.to_string()))
96    }
97
98    /// Decode a message from bytes using postcard.
99    ///
100    /// Rejects payloads larger than [`MAX_WIRE_MESSAGE_SIZE`] before
101    /// attempting deserialization.
102    ///
103    /// # Errors
104    ///
105    /// Returns [`ProtocolError::MessageTooLarge`] if the input exceeds the
106    /// size limit, or [`ProtocolError::DeserializationFailed`] if postcard
107    /// cannot parse the data.
108    pub fn decode(data: &[u8]) -> Result<Self, ProtocolError> {
109        if data.len() > MAX_WIRE_MESSAGE_SIZE {
110            return Err(ProtocolError::MessageTooLarge {
111                size: data.len(),
112                max_size: MAX_WIRE_MESSAGE_SIZE,
113            });
114        }
115        postcard::from_bytes(data).map_err(|e| ProtocolError::DeserializationFailed(e.to_string()))
116    }
117}
118
119// =============================================================================
120// PUT Request/Response
121// =============================================================================
122
123/// Request to store a chunk.
124#[derive(Debug, Clone, Serialize, Deserialize)]
125pub struct ChunkPutRequest {
126    /// The content-addressed identifier (BLAKE3 of content).
127    pub address: XorName,
128    /// The chunk data.
129    pub content: Vec<u8>,
130    /// Optional payment proof (serialized `ProofOfPayment`).
131    /// Required for new chunks unless already verified.
132    pub payment_proof: Option<Vec<u8>>,
133}
134
135impl ChunkPutRequest {
136    /// Create a new PUT request.
137    #[must_use]
138    pub fn new(address: XorName, content: Vec<u8>) -> Self {
139        Self {
140            address,
141            content,
142            payment_proof: None,
143        }
144    }
145
146    /// Create a new PUT request with payment proof.
147    #[must_use]
148    pub fn with_payment(address: XorName, content: Vec<u8>, payment_proof: Vec<u8>) -> Self {
149        Self {
150            address,
151            content,
152            payment_proof: Some(payment_proof),
153        }
154    }
155}
156
157/// Response to a PUT request.
158#[derive(Debug, Clone, Serialize, Deserialize)]
159#[non_exhaustive]
160pub enum ChunkPutResponse {
161    /// Chunk stored successfully.
162    Success {
163        /// The address where the chunk was stored.
164        address: XorName,
165    },
166    /// Chunk already exists (idempotent success).
167    AlreadyExists {
168        /// The existing chunk address.
169        address: XorName,
170    },
171    /// Payment is required to store this chunk.
172    PaymentRequired {
173        /// Error message.
174        message: String,
175    },
176    /// An error occurred.
177    Error(ProtocolError),
178}
179
180// =============================================================================
181// GET Request/Response
182// =============================================================================
183
184/// Request to retrieve a chunk.
185#[derive(Debug, Clone, Serialize, Deserialize)]
186pub struct ChunkGetRequest {
187    /// The content-addressed identifier to retrieve.
188    pub address: XorName,
189}
190
191impl ChunkGetRequest {
192    /// Create a new GET request.
193    #[must_use]
194    pub fn new(address: XorName) -> Self {
195        Self { address }
196    }
197}
198
199/// Response to a GET request.
200#[derive(Debug, Clone, Serialize, Deserialize)]
201#[non_exhaustive]
202pub enum ChunkGetResponse {
203    /// Chunk found and returned.
204    Success {
205        /// The chunk address.
206        address: XorName,
207        /// The chunk data.
208        content: Vec<u8>,
209    },
210    /// Chunk not found.
211    NotFound {
212        /// The requested address.
213        address: XorName,
214    },
215    /// An error occurred.
216    Error(ProtocolError),
217}
218
219// =============================================================================
220// Quote Request/Response
221// =============================================================================
222
223/// Request a storage quote for a chunk.
224#[derive(Debug, Clone, Serialize, Deserialize)]
225pub struct ChunkQuoteRequest {
226    /// The content address of the data to store.
227    pub address: XorName,
228    /// Size of the data in bytes.
229    pub data_size: u64,
230    /// Data type identifier (0 for chunks).
231    pub data_type: u32,
232}
233
234impl ChunkQuoteRequest {
235    /// Create a new quote request.
236    #[must_use]
237    pub fn new(address: XorName, data_size: u64) -> Self {
238        Self {
239            address,
240            data_size,
241            data_type: DATA_TYPE_CHUNK,
242        }
243    }
244}
245
246/// Response with a storage quote.
247#[derive(Debug, Clone, Serialize, Deserialize)]
248#[non_exhaustive]
249pub enum ChunkQuoteResponse {
250    /// Quote generated successfully.
251    ///
252    /// When `already_stored` is `true` the node already holds this chunk and no
253    /// payment is required — the client should skip the pay-then-PUT cycle for
254    /// this address. The quote is still included for informational purposes.
255    Success {
256        /// Serialized `PaymentQuote`.
257        quote: Vec<u8>,
258        /// `true` when the chunk already exists on this node (skip payment).
259        already_stored: bool,
260    },
261    /// Quote generation failed.
262    Error(ProtocolError),
263}
264
265// =============================================================================
266// Merkle Candidate Quote Request/Response
267// =============================================================================
268
269/// Request a merkle candidate quote for batch payments.
270///
271/// Part of the merkle batch payment system where clients collect
272/// signed candidate quotes from 16 closest peers per pool.
273#[derive(Debug, Clone, Serialize, Deserialize)]
274pub struct MerkleCandidateQuoteRequest {
275    /// The candidate pool address (hash of midpoint || root || timestamp).
276    pub address: XorName,
277    /// Data type identifier (0 for chunks).
278    pub data_type: u32,
279    /// Size of the data in bytes.
280    pub data_size: u64,
281    /// Client-provided merkle payment timestamp (unix seconds).
282    pub merkle_payment_timestamp: u64,
283}
284
285/// Response with a merkle candidate quote.
286#[derive(Debug, Clone, Serialize, Deserialize)]
287#[non_exhaustive]
288pub enum MerkleCandidateQuoteResponse {
289    /// Candidate quote generated successfully.
290    /// Contains the serialized `MerklePaymentCandidateNode`.
291    Success {
292        /// Serialized `MerklePaymentCandidateNode`.
293        candidate_node: Vec<u8>,
294    },
295    /// Quote generation failed.
296    Error(ProtocolError),
297}
298
299// =============================================================================
300// Payment Proof Type Tags
301// =============================================================================
302
303/// Version byte prefix for payment proof serialization.
304/// Allows the verifier to detect proof type before deserialization.
305pub const PROOF_TAG_SINGLE_NODE: u8 = 0x01;
306/// Version byte prefix for merkle payment proofs.
307pub const PROOF_TAG_MERKLE: u8 = 0x02;
308
309// =============================================================================
310// Protocol Errors
311// =============================================================================
312
313/// Errors that can occur during protocol operations.
314#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
315#[non_exhaustive]
316pub enum ProtocolError {
317    /// Message serialization failed.
318    SerializationFailed(String),
319    /// Message deserialization failed.
320    DeserializationFailed(String),
321    /// Wire message exceeds the maximum allowed size.
322    MessageTooLarge {
323        /// Actual size of the message in bytes.
324        size: usize,
325        /// Maximum allowed size.
326        max_size: usize,
327    },
328    /// Chunk exceeds maximum size.
329    ChunkTooLarge {
330        /// Size of the chunk in bytes.
331        size: usize,
332        /// Maximum allowed size.
333        max_size: usize,
334    },
335    /// Content address mismatch (hash(content) != address).
336    AddressMismatch {
337        /// Expected address.
338        expected: XorName,
339        /// Actual address computed from content.
340        actual: XorName,
341    },
342    /// Storage operation failed.
343    StorageFailed(String),
344    /// Payment verification failed.
345    PaymentFailed(String),
346    /// Quote generation failed.
347    QuoteFailed(String),
348    /// Internal error.
349    Internal(String),
350}
351
352impl std::fmt::Display for ProtocolError {
353    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
354        match self {
355            Self::SerializationFailed(msg) => write!(f, "serialization failed: {msg}"),
356            Self::DeserializationFailed(msg) => write!(f, "deserialization failed: {msg}"),
357            Self::MessageTooLarge { size, max_size } => {
358                write!(f, "message size {size} exceeds maximum {max_size}")
359            }
360            Self::ChunkTooLarge { size, max_size } => {
361                write!(f, "chunk size {size} exceeds maximum {max_size}")
362            }
363            Self::AddressMismatch { expected, actual } => {
364                write!(
365                    f,
366                    "address mismatch: expected {}, got {}",
367                    hex::encode(expected),
368                    hex::encode(actual)
369                )
370            }
371            Self::StorageFailed(msg) => write!(f, "storage failed: {msg}"),
372            Self::PaymentFailed(msg) => write!(f, "payment failed: {msg}"),
373            Self::QuoteFailed(msg) => write!(f, "quote failed: {msg}"),
374            Self::Internal(msg) => write!(f, "internal error: {msg}"),
375        }
376    }
377}
378
379impl std::error::Error for ProtocolError {}
380
381#[cfg(test)]
382#[allow(clippy::unwrap_used, clippy::expect_used, clippy::panic)]
383mod tests {
384    use super::*;
385
386    #[test]
387    fn test_put_request_encode_decode() {
388        let address = [0xAB; 32];
389        let content = vec![1, 2, 3, 4, 5];
390        let request = ChunkPutRequest::new(address, content.clone());
391        let msg = ChunkMessage {
392            request_id: 42,
393            body: ChunkMessageBody::PutRequest(request),
394        };
395
396        let encoded = msg.encode().expect("encode should succeed");
397        let decoded = ChunkMessage::decode(&encoded).expect("decode should succeed");
398
399        assert_eq!(decoded.request_id, 42);
400        if let ChunkMessageBody::PutRequest(req) = decoded.body {
401            assert_eq!(req.address, address);
402            assert_eq!(req.content, content);
403            assert!(req.payment_proof.is_none());
404        } else {
405            panic!("expected PutRequest");
406        }
407    }
408
409    #[test]
410    fn test_put_request_with_payment() {
411        let address = [0xAB; 32];
412        let content = vec![1, 2, 3, 4, 5];
413        let payment = vec![10, 20, 30];
414        let request = ChunkPutRequest::with_payment(address, content.clone(), payment.clone());
415
416        assert_eq!(request.address, address);
417        assert_eq!(request.content, content);
418        assert_eq!(request.payment_proof, Some(payment));
419    }
420
421    #[test]
422    fn test_get_request_encode_decode() {
423        let address = [0xCD; 32];
424        let request = ChunkGetRequest::new(address);
425        let msg = ChunkMessage {
426            request_id: 7,
427            body: ChunkMessageBody::GetRequest(request),
428        };
429
430        let encoded = msg.encode().expect("encode should succeed");
431        let decoded = ChunkMessage::decode(&encoded).expect("decode should succeed");
432
433        assert_eq!(decoded.request_id, 7);
434        if let ChunkMessageBody::GetRequest(req) = decoded.body {
435            assert_eq!(req.address, address);
436        } else {
437            panic!("expected GetRequest");
438        }
439    }
440
441    #[test]
442    fn test_put_response_success() {
443        let address = [0xEF; 32];
444        let response = ChunkPutResponse::Success { address };
445        let msg = ChunkMessage {
446            request_id: 99,
447            body: ChunkMessageBody::PutResponse(response),
448        };
449
450        let encoded = msg.encode().expect("encode should succeed");
451        let decoded = ChunkMessage::decode(&encoded).expect("decode should succeed");
452
453        assert_eq!(decoded.request_id, 99);
454        if let ChunkMessageBody::PutResponse(ChunkPutResponse::Success { address: addr }) =
455            decoded.body
456        {
457            assert_eq!(addr, address);
458        } else {
459            panic!("expected PutResponse::Success");
460        }
461    }
462
463    #[test]
464    fn test_get_response_not_found() {
465        let address = [0x12; 32];
466        let response = ChunkGetResponse::NotFound { address };
467        let msg = ChunkMessage {
468            request_id: 0,
469            body: ChunkMessageBody::GetResponse(response),
470        };
471
472        let encoded = msg.encode().expect("encode should succeed");
473        let decoded = ChunkMessage::decode(&encoded).expect("decode should succeed");
474
475        assert_eq!(decoded.request_id, 0);
476        if let ChunkMessageBody::GetResponse(ChunkGetResponse::NotFound { address: addr }) =
477            decoded.body
478        {
479            assert_eq!(addr, address);
480        } else {
481            panic!("expected GetResponse::NotFound");
482        }
483    }
484
485    #[test]
486    fn test_quote_request_encode_decode() {
487        let address = [0x34; 32];
488        let request = ChunkQuoteRequest::new(address, 1024);
489        let msg = ChunkMessage {
490            request_id: 1,
491            body: ChunkMessageBody::QuoteRequest(request),
492        };
493
494        let encoded = msg.encode().expect("encode should succeed");
495        let decoded = ChunkMessage::decode(&encoded).expect("decode should succeed");
496
497        assert_eq!(decoded.request_id, 1);
498        if let ChunkMessageBody::QuoteRequest(req) = decoded.body {
499            assert_eq!(req.address, address);
500            assert_eq!(req.data_size, 1024);
501            assert_eq!(req.data_type, DATA_TYPE_CHUNK);
502        } else {
503            panic!("expected QuoteRequest");
504        }
505    }
506
507    #[test]
508    fn test_protocol_error_display() {
509        let err = ProtocolError::ChunkTooLarge {
510            size: 5_000_000,
511            max_size: MAX_CHUNK_SIZE,
512        };
513        assert!(err.to_string().contains("5000000"));
514        assert!(err.to_string().contains(&MAX_CHUNK_SIZE.to_string()));
515
516        let err = ProtocolError::AddressMismatch {
517            expected: [0xAA; 32],
518            actual: [0xBB; 32],
519        };
520        let display = err.to_string();
521        assert!(display.contains("address mismatch"));
522    }
523
524    #[test]
525    fn test_decode_rejects_oversized_payload() {
526        let oversized = vec![0u8; MAX_WIRE_MESSAGE_SIZE + 1];
527        let result = ChunkMessage::decode(&oversized);
528        assert!(result.is_err());
529        let err = result.unwrap_err();
530        assert!(
531            matches!(err, ProtocolError::MessageTooLarge { .. }),
532            "expected MessageTooLarge, got {err:?}"
533        );
534    }
535
536    #[test]
537    fn test_invalid_decode() {
538        let invalid_data = vec![0xFF, 0xFF, 0xFF];
539        let result = ChunkMessage::decode(&invalid_data);
540        assert!(result.is_err());
541    }
542
543    #[test]
544    fn test_constants() {
545        assert_eq!(CHUNK_PROTOCOL_ID, "autonomi.ant.chunk.v1");
546        assert_eq!(PROTOCOL_VERSION, 1);
547        assert_eq!(MAX_CHUNK_SIZE, 4 * 1024 * 1024);
548        assert_eq!(DATA_TYPE_CHUNK, 0);
549    }
550
551    #[test]
552    fn test_proof_tag_constants() {
553        // Tags must be distinct non-zero bytes
554        assert_ne!(PROOF_TAG_SINGLE_NODE, PROOF_TAG_MERKLE);
555        assert_ne!(PROOF_TAG_SINGLE_NODE, 0x00);
556        assert_ne!(PROOF_TAG_MERKLE, 0x00);
557        assert_eq!(PROOF_TAG_SINGLE_NODE, 0x01);
558        assert_eq!(PROOF_TAG_MERKLE, 0x02);
559    }
560
561    #[test]
562    fn test_merkle_candidate_quote_request_encode_decode() {
563        let address = [0x56; 32];
564        let request = MerkleCandidateQuoteRequest {
565            address,
566            data_type: DATA_TYPE_CHUNK,
567            data_size: 2048,
568            merkle_payment_timestamp: 1_700_000_000,
569        };
570        let msg = ChunkMessage {
571            request_id: 500,
572            body: ChunkMessageBody::MerkleCandidateQuoteRequest(request),
573        };
574
575        let encoded = msg.encode().expect("encode should succeed");
576        let decoded = ChunkMessage::decode(&encoded).expect("decode should succeed");
577
578        assert_eq!(decoded.request_id, 500);
579        if let ChunkMessageBody::MerkleCandidateQuoteRequest(req) = decoded.body {
580            assert_eq!(req.address, address);
581            assert_eq!(req.data_type, DATA_TYPE_CHUNK);
582            assert_eq!(req.data_size, 2048);
583            assert_eq!(req.merkle_payment_timestamp, 1_700_000_000);
584        } else {
585            panic!("expected MerkleCandidateQuoteRequest");
586        }
587    }
588
589    #[test]
590    fn test_merkle_candidate_quote_response_success_encode_decode() {
591        let candidate_node_bytes = vec![0xAA, 0xBB, 0xCC, 0xDD];
592        let response = MerkleCandidateQuoteResponse::Success {
593            candidate_node: candidate_node_bytes.clone(),
594        };
595        let msg = ChunkMessage {
596            request_id: 501,
597            body: ChunkMessageBody::MerkleCandidateQuoteResponse(response),
598        };
599
600        let encoded = msg.encode().expect("encode should succeed");
601        let decoded = ChunkMessage::decode(&encoded).expect("decode should succeed");
602
603        assert_eq!(decoded.request_id, 501);
604        if let ChunkMessageBody::MerkleCandidateQuoteResponse(
605            MerkleCandidateQuoteResponse::Success { candidate_node },
606        ) = decoded.body
607        {
608            assert_eq!(candidate_node, candidate_node_bytes);
609        } else {
610            panic!("expected MerkleCandidateQuoteResponse::Success");
611        }
612    }
613
614    #[test]
615    fn test_merkle_candidate_quote_response_error_encode_decode() {
616        let error = ProtocolError::QuoteFailed("no libp2p keypair".to_string());
617        let response = MerkleCandidateQuoteResponse::Error(error.clone());
618        let msg = ChunkMessage {
619            request_id: 502,
620            body: ChunkMessageBody::MerkleCandidateQuoteResponse(response),
621        };
622
623        let encoded = msg.encode().expect("encode should succeed");
624        let decoded = ChunkMessage::decode(&encoded).expect("decode should succeed");
625
626        assert_eq!(decoded.request_id, 502);
627        if let ChunkMessageBody::MerkleCandidateQuoteResponse(
628            MerkleCandidateQuoteResponse::Error(err),
629        ) = decoded.body
630        {
631            assert_eq!(err, error);
632        } else {
633            panic!("expected MerkleCandidateQuoteResponse::Error");
634        }
635    }
636}