Skip to main content

nectar_primitives/chunk/
content.rs

1//! Content-addressed chunk implementation
2//!
3//! This module provides the implementation of content-addressed chunks,
4//! which are chunks whose address is derived from the hash of their content.
5
6use alloy_primitives::hex;
7use bytes::Bytes;
8use std::fmt;
9use std::marker::PhantomData;
10
11use crate::bmt::DEFAULT_BODY_SIZE;
12use crate::cache::OnceCache;
13use crate::error::{PrimitivesError, Result};
14
15use super::bmt_body::BmtBody;
16use super::traits::{BmtChunk, Chunk, ChunkAddress, ChunkHeader, ChunkMetadata};
17
18/// A content-addressed chunk with configurable body size.
19///
20/// This type represents a chunk of data whose address is derived from the hash
21/// of its contents. It is immutable once created.
22#[derive(Debug, Clone)]
23pub struct ContentChunk<const BODY_SIZE: usize = DEFAULT_BODY_SIZE> {
24    /// The header containing type ID, version, and metadata
25    header: ContentChunkHeader,
26    /// The body of the chunk, containing the actual data
27    body: BmtBody<BODY_SIZE>,
28    /// Cache for the chunk's address
29    address_cache: OnceCache<ChunkAddress>,
30}
31
32/// Metadata for a content-addressed chunk
33///
34/// Content chunks don't have any specific metadata, so this is empty.
35#[derive(Debug, Clone)]
36pub struct ContentChunkMetadata;
37
38impl ChunkMetadata for ContentChunkMetadata {
39    fn bytes(&self) -> Bytes {
40        Bytes::new()
41    }
42}
43
44/// Header for a content-addressed chunk
45#[derive(Debug, Clone)]
46pub struct ContentChunkHeader {
47    metadata: ContentChunkMetadata,
48}
49
50impl ContentChunkHeader {
51    /// Create a new header with default metadata
52    pub const fn new() -> Self {
53        Self {
54            metadata: ContentChunkMetadata,
55        }
56    }
57}
58
59impl Default for ContentChunkHeader {
60    fn default() -> Self {
61        Self::new()
62    }
63}
64
65impl ChunkHeader for ContentChunkHeader {
66    type Metadata = ContentChunkMetadata;
67
68    fn id(&self) -> u8 {
69        0
70    }
71
72    fn version(&self) -> u8 {
73        1
74    }
75
76    fn metadata(&self) -> &Self::Metadata {
77        &self.metadata
78    }
79
80    fn bytes(&self) -> Bytes {
81        Bytes::new()
82    }
83}
84
85impl<const BODY_SIZE: usize> ContentChunk<BODY_SIZE> {
86    /// Create a new content chunk with the given data.
87    ///
88    /// This function automatically calculates the span based on the data length.
89    ///
90    /// # Arguments
91    ///
92    /// * `data` - The raw data content to encapsulate in the chunk.
93    ///
94    /// # Returns
95    ///
96    /// A Result containing the new ContentChunk, or an error if creation fails.
97    #[must_use = "this returns a new chunk without modifying the input"]
98    pub fn new(data: impl Into<Bytes>) -> Result<Self> {
99        Ok(ContentChunkBuilderImpl::<BODY_SIZE, _>::default()
100            .auto_from_data(data)?
101            .build())
102    }
103
104    /// Create a new ContentChunk with a pre-computed address.
105    ///
106    /// This function is useful when the address is already known, for example
107    /// when retrieving a chunk from a database.
108    ///
109    /// # Arguments
110    ///
111    /// * `data` - The raw data content to encapsulate in the chunk.
112    /// * `address` - The pre-computed address of the chunk.
113    ///
114    /// # Returns
115    ///
116    /// A Result containing the new ContentChunk, or an error if creation fails.
117    #[must_use = "this returns a new chunk without modifying the input"]
118    pub fn with_address(data: impl Into<Bytes>, address: ChunkAddress) -> Result<Self> {
119        Ok(ContentChunkBuilderImpl::<BODY_SIZE, _>::default()
120            .auto_from_data(data)?
121            .with_address(address)
122            .build())
123    }
124
125    /// Create a ContentChunk from a pre-existing BmtBody.
126    ///
127    /// This is an advanced method for when you already have a BmtBody,
128    /// such as when reconstructing chunks from storage or building
129    /// intermediate nodes in a merkle tree.
130    #[must_use]
131    pub const fn from_body(body: BmtBody<BODY_SIZE>) -> Self {
132        Self {
133            header: ContentChunkHeader::new(),
134            body,
135            address_cache: OnceCache::new(),
136        }
137    }
138
139    /// Borrow the BMT body of this content chunk.
140    ///
141    /// The body carries the chunk's `span`, `payload`, and the `BODY_SIZE`
142    /// const, so this is the zero-copy accessor callers use to feed the body
143    /// into BMT operations (e.g. [`BmtBody::transformed_root`]) without
144    /// re-slicing the span/payload back out of the wire form.
145    pub const fn body(&self) -> &BmtBody<BODY_SIZE> {
146        &self.body
147    }
148
149    /// Create a ContentChunk from a pre-existing BmtBody with a known address.
150    ///
151    /// This is an advanced method for when you already have both the body
152    /// and know the chunk's address (e.g., when reconstructing from storage).
153    #[must_use]
154    pub fn from_body_with_address(body: BmtBody<BODY_SIZE>, address: ChunkAddress) -> Self {
155        Self {
156            header: ContentChunkHeader::new(),
157            body,
158            address_cache: OnceCache::with_value(address),
159        }
160    }
161}
162
163/// Result of encrypting a content chunk.
164#[cfg(feature = "encryption")]
165#[derive(Debug, Clone)]
166pub struct EncryptedContentChunk<const BODY_SIZE: usize = DEFAULT_BODY_SIZE> {
167    chunk: ContentChunk<BODY_SIZE>,
168    encrypted_ref: super::encryption::EncryptedChunkRef,
169}
170
171#[cfg(feature = "encryption")]
172impl<const BODY_SIZE: usize> EncryptedContentChunk<BODY_SIZE> {
173    /// The encrypted chunk (ciphertext hashed to a new address).
174    pub const fn chunk(&self) -> &ContentChunk<BODY_SIZE> {
175        &self.chunk
176    }
177
178    /// The encrypted reference (address + decryption key).
179    pub const fn encrypted_ref(&self) -> &super::encryption::EncryptedChunkRef {
180        &self.encrypted_ref
181    }
182
183    /// Consume and return (chunk, encrypted_ref).
184    pub fn into_parts(
185        self,
186    ) -> (
187        ContentChunk<BODY_SIZE>,
188        super::encryption::EncryptedChunkRef,
189    ) {
190        (self.chunk, self.encrypted_ref)
191    }
192
193    /// Decrypt back to a plaintext `ContentChunk`.
194    pub fn decrypt(&self) -> Result<ContentChunk<BODY_SIZE>> {
195        use super::encryption::transcrypt;
196        use crate::bmt::SPAN_SIZE;
197
198        let encrypted_data: Bytes = self.chunk.clone().into();
199        let key = self.encrypted_ref.key();
200
201        // Decrypt the span to learn the original data length
202        let span_ctr = (BODY_SIZE / super::encryption::EncryptionKey::SIZE) as u32;
203        let mut span_buf = [0u8; SPAN_SIZE];
204        transcrypt(key, span_ctr, &encrypted_data[..SPAN_SIZE], &mut span_buf)?;
205        let data_length = u64::from_le_bytes(span_buf) as usize;
206
207        let decrypted =
208            super::encryption::decrypt_chunk_data::<BODY_SIZE>(&encrypted_data, key, data_length)?;
209        ContentChunk::try_from(Bytes::from(decrypted))
210    }
211}
212
213#[cfg(feature = "encryption")]
214impl<const BODY_SIZE: usize> super::encryption::ChunkEncrypt for ContentChunk<BODY_SIZE> {
215    type Encrypted = EncryptedContentChunk<BODY_SIZE>;
216
217    /// Encrypt this chunk with a caller-provided key.
218    ///
219    /// The returned `EncryptedContentChunk` contains:
220    /// - `chunk`: a new `ContentChunk` whose data is the ciphertext
221    /// - `encrypted_ref`: the 64-byte reference (new address + decryption key)
222    ///
223    /// ```
224    /// # use nectar_primitives::{Chunk, ContentChunk};
225    /// # use nectar_primitives::chunk::encryption::{ChunkEncrypt, EncryptionKey};
226    /// # use nectar_primitives::bmt::DEFAULT_BODY_SIZE;
227    /// let chunk = ContentChunk::<DEFAULT_BODY_SIZE>::new(b"secret data".to_vec()).unwrap();
228    /// let encrypted = chunk.encrypt().unwrap();
229    ///
230    /// // The encrypted chunk has a different address
231    /// assert_ne!(chunk.address(), encrypted.chunk().address());
232    /// ```
233    fn encrypt_with(
234        &self,
235        key: &super::encryption::EncryptionKey,
236    ) -> Result<EncryptedContentChunk<BODY_SIZE>> {
237        let raw: Bytes = self.clone().into(); // span || data
238        let ciphertext = super::encryption::encrypt_chunk::<BODY_SIZE>(&raw, key)?;
239        let encrypted_chunk = Self::try_from(Bytes::from(ciphertext))?;
240        let encrypted_ref =
241            super::encryption::EncryptedChunkRef::new(*encrypted_chunk.address(), key.clone());
242        Ok(EncryptedContentChunk {
243            chunk: encrypted_chunk,
244            encrypted_ref,
245        })
246    }
247    // encrypt() uses default impl — generates random key, calls encrypt_with()
248}
249
250impl<const BODY_SIZE: usize> Chunk for ContentChunk<BODY_SIZE> {
251    type Header = ContentChunkHeader;
252
253    fn address(&self) -> &ChunkAddress {
254        self.address_cache.get_or_compute(|| self.body.hash())
255    }
256
257    fn data(&self) -> &Bytes {
258        self.body.data()
259    }
260
261    fn size(&self) -> usize {
262        self.header().bytes().len() + self.body.size()
263    }
264
265    fn header(&self) -> &Self::Header {
266        &self.header
267    }
268}
269
270impl<const BODY_SIZE: usize> BmtChunk for ContentChunk<BODY_SIZE> {
271    fn span(&self) -> u64 {
272        self.body.span()
273    }
274}
275
276impl<const BODY_SIZE: usize> From<ContentChunk<BODY_SIZE>> for Bytes {
277    fn from(chunk: ContentChunk<BODY_SIZE>) -> Self {
278        chunk.body.into()
279    }
280}
281
282impl<const BODY_SIZE: usize> TryFrom<Bytes> for ContentChunk<BODY_SIZE> {
283    type Error = PrimitivesError;
284
285    fn try_from(bytes: Bytes) -> Result<Self> {
286        Ok(Self {
287            header: ContentChunkHeader::new(),
288            body: BmtBody::try_from(bytes)?,
289            address_cache: OnceCache::new(),
290        })
291    }
292}
293
294impl<const BODY_SIZE: usize> TryFrom<&[u8]> for ContentChunk<BODY_SIZE> {
295    type Error = PrimitivesError;
296
297    fn try_from(bytes: &[u8]) -> Result<Self> {
298        Self::try_from(Bytes::copy_from_slice(bytes))
299    }
300}
301
302impl<const BODY_SIZE: usize> fmt::Display for ContentChunk<BODY_SIZE> {
303    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
304        write!(
305            f,
306            "ContentChunk[{}]",
307            hex::encode(&self.address().as_bytes()[..8])
308        )
309    }
310}
311
312impl<const BODY_SIZE: usize> PartialEq for ContentChunk<BODY_SIZE> {
313    fn eq(&self, other: &Self) -> bool {
314        self.address() == other.address()
315    }
316}
317
318impl<const BODY_SIZE: usize> Eq for ContentChunk<BODY_SIZE> {}
319
320impl<const BODY_SIZE: usize> super::chunk_type::ChunkType for ContentChunk<BODY_SIZE> {
321    const TYPE_ID: super::type_id::ChunkTypeId = super::type_id::ChunkTypeId::CONTENT;
322    const TYPE_NAME: &'static str = "content";
323}
324
325// Internal builder implementation
326trait BuilderState {}
327
328#[derive(Debug, Default)]
329struct Initial;
330impl BuilderState for Initial {}
331
332#[derive(Debug)]
333struct ReadyToBuild;
334impl BuilderState for ReadyToBuild {}
335
336/// Builder for ContentChunk with type state pattern
337#[derive(Debug)]
338struct ContentChunkBuilderImpl<const BODY_SIZE: usize, S: BuilderState = Initial> {
339    /// The body to use for the chunk
340    body: Option<BmtBody<BODY_SIZE>>,
341    /// Pre-computed address for the chunk
342    address: Option<ChunkAddress>,
343    /// Marker for the builder state
344    _state: PhantomData<S>,
345}
346
347impl<const BODY_SIZE: usize> Default for ContentChunkBuilderImpl<BODY_SIZE, Initial> {
348    fn default() -> Self {
349        Self {
350            body: None,
351            address: None,
352            _state: PhantomData,
353        }
354    }
355}
356
357impl<const BODY_SIZE: usize> ContentChunkBuilderImpl<BODY_SIZE, Initial> {
358    /// Initialize from data with automatically calculated span
359    fn auto_from_data(
360        mut self,
361        data: impl Into<Bytes>,
362    ) -> Result<ContentChunkBuilderImpl<BODY_SIZE, ReadyToBuild>> {
363        let body = BmtBody::<BODY_SIZE>::builder()
364            .auto_from_data(data)?
365            .build()?;
366        self.body = Some(body);
367
368        Ok(ContentChunkBuilderImpl {
369            body: self.body,
370            address: self.address,
371            _state: PhantomData,
372        })
373    }
374}
375
376impl<const BODY_SIZE: usize> ContentChunkBuilderImpl<BODY_SIZE, ReadyToBuild> {
377    /// Set a pre-computed address for the chunk
378    const fn with_address(mut self, address: ChunkAddress) -> Self {
379        self.address = Some(address);
380        self
381    }
382
383    /// Build the final ContentChunk
384    fn build(self) -> ContentChunk<BODY_SIZE> {
385        // This is safe as we have already checked that the body is set
386        let body = self.body.unwrap();
387
388        let address_cache = self
389            .address
390            .map_or_else(OnceCache::new, OnceCache::with_value);
391
392        ContentChunk {
393            header: ContentChunkHeader::new(),
394            body,
395            address_cache,
396        }
397    }
398}
399
400#[cfg(any(test, feature = "arbitrary"))]
401impl<'a, const BODY_SIZE: usize> arbitrary::Arbitrary<'a> for ContentChunk<BODY_SIZE> {
402    fn arbitrary(u: &mut arbitrary::Unstructured<'a>) -> arbitrary::Result<Self> {
403        Ok(Self::from_body(BmtBody::<BODY_SIZE>::arbitrary(u)?))
404    }
405}
406
407#[cfg(test)]
408mod tests {
409    use crate::{DEFAULT_BODY_SIZE, chunk::error::ChunkError};
410
411    use super::*;
412    use alloy_primitives::b256;
413    use proptest::prelude::*;
414    use proptest_arbitrary_interop::arb;
415
416    type DefaultContentChunk = ContentChunk<DEFAULT_BODY_SIZE>;
417
418    // Strategy for generating ContentChunk using the Arbitrary implementation
419    fn chunk_strategy() -> impl Strategy<Value = DefaultContentChunk> {
420        arb::<DefaultContentChunk>()
421    }
422
423    proptest! {
424        #[test]
425        fn test_chunk_properties(chunk in chunk_strategy()) {
426            // Test basic properties
427            prop_assert!(chunk.data().len() <= DEFAULT_BODY_SIZE);
428            prop_assert_eq!(chunk.size(), 8 + chunk.data().len());
429
430            // Test round-trip conversion
431            let bytes: Bytes = chunk.clone().into();
432            let decoded = DefaultContentChunk::try_from(bytes).unwrap();
433            prop_assert_eq!(chunk.address(), decoded.address());
434            prop_assert_eq!(chunk.data(), decoded.data());
435            prop_assert_eq!(chunk.span(), decoded.span());
436        }
437
438        #[test]
439        fn test_from_body(chunk in chunk_strategy()) {
440            // Test creating a chunk from an existing body via BmtBody
441            let body_data = chunk.data().clone();
442            let body_span = chunk.span();
443
444            // Create a new chunk using from_body with a fresh BmtBody
445            let new_body = BmtBody::<DEFAULT_BODY_SIZE>::try_from(Bytes::from(chunk.clone())).unwrap();
446            let new_chunk = DefaultContentChunk::from_body(new_body);
447
448            prop_assert_eq!(new_chunk.data(), &body_data);
449            prop_assert_eq!(new_chunk.span(), body_span);
450            prop_assert_eq!(new_chunk.address(), chunk.address());
451        }
452
453        #[test]
454        fn test_new_content_chunk(data in proptest::collection::vec(any::<u8>(), 0..DEFAULT_BODY_SIZE)) {
455            let chunk = DefaultContentChunk::new(data.clone()).unwrap();
456
457            prop_assert_eq!(chunk.data(), &data);
458            prop_assert_eq!(chunk.span(), data.len() as u64);
459            prop_assert!(!chunk.address().is_zero());
460        }
461
462        #[test]
463        fn test_chunk_size_validation(data in proptest::collection::vec(any::<u8>(), DEFAULT_BODY_SIZE + 1..DEFAULT_BODY_SIZE * 2)) {
464            let result = DefaultContentChunk::new(data);
465            prop_assert_eq!(matches!(result, Err(PrimitivesError::Chunk(ChunkError::InvalidSize { .. }))), true);
466        }
467
468        #[test]
469        fn test_empty_and_edge_cases(size in 0usize..=10usize) {
470            // Test with empty or small data
471            let data = vec![0u8; size];
472            let chunk = DefaultContentChunk::new(data).unwrap();
473
474            prop_assert_eq!(chunk.data().len(), size);
475            prop_assert_eq!(chunk.span(), size as u64);
476            prop_assert_eq!(chunk.size(), 8 + size);
477        }
478
479        #[test]
480        fn test_deserialize_invalid_chunks(data in proptest::collection::vec(any::<u8>(), 0..8)) {
481            let result = DefaultContentChunk::try_from(data.as_slice());
482            prop_assert_eq!(matches!(result, Err(PrimitivesError::Chunk(ChunkError::InvalidSize { .. }))), true);
483        }
484    }
485
486    #[test]
487    fn test_new() {
488        let data = b"greaterthanspan";
489        let bmt_hash = b256!("27913f1bdb6e8e52cbd5a5fd4ab577c857287edf6969b41efe926b51de0f4f23");
490
491        let chunk = DefaultContentChunk::new(data.to_vec()).unwrap();
492        assert_eq!(chunk.address().as_ref(), bmt_hash);
493        assert_eq!(chunk.data(), data.as_slice());
494    }
495
496    #[test]
497    fn test_from_bytes() {
498        let data = b"greaterthanspan";
499        let bmt_hash = b256!("95022e6af5c6d6a564ee55a67f8455a3e18c511b5697c932d9e44f07f2fb8c53");
500
501        let chunk = DefaultContentChunk::try_from(data.as_slice()).unwrap();
502        assert_eq!(chunk.address().as_ref(), bmt_hash);
503        assert_eq!(
504            <DefaultContentChunk as Into<Bytes>>::into(chunk),
505            data.as_slice()
506        );
507    }
508
509    #[test]
510    fn test_specific_content_hash() {
511        // Test with known valid data and hash
512        let data = b"foo".to_vec();
513        let expected_hash =
514            b256!("2387e8e7d8a48c2a9339c97c1dc3461a9a7aa07e994c5cb8b38fd7c1b3e6ea48");
515
516        let chunk = DefaultContentChunk::new(data).unwrap();
517        assert_eq!(chunk.address().as_ref(), expected_hash);
518
519        // Test with "Digital Freedom Now"
520        let data = b"Digital Freedom Now".to_vec();
521        let chunk = DefaultContentChunk::new(data).unwrap();
522        assert!(chunk.address().as_ref() != ChunkAddress::default().as_ref()); // Ensure we get a non-zero hash
523    }
524
525    #[test]
526    fn test_exact_span_size() {
527        // Create a valid 8-byte span with no data
528        let mut data = vec![0u8; 8];
529        data.copy_from_slice(&0u64.to_le_bytes());
530
531        let chunk = DefaultContentChunk::try_from(data.as_slice()).unwrap();
532
533        assert_eq!(chunk.span(), 0);
534        assert_eq!(chunk.data(), &[0u8; 0].as_slice());
535        assert_eq!(chunk.size(), 8);
536    }
537}