zeroclaw_runtime/rpc/
tui_identity.rs

1//! TUI session identity — UID generation, HMAC signing, and live
2//! connection registry.
3//!
4//! **Source of truth** for connected TUI state. The `TuiRegistry` lives
5//! on [`super::context::RpcContext`] and is the single canonical location
6//! for "which TUIs are connected right now." Nothing else stores this.
7
8use std::collections::HashMap;
9use std::path::Path;
10use std::sync::Mutex;
11
12use chrono::{DateTime, Utc};
13use hmac::{Hmac, Mac};
14use sha2::Sha256;
15
16type HmacSha256 = Hmac<Sha256>;
17
18// ── TUI entry ────────────────────────────────────────────────────
19
20/// A connected TUI client.
21#[derive(Debug, Clone)]
22pub struct TuiEntry {
23    pub tui_id: String,
24    pub connected_at: DateTime<Utc>,
25    pub peer_label: String,
26    /// Transport protocol: `"unix"` or `"wss"`.
27    pub transport: String,
28    /// Full shell environment captured from the TUI process at connect time.
29    /// Used to pass the user's real env (PATH, SSH_AUTH_SOCK, etc.) through
30    /// to subprocesses spawned by the daemon on their behalf.
31    pub env: HashMap<String, String>,
32}
33
34// ── Registry ─────────────────────────────────────────────────────
35
36/// Daemon-wide registry of connected TUI clients.
37///
38/// **Source of truth** for live TUI connection state. Not persisted —
39/// rebuilt on each daemon start from incoming `initialize` handshakes.
40pub struct TuiRegistry {
41    /// HMAC signing key loaded from `.secret_key`. `None` = signing
42    /// disabled — UIDs are issued unsigned and reconnects trust claimed
43    /// identities without verification.
44    signing_key: Option<Vec<u8>>,
45    /// Connected TUIs keyed by `tui_id`.
46    connected: Mutex<HashMap<String, TuiEntry>>,
47}
48
49impl TuiRegistry {
50    /// Create a registry, attempting to load the signing key from
51    /// `<config_dir>/.secret_key`. If the file is missing or
52    /// unreadable, signing is silently disabled.
53    pub fn new(config_dir: &Path) -> Self {
54        let key_path = config_dir.join(".secret_key");
55        let signing_key = std::fs::read_to_string(&key_path)
56            .ok()
57            .and_then(|hex_str| hex::decode(hex_str.trim()).ok())
58            .filter(|key| !key.is_empty());
59
60        Self {
61            signing_key,
62            connected: Mutex::new(HashMap::new()),
63        }
64    }
65
66    /// Test constructor with no signing key.
67    #[cfg(test)]
68    pub fn new_unsigned() -> Self {
69        Self {
70            signing_key: None,
71            connected: Mutex::new(HashMap::new()),
72        }
73    }
74
75    /// Whether HMAC signing is enabled (`.secret_key` was loaded).
76    pub fn signing_is_enabled(&self) -> bool {
77        self.signing_key.is_some()
78    }
79
80    // ── UID generation ───────────────────────────────────────────
81
82    /// Generate a short TUI ID: `tui_` + 8 hex chars (4 random bytes).
83    pub fn generate_tui_id() -> String {
84        let bytes: [u8; 4] = rand::random();
85        format!("tui_{}", hex::encode(bytes))
86    }
87
88    /// Generate a TUI ID that is not currently in the registry.
89    pub fn generate_unique_tui_id(&self) -> String {
90        let connected = self.connected.lock().unwrap_or_else(|e| e.into_inner());
91        loop {
92            let id = Self::generate_tui_id();
93            if !connected.contains_key(&id) {
94                return id;
95            }
96        }
97    }
98
99    // ── HMAC signing ─────────────────────────────────────────────
100
101    /// Sign a TUI ID with HMAC-SHA256. Returns `None` if signing is
102    /// disabled.
103    pub fn sign(&self, tui_id: &str) -> Option<String> {
104        let key = self.signing_key.as_ref()?;
105        let mut mac = HmacSha256::new_from_slice(key).expect("HMAC accepts any key length");
106        mac.update(tui_id.as_bytes());
107        Some(hex::encode(mac.finalize().into_bytes()))
108    }
109
110    /// Verify a TUI ID + signature. Returns `true` if:
111    /// - Signing is disabled (trust all), OR
112    /// - The signature is valid.
113    pub fn verify(&self, tui_id: &str, sig: &str) -> bool {
114        let Some(ref key) = self.signing_key else {
115            return true;
116        };
117        let Ok(sig_bytes) = hex::decode(sig) else {
118            return false;
119        };
120        let mut mac = HmacSha256::new_from_slice(key).expect("HMAC accepts any key length");
121        mac.update(tui_id.as_bytes());
122        mac.verify_slice(&sig_bytes).is_ok()
123    }
124
125    // ── Registry operations ──────────────────────────────────────
126
127    /// Register a connected TUI.
128    pub fn register(&self, entry: TuiEntry) {
129        self.connected
130            .lock()
131            .unwrap_or_else(|e| e.into_inner())
132            .insert(entry.tui_id.clone(), entry);
133    }
134
135    /// Unregister a disconnected TUI.
136    pub fn unregister(&self, tui_id: &str) {
137        self.connected
138            .lock()
139            .unwrap_or_else(|e| e.into_inner())
140            .remove(tui_id);
141    }
142
143    /// Snapshot of all connected TUIs.
144    pub fn list(&self) -> Vec<TuiEntry> {
145        self.connected
146            .lock()
147            .unwrap_or_else(|e| e.into_inner())
148            .values()
149            .cloned()
150            .collect()
151    }
152
153    /// Return a clone of the environment captured from the TUI identified by
154    /// `tui_id`, or `None` if the TUI is not currently connected.
155    pub fn get_env(&self, tui_id: &str) -> Option<HashMap<String, String>> {
156        self.connected
157            .lock()
158            .unwrap_or_else(|e| e.into_inner())
159            .get(tui_id)
160            .map(|e| e.env.clone())
161    }
162}
163
164// ── Tests ────────────────────────────────────────────────────────
165
166#[cfg(test)]
167mod tests {
168    use super::*;
169
170    #[test]
171    fn generate_tui_id_format() {
172        let id = TuiRegistry::generate_tui_id();
173        assert!(id.starts_with("tui_"), "expected tui_ prefix, got {id}");
174        assert_eq!(id.len(), 12, "tui_ + 8 hex chars = 12, got {}", id.len());
175        // Hex chars only after prefix
176        assert!(
177            id[4..].chars().all(|c| c.is_ascii_hexdigit()),
178            "non-hex chars in {id}"
179        );
180    }
181
182    #[test]
183    fn sign_verify_roundtrip() {
184        let registry = TuiRegistry {
185            signing_key: Some(vec![0xAB; 32]),
186            connected: Mutex::new(HashMap::new()),
187        };
188        let id = "tui_deadbeef";
189        let sig = registry.sign(id).expect("signing should succeed");
190        assert!(registry.verify(id, &sig), "roundtrip verify failed");
191    }
192
193    #[test]
194    fn verify_rejects_tampered_sig() {
195        let registry = TuiRegistry {
196            signing_key: Some(vec![0xAB; 32]),
197            connected: Mutex::new(HashMap::new()),
198        };
199        let id = "tui_deadbeef";
200        let sig = registry.sign(id).unwrap();
201        // Flip a character
202        let mut tampered = sig.clone();
203        let replacement = if tampered.ends_with('0') { 'f' } else { '0' };
204        tampered.pop();
205        tampered.push(replacement);
206        assert!(!registry.verify(id, &tampered), "tampered sig should fail");
207    }
208
209    #[test]
210    fn verify_rejects_wrong_id() {
211        let registry = TuiRegistry {
212            signing_key: Some(vec![0xAB; 32]),
213            connected: Mutex::new(HashMap::new()),
214        };
215        let sig = registry.sign("tui_aaaaaaaa").unwrap();
216        assert!(
217            !registry.verify("tui_bbbbbbbb", &sig),
218            "wrong ID should fail"
219        );
220    }
221
222    #[test]
223    fn verify_without_key_trusts_all() {
224        let registry = TuiRegistry::new_unsigned();
225        assert!(registry.verify("tui_anything", "bogus_sig"));
226    }
227
228    #[test]
229    fn signing_disabled_returns_none() {
230        let registry = TuiRegistry::new_unsigned();
231        assert!(registry.sign("tui_test").is_none());
232        assert!(!registry.signing_is_enabled());
233    }
234
235    #[test]
236    fn register_unregister_lifecycle() {
237        let registry = TuiRegistry::new_unsigned();
238        assert!(registry.list().is_empty());
239
240        registry.register(TuiEntry {
241            tui_id: "tui_aabb0011".to_string(),
242            connected_at: Utc::now(),
243            peer_label: "test".to_string(),
244            transport: "unix".to_string(),
245            env: HashMap::new(),
246        });
247        assert_eq!(registry.list().len(), 1);
248        assert_eq!(registry.list()[0].tui_id, "tui_aabb0011");
249
250        registry.unregister("tui_aabb0011");
251        assert!(registry.list().is_empty());
252    }
253
254    #[test]
255    fn unregister_unknown_is_noop() {
256        let registry = TuiRegistry::new_unsigned();
257        registry.unregister("tui_nonexistent"); // must not panic
258    }
259
260    #[test]
261    fn generate_unique_avoids_existing() {
262        let registry = TuiRegistry::new_unsigned();
263        // Pre-populate with a known ID
264        registry.register(TuiEntry {
265            tui_id: "tui_00000000".to_string(),
266            connected_at: Utc::now(),
267            peer_label: "test".to_string(),
268            transport: "unix".to_string(),
269            env: HashMap::new(),
270        });
271        // generate_unique should return something different
272        let id = registry.generate_unique_tui_id();
273        assert_ne!(id, "tui_00000000");
274    }
275
276    // ── TUI env passthrough tests ─────────────────────────────────
277
278    #[test]
279    fn tui_entry_stores_env() {
280        let registry = TuiRegistry::new_unsigned();
281        let mut env = HashMap::new();
282        env.insert("MY_VAR".to_string(), "my_value".to_string());
283        env.insert("ANTHROPIC_API_KEY".to_string(), "sk-secret".to_string());
284
285        registry.register(TuiEntry {
286            tui_id: "tui_aabbccdd".to_string(),
287            connected_at: Utc::now(),
288            peer_label: "test".to_string(),
289            transport: "unix".to_string(),
290            env,
291        });
292
293        let entries = registry.list();
294        assert_eq!(entries.len(), 1);
295        assert_eq!(
296            entries[0].env.get("MY_VAR").map(|s| s.as_str()),
297            Some("my_value")
298        );
299        assert_eq!(
300            entries[0].env.get("ANTHROPIC_API_KEY").map(|s| s.as_str()),
301            Some("sk-secret"),
302            "full env should be stored without filtering"
303        );
304    }
305
306    #[test]
307    fn tui_entry_env_defaults_to_empty() {
308        // Entries with no env (e.g. old clients) should work fine
309        let registry = TuiRegistry::new_unsigned();
310        registry.register(TuiEntry {
311            tui_id: "tui_11223344".to_string(),
312            connected_at: Utc::now(),
313            peer_label: "test".to_string(),
314            transport: "unix".to_string(),
315            env: HashMap::new(),
316        });
317
318        let entries = registry.list();
319        assert!(entries[0].env.is_empty());
320    }
321
322    #[test]
323    fn tui_entry_env_dropped_on_unregister() {
324        let registry = TuiRegistry::new_unsigned();
325        let mut env = HashMap::new();
326        env.insert("SOME_VAR".to_string(), "some_value".to_string());
327
328        registry.register(TuiEntry {
329            tui_id: "tui_deadbeef".to_string(),
330            connected_at: Utc::now(),
331            peer_label: "test".to_string(),
332            transport: "unix".to_string(),
333            env,
334        });
335        assert_eq!(registry.list().len(), 1);
336
337        registry.unregister("tui_deadbeef");
338        assert!(
339            registry.list().is_empty(),
340            "env should be dropped with entry"
341        );
342    }
343
344    #[test]
345    fn tui_entry_env_survives_clone() {
346        // TuiEntry derives Clone — env must be included
347        let mut env = HashMap::new();
348        env.insert("CLONED_VAR".to_string(), "cloned_value".to_string());
349
350        let entry = TuiEntry {
351            tui_id: "tui_cafebabe".to_string(),
352            connected_at: Utc::now(),
353            peer_label: "test".to_string(),
354            transport: "unix".to_string(),
355            env,
356        };
357        let cloned = entry.clone();
358        assert_eq!(
359            cloned.env.get("CLONED_VAR").map(|s| s.as_str()),
360            Some("cloned_value")
361        );
362    }
363
364    #[test]
365    fn get_env_returns_env_for_connected_tui() {
366        let registry = TuiRegistry::new_unsigned();
367        let mut env = HashMap::new();
368        env.insert("PATH".to_string(), "/usr/bin:/usr/local/bin".to_string());
369        env.insert("SSH_AUTH_SOCK".to_string(), "/tmp/ssh.sock".to_string());
370
371        registry.register(TuiEntry {
372            tui_id: "tui_getenv01".to_string(),
373            connected_at: Utc::now(),
374            peer_label: "test".to_string(),
375            transport: "unix".to_string(),
376            env,
377        });
378
379        let got = registry.get_env("tui_getenv01").expect("should find env");
380        assert_eq!(
381            got.get("PATH").map(|s| s.as_str()),
382            Some("/usr/bin:/usr/local/bin")
383        );
384        assert_eq!(
385            got.get("SSH_AUTH_SOCK").map(|s| s.as_str()),
386            Some("/tmp/ssh.sock")
387        );
388    }
389
390    #[test]
391    fn get_env_returns_none_for_unknown_tui() {
392        let registry = TuiRegistry::new_unsigned();
393        assert!(registry.get_env("tui_nothere").is_none());
394    }
395
396    #[test]
397    fn get_env_returns_none_after_unregister() {
398        let registry = TuiRegistry::new_unsigned();
399        let mut env = HashMap::new();
400        env.insert("SOME_VAR".to_string(), "val".to_string());
401        registry.register(TuiEntry {
402            tui_id: "tui_gone0001".to_string(),
403            connected_at: Utc::now(),
404            peer_label: "test".to_string(),
405            transport: "unix".to_string(),
406            env,
407        });
408        assert!(registry.get_env("tui_gone0001").is_some());
409        registry.unregister("tui_gone0001");
410        assert!(registry.get_env("tui_gone0001").is_none());
411    }
412}
zeroclaw_runtime/rpc/tui_identity.rs

zeroclaw_runtime/rpc/
tui_identity.rs
