From 00c8550c3c8cabe6dacea95e445e883a640ef44c Mon Sep 17 00:00:00 2001 From: Noah Santschi-Cooney Date: Tue, 11 Aug 2026 17:25:38 +0100 Subject: [PATCH] fix: persist total_components count for EI analysis when job terminates --- .../src/endpoints/test.rs | 19 ++- .../src/runner/polling.rs | 11 +- .../exploit-intelligence/src/runner/test.rs | 109 +++++++++++++++--- .../exploit-intelligence/src/runner/worker.rs | 5 +- .../exploit-intelligence/src/service/mod.rs | 57 ++------- .../exploit-intelligence/src/service/test.rs | 7 +- 6 files changed, 131 insertions(+), 77 deletions(-) diff --git a/modules/exploit-intelligence/src/endpoints/test.rs b/modules/exploit-intelligence/src/endpoints/test.rs index 0c9c1cd94..08644a8b9 100644 --- a/modules/exploit-intelligence/src/endpoints/test.rs +++ b/modules/exploit-intelligence/src/endpoints/test.rs @@ -50,6 +50,7 @@ async fn insert_test_job( vulnerability_id: &str, status: ExploitIntelligenceJobStatus, error_message: Option<&str>, + total_components: Option, ) -> exploit_intelligence_job::Model { let now = OffsetDateTime::now_utc(); let model = exploit_intelligence_job::ActiveModel { @@ -59,7 +60,7 @@ async fn insert_test_job( status: Set(status), error_message: Set(error_message.map(|s| s.to_string())), product_id: Set(None), - total_components: Set(None), + total_components: Set(total_components), retry_count: Set(0), created: Set(now), updated: Set(now), @@ -115,6 +116,7 @@ async fn get_job_completed(ctx: &TrustifyContext) -> anyhow::Result<()> { "CVE-2024-1234", ExploitIntelligenceJobStatus::Completed, None, + Some(1), ) .await; insert_test_component( @@ -155,6 +157,7 @@ async fn get_job_completed(ctx: &TrustifyContext) -> anyhow::Result<()> { body.finding, Some(ExploitIntelligenceFinding::NotVulnerable) ); + assert_eq!(body.total_components, Some(1)); Ok(()) } @@ -174,6 +177,7 @@ async fn get_job_pending(ctx: &TrustifyContext) -> anyhow::Result<()> { "CVE-2024-5678", ExploitIntelligenceJobStatus::Pending, None, + None, ) .await; @@ -221,6 +225,7 @@ async fn get_job_failed(ctx: &TrustifyContext) -> anyhow::Result<()> { "CVE-2024-9999", ExploitIntelligenceJobStatus::Failed, Some("analysis timed out"), + None, ) .await; @@ -295,6 +300,7 @@ async fn list_jobs_paginated(ctx: &TrustifyContext) -> anyhow::Result<()> { &format!("CVE-2024-{i:04}"), ExploitIntelligenceJobStatus::Pending, None, + None, ) .await; } @@ -342,6 +348,7 @@ async fn list_jobs_filter_by_sbom_id(ctx: &TrustifyContext) -> anyhow::Result<() "CVE-2024-0001", ExploitIntelligenceJobStatus::Pending, None, + None, ) .await; @@ -352,6 +359,7 @@ async fn list_jobs_filter_by_sbom_id(ctx: &TrustifyContext) -> anyhow::Result<() "CVE-2024-0002", ExploitIntelligenceJobStatus::Pending, None, + None, ) .await; @@ -403,6 +411,7 @@ async fn analyze_deduplicates_active_jobs(ctx: &TrustifyContext) -> anyhow::Resu "CVE-2024-DEDUP", ExploitIntelligenceJobStatus::Pending, None, + None, ) .await; @@ -455,6 +464,7 @@ async fn get_job_timeout_failure(ctx: &TrustifyContext) -> anyhow::Result<()> { "CVE-2024-TIMEOUT", ExploitIntelligenceJobStatus::Running, None, + None, ) .await; @@ -462,7 +472,7 @@ async fn get_job_timeout_failure(ctx: &TrustifyContext) -> anyhow::Result<()> { let ei_service = test_service(); let db_rw = db::ReadWrite::new(ctx.db.clone()); ei_service - .update_job_failed(job_id, "analysis timed out after 1800s", &db_rw) + .update_job_failed(job_id, "analysis timed out after 1800s", None, &db_rw) .await?; // Then fetching the job via GET shows the failed status and timeout message @@ -507,6 +517,7 @@ async fn get_job_finding_without_advisory(ctx: &TrustifyContext) -> anyhow::Resu "CVE-2024-VULN", ExploitIntelligenceJobStatus::Completed, None, + None, ) .await; insert_test_component( @@ -605,6 +616,7 @@ async fn get_job_running(ctx: &TrustifyContext) -> anyhow::Result<()> { "CVE-2024-RUN", ExploitIntelligenceJobStatus::Running, None, + None, ) .await; @@ -652,6 +664,7 @@ async fn list_jobs_filter_by_status(ctx: &TrustifyContext) -> anyhow::Result<()> "CVE-2024-S001", ExploitIntelligenceJobStatus::Completed, None, + None, ) .await; @@ -662,6 +675,7 @@ async fn list_jobs_filter_by_status(ctx: &TrustifyContext) -> anyhow::Result<()> "CVE-2024-S002", ExploitIntelligenceJobStatus::Pending, None, + None, ) .await; @@ -672,6 +686,7 @@ async fn list_jobs_filter_by_status(ctx: &TrustifyContext) -> anyhow::Result<()> "CVE-2024-S003", ExploitIntelligenceJobStatus::Completed, None, + None, ) .await; diff --git a/modules/exploit-intelligence/src/runner/polling.rs b/modules/exploit-intelligence/src/runner/polling.rs index b400e8df9..fcd23b986 100644 --- a/modules/exploit-intelligence/src/runner/polling.rs +++ b/modules/exploit-intelligence/src/runner/polling.rs @@ -131,7 +131,7 @@ pub(crate) async fn poll_for_result( .update_component_completed(component_id, finding, advisory_id, db) .await?; - ei_service.update_job_completed(job_id, db).await?; + ei_service.update_job_completed(job_id, Some(1), db).await?; tracing::info!( job_id = %job_id, @@ -154,7 +154,7 @@ pub(crate) async fn poll_for_result( .update_component_failed(component_id, &error_msg, db) .await?; - ei_service.update_job_completed(job_id, db).await?; + ei_service.update_job_completed(job_id, Some(1), db).await?; Ok(()) } @@ -288,8 +288,8 @@ pub(crate) async fn poll_for_product_result( .await?; } - // Check whether any component actually has a finding. let components = ei_service.fetch_components(job_id, db).await?; + let total = Some(components.len() as i32); let has_finding = components.iter().any(|c| { c.status == ExploitIntelligenceJobStatus::Completed && c.finding.is_some() @@ -297,7 +297,7 @@ pub(crate) async fn poll_for_product_result( let all_excluded = !components.is_empty() && components.iter().all(|c| c.excluded); if has_finding || all_excluded { - ei_service.update_job_completed(job_id, db).await?; + ei_service.update_job_completed(job_id, total, db).await?; tracing::info!( job_id = %job_id, @@ -314,6 +314,7 @@ pub(crate) async fn poll_for_product_result( "no components produced findings ({} total, {failed_not_excluded} failed)", components.len() ), + total, db, ) .await?; @@ -329,7 +330,7 @@ pub(crate) async fn poll_for_product_result( } ProductState::Failed => { let msg = format!("EI product analysis failed for product {product_id}"); - ei_service.update_job_failed(job_id, &msg, db).await?; + ei_service.update_job_failed(job_id, &msg, None, db).await?; Ok(()) } // Still in progress — no DB changes, job stays Running for re-poll next tick. diff --git a/modules/exploit-intelligence/src/runner/test.rs b/modules/exploit-intelligence/src/runner/test.rs index c0ce7f561..80016090c 100644 --- a/modules/exploit-intelligence/src/runner/test.rs +++ b/modules/exploit-intelligence/src/runner/test.rs @@ -79,7 +79,7 @@ async fn run_pipeline( .await { let _ = ei_service - .update_job_failed(claimed.id, &e.to_string(), &claim_tx) + .update_job_failed(claimed.id, &e.to_string(), None, &claim_tx) .await; claim_tx.commit().await.expect("commit tx"); break; @@ -148,6 +148,16 @@ async fn cdx_completed_not_vulnerable(ctx: &TrustifyContext) -> anyhow::Result<( details.finding, Some(ExploitIntelligenceFinding::NotVulnerable) ); + assert_eq!(details.product_id, None); + assert!(details.report_url.is_some()); + assert_eq!(details.total_components, Some(1)); + assert_eq!(details.completed_components, Some(1)); + assert_eq!(details.failed_components, Some(0)); + assert_eq!(details.excluded_components, Some(0)); + assert_eq!(details.vulnerable_components, Some(0)); + assert_eq!(details.not_vulnerable_components, Some(1)); + assert_eq!(details.uncertain_components, Some(0)); + assert_eq!(details.components.len(), 1); Ok(()) } @@ -187,11 +197,19 @@ async fn cdx_ei_analysis_failed(ctx: &TrustifyContext) -> anyhow::Result<()> { let details = ei_service.fetch_job(job_id, &*db_rw).await?; let details = details.expect("job should exist"); assert_eq!(details.status, ExploitIntelligenceJobStatus::Completed); - assert!( - details - .components - .iter() - .any(|c| { c.status == ExploitIntelligenceJobStatus::Failed }) + assert_eq!(details.finding, None); + assert_eq!(details.product_id, None); + assert_eq!(details.total_components, Some(1)); + assert_eq!(details.completed_components, Some(0)); + assert_eq!(details.failed_components, Some(1)); + assert_eq!(details.excluded_components, Some(0)); + assert_eq!(details.vulnerable_components, Some(0)); + assert_eq!(details.not_vulnerable_components, Some(0)); + assert_eq!(details.uncertain_components, Some(0)); + assert_eq!(details.components.len(), 1); + assert_eq!( + details.components[0].status, + ExploitIntelligenceJobStatus::Failed ); Ok(()) } @@ -216,6 +234,10 @@ async fn cdx_upload_permanent_failure(ctx: &TrustifyContext) -> anyhow::Result<( let details = ei_service.fetch_job(job_id, &*db_rw).await?; let details = details.expect("job should exist"); assert_eq!(details.status, ExploitIntelligenceJobStatus::Failed); + assert_eq!(details.finding, None); + assert_eq!(details.product_id, None); + assert_eq!(details.total_components, None); + assert!(details.components.is_empty()); Ok(()) } @@ -239,6 +261,10 @@ async fn cdx_upload_transient_failure(ctx: &TrustifyContext) -> anyhow::Result<( let details = ei_service.fetch_job(job_id, &*db_rw).await?; let details = details.expect("job should exist"); assert_eq!(details.status, ExploitIntelligenceJobStatus::Failed); + assert_eq!(details.finding, None); + assert_eq!(details.product_id, None); + assert_eq!(details.total_components, None); + assert!(details.components.is_empty()); Ok(()) } @@ -322,6 +348,15 @@ async fn spdx_completed_with_finding(ctx: &TrustifyContext) -> anyhow::Result<() details.finding, Some(ExploitIntelligenceFinding::NotVulnerable) ); + assert_eq!(details.product_id.as_deref(), Some("prod-001")); + assert!(details.report_url.is_some()); + assert_eq!(details.total_components, Some(1)); + assert_eq!(details.completed_components, Some(1)); + assert_eq!(details.failed_components, Some(0)); + assert_eq!(details.excluded_components, Some(0)); + assert_eq!(details.vulnerable_components, Some(0)); + assert_eq!(details.not_vulnerable_components, Some(1)); + assert_eq!(details.uncertain_components, Some(0)); assert_eq!(details.components.len(), 1); Ok(()) } @@ -382,7 +417,16 @@ async fn spdx_all_components_excluded(ctx: &TrustifyContext) -> anyhow::Result<( .expect("job should exist"); assert_eq!(details.status, ExploitIntelligenceJobStatus::Completed); assert_eq!(details.finding, None); - assert!(details.components.iter().any(|c| c.excluded)); + assert_eq!(details.product_id.as_deref(), Some("prod-excl")); + assert_eq!(details.total_components, Some(1)); + assert_eq!(details.completed_components, Some(0)); + assert_eq!(details.failed_components, Some(0)); + assert_eq!(details.excluded_components, Some(1)); + assert_eq!(details.vulnerable_components, Some(0)); + assert_eq!(details.not_vulnerable_components, Some(0)); + assert_eq!(details.uncertain_components, Some(0)); + assert_eq!(details.components.len(), 1); + assert!(details.components[0].excluded); Ok(()) } @@ -429,6 +473,10 @@ async fn spdx_product_failed(ctx: &TrustifyContext) -> anyhow::Result<()> { .await? .expect("job should exist"); assert_eq!(details.status, ExploitIntelligenceJobStatus::Failed); + assert_eq!(details.finding, None); + assert_eq!(details.product_id.as_deref(), Some("prod-fail")); + assert_eq!(details.total_components, None); + assert!(details.components.is_empty()); Ok(()) } @@ -476,9 +524,17 @@ async fn cdx_completed_with_vex_ingestion(ctx: &TrustifyContext) -> anyhow::Resu .await? .expect("job should exist"); assert_eq!(details.status, ExploitIntelligenceJobStatus::Completed); + assert_eq!( + details.finding, + Some(ExploitIntelligenceFinding::NotVulnerable) + ); + assert_eq!(details.total_components, Some(1)); + assert_eq!(details.completed_components, Some(1)); + assert_eq!(details.not_vulnerable_components, Some(1)); + assert_eq!(details.components.len(), 1); assert!( - details.components.iter().any(|c| c.advisory_id.is_some()), - "expected a component with an ingested advisory" + details.components[0].advisory_id.is_some(), + "expected component with an ingested advisory" ); Ok(()) } @@ -568,9 +624,16 @@ async fn spdx_mixed_completed_and_excluded(ctx: &TrustifyContext) -> anyhow::Res details.finding, Some(ExploitIntelligenceFinding::NotVulnerable) ); + assert_eq!(details.product_id.as_deref(), Some("prod-mixed")); + assert!(details.report_url.is_some()); + assert_eq!(details.total_components, Some(2)); + assert_eq!(details.completed_components, Some(1)); + assert_eq!(details.failed_components, Some(0)); + assert_eq!(details.excluded_components, Some(1)); + assert_eq!(details.vulnerable_components, Some(0)); + assert_eq!(details.not_vulnerable_components, Some(1)); + assert_eq!(details.uncertain_components, Some(0)); assert_eq!(details.components.len(), 2); - assert_eq!(details.completed_components, Some(1),); - assert_eq!(details.excluded_components, Some(1),); Ok(()) } @@ -661,6 +724,10 @@ async fn spdx_reports_endpoint_failure(ctx: &TrustifyContext) -> anyhow::Result< .await? .expect("job should exist"); assert_eq!(details.status, ExploitIntelligenceJobStatus::Failed); + assert_eq!(details.finding, None); + assert_eq!(details.product_id.as_deref(), Some("prod-reportfail")); + assert_eq!(details.total_components, None); + assert!(details.components.is_empty()); Ok(()) } @@ -711,12 +778,16 @@ async fn spdx_component_expired_report(ctx: &TrustifyContext) -> anyhow::Result< .await? .expect("job should exist"); assert_eq!(details.status, ExploitIntelligenceJobStatus::Failed); - assert!( - details - .components - .iter() - .any(|c| { c.status == ExploitIntelligenceJobStatus::Failed }), - "expected a failed component" + assert_eq!(details.finding, None); + assert_eq!(details.product_id.as_deref(), Some("prod-compexp")); + assert_eq!(details.total_components, Some(1)); + assert_eq!(details.completed_components, Some(0)); + assert_eq!(details.failed_components, Some(1)); + assert_eq!(details.excluded_components, Some(0)); + assert_eq!(details.components.len(), 1); + assert_eq!( + details.components[0].status, + ExploitIntelligenceJobStatus::Failed ); Ok(()) } @@ -761,5 +832,9 @@ async fn spdx_no_findings_marks_job_failed(ctx: &TrustifyContext) -> anyhow::Res .await? .expect("job should exist"); assert_eq!(details.status, ExploitIntelligenceJobStatus::Failed); + assert_eq!(details.finding, None); + assert_eq!(details.product_id.as_deref(), Some("prod-nofind")); + assert_eq!(details.total_components, Some(0)); + assert!(details.components.is_empty()); Ok(()) } diff --git a/modules/exploit-intelligence/src/runner/worker.rs b/modules/exploit-intelligence/src/runner/worker.rs index 8ee3a4326..02592665c 100644 --- a/modules/exploit-intelligence/src/runner/worker.rs +++ b/modules/exploit-intelligence/src/runner/worker.rs @@ -45,7 +45,7 @@ impl EiWorker { if job.retry_count as usize >= max_retries { tracing::warn!(job_id = %job_id, retry_count = job.retry_count, "failing job: max retries exhausted on claim"); self.ei_service - .update_job_failed(job_id, "max retries exhausted", &tx) + .update_job_failed(job_id, "max retries exhausted", None, &tx) .await?; tx.commit().await?; return Ok(()); @@ -67,7 +67,7 @@ impl EiWorker { tracing::error!(job_id = %job_id, error = %e, "permanent failure, failing job"); if let Err(db_err) = self .ei_service - .update_job_failed(job_id, &e.to_string(), &*self.db_rw) + .update_job_failed(job_id, &e.to_string(), None, &*self.db_rw) .await { tracing::warn!(job_id = %job_id, error = %db_err, "failed to mark job as failed"); @@ -88,6 +88,7 @@ impl EiWorker { .update_job_failed( job_id, &format!("max retries exhausted: {e}"), + None, &*self.db_rw, ) .await diff --git a/modules/exploit-intelligence/src/service/mod.rs b/modules/exploit-intelligence/src/service/mod.rs index 04eb9efab..ced2d6fd9 100644 --- a/modules/exploit-intelligence/src/service/mod.rs +++ b/modules/exploit-intelligence/src/service/mod.rs @@ -6,8 +6,10 @@ use crate::{ model::{ComponentCounts, ExploitIntelligenceJobDetails, ExploitIntelligenceJobSummary}, }; use sea_orm::{ - ActiveModelTrait, ColumnTrait, Condition, ConnectionTrait, DatabaseTransaction, EntityTrait, - PaginatorTrait, QueryFilter, QueryOrder, QuerySelect, Set, TransactionTrait, + ActiveModelTrait, + ActiveValue::NotSet, + ColumnTrait, Condition, ConnectionTrait, DatabaseTransaction, EntityTrait, PaginatorTrait, + QueryFilter, QueryOrder, QuerySelect, Set, TransactionTrait, sea_query::{Expr, LockBehavior, LockType, SimpleExpr}, }; use std::{collections::HashMap, sync::Arc, time::Duration}; @@ -392,15 +394,7 @@ impl ExploitIntelligenceService { .iter() .map(|model| { let mut summary = if let Some(c) = counts_map.get(&model.id) { - let mut s = ExploitIntelligenceJobSummary::from_entity_with_counts( - model, - self.ui_url(), - c, - ); - if s.total_components.is_none() { - s.total_components = Some(c.total); - } - s + ExploitIntelligenceJobSummary::from_entity_with_counts(model, self.ui_url(), c) } else { ExploitIntelligenceJobSummary::from_entity(model, self.ui_url()) }; @@ -430,11 +424,13 @@ impl ExploitIntelligenceService { pub async fn update_job_completed( &self, job_id: Uuid, + total_components: Option, db: &(impl ConnectionTrait + Send), ) -> Result<(), Error> { let model = exploit_intelligence_job::ActiveModel { id: Set(job_id), status: Set(ExploitIntelligenceJobStatus::Completed), + total_components: total_components.map_or(NotSet, |n| Set(Some(n))), updated: Set(OffsetDateTime::now_utc()), ..Default::default() }; @@ -456,6 +452,7 @@ impl ExploitIntelligenceService { &self, job_id: Uuid, error_message: &str, + total_components: Option, db: &(impl ConnectionTrait + Send), ) -> Result<(), Error> { let now = OffsetDateTime::now_utc(); @@ -464,6 +461,7 @@ impl ExploitIntelligenceService { id: Set(job_id), status: Set(ExploitIntelligenceJobStatus::Failed), error_message: Set(Some(error_message.to_string())), + total_components: total_components.map_or(NotSet, |n| Set(Some(n))), updated: Set(now), ..Default::default() }; @@ -530,43 +528,6 @@ impl ExploitIntelligenceService { Ok(result.map(|m| ExploitIntelligenceJobDetails::from_entity(&m, self.ui_url()))) } - /// Create an analysis job with product_id and total_components set (multi-component SPDX flow). - #[instrument(skip_all, err(level = tracing::Level::INFO))] - pub async fn create_job_with_product( - &self, - sbom_id: Uuid, - vulnerability_id: &str, - product_id: &str, - total_components: i32, - db: &db::ReadWrite, - ) -> Result { - let _config = self.runtime.as_ref().ok_or(Error::Unavailable)?; - let now = OffsetDateTime::now_utc(); - let job_id = Uuid::now_v7(); - - let model = exploit_intelligence_job::ActiveModel { - id: Set(job_id), - sbom_id: Set(sbom_id), - vulnerability_id: Set(vulnerability_id.to_string()), - status: Set(ExploitIntelligenceJobStatus::Pending), - error_message: Set(None), - product_id: Set(Some(product_id.to_string())), - total_components: Set(Some(total_components)), - retry_count: Set(0), - created: Set(now), - updated: Set(now), - }; - - let tx = db.begin().await?; - let result = model.insert(&tx).await.map_err(Error::from)?; - tx.commit().await?; - - Ok(ExploitIntelligenceJobDetails::from_entity( - &result, - self.ui_url(), - )) - } - /// Create a component record for an analysis job. #[instrument(skip_all, err(level = tracing::Level::INFO))] pub async fn create_component( diff --git a/modules/exploit-intelligence/src/service/test.rs b/modules/exploit-intelligence/src/service/test.rs index a525b9908..80ac0df19 100644 --- a/modules/exploit-intelligence/src/service/test.rs +++ b/modules/exploit-intelligence/src/service/test.rs @@ -243,13 +243,14 @@ async fn update_job_completed_sets_status(ctx: &TrustifyContext) -> anyhow::Resu ) .await; - svc.update_job_completed(job.id, &*db_rw).await?; + svc.update_job_completed(job.id, Some(4), &*db_rw).await?; let updated = exploit_intelligence_job::Entity::find_by_id(job.id) .one(&*db_rw) .await? .expect("job should exist"); assert_eq!(updated.status, ExploitIntelligenceJobStatus::Completed); + assert_eq!(updated.total_components, Some(4)); Ok(()) } @@ -271,7 +272,7 @@ async fn update_job_failed_cascades_to_components(ctx: &TrustifyContext) -> anyh let comp_id = svc.create_component(job.id, "test-comp", &*db_rw).await?; - svc.update_job_failed(job.id, "something broke", &*db_rw) + svc.update_job_failed(job.id, "something broke", None, &*db_rw) .await?; let comp = exploit_intelligence_job_component::Entity::find_by_id(comp_id) @@ -310,7 +311,7 @@ async fn update_job_failed_preserves_completed_components( ) .await?; - svc.update_job_failed(job.id, "partial failure", &*db_rw) + svc.update_job_failed(job.id, "partial failure", None, &*db_rw) .await?; let comp = exploit_intelligence_job_component::Entity::find_by_id(comp_id)