Skip to main content

max / makenotwork

11.1 KB · 262 lines History Blame Raw
1 //! Async scan worker.
2 //!
3 //! Spawned at startup from `main.rs`. Drains `scan_jobs` via
4 //! `db::scan_jobs::claim_next` and runs each job through the pipeline. On
5 //! completion, updates the target entity's `scan_status` (for entities that
6 //! have one) and creates a WAM ticket on `Quarantined`.
7 //!
8 //! See `docs/scan-pipeline-audit.md` § 4.4 for the architecture.
9
10 use std::sync::Arc;
11 use std::time::Duration;
12
13 use sqlx::PgPool;
14 use tokio::sync::Semaphore;
15 use uuid::Uuid;
16
17 use crate::constants;
18 use crate::db::{self, scan_jobs::{ScanJob, ScanTargetKind}, FileScanStatus, ItemId, VersionId};
19 use crate::storage::{FileType, StorageBackend};
20 use crate::wam_client::WamClient;
21
22 use super::{LayerVerdict, ScanPipeline, ScanResult};
23
24 /// Worker poll interval when the queue is empty.
25 const IDLE_POLL_INTERVAL: Duration = Duration::from_millis(500);
26
27 /// How long a `running` job can sit before the reaper resets it to `queued`.
28 /// A scan that legitimately takes longer than this is an outlier and warrants
29 /// admin attention anyway.
30 const STUCK_JOB_SECS: i64 = 300;
31
32 /// Cadence at which any worker tries to reap stuck jobs.
33 const REAPER_INTERVAL: Duration = Duration::from_secs(60);
34
35 /// Shared dependencies the worker pool needs.
36 pub struct WorkerContext {
37 pub db: PgPool,
38 pub s3: Arc<dyn StorageBackend>,
39 pub pipeline: Arc<ScanPipeline>,
40 pub scan_semaphore: Arc<Semaphore>,
41 pub wam: Option<WamClient>,
42 }
43
44 /// Spawn `n` scan workers on the current tokio runtime. Each worker drains
45 /// `scan_jobs` independently with FOR UPDATE SKIP LOCKED. A single reaper
46 /// task per pool resets jobs that get stuck in `running`.
47 ///
48 /// All tasks observe `shutdown_rx`: when the sender is dropped (or the value
49 /// changes), they exit on their next idle cycle.
50 pub fn spawn_pool(n: usize, ctx: Arc<WorkerContext>, shutdown_rx: tokio::sync::watch::Receiver<()>) {
51 for worker_id in 0..n {
52 let ctx = Arc::clone(&ctx);
53 let mut shutdown_rx = shutdown_rx.clone();
54 tokio::spawn(async move {
55 tracing::info!(worker_id, "scan worker started");
56 loop {
57 match db::scan_jobs::claim_next(&ctx.db).await {
58 Ok(Some(job)) => {
59 let job_id = job.id;
60 if let Err(e) = process_job(&ctx, job).await {
61 tracing::error!(worker_id, %job_id, error = %e, "scan job failed");
62 if let Err(e2) = db::scan_jobs::mark_failed(&ctx.db, job_id, &e.to_string()).await {
63 tracing::error!(worker_id, %job_id, error = %e2, "failed to mark job failed");
64 }
65 }
66 }
67 Ok(None) => {
68 tokio::select! {
69 _ = tokio::time::sleep(IDLE_POLL_INTERVAL) => {}
70 res = shutdown_rx.changed() => {
71 if res.is_err() {
72 tracing::info!(worker_id, "scan worker shutting down");
73 break;
74 }
75 }
76 }
77 }
78 Err(e) => {
79 tracing::error!(worker_id, error = %e, "claim_next failed; backing off");
80 tokio::select! {
81 _ = tokio::time::sleep(Duration::from_secs(5)) => {}
82 res = shutdown_rx.changed() => {
83 if res.is_err() {
84 break;
85 }
86 }
87 }
88 }
89 }
90 }
91 });
92 }
93
94 let ctx_reaper = Arc::clone(&ctx);
95 let mut shutdown_rx = shutdown_rx;
96 tokio::spawn(async move {
97 loop {
98 match db::scan_jobs::reap_stuck(&ctx_reaper.db, STUCK_JOB_SECS).await {
99 Ok(n) if n > 0 => {
100 tracing::warn!(reset = n, max_age_secs = STUCK_JOB_SECS, "reset stuck scan jobs");
101 }
102 Ok(_) => {}
103 Err(e) => tracing::error!(error = %e, "scan job reaper failed"),
104 }
105 tokio::select! {
106 _ = tokio::time::sleep(REAPER_INTERVAL) => {}
107 res = shutdown_rx.changed() => {
108 if res.is_err() {
109 break;
110 }
111 }
112 }
113 }
114 });
115 }
116
117 /// Test/dev helper: claim and process at most one queued scan job synchronously.
118 /// Returns `Ok(true)` when a job ran, `Ok(false)` when the queue was empty.
119 /// Mirrors `spawn_pool`'s per-iteration logic without spawning a background
120 /// task, so integration tests can deterministically drain the queue between
121 /// upload-confirm and assertion.
122 pub async fn process_next_for_test(ctx: &WorkerContext) -> Result<bool, Box<dyn std::error::Error + Send + Sync>> {
123 match db::scan_jobs::claim_next(&ctx.db).await? {
124 Some(job) => {
125 let job_id = job.id;
126 if let Err(e) = process_job(ctx, job).await {
127 db::scan_jobs::mark_failed(&ctx.db, job_id, &e.to_string()).await?;
128 return Err(e);
129 }
130 Ok(true)
131 }
132 None => Ok(false),
133 }
134 }
135
136 /// Run a single scan job end-to-end. On success the job is marked done; the
137 /// caller marks failed if this returns an error.
138 ///
139 /// On pipeline error (e.g. S3 download failure), reset the entity from
140 /// Scanning back to HeldForReview before bubbling the error up. Otherwise
141 /// the entity stays stuck at Scanning forever — a real regression we hit
142 /// in production with stale s3_keys.
143 #[tracing::instrument(skip_all, fields(%job_id = job.id, target_kind = %job.target_kind, %target_id = job.target_id, attempts = job.attempts))]
144 async fn process_job(ctx: &WorkerContext, job: ScanJob) -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
145 let job_id = job.id;
146 let kind = job.typed_kind().ok_or_else(|| format!("unknown target_kind: {}", job.target_kind))?;
147 let file_type = job.typed_file_type().ok_or_else(|| format!("unknown file_type: {}", job.file_type))?;
148 let target_id = job.target_id;
149
150 // Mark target as Scanning while the worker is running (only entities with
151 // a scan_status column). This is a visible signal in the admin dashboard
152 // queue panel.
153 update_entity_status(&ctx.db, kind, target_id, FileScanStatus::Scanning).await.ok();
154
155 let entity_status = match run_pipeline_and_decide(ctx, &job, file_type).await {
156 Ok(s) => s,
157 Err(e) => {
158 // Pipeline blew up — most often a stale s3_key. Reset entity to
159 // HeldForReview so admins see it on the dashboard and decide
160 // whether to delete the orphan record.
161 update_entity_status(&ctx.db, kind, target_id, FileScanStatus::HeldForReview).await.ok();
162 return Err(e);
163 }
164 };
165 update_entity_status(&ctx.db, kind, target_id, entity_status).await?;
166
167 db::scan_jobs::mark_done(&ctx.db, job_id).await?;
168 Ok(())
169 }
170
171 /// Run the pipeline against the S3 object and return the entity status to
172 /// apply, honoring the size guard, trust gate, and WAM ticketing.
173 async fn run_pipeline_and_decide(
174 ctx: &WorkerContext,
175 job: &ScanJob,
176 file_type: FileType,
177 ) -> Result<FileScanStatus, Box<dyn std::error::Error + Send + Sync>> {
178 // Two paths, gated on file size. Small files go through the original
179 // buffered `Pipeline::scan(Vec<u8>)`: a single S3 GET into a heap
180 // buffer, then layers walk the slice. Big files (>= SCAN_MAX_MEMORY_BYTES)
181 // stream from S3 into a tempfile under SCAN_SPOOL_DIR, then layers
182 // run against the spooled path (mmap or streamed). The buffered path
183 // stays alive: it's the hot path for tip-jar avatars / small audio /
184 // download files, and avoiding the tempfile syscall + write matters at
185 // that scale. Both branches run S3 IO *outside* the scan_semaphore:
186 // the permit bounds the CPU/clamd-heavy scan phase, not network IO.
187 // Holding it across the GET serializes downloads at SCAN_MAX_CONCURRENT
188 // and lets a scan backlog starve the DB pool.
189 let result: ScanResult = if (job.file_size_bytes as usize) < constants::SCAN_MAX_MEMORY_BYTES {
190 let data = ctx.s3.download_object(&job.s3_key).await?;
191 let _permit = ctx.scan_semaphore.acquire().await?;
192 Arc::clone(&ctx.pipeline).scan(data, file_type).await
193 } else {
194 let stream = ctx.s3.download_stream(&job.s3_key).await?;
195 let spool = super::spool::download_into_tempfile(
196 std::path::Path::new(constants::SCAN_SPOOL_DIR),
197 &job.s3_key,
198 job.file_size_bytes as u64,
199 stream,
200 )
201 .await?;
202 let _permit = ctx.scan_semaphore.acquire().await?;
203 Arc::clone(&ctx.pipeline).scan_stream(spool, file_type).await
204 };
205
206 db::scanning::insert_scan_result(&ctx.db, &job.s3_key, &result).await?;
207
208 if result.status == FileScanStatus::Quarantined {
209 let failed_layers: Vec<&str> = result.layers.iter()
210 .filter(|l| l.verdict == LayerVerdict::Fail)
211 .map(|l| l.layer)
212 .collect();
213 if let Some(ref wam) = ctx.wam {
214 let title = format!("File quarantined: {}", job.s3_key);
215 let body = format!(
216 "Upload by user {} flagged as malicious.\n\
217 Failed layers: {}\nFile type: {file_type:?}\nSize: {}",
218 job.user_id, failed_layers.join(", "), job.file_size_bytes,
219 );
220 wam.create_ticket(&title, Some(&body), "high", "malware-quarantine", Some(&job.s3_key)).await;
221 }
222 return Ok(FileScanStatus::Quarantined);
223 }
224
225 // Pipeline returned Clean or HeldForReview. Apply the uploader-trust
226 // overlay: untrusted users always route to admin review even on a clean
227 // scan. This preserves the pre-async semantics — the architecture changed,
228 // not the policy.
229 let is_trusted = db::users::is_upload_trusted(&ctx.db, job.user_id).await?;
230 Ok(if is_trusted { FileScanStatus::Clean } else { FileScanStatus::HeldForReview })
231 }
232
233 /// Update the per-entity `scan_status` column for the kinds that have one.
234 /// `ProjectImage` / `ContentInsertion` don't carry their own column today —
235 /// the worker still scanned the file and recorded results, but there's no
236 /// status to flip on those entities.
237 async fn update_entity_status(
238 db: &PgPool,
239 kind: ScanTargetKind,
240 target_id: Uuid,
241 status: FileScanStatus,
242 ) -> Result<(), sqlx::Error> {
243 match kind {
244 ScanTargetKind::Version => {
245 db::scanning::update_version_scan_status(db, VersionId::from(target_id), status).await
246 }
247 ScanTargetKind::Item | ScanTargetKind::ItemImage => {
248 db::scanning::update_item_scan_status(db, ItemId::from(target_id), status).await
249 }
250 ScanTargetKind::Media => {
251 db::scanning::update_media_file_scan_status(
252 db,
253 db::MediaFileId::from(target_id),
254 status,
255 )
256 .await
257 }
258 ScanTargetKind::ProjectImage | ScanTargetKind::ContentInsertion => Ok(()),
259 }
260 }
261
262