1use std::collections::BTreeMap;
22use std::path::{Path, PathBuf};
23
24use serde::{Deserialize, Serialize};
25use sha2::{Digest, Sha256};
26
27use crate::provider::ProviderKind;
28
29pub const SESSION_CAP_ENV: &str = "M4A_SESSION_CAP";
31
32#[derive(Debug, Clone, PartialEq, Eq)]
33pub struct SessionIdentity {
34 pub provider: ProviderKind,
35 pub provider_session_id: String,
37}
38
39impl SessionIdentity {
40 pub fn new(provider: ProviderKind, provider_session_id: &str) -> Result<Self, AccountError> {
41 let id = provider_session_id.trim();
42 if id.is_empty() || id.len() > 256 || id.chars().any(char::is_control) {
43 return Err(AccountError::BadIdentity);
44 }
45 Ok(Self { provider, provider_session_id: id.to_string() })
46 }
47
48 pub fn account_ref(&self) -> String {
50 let h = Sha256::digest(format!("m4a-session-account/v1\0{}\0{}", self.provider.id(), self.provider_session_id).as_bytes());
51 format!("{}-{}", self.provider.id(), hex(&h[..8]))
52 }
53}
54
55#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
56#[serde(rename_all = "snake_case")]
57pub enum ServerState {
58 Local,
60 InvitePending,
62 Enrolled,
64 Retired,
66}
67
68#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
69pub struct AccountEntry {
70 pub nick: String,
71 pub provider: String,
72 pub state: ServerState,
73 #[serde(default)]
75 pub cap_sha256: Option<String>,
76 #[serde(default, skip_serializing_if = "Option::is_none")]
79 pub os_uid: Option<u32>,
80}
81
82#[derive(Debug, Default, Serialize, Deserialize)]
83struct RegistryFile {
84 accounts: BTreeMap<String, AccountEntry>,
86}
87
88#[derive(Debug, thiserror::Error, PartialEq, Eq)]
89pub enum AccountError {
90 #[error("session identity is empty or malformed")]
91 BadIdentity,
92 #[error("nick is not a valid slug")]
93 BadNick,
94 #[error("no account for this session")]
95 Unknown,
96 #[error("capability refused")]
97 Refused,
98 #[error("session capability missing")]
99 Missing,
100 #[error("caller is not the session's OS user")]
101 WrongCaller,
102 #[error("account registry io: {0}")]
103 Io(String),
104}
105
106pub struct AccountRegistry {
107 base: PathBuf,
108 file: RegistryFile,
109}
110
111impl AccountRegistry {
112 pub fn open(base: &Path) -> Result<Self, AccountError> {
113 let dir = base.join(".m4a-accounts");
114 create_private_dir(&dir)?;
115 let path = dir.join("registry.json");
116 let file = match std::fs::read(&path) {
117 Ok(b) => serde_json::from_slice(&b).map_err(|e| AccountError::Io(e.to_string()))?,
118 Err(e) if e.kind() == std::io::ErrorKind::NotFound => RegistryFile::default(),
119 Err(e) => return Err(AccountError::Io(e.to_string())),
120 };
121 Ok(Self { base: base.to_path_buf(), file })
122 }
123
124 pub fn provision(&mut self, identity: &SessionIdentity, base_nick: Option<&str>) -> Result<AccountEntry, AccountError> {
127 let r = identity.account_ref();
128 if let Some(e) = self.file.accounts.get(&r) {
129 return Ok(e.clone());
130 }
131 let stem = match base_nick {
132 Some(n) => slug(n).ok_or(AccountError::BadNick)?,
133 None => identity.provider.id().to_string(),
134 };
135 let taken = |n: &str| self.file.accounts.values().any(|e| e.nick == n);
136 let nick = if base_nick.is_some() && !taken(&stem) {
137 stem
138 } else {
139 (1..).map(|i| format!("{stem}-{i}")).find(|n| !taken(n)).expect("free nick")
140 };
141 let entry = AccountEntry { nick, provider: identity.provider.id().to_string(), state: ServerState::Local, cap_sha256: None, os_uid: None };
142 self.file.accounts.insert(r, entry.clone());
143 self.save()?;
144 create_private_dir(&self.store_root(identity))?;
145 Ok(entry)
146 }
147
148 pub fn get(&self, identity: &SessionIdentity) -> Option<&AccountEntry> {
149 self.file.accounts.get(&identity.account_ref())
150 }
151
152 pub fn set_state(&mut self, identity: &SessionIdentity, state: ServerState) -> Result<(), AccountError> {
153 let e = self.file.accounts.get_mut(&identity.account_ref()).ok_or(AccountError::Unknown)?;
154 e.state = state;
155 self.save()
156 }
157
158 pub fn store_root(&self, identity: &SessionIdentity) -> PathBuf {
160 self.base.join("accounts").join(identity.account_ref())
161 }
162
163 pub fn issue_capability(&mut self, identity: &SessionIdentity) -> Result<String, AccountError> {
165 let e = self.file.accounts.get_mut(&identity.account_ref()).ok_or(AccountError::Unknown)?;
166 let raw = m4a_agent::vault::random_bytes(32);
167 let cap = hex(&raw);
168 e.cap_sha256 = Some(hex(&Sha256::digest(cap.as_bytes())));
169 self.save()?;
170 Ok(cap)
171 }
172
173 pub fn authorize(&self, identity: &SessionIdentity, cap: &str) -> Result<&AccountEntry, AccountError> {
175 let e = self.get(identity).ok_or(AccountError::Unknown)?;
176 let want = e.cap_sha256.as_deref().ok_or(AccountError::Refused)?;
177 let got = hex(&Sha256::digest(cap.trim().as_bytes()));
178 if ct_eq(want, &got) && e.state != ServerState::Retired {
179 Ok(e)
180 } else {
181 Err(AccountError::Refused)
182 }
183 }
184
185 pub fn accounts(&self) -> impl Iterator<Item = (&str, &AccountEntry)> {
187 self.file.accounts.iter().map(|(r, e)| (r.as_str(), e))
188 }
189
190 pub fn store_root_of(&self, account_ref: &str) -> PathBuf {
192 self.base.join("accounts").join(account_ref)
193 }
194
195 pub fn set_state_ref(&mut self, account_ref: &str, state: ServerState) -> Result<(), AccountError> {
197 let e = self.file.accounts.get_mut(account_ref).ok_or(AccountError::Unknown)?;
198 e.state = state;
199 self.save()
200 }
201
202 pub fn find_by_nick(&self, nick: &str) -> Option<(&str, &AccountEntry)> {
204 self.file.accounts.iter().find(|(_, e)| e.nick.eq_ignore_ascii_case(nick.trim())).map(|(r, e)| (r.as_str(), e))
205 }
206
207 pub fn set_os_uid(&mut self, identity: &SessionIdentity, uid: Option<u32>) -> Result<(), AccountError> {
209 let e = self.file.accounts.get_mut(&identity.account_ref()).ok_or(AccountError::Unknown)?;
210 e.os_uid = uid;
211 self.save()
212 }
213
214 pub fn gate(&self, nick: &str, cap: Option<&str>, peer_uid: Option<u32>, require: bool) -> Result<(), AccountError> {
219 let Some((_, e)) = self.find_by_nick(nick) else {
220 return if require { Err(AccountError::Unknown) } else { Ok(()) };
221 };
222 if e.state == ServerState::Retired {
223 return Err(AccountError::Refused);
224 }
225 let want = e.cap_sha256.as_deref().ok_or(AccountError::Refused)?;
226 let cap = cap.map(str::trim).filter(|c| !c.is_empty()).ok_or(AccountError::Missing)?;
227 if !ct_eq(want, &hex(&Sha256::digest(cap.as_bytes()))) {
228 return Err(AccountError::Refused);
229 }
230 match (e.os_uid, peer_uid) {
231 (Some(owner), Some(caller)) if owner != caller => Err(AccountError::WrongCaller),
232 _ => Ok(()),
233 }
234 }
235
236 fn save(&self) -> Result<(), AccountError> {
237 let dir = self.base.join(".m4a-accounts");
238 let tmp = dir.join("registry.json.tmp");
239 let bytes = serde_json::to_vec_pretty(&self.file).map_err(|e| AccountError::Io(e.to_string()))?;
240 write_private(&tmp, &bytes)?;
241 std::fs::rename(&tmp, dir.join("registry.json")).map_err(|e| AccountError::Io(e.to_string()))
242 }
243}
244
245pub struct Provisioned {
247 pub account_ref: String,
248 pub entry: AccountEntry,
249 pub store_root: PathBuf,
251 pub cap: Option<zeroize::Zeroizing<String>>,
253}
254
255impl Provisioned {
256 pub fn session_env(&self) -> Vec<(&'static str, String)> {
259 let mut v = vec![
260 ("M4A_SESSION_ID", self.account_ref.clone()),
261 ("M4A_NICK", self.entry.nick.clone()),
262 ("M4A_BOT_NAME", self.entry.nick.clone()),
263 ("M4A_STORE_ROOT", self.store_root.to_string_lossy().into_owned()),
264 ];
265 if let Some(cap) = &self.cap {
266 v.push((SESSION_CAP_ENV, cap.to_string()));
267 }
268 v
269 }
270}
271
272pub fn provision_session(base: &Path, provider: ProviderKind, provider_session_id: &str, nick: Option<&str>, rotate: bool) -> Result<Provisioned, AccountError> {
276 let identity = SessionIdentity::new(provider, provider_session_id)?;
277 let mut reg = AccountRegistry::open(base)?;
278 let mut entry = reg.provision(&identity, nick)?;
279 let cap = if rotate || entry.cap_sha256.is_none() {
280 let c = reg.issue_capability(&identity)?;
281 entry = reg.get(&identity).cloned().ok_or(AccountError::Unknown)?;
282 Some(zeroize::Zeroizing::new(c))
283 } else {
284 None
285 };
286 Ok(Provisioned { account_ref: identity.account_ref(), store_root: reg.store_root(&identity), entry, cap })
287}
288
289pub fn append_env_file(path: &Path, pairs: &[(&str, String)]) -> Result<(), AccountError> {
292 use std::io::Write;
293 let mut o = std::fs::OpenOptions::new();
294 o.create(true).append(true);
295 #[cfg(unix)]
296 {
297 use std::os::unix::fs::OpenOptionsExt;
298 o.mode(0o600);
299 }
300 let mut f = o.open(path).map_err(|e| AccountError::Io(e.to_string()))?;
301 for (k, v) in pairs {
302 if v.contains('\n') {
303 return Err(AccountError::Io("newline in env value".into()));
304 }
305 writeln!(f, "export {k}={v}").map_err(|e| AccountError::Io(e.to_string()))?;
306 }
307 Ok(())
308}
309
310pub const SESSION_FILE_KEYS: [&str; 3] = ["M4A_NICK", "M4A_SESSION_ID", SESSION_CAP_ENV];
312
313pub fn parse_session_file(text: &str) -> Vec<(&'static str, String)> {
317 let mut out = Vec::new();
318 for line in text.lines() {
319 let line = line.trim();
320 let line = line.strip_prefix("export ").unwrap_or(line);
321 let Some((k, v)) = line.split_once('=') else { continue };
322 if let Some(key) = SESSION_FILE_KEYS.iter().find(|key| **key == k.trim()) {
323 let v = v.trim();
324 if !v.is_empty() && !v.chars().any(char::is_control) {
325 out.push((*key, v.to_owned()));
326 }
327 }
328 }
329 out
330}
331
332pub fn apply_session_file(path: &Path) -> Result<(), AccountError> {
334 let text = std::fs::read_to_string(path).map_err(|e| AccountError::Io(e.kind().to_string()))?;
335 let pairs = parse_session_file(&text);
336 if pairs.is_empty() {
337 return Err(AccountError::Io("no session identity in file".into()));
338 }
339 for (k, v) in pairs {
340 if std::env::var_os(k).filter(|x| !x.is_empty()).is_none() {
341 std::env::set_var(k, v);
342 }
343 }
344 Ok(())
345}
346
347pub fn mark_enrolled(store_root: &Path) -> Result<Option<String>, AccountError> {
351 let Some(r) = store_root.file_name().and_then(|n| n.to_str()).map(str::to_string) else { return Ok(None) };
352 let Some(base) = store_root.parent().filter(|p| p.file_name().is_some_and(|n| n == "accounts")).and_then(Path::parent) else { return Ok(None) };
353 if !base.join(".m4a-accounts").join("registry.json").exists() {
354 return Ok(None);
355 }
356 let mut reg = AccountRegistry::open(base)?;
357 if reg.file.accounts.contains_key(&r) {
358 reg.set_state_ref(&r, ServerState::Enrolled)?;
359 Ok(Some(r))
360 } else {
361 Ok(None)
362 }
363}
364
365fn ct_eq(a: &str, b: &str) -> bool {
366 a.len() == b.len() && a.bytes().zip(b.bytes()).fold(0u8, |acc, (x, y)| acc | (x ^ y)) == 0
367}
368
369fn hex(b: &[u8]) -> String {
370 b.iter().map(|x| format!("{x:02x}")).collect()
371}
372
373fn slug(s: &str) -> Option<String> {
374 crate::nick::routine_folder_id(s).filter(|n| !n.is_empty() && n.len() <= 64)
375}
376
377fn create_private_dir(p: &Path) -> Result<(), AccountError> {
378 std::fs::create_dir_all(p).map_err(|e| AccountError::Io(e.to_string()))?;
379 #[cfg(unix)]
380 {
381 use std::os::unix::fs::PermissionsExt;
382 std::fs::set_permissions(p, std::fs::Permissions::from_mode(0o700)).map_err(|e| AccountError::Io(e.to_string()))?;
383 }
384 Ok(())
385}
386
387fn write_private(p: &Path, b: &[u8]) -> Result<(), AccountError> {
388 #[cfg(unix)]
389 {
390 use std::io::Write;
391 use std::os::unix::fs::OpenOptionsExt;
392 let mut f = std::fs::OpenOptions::new().write(true).create(true).truncate(true).mode(0o600).open(p).map_err(|e| AccountError::Io(e.to_string()))?;
393 f.write_all(b).map_err(|e| AccountError::Io(e.to_string()))
394 }
395 #[cfg(not(unix))]
396 {
397 std::fs::write(p, b).map_err(|e| AccountError::Io(e.to_string()))
398 }
399}
400
401#[cfg(test)]
402mod tests {
403 use super::*;
404
405 fn base(tag: &str) -> PathBuf {
406 let p = std::env::temp_dir().join(format!("m4a-acct-{tag}-{}", std::process::id()));
407 let _ = std::fs::remove_dir_all(&p);
408 p
409 }
410 fn id(p: ProviderKind, s: &str) -> SessionIdentity { SessionIdentity::new(p, s).unwrap() }
411
412 #[test]
413 fn one_account_per_session_and_stable_across_reopen() {
414 let b = base("stable");
415 let mut r = AccountRegistry::open(&b).unwrap();
416 let a = r.provision(&id(ProviderKind::Codex, "thread-a"), None).unwrap();
417 let again = r.provision(&id(ProviderKind::Codex, "thread-a"), None).unwrap();
418 assert_eq!(a, again);
419 let other = r.provision(&id(ProviderKind::Codex, "thread-b"), None).unwrap();
420 assert_ne!(a.nick, other.nick, "two sessions of one provider get two accounts");
421 assert_eq!(AccountRegistry::open(&b).unwrap().get(&id(ProviderKind::Codex, "thread-a")).unwrap().nick, a.nick);
422 let _ = std::fs::remove_dir_all(&b);
423 }
424
425 #[test]
426 fn nick_is_never_shared_between_identities() {
427 let b = base("nick");
428 let mut r = AccountRegistry::open(&b).unwrap();
429 let x = r.provision(&id(ProviderKind::Codex, "t1"), Some("Home Codex")).unwrap();
430 let y = r.provision(&id(ProviderKind::Codex, "t2"), Some("home codex")).unwrap();
431 assert_eq!(x.nick, "home-codex");
432 assert_eq!(y.nick, "home-codex-1");
433 let k = r.provision(&id(ProviderKind::KimiCode, "t1"), None).unwrap();
434 assert_ne!(k.nick, x.nick, "same native id under another provider is another account");
435 let _ = std::fs::remove_dir_all(&b);
436 }
437
438 #[test]
439 fn foreign_capability_is_refused() {
440 let b = base("cap");
441 let mut r = AccountRegistry::open(&b).unwrap();
442 let (a, c) = (id(ProviderKind::ClaudeCode, "s-a"), id(ProviderKind::ClaudeCode, "s-c"));
443 r.provision(&a, None).unwrap();
444 r.provision(&c, None).unwrap();
445 let cap_a = r.issue_capability(&a).unwrap();
446 assert!(r.authorize(&a, &cap_a).is_ok());
447 assert_eq!(r.authorize(&c, &cap_a).unwrap_err(), AccountError::Refused, "c has no capability yet");
448 let cap_c = r.issue_capability(&c).unwrap();
449 assert_eq!(r.authorize(&a, &cap_c).unwrap_err(), AccountError::Refused);
450 let rotated = r.issue_capability(&a).unwrap();
451 assert!(r.authorize(&a, &cap_a).is_err() && r.authorize(&a, &rotated).is_ok());
452 r.set_state(&a, ServerState::Retired).unwrap();
453 assert!(r.authorize(&a, &rotated).is_err());
454 let raw = std::fs::read_to_string(b.join(".m4a-accounts/registry.json")).unwrap();
455 assert!(!raw.contains(&rotated) && !raw.contains("s-a"), "no capability or native id in clear");
456 let _ = std::fs::remove_dir_all(&b);
457 }
458
459 #[test]
460 fn each_account_has_its_own_private_store_root() {
461 let b = base("root");
462 let mut r = AccountRegistry::open(&b).unwrap();
463 let (a, c) = (id(ProviderKind::Grok, "g1"), id(ProviderKind::Grok, "g2"));
464 r.provision(&a, None).unwrap();
465 r.provision(&c, None).unwrap();
466 assert_ne!(r.store_root(&a), r.store_root(&c));
467 #[cfg(unix)]
468 {
469 use std::os::unix::fs::PermissionsExt;
470 assert_eq!(std::fs::metadata(r.store_root(&a)).unwrap().permissions().mode() & 0o777, 0o700);
471 assert_eq!(std::fs::metadata(b.join(".m4a-accounts/registry.json")).unwrap().permissions().mode() & 0o777, 0o600);
472 }
473 let _ = std::fs::remove_dir_all(&b);
474 }
475
476 #[test]
477 fn identity_rejects_empty_and_control_chars() {
478 assert!(SessionIdentity::new(ProviderKind::Codex, " ").is_err());
479 assert!(SessionIdentity::new(ProviderKind::Codex, "a\nb").is_err());
480 }
481
482 #[test]
483 fn two_sessions_same_provider_get_distinct_accounts_roots_and_logins() {
484 let dir = base("s2-two");
485 let a = provision_session(&dir, ProviderKind::Codex, "thread-a", Some("codex"), false).unwrap();
486 let b = provision_session(&dir, ProviderKind::Codex, "thread-b", Some("codex"), false).unwrap();
487 assert_ne!(a.entry.nick, b.entry.nick);
488 assert_ne!(a.store_root, b.store_root);
489 assert_ne!(a.account_ref, b.account_ref);
490 let env_a: Vec<_> = a.session_env().into_iter().map(|(k, _)| k).collect();
491 assert!(env_a.contains(&SESSION_CAP_ENV));
492 let a2 = provision_session(&dir, ProviderKind::Codex, "thread-a", Some("codex"), false).unwrap();
494 assert!(a2.cap.is_none());
495 assert_eq!(a2.entry.nick, a.entry.nick);
496 }
497
498 #[test]
499 fn gate_refuses_foreign_missing_and_wrong_caller() {
500 let dir = base("s2-gate");
501 let a = provision_session(&dir, ProviderKind::Codex, "thread-a", None, false).unwrap();
502 let b = provision_session(&dir, ProviderKind::Codex, "thread-b", None, false).unwrap();
503 let (ca, cb) = (a.cap.unwrap(), b.cap.unwrap());
504 let mut reg = AccountRegistry::open(&dir).unwrap();
505 assert_eq!(reg.gate(&a.entry.nick, Some(&ca), None, true), Ok(()));
506 assert_eq!(reg.gate(&a.entry.nick, Some(&cb), None, true), Err(AccountError::Refused));
507 assert_eq!(reg.gate(&a.entry.nick, None, None, false), Err(AccountError::Missing));
508 assert_eq!(reg.gate("operator-bot", None, None, false), Ok(()));
509 assert_eq!(reg.gate("operator-bot", None, None, true), Err(AccountError::Unknown));
510 let ida = SessionIdentity::new(ProviderKind::Codex, "thread-a").unwrap();
511 reg.set_os_uid(&ida, Some(1000)).unwrap();
512 assert_eq!(reg.gate(&a.entry.nick, Some(&ca), Some(1001), true), Err(AccountError::WrongCaller));
513 assert_eq!(reg.gate(&a.entry.nick, Some(&ca), Some(1000), true), Ok(()));
514 reg.set_state(&ida, ServerState::Retired).unwrap();
515 assert_eq!(reg.gate(&a.entry.nick, Some(&ca), Some(1000), true), Err(AccountError::Refused));
516 }
517
518 #[test]
519 fn enroll_marks_only_that_account_and_daemon_lists_it_from_its_own_root() {
520 let dir = base("s3-enroll");
521 let a = provision_session(&dir, ProviderKind::KimiCode, "a", Some("kimi"), false).unwrap();
522 let b = provision_session(&dir, ProviderKind::KimiCode, "b", Some("kimi"), false).unwrap();
523 assert!(crate::account_sessions(&dir).is_empty(), "local accounts are not opened on the server");
524 assert_eq!(mark_enrolled(&a.store_root).unwrap(), Some(a.account_ref.clone()));
525 assert_eq!(mark_enrolled(&dir.join("elsewhere")).unwrap(), None);
526 let listed = crate::account_sessions(&dir);
527 assert_eq!(listed.len(), 1);
528 assert_eq!(listed[0].bot_name, a.entry.nick);
529 assert_eq!(listed[0].session_id, a.account_ref);
530 assert_eq!(listed[0].store_root.as_deref(), Some(a.store_root.as_path()));
531 assert_ne!(listed[0].store_root.as_deref(), Some(b.store_root.as_path()));
532 }
533
534 #[test]
535 fn session_file_keeps_only_identity_keys() {
536 let got = parse_session_file("export M4A_NICK=a-1\nM4A_SESSION_ID=s\nPATH=/x\nM4A_SESSION_CAP=\n# c\n");
537 assert_eq!(got, vec![("M4A_NICK", "a-1".to_owned()), ("M4A_SESSION_ID", "s".to_owned())]);
538 }
539
540 #[test]
541 fn env_file_is_private_and_rejects_newlines() {
542 let dir = base("s2-env");
543 std::fs::create_dir_all(&dir).unwrap();
544 let p = dir.join("env");
545 append_env_file(&p, &[("M4A_NICK", "x".into())]).unwrap();
546 assert!(append_env_file(&p, &[("M4A_NICK", "x\ny".into())]).is_err());
547 assert_eq!(std::fs::read_to_string(&p).unwrap(), "export M4A_NICK=x\n");
548 #[cfg(unix)]
549 {
550 use std::os::unix::fs::PermissionsExt;
551 assert_eq!(std::fs::metadata(&p).unwrap().permissions().mode() & 0o777, 0o600);
552 }
553 }
554}