Skip to main content

superfluid_daemon/
media.rs

1//! The content-addressed media pool.
2
3use 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}