Skip to main content

hashtree_cli/storage/
upload.rs

1use super::*;
2use futures::io::AllowStdIo;
3use std::collections::HashMap;
4use std::io::Read;
5use std::sync::atomic::{AtomicBool, AtomicU64, Ordering};
6
7#[derive(Debug, Default)]
8pub struct AddProgress {
9    bytes_processed: AtomicU64,
10    bytes_total: AtomicU64,
11    files_processed: AtomicU64,
12    files_total: AtomicU64,
13    totals_known: AtomicBool,
14}
15
16#[derive(Clone, Copy, Debug, Eq, PartialEq)]
17pub struct AddProgressSnapshot {
18    pub bytes_processed: u64,
19    pub bytes_total: u64,
20    pub files_processed: u64,
21    pub files_total: u64,
22    pub totals_known: bool,
23}
24
25impl AddProgress {
26    pub fn new() -> Self {
27        Self::default()
28    }
29
30    pub fn record_file_discovered(&self, bytes: u64) {
31        if self.totals_known.load(Ordering::Relaxed) {
32            return;
33        }
34        self.files_total.fetch_add(1, Ordering::Relaxed);
35        self.bytes_total.fetch_add(bytes, Ordering::Relaxed);
36    }
37
38    pub fn mark_totals_known(&self) {
39        self.totals_known.store(true, Ordering::Relaxed);
40    }
41
42    pub fn record_bytes(&self, bytes: u64) {
43        self.bytes_processed.fetch_add(bytes, Ordering::Relaxed);
44    }
45
46    pub fn record_file_finished(&self) {
47        self.files_processed.fetch_add(1, Ordering::Relaxed);
48    }
49
50    pub fn snapshot(&self) -> AddProgressSnapshot {
51        AddProgressSnapshot {
52            bytes_processed: self.bytes_processed.load(Ordering::Relaxed),
53            bytes_total: self.bytes_total.load(Ordering::Relaxed),
54            files_processed: self.files_processed.load(Ordering::Relaxed),
55            files_total: self.files_total.load(Ordering::Relaxed),
56            totals_known: self.totals_known.load(Ordering::Relaxed),
57        }
58    }
59}
60
61impl HashtreeStore {
62    /// Upload a file as raw plaintext and return its CID, with auto-pin
63    pub fn upload_file<P: AsRef<Path>>(&self, file_path: P) -> Result<String> {
64        self.upload_file_with_chunk_size(file_path, None)
65    }
66
67    /// Upload a file without pinning (for blossom uploads that can be evicted)
68    pub fn upload_file_no_pin<P: AsRef<Path>>(&self, file_path: P) -> Result<String> {
69        self.upload_file_internal(file_path, false, None, None)
70    }
71
72    pub fn upload_file_with_chunk_size<P: AsRef<Path>>(
73        &self,
74        file_path: P,
75        chunk_size: Option<usize>,
76    ) -> Result<String> {
77        self.upload_file_internal(file_path, true, chunk_size, None)
78    }
79
80    pub fn upload_file_with_chunk_size_and_progress<P: AsRef<Path>>(
81        &self,
82        file_path: P,
83        chunk_size: Option<usize>,
84        progress: &AddProgress,
85    ) -> Result<String> {
86        self.upload_file_internal(file_path, true, chunk_size, Some(progress))
87    }
88
89    fn upload_file_internal<P: AsRef<Path>>(
90        &self,
91        file_path: P,
92        pin: bool,
93        chunk_size: Option<usize>,
94        progress: Option<&AddProgress>,
95    ) -> Result<String> {
96        let file_path = file_path.as_ref();
97        let file = std::fs::File::open(file_path)
98            .with_context(|| format!("Failed to open file {}", file_path.display()))?;
99        if let Some(progress) = progress {
100            if let Ok(metadata) = file.metadata() {
101                progress.record_file_discovered(metadata.len());
102            }
103        }
104
105        // Store raw plaintext blobs without CHK encryption, streaming from disk.
106        let store = self.store_arc();
107        let mut config = HashTreeConfig::new(store).public();
108        if let Some(chunk_size) = chunk_size {
109            config = config.with_chunk_size(chunk_size);
110        }
111        let tree = HashTree::new(config);
112
113        let (cid, _size) = sync_block_on(async {
114            tree.put_stream_with_progress(AllowStdIo::new(file), |bytes| {
115                if let Some(progress) = progress {
116                    progress.record_bytes(bytes);
117                }
118            })
119            .await
120        })
121        .context("Failed to store file")?;
122        if let Some(progress) = progress {
123            progress.record_file_finished();
124        }
125
126        // Only pin if requested (htree add = pin, blossom upload = no pin)
127        if pin {
128            let mut wtxn = self.env.write_txn()?;
129            self.pins.put(&mut wtxn, cid.hash.as_slice(), &())?;
130            wtxn.commit()?;
131        }
132
133        Ok(to_hex(&cid.hash))
134    }
135
136    /// Upload a file from a stream with progress callbacks
137    pub fn upload_file_stream<R: Read, F>(
138        &self,
139        reader: R,
140        _file_name: impl Into<String>,
141        mut callback: F,
142    ) -> Result<String>
143    where
144        F: FnMut(&str),
145    {
146        // Use HashTree.put_stream for streaming upload without CHK encryption.
147        let store = self.store_arc();
148        let tree = HashTree::new(HashTreeConfig::new(store).public());
149
150        let (cid, _size) = sync_block_on(async { tree.put_stream(AllowStdIo::new(reader)).await })
151            .context("Failed to store file")?;
152
153        let root_hex = to_hex(&cid.hash);
154        callback(&root_hex);
155
156        // Auto-pin on upload
157        let mut wtxn = self.env.write_txn()?;
158        self.pins.put(&mut wtxn, cid.hash.as_slice(), &())?;
159        wtxn.commit()?;
160
161        Ok(root_hex)
162    }
163
164    /// Upload a directory and return its root hash (hex)
165    /// Respects .gitignore and ignores common OS junk files by default.
166    pub fn upload_dir<P: AsRef<Path>>(&self, dir_path: P) -> Result<String> {
167        self.upload_dir_with_options(dir_path, true)
168    }
169
170    /// Upload a directory with options as raw plaintext (no CHK encryption)
171    pub fn upload_dir_with_options<P: AsRef<Path>>(
172        &self,
173        dir_path: P,
174        respect_gitignore: bool,
175    ) -> Result<String> {
176        self.upload_dir_with_options_and_chunk_size(dir_path, respect_gitignore, None)
177    }
178
179    pub fn upload_dir_with_options_and_chunk_size<P: AsRef<Path>>(
180        &self,
181        dir_path: P,
182        respect_gitignore: bool,
183        chunk_size: Option<usize>,
184    ) -> Result<String> {
185        self.upload_dir_with_options_and_chunk_size_internal(
186            dir_path,
187            respect_gitignore,
188            chunk_size,
189            None,
190        )
191    }
192
193    pub fn upload_dir_with_options_and_chunk_size_and_progress<P: AsRef<Path>>(
194        &self,
195        dir_path: P,
196        respect_gitignore: bool,
197        chunk_size: Option<usize>,
198        progress: &AddProgress,
199    ) -> Result<String> {
200        self.upload_dir_with_options_and_chunk_size_internal(
201            dir_path,
202            respect_gitignore,
203            chunk_size,
204            Some(progress),
205        )
206    }
207
208    fn upload_dir_with_options_and_chunk_size_internal<P: AsRef<Path>>(
209        &self,
210        dir_path: P,
211        respect_gitignore: bool,
212        chunk_size: Option<usize>,
213        progress: Option<&AddProgress>,
214    ) -> Result<String> {
215        let dir_path = dir_path.as_ref();
216
217        let store = self.store_arc();
218        let mut config = HashTreeConfig::new(store).public();
219        if let Some(chunk_size) = chunk_size {
220            config = config.with_chunk_size(chunk_size);
221        }
222        let tree = HashTree::new(config);
223
224        let root_cid = sync_block_on(async {
225            self.upload_dir_recursive(&tree, dir_path, dir_path, respect_gitignore, progress, None)
226                .await
227        })
228        .context("Failed to upload directory")?;
229
230        let root_hex = to_hex(&root_cid.hash);
231
232        let mut wtxn = self.env.write_txn()?;
233        self.pins.put(&mut wtxn, root_cid.hash.as_slice(), &())?;
234        wtxn.commit()?;
235
236        Ok(root_hex)
237    }
238
239    async fn upload_dir_recursive<S: Store>(
240        &self,
241        tree: &HashTree<S>,
242        _root_path: &Path,
243        current_path: &Path,
244        respect_gitignore: bool,
245        progress: Option<&AddProgress>,
246        #[cfg(unix)] mut cache: Option<&mut super::add_cache::DirectoryCache>,
247        #[cfg(not(unix))] _cache: Option<()>,
248    ) -> Result<Cid> {
249        // Build directory structure from flat file list - store full Cid with key
250        let mut dir_contents: HashMap<String, Vec<(String, Cid)>> = HashMap::new();
251        dir_contents.insert(String::new(), Vec::new()); // Root
252
253        let walker = crate::ignore_rules::build_content_walker(current_path, respect_gitignore);
254
255        for result in walker {
256            let entry = result?;
257            let path = entry.path();
258
259            // Skip the root directory itself
260            if path == current_path {
261                continue;
262            }
263
264            let relative = path.strip_prefix(current_path).unwrap_or(path);
265
266            #[cfg(unix)]
267            if cache.is_some() {
268                anyhow::ensure!(
269                    entry
270                        .file_type()
271                        .is_some_and(|kind| kind.is_file() || kind.is_dir()),
272                    "cached add requires regular files and directories; symlinks are not followed"
273                );
274            }
275
276            if path.is_file() {
277                #[cfg(unix)]
278                let cached_cid = match cache.as_deref_mut() {
279                    Some(cache) => Some(cache.file(tree, relative, progress).await?),
280                    None => None,
281                };
282                #[cfg(not(unix))]
283                let cached_cid: Option<Cid> = None;
284                let cid = if let Some(cid) = cached_cid {
285                    cid
286                } else {
287                    let file = std::fs::File::open(path)
288                        .with_context(|| format!("Failed to open file {}", path.display()))?;
289                    if let Some(progress) = progress {
290                        if let Ok(metadata) = file.metadata() {
291                            progress.record_file_discovered(metadata.len());
292                        }
293                    }
294                    let (cid, _size) = tree
295                        .put_stream_with_progress(AllowStdIo::new(file), |bytes| {
296                            if let Some(progress) = progress {
297                                progress.record_bytes(bytes);
298                            }
299                        })
300                        .await
301                        .map_err(|e| {
302                            anyhow::anyhow!("Failed to upload file {}: {}", path.display(), e)
303                        })?;
304                    if let Some(progress) = progress {
305                        progress.record_file_finished();
306                    }
307                    cid
308                };
309
310                // Get parent directory path and file name
311                let parent = relative
312                    .parent()
313                    .map(|p| p.to_string_lossy().to_string())
314                    .unwrap_or_default();
315                let name = relative
316                    .file_name()
317                    .map(|n| n.to_string_lossy().to_string())
318                    .unwrap_or_default();
319
320                dir_contents.entry(parent).or_default().push((name, cid));
321            } else if path.is_dir() {
322                // Ensure directory entry exists
323                let dir_path = relative.to_string_lossy().to_string();
324                dir_contents.entry(dir_path).or_default();
325            }
326        }
327
328        // Build directory tree bottom-up
329        self.build_directory_tree(tree, &mut dir_contents).await
330    }
331
332    async fn build_directory_tree<S: Store>(
333        &self,
334        tree: &HashTree<S>,
335        dir_contents: &mut HashMap<String, Vec<(String, Cid)>>,
336    ) -> Result<Cid> {
337        // Sort directories by depth (deepest first) to build bottom-up
338        let mut dirs: Vec<String> = dir_contents.keys().cloned().collect();
339        dirs.sort_by(|a, b| {
340            let depth_a = a.matches('/').count() + if a.is_empty() { 0 } else { 1 };
341            let depth_b = b.matches('/').count() + if b.is_empty() { 0 } else { 1 };
342            depth_b.cmp(&depth_a) // Deepest first
343        });
344
345        let mut dir_cids: HashMap<String, Cid> = HashMap::new();
346
347        for dir_path in dirs {
348            let files = dir_contents.get(&dir_path).cloned().unwrap_or_default();
349
350            let mut entries: Vec<hashtree_core::DirEntry> = files
351                .into_iter()
352                .map(|(name, cid)| hashtree_core::DirEntry::from_cid(name, &cid))
353                .collect();
354
355            // Add subdirectory entries
356            for (subdir_path, cid) in &dir_cids {
357                let parent = Path::new(subdir_path)
358                    .parent()
359                    .map(|p| p.to_string_lossy().to_string())
360                    .unwrap_or_default();
361
362                if parent == dir_path {
363                    let name = Path::new(subdir_path)
364                        .file_name()
365                        .map(|n| n.to_string_lossy().to_string())
366                        .unwrap_or_default();
367                    entries.push(
368                        hashtree_core::DirEntry::from_cid(name, cid)
369                            .with_link_type(hashtree_core::LinkType::Dir),
370                    );
371                }
372            }
373
374            let cid = tree
375                .put_directory(entries)
376                .await
377                .map_err(|e| anyhow::anyhow!("Failed to create directory node: {}", e))?;
378
379            dir_cids.insert(dir_path, cid);
380        }
381
382        // Return root Cid
383        dir_cids
384            .get("")
385            .cloned()
386            .ok_or_else(|| anyhow::anyhow!("No root directory"))
387    }
388
389    /// Upload a file with CHK encryption, returns CID in format "hash:key"
390    pub fn upload_file_encrypted<P: AsRef<Path>>(&self, file_path: P) -> Result<String> {
391        self.upload_file_encrypted_with_chunk_size(file_path, None)
392    }
393
394    pub fn upload_file_encrypted_with_chunk_size<P: AsRef<Path>>(
395        &self,
396        file_path: P,
397        chunk_size: Option<usize>,
398    ) -> Result<String> {
399        self.upload_file_encrypted_with_chunk_size_internal(file_path, chunk_size, None)
400    }
401
402    pub fn upload_file_encrypted_with_chunk_size_and_progress<P: AsRef<Path>>(
403        &self,
404        file_path: P,
405        chunk_size: Option<usize>,
406        progress: &AddProgress,
407    ) -> Result<String> {
408        self.upload_file_encrypted_with_chunk_size_internal(file_path, chunk_size, Some(progress))
409    }
410
411    fn upload_file_encrypted_with_chunk_size_internal<P: AsRef<Path>>(
412        &self,
413        file_path: P,
414        chunk_size: Option<usize>,
415        progress: Option<&AddProgress>,
416    ) -> Result<String> {
417        let file_path = file_path.as_ref();
418        let file = std::fs::File::open(file_path)
419            .with_context(|| format!("Failed to open file {}", file_path.display()))?;
420        if let Some(progress) = progress {
421            if let Ok(metadata) = file.metadata() {
422                progress.record_file_discovered(metadata.len());
423            }
424        }
425
426        // Use unified API with encryption enabled (default), streaming from disk.
427        let store = self.store_arc();
428        let mut config = HashTreeConfig::new(store);
429        if let Some(chunk_size) = chunk_size {
430            config = config.with_chunk_size(chunk_size);
431        }
432        let tree = HashTree::new(config);
433
434        let (cid, _size) = sync_block_on(async {
435            tree.put_stream_with_progress(AllowStdIo::new(file), |bytes| {
436                if let Some(progress) = progress {
437                    progress.record_bytes(bytes);
438                }
439            })
440            .await
441        })
442        .map_err(|e| anyhow::anyhow!("Failed to encrypt file: {}", e))?;
443        if let Some(progress) = progress {
444            progress.record_file_finished();
445        }
446
447        let cid_str = cid.to_string();
448
449        let mut wtxn = self.env.write_txn()?;
450        self.pins.put(&mut wtxn, cid.hash.as_slice(), &())?;
451        wtxn.commit()?;
452
453        Ok(cid_str)
454    }
455
456    /// Upload a directory with CHK encryption, returns CID
457    /// Respects .gitignore and ignores common OS junk files by default.
458    pub fn upload_dir_encrypted<P: AsRef<Path>>(&self, dir_path: P) -> Result<String> {
459        self.upload_dir_encrypted_with_options(dir_path, true)
460    }
461
462    /// Upload a directory with CHK encryption and options
463    /// Returns CID as "hash:key" format for encrypted directories
464    pub fn upload_dir_encrypted_with_options<P: AsRef<Path>>(
465        &self,
466        dir_path: P,
467        respect_gitignore: bool,
468    ) -> Result<String> {
469        self.upload_dir_encrypted_with_options_and_chunk_size(dir_path, respect_gitignore, None)
470    }
471
472    pub fn upload_dir_encrypted_with_options_and_chunk_size<P: AsRef<Path>>(
473        &self,
474        dir_path: P,
475        respect_gitignore: bool,
476        chunk_size: Option<usize>,
477    ) -> Result<String> {
478        self.upload_dir_encrypted_with_options_and_chunk_size_internal(
479            dir_path,
480            respect_gitignore,
481            chunk_size,
482            None,
483        )
484    }
485
486    pub fn upload_dir_encrypted_with_options_and_chunk_size_and_progress<P: AsRef<Path>>(
487        &self,
488        dir_path: P,
489        respect_gitignore: bool,
490        chunk_size: Option<usize>,
491        progress: &AddProgress,
492    ) -> Result<String> {
493        self.upload_dir_encrypted_with_options_and_chunk_size_internal(
494            dir_path,
495            respect_gitignore,
496            chunk_size,
497            Some(progress),
498        )
499    }
500
501    fn upload_dir_encrypted_with_options_and_chunk_size_internal<P: AsRef<Path>>(
502        &self,
503        dir_path: P,
504        respect_gitignore: bool,
505        chunk_size: Option<usize>,
506        progress: Option<&AddProgress>,
507    ) -> Result<String> {
508        let dir_path = dir_path.as_ref();
509        let store = self.store_arc();
510
511        // Use unified API with encryption enabled (default)
512        let mut config = HashTreeConfig::new(store);
513        if let Some(chunk_size) = chunk_size {
514            config = config.with_chunk_size(chunk_size);
515        }
516        let tree = HashTree::new(config);
517
518        let root_cid = sync_block_on(async {
519            self.upload_dir_recursive(&tree, dir_path, dir_path, respect_gitignore, progress, None)
520                .await
521        })
522        .context("Failed to upload encrypted directory")?;
523
524        let cid_str = root_cid.to_string(); // Returns "hash:key" or "hash"
525
526        let mut wtxn = self.env.write_txn()?;
527        // Pin by hash only (the key is for decryption, not identification)
528        self.pins.put(&mut wtxn, root_cid.hash.as_slice(), &())?;
529        wtxn.commit()?;
530
531        Ok(cid_str)
532    }
533
534    /// Cache unchanged source-file CIDs outside the public directory. Unix file
535    /// identity/change metadata and CAS retention are trusted; this is not a
536    /// substitute for a full stored-payload bitrot audit.
537    #[cfg(unix)]
538    pub fn upload_dir_cached_with_progress(
539        &self,
540        source: &Path,
541        respect_ignore: bool,
542        chunk_size: Option<usize>,
543        encrypted: bool,
544        progress: &AddProgress,
545        cache_path: &Path,
546    ) -> Result<String> {
547        let mut config = HashTreeConfig::new(self.store_arc());
548        if !encrypted {
549            config = config.public();
550        }
551        if let Some(chunk_size) = chunk_size {
552            config = config.with_chunk_size(chunk_size);
553        }
554        let tree = HashTree::new(config);
555        let root = sync_block_on(async {
556            let mut cache = super::add_cache::DirectoryCache::open(
557                &tree,
558                source,
559                &self.base_path,
560                respect_ignore,
561                cache_path,
562            )
563            .await?;
564            let root = self
565                .upload_dir_recursive(
566                    &tree,
567                    source,
568                    source,
569                    respect_ignore,
570                    Some(progress),
571                    Some(&mut cache),
572                )
573                .await?;
574            cache.save(&root)?;
575            Ok::<_, anyhow::Error>(root)
576        })?;
577        let mut transaction = self.env.write_txn()?;
578        self.pins.put(&mut transaction, root.hash.as_slice(), &())?;
579        transaction.commit()?;
580        Ok(root.to_string())
581    }
582}