feat(controls): enrich semantic retrieval query with finding intent + C5 live test #222
@@ -234,13 +234,22 @@ pub async fn semantic_stamp_findings(
|
|||||||
let Some(region) = fetch_region(repo_path, &file, line) else {
|
let Some(region) = fetch_region(repo_path, &file, line) else {
|
||||||
continue;
|
continue;
|
||||||
};
|
};
|
||||||
let region_emb = match llm.embed(vec![region.content.clone()]).await {
|
// Retrieve on the finding's intent + the code, not the region alone: two
|
||||||
|
// findings in one file share overlapping windows and otherwise embed alike,
|
||||||
|
// collapsing onto the same controls. The finding's title/description carry
|
||||||
|
// the discriminating signal (e.g. "brute-force protection" vs "weak hash").
|
||||||
|
// The raw `region` still goes to the judge for snippet grounding.
|
||||||
|
let query = format!(
|
||||||
|
"{}\n{}\n\n{}",
|
||||||
|
finding.title, finding.description, region.content
|
||||||
|
);
|
||||||
|
let query_emb = match llm.embed(vec![query]).await {
|
||||||
Ok(mut embs) => match embs.pop() {
|
Ok(mut embs) => match embs.pop() {
|
||||||
Some(v) => v,
|
Some(v) => v,
|
||||||
None => continue,
|
None => continue,
|
||||||
},
|
},
|
||||||
Err(e) => {
|
Err(e) => {
|
||||||
tracing::warn!(error = %e, "region embed failed; skipping finding");
|
tracing::warn!(error = %e, "query embed failed; skipping finding");
|
||||||
continue;
|
continue;
|
||||||
}
|
}
|
||||||
};
|
};
|
||||||
@@ -248,7 +257,7 @@ pub async fn semantic_stamp_findings(
|
|||||||
.check(
|
.check(
|
||||||
&index,
|
&index,
|
||||||
®ion,
|
®ion,
|
||||||
®ion_emb,
|
&query_emb,
|
||||||
SEMANTIC_TOP_K,
|
SEMANTIC_TOP_K,
|
||||||
&finding.repo_id,
|
&finding.repo_id,
|
||||||
)
|
)
|
||||||
|
|||||||
@@ -22,18 +22,20 @@ impl<J: ControlJudge> SemanticControlChecker<J> {
|
|||||||
Self { judge }
|
Self { judge }
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Map a code region to the controls it violates. `region_embedding` is the
|
/// Map a code region to the controls it violates. `query_embedding` is the
|
||||||
/// region's embedding (the caller computes it via the LLM); the top-`k`
|
/// caller-supplied retrieval embedding — typically the finding's intent
|
||||||
/// nearest controls in `index` are judged and grounded.
|
/// (title/description) plus the region, so retrieval keys on what the finding
|
||||||
|
/// is *about*, not just the ambient code. The top-`k` nearest controls in
|
||||||
|
/// `index` are then judged against the raw `region` and grounded.
|
||||||
pub async fn check(
|
pub async fn check(
|
||||||
&self,
|
&self,
|
||||||
index: &ControlIndex,
|
index: &ControlIndex,
|
||||||
region: &CandidateRegion,
|
region: &CandidateRegion,
|
||||||
region_embedding: &[f64],
|
query_embedding: &[f64],
|
||||||
k: usize,
|
k: usize,
|
||||||
repo_id: &str,
|
repo_id: &str,
|
||||||
) -> Vec<Finding> {
|
) -> Vec<Finding> {
|
||||||
let candidates = index.nearest(region_embedding, k);
|
let candidates = index.nearest(query_embedding, k);
|
||||||
let mut findings = Vec::new();
|
let mut findings = Vec::new();
|
||||||
for spec in &candidates {
|
for spec in &candidates {
|
||||||
let verdict = self.judge.judge(spec, region).await;
|
let verdict = self.judge.judge(spec, region).await;
|
||||||
|
|||||||
@@ -0,0 +1,145 @@
|
|||||||
|
//! C5 live verification — the semantic master-controls path end to end against the
|
||||||
|
//! deployed api-dev catalog. Ignored (hits api-dev + LiteLLM). Run explicitly:
|
||||||
|
//!
|
||||||
|
//! set -a; . ./.env; set +a
|
||||||
|
//! BREAKPILOT_BASE_URL=https://api-dev.breakpilot.ai \
|
||||||
|
//! cargo test -p compliance-agent --test c5_semantic_live -- --ignored --nocapture
|
||||||
|
//!
|
||||||
|
//! Pulls the live master-controls catalog, embeds the corpus (chunked), then for a
|
||||||
|
//! couple of real vulnerable findings retrieves the nearest master controls and
|
||||||
|
//! grounded-judges them, stamping master-control refs.
|
||||||
|
|
||||||
|
mod common;
|
||||||
|
|
||||||
|
use std::sync::Arc;
|
||||||
|
|
||||||
|
use compliance_agent::llm::LlmClient;
|
||||||
|
use compliance_core::config::BreakpilotConfig;
|
||||||
|
use compliance_core::models::finding::{Finding, Severity};
|
||||||
|
use compliance_core::models::scan::ScanType;
|
||||||
|
use secrecy::SecretString;
|
||||||
|
|
||||||
|
fn env(k: &str) -> String {
|
||||||
|
std::env::var(k).unwrap_or_else(|_| panic!("env {k} must be set for the live C5 test"))
|
||||||
|
}
|
||||||
|
|
||||||
|
fn mk_finding(file: &str, line: u32, title: &str) -> Finding {
|
||||||
|
let mut f = Finding::new(
|
||||||
|
"repo-c5".into(),
|
||||||
|
format!("{file}:{line}"),
|
||||||
|
"semgrep".into(),
|
||||||
|
ScanType::Sast,
|
||||||
|
title.into(),
|
||||||
|
title.into(),
|
||||||
|
Severity::High,
|
||||||
|
);
|
||||||
|
f.file_path = Some(file.into());
|
||||||
|
f.line_number = Some(line);
|
||||||
|
f
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
#[ignore = "live: requires deployed api-dev master-controls (fetch+parse only, no LLM)"]
|
||||||
|
async fn c5_ingest_master_controls_catalog() {
|
||||||
|
use compliance_agent::controls::OscalControlsProvider;
|
||||||
|
|
||||||
|
let provider = OscalControlsProvider::new(
|
||||||
|
reqwest::Client::new(),
|
||||||
|
env("BREAKPILOT_BASE_URL"),
|
||||||
|
None,
|
||||||
|
std::env::temp_dir().join("c5-ingest-snap"),
|
||||||
|
);
|
||||||
|
let doc = provider
|
||||||
|
.load_master_controls()
|
||||||
|
.await
|
||||||
|
.expect("pull + parse master-controls catalog");
|
||||||
|
let controls = doc.to_controls();
|
||||||
|
println!(
|
||||||
|
"\n=== C5 ingest: {} master controls parsed ===",
|
||||||
|
controls.len()
|
||||||
|
);
|
||||||
|
for c in controls.iter().take(4) {
|
||||||
|
let text: String = c.text.chars().take(90).collect();
|
||||||
|
println!(" {} | {} | {}", c.id, c.title, text);
|
||||||
|
}
|
||||||
|
assert!(
|
||||||
|
!controls.is_empty(),
|
||||||
|
"expected a non-empty master-control corpus"
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
#[ignore = "live: requires deployed api-dev master-controls + LiteLLM"]
|
||||||
|
async fn c5_semantic_stamps_master_control_refs() {
|
||||||
|
let llm = Arc::new(LlmClient::new(
|
||||||
|
env("LITELLM_URL"),
|
||||||
|
SecretString::from(env("LITELLM_API_KEY")),
|
||||||
|
env("LITELLM_MODEL"),
|
||||||
|
env("LITELLM_EMBED_MODEL"),
|
||||||
|
));
|
||||||
|
|
||||||
|
let mut config = common::dev_config("mongodb://unused".into(), "c5".into());
|
||||||
|
let snapshot = std::env::temp_dir().join("c5-oscal-snap");
|
||||||
|
config.breakpilot = BreakpilotConfig {
|
||||||
|
base_url: Some(env("BREAKPILOT_BASE_URL")),
|
||||||
|
token: None,
|
||||||
|
snapshot_dir: snapshot.to_string_lossy().into_owned(),
|
||||||
|
semantic_mapping: true,
|
||||||
|
grounded_control_checks: false,
|
||||||
|
};
|
||||||
|
|
||||||
|
// Fixture repo with recognizable code-checkable surfaces.
|
||||||
|
let repo = std::env::temp_dir().join("c5-fixture-repo");
|
||||||
|
let _ = std::fs::remove_dir_all(&repo);
|
||||||
|
std::fs::create_dir_all(repo.join("app")).expect("mkdir");
|
||||||
|
std::fs::write(
|
||||||
|
repo.join("app/auth.py"),
|
||||||
|
concat!(
|
||||||
|
"import hashlib\n",
|
||||||
|
"\n",
|
||||||
|
"def store_password(user, password):\n",
|
||||||
|
" # weak, unsalted password hashing\n",
|
||||||
|
" digest = hashlib.md5(password.encode()).hexdigest()\n",
|
||||||
|
" db.save(user, digest)\n",
|
||||||
|
"\n",
|
||||||
|
"@app.route('/login', methods=['POST'])\n",
|
||||||
|
"def login():\n",
|
||||||
|
" u = request.form['username']\n",
|
||||||
|
" p = request.form['password']\n",
|
||||||
|
" return 'ok' if check(u, p) else ('bad', 401)\n",
|
||||||
|
),
|
||||||
|
)
|
||||||
|
.expect("write fixture");
|
||||||
|
|
||||||
|
let mut findings = vec![
|
||||||
|
mk_finding("app/auth.py", 5, "Weak password hash (md5, unsalted)"),
|
||||||
|
mk_finding(
|
||||||
|
"app/auth.py",
|
||||||
|
9,
|
||||||
|
"Login endpoint without brute-force protection",
|
||||||
|
),
|
||||||
|
];
|
||||||
|
|
||||||
|
let tagged =
|
||||||
|
compliance_agent::controls::semantic_stamp_findings(&config, llm, &repo, &mut findings)
|
||||||
|
.await;
|
||||||
|
|
||||||
|
println!("\n=== C5 semantic master-controls stamping ===");
|
||||||
|
for f in &findings {
|
||||||
|
println!(
|
||||||
|
" {:50} {}:{:?} -> {:?}",
|
||||||
|
f.title,
|
||||||
|
f.file_path.as_deref().unwrap_or(""),
|
||||||
|
f.line_number,
|
||||||
|
f.control_refs
|
||||||
|
);
|
||||||
|
}
|
||||||
|
println!("findings that gained >=1 master-control ref: {tagged}");
|
||||||
|
let _ = std::fs::remove_dir_all(&repo);
|
||||||
|
|
||||||
|
// Live corpus — assert only that the path runs and stamps at least one ref.
|
||||||
|
assert!(
|
||||||
|
tagged >= 1,
|
||||||
|
"expected at least one finding to gain a master-control ref"
|
||||||
|
);
|
||||||
|
}
|
||||||
Reference in New Issue
Block a user