oxios_kernel/pty/
manager.rs

1//! PtyManager — session registry, GC tick, attach/detach (RFC-038 §5.3, §6.3, §8.3).
2use parking_lot::RwLock;
3use std::collections::{HashMap, HashSet};
4use std::sync::Arc;
5use std::time::Instant;
6
7use crate::config::PtyConfig;
8
9use super::error::PtyError;
10use super::session::{PtySession, PtySessionId, PtySessionInfo, PtySessionState, PtySize};
11
12/// The shared PTY session registry + lifecycle manager.
13pub struct PtyManager {
14    sessions: RwLock<HashMap<PtySessionId, Arc<PtySession>>>,
15    by_principal: parking_lot::Mutex<HashMap<String, HashSet<PtySessionId>>>,
16    config: Arc<RwLock<PtyConfig>>,
17}
18
19impl PtyManager {
20    pub fn new(config: Arc<RwLock<PtyConfig>>) -> Self {
21        Self {
22            sessions: RwLock::new(HashMap::new()),
23            by_principal: parking_lot::Mutex::new(HashMap::new()),
24            config,
25        }
26    }
27
28    /// Read-only snapshot of the current config.
29    pub fn config_snapshot(&self) -> PtyConfig {
30        self.config.read().clone()
31    }
32
33    /// Check if PTY subsystem is enabled.
34    pub fn is_enabled(&self) -> bool {
35        self.config.read().enabled
36    }
37
38    /// Generate a fresh session id (ULID-style: timestamp + random hex).
39    pub fn new_id() -> PtySessionId {
40        let now = std::time::SystemTime::now()
41            .duration_since(std::time::UNIX_EPOCH)
42            .map(|d| d.as_millis())
43            .unwrap_or(0);
44        let mut bytes = [0u8; 26];
45        bytes[0..8].copy_from_slice(&now.to_be_bytes()[..8]);
46        bytes[8] = (now >> 40) as u8;
47        let rand_bytes: [u8; 16] = rand_bytes();
48        bytes[10..26].copy_from_slice(&rand_bytes);
49        hex_encode(&bytes)
50    }
51
52    /// Open a new PTY session. Performs shell allowlist check and per-principal cap.
53    pub fn open(
54        self: &Arc<Self>,
55        principal: &str,
56        shell: Option<String>,
57        size: PtySize,
58    ) -> Result<Arc<PtySession>, PtyError> {
59        let cfg = self.config_snapshot();
60        if !cfg.enabled {
61            return Err(PtyError::Disabled);
62        }
63        let resolved = match shell {
64            Some(s) if !s.is_empty() => s,
65            _ => cfg.default_shell.clone(),
66        };
67        if !cfg.is_shell_allowed(&resolved) {
68            return Err(PtyError::ShellNotAllowed { shell: resolved });
69        }
70        {
71            let map = self.by_principal.lock();
72            let count = map.get(principal).map(|s| s.len()).unwrap_or(0);
73            if count >= cfg.max_sessions as usize {
74                return Err(PtyError::SessionCapReached {
75                    max: cfg.max_sessions,
76                });
77            }
78        }
79        let id = Self::new_id();
80        let session = PtySession::spawn(id.clone(), principal.to_string(), resolved, size, &cfg)?;
81        self.sessions
82            .write()
83            .insert(id.clone(), Arc::clone(&session));
84        self.by_principal
85            .lock()
86            .entry(principal.to_string())
87            .or_default()
88            .insert(id.clone());
89        tracing::info!(
90            session = %id,
91            shell = %session.shell,
92            principal = %principal,
93            "pty.open"
94        );
95        Ok(session)
96    }
97
98    /// Re-attach an existing session by id. Validates principal match.
99    pub fn attach(
100        self: &Arc<Self>,
101        principal: &str,
102        session_id: &str,
103    ) -> Result<Arc<PtySession>, PtyError> {
104        let session = {
105            let map = self.sessions.read();
106            map.get(session_id).cloned()
107        }
108        .ok_or_else(|| PtyError::NotFound(session_id.to_string()))?;
109        if session.principal != principal {
110            return Err(PtyError::NotOwner(session_id.to_string()));
111        }
112        if let PtySessionState::Closed { .. } = &*session.state.lock() {
113            return Err(PtyError::Closed(session_id.to_string()));
114        }
115        Ok(session)
116    }
117
118    /// Write bytes to PTY stdin.
119    pub fn write(&self, session_id: &str, bytes: &[u8]) -> Result<(), PtyError> {
120        let session = self
121            .sessions
122            .read()
123            .get(session_id)
124            .cloned()
125            .ok_or_else(|| PtyError::NotFound(session_id.to_string()))?;
126        let mut guard = session.master.lock();
127        let master = guard
128            .as_mut()
129            .ok_or_else(|| PtyError::Closed(session_id.to_string()))?;
130        master
131            .take_writer()
132            .map_err(|e| PtyError::Io(e.to_string()))?
133            .write_all(bytes)
134            .map_err(|e| PtyError::Io(e.to_string()))?;
135        session.touch_input();
136        Ok(())
137    }
138
139    /// Resize the PTY.
140    pub fn resize(&self, session_id: &str, cols: u16, rows: u16) -> Result<(), PtyError> {
141        let session = self
142            .sessions
143            .read()
144            .get(session_id)
145            .cloned()
146            .ok_or_else(|| PtyError::NotFound(session_id.to_string()))?;
147        let guard = session.master.lock();
148        let master = guard
149            .as_ref()
150            .ok_or_else(|| PtyError::Closed(session_id.to_string()))?;
151        let _ = master.resize(portable_pty::PtySize {
152            rows,
153            cols,
154            pixel_width: 0,
155            pixel_height: 0,
156        });
157        Ok(())
158    }
159
160    /// Take the master reader for streaming to a WS client.
161    pub fn try_clone_reader(
162        &self,
163        session_id: &str,
164    ) -> Result<Box<dyn std::io::Read + Send>, PtyError> {
165        let session = self
166            .sessions
167            .read()
168            .get(session_id)
169            .cloned()
170            .ok_or_else(|| PtyError::NotFound(session_id.to_string()))?;
171        let guard = session.master.lock();
172        let master = guard
173            .as_ref()
174            .ok_or_else(|| PtyError::Closed(session_id.to_string()))?;
175        master
176            .try_clone_reader()
177            .map_err(|e| PtyError::Io(e.to_string()))
178    }
179
180    /// List sessions for a principal.
181    pub fn list_sessions(&self, principal: &str) -> Vec<PtySessionInfo> {
182        let ids: Vec<PtySessionId> = self
183            .by_principal
184            .lock()
185            .get(principal)
186            .cloned()
187            .unwrap_or_default()
188            .into_iter()
189            .collect();
190        let sessions = self.sessions.read();
191        ids.into_iter()
192            .filter_map(|id| sessions.get(&id).map(|s| s.info()))
193            .collect()
194    }
195
196    /// Transition a session to Attached (called by WS handler on Open).
197    pub fn mark_attached(&self, session_id: &str) -> bool {
198        let session = match self.sessions.read().get(session_id).cloned() {
199            Some(s) => s,
200            None => return false,
201        };
202        let mut st = session.state.lock();
203        match &*st {
204            PtySessionState::Closed { .. } => false,
205            _ => {
206                let (tx, _rx) = tokio::sync::mpsc::channel::<Vec<u8>>(64);
207                *st = PtySessionState::Attached { ws_tx: tx };
208                true
209            }
210        }
211    }
212
213    /// Mark a session as Detached (called by WS handler on Close).
214    pub fn mark_detached(&self, session_id: &str) -> bool {
215        let session = match self.sessions.read().get(session_id).cloned() {
216            Some(s) => s,
217            None => return false,
218        };
219        let mut st = session.state.lock();
220        match &*st {
221            PtySessionState::Closed { .. } => false,
222            _ => {
223                *st = PtySessionState::Detached {
224                    orphan_since: Instant::now(),
225                };
226                true
227            }
228        }
229    }
230
231    /// Close + remove a session.
232    pub fn close(&self, session_id: &str) -> Result<(), PtyError> {
233        let session = self
234            .sessions
235            .read()
236            .get(session_id)
237            .cloned()
238            .ok_or_else(|| PtyError::NotFound(session_id.to_string()))?;
239        let _ = session.master.lock().take();
240        let mut st = session.state.lock();
241        *st = PtySessionState::Closed {
242            exit_code: None,
243            signal: None,
244            at: Instant::now(),
245        };
246        self.sessions.write().remove(session_id);
247        if let Some(set) = self.by_principal.lock().get_mut(&session.principal) {
248            set.remove(session_id);
249        }
250        tracing::info!(
251            session = %session_id,
252            principal = %session.principal,
253            "pty.close"
254        );
255        Ok(())
256    }
257
258    /// GC tick — close idle / past-max-lifetime sessions.
259    pub fn gc_tick(&self) {
260        let cfg = self.config_snapshot();
261        let idle_ms_threshold = cfg.idle_timeout_secs.saturating_mul(1000);
262        let max_life_ms = cfg.max_lifetime_secs.saturating_mul(1000);
263        let mut to_close: Vec<PtySessionId> = Vec::new();
264        {
265            let map = self.sessions.read();
266            for (id, s) in map.iter() {
267                if s.lifetime_ms() > max_life_ms {
268                    to_close.push(id.clone());
269                    continue;
270                }
271                if s.idle_ms() > idle_ms_threshold {
272                    let st = s.state.lock();
273                    if !matches!(*st, PtySessionState::Closed { .. }) {
274                        to_close.push(id.clone());
275                    }
276                }
277            }
278        }
279        for id in to_close {
280            let _ = self.close(&id);
281        }
282    }
283
284    /// Spawn the GC tick task.
285    pub fn start_gc(self: Arc<Self>) -> tokio::task::JoinHandle<()> {
286        tokio::spawn(async move {
287            loop {
288                let interval = {
289                    let cfg = self.config_snapshot();
290                    std::cmp::min(cfg.idle_timeout_secs / 4, 60).max(5)
291                };
292                tokio::time::sleep(std::time::Duration::from_secs(interval)).await;
293                self.gc_tick();
294            }
295        })
296    }
297}
298
299fn rand_bytes() -> [u8; 16] {
300    use std::hash::{BuildHasher, Hasher, RandomState};
301    let mut state = RandomState::new().build_hasher();
302    state.write_u128(
303        std::time::SystemTime::now()
304            .duration_since(std::time::UNIX_EPOCH)
305            .map(|d| d.as_nanos())
306            .unwrap_or(0),
307    );
308    let a = state.finish();
309    state.write_u64(std::process::id() as u64);
310    let b = state.finish();
311    let mut out = [0u8; 16];
312    out[..8].copy_from_slice(&a.to_be_bytes());
313    out[8..].copy_from_slice(&b.to_be_bytes());
314    out
315}
316
317fn hex_encode(bytes: &[u8]) -> String {
318    let mut s = String::with_capacity(bytes.len() * 2);
319    for b in bytes {
320        s.push_str(&format!("{:02x}", b));
321    }
322    s
323}
324
325#[cfg(test)]
326mod tests {
327    use super::*;
328
329    #[test]
330    fn new_id_is_unique() {
331        let a = PtyManager::new_id();
332        let b = PtyManager::new_id();
333        assert_ne!(a, b);
334        assert!(a.len() >= 32);
335    }
336
337    #[test]
338    fn disabled_blocks_open() {
339        let cfg = Arc::new(RwLock::new(PtyConfig::default()));
340        let m = Arc::new(PtyManager::new(cfg));
341        let res = m.open("user", None, PtySize::default_80x24());
342        assert!(matches!(res, Err(PtyError::Disabled)));
343    }
344}
oxios_kernel/pty/manager.rs

oxios_kernel/pty/
manager.rs
