Skip to main content

max / makenotwork

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