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    /// Create a ContentChunk from a pre-existing BmtBody with a known address.
140    ///
141    /// This is an advanced method for when you already have both the body
142    /// and know the chunk's address (e.g., when reconstructing from storage).
143    #[must_use]
144    pub fn from_body_with_address(body: BmtBody<BODY_SIZE>, address: ChunkAddress) -> Self {
145        Self {
146            header: ContentChunkHeader::new(),
147            body,
148            address_cache: OnceCache::with_value(address),
149        }
150    }
151}
152
153/// Result of encrypting a content chunk.
154#[cfg(feature = "encryption")]
155#[derive(Debug, Clone)]
156pub struct EncryptedContentChunk<const BODY_SIZE: usize = DEFAULT_BODY_SIZE> {
157    chunk: ContentChunk<BODY_SIZE>,
158    encrypted_ref: super::encryption::EncryptedChunkRef,
159}
160
161#[cfg(feature = "encryption")]
162impl<const BODY_SIZE: usize> EncryptedContentChunk<BODY_SIZE> {
163    /// The encrypted chunk (ciphertext hashed to a new address).
164    pub const fn chunk(&self) -> &ContentChunk<BODY_SIZE> {
165        &self.chunk
166    }
167
168    /// The encrypted reference (address + decryption key).
169    pub const fn encrypted_ref(&self) -> &super::encryption::EncryptedChunkRef {
170        &self.encrypted_ref
171    }
172
173    /// Consume and return (chunk, encrypted_ref).
174    pub fn into_parts(
175        self,
176    ) -> (
177        ContentChunk<BODY_SIZE>,
178        super::encryption::EncryptedChunkRef,
179    ) {
180        (self.chunk, self.encrypted_ref)
181    }
182
183    /// Decrypt back to a plaintext `ContentChunk`.
184    pub fn decrypt(&self) -> Result<ContentChunk<BODY_SIZE>> {
185        use super::encryption::transcrypt;
186        use crate::bmt::SPAN_SIZE;
187
188        let encrypted_data: Bytes = self.chunk.clone().into();
189        let key = self.encrypted_ref.key();
190
191        // Decrypt the span to learn the original data length
192        let span_ctr = (BODY_SIZE / super::encryption::EncryptionKey::SIZE) as u32;
193        let mut span_buf = [0u8; SPAN_SIZE];
194        transcrypt(key, span_ctr, &encrypted_data[..SPAN_SIZE], &mut span_buf)?;
195        let data_length = u64::from_le_bytes(span_buf) as usize;
196
197        let decrypted =
198            super::encryption::decrypt_chunk_data::<BODY_SIZE>(&encrypted_data, key, data_length)?;
199        ContentChunk::try_from(Bytes::from(decrypted))
200    }
201}
202
203#[cfg(feature = "encryption")]
204impl<const BODY_SIZE: usize> super::encryption::ChunkEncrypt for ContentChunk<BODY_SIZE> {
205    type Encrypted = EncryptedContentChunk<BODY_SIZE>;
206
207    /// Encrypt this chunk with a caller-provided key.
208    ///
209    /// The returned `EncryptedContentChunk` contains:
210    /// - `chunk`: a new `ContentChunk` whose data is the ciphertext
211    /// - `encrypted_ref`: the 64-byte reference (new address + decryption key)
212    ///
213    /// ```
214    /// # use nectar_primitives::{Chunk, ContentChunk};
215    /// # use nectar_primitives::chunk::encryption::{ChunkEncrypt, EncryptionKey};
216    /// # use nectar_primitives::bmt::DEFAULT_BODY_SIZE;
217    /// let chunk = ContentChunk::<DEFAULT_BODY_SIZE>::new(b"secret data".to_vec()).unwrap();
218    /// let encrypted = chunk.encrypt().unwrap();
219    ///
220    /// // The encrypted chunk has a different address
221    /// assert_ne!(chunk.address(), encrypted.chunk().address());
222    /// ```
223    fn encrypt_with(
224        &self,
225        key: &super::encryption::EncryptionKey,
226    ) -> Result<EncryptedContentChunk<BODY_SIZE>> {
227        let raw: Bytes = self.clone().into(); // span || data
228        let ciphertext = super::encryption::encrypt_chunk::<BODY_SIZE>(&raw, key)?;
229        let encrypted_chunk = Self::try_from(Bytes::from(ciphertext))?;
230        let encrypted_ref =
231            super::encryption::EncryptedChunkRef::new(*encrypted_chunk.address(), key.clone());
232        Ok(EncryptedContentChunk {
233            chunk: encrypted_chunk,
234            encrypted_ref,
235        })
236    }
237    // encrypt() uses default impl — generates random key, calls encrypt_with()
238}
239
240impl<const BODY_SIZE: usize> Chunk for ContentChunk<BODY_SIZE> {
241    type Header = ContentChunkHeader;
242
243    fn address(&self) -> &ChunkAddress {
244        self.address_cache.get_or_compute(|| self.body.hash())
245    }
246
247    fn data(&self) -> &Bytes {
248        self.body.data()
249    }
250
251    fn size(&self) -> usize {
252        self.header().bytes().len() + self.body.size()
253    }
254
255    fn header(&self) -> &Self::Header {
256        &self.header
257    }
258}
259
260impl<const BODY_SIZE: usize> BmtChunk for ContentChunk<BODY_SIZE> {
261    fn span(&self) -> u64 {
262        self.body.span()
263    }
264}
265
266impl<const BODY_SIZE: usize> From<ContentChunk<BODY_SIZE>> for Bytes {
267    fn from(chunk: ContentChunk<BODY_SIZE>) -> Self {
268        chunk.body.into()
269    }
270}
271
272impl<const BODY_SIZE: usize> TryFrom<Bytes> for ContentChunk<BODY_SIZE> {
273    type Error = PrimitivesError;
274
275    fn try_from(bytes: Bytes) -> Result<Self> {
276        Ok(Self {
277            header: ContentChunkHeader::new(),
278            body: BmtBody::try_from(bytes)?,
279            address_cache: OnceCache::new(),
280        })
281    }
282}
283
284impl<const BODY_SIZE: usize> TryFrom<&[u8]> for ContentChunk<BODY_SIZE> {
285    type Error = PrimitivesError;
286
287    fn try_from(bytes: &[u8]) -> Result<Self> {
288        Self::try_from(Bytes::copy_from_slice(bytes))
289    }
290}
291
292impl<const BODY_SIZE: usize> fmt::Display for ContentChunk<BODY_SIZE> {
293    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
294        write!(
295            f,
296            "ContentChunk[{}]",
297            hex::encode(&self.address().as_bytes()[..8])
298        )
299    }
300}
301
302impl<const BODY_SIZE: usize> PartialEq for ContentChunk<BODY_SIZE> {
303    fn eq(&self, other: &Self) -> bool {
304        self.address() == other.address()
305    }
306}
307
308impl<const BODY_SIZE: usize> Eq for ContentChunk<BODY_SIZE> {}
309
310impl<const BODY_SIZE: usize> super::chunk_type::ChunkType for ContentChunk<BODY_SIZE> {
311    const TYPE_ID: super::type_id::ChunkTypeId = super::type_id::ChunkTypeId::CONTENT;
312    const TYPE_NAME: &'static str = "content";
313}
314
315// Internal builder implementation
316trait BuilderState {}
317
318#[derive(Debug, Default)]
319struct Initial;
320impl BuilderState for Initial {}
321
322#[derive(Debug)]
323struct ReadyToBuild;
324impl BuilderState for ReadyToBuild {}
325
326/// Builder for ContentChunk with type state pattern
327#[derive(Debug)]
328struct ContentChunkBuilderImpl<const BODY_SIZE: usize, S: BuilderState = Initial> {
329    /// The body to use for the chunk
330    body: Option<BmtBody<BODY_SIZE>>,
331    /// Pre-computed address for the chunk
332    address: Option<ChunkAddress>,
333    /// Marker for the builder state
334    _state: PhantomData<S>,
335}
336
337impl<const BODY_SIZE: usize> Default for ContentChunkBuilderImpl<BODY_SIZE, Initial> {
338    fn default() -> Self {
339        Self {
340            body: None,
341            address: None,
342            _state: PhantomData,
343        }
344    }
345}
346
347impl<const BODY_SIZE: usize> ContentChunkBuilderImpl<BODY_SIZE, Initial> {
348    /// Initialize from data with automatically calculated span
349    fn auto_from_data(
350        mut self,
351        data: impl Into<Bytes>,
352    ) -> Result<ContentChunkBuilderImpl<BODY_SIZE, ReadyToBuild>> {
353        let body = BmtBody::<BODY_SIZE>::builder()
354            .auto_from_data(data)?
355            .build()?;
356        self.body = Some(body);
357
358        Ok(ContentChunkBuilderImpl {
359            body: self.body,
360            address: self.address,
361            _state: PhantomData,
362        })
363    }
364}
365
366impl<const BODY_SIZE: usize> ContentChunkBuilderImpl<BODY_SIZE, ReadyToBuild> {
367    /// Set a pre-computed address for the chunk
368    const fn with_address(mut self, address: ChunkAddress) -> Self {
369        self.address = Some(address);
370        self
371    }
372
373    /// Build the final ContentChunk
374    fn build(self) -> ContentChunk<BODY_SIZE> {
375        // This is safe as we have already checked that the body is set
376        let body = self.body.unwrap();
377
378        let address_cache = self
379            .address
380            .map_or_else(OnceCache::new, OnceCache::with_value);
381
382        ContentChunk {
383            header: ContentChunkHeader::new(),
384            body,
385            address_cache,
386        }
387    }
388}
389
390#[cfg(any(test, feature = "arbitrary"))]
391impl<'a, const BODY_SIZE: usize> arbitrary::Arbitrary<'a> for ContentChunk<BODY_SIZE> {
392    fn arbitrary(u: &mut arbitrary::Unstructured<'a>) -> arbitrary::Result<Self> {
393        Ok(Self::from_body(BmtBody::<BODY_SIZE>::arbitrary(u)?))
394    }
395}
396
397#[cfg(test)]
398mod tests {
399    use crate::{DEFAULT_BODY_SIZE, chunk::error::ChunkError};
400
401    use super::*;
402    use alloy_primitives::b256;
403    use proptest::prelude::*;
404    use proptest_arbitrary_interop::arb;
405
406    type DefaultContentChunk = ContentChunk<DEFAULT_BODY_SIZE>;
407
408    // Strategy for generating ContentChunk using the Arbitrary implementation
409    fn chunk_strategy() -> impl Strategy<Value = DefaultContentChunk> {
410        arb::<DefaultContentChunk>()
411    }
412
413    proptest! {
414        #[test]
415        fn test_chunk_properties(chunk in chunk_strategy()) {
416            // Test basic properties
417            prop_assert!(chunk.data().len() <= DEFAULT_BODY_SIZE);
418            prop_assert_eq!(chunk.size(), 8 + chunk.data().len());
419
420            // Test round-trip conversion
421            let bytes: Bytes = chunk.clone().into();
422            let decoded = DefaultContentChunk::try_from(bytes).unwrap();
423            prop_assert_eq!(chunk.address(), decoded.address());
424            prop_assert_eq!(chunk.data(), decoded.data());
425            prop_assert_eq!(chunk.span(), decoded.span());
426        }
427
428        #[test]
429        fn test_from_body(chunk in chunk_strategy()) {
430            // Test creating a chunk from an existing body via BmtBody
431            let body_data = chunk.data().clone();
432            let body_span = chunk.span();
433
434            // Create a new chunk using from_body with a fresh BmtBody
435            let new_body = BmtBody::<DEFAULT_BODY_SIZE>::try_from(Bytes::from(chunk.clone())).unwrap();
436            let new_chunk = DefaultContentChunk::from_body(new_body);
437
438            prop_assert_eq!(new_chunk.data(), &body_data);
439            prop_assert_eq!(new_chunk.span(), body_span);
440            prop_assert_eq!(new_chunk.address(), chunk.address());
441        }
442
443        #[test]
444        fn test_new_content_chunk(data in proptest::collection::vec(any::<u8>(), 0..DEFAULT_BODY_SIZE)) {
445            let chunk = DefaultContentChunk::new(data.clone()).unwrap();
446
447            prop_assert_eq!(chunk.data(), &data);
448            prop_assert_eq!(chunk.span(), data.len() as u64);
449            prop_assert!(!chunk.address().is_zero());
450        }
451
452        #[test]
453        fn test_chunk_size_validation(data in proptest::collection::vec(any::<u8>(), DEFAULT_BODY_SIZE + 1..DEFAULT_BODY_SIZE * 2)) {
454            let result = DefaultContentChunk::new(data);
455            prop_assert_eq!(matches!(result, Err(PrimitivesError::Chunk(ChunkError::InvalidSize { .. }))), true);
456        }
457
458        #[test]
459        fn test_empty_and_edge_cases(size in 0usize..=10usize) {
460            // Test with empty or small data
461            let data = vec![0u8; size];
462            let chunk = DefaultContentChunk::new(data).unwrap();
463
464            prop_assert_eq!(chunk.data().len(), size);
465            prop_assert_eq!(chunk.span(), size as u64);
466            prop_assert_eq!(chunk.size(), 8 + size);
467        }
468
469        #[test]
470        fn test_deserialize_invalid_chunks(data in proptest::collection::vec(any::<u8>(), 0..8)) {
471            let result = DefaultContentChunk::try_from(data.as_slice());
472            prop_assert_eq!(matches!(result, Err(PrimitivesError::Chunk(ChunkError::InvalidSize { .. }))), true);
473        }
474    }
475
476    #[test]
477    fn test_new() {
478        let data = b"greaterthanspan";
479        let bmt_hash = b256!("27913f1bdb6e8e52cbd5a5fd4ab577c857287edf6969b41efe926b51de0f4f23");
480
481        let chunk = DefaultContentChunk::new(data.to_vec()).unwrap();
482        assert_eq!(chunk.address().as_ref(), bmt_hash);
483        assert_eq!(chunk.data(), data.as_slice());
484    }
485
486    #[test]
487    fn test_from_bytes() {
488        let data = b"greaterthanspan";
489        let bmt_hash = b256!("95022e6af5c6d6a564ee55a67f8455a3e18c511b5697c932d9e44f07f2fb8c53");
490
491        let chunk = DefaultContentChunk::try_from(data.as_slice()).unwrap();
492        assert_eq!(chunk.address().as_ref(), bmt_hash);
493        assert_eq!(
494            <DefaultContentChunk as Into<Bytes>>::into(chunk),
495            data.as_slice()
496        );
497    }
498
499    #[test]
500    fn test_specific_content_hash() {
501        // Test with known valid data and hash
502        let data = b"foo".to_vec();
503        let expected_hash =
504            b256!("2387e8e7d8a48c2a9339c97c1dc3461a9a7aa07e994c5cb8b38fd7c1b3e6ea48");
505
506        let chunk = DefaultContentChunk::new(data).unwrap();
507        assert_eq!(chunk.address().as_ref(), expected_hash);
508
509        // Test with "Digital Freedom Now"
510        let data = b"Digital Freedom Now".to_vec();
511        let chunk = DefaultContentChunk::new(data).unwrap();
512        assert!(chunk.address().as_ref() != ChunkAddress::default().as_ref()); // Ensure we get a non-zero hash
513    }
514
515    #[test]
516    fn test_exact_span_size() {
517        // Create a valid 8-byte span with no data
518        let mut data = vec![0u8; 8];
519        data.copy_from_slice(&0u64.to_le_bytes());
520
521        let chunk = DefaultContentChunk::try_from(data.as_slice()).unwrap();
522
523        assert_eq!(chunk.span(), 0);
524        assert_eq!(chunk.data(), &[0u8; 0].as_slice());
525        assert_eq!(chunk.size(), 8);
526    }
527}