1use std::sync::Arc;
9
10use crate::codec::encode_and_hash;
11use crate::directory::{directory_fanout_meta, DirectoryFanoutSpan};
12use crate::hash::sha256;
13use crate::store::Store;
14use crate::types::{Cid, DirEntry, Hash, Link, LinkType, TreeNode};
15
16use crate::crypto::{encrypt_chk, EncryptionKey};
17
18pub const DEFAULT_CHUNK_SIZE: usize = 2 * 1024 * 1024;
20
21pub const BEP52_CHUNK_SIZE: usize = 16 * 1024;
23
24pub const DEFAULT_MAX_LINKS: usize = 174;
26
27#[derive(Clone)]
29pub struct BuilderConfig<S: Store> {
30 pub store: Arc<S>,
31 pub chunk_size: usize,
32 pub max_links: usize,
33 pub encrypted: bool,
35}
36
37impl<S: Store> BuilderConfig<S> {
38 pub fn new(store: Arc<S>) -> Self {
39 Self {
40 store,
41 chunk_size: DEFAULT_CHUNK_SIZE,
42 max_links: DEFAULT_MAX_LINKS,
43 encrypted: true,
44 }
45 }
46
47 pub fn with_chunk_size(mut self, chunk_size: usize) -> Self {
48 self.chunk_size = chunk_size;
49 self
50 }
51
52 pub fn with_max_links(mut self, max_links: usize) -> Self {
53 self.max_links = max_links;
54 self
55 }
56
57 pub fn public(mut self) -> Self {
59 self.encrypted = false;
60 self
61 }
62
63 pub fn encrypted(mut self) -> Self {
65 self.encrypted = true;
66 self
67 }
68}
69
70pub struct TreeBuilder<S: Store> {
72 store: Arc<S>,
73 chunk_size: usize,
74 max_links: usize,
75 encrypted: bool,
76}
77
78impl<S: Store> TreeBuilder<S> {
79 pub fn new(config: BuilderConfig<S>) -> Self {
80 Self {
81 store: config.store,
82 chunk_size: config.chunk_size,
83 max_links: config.max_links,
84 encrypted: config.encrypted,
85 }
86 }
87
88 pub fn is_encrypted(&self) -> bool {
90 self.encrypted
91 }
92
93 pub async fn put_blob(&self, data: &[u8]) -> Result<Hash, BuilderError> {
96 let hash = sha256(data);
97 self.store
98 .put(hash, data.to_vec())
99 .await
100 .map_err(|e| BuilderError::Store(e.to_string()))?;
101 Ok(hash)
102 }
103
104 async fn put_chunk_internal(
107 &self,
108 data: &[u8],
109 ) -> Result<(Hash, Option<EncryptionKey>), BuilderError> {
110 if self.encrypted {
111 let (encrypted, key) =
112 encrypt_chk(data).map_err(|e| BuilderError::Encryption(e.to_string()))?;
113 let hash = sha256(&encrypted);
114 self.store
115 .put(hash, encrypted)
116 .await
117 .map_err(|e| BuilderError::Store(e.to_string()))?;
118 Ok((hash, Some(key)))
119 } else {
120 let hash = self.put_blob(data).await?;
121 Ok((hash, None))
122 }
123 }
124
125 pub async fn put(&self, data: &[u8]) -> Result<(Cid, u64), BuilderError> {
131 let size = data.len() as u64;
132
133 if data.len() <= self.chunk_size {
135 let (hash, key) = self.put_chunk_internal(data).await?;
136 return Ok((Cid { hash, key }, size));
137 }
138
139 let mut links: Vec<Link> = Vec::new();
141 let mut offset = 0;
142
143 while offset < data.len() {
144 let end = (offset + self.chunk_size).min(data.len());
145 let chunk = &data[offset..end];
146 let chunk_size = chunk.len() as u64;
147 let (hash, key) = self.put_chunk_internal(chunk).await?;
148 links.push(Link {
149 hash,
150 name: None,
151 size: chunk_size,
152 key,
153 link_type: LinkType::Blob, meta: None,
155 });
156 offset = end;
157 }
158
159 let (root_hash, root_key, _) = self.build_tree_internal(links, Some(size)).await?;
161
162 Ok((
163 Cid {
164 hash: root_hash,
165 key: root_key,
166 },
167 size,
168 ))
169 }
170
171 async fn build_tree_internal(
174 &self,
175 links: Vec<Link>,
176 total_size: Option<u64>,
177 ) -> Result<(Hash, Option<[u8; 32]>, LinkType), BuilderError> {
178 if links.len() == 1 {
180 if let Some(ts) = total_size {
181 if links[0].size == ts {
182 return Ok((links[0].hash, links[0].key, links[0].link_type));
183 }
184 }
185 }
186
187 if links.len() <= self.max_links {
188 let node = TreeNode {
189 node_type: LinkType::File,
190 links,
191 };
192 let (data, _) = encode_and_hash(&node)?;
193
194 if self.encrypted {
195 let (encrypted, key) =
196 encrypt_chk(&data).map_err(|e| BuilderError::Encryption(e.to_string()))?;
197 let hash = sha256(&encrypted);
198 self.store
199 .put(hash, encrypted)
200 .await
201 .map_err(|e| BuilderError::Store(e.to_string()))?;
202 return Ok((hash, Some(key), LinkType::File));
203 }
204
205 let hash = sha256(&data);
207 self.store
208 .put(hash, data)
209 .await
210 .map_err(|e| BuilderError::Store(e.to_string()))?;
211 return Ok((hash, None, LinkType::File));
212 }
213
214 let mut sub_links = Vec::new();
216 for batch in links.chunks(self.max_links) {
217 let batch_size: u64 = batch.iter().map(|l| l.size).sum();
218 let (hash, key, link_type) =
219 Box::pin(self.build_tree_internal(batch.to_vec(), Some(batch_size))).await?;
220 sub_links.push(Link {
221 hash,
222 name: None,
223 size: batch_size,
224 key,
225 link_type,
226 meta: None,
227 });
228 }
229
230 Box::pin(self.build_tree_internal(sub_links, total_size)).await
231 }
232
233 pub async fn put_directory(&self, entries: Vec<DirEntry>) -> Result<Hash, BuilderError> {
236 let mut sorted = entries;
238 sorted.sort_by(|a, b| a.name.cmp(&b.name));
239
240 let links: Vec<Link> = sorted
241 .into_iter()
242 .map(|e| Link {
243 hash: e.hash,
244 name: Some(e.name),
245 size: e.size,
246 key: e.key,
247 link_type: e.link_type,
248 meta: e.meta,
249 })
250 .collect();
251
252 if links.len() <= self.max_links {
254 let node = TreeNode {
255 node_type: LinkType::Dir,
256 links,
257 };
258 let (data, hash) = encode_and_hash(&node)?;
259 self.store
260 .put(hash, data)
261 .await
262 .map_err(|e| BuilderError::Store(e.to_string()))?;
263 return Ok(hash);
264 }
265
266 self.build_directory_by_chunks(links).await
267 }
268
269 async fn build_directory_by_chunks(&self, links: Vec<Link>) -> Result<Hash, BuilderError> {
271 let spans = links
272 .into_iter()
273 .map(|link| {
274 let name = link
275 .name
276 .clone()
277 .unwrap_or_else(|| crate::types::to_hex(&link.hash));
278 DirectoryFanoutSpan {
279 link,
280 count: 1,
281 first: name.clone(),
282 last: name,
283 }
284 })
285 .collect();
286 self.build_directory_fanout_level(spans, LinkType::Dir)
287 .await
288 }
289
290 async fn build_directory_fanout_level(
291 &self,
292 spans: Vec<DirectoryFanoutSpan>,
293 child_node_type: LinkType,
294 ) -> Result<Hash, BuilderError> {
295 let mut sub_trees: Vec<DirectoryFanoutSpan> = Vec::new();
296
297 for batch in spans.chunks(self.max_links) {
298 let Some(first_span) = batch.first() else {
299 continue;
300 };
301 let last_span = batch.last().expect("non-empty fanout batch");
302 let count: usize = batch.iter().map(|span| span.count).sum();
303 let batch_size: u64 = batch.iter().map(|span| span.link.size).sum();
304
305 let node = TreeNode {
306 node_type: child_node_type,
307 links: batch.iter().map(|span| span.link.clone()).collect(),
308 };
309 let (data, hash) = encode_and_hash(&node)?;
310 self.store
311 .put(hash, data)
312 .await
313 .map_err(|e| BuilderError::Store(e.to_string()))?;
314
315 sub_trees.push(DirectoryFanoutSpan {
316 link: Link {
317 hash,
318 name: None,
319 size: batch_size,
320 key: None,
321 link_type: child_node_type,
322 meta: Some(directory_fanout_meta(
323 count,
324 &first_span.first,
325 &last_span.last,
326 )),
327 },
328 count,
329 first: first_span.first.clone(),
330 last: last_span.last.clone(),
331 });
332 }
333
334 if sub_trees.len() <= self.max_links {
335 let node = TreeNode {
336 node_type: LinkType::Fanout,
337 links: sub_trees.into_iter().map(|span| span.link).collect(),
338 };
339 let (data, hash) = encode_and_hash(&node)?;
340 self.store
341 .put(hash, data)
342 .await
343 .map_err(|e| BuilderError::Store(e.to_string()))?;
344 return Ok(hash);
345 }
346
347 Box::pin(self.build_directory_fanout_level(sub_trees, LinkType::Fanout)).await
349 }
350
351 pub async fn put_tree_node(&self, links: Vec<Link>) -> Result<Hash, BuilderError> {
353 let node = TreeNode {
354 node_type: LinkType::Dir,
355 links,
356 };
357
358 let (data, hash) = encode_and_hash(&node)?;
359 self.store
360 .put(hash, data)
361 .await
362 .map_err(|e| BuilderError::Store(e.to_string()))?;
363 Ok(hash)
364 }
365}
366
367pub struct StreamBuilder<S: Store> {
369 store: Arc<S>,
370 chunk_size: usize,
371 max_links: usize,
372
373 buffer: Vec<u8>,
375
376 chunks: Vec<Link>,
378 total_size: u64,
379}
380
381impl<S: Store> StreamBuilder<S> {
382 pub fn new(config: BuilderConfig<S>) -> Self {
383 Self {
384 store: config.store,
385 chunk_size: config.chunk_size,
386 max_links: config.max_links,
387 buffer: Vec::with_capacity(config.chunk_size),
388 chunks: Vec::new(),
389 total_size: 0,
390 }
391 }
392
393 pub async fn append(&mut self, data: &[u8]) -> Result<(), BuilderError> {
395 let mut offset = 0;
396
397 while offset < data.len() {
398 let space = self.chunk_size - self.buffer.len();
399 let to_write = space.min(data.len() - offset);
400
401 self.buffer
402 .extend_from_slice(&data[offset..offset + to_write]);
403 offset += to_write;
404
405 if self.buffer.len() == self.chunk_size {
407 self.flush_chunk().await?;
408 }
409 }
410
411 self.total_size += data.len() as u64;
412 Ok(())
413 }
414
415 async fn flush_chunk(&mut self) -> Result<(), BuilderError> {
417 if self.buffer.is_empty() {
418 return Ok(());
419 }
420
421 let chunk = std::mem::take(&mut self.buffer);
422 let hash = sha256(&chunk);
423 self.store
424 .put(hash, chunk.clone())
425 .await
426 .map_err(|e| BuilderError::Store(e.to_string()))?;
427
428 self.chunks.push(Link {
429 hash,
430 name: None,
431 size: chunk.len() as u64,
432 key: None,
433 link_type: LinkType::Blob, meta: None,
435 });
436
437 self.buffer = Vec::with_capacity(self.chunk_size);
438 Ok(())
439 }
440
441 pub async fn current_root(&mut self) -> Result<Option<Hash>, BuilderError> {
444 if self.chunks.is_empty() && self.buffer.is_empty() {
445 return Ok(None);
446 }
447
448 let mut temp_chunks = self.chunks.clone();
450 if !self.buffer.is_empty() {
451 let chunk = self.buffer.clone();
452 let hash = sha256(&chunk);
453 self.store
454 .put(hash, chunk.clone())
455 .await
456 .map_err(|e| BuilderError::Store(e.to_string()))?;
457 temp_chunks.push(Link {
458 hash,
459 name: None,
460 size: chunk.len() as u64,
461 key: None,
462 link_type: LinkType::Blob, meta: None,
464 });
465 }
466
467 let hash = self
468 .build_tree_from_chunks(&temp_chunks, self.total_size)
469 .await?;
470 Ok(Some(hash))
471 }
472
473 pub async fn finalize(mut self) -> Result<(Hash, u64), BuilderError> {
475 self.flush_chunk().await?;
477
478 if self.chunks.is_empty() {
479 let empty_hash = sha256(&[]);
481 self.store
482 .put(empty_hash, vec![])
483 .await
484 .map_err(|e| BuilderError::Store(e.to_string()))?;
485 return Ok((empty_hash, 0));
486 }
487
488 let hash = self
489 .build_tree_from_chunks(&self.chunks, self.total_size)
490 .await?;
491 Ok((hash, self.total_size))
492 }
493
494 async fn build_tree_from_chunks(
496 &self,
497 chunks: &[Link],
498 total_size: u64,
499 ) -> Result<Hash, BuilderError> {
500 if chunks.len() == 1 {
501 return Ok(chunks[0].hash);
502 }
503
504 if chunks.len() <= self.max_links {
505 let node = TreeNode {
506 node_type: LinkType::File,
507 links: chunks.to_vec(),
508 };
509 let (data, hash) = encode_and_hash(&node)?;
510 self.store
511 .put(hash, data)
512 .await
513 .map_err(|e| BuilderError::Store(e.to_string()))?;
514 return Ok(hash);
515 }
516
517 let mut sub_trees: Vec<Link> = Vec::new();
519 for batch in chunks.chunks(self.max_links) {
520 let batch_size: u64 = batch.iter().map(|l| l.size).sum();
521
522 let node = TreeNode {
523 node_type: LinkType::File,
524 links: batch.to_vec(),
525 };
526 let (data, hash) = encode_and_hash(&node)?;
527 self.store
528 .put(hash, data)
529 .await
530 .map_err(|e| BuilderError::Store(e.to_string()))?;
531
532 sub_trees.push(Link {
533 hash,
534 name: None,
535 size: batch_size,
536 key: None,
537 link_type: LinkType::File, meta: None,
539 });
540 }
541
542 Box::pin(self.build_tree_from_chunks(&sub_trees, total_size)).await
543 }
544
545 pub fn stats(&self) -> StreamStats {
547 StreamStats {
548 chunks: self.chunks.len(),
549 buffered: self.buffer.len(),
550 total_size: self.total_size,
551 }
552 }
553}
554
555#[derive(Debug, Clone, PartialEq)]
556pub struct StreamStats {
557 pub chunks: usize,
558 pub buffered: usize,
559 pub total_size: u64,
560}
561
562#[derive(Debug, thiserror::Error)]
564pub enum BuilderError {
565 #[error("Store error: {0}")]
566 Store(String),
567 #[error("Codec error: {0}")]
568 Codec(#[from] crate::codec::CodecError),
569 #[error("Encryption error: {0}")]
570 Encryption(String),
571}
572
573#[cfg(test)]
574mod tests {
575 use super::*;
576 use crate::store::MemoryStore;
577 use crate::types::to_hex;
578 use std::collections::HashMap;
579
580 fn make_store() -> Arc<MemoryStore> {
581 Arc::new(MemoryStore::new())
582 }
583
584 #[tokio::test]
585 async fn test_put_blob() {
586 let store = make_store();
587 let builder = TreeBuilder::new(BuilderConfig::new(store.clone()));
588
589 let data = vec![1u8, 2, 3, 4, 5];
590 let hash = builder.put_blob(&data).await.unwrap();
591
592 assert_eq!(hash.len(), 32);
593 assert!(store.has(&hash).await.unwrap());
594
595 let retrieved = store.get(&hash).await.unwrap();
596 assert_eq!(retrieved, Some(data));
597 }
598
599 #[tokio::test]
600 async fn test_put_blob_correct_hash() {
601 let store = make_store();
602 let builder = TreeBuilder::new(BuilderConfig::new(store));
603
604 let data = vec![1u8, 2, 3];
605 let hash = builder.put_blob(&data).await.unwrap();
606 let expected_hash = sha256(&data);
607
608 assert_eq!(to_hex(&hash), to_hex(&expected_hash));
609 }
610
611 #[tokio::test]
612 async fn test_put_small() {
613 let store = make_store();
614 let builder = TreeBuilder::new(BuilderConfig::new(store.clone()).public());
616
617 let data = vec![1u8, 2, 3, 4, 5];
618 let (cid, size) = builder.put(&data).await.unwrap();
619
620 assert_eq!(size, 5);
621 assert!(cid.key.is_none()); let retrieved = store.get(&cid.hash).await.unwrap();
623 assert_eq!(retrieved, Some(data));
624 }
625
626 #[tokio::test]
627 async fn test_put_chunked() {
628 let store = make_store();
629 let config = BuilderConfig::new(store.clone())
630 .with_chunk_size(1024)
631 .public();
632 let builder = TreeBuilder::new(config);
633
634 let mut data = vec![0u8; 1024 * 2 + 100];
635 for i in 0..data.len() {
636 data[i] = (i % 256) as u8;
637 }
638
639 let (_cid, size) = builder.put(&data).await.unwrap();
640 assert_eq!(size, data.len() as u64);
641
642 assert!(store.size() > 1);
644 }
645
646 #[tokio::test]
647 async fn test_put_directory() {
648 let store = make_store();
649 let builder = TreeBuilder::new(BuilderConfig::new(store.clone()));
650
651 let file1 = vec![1u8, 2, 3];
652 let file2 = vec![4u8, 5, 6, 7];
653
654 let hash1 = builder.put_blob(&file1).await.unwrap();
655 let hash2 = builder.put_blob(&file2).await.unwrap();
656
657 let dir_hash = builder
658 .put_directory(vec![
659 DirEntry::new("a.txt", hash1).with_size(file1.len() as u64),
660 DirEntry::new("b.txt", hash2).with_size(file2.len() as u64),
661 ])
662 .await
663 .unwrap();
664
665 assert!(store.has(&dir_hash).await.unwrap());
666 }
667
668 #[tokio::test]
669 async fn test_put_directory_sorted() {
670 let store = make_store();
671 let builder = TreeBuilder::new(BuilderConfig::new(store.clone()));
672
673 let hash = builder.put_blob(&[1u8]).await.unwrap();
674
675 let dir_hash = builder
676 .put_directory(vec![
677 DirEntry::new("zebra", hash),
678 DirEntry::new("apple", hash),
679 DirEntry::new("mango", hash),
680 ])
681 .await
682 .unwrap();
683
684 let data = store.get(&dir_hash).await.unwrap().unwrap();
685 let node = crate::codec::decode_tree_node(&data).unwrap();
686
687 let names: Vec<_> = node.links.iter().filter_map(|l| l.name.clone()).collect();
688 assert_eq!(names, vec!["apple", "mango", "zebra"]);
689 }
690
691 #[tokio::test]
692 async fn test_put_tree_node_with_link_meta() {
693 let store = make_store();
694 let builder = TreeBuilder::new(BuilderConfig::new(store.clone()));
695
696 let hash = builder.put_blob(&[1u8]).await.unwrap();
697
698 let mut meta = HashMap::new();
699 meta.insert("version".to_string(), serde_json::json!(2));
700 meta.insert("created".to_string(), serde_json::json!("2024-01-01"));
701
702 let node_hash = builder
703 .put_tree_node(vec![Link {
704 hash,
705 name: Some("test".to_string()),
706 size: 1,
707 key: None,
708 link_type: LinkType::Blob,
709 meta: Some(meta.clone()),
710 }])
711 .await
712 .unwrap();
713
714 let data = store.get(&node_hash).await.unwrap().unwrap();
715 let node = crate::codec::decode_tree_node(&data).unwrap();
716
717 assert!(node.links[0].meta.is_some());
718 let m = node.links[0].meta.as_ref().unwrap();
719 assert_eq!(m.get("version"), Some(&serde_json::json!(2)));
720 }
721
722 #[tokio::test]
723 async fn test_stream_builder() {
724 let store = make_store();
725 let config = BuilderConfig::new(store.clone()).with_chunk_size(100);
726 let mut stream = StreamBuilder::new(config);
727
728 stream.append(&[1u8, 2, 3]).await.unwrap();
729 stream.append(&[4u8, 5]).await.unwrap();
730 stream.append(&[6u8, 7, 8, 9]).await.unwrap();
731
732 let (hash, size) = stream.finalize().await.unwrap();
733
734 assert_eq!(size, 9);
735 assert!(store.has(&hash).await.unwrap());
736 }
737
738 #[tokio::test]
739 async fn test_stream_stats() {
740 let store = make_store();
741 let config = BuilderConfig::new(store).with_chunk_size(100);
742 let mut stream = StreamBuilder::new(config);
743
744 assert_eq!(stream.stats().chunks, 0);
745 assert_eq!(stream.stats().buffered, 0);
746 assert_eq!(stream.stats().total_size, 0);
747
748 stream.append(&[0u8; 50]).await.unwrap();
749 assert_eq!(stream.stats().buffered, 50);
750 assert_eq!(stream.stats().total_size, 50);
751
752 stream.append(&[0u8; 60]).await.unwrap(); assert_eq!(stream.stats().chunks, 1);
754 assert_eq!(stream.stats().buffered, 10);
755 assert_eq!(stream.stats().total_size, 110);
756 }
757
758 #[tokio::test]
759 async fn test_stream_current_root() {
760 let store = make_store();
761 let config = BuilderConfig::new(store).with_chunk_size(100);
762 let mut stream = StreamBuilder::new(config);
763
764 stream.append(&[1u8, 2, 3]).await.unwrap();
765 let root1 = stream.current_root().await.unwrap();
766
767 stream.append(&[4u8, 5, 6]).await.unwrap();
768 let root2 = stream.current_root().await.unwrap();
769
770 assert_ne!(to_hex(&root1.unwrap()), to_hex(&root2.unwrap()));
772 }
773
774 #[tokio::test]
775 async fn test_stream_empty() {
776 let store = make_store();
777 let config = BuilderConfig::new(store.clone());
778 let stream = StreamBuilder::new(config);
779
780 let (hash, size) = stream.finalize().await.unwrap();
781 assert_eq!(size, 0);
782 assert!(store.has(&hash).await.unwrap());
783 }
784
785 #[tokio::test]
786 async fn test_unified_put_public() {
787 let store = make_store();
788 let config = BuilderConfig::new(store.clone()).public();
790 let builder = TreeBuilder::new(config);
791
792 let data = b"Hello, World!";
793 let (cid, size) = builder.put(data).await.unwrap();
794
795 assert_eq!(size, data.len() as u64);
796 assert!(cid.key.is_none()); assert!(store.has(&cid.hash).await.unwrap());
798 }
799
800 #[tokio::test]
801 async fn test_unified_put_encrypted() {
802 use crate::reader::TreeReader;
803
804 let store = make_store();
805 let config = BuilderConfig::new(store.clone());
807 let builder = TreeBuilder::new(config);
808
809 let data = b"Hello, encrypted world!";
810 let (cid, size) = builder.put(data).await.unwrap();
811
812 assert_eq!(size, data.len() as u64);
813 assert!(cid.key.is_some()); let reader = TreeReader::new(store);
817 let retrieved = reader.get(&cid).await.unwrap().unwrap();
818 assert_eq!(retrieved, data);
819 }
820
821 #[tokio::test]
822 async fn test_unified_put_encrypted_chunked() {
823 use crate::reader::TreeReader;
824
825 let store = make_store();
826 let config = BuilderConfig::new(store.clone()).with_chunk_size(100);
827 let builder = TreeBuilder::new(config);
828
829 let data: Vec<u8> = (0..500).map(|i| (i % 256) as u8).collect();
831 let (cid, size) = builder.put(&data).await.unwrap();
832
833 assert_eq!(size, data.len() as u64);
834 assert!(cid.key.is_some());
835
836 let reader = TreeReader::new(store);
838 let retrieved = reader.get(&cid).await.unwrap().unwrap();
839 assert_eq!(retrieved, data);
840 }
841
842 #[tokio::test]
843 async fn test_cid_deterministic() {
844 let store = make_store();
845 let config = BuilderConfig::new(store.clone());
846 let builder = TreeBuilder::new(config);
847
848 let data = b"Same content produces same CID";
849
850 let (cid1, _) = builder.put(data).await.unwrap();
851 let (cid2, _) = builder.put(data).await.unwrap();
852
853 assert_eq!(cid1.hash, cid2.hash);
855 assert_eq!(cid1.key, cid2.key);
856 assert_eq!(cid1.to_string(), cid2.to_string());
857 }
858}