hashtree_core/
builder.rs

1//! Tree builder with chunking and fanout support
2//!
3//! - Large files are split into chunks
4//! - Large directories are split into sub-trees
5//! - Supports streaming appends
6//! - Encryption enabled by default (CHK - Content Hash Key)
7
8use 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
18/// Default chunk size: 2MB (optimized for blossom uploads, matches hashtree-ts)
19pub const DEFAULT_CHUNK_SIZE: usize = 2 * 1024 * 1024;
20
21/// BEP52 chunk size: 16KB
22pub const BEP52_CHUNK_SIZE: usize = 16 * 1024;
23
24/// Default max links per tree node (fanout)
25pub const DEFAULT_MAX_LINKS: usize = 174;
26
27/// Builder configuration
28#[derive(Clone)]
29pub struct BuilderConfig<S: Store> {
30    pub store: Arc<S>,
31    pub chunk_size: usize,
32    pub max_links: usize,
33    /// Whether to encrypt content (default: true when encryption feature enabled)
34    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    /// Disable encryption (store content publicly)
58    pub fn public(mut self) -> Self {
59        self.encrypted = false;
60        self
61    }
62
63    /// Enable encryption (CHK - Content Hash Key)
64    pub fn encrypted(mut self) -> Self {
65        self.encrypted = true;
66        self
67    }
68}
69
70/// TreeBuilder - builds content-addressed merkle trees
71pub 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    /// Check if encryption is enabled
89    pub fn is_encrypted(&self) -> bool {
90        self.encrypted
91    }
92
93    /// Store a blob directly (small data, no encryption)
94    /// Returns the content hash
95    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    /// Store a chunk with optional encryption
105    /// Returns (hash, optional_key) where hash is of stored data
106    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    /// Store a file, chunking if necessary
126    /// Returns (Cid, size) where Cid contains hash and optional encryption key
127    ///
128    /// When encryption is enabled (default), each chunk is CHK encrypted
129    /// and the result contains the decryption key.
130    pub async fn put(&self, data: &[u8]) -> Result<(Cid, u64), BuilderError> {
131        let size = data.len() as u64;
132
133        // Small file - store as single chunk
134        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        // Large file - chunk it
140        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, // leaf chunk
154                meta: None,
155            });
156            offset = end;
157        }
158
159        // Build tree from chunks
160        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    /// Build tree and return (hash, optional_key)
172    /// When encrypted, tree nodes are also CHK encrypted
173    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        // Single link with matching size - return directly
179        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            // Unencrypted path
206            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        // Too many links - create subtrees
215        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    /// Build a directory from entries
234    /// Entries can be files or subdirectories
235    pub async fn put_directory(&self, entries: Vec<DirEntry>) -> Result<Hash, BuilderError> {
236        // Sort entries by name for deterministic hashing
237        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        // Fits in one node
253        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    /// Split directories into canonical BUD-17 fanout nodes.
270    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        // Recursively build more levels
348        Box::pin(self.build_directory_fanout_level(sub_trees, LinkType::Fanout)).await
349    }
350
351    /// Create a tree node
352    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
367/// StreamBuilder - supports incremental appends
368pub struct StreamBuilder<S: Store> {
369    store: Arc<S>,
370    chunk_size: usize,
371    max_links: usize,
372
373    // Current partial chunk being built
374    buffer: Vec<u8>,
375
376    // Completed chunks
377    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    /// Append data to the stream
394    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            // Flush full chunk
406            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    /// Flush current buffer as a chunk
416    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, // Leaf chunk (raw blob)
434            meta: None,
435        });
436
437        self.buffer = Vec::with_capacity(self.chunk_size);
438        Ok(())
439    }
440
441    /// Get current root hash without finalizing
442    /// Useful for checkpoints
443    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        // Temporarily include buffer
449        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, // Leaf chunk (raw blob)
463                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    /// Finalize the stream and return root hash
474    pub async fn finalize(mut self) -> Result<(Hash, u64), BuilderError> {
475        // Flush remaining buffer
476        self.flush_chunk().await?;
477
478        if self.chunks.is_empty() {
479            // Empty stream - return hash of empty data
480            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    /// Build balanced tree from chunks
495    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        // Build intermediate level
518        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, // Internal tree node
538                meta: None,
539            });
540        }
541
542        Box::pin(self.build_tree_from_chunks(&sub_trees, total_size)).await
543    }
544
545    /// Get stats
546    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/// Builder error type
563#[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        // Use public() to disable encryption for this test
615        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()); // public content
622        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, byte) in data.iter_mut().enumerate() {
636            *byte = (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        // Verify store has multiple items (chunks + tree node)
643        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(); // Crosses boundary
753        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        // Roots should be different
771        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        // Use .public() to disable encryption
789        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()); // No encryption key for public content
797        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        // Default config has encryption enabled
806        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()); // Has encryption key
814
815        // Verify we can read it back
816        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        // Data larger than chunk size
830        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        // Verify roundtrip
837        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        // CHK: same content = same hash AND same key
854        assert_eq!(cid1.hash, cid2.hash);
855        assert_eq!(cid1.key, cid2.key);
856        assert_eq!(cid1.to_string(), cid2.to_string());
857    }
858}
hashtree_core/builder.rs

hashtree_core/
builder.rs
