use axiom_engine::session_recorder::{ExchangeRecord, read_session, scrub, sessions_dir};
use serde_json::Value;
use sha2::{Digest, Sha256};
use std::path::{Path, PathBuf};
pub fn default_corpus_dir() -> PathBuf {
PathBuf::from(concat!(env!("CARGO_MANIFEST_DIR"), "/../bench/corpus"))
}
fn session_source(input: &str) -> PathBuf {
let candidate = PathBuf::from(input);
if candidate.exists() {
return candidate;
}
let file = if input.ends_with(".jsonl") {
input.to_string()
} else {
format!("{input}.jsonl")
};
sessions_dir().join(file)
}
fn looks_secret(s: &str) -> bool {
let lower = s.to_ascii_lowercase();
let upper = s.to_ascii_uppercase();
lower.starts_with("bearer ")
|| lower.starts_with("-----begin")
|| upper.starts_with("AKIA")
|| s.starts_with("AIza")
|| s.len() >= 20
&& ["sk-", "sk_live_", "sk_test_", "ghp_", "gho_", "xoxb-", "eyj"]
.iter()
.any(|prefix| lower.contains(prefix))
}
fn contains_secret_miss(value: &Value) -> bool {
match value {
Value::String(s) => looks_secret(s),
Value::Array(items) => items.iter().any(contains_secret_miss),
Value::Object(map) => map.values().any(contains_secret_miss),
_ => false,
}
}
fn scrub_record(record: &ExchangeRecord) -> ExchangeRecord {
ExchangeRecord {
request: scrub(&record.request),
response: scrub(&record.response),
..record.clone()
}
}
fn scrubbed_jsonl(records: &[ExchangeRecord]) -> Result<Vec<u8>, String> {
let mut out = Vec::new();
for record in records {
let scrubbed = scrub_record(record);
if contains_secret_miss(&scrubbed.request) || contains_secret_miss(&scrubbed.response) {
return Err(format!(
"secret-looking value survived scrub in session {}",
scrubbed.session_id
));
}
let line = serde_json::to_string(&scrubbed)
.map_err(|error| format!("serialize scrubbed record failed: {error}"))?;
out.extend_from_slice(line.as_bytes());
out.push(b'\n');
}
Ok(out)
}
fn sha256_hex(bytes: &[u8]) -> String {
format!("{:x}", Sha256::digest(bytes))
}
pub fn promote_to(source: &Path, corpus_dir: &Path) -> Result<PathBuf, String> {
let records = read_session(source)
.map_err(|error| format!("read {} failed: {error}", source.display()))?;
if records.is_empty() {
return Err(format!("{} contains no exchange records", source.display()));
}
let bytes = scrubbed_jsonl(&records)?;
let hash = sha256_hex(&bytes);
std::fs::create_dir_all(corpus_dir)
.map_err(|error| format!("create {} failed: {error}", corpus_dir.display()))?;
let path = corpus_dir.join(format!("{}.jsonl", &hash[..16]));
std::fs::write(&path, bytes)
.map_err(|error| format!("write {} failed: {error}", path.display()))?;
Ok(path)
}
pub fn list_corpus(corpus_dir: &Path) -> Result<Vec<(PathBuf, usize, String)>, String> {
let mut rows = Vec::new();
let entries = match std::fs::read_dir(corpus_dir) {
Ok(entries) => entries,
Err(error) if error.kind() == std::io::ErrorKind::NotFound => return Ok(rows),
Err(error) => return Err(format!("read {} failed: {error}", corpus_dir.display())),
};
for entry in entries {
let entry = entry.map_err(|error| format!("read corpus entry failed: {error}"))?;
let path = entry.path();
if path.extension().and_then(|value| value.to_str()) != Some("jsonl") {
continue;
}
let bytes = std::fs::read(&path)
.map_err(|error| format!("read {} failed: {error}", path.display()))?;
let records = read_session(&path)
.map_err(|error| format!("parse {} failed: {error}", path.display()))?;
rows.push((path, records.len(), sha256_hex(&bytes)));
}
rows.sort_by(|a, b| a.0.cmp(&b.0));
Ok(rows)
}
pub fn run(args: &[String]) -> Result<(), String> {
match args.first().map(String::as_str) {
Some("promote") => {
let Some(input) = args.get(1) else {
return Err("usage: axiombench corpus promote <session-id|path>".into());
};
let source = session_source(input);
let path = promote_to(&source, &default_corpus_dir())?;
let records = read_session(&path)
.map_err(|error| format!("read promoted {} failed: {error}", path.display()))?;
let bytes = std::fs::read(&path)
.map_err(|error| format!("read {} failed: {error}", path.display()))?;
println!(
"promoted source={} records={} sha256={} -> {}",
source.display(),
records.len(),
sha256_hex(&bytes),
path.display()
);
Ok(())
}
Some("list") => {
println!("file\trecords\tsha256");
for (path, records, hash) in list_corpus(&default_corpus_dir())? {
println!("{}\t{}\t{}", path.display(), records, hash);
}
Ok(())
}
_ => Err("usage: axiombench corpus <promote|list> ...".into()),
}
}
#[cfg(test)]
mod tests {
use super::*;
use serde_json::json;
fn record(session_id: &str, request: Value) -> ExchangeRecord {
ExchangeRecord {
ts: 1,
endpoint: "/v1/responses".into(),
session_id: session_id.into(),
request,
response: json!({"status": 200}),
compressed: false,
}
}
#[test]
fn promote_and_list_roundtrip() {
let temp = tempfile::tempdir().unwrap();
let source_dir = temp.path().join("source");
let corpus_dir = temp.path().join("corpus");
std::fs::create_dir_all(&source_dir).unwrap();
let source = source_dir.join("source.jsonl");
let rec = record(
"s1",
json!({"input": "hello", "authorization": "Bearer secret"}),
);
std::fs::write(
&source,
format!("{}\n", serde_json::to_string(&rec).unwrap()),
)
.unwrap();
let out = promote_to(&source, &corpus_dir).unwrap();
assert!(out.exists());
let promoted = std::fs::read_to_string(&out).unwrap();
assert!(promoted.contains("[REDACTED]"));
let rows = list_corpus(&corpus_dir).unwrap();
assert_eq!(rows.len(), 1);
assert!(
rows.iter()
.any(|(path, records, _)| path == &out && *records == 1)
);
}
#[test]
fn promotion_rejects_scrub_miss() {
let temp = tempfile::tempdir().unwrap();
let source = temp.path().join("source.jsonl");
let rec = record(
"s1",
json!({"input": "token xoxb-abcdefghijklmnopqrstuvwxyz"}),
);
std::fs::write(
&source,
format!("{}\n", serde_json::to_string(&rec).unwrap()),
)
.unwrap();
let error = promote_to(&source, temp.path()).unwrap_err();
assert!(error.contains("secret-looking value survived scrub"));
}
}