Skip to main content

max / makenotwork

11.9 KB · 316 lines History Blame Raw
1 //! Tempfile spool for large scan jobs.
2 //!
3 //! Objects above `SCAN_MAX_MEMORY_BYTES` stream from S3 into a tempfile
4 //! under `SCAN_SPOOL_DIR` rather than buffering into a `Vec<u8>`. The
5 //! resulting `SpoolHandle` exposes both a stable filesystem path (for
6 //! layers that need random access, ZIP central directory, YARA) and an
7 //! `AsyncRead` (for clamd INSTREAM). The file is removed on drop.
8
9 use crate::constants::{
10 SCAN_SPOOL_FREE_RESERVE_BYTES, SCAN_SPOOL_MAX_BYTES, SCAN_SPOOL_ORPHAN_AGE_SECS,
11 SCAN_SPOOL_SLACK_BYTES, SCAN_WORKER_COUNT,
12 };
13 use s3_storage::ByteStream;
14 use std::path::{Path, PathBuf};
15 use tokio::fs::{File, OpenOptions};
16 use tokio::io::AsyncWriteExt;
17
18 /// Memory-map a spooled file for the byte-slice-only layers (structural,
19 /// signing_*). Mmap gives us `&[u8]` backed by the OS page cache: the
20 /// pages get demand-paged in as the parser walks them, so a 500 MB file
21 /// doesn't allocate 500 MB of RSS.
22 ///
23 /// Safety: the caller must guarantee no other process is writing to the
24 /// file while the mapping is alive. Scan-spool tempfiles are written
25 /// exclusively by the scanner before the layer runs and unlinked on
26 /// drop, so no other writer exists.
27 pub fn mmap_read(path: &Path) -> Result<memmap2::Mmap, String> {
28 let file =
29 std::fs::File::open(path).map_err(|e| format!("open spool {}: {e}", path.display()))?;
30 unsafe { memmap2::Mmap::map(&file) }.map_err(|e| format!("mmap spool {}: {e}", path.display()))
31 }
32
33 /// Owned handle to a spooled tempfile. The file is unlinked when the
34 /// handle drops, even if a layer panics mid-scan.
35 pub struct SpoolHandle {
36 path: PathBuf,
37 }
38
39 impl SpoolHandle {
40 /// On-disk path. Layers that want random access (ZIP, YARA) open this
41 /// directly. Layers that want a stream open it and pass the `File`.
42 pub fn path(&self) -> &Path {
43 &self.path
44 }
45
46 /// Open a fresh read handle at offset 0. Cheap; each consumer gets
47 /// its own cursor.
48 pub async fn reader(&self) -> std::io::Result<File> {
49 OpenOptions::new().read(true).open(&self.path).await
50 }
51 }
52
53 impl Drop for SpoolHandle {
54 fn drop(&mut self) {
55 let _ = std::fs::remove_file(&self.path);
56 }
57 }
58
59 /// Refuse the job if writing `expected_size` would leave the spool volume
60 /// with less than `SCAN_SPOOL_FREE_RESERVE_BYTES` free, or if it exceeds
61 /// the per-object cap.
62 pub fn check_free_space(spool_dir: &Path, expected_size: u64) -> Result<(), String> {
63 if expected_size > SCAN_SPOOL_MAX_BYTES {
64 return Err(format!(
65 "object exceeds scan spool cap ({expected_size} > {SCAN_SPOOL_MAX_BYTES} bytes)"
66 ));
67 }
68 let probe = if spool_dir.exists() {
69 spool_dir
70 } else {
71 spool_dir.parent().unwrap_or(Path::new("/"))
72 };
73 let free =
74 fs2::available_space(probe).map_err(|e| format!("statvfs {}: {e}", probe.display()))?;
75 // The check and the write aren't atomic: with `SCAN_WORKER_COUNT` workers, up
76 // to that many spool downloads can pass this check and then write at once.
77 // Reserve headroom for ALL of them rather than just this one, so concurrent
78 // workers can't collectively overrun the free-space floor (ultra-fuzz Run 4
79 // Perf TOCTOU). Conservative by at most `(workers-1) * expected`.
80 let concurrent_reserve = expected_size.saturating_mul(SCAN_WORKER_COUNT as u64);
81 if free.saturating_sub(concurrent_reserve) < SCAN_SPOOL_FREE_RESERVE_BYTES {
82 return Err(format!(
83 "scan spool volume short on space: free={free} expected={expected_size} workers={SCAN_WORKER_COUNT} reserve={SCAN_SPOOL_FREE_RESERVE_BYTES}"
84 ));
85 }
86 Ok(())
87 }
88
89 /// Stream the entire `ByteStream` into a fresh tempfile under `spool_dir`
90 /// and return a handle. `expected_size` is used for the pre-flight cap
91 /// check; the actual write enforces nothing beyond the cap (the upload
92 /// presign already bounded the object size).
93 pub async fn download_into_tempfile(
94 spool_dir: &Path,
95 unique: &str,
96 s3_key: &str,
97 expected_size: u64,
98 mut stream: ByteStream,
99 ) -> Result<SpoolHandle, String> {
100 check_free_space(spool_dir, expected_size)?;
101
102 tokio::fs::create_dir_all(spool_dir)
103 .await
104 .map_err(|e| format!("create spool dir {}: {e}", spool_dir.display()))?;
105
106 // Spooled uploads can be paid/private content; keep the dir owner-only so a
107 // shared spool path can't be listed by other local users (Run #2 Security
108 // MINOR). The per-file 0o600 below is the primary guard.
109 {
110 use std::os::unix::fs::PermissionsExt;
111 let _ = tokio::fs::set_permissions(spool_dir, std::fs::Permissions::from_mode(0o700)).await;
112 }
113
114 // `unique` (the scan job's UUID) makes the path collision-free even when two
115 // workers scan the same s3_key concurrently (an item and its version, or a
116 // re-queued job). Keying solely on pid+s3_key let the second `create_new`
117 // fail `AlreadyExists`, failing a legitimately-running job, and let one
118 // scan's `SpoolHandle::drop` unlink the other's live file. The s3_key suffix
119 // is retained only for human-readable debugging.
120 //
121 // Sanitize it to a bounded, safe charset: the UUID already guarantees
122 // uniqueness, so the suffix is cosmetic, but a raw s3_key could carry path
123 // separators, control chars, or exceed NAME_MAX. Map anything outside
124 // [A-Za-z0-9._-] to `_` (preserving the old `/`→`_` readability) and cap length.
125 let suffix: String = s3_key
126 .chars()
127 .map(|c| {
128 if c.is_ascii_alphanumeric() || matches!(c, '.' | '_' | '-') {
129 c
130 } else {
131 '_'
132 }
133 })
134 .take(64)
135 .collect();
136 let path = spool_dir.join(format!(
137 "scan-{}-{}-{}.tmp",
138 std::process::id(),
139 unique,
140 suffix
141 ));
142
143 // 0o600: the spooled file holds the (possibly paid/private) upload bytes for
144 // the scan window; no other local user should be able to read it (Run #2
145 // Security MINOR). `create_new` already blocks symlink pre-placement.
146 let mut file = OpenOptions::new()
147 .create_new(true)
148 .write(true)
149 .mode(0o600)
150 .open(&path)
151 .await
152 .map_err(|e| format!("open spool tempfile {}: {e}", path.display()))?;
153
154 // Hard byte ceiling independent of `expected_size`. `check_free_space`
155 // reserved disk for the *presigned* object size, but a misbehaving or
156 // abusive upstream could stream more than that; without a running cap the
157 // write loop fills the disk down to `SCAN_SPOOL_FREE_RESERVE_BYTES`. Abort
158 // past the global max and remove the partial file (the orphan reaper would
159 // catch it eventually, but don't wait).
160 // Stream into the tempfile. Run the whole copy in one fallible block so that
161 // *every* failure path, stream read error, over-cap abort, write, or
162 // flush, funnels through a single cleanup that unlinks the partial file. The
163 // earlier code only removed it on the over-cap branch, leaking a partial
164 // tempfile on a write/flush/stream error until the orphan reaper swept it.
165 // Bound the streamed bytes by the *claimed* object size plus a small slack,
166 // not just the global ceiling: an object that under-reports its size must not
167 // be allowed to stream the full `SCAN_SPOOL_MAX_BYTES` to scratch before we
168 // abort (Run 6 R6-Perf-M3). `check_free_space`/worker already cap
169 // `expected_size` at the global max, so this is always the tighter bound.
170 let write_cap = expected_size
171 .saturating_add(SCAN_SPOOL_SLACK_BYTES)
172 .min(SCAN_SPOOL_MAX_BYTES);
173 let write_result: Result<(), String> = async {
174 let mut written: u64 = 0;
175 while let Some(chunk) = stream
176 .try_next()
177 .await
178 .map_err(|e| format!("read S3 stream: {e}"))?
179 {
180 written += chunk.len() as u64;
181 if written > write_cap {
182 return Err(format!(
183 "S3 stream exceeded expected size bound ({write_cap} bytes, claimed {expected_size}); aborting scan spool"
184 ));
185 }
186 file.write_all(&chunk)
187 .await
188 .map_err(|e| format!("write spool tempfile: {e}"))?;
189 }
190 file.flush()
191 .await
192 .map_err(|e| format!("flush spool tempfile: {e}"))?;
193 Ok(())
194 }
195 .await;
196 drop(file);
197
198 if let Err(e) = write_result {
199 let _ = tokio::fs::remove_file(&path).await;
200 return Err(e);
201 }
202
203 Ok(SpoolHandle { path })
204 }
205
206 /// Reaper outcome, surfaced for logs and metrics.
207 #[derive(Debug, Default, Clone, Copy)]
208 pub struct ReaperReport {
209 pub deleted: u64,
210 pub kept: u64,
211 pub errors: u64,
212 }
213
214 /// Walk the spool directory and delete every regular file present,
215 /// intended for startup, when no live scan can own anything on disk.
216 /// Missing directory is not an error (a fresh box hasn't created it yet).
217 pub fn reap_all(spool_dir: &Path) -> ReaperReport {
218 reap_predicate(spool_dir, |_meta| true)
219 }
220
221 /// Walk the spool directory and delete regular files older than the
222 /// orphan-age threshold. Intended for the scheduler's 5-minute tick.
223 pub fn reap_orphans(spool_dir: &Path) -> ReaperReport {
224 let threshold = std::time::Duration::from_secs(SCAN_SPOOL_ORPHAN_AGE_SECS);
225 reap_predicate(spool_dir, |meta| {
226 meta.modified()
227 .ok()
228 .and_then(|m| m.elapsed().ok())
229 .is_some_and(|age| age > threshold)
230 })
231 }
232
233 fn reap_predicate<F: Fn(&std::fs::Metadata) -> bool>(
234 spool_dir: &Path,
235 should_delete: F,
236 ) -> ReaperReport {
237 let mut report = ReaperReport::default();
238 let entries = match std::fs::read_dir(spool_dir) {
239 Ok(rd) => rd,
240 Err(e) if e.kind() == std::io::ErrorKind::NotFound => return report,
241 Err(e) => {
242 tracing::warn!(dir = %spool_dir.display(), error = %e, "spool reaper: read_dir failed");
243 report.errors += 1;
244 return report;
245 }
246 };
247 for entry in entries.flatten() {
248 let path = entry.path();
249 let Ok(meta) = entry.metadata() else {
250 report.errors += 1;
251 continue;
252 };
253 if !meta.is_file() {
254 continue;
255 }
256 if !should_delete(&meta) {
257 report.kept += 1;
258 continue;
259 }
260 match std::fs::remove_file(&path) {
261 Ok(()) => {
262 report.deleted += 1;
263 tracing::info!(path = %path.display(), "spool reaper: deleted orphan");
264 }
265 Err(e) => {
266 report.errors += 1;
267 tracing::warn!(path = %path.display(), error = %e, "spool reaper: delete failed");
268 }
269 }
270 }
271 report
272 }
273
274 #[cfg(test)]
275 mod tests {
276 use super::*;
277
278 #[test]
279 fn reap_all_deletes_present_files() {
280 let dir = tempfile::tempdir().unwrap();
281 std::fs::write(dir.path().join("scan-1.tmp"), b"x").unwrap();
282 std::fs::write(dir.path().join("scan-2.tmp"), b"y").unwrap();
283 let report = reap_all(dir.path());
284 assert_eq!(report.deleted, 2);
285 assert_eq!(report.errors, 0);
286 assert_eq!(std::fs::read_dir(dir.path()).unwrap().count(), 0);
287 }
288
289 #[test]
290 fn reap_orphans_keeps_recent_files() {
291 let dir = tempfile::tempdir().unwrap();
292 std::fs::write(dir.path().join("scan-1.tmp"), b"x").unwrap();
293 let report = reap_orphans(dir.path());
294 assert_eq!(report.deleted, 0);
295 assert_eq!(report.kept, 1);
296 }
297
298 #[test]
299 fn reap_handles_missing_directory() {
300 let report = reap_all(std::path::Path::new("/nonexistent/spool/dir/xyz123"));
301 assert_eq!(report.deleted, 0);
302 assert_eq!(report.errors, 0);
303 }
304
305 #[test]
306 fn spool_handle_drops_file() {
307 let dir = tempfile::tempdir().unwrap();
308 let path = dir.path().join("scan-drop.tmp");
309 std::fs::write(&path, b"data").unwrap();
310 let handle = SpoolHandle { path: path.clone() };
311 assert!(path.exists());
312 drop(handle);
313 assert!(!path.exists());
314 }
315 }
316