Skip to main content

koh/
lib.rs

1//! The library behind the `koh` remote shell: `koh serve` hosts PTYs for allowlisted peers and
2//! `koh connect` relays one to the user's terminal with predictive local echo, reconnecting
3//! transparently; `koh exec` runs a command there, `koh run` a job that outlives the connection,
4//! and `koh cp` brings files into line.
5//!
6//! [`transport_iroh`] sets up endpoints; [`identity`] loads the keys; [`keycmd`] is `koh id` and
7//! `koh key reset`; [`names`] names the servers and clients a machine knows. [`request`] is what a
8//! connection opens with, [`proto`] a shell's wire protocol, [`stream`] what of a program's output
9//! reaches the user's terminal, [`events`] the input it carries, [`sync`] a copy's; [`terminal`],
10//! [`predict`] and [`pty`] the server's emulator, the prediction and the PTY; [`server`] hosts
11//! sessions, commands and jobs and [`client`] drives them. [`bounded`] holds every queue.
12//!
13//! Production code is panic-free by construction: `Cargo.toml` forbids the panic lints, lossy casts
14//! and unchecked arithmetic, and a `forbid` cannot be lifted locally. Tests may panic.
15//!
16//! The library exists so the binary, tests and fuzz targets share code; all of it is internal.
17
18pub mod bounded;
19pub mod client;
20pub mod events;
21pub mod identity;
22pub mod keycmd;
23pub mod log;
24pub mod menu;
25pub mod names;
26pub mod predict;
27pub mod proto;
28pub mod pty;
29pub mod request;
30pub mod server;
31pub mod stream;
32pub mod sync;
33pub mod terminal;
34pub mod transport_iroh;
35
36#[cfg(test)]
37mod test_runtime;
38
39/// Cancel `token` on the first SIGTERM, SIGINT or SIGHUP, so `koh serve` drains and `koh connect`
40/// restores the terminal instead of the process dying where it stands; the token returned is
41/// cancelled on a later one ([`Shutdown`]), when waiting longer is not wanted. A SIGHUP ignored as
42/// the process started (`nohup`) stays ignored. Fails only if a handler cannot be installed.
43pub(crate) fn cancel_on_signals(
44    token: &tokio_util::sync::CancellationToken,
45) -> std::io::Result<tokio_util::sync::CancellationToken> {
46    use tokio::signal::unix::{signal, SignalKind};
47    // Read before the handler replaces it. The handler is installed even so: a program started
48    // later gets SIGHUP's default at its `exec`, as a hosted shell must to hang up with its session.
49    let nohup = ignored_at_start() & 1 == 1;
50    let mut term = signal(SignalKind::terminate())?;
51    let mut int = signal(SignalKind::interrupt())?;
52    let mut hup = signal(SignalKind::hangup())?;
53    let (draining, forced) = (token.clone(), tokio_util::sync::CancellationToken::new());
54    let force = forced.clone();
55    #[expect(clippy::disallowed_methods, reason = "the one signal task")]
56    tokio::spawn(async move {
57        let mut shutdown = Shutdown::default();
58        loop {
59            let name = tokio::select! {
60                Some(()) = term.recv() => "SIGTERM",
61                Some(()) = int.recv() => "SIGINT",
62                Some(()) = hup.recv() => if nohup {
63                    tracing::info!("SIGHUP ignored, as it was when koh started (nohup)");
64                    continue;
65                } else {
66                    "SIGHUP"
67                },
68                else => return,
69            };
70            tracing::info!(signal = name, "shutting down");
71            let _ = STOPPED_BY.set(name);
72            match shutdown.signal(std::time::Instant::now()) {
73                Shutdown::Draining(_) => draining.cancel(),
74                Shutdown::Forced => force.cancel(),
75                Shutdown::Running => {}
76            }
77        }
78    });
79    Ok(forced)
80}
81
82/// The first signal [`cancel_on_signals`] took, by name.
83static STOPPED_BY: std::sync::OnceLock<&'static str> = std::sync::OnceLock::new();
84
85/// The signal that stopped this process, by name, if one did ([`cancel_on_signals`]).
86pub(crate) fn stopped_by() -> Option<&'static str> {
87    STOPPED_BY.get().copied()
88}
89
90/// The signals this process was started ignoring, a bit each (signal `n` is bit `n - 1`): the
91/// `SigIgn` mask Linux and Android show; elsewhere none is known.
92fn ignored_at_start() -> u64 {
93    let status = std::fs::read_to_string("/proc/self/status").unwrap_or_default();
94    let mask = status.lines().find_map(|line| line.strip_prefix("SigIgn:"));
95    mask.and_then(|mask| u64::from_str_radix(mask.trim(), 16).ok())
96        .unwrap_or(0)
97}
98
99/// Have the programs this process starts begin with every signal at its default, as from a login,
100/// though it was started ignoring some: a script's `&` ignores SIGINT and SIGQUIT, and an ignored
101/// signal stays ignored across `exec`, so a hosted shell's `Ctrl-\` would do nothing. Each is
102/// caught instead, by a handler that does nothing: this process goes on as if it ignored it, and
103/// `exec` gives a caught signal its default. Not SIGPIPE, which Rust ignores itself and gives its
104/// default back to what it starts, nor those [`cancel_on_signals`] catches, SIGHUP among them.
105/// Only where the mask is known ([`ignored_at_start`]).
106pub(crate) fn default_signals_for_children() {
107    use tokio::signal::unix::{signal, SignalKind};
108    const CAUGHT: [i32; 4] = [1, 2, 13, 15];
109    let ignored = ignored_at_start();
110    for (number, bit) in (1..=31_i32).zip(0..31_u32) {
111        if ignored.checked_shr(bit).unwrap_or(0) & 1 == 0 || CAUGHT.contains(&number) {
112            continue;
113        }
114        // Once installed, the handler stays when its stream is dropped.
115        if let Err(error) = signal(SignalKind::from_raw(number)) {
116            tracing::debug!(number, %error, "an ignored signal is left ignored");
117        }
118    }
119}
120
121/// Where a process is in shutting down, by the signals it got. Counted apart from what they
122/// cancel: a drain the process began on its own (its endpoint closed) is not one a signal began.
123#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
124enum Shutdown {
125    #[default]
126    Running,
127    /// The first signal came, at this instant: drain.
128    Draining(std::time::Instant),
129    /// Another came, [`Shutdown::DEBOUNCE`] or more after the first: stop waiting.
130    Forced,
131}
132
133impl Shutdown {
134    /// Signals this close after the first are one event: a terminal closing sends its session a
135    /// SIGHUP, and its shell sends another, and that must not skip the drain.
136    const DEBOUNCE: std::time::Duration = std::time::Duration::from_secs(1);
137
138    /// A signal came at `now`: where that leaves the process.
139    fn signal(&mut self, now: std::time::Instant) -> Self {
140        *self = match *self {
141            Self::Running => Self::Draining(now),
142            Self::Draining(first) if now.saturating_duration_since(first) < Self::DEBOUNCE => {
143                Self::Draining(first)
144            }
145            Self::Draining(_) | Self::Forced => Self::Forced,
146        };
147        *self
148    }
149}
150
151#[cfg(test)]
152mod tests {
153    use super::Shutdown;
154
155    #[test]
156    fn a_double_sighup_still_drains_and_a_later_signal_forces() {
157        let t0 = std::time::Instant::now();
158        let mut s = Shutdown::default();
159        assert_eq!(s.signal(t0), Shutdown::Draining(t0));
160        let echo = t0 + std::time::Duration::from_millis(5);
161        assert_eq!(s.signal(echo), Shutdown::Draining(t0), "the same event");
162        assert_eq!(s.signal(t0 + Shutdown::DEBOUNCE), Shutdown::Forced);
163    }
164}