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