use std::collections::HashMap;
use std::fs::OpenOptions;
use std::io::Write;
use std::sync::{Mutex, OnceLock};
use serde::{Deserialize, Serialize};
use crate::bash_rewrite::catalog::ControlRole;
use crate::bash_rewrite::rules::{
CatAppendRule, CatRule, FindRule, GrepRule, LsRule, RgRule, SedRule,
};
use crate::bash_rewrite::RewriteRule;
use crate::context::AppContext;
use crate::protocol::Response;
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
#[serde(tag = "kind", rename_all = "snake_case")]
pub enum DispatchRoute {
Rewritten {
rule_id: String,
branch_id: String,
decision_class_id: String,
},
Native {
role: String,
branch_id: String,
reason: String,
},
}
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
pub struct DispatchRecord {
pub request_id: String,
pub route: DispatchRoute,
}
static ROUTE_RECORDS: OnceLock<Mutex<HashMap<String, DispatchRecord>>> = OnceLock::new();
fn route_records() -> &'static Mutex<HashMap<String, DispatchRecord>> {
ROUTE_RECORDS.get_or_init(|| Mutex::new(HashMap::new()))
}
fn store_record(record: DispatchRecord) {
if let Ok(mut records) = route_records().lock() {
records.insert(record.request_id.clone(), record.clone());
}
if let Some(path) = std::env::var_os("AFT_BASH_REWRITE_ROUTE_RECORD") {
if let Ok(mut file) = OpenOptions::new().create(true).append(true).open(path) {
if let Ok(line) = serde_json::to_string(&record) {
let _ = writeln!(file, "{line}");
}
}
}
}
pub fn route_record(request_id: &str) -> Option<DispatchRecord> {
route_records()
.lock()
.ok()
.and_then(|records| records.get(request_id).cloned())
}
pub fn take_route_record(request_id: &str) -> Option<DispatchRecord> {
route_records()
.lock()
.ok()
.and_then(|mut records| records.remove(request_id))
}
pub fn record_native(request_id: &str, role: ControlRole, branch_id: &str, reason: &str) {
store_record(DispatchRecord {
request_id: request_id.to_string(),
route: DispatchRoute::Native {
role: role.id().to_string(),
branch_id: branch_id.to_string(),
reason: reason.to_string(),
},
});
}
pub fn dispatch(command: &str, session_id: Option<&str>, ctx: &AppContext) -> Option<Response> {
dispatch_for_request(command, "bash_rewrite", session_id, ctx)
}
pub fn dispatch_for_request(
command: &str,
request_id: &str,
session_id: Option<&str>,
ctx: &AppContext,
) -> Option<Response> {
if !ctx.config().experimental_bash_rewrite {
record_native(
request_id,
ControlRole::Native,
"dispatch.native.no_rule",
"experimental bash rewriting is disabled",
);
return None;
}
let rules: [&dyn RewriteRule; 7] = [
&GrepRule,
&RgRule,
&FindRule,
&CatRule,
&CatAppendRule,
&SedRule,
&LsRule,
];
for rule in rules {
let decision = rule.decide(command, request_id, session_id, ctx);
match decision {
crate::bash_rewrite::RewriteDecision::Accept(request) => {
store_record(DispatchRecord {
request_id: request_id.to_string(),
route: DispatchRoute::Rewritten {
rule_id: request.rule_id.to_string(),
branch_id: request.branch_id.to_string(),
decision_class_id: request.decision_class_id.to_string(),
},
});
return Some(rule.execute(&request, ctx));
}
crate::bash_rewrite::RewriteDecision::Decline(reason) => {
if reason.rule_id.is_some() {
crate::slog_debug!(
"bash rewrite rule {} declined before execution: {}",
reason.rule_id.unwrap_or("unknown"),
reason.reason
);
}
}
}
}
record_native(
request_id,
ControlRole::Native,
"dispatch.native.no_rule",
"no rewrite rule accepted the command shape",
);
None
}
#[cfg(test)]
mod tests {
use super::*;
use crate::bash_rewrite::catalog::ControlRole;
#[test]
fn route_records_are_request_correlated_and_replaceable() {
record_native(
"route-test",
ControlRole::Native,
"dispatch.native.no_rule",
"test",
);
assert_eq!(
route_record("route-test").expect("route record").request_id,
"route-test"
);
record_native(
"route-test",
ControlRole::Sandbox,
"dispatch.native.sandbox",
"test",
);
assert_eq!(
route_record("route-test").expect("replacement").route,
DispatchRoute::Native {
role: "sandbox".to_string(),
branch_id: "dispatch.native.sandbox".to_string(),
reason: "test".to_string(),
}
);
let _ = take_route_record("route-test");
}
}