compliance-scanner-agent/compliance-agent/src/api/handlers/werkbank_jobs.rs

//! 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<LeaseRequest>,
) -> Result<Response, StatusCode> {
    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<HeartbeatRequest>,
) -> Result<Response, StatusCode> {
    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<CompleteRequest>,
) -> Result<Json<CompleteResponse>, 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<String>,
) -> Result<Response, StatusCode> {
    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<EnqueueRequest>,
) -> Result<Json<EnqueueResponse>, 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<Database, StatusCode> {
    agent.db_pool.for_tenant_id(tenant).await.map_err(internal)
}

/// Map any internal error to a 500.
fn internal<E: std::fmt::Display>(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"));
    }
}