1use 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#[derive(Debug, Clone)]
23pub struct ContentChunk<const BODY_SIZE: usize = DEFAULT_BODY_SIZE> {
24 header: ContentChunkHeader,
26 body: BmtBody<BODY_SIZE>,
28 address_cache: OnceCache<ChunkAddress>,
30}
31
32#[derive(Debug, Clone)]
36pub struct ContentChunkMetadata;
37
38impl ChunkMetadata for ContentChunkMetadata {
39 fn bytes(&self) -> Bytes {
40 Bytes::new()
41 }
42}
43
44#[derive(Debug, Clone)]
46pub struct ContentChunkHeader {
47 metadata: ContentChunkMetadata,
48}
49
50impl ContentChunkHeader {
51 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 #[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 #[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 #[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 #[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#[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 pub const fn chunk(&self) -> &ContentChunk<BODY_SIZE> {
165 &self.chunk
166 }
167
168 pub const fn encrypted_ref(&self) -> &super::encryption::EncryptedChunkRef {
170 &self.encrypted_ref
171 }
172
173 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 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 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 fn encrypt_with(
224 &self,
225 key: &super::encryption::EncryptionKey,
226 ) -> Result<EncryptedContentChunk<BODY_SIZE>> {
227 let raw: Bytes = self.clone().into(); 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 }
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
315trait BuilderState {}
317
318#[derive(Debug, Default)]
319struct Initial;
320impl BuilderState for Initial {}
321
322#[derive(Debug)]
323struct ReadyToBuild;
324impl BuilderState for ReadyToBuild {}
325
326#[derive(Debug)]
328struct ContentChunkBuilderImpl<const BODY_SIZE: usize, S: BuilderState = Initial> {
329 body: Option<BmtBody<BODY_SIZE>>,
331 address: Option<ChunkAddress>,
333 _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 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 const fn with_address(mut self, address: ChunkAddress) -> Self {
369 self.address = Some(address);
370 self
371 }
372
373 fn build(self) -> ContentChunk<BODY_SIZE> {
375 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 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 prop_assert!(chunk.data().len() <= DEFAULT_BODY_SIZE);
418 prop_assert_eq!(chunk.size(), 8 + chunk.data().len());
419
420 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 let body_data = chunk.data().clone();
432 let body_span = chunk.span();
433
434 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 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 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 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()); }
514
515 #[test]
516 fn test_exact_span_size() {
517 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}