1use std::{
20 path::{Path, PathBuf},
21 time::{Duration, SystemTime},
22};
23
24use sha2::{Digest, Sha256};
25use tokio::{fs, io::AsyncWriteExt};
26use uuid::Uuid;
27
28use crate::error::{AppError, AppResult};
29
30#[derive(Clone, Debug)]
31pub struct Vault {
32 root: PathBuf,
33 scanner: Option<std::sync::Arc<crate::malware::ScannerConfig>>,
36}
37
38#[derive(Clone, Debug)]
40pub struct StoredBlob {
41 pub sha256: String,
42 pub size: u64,
43}
44
45#[derive(Clone, Debug)]
47pub struct StoredObject {
48 pub sha256: String,
49 pub size: u64,
50 pub modified: SystemTime,
52}
53
54fn is_hex_prefix(name: &str) -> bool {
56 name.len() == 2 && name.bytes().all(|b| matches!(b, b'0'..=b'9' | b'a'..=b'f'))
57}
58
59async fn put_in_place(source: &Path, target: &Path) -> AppResult<()> {
65 if fs::try_exists(target).await? {
66 let path = target.to_path_buf();
67 let renewed = tokio::task::spawn_blocking(move || {
68 std::fs::OpenOptions::new()
69 .write(true)
70 .open(&path)
71 .and_then(|f| f.set_modified(SystemTime::now()))
72 })
73 .await;
74 if matches!(renewed, Ok(Ok(()))) && fs::try_exists(target).await? {
77 fs::remove_file(source).await?;
78 return Ok(());
79 }
80 }
81 if let Some(parent) = target.parent() {
82 fs::create_dir_all(parent).await?;
83 }
84 fs::rename(source, target).await?;
85 Ok(())
86}
87
88impl Vault {
89 pub async fn open(root: impl Into<PathBuf>) -> std::io::Result<Vault> {
91 let vault = Vault {
92 root: root.into(),
93 scanner: None,
94 };
95 fs::create_dir_all(vault.objects_dir()).await?;
96 fs::create_dir_all(vault.tmp_dir()).await?;
97 fs::create_dir_all(vault.uploads_dir()).await?;
98 vault.clean_tmp(Duration::from_secs(24 * 3600)).await?;
99 Ok(vault)
100 }
101
102 pub fn root(&self) -> &Path {
103 &self.root
104 }
105
106 pub fn with_scanner(mut self, scanner: Option<crate::malware::ScannerConfig>) -> Vault {
109 self.scanner = scanner.map(std::sync::Arc::new);
110 self
111 }
112
113 async fn ensure_clean(&self, path: &Path, size: u64) -> AppResult<()> {
116 match &self.scanner {
117 Some(scanner) => crate::malware::ensure_clean(scanner, path, size).await,
118 None => Ok(()),
119 }
120 }
121
122 fn objects_dir(&self) -> PathBuf {
123 self.root.join("objects")
124 }
125
126 pub fn tmp_dir(&self) -> PathBuf {
127 self.root.join("tmp")
128 }
129
130 pub fn object_path(&self, sha256: &str) -> AppResult<PathBuf> {
132 if sha256.len() != 64 || !sha256.bytes().all(|b| b.is_ascii_hexdigit()) {
133 return Err(AppError::BadRequest("invalid content hash".into()));
134 }
135 let sha = sha256.to_ascii_lowercase();
136 Ok(self
137 .objects_dir()
138 .join(&sha[0..2])
139 .join(&sha[2..4])
140 .join(&sha))
141 }
142
143 pub async fn begin_upload(&self, max_bytes: u64) -> AppResult<Upload> {
145 let tmp_path = self.tmp_dir().join(format!("upload-{}", Uuid::new_v4()));
146 let file = fs::File::create(&tmp_path).await?;
147 Ok(Upload {
148 vault: self.clone(),
149 tmp_path,
150 file: Some(file),
151 hasher: Sha256::new(),
152 size: 0,
153 max_bytes,
154 finished: false,
155 })
156 }
157
158 pub async fn store_bytes(&self, data: &[u8]) -> AppResult<StoredBlob> {
160 let mut upload = self.begin_upload(u64::MAX).await?;
161 upload.write(data).await?;
162 upload.finish().await
163 }
164
165 pub async fn open_blob(&self, sha256: &str) -> AppResult<fs::File> {
166 let path = self.object_path(sha256)?;
167 fs::File::open(&path).await.map_err(|e| {
168 if e.kind() == std::io::ErrorKind::NotFound {
169 tracing::error!(sha256, "blob missing from vault");
170 AppError::NotFound(
171 "The file content is missing from the vault. Please contact an administrator."
172 .into(),
173 )
174 } else {
175 AppError::Io(e)
176 }
177 })
178 }
179
180 pub async fn exists(&self, sha256: &str) -> bool {
181 match self.object_path(sha256) {
182 Ok(p) => fs::try_exists(p).await.unwrap_or(false),
183 Err(_) => false,
184 }
185 }
186
187 pub async fn verify(&self, sha256: &str) -> AppResult<bool> {
189 use tokio::io::AsyncReadExt;
190 let mut file = self.open_blob(sha256).await?;
191 let mut hasher = Sha256::new();
192 let mut buf = vec![0u8; 256 * 1024];
193 loop {
194 let n = file.read(&mut buf).await?;
195 if n == 0 {
196 break;
197 }
198 hasher.update(&buf[..n]);
199 }
200 Ok(hex::encode(hasher.finalize()) == sha256.to_ascii_lowercase())
201 }
202
203 pub async fn check_writable(&self) -> AppResult<()> {
205 let probe = self.tmp_dir().join(format!(".probe-{}", Uuid::new_v4()));
206 fs::write(&probe, b"ok").await?;
207 fs::remove_file(&probe).await?;
208 Ok(())
209 }
210
211 pub fn uploads_dir(&self) -> PathBuf {
214 self.root.join("uploads")
215 }
216
217 fn staging_path(&self, session: Uuid) -> PathBuf {
218 self.uploads_dir().join(session.simple().to_string())
219 }
220
221 pub async fn create_staging(&self, session: Uuid) -> AppResult<()> {
223 fs::File::create(self.staging_path(session)).await?;
224 Ok(())
225 }
226
227 pub async fn write_staged(&self, session: Uuid, offset: u64, chunk: &[u8]) -> AppResult<()> {
232 use tokio::io::AsyncSeekExt;
233 let mut file = fs::OpenOptions::new()
234 .write(true)
235 .open(self.staging_path(session))
236 .await
237 .map_err(|e| {
238 if e.kind() == std::io::ErrorKind::NotFound {
239 AppError::Conflict(
240 "The data of this upload is gone – start the upload again.".into(),
241 )
242 } else {
243 AppError::Io(e)
244 }
245 })?;
246 let len = file.metadata().await?.len();
247 if len < offset {
248 return Err(AppError::Conflict(
249 "The data of this upload is incomplete – start the upload again.".into(),
250 ));
251 }
252 file.set_len(offset).await?;
253 file.seek(std::io::SeekFrom::Start(offset)).await?;
254 file.write_all(chunk).await?;
255 file.flush().await?;
256 file.sync_data().await?;
257 Ok(())
258 }
259
260 pub async fn hash_staged(&self, session: Uuid) -> AppResult<StoredBlob> {
263 use tokio::io::AsyncReadExt;
264 let mut file = fs::File::open(self.staging_path(session)).await?;
265 let mut hasher = Sha256::new();
266 let mut buf = vec![0u8; 1024 * 1024];
267 let mut size = 0u64;
268 loop {
269 let n = file.read(&mut buf).await?;
270 if n == 0 {
271 break;
272 }
273 size += n as u64;
274 hasher.update(&buf[..n]);
275 }
276 Ok(StoredBlob {
277 sha256: hex::encode(hasher.finalize()),
278 size,
279 })
280 }
281
282 pub async fn adopt_staged(&self, session: Uuid, sha256: &str) -> AppResult<()> {
285 let source = self.staging_path(session);
286 let target = self.object_path(sha256)?;
287 if self.scanner.is_some() {
288 let size = fs::metadata(&source).await?.len();
289 self.ensure_clean(&source, size).await?;
290 }
291 put_in_place(&source, &target).await
292 }
293
294 pub async fn remove_staged(&self, session: Uuid) -> AppResult<()> {
296 match fs::remove_file(self.staging_path(session)).await {
297 Ok(()) => Ok(()),
298 Err(e) if e.kind() == std::io::ErrorKind::NotFound => Ok(()),
299 Err(e) => Err(AppError::Io(e)),
300 }
301 }
302
303 pub async fn stale_staged(&self, older_than: Duration) -> AppResult<Vec<Uuid>> {
305 let mut stale = Vec::new();
306 let mut entries = fs::read_dir(self.uploads_dir()).await?;
307 let now = SystemTime::now();
308 while let Some(entry) = entries.next_entry().await? {
309 let Some(id) = entry
310 .file_name()
311 .to_str()
312 .and_then(|n| Uuid::try_parse(n).ok())
313 else {
314 continue;
315 };
316 let old = entry
317 .metadata()
318 .await
319 .and_then(|m| m.modified())
320 .map(|modified| now.duration_since(modified).unwrap_or_default() > older_than)
321 .unwrap_or(false);
322 if old {
323 stale.push(id);
324 }
325 }
326 Ok(stale)
327 }
328
329 pub async fn object_prefixes(&self) -> AppResult<Vec<String>> {
332 let mut prefixes = Vec::new();
333 let mut entries = fs::read_dir(self.objects_dir()).await?;
334 while let Some(entry) = entries.next_entry().await? {
335 if let Some(name) = entry.file_name().to_str() {
336 if is_hex_prefix(name) && entry.file_type().await?.is_dir() {
337 prefixes.push(name.to_string());
338 }
339 }
340 }
341 prefixes.sort();
342 Ok(prefixes)
343 }
344
345 pub async fn objects_in(&self, prefix: &str) -> AppResult<Vec<StoredObject>> {
349 let mut found = Vec::new();
350 if !is_hex_prefix(prefix) {
351 return Ok(found);
352 }
353 let top = self.objects_dir().join(prefix);
354 let mut dirs = fs::read_dir(&top).await?;
355 while let Some(dir) = dirs.next_entry().await? {
356 let Some(second) = dir.file_name().to_str().map(str::to_string) else {
357 continue;
358 };
359 if !is_hex_prefix(&second) || !dir.file_type().await?.is_dir() {
360 continue;
361 }
362 let mut files = fs::read_dir(dir.path()).await?;
363 while let Some(file) = files.next_entry().await? {
364 let Some(name) = file.file_name().to_str().map(str::to_string) else {
365 continue;
366 };
367 let at_address = name.len() == 64
368 && name.bytes().all(|b| matches!(b, b'0'..=b'9' | b'a'..=b'f'))
369 && name[0..2] == *prefix
370 && name[2..4] == second;
371 if !at_address {
372 continue;
373 }
374 let meta = file.metadata().await?;
375 if meta.is_file() {
376 found.push(StoredObject {
377 sha256: name,
378 size: meta.len(),
379 modified: meta.modified().unwrap_or(SystemTime::UNIX_EPOCH),
380 });
381 }
382 }
383 }
384 Ok(found)
385 }
386
387 pub async fn object_age(&self, sha256: &str) -> AppResult<Option<Duration>> {
390 match fs::metadata(self.object_path(sha256)?).await {
391 Ok(meta) => {
392 let modified = meta.modified().unwrap_or(SystemTime::UNIX_EPOCH);
393 Ok(Some(
394 SystemTime::now()
395 .duration_since(modified)
396 .unwrap_or_default(),
397 ))
398 }
399 Err(e) if e.kind() == std::io::ErrorKind::NotFound => Ok(None),
400 Err(e) => Err(AppError::Io(e)),
401 }
402 }
403
404 pub async fn remove_object_older_than(&self, sha256: &str, grace: Duration) -> AppResult<bool> {
408 match self.object_age(sha256).await? {
409 Some(age) if age >= grace => {}
410 _ => return Ok(false),
411 }
412 match fs::remove_file(self.object_path(sha256)?).await {
413 Ok(()) => Ok(true),
414 Err(e) if e.kind() == std::io::ErrorKind::NotFound => Ok(false),
415 Err(e) => Err(AppError::Io(e)),
416 }
417 }
418
419 async fn clean_tmp(&self, older_than: Duration) -> std::io::Result<()> {
420 let mut entries = fs::read_dir(self.tmp_dir()).await?;
421 let now = SystemTime::now();
422 while let Some(entry) = entries.next_entry().await? {
423 let stale = entry
424 .metadata()
425 .await
426 .and_then(|m| m.modified())
427 .map(|modified| now.duration_since(modified).unwrap_or_default() > older_than)
428 .unwrap_or(false);
429 if stale {
430 let _ = fs::remove_file(entry.path()).await;
431 }
432 }
433 Ok(())
434 }
435}
436
437pub struct Upload {
440 vault: Vault,
441 tmp_path: PathBuf,
442 file: Option<fs::File>,
443 hasher: Sha256,
444 size: u64,
445 max_bytes: u64,
446 finished: bool,
447}
448
449impl Upload {
450 pub async fn write(&mut self, chunk: &[u8]) -> AppResult<()> {
451 self.size += chunk.len() as u64;
452 if self.size > self.max_bytes {
453 return Err(AppError::PayloadTooLarge(format!(
454 "The file exceeds the upload limit of {} MiB.",
455 self.max_bytes / (1024 * 1024)
456 )));
457 }
458 self.hasher.update(chunk);
459 let file = self
460 .file
461 .as_mut()
462 .ok_or_else(|| AppError::Other(anyhow::anyhow!("upload already finished")))?;
463 file.write_all(chunk).await?;
464 Ok(())
465 }
466
467 pub fn size(&self) -> u64 {
468 self.size
469 }
470
471 pub async fn finish(&mut self) -> AppResult<StoredBlob> {
473 let mut file = self
474 .file
475 .take()
476 .ok_or_else(|| AppError::Other(anyhow::anyhow!("upload already finished")))?;
477 file.flush().await?;
478 file.sync_all().await?;
479 drop(file);
480 let sha256 = hex::encode(std::mem::take(&mut self.hasher).finalize());
481 let target = self.vault.object_path(&sha256)?;
482 self.vault.ensure_clean(&self.tmp_path, self.size).await?;
485 put_in_place(&self.tmp_path, &target).await?;
487 self.finished = true;
488 Ok(StoredBlob {
489 sha256,
490 size: self.size,
491 })
492 }
493}
494
495impl Drop for Upload {
496 fn drop(&mut self) {
497 if !self.finished {
498 let _ = std::fs::remove_file(&self.tmp_path);
499 }
500 }
501}
502
503#[cfg(test)]
504mod tests {
505 use super::*;
506
507 fn temp_root() -> PathBuf {
508 std::env::temp_dir().join(format!("teamcraft-vault-test-{}", Uuid::new_v4()))
509 }
510
511 #[tokio::test]
512 async fn stores_content_addressed_and_deduplicates() {
513 let root = temp_root();
514 let vault = Vault::open(&root).await.unwrap();
515 let a = vault.store_bytes(b"hello vault").await.unwrap();
516 let b = vault.store_bytes(b"hello vault").await.unwrap();
517 assert_eq!(a.sha256, b.sha256);
518 assert_eq!(a.size, 11);
519 assert_eq!(a.sha256, hex::encode(Sha256::digest(b"hello vault")));
520 assert!(vault.exists(&a.sha256).await);
521 assert!(vault.verify(&a.sha256).await.unwrap());
522 let mut tmp = fs::read_dir(root.join("tmp")).await.unwrap();
524 assert!(tmp.next_entry().await.unwrap().is_none());
525 fs::remove_dir_all(root).await.ok();
526 }
527
528 #[tokio::test]
529 async fn enforces_size_limit_and_cleans_up() {
530 let root = temp_root();
531 let vault = Vault::open(&root).await.unwrap();
532 {
533 let mut upload = vault.begin_upload(4).await.unwrap();
534 upload.write(b"1234").await.unwrap();
535 let err = upload.write(b"5").await.unwrap_err();
536 assert!(matches!(err, AppError::PayloadTooLarge(_)));
537 }
538 let mut tmp = fs::read_dir(root.join("tmp")).await.unwrap();
539 assert!(
540 tmp.next_entry().await.unwrap().is_none(),
541 "temp file removed on drop"
542 );
543 fs::remove_dir_all(root).await.ok();
544 }
545
546 #[tokio::test]
547 async fn staged_chunks_resume_and_move_into_the_vault() {
548 let root = temp_root();
549 let vault = Vault::open(&root).await.unwrap();
550 let id = Uuid::new_v4();
551 vault.create_staging(id).await.unwrap();
552 vault.write_staged(id, 0, b"hello ").await.unwrap();
553 vault.write_staged(id, 6, b"wor").await.unwrap();
555 vault.write_staged(id, 6, b"vault").await.unwrap();
557 assert!(matches!(
558 vault.write_staged(id, 99, b"gap").await,
559 Err(AppError::Conflict(_))
560 ));
561 let blob = vault.hash_staged(id).await.unwrap();
562 assert_eq!(blob.size, 11);
563 assert_eq!(blob.sha256, hex::encode(Sha256::digest(b"hello vault")));
564 vault.adopt_staged(id, &blob.sha256).await.unwrap();
565 assert!(vault.verify(&blob.sha256).await.unwrap());
566 let twin = Uuid::new_v4();
568 vault.create_staging(twin).await.unwrap();
569 vault.write_staged(twin, 0, b"hello vault").await.unwrap();
570 vault.adopt_staged(twin, &blob.sha256).await.unwrap();
571 assert!(vault.stale_staged(Duration::ZERO).await.unwrap().is_empty());
572 vault.remove_staged(twin).await.unwrap();
573 fs::remove_dir_all(root).await.ok();
574 }
575
576 #[tokio::test]
577 async fn infected_content_never_enters_the_vault() {
578 let root = temp_root();
579 let scanner = crate::malware::ScannerConfig {
580 addr: crate::malware::tests::fake_clamd().await,
581 max_bytes: 1024 * 1024,
582 timeout: Duration::from_secs(5),
583 };
584 let vault = Vault::open(&root)
585 .await
586 .unwrap()
587 .with_scanner(Some(scanner));
588 let bad = b"X5O!P%@AP EICAR test file";
590 assert!(matches!(
591 vault.store_bytes(bad).await,
592 Err(AppError::BadRequest(_))
593 ));
594 let sha = hex::encode(Sha256::digest(bad));
595 assert!(!vault.exists(&sha).await);
596 assert_eq!(std::fs::read_dir(vault.tmp_dir()).unwrap().count(), 0);
597 let good = vault.store_bytes(b"clean content").await.unwrap();
598 assert!(vault.exists(&good.sha256).await);
599 let session = Uuid::new_v4();
601 vault.create_staging(session).await.unwrap();
602 vault.write_staged(session, 0, bad).await.unwrap();
603 let staged = vault.hash_staged(session).await.unwrap();
604 assert!(vault.adopt_staged(session, &staged.sha256).await.is_err());
605 assert!(!vault.exists(&staged.sha256).await);
606 fs::remove_dir_all(root).await.ok();
607 }
608
609 #[test]
610 fn rejects_bad_hashes() {
611 let vault = Vault {
612 root: PathBuf::from("/nonexistent"),
613 scanner: None,
614 };
615 assert!(vault.object_path("../../etc/passwd").is_err());
616 assert!(vault.object_path(&"g".repeat(64)).is_err());
617 let ok = vault.object_path(&"ab".repeat(32)).unwrap();
618 assert!(ok.ends_with(format!("objects/ab/ab/{}", "ab".repeat(32))));
619 }
620
621 fn age(path: &Path, ago: Duration) {
623 std::fs::OpenOptions::new()
624 .write(true)
625 .open(path)
626 .unwrap()
627 .set_modified(SystemTime::now() - ago)
628 .unwrap();
629 }
630
631 #[tokio::test]
632 async fn lists_blobs_and_removes_only_old_ones() {
633 let root = temp_root();
634 let vault = Vault::open(&root).await.unwrap();
635 let a = vault.store_bytes(b"first").await.unwrap();
636 let b = vault.store_bytes(b"second").await.unwrap();
637 let a_dir = vault.object_path(&a.sha256).unwrap();
639 let a_dir = a_dir.parent().unwrap();
640 std::fs::write(a_dir.join("notes.txt"), b"x").unwrap();
641 let misplaced = if a.sha256.starts_with("ee") {
642 "dd"
643 } else {
644 "ee"
645 }
646 .repeat(32);
647 std::fs::write(a_dir.join(&misplaced), b"misplaced").unwrap();
648 std::fs::create_dir_all(root.join("objects/zz")).unwrap();
649
650 let mut listed = Vec::new();
651 for prefix in vault.object_prefixes().await.unwrap() {
652 listed.extend(vault.objects_in(&prefix).await.unwrap());
653 }
654 let mut names: Vec<_> = listed.iter().map(|o| o.sha256.clone()).collect();
655 names.sort();
656 let mut expected = vec![a.sha256.clone(), b.sha256.clone()];
657 expected.sort();
658 assert_eq!(names, expected);
659 assert!(vault.objects_in("../..").await.unwrap().is_empty());
660
661 let hour = Duration::from_secs(3600);
663 assert!(!vault
664 .remove_object_older_than(&a.sha256, hour)
665 .await
666 .unwrap());
667 age(&vault.object_path(&a.sha256).unwrap(), 2 * hour);
668 assert!(vault.object_age(&a.sha256).await.unwrap().unwrap() >= hour);
669 vault.store_bytes(b"first").await.unwrap();
671 assert!(vault.object_age(&a.sha256).await.unwrap().unwrap() < hour);
672 assert!(!vault
673 .remove_object_older_than(&a.sha256, hour)
674 .await
675 .unwrap());
676 age(&vault.object_path(&a.sha256).unwrap(), 2 * hour);
678 assert!(vault
679 .remove_object_older_than(&a.sha256, hour)
680 .await
681 .unwrap());
682 assert!(!vault.exists(&a.sha256).await);
683 assert!(vault.object_age(&a.sha256).await.unwrap().is_none());
684 assert!(!vault
685 .remove_object_older_than(&a.sha256, hour)
686 .await
687 .unwrap());
688 age(&vault.object_path(&b.sha256).unwrap(), 2 * hour);
690 let id = Uuid::new_v4();
691 vault.create_staging(id).await.unwrap();
692 vault.write_staged(id, 0, b"second").await.unwrap();
693 vault.adopt_staged(id, &b.sha256).await.unwrap();
694 assert!(vault.object_age(&b.sha256).await.unwrap().unwrap() < hour);
695 assert!(vault.verify(&b.sha256).await.unwrap());
696 fs::remove_dir_all(root).await.ok();
697 }
698}