//! Werkbank runner endpoints (`/api/v1/werkbank/jobs/*`). //! //! The pull API a Werkbank runner talks to: lease a job, heartbeat while it runs, //! and post the result back. Machine auth is a **static bearer token** //! (`WERKBANK_RUNNER_TOKEN`) — not a Keycloak JWT, because a runner acts across //! tenants (each request names its `tenant`). Routes are only mounted when the //! token is configured; with none set they don't exist (404). //! //! On completion the runner's findings are persisted against the job's target, //! so a job run by a remote runner lands the same findings an in-process run //! would (WB-05, the control-plane cut-over). use axum::extract::{Extension, Path, Request}; use axum::http::{header, StatusCode}; use axum::middleware::Next; use axum::response::{IntoResponse, Response}; use axum::Json; use mongodb::bson::{doc, oid::ObjectId}; use secrecy::ExposeSecret; use serde::{Deserialize, Serialize}; use std::time::Duration; use compliance_core::models::werkbank::{ CompleteRequest, CompleteResponse, HeartbeatRequest, InputRef, Job, JobResult, LeaseRequest, }; use compliance_core::models::ArtifactKind; use super::dto::AgentExt; use crate::database::Database; use crate::werkbank::JobQueue; /// Gate the runner endpoints behind the static runner bearer token. pub async fn require_runner_token( Extension(agent): AgentExt, request: Request, next: Next, ) -> Response { let Some(expected) = agent.config.werkbank_runner_token.as_ref() else { return (StatusCode::NOT_FOUND, "werkbank runner API disabled").into_response(); }; let presented = request .headers() .get(header::AUTHORIZATION) .and_then(|v| v.to_str().ok()) .and_then(|s| s.strip_prefix("Bearer ")) .map(str::trim) .filter(|s| !s.is_empty()); let Some(presented) = presented else { return (StatusCode::UNAUTHORIZED, "Missing bearer token").into_response(); }; if !constant_time_eq(presented, expected.expose_secret()) { return (StatusCode::UNAUTHORIZED, "Invalid runner token").into_response(); } next.run(request).await } /// `POST /api/v1/werkbank/jobs/lease` — lease the oldest runnable job, or `204`. #[tracing::instrument(skip_all, fields(tenant = %req.tenant, runner = %req.runner_id))] pub async fn lease( Extension(agent): AgentExt, Json(req): Json, ) -> Result { let queue = JobQueue::new(&tenant_db(&agent, &req.tenant).await?); let leased = queue .lease( &req.runner_id, req.executor, &req.labels, Duration::from_secs(req.lease_ttl_secs), chrono::Utc::now(), ) .await .map_err(internal)?; Ok(match leased { Some(job) => Json(job).into_response(), None => StatusCode::NO_CONTENT.into_response(), }) } /// `POST /api/v1/werkbank/jobs/heartbeat` — extend the lease; `409` if it's lost. #[tracing::instrument(skip_all, fields(tenant = %req.tenant, job = %req.job_id))] pub async fn heartbeat( Extension(agent): AgentExt, Json(req): Json, ) -> Result { let queue = JobQueue::new(&tenant_db(&agent, &req.tenant).await?); let ack = queue .heartbeat( &req.job_id, &req.lease_token, Duration::from_secs(req.lease_ttl_secs), chrono::Utc::now(), ) .await .map_err(internal)?; Ok(match ack { Some(ack) => Json(ack).into_response(), // Lease lost — the runner should abandon the job. None => StatusCode::CONFLICT.into_response(), }) } /// `POST /api/v1/werkbank/jobs/complete` — record the result and persist findings. #[tracing::instrument(skip_all, fields(tenant = %req.tenant, job = %req.job_id))] pub async fn complete( Extension(agent): AgentExt, Json(req): Json, ) -> Result, StatusCode> { let db = tenant_db(&agent, &req.tenant).await?; let queue = JobQueue::new(&db); let now = chrono::Utc::now(); let recorded = queue .complete(&req.job_id, &req.lease_token, &req.result, now) .await .map_err(internal)?; // Only persist findings for the run that actually recorded the result, so a // duplicate/late completion can't double-insert. if recorded { if let Some(record) = queue.get(&req.job_id).await.map_err(internal)? { persist_findings(&db, &record.job.target_id, &req.result).await; } } Ok(Json(CompleteResponse { recorded })) } /// `GET /api/v1/werkbank/artifacts/{hash}` — serve a content-addressed blob (the /// program a runner needs to load). The hash is validated against traversal by /// [`crate::ingest::blob::read_blob`]; a runner fetches this for a job's `blob` /// input. #[tracing::instrument(skip_all, fields(hash = %hash))] pub async fn serve_artifact( Extension(agent): AgentExt, Path(hash): Path, ) -> Result { let base = std::path::Path::new(&agent.config.artifact_store_base_path); match crate::ingest::blob::read_blob(base, &hash) { Ok(bytes) => { Ok(([(header::CONTENT_TYPE, "application/octet-stream")], bytes).into_response()) } Err(_) => Err(StatusCode::NOT_FOUND), } } /// Enqueue a `plc-provision` job for a target: extract its control-logic program, /// stash it as a content-addressed blob (which the runner fetches via /// [`serve_artifact`]), and queue the job. This is the control-plane "enqueue" /// half of the loop — a runner then leases it, provisions, and posts results. #[derive(Debug, Deserialize)] pub struct EnqueueRequest { /// The tenant whose queue to enqueue into. pub tenant: String, /// The onboarded target to test. pub target_id: String, } /// The enqueued job's id. #[derive(Debug, Serialize)] pub struct EnqueueResponse { /// The new job id. pub job_id: String, /// Whether this call inserted it (false = already queued). pub enqueued: bool, } #[tracing::instrument(skip_all, fields(tenant = %req.tenant, target = %req.target_id))] pub async fn enqueue( Extension(agent): AgentExt, Json(req): Json, ) -> Result, StatusCode> { let db = tenant_db(&agent, &req.tenant).await?; let oid = ObjectId::parse_str(&req.target_id).map_err(|_| StatusCode::BAD_REQUEST)?; let target = db .onboarded_targets() .find_one(doc! { "_id": oid }) .await .map_err(internal)? .ok_or(StatusCode::NOT_FOUND)?; // Extract the control-logic program from the target's PLC-source artifacts // (same selection as the in-process PLC scan). let ctx = crate::ingest::IngestContext::from_config(&agent.config, &req.target_id); let ingest_set = crate::ingest::ingest_all(&target, &ctx).map_err(internal)?; let program = target .artifacts .iter() .filter(|a| { matches!( a.kind, ArtifactKind::PlcProject | ArtifactKind::GitRepo | ArtifactKind::SourceArchive ) }) .find_map(|a| { let path = ingest_set .get(&a.id) .and_then(|ia| ia.working_path.clone())?; werkbank_exec::plc::extract_program(&path) }) .ok_or(StatusCode::UNPROCESSABLE_ENTITY)?; // Stash the program source so the runner can fetch it by hash. let base = std::path::Path::new(&agent.config.artifact_store_base_path); let hash = crate::ingest::blob::store_bytes(base, program.source.as_bytes()).map_err(internal)?; let job_id = format!("job_{}", uuid::Uuid::new_v4().simple()); let job = Job::plc_provision( &job_id, &req.tenant, &req.target_id, InputRef::blob(hash), agent.config.plc_runtime.max_lifetime_secs, ); let enqueued = JobQueue::new(&db) .enqueue(job, chrono::Utc::now()) .await .map_err(internal)?; Ok(Json(EnqueueResponse { job_id, enqueued })) } /// Persist a job result's findings against its target: general findings /// (dedup'd by fingerprint) and DAST findings. Best-effort — a persistence hiccup /// is logged, not surfaced to the runner (its result is already recorded). async fn persist_findings(db: &Database, target_id: &str, result: &JobResult) { for finding in &result.findings { let exists = db .findings() .find_one(doc! { "fingerprint": &finding.fingerprint }) .await .ok() .flatten() .is_some(); if !exists { if let Err(e) = db.findings().insert_one(finding).await { tracing::warn!(target_id, error = %e, "werkbank: persist finding failed"); } } } for finding in &result.dast_findings { if let Err(e) = db.dast_findings().insert_one(finding).await { tracing::warn!(target_id, error = %e, "werkbank: persist DAST finding failed"); } } tracing::info!( target_id, findings = result.findings.len(), dast = result.dast_findings.len(), "werkbank: persisted runner results" ); } /// Resolve the tenant-scoped database for a request. async fn tenant_db( agent: &crate::agent::ComplianceAgent, tenant: &str, ) -> Result { agent.db_pool.for_tenant_id(tenant).await.map_err(internal) } /// Map any internal error to a 500. fn internal(e: E) -> StatusCode { tracing::error!("werkbank endpoint error: {e}"); StatusCode::INTERNAL_SERVER_ERROR } /// Length-checked, constant-time-ish token comparison. fn constant_time_eq(a: &str, b: &str) -> bool { if a.len() != b.len() { return false; } let mut diff = 0u8; for (x, y) in a.bytes().zip(b.bytes()) { diff |= x ^ y; } diff == 0 } #[cfg(test)] mod tests { use super::constant_time_eq; #[test] fn token_compare() { assert!(constant_time_eq("secret", "secret")); assert!(!constant_time_eq("secret", "secrex")); assert!(!constant_time_eq("secret", "secretx")); assert!(!constant_time_eq("", "x")); } }