Compare commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
7cd4ffdaab |
@@ -234,22 +234,13 @@ 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;
|
||||||
};
|
};
|
||||||
// Retrieve on the finding's intent + the code, not the region alone: two
|
let region_emb = match llm.embed(vec![region.content.clone()]).await {
|
||||||
// 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, "query embed failed; skipping finding");
|
tracing::warn!(error = %e, "region embed failed; skipping finding");
|
||||||
continue;
|
continue;
|
||||||
}
|
}
|
||||||
};
|
};
|
||||||
@@ -257,7 +248,7 @@ pub async fn semantic_stamp_findings(
|
|||||||
.check(
|
.check(
|
||||||
&index,
|
&index,
|
||||||
®ion,
|
®ion,
|
||||||
&query_emb,
|
®ion_emb,
|
||||||
SEMANTIC_TOP_K,
|
SEMANTIC_TOP_K,
|
||||||
&finding.repo_id,
|
&finding.repo_id,
|
||||||
)
|
)
|
||||||
|
|||||||
@@ -22,20 +22,18 @@ impl<J: ControlJudge> SemanticControlChecker<J> {
|
|||||||
Self { judge }
|
Self { judge }
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Map a code region to the controls it violates. `query_embedding` is the
|
/// Map a code region to the controls it violates. `region_embedding` is the
|
||||||
/// caller-supplied retrieval embedding — typically the finding's intent
|
/// region's embedding (the caller computes it via the LLM); the top-`k`
|
||||||
/// (title/description) plus the region, so retrieval keys on what the finding
|
/// nearest controls in `index` are judged and grounded.
|
||||||
/// 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,
|
||||||
query_embedding: &[f64],
|
region_embedding: &[f64],
|
||||||
k: usize,
|
k: usize,
|
||||||
repo_id: &str,
|
repo_id: &str,
|
||||||
) -> Vec<Finding> {
|
) -> Vec<Finding> {
|
||||||
let candidates = index.nearest(query_embedding, k);
|
let candidates = index.nearest(region_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;
|
||||||
|
|||||||
@@ -22,11 +22,6 @@ struct EmbeddingData {
|
|||||||
index: usize,
|
index: usize,
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Max inputs per embedding request. The bge/OpenAI-like backends cap the input
|
|
||||||
/// array (bge-multilingual-gemma2 rejects >25 with "batch size overflow"), so we
|
|
||||||
/// chunk larger corpora — a whole control catalog (~1.8k) would otherwise 500.
|
|
||||||
const EMBED_BATCH_SIZE: usize = 16;
|
|
||||||
|
|
||||||
// ── Embedding implementation ───────────────────────────────────
|
// ── Embedding implementation ───────────────────────────────────
|
||||||
|
|
||||||
impl LlmClient {
|
impl LlmClient {
|
||||||
@@ -34,21 +29,8 @@ impl LlmClient {
|
|||||||
&self.embed_model
|
&self.embed_model
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Generate embeddings for a batch of texts, chunking into backend-sized
|
/// Generate embeddings for a batch of texts
|
||||||
/// requests and preserving input order across chunks.
|
|
||||||
pub async fn embed(&self, texts: Vec<String>) -> Result<Vec<Vec<f64>>, AgentError> {
|
pub async fn embed(&self, texts: Vec<String>) -> Result<Vec<Vec<f64>>, AgentError> {
|
||||||
if texts.is_empty() {
|
|
||||||
return Ok(Vec::new());
|
|
||||||
}
|
|
||||||
let mut out = Vec::with_capacity(texts.len());
|
|
||||||
for chunk in texts.chunks(EMBED_BATCH_SIZE) {
|
|
||||||
out.extend(self.embed_batch(chunk.to_vec()).await?);
|
|
||||||
}
|
|
||||||
Ok(out)
|
|
||||||
}
|
|
||||||
|
|
||||||
/// Embed one backend-sized batch (≤ [`EMBED_BATCH_SIZE`]) in a single request.
|
|
||||||
async fn embed_batch(&self, texts: Vec<String>) -> Result<Vec<Vec<f64>>, AgentError> {
|
|
||||||
let url = format!("{}/v1/embeddings", self.base_url.trim_end_matches('/'));
|
let url = format!("{}/v1/embeddings", self.base_url.trim_end_matches('/'));
|
||||||
|
|
||||||
let request_body = EmbeddingRequest {
|
let request_body = EmbeddingRequest {
|
||||||
@@ -90,33 +72,3 @@ impl LlmClient {
|
|||||||
Ok(data.into_iter().map(|d| d.embedding).collect())
|
Ok(data.into_iter().map(|d| d.embedding).collect())
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
#[cfg(test)]
|
|
||||||
mod tests {
|
|
||||||
use super::*;
|
|
||||||
use secrecy::SecretString;
|
|
||||||
|
|
||||||
fn client() -> LlmClient {
|
|
||||||
LlmClient::new(
|
|
||||||
"http://unused".into(),
|
|
||||||
SecretString::from(String::new()),
|
|
||||||
"m".into(),
|
|
||||||
"e".into(),
|
|
||||||
)
|
|
||||||
}
|
|
||||||
|
|
||||||
#[tokio::test]
|
|
||||||
async fn empty_input_makes_no_request() {
|
|
||||||
// Must short-circuit before any HTTP call (base_url is unroutable).
|
|
||||||
let out = client().embed(Vec::new()).await.unwrap();
|
|
||||||
assert!(out.is_empty());
|
|
||||||
}
|
|
||||||
|
|
||||||
#[test]
|
|
||||||
fn batch_size_is_within_backend_cap() {
|
|
||||||
assert!(
|
|
||||||
EMBED_BATCH_SIZE <= 25,
|
|
||||||
"must stay under the bge 25-input cap"
|
|
||||||
);
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|||||||
@@ -1,145 +0,0 @@
|
|||||||
//! 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