koh 0.1.0

//! # koh-pty — PTY allocation, shell spawn, resize, reaping
//!
//! The server side's plumbing to the real shell. Allocates a pseudo-terminal, spawns the
//! user's login shell under it, pumps the child's output to an async channel from a dedicated
//! blocking thread (portable-pty's reader is blocking-only), forwards input bytes to the
//! child via a second dedicated thread (so a slow child never blocks a tokio worker), and
//! propagates window-size changes (which `ioctl(TIOCSWINSZ)` turns into `SIGWINCH`).

use std::io::{self, Read, Write};
use std::sync::mpsc::{sync_channel, SyncSender, TrySendError};

use portable_pty::{
    native_pty_system, ChildKiller, CommandBuilder, ExitStatus, MasterPty, PtySize,
};
use tokio::sync::mpsc;

/// Size of each output chunk read from the PTY master.
const READ_CHUNK: usize = 8192;
/// Bound on the output channel (chunks). Backpressure here naturally slows the reader thread.
const OUTPUT_CHANNEL_DEPTH: usize = 512;
/// Bound on the input channel (chunks) feeding the writer thread. Generous, because under normal
/// interactive use the child drains its input promptly; a full queue means the child has stopped
/// reading (flow-controlled or hung), which [`Pty::write_input`] surfaces rather than blocking on.
const WRITE_CHANNEL_DEPTH: usize = 1024;

