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 pub const fn body(&self) -> &BmtBody<BODY_SIZE> {
146 &self.body
147 }
148
149 #[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#[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 pub const fn chunk(&self) -> &ContentChunk<BODY_SIZE> {
175 &self.chunk
176 }
177
178 pub const fn encrypted_ref(&self) -> &super::encryption::EncryptedChunkRef {
180 &self.encrypted_ref
181 }
182
183 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 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 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 fn encrypt_with(
234 &self,
235 key: &super::encryption::EncryptionKey,
236 ) -> Result<EncryptedContentChunk<BODY_SIZE>> {
237 let raw: Bytes = self.clone().into(); 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 }
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
325trait BuilderState {}
327
328#[derive(Debug, Default)]
329struct Initial;
330impl BuilderState for Initial {}
331
332#[derive(Debug)]
333struct ReadyToBuild;
334impl BuilderState for ReadyToBuild {}
335
336#[derive(Debug)]
338struct ContentChunkBuilderImpl<const BODY_SIZE: usize, S: BuilderState = Initial> {
339 body: Option<BmtBody<BODY_SIZE>>,
341 address: Option<ChunkAddress>,
343 _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 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 const fn with_address(mut self, address: ChunkAddress) -> Self {
379 self.address = Some(address);
380 self
381 }
382
383 fn build(self) -> ContentChunk<BODY_SIZE> {
385 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 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 prop_assert!(chunk.data().len() <= DEFAULT_BODY_SIZE);
428 prop_assert_eq!(chunk.size(), 8 + chunk.data().len());
429
430 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 let body_data = chunk.data().clone();
442 let body_span = chunk.span();
443
444 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 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 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 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()); }
524
525 #[test]
526 fn test_exact_span_size() {
527 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}