lios 0.1.89

use super::cli::{Command as Launch, Options};
use std::collections::VecDeque;
use std::fs::File;
use std::io::{self, Read, Write};
use std::os::fd::{AsRawFd, FromRawFd, OwnedFd};
use std::os::unix::process::CommandExt;
use std::process::{Child, Command, ExitStatus, Stdio};
use std::time::{Duration, Instant};

const MAX_PENDING_INPUT: usize = 64 * 1024;
pub const READ_BUDGET: usize = 64 * 1024;

pub struct Pty {
    master: File,
    child: Child,
    pending: VecDeque<u8>,
    pub eof: bool,
    status: Option<ExitStatus>,
}

impl Pty {
    pub fn spawn(options: &Options, rows: u16, cols: u16) -> io::Result<Self> {
        let mut master = -1;
        let mut slave = -1;
        let size = winsize(rows, cols);
        // SAFETY: output pointers and the initialized size remain valid for this call.
        if unsafe {
            libc::openpty(
                &mut master,
                &mut slave,
                std::ptr::null_mut(),
                std::ptr::null(),
                &size,
            )
        } < 0
        {
            return Err(io::Error::last_os_error());
        }
        // SAFETY: openpty returned two newly owned file descriptors.
        let master = unsafe { OwnedFd::from_raw_fd(master) };
        let slave = unsafe { OwnedFd::from_raw_fd(slave) };
        for fd in [&master, &slave] {
            // SAFETY: fd is an owned, live descriptor; only descriptor flags change.
            if unsafe { libc::fcntl(fd.as_raw_fd(), libc::F_SETFD, libc::FD_CLOEXEC) } < 0 {
                return Err(io::Error::last_os_error());
            }
        }
        let shell = std::env::var_os("SHELL")
            .filter(|s| !s.is_empty())
            .unwrap_or_else(|| "/bin/sh".into());
        let mut command = match &options.command {
            Launch::Shell => {
                let mut c = Command::new(&shell);
                c.arg("-i");
                c
            }
            Launch::Script(script) => {
                let mut c = Command::new(&shell);
                c.arg("-c").arg(script);
                c
            }
            Launch::Argv(argv) => {
                let mut c = Command::new(&argv[0]);
                c.args(&argv[1..]);
                c
            }
        };
        if let Some(path) = &options.directory {
            command.current_dir(path);
        }
        command
            .env("TERM", "xterm-256color")
            .env("COLORTERM", "truecolor")
            .env("TERM_PROGRAM", "lios")
            .env("TERM_PROGRAM_VERSION", env!("CARGO_PKG_VERSION"))
            .env_remove("DISPLAY")
            .env_remove("WAYLAND_DISPLAY")
            .env_remove("WAYLAND_SOCKET")
            .stdin(Stdio::from(slave.try_clone()?))
            .stdout(Stdio::from(slave.try_clone()?))
            .stderr(Stdio::from(slave));
        // SAFETY: this callback executes only async-signal-safe libc calls after
        // fork. Stdio has already installed the slave on descriptors 0, 1 and 2.
        unsafe {
            command.pre_exec(|| {
                if libc::setsid() < 0 || libc::ioctl(0, libc::TIOCSCTTY, 0) < 0 {
                    return Err(io::Error::last_os_error());
                }
                Ok(())
            });
        }
        // Set nonblocking mode on the master only, never on the child's slave.
        // SAFETY: master is live and both fcntl commands use integer flags.
        let flags = unsafe { libc::fcntl(master.as_raw_fd(), libc::F_GETFL) };
        if flags < 0
            || unsafe { libc::fcntl(master.as_raw_fd(), libc::F_SETFL, flags | libc::O_NONBLOCK) }
                < 0
        {
            return Err(io::Error::last_os_error());
        }
        let child = command.spawn()?;
        Ok(Self {
            master: File::from(master),
            child,
            pending: VecDeque::new(),
            eof: false,
            status: None,
        })
    }

