Skip to main content

max / makenotwork

10.8 KB · 348 lines History Blame Raw
1 //! Async scan-job queue for the malware pipeline.
2 //!
3 //! Upload routes call [`enqueue`] to register a scan job and return their
4 //! request to the client. A pool of scan workers (`crate::scanning::worker`)
5 //! drains the queue with `FOR UPDATE SKIP LOCKED` via [`claim_next`], runs the
6 //! pipeline against the S3 object, and finalizes the job with [`mark_done`] or
7 //! [`mark_failed`].
8 //!
9 //! See `docs/scan-pipeline-audit.md` for the wider architecture.
10
11 use chrono::{DateTime, Utc};
12 use sqlx::{FromRow, PgPool};
13 use uuid::Uuid;
14
15 use crate::storage::FileType;
16
17 use super::UserId;
18
19 /// The entity whose `scan_status` the worker should update when the scan
20 /// completes. `Item`, `Version`, and `Media` have `scan_status` columns;
21 /// `ProjectImage` and `ContentInsertion` do not — for those, the worker
22 /// still scans (recording results in `file_scan_results`) but only acts on
23 /// `Quarantined` by creating a WAM ticket for admin follow-up.
24 #[derive(Debug, Clone, Copy, PartialEq, Eq)]
25 pub enum ScanTargetKind {
26 Item,
27 Version,
28 Media,
29 ProjectImage,
30 ItemImage,
31 ContentInsertion,
32 }
33
34 impl ScanTargetKind {
35 pub fn as_str(&self) -> &'static str {
36 match self {
37 ScanTargetKind::Item => "item",
38 ScanTargetKind::Version => "version",
39 ScanTargetKind::Media => "media",
40 ScanTargetKind::ProjectImage => "project_image",
41 ScanTargetKind::ItemImage => "item_image",
42 ScanTargetKind::ContentInsertion => "content_insertion",
43 }
44 }
45
46 pub fn from_str(s: &str) -> Option<Self> {
47 Some(match s {
48 "item" => ScanTargetKind::Item,
49 "version" => ScanTargetKind::Version,
50 "media" => ScanTargetKind::Media,
51 "project_image" => ScanTargetKind::ProjectImage,
52 "item_image" => ScanTargetKind::ItemImage,
53 "content_insertion" => ScanTargetKind::ContentInsertion,
54 _ => return None,
55 })
56 }
57 }
58
59 /// A queued or running scan job, as claimed by a worker. Most fields are
60 /// populated via `FromRow` from sqlx; fields not consumed by the worker
61 /// today are kept for the admin dashboard (Phase 2 of the audit).
62 #[allow(dead_code)]
63 #[derive(Debug, Clone, FromRow)]
64 pub struct ScanJob {
65 pub id: Uuid,
66 pub target_kind: String,
67 pub target_id: Uuid,
68 pub s3_key: String,
69 pub file_type: String,
70 pub user_id: UserId,
71 pub file_size_bytes: i64,
72 pub status: String,
73 pub attempts: i32,
74 pub enqueued_at: DateTime<Utc>,
75 pub started_at: Option<DateTime<Utc>>,
76 pub completed_at: Option<DateTime<Utc>>,
77 pub last_error: Option<String>,
78 }
79
80 impl ScanJob {
81 pub fn typed_kind(&self) -> Option<ScanTargetKind> {
82 ScanTargetKind::from_str(&self.target_kind)
83 }
84
85 pub fn typed_file_type(&self) -> Option<FileType> {
86 self.file_type.parse().ok()
87 }
88 }
89
90 /// Enqueue a scan job. Returns the job id.
91 #[tracing::instrument(skip_all, fields(target_kind = target_kind.as_str(), %target_id, s3_key))]
92 pub async fn enqueue(
93 db: &PgPool,
94 target_kind: ScanTargetKind,
95 target_id: Uuid,
96 s3_key: &str,
97 file_type: FileType,
98 user_id: UserId,
99 file_size_bytes: i64,
100 ) -> Result<Uuid, sqlx::Error> {
101 let id = sqlx::query_scalar::<_, Uuid>(
102 r#"
103 INSERT INTO scan_jobs
104 (target_kind, target_id, s3_key, file_type, user_id, file_size_bytes)
105 VALUES ($1, $2, $3, $4, $5, $6)
106 RETURNING id
107 "#,
108 )
109 .bind(target_kind.as_str())
110 .bind(target_id)
111 .bind(s3_key)
112 .bind(file_type.as_str())
113 .bind(user_id)
114 .bind(file_size_bytes)
115 .fetch_one(db)
116 .await?;
117
118 Ok(id)
119 }
120
121 /// Atomically claim the next queued scan job for processing.
122 ///
123 /// Uses `FOR UPDATE SKIP LOCKED` so multiple workers can drain the queue in
124 /// parallel without contention. Sets `status='running'`, increments `attempts`,
125 /// and stamps `started_at`. Returns `Ok(None)` if the queue is empty.
126 #[tracing::instrument(skip_all)]
127 pub async fn claim_next(db: &PgPool) -> Result<Option<ScanJob>, sqlx::Error> {
128 let job = sqlx::query_as::<_, ScanJob>(
129 r#"
130 WITH next AS (
131 SELECT id FROM scan_jobs
132 WHERE status = 'queued'
133 ORDER BY enqueued_at ASC
134 FOR UPDATE SKIP LOCKED
135 LIMIT 1
136 )
137 UPDATE scan_jobs
138 SET status = 'running',
139 attempts = attempts + 1,
140 started_at = NOW()
141 WHERE id = (SELECT id FROM next)
142 RETURNING *
143 "#,
144 )
145 .fetch_optional(db)
146 .await?;
147
148 Ok(job)
149 }
150
151 /// Mark a job as completed successfully.
152 #[tracing::instrument(skip_all, fields(%job_id))]
153 pub async fn mark_done(db: &PgPool, job_id: Uuid) -> Result<(), sqlx::Error> {
154 sqlx::query(
155 r#"
156 UPDATE scan_jobs
157 SET status = 'done', completed_at = NOW(), last_error = NULL
158 WHERE id = $1
159 "#,
160 )
161 .bind(job_id)
162 .execute(db)
163 .await?;
164 Ok(())
165 }
166
167 /// Mark a job as failed (worker exception, S3 fetch failure, etc.).
168 ///
169 /// Records `last_error` for admin inspection. The job stays in `failed`
170 /// status; an admin or operator can manually re-enqueue via the dashboard
171 /// (Phase 2) by inserting a fresh row.
172 #[tracing::instrument(skip_all, fields(%job_id))]
173 pub async fn mark_failed(db: &PgPool, job_id: Uuid, err: &str) -> Result<(), sqlx::Error> {
174 sqlx::query(
175 r#"
176 UPDATE scan_jobs
177 SET status = 'failed', completed_at = NOW(), last_error = $1
178 WHERE id = $2
179 "#,
180 )
181 .bind(err)
182 .bind(job_id)
183 .execute(db)
184 .await?;
185 Ok(())
186 }
187
188 /// Reset jobs that have been stuck in `running` longer than `max_age_secs`.
189 ///
190 /// Run on worker startup to recover from a previous-process crash mid-scan:
191 /// the row would otherwise stay `running` forever and never be re-claimed.
192 /// Increments `attempts` so a perpetually-crashing scan eventually trips the
193 /// retry budget.
194 #[tracing::instrument(skip_all)]
195 pub async fn reap_stuck(db: &PgPool, max_age_secs: i64) -> Result<u64, sqlx::Error> {
196 let affected = sqlx::query(
197 r#"
198 UPDATE scan_jobs
199 SET status = 'queued', started_at = NULL
200 WHERE status = 'running'
201 AND started_at < NOW() - ($1 || ' seconds')::interval
202 "#,
203 )
204 .bind(max_age_secs.to_string())
205 .execute(db)
206 .await?
207 .rows_affected();
208 Ok(affected)
209 }
210
211 /// Delete terminal-state rows older than `older_than`. Returns the count.
212 ///
213 /// Only touches `done`/`failed` rows — operational state (`queued`,
214 /// `running`) is owned by the worker loop and `reap_stuck`. The verdict
215 /// (Clean / Quarantined / HeldForReview) lives on the entity's
216 /// `scan_status` column, not here, so dropping a `done` row loses queue
217 /// history only, not malware-detection state.
218 ///
219 /// No supporting index today: at soft-launch volume Postgres seq-scans the
220 /// table fine. Revisit once `EXPLAIN ANALYZE` shows it as a bottleneck.
221 #[tracing::instrument(skip_all)]
222 pub async fn purge_old_terminal(
223 db: &PgPool,
224 older_than: chrono::Duration,
225 ) -> Result<u64, sqlx::Error> {
226 let cutoff = chrono::Utc::now() - older_than;
227 let n = sqlx::query(
228 r#"
229 DELETE FROM scan_jobs
230 WHERE status IN ('done', 'failed')
231 AND COALESCE(completed_at, started_at, enqueued_at) < $1
232 "#,
233 )
234 .bind(cutoff)
235 .execute(db)
236 .await?
237 .rows_affected();
238 Ok(n)
239 }
240
241 /// Count of currently-queued jobs. Used by the admin dashboard health panel
242 /// (Phase 2 of the audit). Allowed dead code until that route lands.
243 #[allow(dead_code)]
244 pub async fn queued_count(db: &PgPool) -> Result<i64, sqlx::Error> {
245 sqlx::query_scalar::<_, i64>("SELECT COUNT(*) FROM scan_jobs WHERE status = 'queued'")
246 .fetch_one(db)
247 .await
248 }
249
250 /// Count of currently-running jobs. Phase 2 dashboard consumer.
251 #[allow(dead_code)]
252 pub async fn running_count(db: &PgPool) -> Result<i64, sqlx::Error> {
253 sqlx::query_scalar::<_, i64>("SELECT COUNT(*) FROM scan_jobs WHERE status = 'running'")
254 .fetch_one(db)
255 .await
256 }
257
258 /// Count of running jobs that have been in flight longer than `max_age_secs`.
259 /// Used by PoM to alert on stuck workers.
260 pub async fn stuck_count(db: &PgPool, max_age_secs: i64) -> Result<i64, sqlx::Error> {
261 sqlx::query_scalar::<_, i64>(
262 "SELECT COUNT(*) FROM scan_jobs WHERE status = 'running' AND started_at < NOW() - ($1 || ' seconds')::interval",
263 )
264 .bind(max_age_secs.to_string())
265 .fetch_one(db)
266 .await
267 }
268
269 /// A held version with enough context to re-enqueue it for scanning.
270 #[allow(dead_code)]
271 #[derive(Debug, Clone, FromRow)]
272 pub struct RescanCandidateVersion {
273 pub version_id: Uuid,
274 pub s3_key: String,
275 pub file_size_bytes: i64,
276 pub user_id: UserId,
277 }
278
279 /// A held item (audio or cover) with re-enqueue context.
280 #[allow(dead_code)]
281 #[derive(Debug, Clone, FromRow)]
282 pub struct RescanCandidateItem {
283 pub item_id: Uuid,
284 pub s3_key: String,
285 pub file_size_bytes: i64,
286 pub user_id: UserId,
287 /// Audio / cover, selected by the query.
288 pub file_type: String,
289 }
290
291 /// Find currently-held versions with enough context to be re-scanned.
292 pub async fn rescan_candidates_versions(db: &PgPool) -> Result<Vec<RescanCandidateVersion>, sqlx::Error> {
293 sqlx::query_as::<_, RescanCandidateVersion>(
294 r#"
295 SELECT v.id AS version_id,
296 v.s3_key,
297 COALESCE(v.file_size_bytes, 0) AS file_size_bytes,
298 p.user_id
299 FROM versions v
300 JOIN items i ON i.id = v.item_id
301 JOIN projects p ON p.id = i.project_id
302 WHERE v.scan_status = 'held_for_review'
303 AND v.s3_key IS NOT NULL
304 "#,
305 )
306 .fetch_all(db)
307 .await
308 }
309
310 /// Find currently-held items (audio or cover) with re-enqueue context.
311 pub async fn rescan_candidates_items(db: &PgPool) -> Result<Vec<RescanCandidateItem>, sqlx::Error> {
312 sqlx::query_as::<_, RescanCandidateItem>(
313 r#"
314 SELECT i.id AS item_id,
315 COALESCE(i.audio_s3_key, i.cover_s3_key) AS s3_key,
316 COALESCE(i.audio_file_size_bytes, i.cover_file_size_bytes, 0) AS file_size_bytes,
317 p.user_id,
318 CASE WHEN i.audio_s3_key IS NOT NULL THEN 'audio' ELSE 'cover' END AS file_type
319 FROM items i
320 JOIN projects p ON p.id = i.project_id
321 WHERE i.scan_status = 'held_for_review'
322 AND COALESCE(i.audio_s3_key, i.cover_s3_key) IS NOT NULL
323 "#,
324 )
325 .fetch_all(db)
326 .await
327 }
328
329 #[cfg(test)]
330 mod tests {
331 use super::*;
332
333 #[test]
334 fn target_kind_round_trip() {
335 for kind in [
336 ScanTargetKind::Item,
337 ScanTargetKind::Version,
338 ScanTargetKind::Media,
339 ScanTargetKind::ProjectImage,
340 ScanTargetKind::ItemImage,
341 ScanTargetKind::ContentInsertion,
342 ] {
343 assert_eq!(ScanTargetKind::from_str(kind.as_str()), Some(kind));
344 }
345 assert_eq!(ScanTargetKind::from_str("bogus"), None);
346 }
347 }
348