diff --git a/src/cli/commands.rs b/src/cli/commands.rs index d975eeb..5a0b0da 100644 --- a/src/cli/commands.rs +++ b/src/cli/commands.rs @@ -59,7 +59,13 @@ pub fn cmd_wrap(capture: &str) -> Result<()> { if slave_raw > libc::STDERR_FILENO { let _ = close(slave_raw); } - let shell_cstr = CString::new(shell.as_str()).unwrap(); + let shell_cstr = match CString::new(shell.as_str()) { + Ok(c) => c, + Err(_) => { + eprintln!("error: SHELL contains invalid characters"); + unsafe { libc::_exit(1) }; + } + }; let _ = execvp(&shell_cstr, std::slice::from_ref(&shell_cstr)); unsafe { libc::_exit(127); @@ -403,7 +409,12 @@ pub async fn cmd_info(config: &Config, target: String, json: bool) -> Result<()> } } -pub async fn cmd_attach(config: &Config, target: String, read_only: bool) -> Result<()> { +pub async fn cmd_attach( + config: &Config, + target: String, + read_only: bool, + force: bool, +) -> Result<()> { use crossterm::event::{Event, EventStream, KeyCode, KeyModifiers}; use crossterm::terminal; use futures_lite::StreamExt; @@ -415,6 +426,7 @@ pub async fn cmd_attach(config: &Config, target: String, read_only: bool) -> Res .request(&Request::SessionAttach { target: target.clone(), read_only, + force, }) .await?; @@ -471,6 +483,18 @@ pub async fn cmd_attach(config: &Config, target: String, read_only: bool) -> Res let _ = std::io::Write::write_all(&mut tty_out, &payload); let _ = std::io::Write::flush(&mut tty_out); } else if msg_type == MSG_SESSION_EVENT { + let msg = if let Ok(resp) = decode_response(msg_type, &payload) { + match resp { + Response::SessionEvent { event, .. } if event == "stolen" => { + "\r\n\x1b[33m[Session stolen by another client]\x1b[0m\r\n" + } + _ => "\r\n\x1b[33m[Session killed by snag]\x1b[0m\r\n", + } + } else { + "\r\n\x1b[33m[Session ended]\x1b[0m\r\n" + }; + let _ = std::io::Write::write_all(&mut tty_out, msg.as_bytes()); + let _ = std::io::Write::flush(&mut tty_out); break Ok(()); } // MSG_OK/MSG_ERROR from control messages — ignore diff --git a/src/cli/mod.rs b/src/cli/mod.rs index 1889de3..0aab09a 100644 --- a/src/cli/mod.rs +++ b/src/cli/mod.rs @@ -66,6 +66,9 @@ pub enum Command { /// Read-only mode #[arg(long)] read_only: bool, + /// Force-steal from another attached client + #[arg(long)] + force: bool, }, /// Send a command to a session Send { diff --git a/src/client.rs b/src/client.rs index 0905610..e1cb796 100644 --- a/src/client.rs +++ b/src/client.rs @@ -90,7 +90,10 @@ fn start_daemon(config: &Config) -> Result<()> { let _ = nix::unistd::dup2(devnull.as_raw_fd(), nix::libc::STDIN_FILENO); let _ = nix::unistd::dup2(devnull.as_raw_fd(), nix::libc::STDOUT_FILENO); // Keep stderr for daemon logging - let log_path = socket_path.parent().unwrap().join("snagd.log"); + let log_path = socket_path + .parent() + .unwrap_or(std::path::Path::new("/tmp")) + .join("snagd.log"); if let Ok(log) = std::fs::File::create(&log_path) { let _ = nix::unistd::dup2(log.as_raw_fd(), nix::libc::STDERR_FILENO); } diff --git a/src/daemon/adopt.rs b/src/daemon/adopt.rs index ddde966..a819220 100644 --- a/src/daemon/adopt.rs +++ b/src/daemon/adopt.rs @@ -97,8 +97,24 @@ fn get_pts_number_for_fd(pid: u32, fd: i32) -> Option { } } - // Fallback: try to figure out from /proc//fd/ -> /dev/pts/ mapping - // by checking all processes for matching tty + // Fallback: try to resolve via /proc//fd/ -> /dev/pts/ symlink. + // Some PTY masters (e.g., opened via /dev/pts/ptmx) have a sibling slave fd + // pointing to /dev/pts/. Scan adjacent fds for one linked to /dev/pts/. + let fd_dir = format!("/proc/{pid}/fd"); + if let Ok(entries) = std::fs::read_dir(&fd_dir) { + for entry in entries.flatten() { + if let Ok(target) = std::fs::read_link(entry.path()) { + let s = target.to_string_lossy(); + if let Some(rest) = s.strip_prefix("/dev/pts/") { + if rest != "ptmx" { + if let Ok(n) = rest.parse::() { + return Some(n); + } + } + } + } + } + } None } diff --git a/src/daemon/pty.rs b/src/daemon/pty.rs index 6007c7b..ccc30c0 100644 --- a/src/daemon/pty.rs +++ b/src/daemon/pty.rs @@ -47,6 +47,11 @@ pub fn spawn_shell(shell: &str, cwd: &Path) -> Result { let _ = std::env::set_current_dir("/"); } + // Ensure terminal environment is set for proper color and feature support. + // The daemon may have been started from a non-terminal context. + std::env::set_var("TERM", "xterm-256color"); + std::env::set_var("COLORTERM", "truecolor"); + // Exec shell let shell_cstr = CString::new(shell).unwrap_or_else(|_| CString::new("/bin/sh").unwrap()); diff --git a/src/daemon/registry.rs b/src/daemon/registry.rs index fc347c0..9c7d5ad 100644 --- a/src/daemon/registry.rs +++ b/src/daemon/registry.rs @@ -140,6 +140,14 @@ impl SessionRegistry { pub fn session_ids(&self) -> Vec { self.sessions.keys().cloned().collect() } + + /// Find a session whose pts_path matches the given path. + pub fn find_by_pts(&self, pts_path: &std::path::Path) -> Option<&str> { + self.sessions + .iter() + .find(|(_, s)| s.pts_path == pts_path) + .map(|(id, _)| id.as_str()) + } } #[cfg(test)] @@ -168,6 +176,7 @@ mod tests { registered: false, capture_path: None, capture_abort: None, + in_alternate_screen: false, } } diff --git a/src/daemon/server.rs b/src/daemon/server.rs index 31f5bdc..7742d6a 100644 --- a/src/daemon/server.rs +++ b/src/daemon/server.rs @@ -32,6 +32,8 @@ struct AttachedClient { tx: mpsc::Sender>, read_only: bool, session_id: Option, + peer_pid: Option, + last_size: Option<(u16, u16)>, // (cols, rows) } pub async fn run_daemon(config: Config) -> Result<()> { @@ -101,11 +103,14 @@ pub async fn run_daemon(config: Config) -> Result<()> { match event { DaemonEvent::NewConnection(stream) => { let client_id = NEXT_CLIENT_ID.fetch_add(1, Ordering::Relaxed); + let peer_pid = stream.peer_cred().ok().and_then(|c| c.pid().map(|p| p as u32)); let (tx, rx) = mpsc::channel(64); clients.insert(client_id, AttachedClient { tx, read_only: false, session_id: None, + peer_pid, + last_size: None, }); let event_tx = event_tx.clone(); tokio::spawn(handle_client_connection(client_id, stream, rx, event_tx)); @@ -150,6 +155,7 @@ pub async fn run_daemon(config: Config) -> Result<()> { } DaemonEvent::PtyData(session_id, data) => { if let Some(session) = registry.get_mut(&session_id) { + update_alternate_screen_state(session, &data); session.scrollback.write(&data); let attached: Vec = session.attached_clients.clone(); let output_frame = encode_response(&Response::PtyOutput(data.clone())).unwrap_or_default(); @@ -309,11 +315,15 @@ async fn handle_request( Request::SessionRename { target, new_name } => { handle_session_rename(registry, &target, new_name) } - Request::SessionList => handle_session_list(registry), + Request::SessionList => handle_session_list(registry, clients), Request::SessionInfo { target } => handle_session_info(registry, &target), - Request::SessionAttach { target, read_only } => { - handle_session_attach(registry, clients, client_id, &target, read_only, event_tx) - } + Request::SessionAttach { + target, + read_only, + force, + } => handle_session_attach( + registry, clients, client_id, &target, read_only, force, event_tx, + ), Request::SessionDetach => handle_session_detach(registry, clients, client_id), Request::SessionSend { target, input } => handle_session_send(registry, &target, &input), Request::SessionOutput { @@ -471,12 +481,12 @@ fn handle_session_kill(registry: &mut SessionRegistry, target: &str) -> Response } if let Some(session) = registry.remove(&id) { - // For registered sessions, just drop the fd — don't kill the shell. - // The shell continues running in the original terminal. - if !session.registered { - if let Some(pid) = session.child_pid { - pty::kill_session(pid); - } + if session.registered { + // Write notification to the wrapped session's terminal + let msg = b"\r\n\x1b[33m[Session unregistered by snag]\x1b[0m\r\n"; + let _ = nix::unistd::write(&session.master_fd, msg); + } else if let Some(pid) = session.child_pid { + pty::kill_session(pid); } Response::Ok(ResponseData::Empty) } else { @@ -513,8 +523,40 @@ fn handle_session_rename( } } -fn handle_session_list(registry: &SessionRegistry) -> Response { - let sessions: Vec<_> = registry.iter().map(|s| s.to_info()).collect(); +fn handle_session_list( + registry: &SessionRegistry, + clients: &HashMap, +) -> Response { + let sessions: Vec<_> = registry + .iter() + .map(|s| { + let mut info = s.to_info(); + // Find which session the active attacher is coming from + if !s.attached_clients.is_empty() { + let snagged_by = s + .attached_clients + .iter() + .filter_map(|&cid| { + let client = clients.get(&cid)?; + if client.read_only { + return None; + } + let pid = client.peer_pid?; + let tty = std::fs::read_link(format!("/proc/{pid}/fd/0")).ok()?; + let src_id = registry.find_by_pts(&tty)?; + let src = registry.get(src_id)?; + Some( + src.name + .clone() + .unwrap_or_else(|| src.id[..8.min(src.id.len())].to_string()), + ) + }) + .next(); + info.snagged_by = snagged_by; + } + info + }) + .collect(); Response::Ok(ResponseData::SessionList(sessions)) } @@ -534,11 +576,52 @@ fn handle_session_attach( client_id: ClientId, target: &str, read_only: bool, + force: bool, _event_tx: &mpsc::Sender, ) -> Response { match registry.resolve(target) { Ok(id) => { + // Detect attach loops (including self-attach) before taking mutable borrow + if let Some(cycle) = detect_attach_cycle(&id, client_id, registry, clients) { + return Response::Error { + code: 5, + message: format!("attach would create a loop: {cycle}"), + }; + } + if let Some(session) = registry.get_mut(&id) { + // Check for existing non-read-only attached client + if !read_only { + let active_client = session + .attached_clients + .iter() + .find(|&&cid| clients.get(&cid).map(|c| !c.read_only).unwrap_or(false)) + .copied(); + + if let Some(existing_cid) = active_client { + if !force { + return Response::Error { + code: 10, + message: "session already attached by another client (use --force to steal)".to_string(), + }; + } + // Force-steal: notify and detach the existing client + if let Some(evicted) = clients.get(&existing_cid) { + let stolen_msg = Response::SessionEvent { + event: "stolen".to_string(), + session_id: id.clone(), + }; + if let Ok(frame) = encode_response(&stolen_msg) { + let _ = evicted.tx.try_send(frame); + } + } + session.attached_clients.retain(|&cid| cid != existing_cid); + if let Some(evicted) = clients.get_mut(&existing_cid) { + evicted.session_id = None; + } + } + } + session.attached_clients.push(client_id); if let Some(client) = clients.get_mut(&client_id) { @@ -580,10 +663,23 @@ fn handle_session_detach( if let Some(ref session_id) = client.session_id.take() { if let Some(session) = registry.get_mut(session_id) { session.attached_clients.retain(|&id| id != client_id); - // If no more attached clients, signal snag wrap to resume - if session.registered && session.attached_clients.is_empty() { - if let Some(pid) = session.child_pid { - let _ = nix::sys::signal::kill(pid, nix::sys::signal::Signal::SIGUSR2); + if session.attached_clients.is_empty() { + // No more clients — signal snag wrap to resume and re-apply terminal size + if session.registered { + if let Some(pid) = session.child_pid { + let _ = nix::sys::signal::kill(pid, nix::sys::signal::Signal::SIGUSR2); + } + } + } else { + // Remaining clients exist — apply a remaining client's terminal size + // so the session PTY matches the active viewer's dimensions + let remaining_size = session + .attached_clients + .iter() + .filter_map(|cid| clients.get(cid).and_then(|c| c.last_size)) + .next(); + if let Some((cols, rows)) = remaining_size { + let _ = pty::set_winsize(session.raw_fd(), rows, cols); } } } @@ -642,6 +738,24 @@ fn handle_session_output( }; } + // When in alternate screen (vim, htop, etc.), show a placeholder + if session.in_alternate_screen && !follow { + let fg = pty::fg_process(&session.pts_path); + let app_name = fg + .iter() + .find(|(pid, _)| { + session + .child_pid + .map(|cp| *pid != cp.as_raw() as u32) + .unwrap_or(true) + }) + .map(|(_, cmd)| cmd.as_str()) + .unwrap_or("application"); + return Response::Ok(ResponseData::Output(format!( + "[TUI running: {app_name}]" + ))); + } + let output = if let Some(n) = lines { session.scrollback.last_n_lines(n as usize) } else { @@ -1017,11 +1131,14 @@ fn capture_dir() -> PathBuf { fn handle_resize( registry: &mut SessionRegistry, - clients: &HashMap, + clients: &mut HashMap, client_id: ClientId, cols: u16, rows: u16, ) -> Response { + if let Some(client) = clients.get_mut(&client_id) { + client.last_size = Some((cols, rows)); + } let session_id = clients.get(&client_id).and_then(|c| c.session_id.clone()); if let Some(id) = session_id { @@ -1068,3 +1185,95 @@ fn handle_daemon_status(registry: &SessionRegistry, start_time: Instant) -> Resp session_count: registry.len(), }) } + +/// Detect if attaching client_id to target_session_id would create a cycle. +/// Returns Some(chain_description) if a cycle is found, None otherwise. +/// +/// A cycle occurs when the target session's PTY is the terminal of a client +/// that is attached to a session whose PTY is the terminal of another client... +/// eventually reaching back to the session that the requesting client is in. +fn detect_attach_cycle( + target_session_id: &str, + client_id: ClientId, + registry: &SessionRegistry, + clients: &HashMap, +) -> Option { + // Find which session the requesting client is running in + let client_tty = clients + .get(&client_id) + .and_then(|c| c.peer_pid) + .and_then(|pid| std::fs::read_link(format!("/proc/{pid}/fd/0")).ok())?; + + let source_session = registry.find_by_pts(&client_tty)?; + + // Walk the attach chain from target_session_id + let mut visited = vec![source_session.to_string()]; + let mut current = target_session_id.to_string(); + + loop { + if current == source_session { + // Cycle detected + visited.push(current); + return Some(visited.join(" → ")); + } + + if visited.contains(¤t) { + break; // Already visited, no cycle back to source + } + visited.push(current.clone()); + + // Find non-read-only clients attached to this session, and which session they're in + let session = registry.get(¤t)?; + let next = session + .attached_clients + .iter() + .filter_map(|&cid| { + let c = clients.get(&cid)?; + if c.read_only { + return None; + } + let pid = c.peer_pid?; + let tty = std::fs::read_link(format!("/proc/{pid}/fd/0")).ok()?; + registry.find_by_pts(&tty).map(|s| s.to_string()) + }) + .next(); + + match next { + Some(next_id) => current = next_id, + None => break, // Chain ends, no cycle + } + } + + None +} + +/// Scan PTY output for alternate screen buffer transitions. +/// Tracks ESC[?1049h/l (xterm) and ESC[?47h/l (older terminals). +fn update_alternate_screen_state(session: &mut Session, data: &[u8]) { + // Look for \x1b[?1049h, \x1b[?1049l, \x1b[?47h, \x1b[?47l + let len = data.len(); + let mut i = 0; + while i < len { + if data[i] == 0x1b && i + 2 < len && data[i + 1] == b'[' && data[i + 2] == b'?' { + // Parse the numeric parameter + let start = i + 3; + let mut j = start; + while j < len && data[j].is_ascii_digit() { + j += 1; + } + if j < len && j > start { + let param = &data[start..j]; + if param == b"1049" || param == b"47" { + match data[j] { + b'h' => session.in_alternate_screen = true, + b'l' => session.in_alternate_screen = false, + _ => {} + } + i = j + 1; + continue; + } + } + } + i += 1; + } +} diff --git a/src/daemon/session.rs b/src/daemon/session.rs index db95c1f..dd9c024 100644 --- a/src/daemon/session.rs +++ b/src/daemon/session.rs @@ -32,6 +32,7 @@ pub struct Session { pub registered: bool, pub capture_path: Option, pub capture_abort: Option, + pub in_alternate_screen: bool, } impl Session { @@ -59,6 +60,7 @@ impl Session { registered: false, capture_path: None, capture_abort: None, + in_alternate_screen: false, } } @@ -88,6 +90,7 @@ impl Session { registered: true, capture_path, capture_abort: None, + in_alternate_screen: false, } } @@ -127,6 +130,7 @@ impl Session { attached: self.attached_clients.len(), registered: self.registered, created_at: self.created_at_utc.clone(), + snagged_by: None, } } diff --git a/src/main.rs b/src/main.rs index fd7ab5e..6be2d6d 100644 --- a/src/main.rs +++ b/src/main.rs @@ -36,9 +36,11 @@ async fn main() { Some(Command::Info { target, json }) => { cli::commands::cmd_info(&config, target, json).await } - Some(Command::Attach { target, read_only }) => { - cli::commands::cmd_attach(&config, target, read_only).await - } + Some(Command::Attach { + target, + read_only, + force, + }) => cli::commands::cmd_attach(&config, target, read_only, force).await, Some(Command::Send { target, command }) => { cli::commands::cmd_send(&config, target, command).await } diff --git a/src/protocol/codec.rs b/src/protocol/codec.rs index 760aebe..bcc8091 100644 --- a/src/protocol/codec.rs +++ b/src/protocol/codec.rs @@ -164,14 +164,20 @@ mod tests { let req = Request::SessionAttach { target: "abc123".to_string(), read_only: true, + force: false, }; let frame = encode_request(&req).unwrap(); assert_eq!(frame[0], MSG_SESSION_ATTACH); let decoded = decode_request(frame[0], &frame[5..]).unwrap(); match decoded { - Request::SessionAttach { target, read_only } => { + Request::SessionAttach { + target, + read_only, + force, + } => { assert_eq!(target, "abc123"); assert!(read_only); + assert!(!force); } _ => panic!("wrong request type"), } @@ -376,6 +382,7 @@ mod tests { attached: 1, registered: false, created_at: "2026-03-22T10:00:00Z".to_string(), + snagged_by: None, }])); let frame = encode_response(&resp).unwrap(); let decoded = decode_response(frame[0], &frame[5..]).unwrap(); diff --git a/src/protocol/types.rs b/src/protocol/types.rs index fd2fb3a..0f9d0fe 100644 --- a/src/protocol/types.rs +++ b/src/protocol/types.rs @@ -45,6 +45,8 @@ pub enum Request { SessionAttach { target: String, read_only: bool, + #[serde(default)] + force: bool, }, SessionDetach, SessionSend { @@ -131,6 +133,8 @@ pub struct SessionInfo { pub attached: usize, pub registered: bool, pub created_at: String, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub snagged_by: Option, } #[derive(Debug, Clone, Serialize, Deserialize)] diff --git a/src/tui/app.rs b/src/tui/app.rs index 7442f6c..796c428 100644 --- a/src/tui/app.rs +++ b/src/tui/app.rs @@ -14,6 +14,7 @@ pub struct App { pub preview_raw: String, pub input_mode: InputMode, pub input_buffer: String, + pub hide_snagged: bool, } impl App { @@ -25,32 +26,49 @@ impl App { preview_raw: String::new(), input_mode: InputMode::Normal, input_buffer: String::new(), + hide_snagged: false, } } + /// Returns the filtered list of sessions based on hide_snagged toggle. + pub fn visible_sessions(&self) -> Vec<&SessionInfo> { + self.sessions + .iter() + .filter(|s| !self.hide_snagged || s.snagged_by.is_none()) + .collect() + } + pub fn select_next(&mut self) { - if !self.sessions.is_empty() { - self.selected = (self.selected + 1) % self.sessions.len(); + let count = self.visible_sessions().len(); + if count > 0 { + self.selected = (self.selected + 1) % count; } } pub fn select_prev(&mut self) { - if !self.sessions.is_empty() { - self.selected = self - .selected - .checked_sub(1) - .unwrap_or(self.sessions.len() - 1); + let count = self.visible_sessions().len(); + if count > 0 { + self.selected = self.selected.checked_sub(1).unwrap_or(count - 1); } } pub fn selected_session(&self) -> Option<&SessionInfo> { - self.sessions.get(self.selected) + let visible = self.visible_sessions(); + visible.get(self.selected).copied() } pub fn selected_id(&self) -> Option { self.selected_session().map(|s| s.id.clone()) } + pub fn toggle_hide_snagged(&mut self) { + self.hide_snagged = !self.hide_snagged; + let count = self.visible_sessions().len(); + if self.selected >= count && count > 0 { + self.selected = count - 1; + } + } + pub fn enter_input_mode(&mut self, mode: InputMode) { self.input_mode = mode; self.input_buffer.clear(); diff --git a/src/tui/mod.rs b/src/tui/mod.rs index 9d5420b..3e4a0e4 100644 --- a/src/tui/mod.rs +++ b/src/tui/mod.rs @@ -95,7 +95,7 @@ async fn handle_normal_key( if let Some(id) = app.selected_id() { // Exit TUI, attach to session cleanup_terminal(terminal)?; - crate::cli::commands::cmd_attach(config, id, false).await?; + crate::cli::commands::cmd_attach(config, id, false, false).await?; // Re-enter TUI after detach terminal::enable_raw_mode()?; crossterm::execute!(io::stdout(), crossterm::terminal::EnterAlternateScreen)?; @@ -135,6 +135,9 @@ async fn handle_normal_key( app.enter_input_mode(InputMode::Send); } } + KeyCode::Char('h') => { + app.toggle_hide_snagged(); + } _ => {} } Ok(()) diff --git a/src/tui/ui.rs b/src/tui/ui.rs index 7116c78..6376837 100644 --- a/src/tui/ui.rs +++ b/src/tui/ui.rs @@ -16,9 +16,9 @@ pub fn draw(frame: &mut Frame, app: &App) { ]) .split(frame.area()); - // Session list - let items: Vec = app - .sessions + // Session list (filtered by hide_snagged toggle) + let visible = app.visible_sessions(); + let items: Vec = visible .iter() .enumerate() .map(|(i, s)| { @@ -29,11 +29,18 @@ pub fn draw(frame: &mut Frame, app: &App) { let marker = if i == app.selected { "▸ " } else { " " }; let type_marker = if s.registered { " [R]" } else { "" }; + let snagged = s + .snagged_by + .as_deref() + .map(|by| format!(" ← {by}")) + .unwrap_or_default(); let style = if i == app.selected { Style::default() .fg(Color::Yellow) .add_modifier(Modifier::BOLD) + } else if s.snagged_by.is_some() { + Style::default().fg(Color::Magenta) } else if s.state != "running" { Style::default().fg(Color::DarkGray) } else { @@ -42,7 +49,7 @@ pub fn draw(frame: &mut Frame, app: &App) { let id_short = &s.id[..8.min(s.id.len())]; ListItem::new(Line::from(vec![Span::styled( - format!("{marker}{id_short} {name:<10} {shell:<6} {cwd:<25} {fg}{type_marker}"), + format!("{marker}{id_short} {name:<10} {shell:<6} {cwd:<25} {fg}{type_marker}{snagged}"), style, )])) }) @@ -110,7 +117,14 @@ pub fn draw(frame: &mut Frame, app: &App) { // Status bar let status_text = match app.input_mode { - InputMode::Normal => " [n]ew [x]kill [r]ename [s]end [Enter]attach [q]uit".to_string(), + InputMode::Normal => { + let hide_label = if app.hide_snagged { + "[h]show snagged" + } else { + "[h]hide snagged" + }; + format!(" [n]ew [x]kill [r]ename [s]end [Enter]attach {hide_label} [q]uit") + } InputMode::Rename => { format!(" Rename: {} [Enter]confirm [Esc]cancel", app.input_buffer) }