    pub fn resize(&self, rows: u16, cols: u16) -> io::Result<()> {
        let size = winsize(rows, cols);
        // SAFETY: the master and the initialized size pointer are live.
        if unsafe { libc::ioctl(self.master.as_raw_fd(), libc::TIOCSWINSZ, &size) } < 0 {
            return Err(io::Error::last_os_error());
        }
        Ok(())
    }
    pub fn read(&mut self, buffer: &mut [u8]) -> io::Result<usize> {
        if self.eof {
            return Ok(0);
        }
        match self.master.read(buffer) {
            Ok(0) => {
                self.eof = true;
                Ok(0)
            }
            Err(error) if error.raw_os_error() == Some(libc::EIO) => {
                self.eof = true;
                Ok(0)
            }
            other => other,
        }
    }
    pub fn enqueue(&mut self, bytes: &[u8]) -> io::Result<()> {
        if self.eof {
            return Ok(());
        }
        if bytes.len() > MAX_PENDING_INPUT.saturating_sub(self.pending.len()) {
            return Err(io::Error::other("Lios PTY input queue exceeded its limit"));
        }
        self.pending.extend(bytes);
        Ok(())
    }
    pub fn flush(&mut self) -> io::Result<()> {
        while !self.pending.is_empty() {
            let (bytes, _) = self.pending.as_slices();
            match self.master.write(bytes) {
                Ok(0) => return Err(io::ErrorKind::WriteZero.into()),
                Ok(count) => {
                    self.pending.drain(..count);
                }
                Err(error) if error.kind() == io::ErrorKind::WouldBlock => break,
                Err(error) if error.kind() == io::ErrorKind::Interrupted => continue,
                Err(error) if error.raw_os_error() == Some(libc::EIO) => {
                    self.pending.clear();
                    self.eof = true;
                    break;
                }
                Err(error) => return Err(error),
            }
        }
        Ok(())
    }
    pub fn available_input_bytes(&self) -> usize {
        if self.eof {
            0
        } else {
            MAX_PENDING_INPUT.saturating_sub(self.pending.len())
        }
    }
    pub fn wants_write(&self) -> bool {
        !self.pending.is_empty() && !self.eof
    }
    pub fn status(&mut self) -> io::Result<Option<ExitStatus>> {
        if self.status.is_none() {
            self.status = self.child.try_wait()?;
        }
        Ok(self.status)
    }
}
impl AsRawFd for Pty {
    fn as_raw_fd(&self) -> std::os::fd::RawFd {
        self.master.as_raw_fd()
    }
}
impl Drop for Pty {
    fn drop(&mut self) {
        if self.status.is_some() {
            return;
        }
        let group = self.child.id() as i32;
        // SAFETY: the child owns a new session with this PTY. It has not been
        // reaped, so its PID cannot be reused. Never signal our own process group.
        unsafe {
            let foreground = libc::tcgetpgrp(self.master.as_raw_fd());
            if foreground > 0 && foreground != libc::getpgrp() {
                libc::kill(-foreground, libc::SIGHUP);
            }
            libc::kill(-group, libc::SIGHUP);
        }
        let deadline = Instant::now() + Duration::from_millis(100);
        while Instant::now() < deadline {
            if self.child.try_wait().ok().flatten().is_some() {
                return;
            }
            std::thread::sleep(Duration::from_millis(5));
        }
        // SAFETY: only our still-unreaped child's process group is targeted.
        unsafe {
            libc::kill(-group, libc::SIGKILL);
        }
        let _ = self.child.kill();
        let _ = self.child.wait();
    }
}
fn winsize(rows: u16, cols: u16) -> libc::winsize {
    libc::winsize {
        ws_row: rows.max(1),
        ws_col: cols.max(1),
        ws_xpixel: 0,
        ws_ypixel: 0,
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    #[test]
    fn child_has_controlling_pty_correct_size_and_native_environment() {
        let super::super::cli::Action::Run(options) = super::super::cli::parse([
            "--", "/bin/sh", "-c", "test -t 0 && test -t 1 && test -t 2 && stty size && printf '%s|%s|%s\\n' \"$TERM\" \"$TERM_PROGRAM\" \"${WAYLAND_DISPLAY-unset}\""
        ].map(Into::into)).unwrap() else { panic!() };
        let mut pty = Pty::spawn(&options, 31, 91).unwrap();
        let deadline = Instant::now() + Duration::from_secs(5);
        let mut output = Vec::new();
        let mut buffer = [0; 1024];
        while !pty.eof {
            match pty.read(&mut buffer) {
                Ok(n) => output.extend_from_slice(&buffer[..n]),
                Err(e) if e.kind() == io::ErrorKind::WouldBlock => {
                    std::thread::sleep(Duration::from_millis(2))
                }
                Err(e) => panic!("{e}"),
            }
            assert!(Instant::now() < deadline);
        }
        let text = String::from_utf8(output).unwrap();
        assert!(text.contains("31 91"), "{text:?}");
        assert!(text.contains("xterm-256color|lios|unset"), "{text:?}");
        pty.resize(20, 60).unwrap();
    }
    #[test]
    fn close_reaps_a_running_child_and_queue_is_bounded() {
        let super::super::cli::Action::Run(options) =
            super::super::cli::parse(["--", "/bin/sh", "-c", "exec sleep 20"].map(Into::into))
                .unwrap()
        else {
            panic!()
        };
        let mut pty = Pty::spawn(&options, 24, 80).unwrap();
        let pid = pty.child.id() as i32;
        assert!(pty.enqueue(&vec![0; MAX_PENDING_INPUT + 1]).is_err());
        drop(pty);
        // SAFETY: signal 0 only inspects whether the private fixture child exists.
        assert_eq!(unsafe { libc::kill(pid, 0) }, -1);
        assert_eq!(io::Error::last_os_error().raw_os_error(), Some(libc::ESRCH));
    }
}