Skip to main content

max / makenotwork

20.6 KB · 554 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 GalleryImage,
32 ContentInsertion,
33 OtaArtifact,
34 }
35
36 impl ScanTargetKind {
37 pub fn as_str(&self) -> &'static str {
38 match self {
39 ScanTargetKind::Item => "item",
40 ScanTargetKind::Version => "version",
41 ScanTargetKind::Media => "media",
42 ScanTargetKind::ProjectImage => "project_image",
43 ScanTargetKind::ItemImage => "item_image",
44 ScanTargetKind::GalleryImage => "gallery_image",
45 ScanTargetKind::ContentInsertion => "content_insertion",
46 ScanTargetKind::OtaArtifact => "ota_artifact",
47 }
48 }
49
50 /// Decode the DB `target_kind` text. Named `from_db_str` (not `from_str`) so
51 /// it doesn't shadow the `FromStr` trait method, it returns `Option`, not the
52 /// `Result` that trait requires, and pairs with `as_str` above.
53 pub fn from_db_str(s: &str) -> Option<Self> {
54 Some(match s {
55 "item" => ScanTargetKind::Item,
56 "version" => ScanTargetKind::Version,
57 "media" => ScanTargetKind::Media,
58 "project_image" => ScanTargetKind::ProjectImage,
59 "item_image" => ScanTargetKind::ItemImage,
60 "gallery_image" => ScanTargetKind::GalleryImage,
61 "content_insertion" => ScanTargetKind::ContentInsertion,
62 "ota_artifact" => ScanTargetKind::OtaArtifact,
63 _ => return None,
64 })
65 }
66
67 /// The S3 bucket this kind's objects live in. Single source of truth for
68 /// the scan worker's client selection, quarantine delete, and durable-delete
69 /// enqueue, all three MUST agree, or a quarantine deletes from the wrong
70 /// bucket and idempotently "succeeds" while the malware persists (ultra-fuzz
71 /// Run #24 Storage CRITICAL). OTA artifacts are the only synckit-bucket kind
72 /// today; this mirrors the `S3_KEY_REFS` registry in `pending_s3_deletions`.
73 pub fn storage_bucket(&self) -> crate::storage::S3Bucket {
74 match self {
75 ScanTargetKind::OtaArtifact => crate::storage::S3Bucket::Synckit,
76 _ => crate::storage::S3Bucket::Main,
77 }
78 }
79
80 /// Whether this kind carries no `scan_status` column and is rendered
81 /// straight from `cdn.makenot.work/{key}`, so a `Quarantined` verdict can
82 /// only be enforced by deleting the *DB row* (which stops the URL from
83 /// rendering and makes the key non-live for the deletion queue).
84 ///
85 /// This governs the DB-row step only. The S3-object purge on quarantine is
86 /// unconditional (see [`quarantine_purges_object`]); these kinds
87 /// additionally need their row removed (or, for `ItemImage`, the cover
88 /// columns NULLed) because they have no status column to gate on.
89 ///
90 /// `ItemImage` is here because the item *cover* (`items.cover_image_url`) is
91 /// rendered straight from the CDN with no per-request gate, `items.scan_status`
92 /// gates only the audio/video. On quarantine the purge NULLs the cover columns
93 /// (keeping the track), it does NOT delete the item row.
94 pub fn is_cdn_served_without_gate(&self) -> bool {
95 matches!(
96 self,
97 ScanTargetKind::ItemImage
98 | ScanTargetKind::ProjectImage
99 | ScanTargetKind::GalleryImage
100 | ScanTargetKind::ContentInsertion
101 )
102 }
103
104 /// Whether this kind's *content* object (post-promote) is served UNSIGNED
105 /// from the public CDN bucket, and so must be promoted cross-bucket into
106 /// [`crate::storage::S3Bucket::Public`]. This is the three image cover/gallery
107 /// kinds, the `is_cdn_served_without_gate` set MINUS `ContentInsertion`,
108 /// which, despite carrying no per-request scan gate, is served *presigned*
109 /// from the private bucket (see `promote_cdn_image_by_key`: it alone has no
110 /// materialized public URL). The matching `S3_KEY_REFS` entries
111 /// (`item_images`, `project_images`, `items.cover_s3_key`,
112 /// `projects.cover_s3_key`) therefore live under bucket `public`, while
113 /// `content_insertions.storage_key` stays `main`. Staging is ALWAYS private
114 /// ([`storage_bucket`] is unchanged), only the promoted content object moves.
115 pub fn content_served_from_public_bucket(&self) -> bool {
116 matches!(
117 self,
118 ScanTargetKind::ItemImage | ScanTargetKind::ProjectImage | ScanTargetKind::GalleryImage
119 )
120 }
121
122 /// Whether a `Quarantined` verdict purges the underlying S3 object.
123 ///
124 /// Always true: a confirmed-malicious object has no reason to remain in
125 /// storage. Crucially, the per-request `scan_status` gate on
126 /// `Item`/`Version`/`Media` is NOT sufficient on its own, free downloadable
127 /// content is handed out as a *permanent* `cdn.makenot.work/{key}` URL with
128 /// no per-request gate (see `routes/storage/downloads.rs::resolve_content_url`),
129 /// so a leaked or edge-cached URL keeps serving the malware from origin until
130 /// the object itself is gone. Removing the object closes that hole for
131 /// downloadable content the same way it does for the gate-less image kinds.
132 /// The decision lives here, beside `is_cdn_served_without_gate`, so the two
133 /// halves of quarantine enforcement can't silently diverge per kind.
134 pub fn quarantine_purges_object(&self) -> bool {
135 true
136 }
137 }
138
139 /// A queued or running scan job, as claimed by a worker. Most fields are
140 /// populated via `FromRow` from sqlx; fields not consumed by the worker
141 /// today are kept for the admin dashboard (Phase 2 of the audit).
142 #[allow(dead_code)]
143 #[derive(Debug, Clone, FromRow)]
144 pub struct ScanJob {
145 pub id: Uuid,
146 pub target_kind: String,
147 pub target_id: Uuid,
148 pub s3_key: String,
149 pub file_type: String,
150 pub user_id: UserId,
151 pub file_size_bytes: i64,
152 pub status: String,
153 pub attempts: i32,
154 pub enqueued_at: DateTime<Utc>,
155 pub started_at: Option<DateTime<Utc>>,
156 /// Last liveness beat from the worker running this job. Bumped on a cadence
157 /// well under STUCK_JOB_SECS; the reaper keys off this (falling back to
158 /// `started_at`) so a slow-but-progressing scan isn't mistaken for a
159 /// crashed worker. See `bump_heartbeat` and migration 167.
160 pub heartbeat_at: Option<DateTime<Utc>>,
161 pub completed_at: Option<DateTime<Utc>>,
162 pub last_error: Option<String>,
163 }
164
165 impl ScanJob {
166 pub fn typed_kind(&self) -> Option<ScanTargetKind> {
167 ScanTargetKind::from_db_str(&self.target_kind)
168 }
169
170 pub fn typed_file_type(&self) -> Option<FileType> {
171 self.file_type.parse().ok()
172 }
173 }
174
175 /// Enqueue a scan job. Returns the job id.
176 #[tracing::instrument(skip_all, fields(target_kind = target_kind.as_str(), %target_id, s3_key))]
177 pub async fn enqueue(
178 db: &PgPool,
179 target_kind: ScanTargetKind,
180 target_id: Uuid,
181 s3_key: &str,
182 file_type: FileType,
183 user_id: UserId,
184 file_size_bytes: i64,
185 ) -> Result<Uuid, sqlx::Error> {
186 let id = sqlx::query_scalar::<_, Uuid>(
187 r"
188 INSERT INTO scan_jobs
189 (target_kind, target_id, s3_key, file_type, user_id, file_size_bytes)
190 VALUES ($1, $2, $3, $4, $5, $6)
191 RETURNING id
192 ",
193 )
194 .bind(target_kind.as_str())
195 .bind(target_id)
196 .bind(s3_key)
197 .bind(file_type.as_str())
198 .bind(user_id)
199 .bind(file_size_bytes)
200 .fetch_one(db)
201 .await?;
202
203 Ok(id)
204 }
205
206 /// Maximum times a scan job may be claimed before it is given up as `failed`.
207 /// Bounds the crash-loop retry budget: a job that reliably wedges the OS process
208 /// (so it's reaped from `running` rather than reaching `mark_failed`) can only
209 /// be re-attempted this many times before `reap_stuck` retires it.
210 pub const MAX_SCAN_ATTEMPTS: i32 = 5;
211
212 /// Atomically claim the next queued scan job for processing.
213 ///
214 /// Uses `FOR UPDATE SKIP LOCKED` so multiple workers can drain the queue in
215 /// parallel without contention. Sets `status='running'`, increments `attempts`,
216 /// and stamps `started_at`. Skips jobs that have already hit `MAX_SCAN_ATTEMPTS`
217 /// (those are retired to `failed` by `reap_stuck`). Returns `Ok(None)` if the
218 /// queue is empty.
219 #[tracing::instrument(skip_all)]
220 pub async fn claim_next(db: &PgPool) -> Result<Option<ScanJob>, sqlx::Error> {
221 let job = sqlx::query_as::<_, ScanJob>(
222 r"
223 WITH next AS (
224 SELECT id FROM scan_jobs
225 WHERE status = 'queued' AND attempts < $1
226 ORDER BY enqueued_at ASC
227 FOR UPDATE SKIP LOCKED
228 LIMIT 1
229 )
230 UPDATE scan_jobs
231 SET status = 'running',
232 attempts = attempts + 1,
233 started_at = NOW(),
234 heartbeat_at = NOW()
235 WHERE id = (SELECT id FROM next)
236 RETURNING *
237 ",
238 )
239 .bind(MAX_SCAN_ATTEMPTS)
240 .fetch_optional(db)
241 .await?;
242
243 Ok(job)
244 }
245
246 /// Mark a job as completed successfully.
247 #[tracing::instrument(skip_all, fields(%job_id))]
248 pub async fn mark_done(db: &PgPool, job_id: Uuid) -> Result<(), sqlx::Error> {
249 sqlx::query(
250 r"
251 UPDATE scan_jobs
252 SET status = 'done', completed_at = NOW(), last_error = NULL
253 WHERE id = $1
254 ",
255 )
256 .bind(job_id)
257 .execute(db)
258 .await?;
259 Ok(())
260 }
261
262 /// Mark a job as failed (worker exception, S3 fetch failure, etc.).
263 ///
264 /// Records `last_error` for admin inspection. The job stays in `failed`
265 /// status; an admin or operator can manually re-enqueue via the dashboard
266 /// (Phase 2) by inserting a fresh row.
267 #[tracing::instrument(skip_all, fields(%job_id))]
268 pub async fn mark_failed(db: &PgPool, job_id: Uuid, err: &str) -> Result<(), sqlx::Error> {
269 sqlx::query(
270 r"
271 UPDATE scan_jobs
272 SET status = 'failed', completed_at = NOW(), last_error = $1
273 WHERE id = $2
274 ",
275 )
276 .bind(err)
277 .bind(job_id)
278 .execute(db)
279 .await?;
280 Ok(())
281 }
282
283 /// Refresh a running job's liveness heartbeat.
284 ///
285 /// Called periodically by the worker that owns the job so `reap_stuck` (and the
286 /// PoM stuck-count) can distinguish a slow-but-progressing scan (fresh beat)
287 /// from a crashed or hung worker (stale beat). The `status = 'running'` guard
288 /// makes a late beat a no-op once the job has already finished or been reaped,
289 /// it never resurrects a terminal row's timestamp.
290 #[tracing::instrument(skip_all, fields(%job_id))]
291 pub async fn bump_heartbeat(db: &PgPool, job_id: Uuid) -> Result<(), sqlx::Error> {
292 sqlx::query("UPDATE scan_jobs SET heartbeat_at = NOW() WHERE id = $1 AND status = 'running'")
293 .bind(job_id)
294 .execute(db)
295 .await?;
296 Ok(())
297 }
298
299 /// Reset jobs whose worker has gone silent for longer than `max_age_secs`.
300 ///
301 /// Run on worker startup and on a timer to recover from a previous-process crash
302 /// mid-scan: the row would otherwise stay `running` forever and never be
303 /// re-claimed. Liveness is measured from the job's last heartbeat
304 /// (`COALESCE(heartbeat_at, started_at)`), NOT from `started_at` alone, a
305 /// legitimately slow scan (large object over a slow S3 link) can run past
306 /// `max_age_secs` while still making progress, and reaping *that* let a second
307 /// worker double-process the same object and inflate `attempts` toward
308 /// MAX_SCAN_ATTEMPTS until a valid file was force-retired (audit Run 22). The
309 /// running worker bumps `heartbeat_at` on a cadence well under `max_age_secs`,
310 /// so only a crashed/hung process (no fresh beat) crosses the threshold. A job
311 /// whose `attempts` has already reached `MAX_SCAN_ATTEMPTS` is retired to
312 /// `failed`; otherwise it returns to `queued` for another attempt.
313 #[tracing::instrument(skip_all)]
314 pub async fn reap_stuck(db: &PgPool, max_age_secs: i64) -> Result<u64, sqlx::Error> {
315 let affected = sqlx::query(
316 r"
317 UPDATE scan_jobs
318 SET status = CASE WHEN attempts >= $2 THEN 'failed' ELSE 'queued' END,
319 started_at = NULL,
320 heartbeat_at = NULL,
321 completed_at = CASE WHEN attempts >= $2 THEN NOW() ELSE completed_at END,
322 last_error = CASE WHEN attempts >= $2
323 THEN 'exceeded max scan attempts (reaped from running)'
324 ELSE last_error END
325 WHERE status = 'running'
326 AND COALESCE(heartbeat_at, started_at) < NOW() - ($1 || ' seconds')::interval
327 ",
328 )
329 .bind(max_age_secs.to_string())
330 .bind(MAX_SCAN_ATTEMPTS)
331 .execute(db)
332 .await?
333 .rows_affected();
334 Ok(affected)
335 }
336
337 /// Delete terminal-state rows older than `older_than`. Returns the count.
338 ///
339 /// Only touches `done`/`failed` rows, operational state (`queued`,
340 /// `running`) is owned by the worker loop and `reap_stuck`. The verdict
341 /// (Clean / Quarantined / HeldForReview) lives on the entity's
342 /// `scan_status` column, not here, so dropping a `done` row loses queue
343 /// history only, not malware-detection state.
344 ///
345 /// No supporting index today: at soft-launch volume Postgres seq-scans the
346 /// table fine. Revisit once `EXPLAIN ANALYZE` shows it as a bottleneck.
347 #[tracing::instrument(skip_all)]
348 pub async fn purge_old_terminal(
349 db: &PgPool,
350 older_than: chrono::Duration,
351 ) -> Result<u64, sqlx::Error> {
352 let cutoff = chrono::Utc::now() - older_than;
353 let n = sqlx::query(
354 r"
355 DELETE FROM scan_jobs
356 WHERE status IN ('done', 'failed')
357 AND COALESCE(completed_at, started_at, enqueued_at) < $1
358 ",
359 )
360 .bind(cutoff)
361 .execute(db)
362 .await?
363 .rows_affected();
364 Ok(n)
365 }
366
367 /// Count of currently-queued jobs. Used by the admin dashboard health panel
368 /// (Phase 2 of the audit). Allowed dead code until that route lands.
369 #[allow(dead_code)]
370 pub async fn queued_count(db: &PgPool) -> Result<i64, sqlx::Error> {
371 sqlx::query_scalar::<_, i64>("SELECT COUNT(*) FROM scan_jobs WHERE status = 'queued'")
372 .fetch_one(db)
373 .await
374 }
375
376 /// Count of currently-running jobs. Phase 2 dashboard consumer.
377 #[allow(dead_code)]
378 pub async fn running_count(db: &PgPool) -> Result<i64, sqlx::Error> {
379 sqlx::query_scalar::<_, i64>("SELECT COUNT(*) FROM scan_jobs WHERE status = 'running'")
380 .fetch_one(db)
381 .await
382 }
383
384 /// Count of running jobs whose worker has gone silent longer than
385 /// `max_age_secs`. Used by PoM to alert on stuck workers. Keyed off the
386 /// heartbeat (falling back to `started_at`) so a slow-but-alive scan isn't
387 /// counted as stuck, matching `reap_stuck`'s liveness definition.
388 pub async fn stuck_count(db: &PgPool, max_age_secs: i64) -> Result<i64, sqlx::Error> {
389 sqlx::query_scalar::<_, i64>(
390 "SELECT COUNT(*) FROM scan_jobs WHERE status = 'running' AND COALESCE(heartbeat_at, started_at) < NOW() - ($1 || ' seconds')::interval",
391 )
392 .bind(max_age_secs.to_string())
393 .fetch_one(db)
394 .await
395 }
396
397 /// A held version with enough context to re-enqueue it for scanning.
398 #[allow(dead_code)]
399 #[derive(Debug, Clone, FromRow)]
400 pub struct RescanCandidateVersion {
401 pub version_id: Uuid,
402 pub s3_key: String,
403 pub file_size_bytes: i64,
404 pub user_id: UserId,
405 }
406
407 /// A held item (audio or cover) with re-enqueue context.
408 #[allow(dead_code)]
409 #[derive(Debug, Clone, FromRow)]
410 pub struct RescanCandidateItem {
411 pub item_id: Uuid,
412 pub s3_key: String,
413 pub file_size_bytes: i64,
414 pub user_id: UserId,
415 /// Audio / cover, selected by the query.
416 pub file_type: String,
417 }
418
419 /// Find currently-held versions with enough context to be re-scanned.
420 pub async fn rescan_candidates_versions(
421 db: &PgPool,
422 ) -> Result<Vec<RescanCandidateVersion>, sqlx::Error> {
423 sqlx::query_as::<_, RescanCandidateVersion>(
424 r"
425 SELECT v.id AS version_id,
426 v.s3_key,
427 COALESCE(v.file_size_bytes, 0) AS file_size_bytes,
428 p.user_id
429 FROM versions v
430 JOIN items i ON i.id = v.item_id
431 JOIN projects p ON p.id = i.project_id
432 WHERE v.scan_status = 'held_for_review'
433 AND v.s3_key IS NOT NULL
434 ",
435 )
436 .fetch_all(db)
437 .await
438 }
439
440 /// Find currently-held items (audio or cover) with re-enqueue context.
441 pub async fn rescan_candidates_items(db: &PgPool) -> Result<Vec<RescanCandidateItem>, sqlx::Error> {
442 sqlx::query_as::<_, RescanCandidateItem>(
443 r"
444 SELECT i.id AS item_id,
445 COALESCE(i.audio_s3_key, i.cover_s3_key) AS s3_key,
446 COALESCE(i.audio_file_size_bytes, i.cover_file_size_bytes, 0) AS file_size_bytes,
447 p.user_id,
448 CASE WHEN i.audio_s3_key IS NOT NULL THEN 'audio' ELSE 'cover' END AS file_type
449 FROM items i
450 JOIN projects p ON p.id = i.project_id
451 WHERE i.scan_status = 'held_for_review'
452 AND COALESCE(i.audio_s3_key, i.cover_s3_key) IS NOT NULL
453 ",
454 )
455 .fetch_all(db)
456 .await
457 }
458
459 #[cfg(test)]
460 mod tests {
461 use super::*;
462
463 #[test]
464 fn target_kind_round_trip() {
465 for kind in [
466 ScanTargetKind::Item,
467 ScanTargetKind::Version,
468 ScanTargetKind::Media,
469 ScanTargetKind::ProjectImage,
470 ScanTargetKind::ItemImage,
471 ScanTargetKind::GalleryImage,
472 ScanTargetKind::ContentInsertion,
473 ] {
474 assert_eq!(ScanTargetKind::from_db_str(kind.as_str()), Some(kind));
475 }
476 assert_eq!(ScanTargetKind::from_db_str("bogus"), None);
477 }
478
479 #[test]
480 fn all_kinds_purge_object_on_quarantine() {
481 // The S3-object purge is unconditional, a confirmed-malicious object
482 // must never survive in storage, including downloadable Item/Version/
483 // Media content reachable via a permanent CDN URL.
484 for kind in [
485 ScanTargetKind::Item,
486 ScanTargetKind::Version,
487 ScanTargetKind::Media,
488 ScanTargetKind::ProjectImage,
489 ScanTargetKind::ItemImage,
490 ScanTargetKind::GalleryImage,
491 ScanTargetKind::ContentInsertion,
492 ScanTargetKind::OtaArtifact,
493 ] {
494 assert!(
495 kind.quarantine_purges_object(),
496 "{} must purge object",
497 kind.as_str()
498 );
499 }
500 }
501
502 #[test]
503 fn storage_bucket_matches_download_and_delete_path() {
504 // Regression for ultra-fuzz Run #24 Storage CRITICAL: the quarantine
505 // worker deleted OTA objects from the main bucket (idempotent success)
506 // while the object persisted in the synckit bucket. `storage_bucket()`
507 // is now the one source of truth the worker uses for the download
508 // client, the delete, and the durable-delete enqueue, pin it here so a
509 // new synckit-bucket kind can't silently default to Main.
510 use crate::storage::S3Bucket;
511 use ScanTargetKind::*;
512 assert_eq!(OtaArtifact.storage_bucket(), S3Bucket::Synckit);
513 for kind in [
514 Item,
515 Version,
516 Media,
517 ProjectImage,
518 ItemImage,
519 GalleryImage,
520 ContentInsertion,
521 ] {
522 assert_eq!(
523 kind.storage_bucket(),
524 S3Bucket::Main,
525 "{} is a main-bucket kind",
526 kind.as_str()
527 );
528 }
529 }
530
531 #[test]
532 fn only_gateless_image_kinds_delete_their_row() {
533 // The DB-row deletion is the gate-less-image-only half of enforcement.
534 use ScanTargetKind::*;
535 // ItemImage is gate-less: the item *cover* is CDN-served with no
536 // per-request gate; `items.scan_status` gates the audio/video, not the
537 // cover (Run #20 Storage SERIOUS).
538 for kind in [ItemImage, ProjectImage, GalleryImage, ContentInsertion] {
539 assert!(
540 kind.is_cdn_served_without_gate(),
541 "{} is gate-less",
542 kind.as_str()
543 );
544 }
545 for kind in [Item, Version, Media] {
546 assert!(
547 !kind.is_cdn_served_without_gate(),
548 "{} has a gate/status",
549 kind.as_str()
550 );
551 }
552 }
553 }
554