1use serde_json::{json, Value};
16
17pub const POLICY_ENV: &str = "HATCHERY_M4A_ATTACH";
18pub const PROVIDERS_ENV: &str = "HATCHERY_M4A_ATTACH_PROVIDERS";
19const MAX_GENERATIONS: usize = 12;
20
21#[derive(Debug, Clone, PartialEq, Eq)]
22pub struct Proc {
23 pub pid: u32,
24 pub ppid: u32,
25 pub name: String,
26}
27
28pub fn provider_images(provider: &str) -> &'static [&'static str] {
30 match provider {
31 "grok" => &["grok"],
32 "codex" => &["codex"],
33 "kimi" => &["kimi", "kimi-code"],
34 "claude" => &["claude"],
35 _ => &[],
36 }
37}
38
39fn base_name(name: &str) -> String {
40 let n = name.rsplit(['/', '\\']).next().unwrap_or(name).to_ascii_lowercase();
41 n.strip_suffix(".exe").map(str::to_owned).unwrap_or(n)
42}
43
44pub fn provider_ancestor(table: &[Proc], pid: u32, provider: &str) -> Option<u32> {
46 let images = provider_images(provider);
47 let mut current = pid;
48 for _ in 0..MAX_GENERATIONS {
49 let p = table.iter().find(|p| p.pid == current)?;
50 let name = base_name(&p.name);
51 if images.iter().any(|i| name == *i || name.starts_with(&format!("{i}-"))) {
52 return Some(p.pid);
53 }
54 if p.ppid == 0 || p.ppid == p.pid {
55 return None;
56 }
57 current = p.ppid;
58 }
59 None
60}
61
62pub fn policy_allows(provider: &str) -> Result<(), &'static str> {
63 let mode = std::env::var(POLICY_ENV).unwrap_or_default().trim().to_ascii_lowercase();
64 if mode != "auto" {
65 return Err("adoption policy is off on this node");
66 }
67 if let Ok(list) = std::env::var(PROVIDERS_ENV) {
68 if !list.split(',').map(str::trim).any(|p| p == provider) {
69 return Err("provider is not allowed by the adoption policy");
70 }
71 }
72 Ok(())
73}
74
75pub fn query_get(query: &str, key: &str) -> Option<String> {
76 query.split('&').find_map(|kv| {
77 let (k, v) = kv.split_once('=')?;
78 (k == key).then(|| percent_decode(v))
79 })
80}
81
82fn percent_decode(v: &str) -> String {
83 let b = v.as_bytes();
84 let mut out = Vec::with_capacity(b.len());
85 let mut i = 0;
86 while i < b.len() {
87 if b[i] == b'%' && i + 2 < b.len() {
88 if let Ok(x) = u8::from_str_radix(std::str::from_utf8(&b[i + 1..i + 3]).unwrap_or(""), 16) {
89 out.push(x);
90 i += 3;
91 continue;
92 }
93 }
94 out.push(if b[i] == b'+' { b' ' } else { b[i] });
95 i += 1;
96 }
97 String::from_utf8_lossy(&out).into_owned()
98}
99
100pub fn valid_session(s: &str) -> bool {
101 !s.is_empty() && s.len() <= 128 && s.bytes().all(|b| b.is_ascii_alphanumeric() || b"-_.:".contains(&b))
102}
103
104pub fn handle(
106 query: &str,
107 accounts: Option<&crate::m4a_accounts::M4aSessionAccounts>,
108 table: &[Proc],
109) -> (u16, Value) {
110 let Some(provider) = query_get(query, "provider") else {
111 return (400, json!({"error": "provider is required"}));
112 };
113 let Some(m4a_provider) = crate::m4a_accounts::m4a_provider(&provider) else {
114 return (400, json!({"error": "unknown provider"}));
115 };
116 if let Err(reason) = policy_allows(&provider) {
117 tracing::info!(%provider, "m4a attach refused by policy");
118 return (403, json!({"error": reason, "state": "pending"}));
119 }
120 let Some(pid) = query_get(query, "pid").and_then(|p| p.parse::<u32>().ok()) else {
121 return (400, json!({"error": "pid is required"}));
122 };
123 let Some(root) = provider_ancestor(table, pid, &provider) else {
124 tracing::warn!(%provider, pid, "m4a attach refused: caller is not inside a live provider process");
125 return (403, json!({"error": "caller is not inside a live process of that provider"}));
126 };
127 let session = query_get(query, "session");
128 if session.as_deref().is_some_and(|s| !valid_session(s)) {
129 return (400, json!({"error": "invalid session id"}));
130 }
131 let Some(accounts) = accounts else {
132 return (503, json!({"error": "m4a accounts are not configured on this node"}));
133 };
134 let session = crate::m4a_adopt::resolve_session_in(session, &provider, root, crate::m4a_adopt::process_cmdline, Some(accounts.base()))
137 .filter(|s| valid_session(s));
138 let key = match &session {
139 Some(s) => format!("attach-{provider}-{s}"),
140 None => format!("attach-{provider}-pid{root}"),
141 };
142 match accounts.provision(m4a_provider, &key) {
143 Ok(account) => {
144 tracing::info!(%provider, provider_pid = root, nick = %account.nick, "m4a account attached to a live session");
145 (200, account.attach_reply())
146 }
147 Err(error) => {
148 tracing::warn!(%provider, %error, "m4a attach provision failed");
149 (500, json!({"error": "account provision failed"}))
150 }
151 }
152}
153
154#[cfg(target_os = "linux")]
156pub fn process_table() -> Vec<Proc> {
157 let mut out = Vec::new();
158 let Ok(dir) = std::fs::read_dir("/proc") else { return out };
159 for e in dir.flatten() {
160 let Some(pid) = e.file_name().to_str().and_then(|s| s.parse::<u32>().ok()) else { continue };
161 let Ok(stat) = std::fs::read_to_string(format!("/proc/{pid}/stat")) else { continue };
162 let (Some(l), Some(r)) = (stat.find('('), stat.rfind(')')) else { continue };
163 let ppid = stat[r + 1..].split_whitespace().nth(1).and_then(|s| s.parse().ok()).unwrap_or(0);
164 let name = std::fs::read_link(format!("/proc/{pid}/exe"))
166 .ok()
167 .and_then(|p| p.file_name().map(|n| n.to_string_lossy().into_owned()))
168 .unwrap_or_else(|| stat[l + 1..r].to_owned());
169 out.push(Proc { pid, ppid, name });
170 }
171 out
172}
173
174#[cfg(windows)]
175pub fn process_table() -> Vec<Proc> {
176 use windows_sys::Win32::Foundation::{CloseHandle, INVALID_HANDLE_VALUE};
177 use windows_sys::Win32::System::Diagnostics::ToolHelp::{
178 CreateToolhelp32Snapshot, Process32FirstW, Process32NextW, PROCESSENTRY32W, TH32CS_SNAPPROCESS,
179 };
180 let mut out = Vec::new();
181 unsafe {
182 let snap = CreateToolhelp32Snapshot(TH32CS_SNAPPROCESS, 0);
183 if snap == INVALID_HANDLE_VALUE {
184 return out;
185 }
186 let mut entry: PROCESSENTRY32W = std::mem::zeroed();
187 entry.dwSize = std::mem::size_of::<PROCESSENTRY32W>() as u32;
188 let mut ok = Process32FirstW(snap, &mut entry) != 0;
189 while ok {
190 let len = entry.szExeFile.iter().position(|c| *c == 0).unwrap_or(entry.szExeFile.len());
191 out.push(Proc {
192 pid: entry.th32ProcessID,
193 ppid: entry.th32ParentProcessID,
194 name: String::from_utf16_lossy(&entry.szExeFile[..len]),
195 });
196 ok = Process32NextW(snap, &mut entry) != 0;
197 }
198 CloseHandle(snap);
199 }
200 out
201}
202
203#[cfg(not(any(target_os = "linux", windows)))]
204pub fn process_table() -> Vec<Proc> {
205 Vec::new()
206}
207
208#[cfg(test)]
209mod tests {
210 use super::*;
211
212 fn table() -> Vec<Proc> {
213 vec![
214 Proc { pid: 1, ppid: 0, name: "init".into() },
215 Proc { pid: 10, ppid: 1, name: "Grok.EXE".into() },
216 Proc { pid: 11, ppid: 10, name: "cmd.exe".into() },
217 Proc { pid: 12, ppid: 11, name: "m4a-agent.exe".into() },
218 Proc { pid: 20, ppid: 1, name: "bash".into() },
219 Proc { pid: 21, ppid: 20, name: "m4a-agent".into() },
220 ]
221 }
222
223 #[test]
224 fn ancestry_finds_the_provider_and_refuses_strangers() {
225 let t = table();
226 assert_eq!(provider_ancestor(&t, 12, "grok"), Some(10));
227 assert_eq!(provider_ancestor(&t, 12, "codex"), None);
228 assert_eq!(provider_ancestor(&t, 21, "grok"), None);
229 assert_eq!(provider_ancestor(&t, 999, "grok"), None);
230 }
231
232 #[test]
233 fn query_parsing_decodes_and_validates() {
234 assert_eq!(query_get("provider=grok&session=a%3Ab", "session").as_deref(), Some("a:b"));
235 assert!(valid_session("0aa-b_c.d:e"));
236 assert!(!valid_session("../x"));
237 assert!(!valid_session(""));
238 }
239
240 #[test]
242 fn policy_and_process_gates_come_before_any_account() {
243 std::env::remove_var(POLICY_ENV);
244 let (s, v) = handle("provider=grok&pid=12", None, &table());
245 assert_eq!((s, v["state"].as_str()), (403, Some("pending")));
246 std::env::set_var(POLICY_ENV, "auto");
247 std::env::set_var(PROVIDERS_ENV, "codex");
248 assert_eq!(handle("provider=grok&pid=12", None, &table()).0, 403);
249 std::env::remove_var(PROVIDERS_ENV);
250 assert_eq!(handle("provider=grok&pid=21", None, &table()).0, 403);
251 assert_eq!(handle("provider=grok&pid=12&session=../x", None, &table()).0, 400);
252 assert_eq!(handle("provider=grok&pid=12", None, &table()).0, 503);
253 assert_eq!(handle("provider=nope&pid=12", None, &table()).0, 400);
254 std::env::remove_var(POLICY_ENV);
255 }
256
257 #[cfg(target_os = "linux")]
258 #[test]
259 fn live_table_contains_this_process() {
260 let me = std::process::id();
261 assert!(process_table().iter().any(|p| p.pid == me));
262 }
263}