use std::collections::VecDeque;
use std::sync::{Mutex, OnceLock};
use std::time::Instant;
use regex::Regex;
use serde_json::{Map, Value};
#[must_use]
pub fn write_payload(
tool: &str,
args: Option<&Map<String, Value>>,
) -> Option<(String, &'static str)> {
let get = |k: &str| args?.get(k)?.as_str().map(String::from);
match tool {
"ctx_edit" => get("new_string").map(|s| (s, "Write")),
"ctx_patch" => patch_payload(args).map(|s| (s, "Write")),
"ctx_shell" | "ctx_execute" => get("command").map(|s| (s, "Action")),
_ => None,
}
}
fn patch_payload(args: Option<&Map<String, Value>>) -> Option<String> {
let map = args?;
let mut parts: Vec<&str> = Vec::new();
if let Some(s) = map.get("new_text").and_then(Value::as_str) {
parts.push(s);
}
if let Some(ops) = map.get("ops").and_then(Value::as_array) {
for op in ops {
if let Some(s) = op.get("new_text").and_then(Value::as_str) {
parts.push(s);
}
}
}
if parts.is_empty() {
None
} else {
Some(parts.join("\n"))
}
}
pub struct EgressConfig {
forbidden: Vec<(String, Regex)>,
block_secrets: bool,
pub max_writes_per_min: Option<u32>,
}
impl Default for EgressConfig {
fn default() -> Self {
Self::off()
}
}
impl EgressConfig {
#[must_use]
pub fn off() -> Self {
Self {
forbidden: Vec::new(),
block_secrets: false,
max_writes_per_min: None,
}
}
#[must_use]
pub fn new(
forbidden_patterns: &[String],
block_secrets: bool,
max_writes_per_min: Option<u32>,
) -> Self {
let forbidden = forbidden_patterns
.iter()
.filter_map(|p| Regex::new(p).ok().map(|re| (p.clone(), re)))
.collect();
Self {
forbidden,
block_secrets,
max_writes_per_min,
}
}
#[must_use]
pub fn is_active(&self) -> bool {
!self.forbidden.is_empty() || self.block_secrets || self.max_writes_per_min.is_some()
}
#[must_use]
pub fn check_content(&self, content: &str, redaction: &[(String, Regex)]) -> Option<String> {
for (source, re) in &self.forbidden {
if re.is_match(content) {
return Some(format!("forbidden-pattern:{source}"));
}
}
if self.block_secrets {
let (_, hits) = crate::core::redaction::redact_with_patterns(content, redaction);
if hits > 0 {
return Some("secret".to_string());
}
if let Some((class, _)) = crate::core::input_filters::pii::detect(content).first() {
return Some(format!("pii:{class}"));
}
}
None
}
}
#[must_use]
pub fn check_rate(max_per_min: u32) -> bool {
check_rate_at(rate_state(), Instant::now(), max_per_min)
}
fn check_rate_at(state: &Mutex<VecDeque<Instant>>, now: Instant, max_per_min: u32) -> bool {
let Ok(mut q) = state.lock() else {
return false;
};
within_limit(&mut q, now, max_per_min)
}
fn rate_state() -> &'static Mutex<VecDeque<Instant>> {
static STATE: OnceLock<Mutex<VecDeque<Instant>>> = OnceLock::new();
STATE.get_or_init(|| Mutex::new(VecDeque::new()))
}
fn within_limit(q: &mut VecDeque<Instant>, now: Instant, max_per_min: u32) -> bool {
while let Some(&front) = q.front() {
if now.duration_since(front).as_secs() >= 60 {
q.pop_front();
} else {
break;
}
}
if u32::try_from(q.len()).unwrap_or(u32::MAX) >= max_per_min {
return false;
}
q.push_back(now);
true
}
#[cfg(test)]
mod tests {
use super::*;
use std::time::Duration;
fn cfg(patterns: &[&str], block_secrets: bool) -> EgressConfig {
let pats: Vec<String> = patterns.iter().map(|s| (*s).to_string()).collect();
EgressConfig::new(&pats, block_secrets, None)
}
#[test]
fn off_config_is_inactive() {
assert!(!EgressConfig::off().is_active());
}
#[test]
fn forbidden_pattern_blocks_action() {
let c = cfg(&[r"prod\.db\.internal"], false);
let reason = c.check_content("psql postgres://prod.db.internal/main", &[]);
assert_eq!(
reason.as_deref(),
Some("forbidden-pattern:prod\\.db\\.internal")
);
}
#[test]
fn clean_content_is_allowed() {
let c = cfg(&[r"prod\.db\.internal"], true);
assert!(
c.check_content("fn main() { println!(\"hi\"); }", &[])
.is_none()
);
}
#[test]
fn block_secrets_catches_pii() {
let c = cfg(&[], true);
let reason = c.check_content("email jane@example.com into config", &[]);
assert_eq!(reason.as_deref(), Some("pii:email"));
}
#[test]
fn block_secrets_catches_redaction_pattern() {
let c = cfg(&[], true);
let redaction = vec![("employee_id".to_string(), Regex::new(r"EMP-\d{4}").unwrap())];
let reason = c.check_content("commit by EMP-1234", &redaction);
assert_eq!(reason.as_deref(), Some("secret"));
}
#[test]
fn rate_limit_triggers_after_max() {
let mut q = VecDeque::new();
let now = Instant::now();
assert!(within_limit(&mut q, now, 2));
assert!(within_limit(&mut q, now, 2));
assert!(!within_limit(&mut q, now, 2));
}
#[test]
fn rate_limit_window_slides() {
let mut q = VecDeque::new();
let base = Instant::now();
assert!(within_limit(&mut q, base, 1));
assert!(!within_limit(&mut q, base, 1));
assert!(within_limit(&mut q, base + Duration::from_secs(61), 1));
}
#[test]
fn poisoned_rate_state_fails_closed_without_panicking() {
let state = std::sync::Arc::new(Mutex::new(VecDeque::new()));
let poisoner = std::sync::Arc::clone(&state);
let _ = std::thread::spawn(move || {
let _guard = poisoner.lock().expect("fresh mutex");
panic!("poison rate state");
})
.join();
assert!(!check_rate_at(&state, Instant::now(), 1));
}
fn args(v: Value) -> Map<String, Value> {
match v {
Value::Object(m) => m,
_ => panic!("expected object"),
}
}
#[test]
fn write_payload_covers_edit_shell_and_execute() {
let edit = args(serde_json::json!({"new_string": "body"}));
assert_eq!(
write_payload("ctx_edit", Some(&edit)),
Some(("body".to_string(), "Write"))
);
let sh = args(serde_json::json!({"command": "rm -rf /tmp/x"}));
assert_eq!(
write_payload("ctx_shell", Some(&sh)),
Some(("rm -rf /tmp/x".to_string(), "Action"))
);
assert_eq!(
write_payload("ctx_execute", Some(&sh)),
Some(("rm -rf /tmp/x".to_string(), "Action"))
);
assert_eq!(write_payload("ctx_read", Some(&sh)), None);
}
#[test]
fn write_payload_collects_every_patch_body() {
let single = args(serde_json::json!({"op": "set_line", "new_text": "top"}));
assert_eq!(
write_payload("ctx_patch", Some(&single)),
Some(("top".to_string(), "Write"))
);
let symbol = args(serde_json::json!({"op": "replace_symbol", "new_text": "fn x() {}"}));
assert_eq!(
write_payload("ctx_patch", Some(&symbol)),
Some(("fn x() {}".to_string(), "Write"))
);
let batch = args(serde_json::json!({"ops": [
{"op": "set_line", "line": 1, "hash": "aa", "new_text": "first"},
{"op": "insert_after", "line": 2, "hash": "bb", "new_text": "second"}
]}));
let (payload, kind) = write_payload("ctx_patch", Some(&batch)).unwrap();
assert_eq!(kind, "Write");
assert!(payload.contains("first") && payload.contains("second"));
let empty = args(serde_json::json!({"op": "delete", "line": 3, "hash": "cc"}));
assert_eq!(write_payload("ctx_patch", Some(&empty)), None);
}
#[test]
fn patch_payload_is_checkable_content() {
let c = cfg(&[r"prod\.db\.internal"], false);
let batch = args(serde_json::json!({"ops": [
{"op": "set_line", "line": 1, "hash": "aa", "new_text": "safe"},
{"op": "set_line", "line": 9, "hash": "bb", "new_text": "url = prod.db.internal"}
]}));
let (payload, _) = write_payload("ctx_patch", Some(&batch)).unwrap();
assert!(c.check_content(&payload, &[]).is_some());
}
}