122 lines
4.1 KiB
Rust
122 lines
4.1 KiB
Rust
//! Semantic control mapping: retrieve the top-K controls nearest a code region,
|
|
//! then confirm each with the grounded judge.
|
|
//!
|
|
//! The `region → controls` direction (vs. the CWE-LUT's `finding → control`) is
|
|
//! what scales to the full master-control corpus: the LLM only ever judges a
|
|
//! handful of retrieved candidates, and every surviving verdict is still anchored
|
|
//! to real code by the grounding gate.
|
|
|
|
use compliance_core::control_check::{ground, CandidateRegion};
|
|
use compliance_core::models::Finding;
|
|
|
|
use super::index::ControlIndex;
|
|
use super::judge::ControlJudge;
|
|
|
|
/// Retrieve → judge → ground, generic over the judge so tests use a stub.
|
|
pub struct SemanticControlChecker<J> {
|
|
judge: J,
|
|
}
|
|
|
|
impl<J: ControlJudge> SemanticControlChecker<J> {
|
|
pub fn new(judge: J) -> Self {
|
|
Self { judge }
|
|
}
|
|
|
|
/// Map a code region to the controls it violates. `query_embedding` is the
|
|
/// caller-supplied retrieval embedding — typically the finding's intent
|
|
/// (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(
|
|
&self,
|
|
index: &ControlIndex,
|
|
region: &CandidateRegion,
|
|
query_embedding: &[f64],
|
|
k: usize,
|
|
repo_id: &str,
|
|
) -> Vec<Finding> {
|
|
let candidates = index.nearest(query_embedding, k);
|
|
let mut findings = Vec::new();
|
|
for spec in &candidates {
|
|
let verdict = self.judge.judge(spec, region).await;
|
|
if let Some(finding) = ground(spec, region, &verdict, repo_id) {
|
|
findings.push(finding);
|
|
}
|
|
}
|
|
findings
|
|
}
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
use compliance_core::control_check::{ControlCheckSpec, LlmVerdict};
|
|
use compliance_core::models::finding::Severity;
|
|
|
|
struct StubJudge {
|
|
verdict: LlmVerdict,
|
|
}
|
|
impl ControlJudge for StubJudge {
|
|
async fn judge(&self, _s: &ControlCheckSpec, _r: &CandidateRegion) -> LlmVerdict {
|
|
self.verdict.clone()
|
|
}
|
|
}
|
|
|
|
fn spec(id: &str) -> ControlCheckSpec {
|
|
ControlCheckSpec {
|
|
control_id: id.into(),
|
|
title: id.into(),
|
|
requirement: id.into(),
|
|
default_cwe: None,
|
|
severity: Severity::Medium,
|
|
}
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn retrieves_then_grounds_the_nearest_control() {
|
|
let index = ControlIndex::from_embeddings(vec![
|
|
(spec("mc-near"), vec![1.0, 0.0]),
|
|
(spec("mc-far"), vec![0.0, 1.0]),
|
|
]);
|
|
let checker = SemanticControlChecker::new(StubJudge {
|
|
verdict: LlmVerdict {
|
|
violates: true,
|
|
snippet: "PASSWORD = \"admin\"".into(),
|
|
cwe: None,
|
|
confidence: 0.9,
|
|
},
|
|
});
|
|
let region = CandidateRegion {
|
|
file: "src/auth.py".into(),
|
|
start_line: 1,
|
|
content: "PASSWORD = \"admin\"\n".into(),
|
|
};
|
|
// Query embedding nearest to mc-near; k=1 → only mc-near is judged.
|
|
let findings = checker
|
|
.check(&index, ®ion, &[0.95, 0.05], 1, "repo")
|
|
.await;
|
|
assert_eq!(findings.len(), 1);
|
|
assert_eq!(findings[0].control_refs, vec!["mc-near".to_string()]);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn ungrounded_verdict_is_dropped() {
|
|
let index = ControlIndex::from_embeddings(vec![(spec("mc-near"), vec![1.0, 0.0])]);
|
|
let checker = SemanticControlChecker::new(StubJudge {
|
|
verdict: LlmVerdict {
|
|
violates: true,
|
|
snippet: "not in the region".into(),
|
|
cwe: None,
|
|
confidence: 0.9,
|
|
},
|
|
});
|
|
let region = CandidateRegion {
|
|
file: "f".into(),
|
|
start_line: 1,
|
|
content: "real code\n".into(),
|
|
};
|
|
let findings = checker.check(&index, ®ion, &[1.0, 0.0], 1, "repo").await;
|
|
assert!(findings.is_empty());
|
|
}
|
|
}
|