superfluid_daemon/
media.rs1use std::collections::HashSet;
4use std::io::Write;
5use std::path::{Path, PathBuf};
6
7use sha2::{Digest, Sha256};
8
9pub fn content_hash(bytes: &[u8]) -> String {
10 let d = Sha256::digest(bytes);
11 d.iter().map(|b| format!("{b:02x}")).collect()
12}
13
14pub struct MediaPool {
15 dir: PathBuf,
16}
17
18impl MediaPool {
19 pub fn new(dir: PathBuf) -> MediaPool {
20 MediaPool { dir }
21 }
22
23 pub fn dir(&self) -> &Path {
24 &self.dir
25 }
26
27 pub fn is_valid_id(id: &str) -> bool {
28 id.len() == 64 && id.bytes().all(|b| b.is_ascii_digit() || (b'a'..=b'f').contains(&b))
29 }
30
31 pub fn path_of(&self, hash: &str) -> PathBuf {
32 self.dir.join(hash)
33 }
34
35 pub fn resolve(&self, id: &str) -> Option<PathBuf> {
36 if !Self::is_valid_id(id) {
37 return None;
38 }
39 let p = self.path_of(id);
40 p.is_file().then_some(p)
41 }
42
43 pub fn put(&self, bytes: &[u8]) -> std::io::Result<String> {
44 let hash = content_hash(bytes);
45 let path = self.path_of(&hash);
46 if path.exists() {
47 return Ok(hash);
48 }
49 std::fs::create_dir_all(&self.dir)?;
50 let tmp = self.dir.join(format!(".{hash}.tmp{}", std::process::id()));
51 {
52 let mut o = std::fs::OpenOptions::new();
53 o.create(true).write(true).truncate(true);
54 #[cfg(unix)]
55 {
56 use std::os::unix::fs::OpenOptionsExt;
57 o.mode(0o600);
58 }
59 let mut f = o.open(&tmp)?;
60 f.write_all(bytes)?;
61 f.sync_all()?;
62 }
63 std::fs::rename(&tmp, &path)?;
64 if let Ok(d) = std::fs::File::open(&self.dir) {
65 let _ = d.sync_all();
66 }
67 Ok(hash)
68 }
69
70 pub fn exists(&self, hash: &str) -> bool {
71 Self::is_valid_id(hash) && self.path_of(hash).is_file()
72 }
73
74 pub fn sweep(&self, live: &HashSet<String>) -> std::io::Result<(usize, usize)> {
75 let mut kept = 0;
76 let mut removed = 0;
77 let Ok(rd) = std::fs::read_dir(&self.dir) else {
78 return Ok((0, 0));
79 };
80 for e in rd.flatten() {
81 let name = e.file_name().to_string_lossy().into_owned();
82 if live.contains(&name) {
83 kept += 1;
84 } else {
85 let _ = std::fs::remove_file(e.path());
86 removed += 1;
87 }
88 }
89 Ok((kept, removed))
90 }
91}
92
93#[cfg(test)]
94mod tests {
95 use super::MediaPool;
96
97 #[test]
98 fn valid_id_is_64_lowercase_hex() {
99 assert!(MediaPool::is_valid_id(&"a".repeat(64)));
100 assert!(MediaPool::is_valid_id(&"0123456789abcdef".repeat(4)));
101 assert!(!MediaPool::is_valid_id(&"A".repeat(64)));
102 assert!(!MediaPool::is_valid_id(&"a".repeat(63)));
103 assert!(!MediaPool::is_valid_id("../../etc/passwd"));
104 assert!(!MediaPool::is_valid_id("/etc/passwd"));
105 assert!(!MediaPool::is_valid_id(&format!("{}g", "a".repeat(63))));
106 }
107
108 #[test]
109 fn exists_and_resolve_reject_traversal_ids() {
110 let dir = std::env::temp_dir().join(format!("superfluid-media-test-{}", std::process::id()));
111 std::fs::create_dir_all(&dir).unwrap();
112 let pool = MediaPool::new(dir);
113 assert!(!pool.exists("../../../../etc/hosts"));
114 assert!(!pool.exists("/etc/hosts"));
115 assert!(pool.resolve("/etc/hosts").is_none());
116 assert!(pool.resolve("../secrets").is_none());
117 }
118}