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::hash::sha256;
12use crate::store::Store;
13use crate::types::{Cid, DirEntry, Hash, Link, LinkType, TreeNode};
14
15use crate::crypto::{encrypt_chk, EncryptionKey};
16
17/// Default chunk size: 2MB (optimized for blossom uploads, matches hashtree-ts)
18pub const DEFAULT_CHUNK_SIZE: usize = 2 * 1024 * 1024;
19
20/// BEP52 chunk size: 16KB
21pub const BEP52_CHUNK_SIZE: usize = 16 * 1024;
22
23/// Default max links per tree node (fanout)
24pub const DEFAULT_MAX_LINKS: usize = 174;
25
26/// Builder configuration
27#[derive(Clone)]
28pub struct BuilderConfig<S: Store> {
29    pub store: Arc<S>,
30    pub chunk_size: usize,
31    pub max_links: usize,
32    /// Whether to encrypt content (default: true when encryption feature enabled)
33    pub encrypted: bool,
34}
35
36impl<S: Store> BuilderConfig<S> {
37    pub fn new(store: Arc<S>) -> Self {
38        Self {
39            store,
40            chunk_size: DEFAULT_CHUNK_SIZE,
41            max_links: DEFAULT_MAX_LINKS,
42            encrypted: true,
43        }
44    }
45
46    pub fn with_chunk_size(mut self, chunk_size: usize) -> Self {
47        self.chunk_size = chunk_size;
48        self
49    }
50
51    pub fn with_max_links(mut self, max_links: usize) -> Self {
52        self.max_links = max_links;
53        self
54    }
55
56    /// Disable encryption (store content publicly)
57    pub fn public(mut self) -> Self {
58        self.encrypted = false;
59        self
60    }
61
62    /// Enable encryption (CHK - Content Hash Key)
63    pub fn encrypted(mut self) -> Self {
64        self.encrypted = true;
65        self
66    }
67}
68
69/// TreeBuilder - builds content-addressed merkle trees
70pub struct TreeBuilder<S: Store> {
71    store: Arc<S>,
72    chunk_size: usize,
73    max_links: usize,
74    encrypted: bool,
75}
76
77impl<S: Store> TreeBuilder<S> {
78    pub fn new(config: BuilderConfig<S>) -> Self {
79        Self {
80            store: config.store,
81            chunk_size: config.chunk_size,
82            max_links: config.max_links,
83            encrypted: config.encrypted,
84        }
85    }
86
87    /// Check if encryption is enabled
88    pub fn is_encrypted(&self) -> bool {
89        self.encrypted
90    }
91
92    /// Store a blob directly (small data, no encryption)
93    /// Returns the content hash
94    pub async fn put_blob(&self, data: &[u8]) -> Result<Hash, BuilderError> {
95        let hash = sha256(data);
96        self.store
97            .put(hash, data.to_vec())
98            .await
99            .map_err(|e| BuilderError::Store(e.to_string()))?;
100        Ok(hash)
101    }
102
103    /// Store a chunk with optional encryption
104    /// Returns (hash, optional_key) where hash is of stored data
105    async fn put_chunk_internal(
106        &self,
107        data: &[u8],
108    ) -> Result<(Hash, Option<EncryptionKey>), BuilderError> {
109        if self.encrypted {
110            let (encrypted, key) =
111                encrypt_chk(data).map_err(|e| BuilderError::Encryption(e.to_string()))?;
112            let hash = sha256(&encrypted);
113            self.store
114                .put(hash, encrypted)
115                .await
116                .map_err(|e| BuilderError::Store(e.to_string()))?;
117            Ok((hash, Some(key)))
118        } else {
119            let hash = self.put_blob(data).await?;
120            Ok((hash, None))
121        }
122    }
123
124    /// Store a file, chunking if necessary
125    /// Returns (Cid, size) where Cid contains hash and optional encryption key
126    ///
127    /// When encryption is enabled (default), each chunk is CHK encrypted
128    /// and the result contains the decryption key.
129    pub async fn put(&self, data: &[u8]) -> Result<(Cid, u64), BuilderError> {
130        let size = data.len() as u64;
131
132        // Small file - store as single chunk
133        if data.len() <= self.chunk_size {
134            let (hash, key) = self.put_chunk_internal(data).await?;
135            return Ok((Cid { hash, key }, size));
136        }
137
138        // Large file - chunk it
139        let mut links: Vec<Link> = Vec::new();
140        let mut offset = 0;
141
142        while offset < data.len() {
143            let end = (offset + self.chunk_size).min(data.len());
144            let chunk = &data[offset..end];
145            let chunk_size = chunk.len() as u64;
146            let (hash, key) = self.put_chunk_internal(chunk).await?;
147            links.push(Link {
148                hash,
149                name: None,
150                size: chunk_size,
151                key,
152                link_type: LinkType::Blob, // leaf chunk
153                meta: None,
154            });
155            offset = end;
156        }
157
158        // Build tree from chunks
159        let (root_hash, root_key) = self.build_tree_internal(links, Some(size)).await?;
160
161        Ok((
162            Cid {
163                hash: root_hash,
164                key: root_key,
165            },
166            size,
167        ))
168    }
169
170    /// Build tree and return (hash, optional_key)
171    /// When encrypted, tree nodes are also CHK encrypted
172    async fn build_tree_internal(
173        &self,
174        links: Vec<Link>,
175        total_size: Option<u64>,
176    ) -> Result<(Hash, Option<[u8; 32]>), BuilderError> {
177        // Single link with matching size - return directly
178        if links.len() == 1 {
179            if let Some(ts) = total_size {
180                if links[0].size == ts {
181                    return Ok((links[0].hash, links[0].key));
182                }
183            }
184        }
185
186        if links.len() <= self.max_links {
187            let node = TreeNode {
188                node_type: LinkType::File,
189                links,
190            };
191            let (data, _) = encode_and_hash(&node)?;
192
193            if self.encrypted {
194                let (encrypted, key) =
195                    encrypt_chk(&data).map_err(|e| BuilderError::Encryption(e.to_string()))?;
196                let hash = sha256(&encrypted);
197                self.store
198                    .put(hash, encrypted)
199                    .await
200                    .map_err(|e| BuilderError::Store(e.to_string()))?;
201                return Ok((hash, Some(key)));
202            }
203
204            // Unencrypted path
205            let hash = sha256(&data);
206            self.store
207                .put(hash, data)
208                .await
209                .map_err(|e| BuilderError::Store(e.to_string()))?;
210            return Ok((hash, None));
211        }
212
213        // Too many links - create subtrees
214        let mut sub_links = Vec::new();
215        for batch in links.chunks(self.max_links) {
216            let batch_size: u64 = batch.iter().map(|l| l.size).sum();
217            let (hash, key) =
218                Box::pin(self.build_tree_internal(batch.to_vec(), Some(batch_size))).await?;
219            sub_links.push(Link {
220                hash,
221                name: None,
222                size: batch_size,
223                key,
224                link_type: LinkType::File, // subtree
225                meta: None,
226            });
227        }
228
229        Box::pin(self.build_tree_internal(sub_links, total_size)).await
230    }
231
232    /// Build a balanced tree from links
233    /// Handles fanout by creating intermediate nodes
234    #[allow(dead_code)]
235    async fn build_tree(
236        &self,
237        links: Vec<Link>,
238        total_size: Option<u64>,
239    ) -> Result<Hash, BuilderError> {
240        // Single link with matching size - return it directly
241        if links.len() == 1 {
242            if let Some(ts) = total_size {
243                if links[0].size == ts {
244                    return Ok(links[0].hash);
245                }
246            }
247        }
248
249        // Fits in one node
250        if links.len() <= self.max_links {
251            let node = TreeNode {
252                node_type: LinkType::File,
253                links,
254            };
255            let (data, hash) = encode_and_hash(&node)?;
256            self.store
257                .put(hash, data)
258                .await
259                .map_err(|e| BuilderError::Store(e.to_string()))?;
260            return Ok(hash);
261        }
262
263        // Need to split into sub-trees
264        let mut sub_trees: Vec<Link> = Vec::new();
265
266        for batch in links.chunks(self.max_links) {
267            let batch_size: u64 = batch.iter().map(|l| l.size).sum();
268
269            let node = TreeNode {
270                node_type: LinkType::File,
271                links: batch.to_vec(),
272            };
273            let (data, hash) = encode_and_hash(&node)?;
274            self.store
275                .put(hash, data)
276                .await
277                .map_err(|e| BuilderError::Store(e.to_string()))?;
278
279            sub_trees.push(Link {
280                hash,
281                name: None,
282                size: batch_size,
283                key: None,
284                link_type: LinkType::File, // subtree
285                meta: None,
286            });
287        }
288
289        // Recursively build parent level
290        Box::pin(self.build_tree(sub_trees, total_size)).await
291    }
292
293    /// Build a directory from entries
294    /// Entries can be files or subdirectories
295    pub async fn put_directory(&self, entries: Vec<DirEntry>) -> Result<Hash, BuilderError> {
296        // Sort entries by name for deterministic hashing
297        let mut sorted = entries;
298        sorted.sort_by(|a, b| a.name.cmp(&b.name));
299
300        let links: Vec<Link> = sorted
301            .into_iter()
302            .map(|e| Link {
303                hash: e.hash,
304                name: Some(e.name),
305                size: e.size,
306                key: e.key,
307                link_type: e.link_type,
308                meta: e.meta,
309            })
310            .collect();
311
312        // Fits in one node
313        if links.len() <= self.max_links {
314            let node = TreeNode {
315                node_type: LinkType::Dir,
316                links,
317            };
318            let (data, hash) = encode_and_hash(&node)?;
319            self.store
320                .put(hash, data)
321                .await
322                .map_err(|e| BuilderError::Store(e.to_string()))?;
323            return Ok(hash);
324        }
325
326        self.build_directory_by_chunks(links).await
327    }
328
329    /// Split directories into canonical BUD-17 `_chunk_<start>` fanout nodes.
330    async fn build_directory_by_chunks(&self, links: Vec<Link>) -> Result<Hash, BuilderError> {
331        let indexed_links: Vec<(usize, Link)> = links.into_iter().enumerate().collect();
332        self.build_indexed_directory_chunks(indexed_links).await
333    }
334
335    async fn build_indexed_directory_chunks(
336        &self,
337        links: Vec<(usize, Link)>,
338    ) -> Result<Hash, BuilderError> {
339        let mut sub_trees: Vec<(usize, Link)> = Vec::new();
340
341        for (i, batch) in links.chunks(self.max_links).enumerate() {
342            let start = batch
343                .first()
344                .map(|(start, _)| *start)
345                .unwrap_or(i * self.max_links);
346            let batch_size: u64 = batch.iter().map(|(_, link)| link.size).sum();
347
348            let node = TreeNode {
349                node_type: LinkType::Dir,
350                links: batch.iter().map(|(_, link)| link.clone()).collect(),
351            };
352            let (data, hash) = encode_and_hash(&node)?;
353            self.store
354                .put(hash, data)
355                .await
356                .map_err(|e| BuilderError::Store(e.to_string()))?;
357
358            sub_trees.push((
359                start,
360                Link {
361                    hash,
362                    name: Some(format!("_chunk_{start}")),
363                    size: batch_size,
364                    key: None,
365                    link_type: LinkType::Dir,
366                    meta: None,
367                },
368            ));
369        }
370
371        if sub_trees.len() <= self.max_links {
372            let node = TreeNode {
373                node_type: LinkType::Dir,
374                links: sub_trees.into_iter().map(|(_, link)| link).collect(),
375            };
376            let (data, hash) = encode_and_hash(&node)?;
377            self.store
378                .put(hash, data)
379                .await
380                .map_err(|e| BuilderError::Store(e.to_string()))?;
381            return Ok(hash);
382        }
383
384        // Recursively build more levels
385        Box::pin(self.build_indexed_directory_chunks(sub_trees)).await
386    }
387
388    /// Create a tree node
389    pub async fn put_tree_node(&self, links: Vec<Link>) -> Result<Hash, BuilderError> {
390        let node = TreeNode {
391            node_type: LinkType::Dir,
392            links,
393        };
394
395        let (data, hash) = encode_and_hash(&node)?;
396        self.store
397            .put(hash, data)
398            .await
399            .map_err(|e| BuilderError::Store(e.to_string()))?;
400        Ok(hash)
401    }
402}
403
404/// StreamBuilder - supports incremental appends
405pub struct StreamBuilder<S: Store> {
406    store: Arc<S>,
407    chunk_size: usize,
408    max_links: usize,
409
410    // Current partial chunk being built
411    buffer: Vec<u8>,
412
413    // Completed chunks
414    chunks: Vec<Link>,
415    total_size: u64,
416}
417
418impl<S: Store> StreamBuilder<S> {
419    pub fn new(config: BuilderConfig<S>) -> Self {
420        Self {
421            store: config.store,
422            chunk_size: config.chunk_size,
423            max_links: config.max_links,
424            buffer: Vec::with_capacity(config.chunk_size),
425            chunks: Vec::new(),
426            total_size: 0,
427        }
428    }
429
430    /// Append data to the stream
431    pub async fn append(&mut self, data: &[u8]) -> Result<(), BuilderError> {
432        let mut offset = 0;
433
434        while offset < data.len() {
435            let space = self.chunk_size - self.buffer.len();
436            let to_write = space.min(data.len() - offset);
437
438            self.buffer
439                .extend_from_slice(&data[offset..offset + to_write]);
440            offset += to_write;
441
442            // Flush full chunk
443            if self.buffer.len() == self.chunk_size {
444                self.flush_chunk().await?;
445            }
446        }
447
448        self.total_size += data.len() as u64;
449        Ok(())
450    }
451
452    /// Flush current buffer as a chunk
453    async fn flush_chunk(&mut self) -> Result<(), BuilderError> {
454        if self.buffer.is_empty() {
455            return Ok(());
456        }
457
458        let chunk = std::mem::take(&mut self.buffer);
459        let hash = sha256(&chunk);
460        self.store
461            .put(hash, chunk.clone())
462            .await
463            .map_err(|e| BuilderError::Store(e.to_string()))?;
464
465        self.chunks.push(Link {
466            hash,
467            name: None,
468            size: chunk.len() as u64,
469            key: None,
470            link_type: LinkType::Blob, // Leaf chunk (raw blob)
471            meta: None,
472        });
473
474        self.buffer = Vec::with_capacity(self.chunk_size);
475        Ok(())
476    }
477
478    /// Get current root hash without finalizing
479    /// Useful for checkpoints
480    pub async fn current_root(&mut self) -> Result<Option<Hash>, BuilderError> {
481        if self.chunks.is_empty() && self.buffer.is_empty() {
482            return Ok(None);
483        }
484
485        // Temporarily include buffer
486        let mut temp_chunks = self.chunks.clone();
487        if !self.buffer.is_empty() {
488            let chunk = self.buffer.clone();
489            let hash = sha256(&chunk);
490            self.store
491                .put(hash, chunk.clone())
492                .await
493                .map_err(|e| BuilderError::Store(e.to_string()))?;
494            temp_chunks.push(Link {
495                hash,
496                name: None,
497                size: chunk.len() as u64,
498                key: None,
499                link_type: LinkType::Blob, // Leaf chunk (raw blob)
500                meta: None,
501            });
502        }
503
504        let hash = self
505            .build_tree_from_chunks(&temp_chunks, self.total_size)
506            .await?;
507        Ok(Some(hash))
508    }
509
510    /// Finalize the stream and return root hash
511    pub async fn finalize(mut self) -> Result<(Hash, u64), BuilderError> {
512        // Flush remaining buffer
513        self.flush_chunk().await?;
514
515        if self.chunks.is_empty() {
516            // Empty stream - return hash of empty data
517            let empty_hash = sha256(&[]);
518            self.store
519                .put(empty_hash, vec![])
520                .await
521                .map_err(|e| BuilderError::Store(e.to_string()))?;
522            return Ok((empty_hash, 0));
523        }
524
525        let hash = self
526            .build_tree_from_chunks(&self.chunks, self.total_size)
527            .await?;
528        Ok((hash, self.total_size))
529    }
530
531    /// Build balanced tree from chunks
532    async fn build_tree_from_chunks(
533        &self,
534        chunks: &[Link],
535        total_size: u64,
536    ) -> Result<Hash, BuilderError> {
537        if chunks.len() == 1 {
538            return Ok(chunks[0].hash);
539        }
540
541        if chunks.len() <= self.max_links {
542            let node = TreeNode {
543                node_type: LinkType::File,
544                links: chunks.to_vec(),
545            };
546            let (data, hash) = encode_and_hash(&node)?;
547            self.store
548                .put(hash, data)
549                .await
550                .map_err(|e| BuilderError::Store(e.to_string()))?;
551            return Ok(hash);
552        }
553
554        // Build intermediate level
555        let mut sub_trees: Vec<Link> = Vec::new();
556        for batch in chunks.chunks(self.max_links) {
557            let batch_size: u64 = batch.iter().map(|l| l.size).sum();
558
559            let node = TreeNode {
560                node_type: LinkType::File,
561                links: batch.to_vec(),
562            };
563            let (data, hash) = encode_and_hash(&node)?;
564            self.store
565                .put(hash, data)
566                .await
567                .map_err(|e| BuilderError::Store(e.to_string()))?;
568
569            sub_trees.push(Link {
570                hash,
571                name: None,
572                size: batch_size,
573                key: None,
574                link_type: LinkType::File, // Internal tree node
575                meta: None,
576            });
577        }
578
579        Box::pin(self.build_tree_from_chunks(&sub_trees, total_size)).await
580    }
581
582    /// Get stats
583    pub fn stats(&self) -> StreamStats {
584        StreamStats {
585            chunks: self.chunks.len(),
586            buffered: self.buffer.len(),
587            total_size: self.total_size,
588        }
589    }
590}
591
592#[derive(Debug, Clone, PartialEq)]
593pub struct StreamStats {
594    pub chunks: usize,
595    pub buffered: usize,
596    pub total_size: u64,
597}
598
599/// Builder error type
600#[derive(Debug, thiserror::Error)]
601pub enum BuilderError {
602    #[error("Store error: {0}")]
603    Store(String),
604    #[error("Codec error: {0}")]
605    Codec(#[from] crate::codec::CodecError),
606    #[error("Encryption error: {0}")]
607    Encryption(String),
608}
609
610#[cfg(test)]
611mod tests {
612    use super::*;
613    use crate::store::MemoryStore;
614    use crate::types::to_hex;
615    use std::collections::HashMap;
616
617    fn make_store() -> Arc<MemoryStore> {
618        Arc::new(MemoryStore::new())
619    }
620
621    #[tokio::test]
622    async fn test_put_blob() {
623        let store = make_store();
624        let builder = TreeBuilder::new(BuilderConfig::new(store.clone()));
625
626        let data = vec![1u8, 2, 3, 4, 5];
627        let hash = builder.put_blob(&data).await.unwrap();
628
629        assert_eq!(hash.len(), 32);
630        assert!(store.has(&hash).await.unwrap());
631
632        let retrieved = store.get(&hash).await.unwrap();
633        assert_eq!(retrieved, Some(data));
634    }
635
636    #[tokio::test]
637    async fn test_put_blob_correct_hash() {
638        let store = make_store();
639        let builder = TreeBuilder::new(BuilderConfig::new(store));
640
641        let data = vec![1u8, 2, 3];
642        let hash = builder.put_blob(&data).await.unwrap();
643        let expected_hash = sha256(&data);
644
645        assert_eq!(to_hex(&hash), to_hex(&expected_hash));
646    }
647
648    #[tokio::test]
649    async fn test_put_small() {
650        let store = make_store();
651        // Use public() to disable encryption for this test
652        let builder = TreeBuilder::new(BuilderConfig::new(store.clone()).public());
653
654        let data = vec![1u8, 2, 3, 4, 5];
655        let (cid, size) = builder.put(&data).await.unwrap();
656
657        assert_eq!(size, 5);
658        assert!(cid.key.is_none()); // public content
659        let retrieved = store.get(&cid.hash).await.unwrap();
660        assert_eq!(retrieved, Some(data));
661    }
662
663    #[tokio::test]
664    async fn test_put_chunked() {
665        let store = make_store();
666        let config = BuilderConfig::new(store.clone())
667            .with_chunk_size(1024)
668            .public();
669        let builder = TreeBuilder::new(config);
670
671        let mut data = vec![0u8; 1024 * 2 + 100];
672        for i in 0..data.len() {
673            data[i] = (i % 256) as u8;
674        }
675
676        let (_cid, size) = builder.put(&data).await.unwrap();
677        assert_eq!(size, data.len() as u64);
678
679        // Verify store has multiple items (chunks + tree node)
680        assert!(store.size() > 1);
681    }
682
683    #[tokio::test]
684    async fn test_put_directory() {
685        let store = make_store();
686        let builder = TreeBuilder::new(BuilderConfig::new(store.clone()));
687
688        let file1 = vec![1u8, 2, 3];
689        let file2 = vec![4u8, 5, 6, 7];
690
691        let hash1 = builder.put_blob(&file1).await.unwrap();
692        let hash2 = builder.put_blob(&file2).await.unwrap();
693
694        let dir_hash = builder
695            .put_directory(vec![
696                DirEntry::new("a.txt", hash1).with_size(file1.len() as u64),
697                DirEntry::new("b.txt", hash2).with_size(file2.len() as u64),
698            ])
699            .await
700            .unwrap();
701
702        assert!(store.has(&dir_hash).await.unwrap());
703    }
704
705    #[tokio::test]
706    async fn test_put_directory_sorted() {
707        let store = make_store();
708        let builder = TreeBuilder::new(BuilderConfig::new(store.clone()));
709
710        let hash = builder.put_blob(&[1u8]).await.unwrap();
711
712        let dir_hash = builder
713            .put_directory(vec![
714                DirEntry::new("zebra", hash),
715                DirEntry::new("apple", hash),
716                DirEntry::new("mango", hash),
717            ])
718            .await
719            .unwrap();
720
721        let data = store.get(&dir_hash).await.unwrap().unwrap();
722        let node = crate::codec::decode_tree_node(&data).unwrap();
723
724        let names: Vec<_> = node.links.iter().filter_map(|l| l.name.clone()).collect();
725        assert_eq!(names, vec!["apple", "mango", "zebra"]);
726    }
727
728    #[tokio::test]
729    async fn test_put_tree_node_with_link_meta() {
730        let store = make_store();
731        let builder = TreeBuilder::new(BuilderConfig::new(store.clone()));
732
733        let hash = builder.put_blob(&[1u8]).await.unwrap();
734
735        let mut meta = HashMap::new();
736        meta.insert("version".to_string(), serde_json::json!(2));
737        meta.insert("created".to_string(), serde_json::json!("2024-01-01"));
738
739        let node_hash = builder
740            .put_tree_node(vec![Link {
741                hash,
742                name: Some("test".to_string()),
743                size: 1,
744                key: None,
745                link_type: LinkType::Blob,
746                meta: Some(meta.clone()),
747            }])
748            .await
749            .unwrap();
750
751        let data = store.get(&node_hash).await.unwrap().unwrap();
752        let node = crate::codec::decode_tree_node(&data).unwrap();
753
754        assert!(node.links[0].meta.is_some());
755        let m = node.links[0].meta.as_ref().unwrap();
756        assert_eq!(m.get("version"), Some(&serde_json::json!(2)));
757    }
758
759    #[tokio::test]
760    async fn test_stream_builder() {
761        let store = make_store();
762        let config = BuilderConfig::new(store.clone()).with_chunk_size(100);
763        let mut stream = StreamBuilder::new(config);
764
765        stream.append(&[1u8, 2, 3]).await.unwrap();
766        stream.append(&[4u8, 5]).await.unwrap();
767        stream.append(&[6u8, 7, 8, 9]).await.unwrap();
768
769        let (hash, size) = stream.finalize().await.unwrap();
770
771        assert_eq!(size, 9);
772        assert!(store.has(&hash).await.unwrap());
773    }
774
775    #[tokio::test]
776    async fn test_stream_stats() {
777        let store = make_store();
778        let config = BuilderConfig::new(store).with_chunk_size(100);
779        let mut stream = StreamBuilder::new(config);
780
781        assert_eq!(stream.stats().chunks, 0);
782        assert_eq!(stream.stats().buffered, 0);
783        assert_eq!(stream.stats().total_size, 0);
784
785        stream.append(&[0u8; 50]).await.unwrap();
786        assert_eq!(stream.stats().buffered, 50);
787        assert_eq!(stream.stats().total_size, 50);
788
789        stream.append(&[0u8; 60]).await.unwrap(); // Crosses boundary
790        assert_eq!(stream.stats().chunks, 1);
791        assert_eq!(stream.stats().buffered, 10);
792        assert_eq!(stream.stats().total_size, 110);
793    }
794
795    #[tokio::test]
796    async fn test_stream_current_root() {
797        let store = make_store();
798        let config = BuilderConfig::new(store).with_chunk_size(100);
799        let mut stream = StreamBuilder::new(config);
800
801        stream.append(&[1u8, 2, 3]).await.unwrap();
802        let root1 = stream.current_root().await.unwrap();
803
804        stream.append(&[4u8, 5, 6]).await.unwrap();
805        let root2 = stream.current_root().await.unwrap();
806
807        // Roots should be different
808        assert_ne!(to_hex(&root1.unwrap()), to_hex(&root2.unwrap()));
809    }
810
811    #[tokio::test]
812    async fn test_stream_empty() {
813        let store = make_store();
814        let config = BuilderConfig::new(store.clone());
815        let stream = StreamBuilder::new(config);
816
817        let (hash, size) = stream.finalize().await.unwrap();
818        assert_eq!(size, 0);
819        assert!(store.has(&hash).await.unwrap());
820    }
821
822    #[tokio::test]
823    async fn test_unified_put_public() {
824        let store = make_store();
825        // Use .public() to disable encryption
826        let config = BuilderConfig::new(store.clone()).public();
827        let builder = TreeBuilder::new(config);
828
829        let data = b"Hello, World!";
830        let (cid, size) = builder.put(data).await.unwrap();
831
832        assert_eq!(size, data.len() as u64);
833        assert!(cid.key.is_none()); // No encryption key for public content
834        assert!(store.has(&cid.hash).await.unwrap());
835    }
836
837    #[tokio::test]
838    async fn test_unified_put_encrypted() {
839        use crate::reader::TreeReader;
840
841        let store = make_store();
842        // Default config has encryption enabled
843        let config = BuilderConfig::new(store.clone());
844        let builder = TreeBuilder::new(config);
845
846        let data = b"Hello, encrypted world!";
847        let (cid, size) = builder.put(data).await.unwrap();
848
849        assert_eq!(size, data.len() as u64);
850        assert!(cid.key.is_some()); // Has encryption key
851
852        // Verify we can read it back
853        let reader = TreeReader::new(store);
854        let retrieved = reader.get(&cid).await.unwrap().unwrap();
855        assert_eq!(retrieved, data);
856    }
857
858    #[tokio::test]
859    async fn test_unified_put_encrypted_chunked() {
860        use crate::reader::TreeReader;
861
862        let store = make_store();
863        let config = BuilderConfig::new(store.clone()).with_chunk_size(100);
864        let builder = TreeBuilder::new(config);
865
866        // Data larger than chunk size
867        let data: Vec<u8> = (0..500).map(|i| (i % 256) as u8).collect();
868        let (cid, size) = builder.put(&data).await.unwrap();
869
870        assert_eq!(size, data.len() as u64);
871        assert!(cid.key.is_some());
872
873        // Verify roundtrip
874        let reader = TreeReader::new(store);
875        let retrieved = reader.get(&cid).await.unwrap().unwrap();
876        assert_eq!(retrieved, data);
877    }
878
879    #[tokio::test]
880    async fn test_cid_deterministic() {
881        let store = make_store();
882        let config = BuilderConfig::new(store.clone());
883        let builder = TreeBuilder::new(config);
884
885        let data = b"Same content produces same CID";
886
887        let (cid1, _) = builder.put(data).await.unwrap();
888        let (cid2, _) = builder.put(data).await.unwrap();
889
890        // CHK: same content = same hash AND same key
891        assert_eq!(cid1.hash, cid2.hash);
892        assert_eq!(cid1.key, cid2.key);
893        assert_eq!(cid1.to_string(), cid2.to_string());
894    }
895}
hashtree_core/builder.rs

hashtree_core/
builder.rs