/// Typed errors from PTY allocation, shell spawn, and resize (mirrors the
/// `transport-iroh::SetupError` pattern so callers can match on the failure stage).
///
/// `portable-pty` surfaces its failures as `anyhow::Error`; we fold those into `io::Error`
/// (via [`io::Error::other`]) so every variant carries one concrete payload. Only the reader
/// thread's `Builder::spawn` is natively an `io::Error`, so it is the single `#[from]` source.
/// Binaries keep `anyhow` internally — their `?`/`.context()` absorb `PtyError` via anyhow's
/// blanket `From<E: Error + Send + Sync>`.
#[derive(Debug, thiserror::Error)]
pub enum PtyError {
    /// Allocating the pseudo-terminal pair (`openpty`) failed.
    #[error("opening pty: {0}")]
    OpenPty(#[source] io::Error),
    /// Spawning the shell under the slave side (`spawn_command`) failed.
    #[error("spawning shell: {0}")]
    Spawn(#[source] io::Error),
    /// Wiring up the master read/write pumps failed: cloning the reader, taking the writer, or
    /// starting the blocking reader thread (`Builder::spawn`, the native `io::Error` source).
    #[error("starting pty reader: {0}")]
    Reader(#[from] io::Error),
    /// Propagating a window-size change to the kernel (`master.resize`) failed.
    #[error("resizing pty: {0}")]
    Resize(#[source] io::Error),
}

/// A running shell behind a PTY.
///
/// Construct with [`Pty::spawn`], which also returns the receiver of the child's output.
/// Hold the `Pty` for the life of the session: dropping it drops `writer_tx`, which lets the
/// writer thread finish and drop the PTY's write handle — and `portable-pty` writes an EOT
/// (Ctrl-D) on that drop, so the child sees EOF on its stdin.
pub struct Pty {
    master: Box<dyn MasterPty + Send>,
    /// Bounded sender to the dedicated writer thread (which owns the blocking `Box<dyn Write>`).
    /// Shared by both input producers (keystrokes + host query replies), so writes stay FIFO.
    writer_tx: SyncSender<Vec<u8>>,
    child: Box<dyn portable_pty::Child + Send + Sync>,
    killer: Box<dyn ChildKiller + Send + Sync>,
    /// Join handles for the reader/writer pump threads, kept so a graceful [`Pty::shutdown`] can
    /// join them rather than leaking detached threads. `None` only after `shutdown` takes them.
    reader_handle: Option<std::thread::JoinHandle<()>>,
    writer_handle: Option<std::thread::JoinHandle<()>>,
}

impl Pty {
    /// Allocate a PTY of `rows`×`cols`, spawn `shell` (or the user's default login shell when
    /// `None`) with `TERM` set, and start streaming its output.
    ///
    /// Returns the [`Pty`] handle plus an async receiver of raw output chunks. The reader runs
    /// on a dedicated OS thread; when the child closes the PTY the channel ends.
    pub fn spawn(
        rows: u16,
        cols: u16,
        shell: Option<&str>,
        term: &str,
    ) -> Result<(Self, mpsc::Receiver<Vec<u8>>), PtyError> {
        let pty_system = native_pty_system();
        let pair = pty_system
            .openpty(PtySize {
                rows,
                cols,
                pixel_width: 0,
                pixel_height: 0,
            })
            .map_err(|e| PtyError::OpenPty(io::Error::other(e)))?;

        let mut cmd = match shell {
            Some(prog) => CommandBuilder::new(prog),
            None => CommandBuilder::new_default_prog(),
        };
        // A real terminal type so curses apps behave; the env is otherwise inherited.
        cmd.env("TERM", term);

        let child = pair
            .slave
            .spawn_command(cmd)
            .map_err(|e| PtyError::Spawn(io::Error::other(e)))?;
        let killer = child.clone_killer();
        // The slave fd is now owned by the child; drop our handle so EOF propagates correctly.
        drop(pair.slave);

        let mut reader = pair
            .master
            .try_clone_reader()
            .map_err(|e| PtyError::Reader(io::Error::other(e)))?;
        let mut writer = pair
            .master
            .take_writer()
            .map_err(|e| PtyError::Reader(io::Error::other(e)))?;

        let (tx, rx) = mpsc::channel::<Vec<u8>>(OUTPUT_CHANNEL_DEPTH);
        let reader_handle = std::thread::Builder::new()
            .name("koh-pty-reader".into())
            .spawn(move || {
                let mut buf = [0u8; READ_CHUNK];
                loop {
                    match reader.read(&mut buf) {
                        Ok(0) => break, // EOF: slave closed (EIO is mapped to 0 on unix)
                        // `Read::read` guarantees `n <= buf.len()`, so `get(..n)` is always
                        // `Some`; the `else` is a panic-free fallback that can't actually run.
                        Ok(n) => {
                            let Some(chunk) = buf.get(..n) else { break };
                            if tx.blocking_send(chunk.to_vec()).is_err() {
                                break; // receiver dropped: session over
                            }
                        }
                        Err(e) => {
                            tracing::debug!(error = %e, "pty reader stopping");
                            break;
                        }
                    }
                }
            })?;

        // Dedicated writer thread: it owns the blocking `Box<dyn Write>` and drains the bounded
        // input channel, so `write_input` never blocks a tokio worker. `recv()` yields every
        // buffered chunk before it observes the senders being dropped, so pending writes flush
        // before the writer is dropped (and `portable-pty` then writes the EOT that EOFs the
        // child). The thread exits as soon as the last sender (held in `Pty`) drops.
        let (writer_tx, writer_rx) = sync_channel::<Vec<u8>>(WRITE_CHANNEL_DEPTH);
        let writer_handle = std::thread::Builder::new()
            .name("koh-pty-writer".into())
            .spawn(move || {
                while let Ok(chunk) = writer_rx.recv() {
                    if writer
                        .write_all(&chunk)
                        .and_then(|()| writer.flush())
                        .is_err()
                    {
                        break; // master closed / child gone
                    }
                }
                // `writer` drops here -> portable-pty sends EOT -> child sees EOF on stdin.
            })?;

        Ok((
            Self {
                master: pair.master,
                writer_tx,
                child,
                killer,
                reader_handle: Some(reader_handle),
                writer_handle: Some(writer_handle),
            },
            rx,
        ))
    }

    /// Gracefully tear down the session and join both I/O pump threads (rather than leaking them
    /// as detached threads). Consumes the `Pty`. It first kills the child — so the reader's
    /// blocking `read` returns EOF — then drops the writer sender — so the writer's `recv` returns
    /// — guaranteeing both threads unblock before we join them, so this never deadlocks.
    pub fn shutdown(mut self) {
        // A failed kill is logged, not ignored: if the child somehow survives it keeps the slave
        // fd open, the reader stays blocked on read(), and the join below would hang — so a warning
        // is the breadcrumb for that (otherwise impossible-looking) stall.
        if let Err(e) = self.killer.kill() {
            tracing::warn!(error = %e, "pty kill on shutdown failed; reader join may stall");
        }
        let reader = self.reader_handle.take();
        let writer = self.writer_handle.take();
        // Dropping `self` drops `writer_tx`, which lets the writer thread observe the channel
        // close and exit; the child kill above lets the reader thread hit EOF and exit.
        drop(self);
        if let Some(h) = writer {
            let _ = h.join();
        }
        if let Some(h) = reader {
            let _ = h.join();
        }
    }

    /// Forward input bytes to the child (verbatim — keystrokes or host query replies).
    ///
    /// Takes `&self` and never blocks: it enqueues `data` onto the bounded channel feeding the
    /// writer thread. Both producers share one sender, and callers enqueue while holding the
    /// session lock, so bytes stay FIFO (a DSR reply can't overtake the keystroke that triggered
    /// it). Returns [`io::ErrorKind::BrokenPipe`] if the writer thread is gone, and
    /// [`io::ErrorKind::WouldBlock`] if the queue is full — the defined over-limit policy: surface
    /// backpressure rather than block a tokio worker or silently drop input (a full 1024-deep
    /// queue means the child has stopped reading, i.e. the session is effectively dead).
    pub fn write_input(&self, data: &[u8]) -> io::Result<()> {
        match self.writer_tx.try_send(data.to_vec()) {
            Ok(()) => Ok(()),
            Err(TrySendError::Full(_)) => Err(io::Error::new(
                io::ErrorKind::WouldBlock,
                "pty writer queue full (child not draining its input)",
            )),
            Err(TrySendError::Disconnected(_)) => Err(io::Error::from(io::ErrorKind::BrokenPipe)),
        }
    }

    /// Propagate a window-size change; the kernel raises `SIGWINCH` in the child.
    pub fn resize(&self, rows: u16, cols: u16) -> Result<(), PtyError> {
        self.master
            .resize(PtySize {
                rows,
                cols,
                pixel_width: 0,
                pixel_height: 0,
            })
            .map_err(|e| PtyError::Resize(io::Error::other(e)))
    }

    /// Non-blocking check for child exit.
    pub fn try_wait(&mut self) -> std::io::Result<Option<ExitStatus>> {
        self.child.try_wait()
    }

    /// Block until the child exits, returning its status.
    pub fn wait(&mut self) -> std::io::Result<ExitStatus> {
        self.child.wait()
    }

    /// A standalone killer that can be moved to another thread/task to terminate the child.
    pub fn killer(&self) -> Box<dyn ChildKiller + Send + Sync> {
        self.killer.clone_killer()
    }

    /// Terminate the child (SIGHUP, then a hard kill if it lingers).
    pub fn kill(&mut self) -> std::io::Result<()> {
        self.killer.kill()
    }

    /// The child's process id, if known.
    pub fn process_id(&self) -> Option<u32> {
        self.child.process_id()
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::time::Duration;

    #[test]
    #[allow(
        clippy::items_after_statements,
        reason = "`_assert_typed` is a deliberate compile-time signature assertion kept beside the runtime checks it documents"
    )]
    fn pty_error_variants_are_constructible_and_reachable() {
        let mk = || io::Error::other("boom");
        // Each stage variant is constructible and renders a non-empty message.
        for e in [
            PtyError::OpenPty(mk()),
            PtyError::Spawn(mk()),
            PtyError::Reader(mk()),
            PtyError::Resize(mk()),
        ] {
            assert!(!e.to_string().is_empty(), "variant must Display");
        }
        // The `#[from] io::Error` source (the reader-thread spawn path) yields `Reader`.
        let from_io: PtyError = mk().into();
        assert!(matches!(from_io, PtyError::Reader(_)));
        // A binary's `?`/`.context()` absorbs PtyError via anyhow's blanket `From` — the
        // typed error stays internal to the lib but composes with anyhow at the edges.
        let absorbed: anyhow::Error = PtyError::OpenPty(mk()).into();
        assert!(absorbed.to_string().contains("opening pty"));
        // The public spawn signature now carries the typed error.
        fn _assert_typed(r: Result<(), PtyError>) -> Result<(), PtyError> {
            r
        }
    }

    #[tokio::test]
    #[allow(
        clippy::match_wild_err_arm,
        reason = "a timeout in this test IS the test failing; panicking on the `Err(_)` deadline arm is the intended assertion"
    )]
    async fn spawns_and_streams_output() {
        // Run a one-shot command in the PTY and confirm we receive its output + reap it.
        let (mut pty, mut rx) =
            Pty::spawn(24, 80, Some("echo"), "xterm-256color").expect("spawn echo");
        // `CommandBuilder::new("echo")` then arg is awkward here (we only take a program),
        // so instead drive a tiny shell snippet via the default shell path below if needed.
        // `echo` with no args prints just a newline; assert we get *something* and EOF.
        let mut collected = Vec::new();
        let deadline = tokio::time::Instant::now() + Duration::from_secs(20);
        loop {
            match tokio::time::timeout_at(deadline, rx.recv()).await {
                Ok(Some(chunk)) => collected.extend_from_slice(&chunk),
                Ok(None) => break, // channel closed: child exited and reader finished
                Err(_) => panic!("timed out waiting for pty output"),
            }
        }
        // `echo` prints a newline (CR/LF in a pty).
        assert!(
            collected.contains(&b'\n'),
            "expected a newline from echo, got {collected:?}"
        );
        // Child should be reapable.
        let status = pty.wait().expect("wait");
        assert!(status.success() || status.exit_code() == 0);
    }

    #[tokio::test]
    #[allow(
        clippy::match_same_arms,
        reason = "channel-close (`Ok(None)`) and deadline (`Err(_)`) are conceptually distinct outcomes kept as separate arms for readability, even though both set `found = false`"
    )]
    async fn interactive_shell_echoes_input() {
        // Spawn the default shell, send a command, and verify the echoed output comes back.
        let (mut pty, mut rx) = Pty::spawn(24, 80, None, "xterm-256color").expect("spawn shell");
        // Give the shell a moment to start, then type a command that prints a marker.
        tokio::time::sleep(Duration::from_millis(300)).await;
        pty.write_input(b"printf KOH_MARKER_OK\n").expect("write");

        let mut collected = Vec::new();
        let deadline = tokio::time::Instant::now() + Duration::from_secs(20);
        let found = loop {
            match tokio::time::timeout_at(deadline, rx.recv()).await {
                Ok(Some(chunk)) => {
                    collected.extend_from_slice(&chunk);
                    if String::from_utf8_lossy(&collected).contains("KOH_MARKER_OK") {
                        break true;
                    }
                }
                Ok(None) => break false,
                Err(_) => break false,
            }
        };
        // Resize should not error while the shell is live.
        let _ = pty.resize(40, 120);
        let _ = pty.kill();
        assert!(
            found,
            "did not observe the marker in shell output: {}",
            String::from_utf8_lossy(&collected)
        );
    }

    #[tokio::test]
    #[allow(
        clippy::match_same_arms,
        reason = "channel-close (`Ok(None)`) and deadline (`Err(_)`) are conceptually distinct outcomes kept as separate arms for readability, even though both set `in_order = false`"
    )]
    async fn write_input_takes_shared_ref_and_preserves_order() {
        // `pty` is bound WITHOUT `mut`, proving write_input takes `&self`. Two separate enqueues
        // must reach the child in FIFO order: the concatenated marker only appears if the second
        // chunk did not overtake the first.
        let (pty, mut rx) = Pty::spawn(24, 80, None, "xterm-256color").expect("spawn shell");
        tokio::time::sleep(Duration::from_millis(300)).await;
        pty.write_input(b"printf ORDER_").expect("first enqueue");
        pty.write_input(b"AB_CD\n").expect("second enqueue");

        let mut collected = Vec::new();
        let deadline = tokio::time::Instant::now() + Duration::from_secs(20);
        let in_order = loop {
            match tokio::time::timeout_at(deadline, rx.recv()).await {
                Ok(Some(chunk)) => {
                    collected.extend_from_slice(&chunk);
                    if String::from_utf8_lossy(&collected).contains("ORDER_AB_CD") {
                        break true;
                    }
                }
                Ok(None) => break false,
                Err(_) => break false,
            }
        };
        drop(pty);
        assert!(
            in_order,
            "FIFO ordering of two enqueues should yield ORDER_AB_CD; got: {}",
            String::from_utf8_lossy(&collected)
        );
    }

    #[tokio::test]
    #[allow(
        clippy::needless_continue,
        clippy::match_wild_err_arm,
        reason = "the explicit `continue` documents the drain-and-keep-reading intent; the `Err(_)` deadline arm panics because a timeout here IS the test failing"
    )]
    async fn dropping_pty_eofs_child_and_stops_writer() {
        // `cat` blocks reading stdin. Dropping the Pty drops the writer-thread sender; the writer
        // thread then finishes and drops the PTY write handle, on which portable-pty sends EOT —
        // so the child sees EOF, exits, the slave closes, and the output channel ends. If the
        // writer thread were stuck (or never dropped its handle), the channel would never close.
        let (pty, mut rx) = Pty::spawn(24, 80, Some("cat"), "xterm-256color").expect("spawn cat");
        tokio::time::sleep(Duration::from_millis(200)).await;
        drop(pty); // no kill(): EOF must come purely from the writer handle being dropped

        let deadline = tokio::time::Instant::now() + Duration::from_secs(20);
        loop {
            match tokio::time::timeout_at(deadline, rx.recv()).await {
                Ok(Some(_)) => continue, // drain any echoed bytes
                Ok(None) => break, // channel closed: child EOF'd + exited; writer thread ended
                Err(_) => panic!("dropping Pty did not EOF the child within 5s (writer stuck?)"),
            }
        }
    }

    #[tokio::test]
    async fn shutdown_joins_both_io_threads_without_deadlock() {
        // Graceful teardown: shutdown() kills the child (so the reader's blocking read returns
        // EOF) and drops the writer sender (so the writer's recv returns), then joins BOTH pump
        // threads. It must return promptly — a hang would mean a thread never unblocked.
        let (pty, mut rx) = Pty::spawn(24, 80, Some("sh"), "xterm-256color").expect("spawn shell");
        // Keep the output channel drained so the reader thread never blocks on a full channel.
        let drain = tokio::spawn(async move { while rx.recv().await.is_some() {} });
        tokio::time::sleep(Duration::from_millis(200)).await;

        tokio::time::timeout(
            Duration::from_secs(20),
            tokio::task::spawn_blocking(move || pty.shutdown()),
        )
        .await
        .expect("shutdown must not deadlock (both threads must unblock and join)")
        .expect("shutdown task panicked");
        let _ = drain.await;
    }
}