zterm/src/pty.rs

//! PTY (pseudo-terminal) handling for shell communication.

use rustix::fs::{fcntl_setfl, OFlags};
use rustix::io::{read, write, Errno};
use rustix::pty::{grantpt, openpt, ptsname, unlockpt, OpenptFlags};
use std::ffi::CString;
use std::os::fd::{AsFd, AsRawFd, BorrowedFd, OwnedFd, RawFd};
use thiserror::Error;

#[derive(Error, Debug)]
pub enum PtyError {
    #[error("Failed to open PTY master: {0}")]
    OpenMaster(#[source] rustix::io::Errno),
    #[error("Failed to grant PTY: {0}")]
    Grant(#[source] rustix::io::Errno),
    #[error("Failed to unlock PTY: {0}")]
    Unlock(#[source] rustix::io::Errno),
    #[error("Failed to get PTS name: {0}")]
    PtsName(#[source] rustix::io::Errno),
    #[error("Failed to open PTS: {0}")]
    OpenSlave(#[source] rustix::io::Errno),
    #[error("Failed to fork: {0}")]
    Fork(#[source] std::io::Error),
    #[error("Failed to execute shell: {0}")]
    Exec(#[source] std::io::Error),
    #[error("I/O error: {0}")]
    Io(#[source] std::io::Error),
}

/// Represents the master side of a PTY pair.
pub struct Pty {
    master: OwnedFd,
    child_pid: rustix::process::Pid,
}

impl Pty {
    /// Creates a new PTY and spawns a shell process.
    pub fn spawn(shell: Option<&str>) -> Result<Self, PtyError> {
        // Open the PTY master
        let master = openpt(OpenptFlags::RDWR | OpenptFlags::NOCTTY | OpenptFlags::CLOEXEC)
            .map_err(PtyError::OpenMaster)?;

        // Set non-blocking mode on master
        fcntl_setfl(&master, OFlags::NONBLOCK).map_err(|e| PtyError::Io(e.into()))?;

        // Grant and unlock the PTY
        grantpt(&master).map_err(PtyError::Grant)?;
        unlockpt(&master).map_err(PtyError::Unlock)?;

        // Get the slave name
        let slave_name = ptsname(&master, Vec::new()).map_err(PtyError::PtsName)?;

        // Fork the process
        // SAFETY: We're careful to only use async-signal-safe functions in the child
        let fork_result = unsafe { libc::fork() };

        match fork_result {
            -1 => Err(PtyError::Fork(std::io::Error::last_os_error())),
            0 => {
                // Child process
                Self::setup_child(&slave_name, shell);
            }
            pid => {
                // Parent process
                let child_pid = unsafe { rustix::process::Pid::from_raw_unchecked(pid) };
                Ok(Self { master, child_pid })
            }
        }
    }

    /// Sets up the child process (runs in forked child).
    fn setup_child(slave_name: &CString, shell: Option<&str>) -> ! {
        // Create a new session
        unsafe { libc::setsid() };

        // Open the slave PTY using libc for async-signal-safety
        let slave_fd = unsafe { libc::open(slave_name.as_ptr(), libc::O_RDWR) };
        if slave_fd < 0 {
            unsafe { libc::_exit(1) };
        }

        // Set as controlling terminal
        unsafe { libc::ioctl(slave_fd, libc::TIOCSCTTY, 0) };

        // Duplicate slave to stdin/stdout/stderr
        unsafe {
            libc::dup2(slave_fd, 0);
            libc::dup2(slave_fd, 1);
            libc::dup2(slave_fd, 2);
        }

        // Close the original slave fd if it's not 0, 1, or 2
        if slave_fd > 2 {
            unsafe { libc::close(slave_fd) };
        }

        // Set TERM environment variable
        // Try zterm first, fall back to xterm-256color if zterm terminfo isn't installed
        // SAFETY: We're in a forked child process before exec, single-threaded
        unsafe {
            std::env::set_var("TERM", "zterm");
            // Also set COLORTERM to indicate true color support
            std::env::set_var("COLORTERM", "truecolor");
        }

        // Determine which shell to use
        let shell_path = shell
            .map(String::from)
            .or_else(|| std::env::var("SHELL").ok())
            .unwrap_or_else(|| "/bin/sh".to_string());

        let shell_cstr = CString::new(shell_path.clone()).expect("Invalid shell path");
        let shell_name = std::path::Path::new(&shell_path)
            .file_name()
            .and_then(|n| n.to_str())
            .unwrap_or("sh");

        // Login shell (prepend with -)
        let login_shell = CString::new(format!("-{}", shell_name)).expect("Invalid shell name");

        // Execute the shell
        let args = [login_shell.as_ptr(), std::ptr::null()];

        unsafe {
            libc::execvp(shell_cstr.as_ptr(), args.as_ptr());
        }

        // If exec fails, exit
        std::process::exit(1);
    }

    /// Returns a reference to the master file descriptor.
    pub fn master_fd(&self) -> BorrowedFd<'_> {
        self.master.as_fd()
    }

    /// Reads data from the PTY master.
    /// Returns Ok(0) if no data is available (non-blocking).
    pub fn read(&self, buf: &mut [u8]) -> Result<usize, PtyError> {
        match read(&self.master, buf) {
            Ok(n) => Ok(n),
            Err(Errno::AGAIN) => Ok(0), // WOULDBLOCK is same as AGAIN on Linux
            Err(e) => Err(PtyError::Io(e.into())),
        }
    }

    /// Writes data to the PTY master.
    pub fn write(&self, buf: &[u8]) -> Result<usize, PtyError> {
        write(&self.master, buf).map_err(|e| PtyError::Io(e.into()))
    }

    /// Resizes the PTY window.
    pub fn resize(&self, cols: u16, rows: u16, xpixel: u16, ypixel: u16) -> Result<(), PtyError> {
        let winsize = libc::winsize {
            ws_row: rows,
            ws_col: cols,
            ws_xpixel: xpixel,
            ws_ypixel: ypixel,
        };

        let fd = std::os::fd::AsRawFd::as_raw_fd(&self.master);
        let result = unsafe { libc::ioctl(fd, libc::TIOCSWINSZ, &winsize) };

        if result == -1 {
            Err(PtyError::Io(std::io::Error::last_os_error()))
        } else {
            Ok(())
        }
    }

    /// Returns the child process ID.
    pub fn child_pid(&self) -> rustix::process::Pid {
        self.child_pid
    }
    
    /// Check if the child process has exited.
    pub fn child_exited(&self) -> bool {
        let mut status: libc::c_int = 0;
        let result = unsafe {
            libc::waitpid(
                self.child_pid.as_raw_nonzero().get(),
                &mut status,
                libc::WNOHANG,
            )
        };
        // If waitpid returns the child PID, the child has exited
        // If it returns 0, the child is still running
        // If it returns -1, there was an error (child might have already been reaped)
        result != 0
    }
    
    /// Get the foreground process group ID of this PTY.
    /// Returns None if the query fails.
    pub fn foreground_pgid(&self) -> Option<i32> {
        let fd = self.master.as_raw_fd();
        let pgid = unsafe { libc::tcgetpgrp(fd) };
        if pgid > 0 {
            Some(pgid)
        } else {
            None
        }
    }
    
    /// Get the name of the foreground process running in this PTY.
    /// Returns the process name (e.g., "nvim", "zsh") or None if unavailable.
    pub fn foreground_process_name(&self) -> Option<String> {
        let pgid = self.foreground_pgid()?;
        
        // Read the command line from /proc/<pid>/comm
        // (comm gives just the process name, cmdline gives full command)
        let comm_path = format!("/proc/{}/comm", pgid);
        std::fs::read_to_string(&comm_path)
            .ok()
            .map(|s| s.trim().to_string())
    }
}

impl AsRawFd for Pty {
    fn as_raw_fd(&self) -> RawFd {
        self.master.as_raw_fd()
    }
}

impl Drop for Pty {
    fn drop(&mut self) {
        // Send SIGHUP to the child process
        unsafe {
            libc::kill(self.child_pid.as_raw_nonzero().get(), libc::SIGHUP);
        }
    }
}