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

ant_node/ant_protocol/
chunk.rs
