Skip to main content

max / makenotwork

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