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)
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    ) -> Result<Cid> {
247        // Build directory structure from flat file list - store full Cid with key
248        let mut dir_contents: HashMap<String, Vec<(String, Cid)>> = HashMap::new();
249        dir_contents.insert(String::new(), Vec::new()); // Root
250
251        let walker = crate::ignore_rules::build_content_walker(current_path, respect_gitignore);
252
253        for result in walker {
254            let entry = result?;
255            let path = entry.path();
256
257            // Skip the root directory itself
258            if path == current_path {
259                continue;
260            }
261
262            let relative = path.strip_prefix(current_path).unwrap_or(path);
263
264            if path.is_file() {
265                let file = std::fs::File::open(path)
266                    .with_context(|| format!("Failed to open file {}", path.display()))?;
267                if let Some(progress) = progress {
268                    if let Ok(metadata) = file.metadata() {
269                        progress.record_file_discovered(metadata.len());
270                    }
271                }
272                let (cid, _size) = tree
273                    .put_stream_with_progress(AllowStdIo::new(file), |bytes| {
274                        if let Some(progress) = progress {
275                            progress.record_bytes(bytes);
276                        }
277                    })
278                    .await
279                    .map_err(|e| {
280                        anyhow::anyhow!("Failed to upload file {}: {}", path.display(), e)
281                    })?;
282                if let Some(progress) = progress {
283                    progress.record_file_finished();
284                }
285
286                // Get parent directory path and file name
287                let parent = relative
288                    .parent()
289                    .map(|p| p.to_string_lossy().to_string())
290                    .unwrap_or_default();
291                let name = relative
292                    .file_name()
293                    .map(|n| n.to_string_lossy().to_string())
294                    .unwrap_or_default();
295
296                dir_contents.entry(parent).or_default().push((name, cid));
297            } else if path.is_dir() {
298                // Ensure directory entry exists
299                let dir_path = relative.to_string_lossy().to_string();
300                dir_contents.entry(dir_path).or_default();
301            }
302        }
303
304        // Build directory tree bottom-up
305        self.build_directory_tree(tree, &mut dir_contents).await
306    }
307
308    async fn build_directory_tree<S: Store>(
309        &self,
310        tree: &HashTree<S>,
311        dir_contents: &mut HashMap<String, Vec<(String, Cid)>>,
312    ) -> Result<Cid> {
313        // Sort directories by depth (deepest first) to build bottom-up
314        let mut dirs: Vec<String> = dir_contents.keys().cloned().collect();
315        dirs.sort_by(|a, b| {
316            let depth_a = a.matches('/').count() + if a.is_empty() { 0 } else { 1 };
317            let depth_b = b.matches('/').count() + if b.is_empty() { 0 } else { 1 };
318            depth_b.cmp(&depth_a) // Deepest first
319        });
320
321        let mut dir_cids: HashMap<String, Cid> = HashMap::new();
322
323        for dir_path in dirs {
324            let files = dir_contents.get(&dir_path).cloned().unwrap_or_default();
325
326            let mut entries: Vec<hashtree_core::DirEntry> = files
327                .into_iter()
328                .map(|(name, cid)| hashtree_core::DirEntry::from_cid(name, &cid))
329                .collect();
330
331            // Add subdirectory entries
332            for (subdir_path, cid) in &dir_cids {
333                let parent = Path::new(subdir_path)
334                    .parent()
335                    .map(|p| p.to_string_lossy().to_string())
336                    .unwrap_or_default();
337
338                if parent == dir_path {
339                    let name = Path::new(subdir_path)
340                        .file_name()
341                        .map(|n| n.to_string_lossy().to_string())
342                        .unwrap_or_default();
343                    entries.push(
344                        hashtree_core::DirEntry::from_cid(name, cid)
345                            .with_link_type(hashtree_core::LinkType::Dir),
346                    );
347                }
348            }
349
350            let cid = tree
351                .put_directory(entries)
352                .await
353                .map_err(|e| anyhow::anyhow!("Failed to create directory node: {}", e))?;
354
355            dir_cids.insert(dir_path, cid);
356        }
357
358        // Return root Cid
359        dir_cids
360            .get("")
361            .cloned()
362            .ok_or_else(|| anyhow::anyhow!("No root directory"))
363    }
364
365    /// Upload a file with CHK encryption, returns CID in format "hash:key"
366    pub fn upload_file_encrypted<P: AsRef<Path>>(&self, file_path: P) -> Result<String> {
367        self.upload_file_encrypted_with_chunk_size(file_path, None)
368    }
369
370    pub fn upload_file_encrypted_with_chunk_size<P: AsRef<Path>>(
371        &self,
372        file_path: P,
373        chunk_size: Option<usize>,
374    ) -> Result<String> {
375        self.upload_file_encrypted_with_chunk_size_internal(file_path, chunk_size, None)
376    }
377
378    pub fn upload_file_encrypted_with_chunk_size_and_progress<P: AsRef<Path>>(
379        &self,
380        file_path: P,
381        chunk_size: Option<usize>,
382        progress: &AddProgress,
383    ) -> Result<String> {
384        self.upload_file_encrypted_with_chunk_size_internal(file_path, chunk_size, Some(progress))
385    }
386
387    fn upload_file_encrypted_with_chunk_size_internal<P: AsRef<Path>>(
388        &self,
389        file_path: P,
390        chunk_size: Option<usize>,
391        progress: Option<&AddProgress>,
392    ) -> Result<String> {
393        let file_path = file_path.as_ref();
394        let file = std::fs::File::open(file_path)
395            .with_context(|| format!("Failed to open file {}", file_path.display()))?;
396        if let Some(progress) = progress {
397            if let Ok(metadata) = file.metadata() {
398                progress.record_file_discovered(metadata.len());
399            }
400        }
401
402        // Use unified API with encryption enabled (default), streaming from disk.
403        let store = self.store_arc();
404        let mut config = HashTreeConfig::new(store);
405        if let Some(chunk_size) = chunk_size {
406            config = config.with_chunk_size(chunk_size);
407        }
408        let tree = HashTree::new(config);
409
410        let (cid, _size) = sync_block_on(async {
411            tree.put_stream_with_progress(AllowStdIo::new(file), |bytes| {
412                if let Some(progress) = progress {
413                    progress.record_bytes(bytes);
414                }
415            })
416            .await
417        })
418        .map_err(|e| anyhow::anyhow!("Failed to encrypt file: {}", e))?;
419        if let Some(progress) = progress {
420            progress.record_file_finished();
421        }
422
423        let cid_str = cid.to_string();
424
425        let mut wtxn = self.env.write_txn()?;
426        self.pins.put(&mut wtxn, cid.hash.as_slice(), &())?;
427        wtxn.commit()?;
428
429        Ok(cid_str)
430    }
431
432    /// Upload a directory with CHK encryption, returns CID
433    /// Respects .gitignore and ignores common OS junk files by default.
434    pub fn upload_dir_encrypted<P: AsRef<Path>>(&self, dir_path: P) -> Result<String> {
435        self.upload_dir_encrypted_with_options(dir_path, true)
436    }
437
438    /// Upload a directory with CHK encryption and options
439    /// Returns CID as "hash:key" format for encrypted directories
440    pub fn upload_dir_encrypted_with_options<P: AsRef<Path>>(
441        &self,
442        dir_path: P,
443        respect_gitignore: bool,
444    ) -> Result<String> {
445        self.upload_dir_encrypted_with_options_and_chunk_size(dir_path, respect_gitignore, None)
446    }
447
448    pub fn upload_dir_encrypted_with_options_and_chunk_size<P: AsRef<Path>>(
449        &self,
450        dir_path: P,
451        respect_gitignore: bool,
452        chunk_size: Option<usize>,
453    ) -> Result<String> {
454        self.upload_dir_encrypted_with_options_and_chunk_size_internal(
455            dir_path,
456            respect_gitignore,
457            chunk_size,
458            None,
459        )
460    }
461
462    pub fn upload_dir_encrypted_with_options_and_chunk_size_and_progress<P: AsRef<Path>>(
463        &self,
464        dir_path: P,
465        respect_gitignore: bool,
466        chunk_size: Option<usize>,
467        progress: &AddProgress,
468    ) -> Result<String> {
469        self.upload_dir_encrypted_with_options_and_chunk_size_internal(
470            dir_path,
471            respect_gitignore,
472            chunk_size,
473            Some(progress),
474        )
475    }
476
477    fn upload_dir_encrypted_with_options_and_chunk_size_internal<P: AsRef<Path>>(
478        &self,
479        dir_path: P,
480        respect_gitignore: bool,
481        chunk_size: Option<usize>,
482        progress: Option<&AddProgress>,
483    ) -> Result<String> {
484        let dir_path = dir_path.as_ref();
485        let store = self.store_arc();
486
487        // Use unified API with encryption enabled (default)
488        let mut config = HashTreeConfig::new(store);
489        if let Some(chunk_size) = chunk_size {
490            config = config.with_chunk_size(chunk_size);
491        }
492        let tree = HashTree::new(config);
493
494        let root_cid = sync_block_on(async {
495            self.upload_dir_recursive(&tree, dir_path, dir_path, respect_gitignore, progress)
496                .await
497        })
498        .context("Failed to upload encrypted directory")?;
499
500        let cid_str = root_cid.to_string(); // Returns "hash:key" or "hash"
501
502        let mut wtxn = self.env.write_txn()?;
503        // Pin by hash only (the key is for decryption, not identification)
504        self.pins.put(&mut wtxn, root_cid.hash.as_slice(), &())?;
505        wtxn.commit()?;
506
507        Ok(cid_str)
508    }
509}