Skip to main content

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