aft/bash_rewrite/
dispatch.rs

1use std::collections::HashMap;
2use std::fs::OpenOptions;
3use std::io::Write;
4use std::sync::{Mutex, OnceLock};
5
6use serde::{Deserialize, Serialize};
7
8use crate::bash_rewrite::catalog::ControlRole;
9use crate::bash_rewrite::rules::{
10    CatAppendRule, CatRule, FindRule, GrepRule, HeadRule, LsRule, RgRule, SedRule, TailRule,
11};
12use crate::bash_rewrite::RewriteRule;
13use crate::context::AppContext;
14use crate::protocol::Response;
15
16#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
17#[serde(tag = "kind", rename_all = "snake_case")]
18pub enum DispatchRoute {
19    Rewritten {
20        rule_id: String,
21        branch_id: String,
22        decision_class_id: String,
23    },
24    Native {
25        role: String,
26        branch_id: String,
27        reason: String,
28    },
29}
30
31#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
32pub struct DispatchRecord {
33    pub request_id: String,
34    pub route: DispatchRoute,
35}
36
37static ROUTE_RECORDS: OnceLock<Mutex<HashMap<String, DispatchRecord>>> = OnceLock::new();
38
39fn route_records() -> &'static Mutex<HashMap<String, DispatchRecord>> {
40    ROUTE_RECORDS.get_or_init(|| Mutex::new(HashMap::new()))
41}
42
43fn store_record(record: DispatchRecord) {
44    if let Ok(mut records) = route_records().lock() {
45        records.insert(record.request_id.clone(), record.clone());
46    }
47
48    // The sidecar is opt-in for the differential child process. Normal AFT
49    // responses remain byte-for-byte free of test metadata.
50    if let Some(path) = std::env::var_os("AFT_BASH_REWRITE_ROUTE_RECORD") {
51        if let Ok(mut file) = OpenOptions::new().create(true).append(true).open(path) {
52            if let Ok(line) = serde_json::to_string(&record) {
53                let _ = writeln!(file, "{line}");
54            }
55        }
56    }
57}
58
59pub fn route_record(request_id: &str) -> Option<DispatchRecord> {
60    route_records()
61        .lock()
62        .ok()
63        .and_then(|records| records.get(request_id).cloned())
64}
65
66pub fn take_route_record(request_id: &str) -> Option<DispatchRecord> {
67    route_records()
68        .lock()
69        .ok()
70        .and_then(|mut records| records.remove(request_id))
71}
72
73pub fn record_native(request_id: &str, role: ControlRole, branch_id: &str, reason: &str) {
74    store_record(DispatchRecord {
75        request_id: request_id.to_string(),
76        route: DispatchRoute::Native {
77            role: role.id().to_string(),
78            branch_id: branch_id.to_string(),
79            reason: reason.to_string(),
80        },
81    });
82}
83
84pub fn dispatch(command: &str, session_id: Option<&str>, ctx: &AppContext) -> Option<Response> {
85    dispatch_for_request(command, "bash_rewrite", session_id, ctx)
86}
87
88pub fn dispatch_for_request(
89    command: &str,
90    request_id: &str,
91    session_id: Option<&str>,
92    ctx: &AppContext,
93) -> Option<Response> {
94    if !ctx.config().experimental_bash_rewrite {
95        record_native(
96            request_id,
97            ControlRole::Native,
98            "dispatch.native.no_rule",
99            "experimental bash rewriting is disabled",
100        );
101        return None;
102    }
103
104    let rules: [&dyn RewriteRule; 9] = [
105        &GrepRule,
106        &RgRule,
107        &FindRule,
108        &CatRule,
109        &HeadRule,
110        &TailRule,
111        &CatAppendRule,
112        &SedRule,
113        &LsRule,
114    ];
115
116    for rule in rules {
117        let decision = rule.decide(command, request_id, session_id, ctx);
118        match decision {
119            crate::bash_rewrite::RewriteDecision::Accept(request) => {
120                store_record(DispatchRecord {
121                    request_id: request_id.to_string(),
122                    route: DispatchRoute::Rewritten {
123                        rule_id: request.rule_id.to_string(),
124                        branch_id: request.branch_id.to_string(),
125                        decision_class_id: request.decision_class_id.to_string(),
126                    },
127                });
128                // Do not turn a handler failure into a native execution. The
129                // handler has already begun and a second execution could
130                // duplicate a mutation such as cat_append.
131                return Some(rule.execute(&request, ctx));
132            }
133            crate::bash_rewrite::RewriteDecision::Decline(reason) => {
134                if reason.rule_id.is_some() {
135                    crate::slog_debug!(
136                        "bash rewrite rule {} declined before execution: {}",
137                        reason.rule_id.unwrap_or("unknown"),
138                        reason.reason
139                    );
140                }
141            }
142        }
143    }
144
145    record_native(
146        request_id,
147        ControlRole::Native,
148        "dispatch.native.no_rule",
149        "no rewrite rule accepted the command shape",
150    );
151    None
152}
153
154#[cfg(test)]
155mod tests {
156    use super::*;
157    use crate::bash_rewrite::catalog::ControlRole;
158
159    #[test]
160    fn route_records_are_request_correlated_and_replaceable() {
161        record_native(
162            "route-test",
163            ControlRole::Native,
164            "dispatch.native.no_rule",
165            "test",
166        );
167        assert_eq!(
168            route_record("route-test").expect("route record").request_id,
169            "route-test"
170        );
171        record_native(
172            "route-test",
173            ControlRole::Sandbox,
174            "dispatch.native.sandbox",
175            "test",
176        );
177        assert_eq!(
178            route_record("route-test").expect("replacement").route,
179            DispatchRoute::Native {
180                role: "sandbox".to_string(),
181                branch_id: "dispatch.native.sandbox".to_string(),
182                reason: "test".to_string(),
183            }
184        );
185        let _ = take_route_record("route-test");
186    }
187}
aft/bash_rewrite/dispatch.rs

aft/bash_rewrite/
dispatch.rs
