superfluid_daemon/
media.rs

1//! The content-addressed media pool.
2
3use std::collections::HashSet;
4use std::io::Write;
5use std::path::{Path, PathBuf};
6use std::sync::atomic::{AtomicU64, Ordering};
7
8use sha2::{Digest, Sha256};
9
10pub fn content_hash(bytes: &[u8]) -> String {
11    let d = Sha256::digest(bytes);
12    d.iter().map(|b| format!("{b:02x}")).collect()
13}
14
15pub struct MediaPool {
16    dir: PathBuf,
17}
18
19impl MediaPool {
20    pub fn new(dir: PathBuf) -> MediaPool {
21        MediaPool { dir }
22    }
23
24    pub fn dir(&self) -> &Path {
25        &self.dir
26    }
27
28    pub fn is_valid_id(id: &str) -> bool {
29        id.len() == 64 && id.bytes().all(|b| b.is_ascii_digit() || (b'a'..=b'f').contains(&b))
30    }
31
32    pub fn path_of(&self, hash: &str) -> PathBuf {
33        self.dir.join(hash)
34    }
35
36    pub fn resolve(&self, id: &str) -> Option<PathBuf> {
37        if !Self::is_valid_id(id) {
38            return None;
39        }
40        let p = self.path_of(id);
41        p.is_file().then_some(p)
42    }
43
44    pub fn put(&self, bytes: &[u8]) -> std::io::Result<String> {
45        // Several requests can bring the same file at once, so each write has a file of its own.
46        static WRITES: AtomicU64 = AtomicU64::new(0);
47        let hash = content_hash(bytes);
48        let path = self.path_of(&hash);
49        if path.exists() {
50            return Ok(hash);
51        }
52        std::fs::create_dir_all(&self.dir)?;
53        let tmp = self.dir.join(format!(".{hash}.tmp{}-{}", std::process::id(), WRITES.fetch_add(1, Ordering::Relaxed)));
54        {
55            let mut o = std::fs::OpenOptions::new();
56            o.create(true).write(true).truncate(true);
57            #[cfg(unix)]
58            {
59                use std::os::unix::fs::OpenOptionsExt;
60                o.mode(0o600);
61            }
62            let mut f = o.open(&tmp)?;
63            f.write_all(bytes)?;
64            f.sync_all()?;
65        }
66        std::fs::rename(&tmp, &path)?;
67        if let Ok(d) = std::fs::File::open(&self.dir) {
68            let _ = d.sync_all();
69        }
70        Ok(hash)
71    }
72
73    pub fn exists(&self, hash: &str) -> bool {
74        Self::is_valid_id(hash) && self.path_of(hash).is_file()
75    }
76
77    pub fn sweep(&self, live: &HashSet<String>) -> std::io::Result<(usize, usize)> {
78        let mut kept = 0;
79        let mut removed = 0;
80        let Ok(rd) = std::fs::read_dir(&self.dir) else {
81            return Ok((0, 0));
82        };
83        for e in rd.flatten() {
84            let name = e.file_name().to_string_lossy().into_owned();
85            if live.contains(&name) {
86                kept += 1;
87            } else {
88                let _ = std::fs::remove_file(e.path());
89                removed += 1;
90            }
91        }
92        Ok((kept, removed))
93    }
94}
95
96#[cfg(test)]
97mod tests {
98    use super::{content_hash, MediaPool};
99
100    #[test]
101    fn valid_id_is_64_lowercase_hex() {
102        assert!(MediaPool::is_valid_id(&"a".repeat(64)));
103        assert!(MediaPool::is_valid_id(&"0123456789abcdef".repeat(4)));
104        assert!(!MediaPool::is_valid_id(&"A".repeat(64)));
105        assert!(!MediaPool::is_valid_id(&"a".repeat(63)));
106        assert!(!MediaPool::is_valid_id("../../etc/passwd"));
107        assert!(!MediaPool::is_valid_id("/etc/passwd"));
108        assert!(!MediaPool::is_valid_id(&format!("{}g", "a".repeat(63))));
109    }
110
111    #[test]
112    fn exists_and_resolve_reject_traversal_ids() {
113        let dir = std::env::temp_dir().join(format!("superfluid-media-test-{}", std::process::id()));
114        std::fs::create_dir_all(&dir).unwrap();
115        let pool = MediaPool::new(dir);
116        assert!(!pool.exists("../../../../etc/hosts"));
117        assert!(!pool.exists("/etc/hosts"));
118        assert!(pool.resolve("/etc/hosts").is_none());
119        assert!(pool.resolve("../secrets").is_none());
120    }
121
122    #[test]
123    fn the_same_file_put_at_once_is_kept_once() {
124        let dir = std::env::temp_dir().join(format!("superfluid-media-race-{}", std::process::id()));
125        let _ = std::fs::remove_dir_all(&dir);
126        let pool = std::sync::Arc::new(MediaPool::new(dir.clone()));
127        let bytes = vec![7u8; 1 << 20];
128        let puts: Vec<_> = (0..16)
129            .map(|_| {
130                let (pool, bytes) = (std::sync::Arc::clone(&pool), bytes.clone());
131                std::thread::spawn(move || pool.put(&bytes))
132            })
133            .collect();
134        let hashes: Vec<String> = puts.into_iter().map(|p| p.join().unwrap().expect("every put succeeds")).collect();
135        assert!(hashes.iter().all(|h| *h == content_hash(&bytes)));
136        assert_eq!(std::fs::read_dir(&dir).unwrap().count(), 1, "one file, no temporary left behind");
137    }
138}
superfluid_daemon/media.rs

superfluid_daemon/
media.rs
