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 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 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 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 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 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 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 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 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 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 let mut dir_contents: HashMap<String, Vec<(String, Cid)>> = HashMap::new();
249 dir_contents.insert(String::new(), Vec::new()); 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 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 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 let dir_path = relative.to_string_lossy().to_string();
300 dir_contents.entry(dir_path).or_default();
301 }
302 }
303
304 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 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) });
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 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 dir_cids
360 .get("")
361 .cloned()
362 .ok_or_else(|| anyhow::anyhow!("No root directory"))
363 }
364
365 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 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 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 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 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(); let mut wtxn = self.env.write_txn()?;
503 self.pins.put(&mut wtxn, root_cid.hash.as_slice(), &())?;
505 wtxn.commit()?;
506
507 Ok(cid_str)
508 }
509}