Skip to main content

max / makenotwork

8.1 KB · 243 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::{SCAN_SPOOL_FREE_RESERVE_BYTES, SCAN_SPOOL_MAX_BYTES, SCAN_SPOOL_ORPHAN_AGE_SECS};
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 = std::fs::File::open(path)
29 .map_err(|e| format!("open spool {}: {e}", path.display()))?;
30 unsafe { memmap2::Mmap::map(&file) }
31 .map_err(|e| format!("mmap spool {}: {e}", path.display()))
32 }
33
34 /// Owned handle to a spooled tempfile. The file is unlinked when the
35 /// handle drops, even if a layer panics mid-scan.
36 pub struct SpoolHandle {
37 path: PathBuf,
38 }
39
40 impl SpoolHandle {
41 /// On-disk path. Layers that want random access (ZIP, YARA) open this
42 /// directly. Layers that want a stream open it and pass the `File`.
43 pub fn path(&self) -> &Path {
44 &self.path
45 }
46
47 /// Open a fresh read handle at offset 0. Cheap; each consumer gets
48 /// its own cursor.
49 pub async fn reader(&self) -> std::io::Result<File> {
50 OpenOptions::new().read(true).open(&self.path).await
51 }
52 }
53
54 impl Drop for SpoolHandle {
55 fn drop(&mut self) {
56 let _ = std::fs::remove_file(&self.path);
57 }
58 }
59
60 /// Refuse the job if writing `expected_size` would leave the spool volume
61 /// with less than `SCAN_SPOOL_FREE_RESERVE_BYTES` free, or if it exceeds
62 /// the per-object cap.
63 pub fn check_free_space(spool_dir: &Path, expected_size: u64) -> Result<(), String> {
64 if expected_size > SCAN_SPOOL_MAX_BYTES {
65 return Err(format!(
66 "object exceeds scan spool cap ({} > {} bytes)",
67 expected_size, SCAN_SPOOL_MAX_BYTES
68 ));
69 }
70 let probe = if spool_dir.exists() {
71 spool_dir
72 } else {
73 spool_dir.parent().unwrap_or(Path::new("/"))
74 };
75 let free = fs2::available_space(probe)
76 .map_err(|e| format!("statvfs {}: {e}", probe.display()))?;
77 if free.saturating_sub(expected_size) < SCAN_SPOOL_FREE_RESERVE_BYTES {
78 return Err(format!(
79 "scan spool volume short on space: free={} expected={} reserve={}",
80 free, expected_size, SCAN_SPOOL_FREE_RESERVE_BYTES
81 ));
82 }
83 Ok(())
84 }
85
86 /// Stream the entire `ByteStream` into a fresh tempfile under `spool_dir`
87 /// and return a handle. `expected_size` is used for the pre-flight cap
88 /// check; the actual write enforces nothing beyond the cap (the upload
89 /// presign already bounded the object size).
90 pub async fn download_into_tempfile(
91 spool_dir: &Path,
92 s3_key: &str,
93 expected_size: u64,
94 mut stream: ByteStream,
95 ) -> Result<SpoolHandle, String> {
96 check_free_space(spool_dir, expected_size)?;
97
98 tokio::fs::create_dir_all(spool_dir)
99 .await
100 .map_err(|e| format!("create spool dir {}: {e}", spool_dir.display()))?;
101
102 let suffix = s3_key.replace('/', "_");
103 let path = spool_dir.join(format!("scan-{}-{}.tmp", std::process::id(), suffix));
104
105 let mut file = OpenOptions::new()
106 .create_new(true)
107 .write(true)
108 .open(&path)
109 .await
110 .map_err(|e| format!("open spool tempfile {}: {e}", path.display()))?;
111
112 while let Some(chunk) = stream
113 .try_next()
114 .await
115 .map_err(|e| format!("read S3 stream: {e}"))?
116 {
117 file.write_all(&chunk)
118 .await
119 .map_err(|e| format!("write spool tempfile: {e}"))?;
120 }
121 file.flush()
122 .await
123 .map_err(|e| format!("flush spool tempfile: {e}"))?;
124 drop(file);
125
126 Ok(SpoolHandle { path })
127 }
128
129 /// Reaper outcome — surfaced for logs and metrics.
130 #[derive(Debug, Default, Clone, Copy)]
131 pub struct ReaperReport {
132 pub deleted: u64,
133 pub kept: u64,
134 pub errors: u64,
135 }
136
137 /// Walk the spool directory and delete every regular file present —
138 /// intended for startup, when no live scan can own anything on disk.
139 /// Missing directory is not an error (a fresh box hasn't created it yet).
140 pub fn reap_all(spool_dir: &Path) -> ReaperReport {
141 reap_predicate(spool_dir, |_meta| true)
142 }
143
144 /// Walk the spool directory and delete regular files older than the
145 /// orphan-age threshold. Intended for the scheduler's 5-minute tick.
146 pub fn reap_orphans(spool_dir: &Path) -> ReaperReport {
147 let threshold = std::time::Duration::from_secs(SCAN_SPOOL_ORPHAN_AGE_SECS);
148 reap_predicate(spool_dir, |meta| {
149 meta.modified()
150 .ok()
151 .and_then(|m| m.elapsed().ok())
152 .map(|age| age > threshold)
153 .unwrap_or(false)
154 })
155 }
156
157 fn reap_predicate<F: Fn(&std::fs::Metadata) -> bool>(
158 spool_dir: &Path,
159 should_delete: F,
160 ) -> ReaperReport {
161 let mut report = ReaperReport::default();
162 let entries = match std::fs::read_dir(spool_dir) {
163 Ok(rd) => rd,
164 Err(e) if e.kind() == std::io::ErrorKind::NotFound => return report,
165 Err(e) => {
166 tracing::warn!(dir = %spool_dir.display(), error = %e, "spool reaper: read_dir failed");
167 report.errors += 1;
168 return report;
169 }
170 };
171 for entry in entries.flatten() {
172 let path = entry.path();
173 let meta = match entry.metadata() {
174 Ok(m) => m,
175 Err(_) => {
176 report.errors += 1;
177 continue;
178 }
179 };
180 if !meta.is_file() {
181 continue;
182 }
183 if !should_delete(&meta) {
184 report.kept += 1;
185 continue;
186 }
187 match std::fs::remove_file(&path) {
188 Ok(()) => {
189 report.deleted += 1;
190 tracing::info!(path = %path.display(), "spool reaper: deleted orphan");
191 }
192 Err(e) => {
193 report.errors += 1;
194 tracing::warn!(path = %path.display(), error = %e, "spool reaper: delete failed");
195 }
196 }
197 }
198 report
199 }
200
201 #[cfg(test)]
202 mod tests {
203 use super::*;
204
205 #[test]
206 fn reap_all_deletes_present_files() {
207 let dir = tempfile::tempdir().unwrap();
208 std::fs::write(dir.path().join("scan-1.tmp"), b"x").unwrap();
209 std::fs::write(dir.path().join("scan-2.tmp"), b"y").unwrap();
210 let report = reap_all(dir.path());
211 assert_eq!(report.deleted, 2);
212 assert_eq!(report.errors, 0);
213 assert_eq!(std::fs::read_dir(dir.path()).unwrap().count(), 0);
214 }
215
216 #[test]
217 fn reap_orphans_keeps_recent_files() {
218 let dir = tempfile::tempdir().unwrap();
219 std::fs::write(dir.path().join("scan-1.tmp"), b"x").unwrap();
220 let report = reap_orphans(dir.path());
221 assert_eq!(report.deleted, 0);
222 assert_eq!(report.kept, 1);
223 }
224
225 #[test]
226 fn reap_handles_missing_directory() {
227 let report = reap_all(std::path::Path::new("/nonexistent/spool/dir/xyz123"));
228 assert_eq!(report.deleted, 0);
229 assert_eq!(report.errors, 0);
230 }
231
232 #[test]
233 fn spool_handle_drops_file() {
234 let dir = tempfile::tempdir().unwrap();
235 let path = dir.path().join("scan-drop.tmp");
236 std::fs::write(&path, b"data").unwrap();
237 let handle = SpoolHandle { path: path.clone() };
238 assert!(path.exists());
239 drop(handle);
240 assert!(!path.exists());
241 }
242 }
243