termtx

lib.rs

//! Lorem ipsum.

#![deny(missing_docs)]

use rustix::fd::{AsFd, AsRawFd, BorrowedFd, IntoFd, IntoRawFd, OwnedFd, RawFd};
use rustix::fs::{cwd, fcntl_dupfd_cloexec, fcntl_getfl, fstat, openat, Mode, OFlags, Stat};
use rustix::io::OwnedFd as RustixFd;
use rustix::io::{poll, read, write, Errno, PollFd, PollFlags};
use rustix::path::Arg;
use rustix::termios::{
    isatty, tcgetattr, tcsetattr, ttyname, OptionalActions, Termios, ECHO, ICANON,
};

use std::io::{self, Read, StderrLock, StdinLock, StdoutLock, Write};
use std::mem::ManuallyDrop;
use std::time::Duration;

/// A type that can be fallibly converted into an owned file descriptor.
trait TryIntoFd: AsFd {
    /// Convert this object into an owned file descriptor.
    fn try_into_fd(self) -> io::Result<RustixFd>;
}

impl TryIntoFd for BorrowedFd<'_> {
    fn try_into_fd(self) -> io::Result<RustixFd> {
        Ok(fcntl_dupfd_cloexec(self, 0)?)
    }
}

impl TryIntoFd for OwnedFd {
    fn try_into_fd(self) -> io::Result<RustixFd> {
        Ok(self.into())
    }
}

impl TryIntoFd for RustixFd {
    fn try_into_fd(self) -> io::Result<RustixFd> {
        Ok(self)
    }
}

/// Transactional read/write operations.
pub trait Transceiver: Sized {
    /// The type used for write operations.
    type Write: Write;
    /// The type used for read operations.
    type Read: Read;

    /// Write a message and read the response.
    fn transceive<W, R, T>(self, write: W, read: R) -> io::Result<T>
    where
        W: FnOnce(&mut Self::Write) -> io::Result<()>,
        R: FnOnce(&mut Self::Read) -> io::Result<T>;

    /// Write a message and read the response, with a timeout.
    fn transceive_timeout<W, R, T>(self, write: W, read: R, timeout: Duration) -> io::Result<T>
    where
        W: FnOnce(&mut Self::Write) -> io::Result<()>,
        R: FnOnce(&mut Self::Read) -> io::Result<T>;

    /// Write a message and return the response.
    fn transceive_buf(self, message: &[u8]) -> io::Result<Vec<u8>> {
        let mut buf = vec![0; 1024];

        let len = self.transceive(
            |w| w.write_all(message),
            |r| r.read(&mut buf),
        )?;

        buf.truncate(len);
        Ok(buf)
    }

    /// Write a message and return the response, with a timeout.
    fn transceive_buf_timeout(self, message: &[u8], timeout: Duration) -> io::Result<Vec<u8>> {
        let mut buf = vec![0; 1024];

        let len = self.transceive_timeout(
            |w| w.write_all(message),
            |r| r.read(&mut buf),
            timeout,
        )?;

        buf.truncate(len);
        Ok(buf)
    }
}

/// Open a terminal.
fn open(path: impl Arg) -> io::Result<RustixFd> {
    Ok(openat(cwd(), path, OFlags::RDWR | OFlags::CLOEXEC, Mode::empty())?)
}

/// Check if two file descriptors refer to the same tty.
fn is_same_tty(stat: &Stat, other: impl AsFd) -> io::Result<bool> {
    let stat2 = fstat(other)?;
    Ok((stat.st_dev, stat.st_ino) == (stat2.st_dev, stat2.st_ino))
}

/// A handle to a [Unix terminal].
///
/// [Unix terminal]: https://pubs.opengroup.org/onlinepubs/9699919799/basedefs/V1_chap11.html#tag_11
pub struct Terminal {
    fd: RustixFd,
    lock_stdin: bool,
    lock_stdout: bool,
    lock_stderr: bool,
}

impl Terminal {
    /// Connect to the terminal attached to the standard input stream.
    pub fn stdin() -> io::Result<Self> {
        Self::new(io::stdin().as_fd())
    }

    /// Connect to the terminal attached to the standard output stream.
    pub fn stdout() -> io::Result<Self> {
        Self::new(io::stdout().as_fd())
    }

    /// Connect to the terminal attached to the standard error stream.
    pub fn stderr() -> io::Result<Self> {
        Self::new(io::stderr().as_fd())
    }

    /// Connect to the terminal attached to one of the standard I/O streams.
    pub fn stdio() -> io::Result<Self> {
        Self::stderr()
            .or_else(|_| Self::stdout())
            .or_else(|_| Self::stdin())
    }

    /// Connect to this process's controlling terminal.
    pub fn controlling() -> io::Result<Self> {
        Self::new(open("/dev/tty")?)
    }

    /// Connect to a terminal that's already open.
    pub fn from_file(file: impl IntoFd) -> io::Result<Self> {
        Self::new(file.into_fd())
    }

    /// Create a new Terminal.
    fn new(file: impl TryIntoFd) -> io::Result<Self> {
        let fd = file.as_fd();
        if !isatty(fd) {
            return Err(Errno::NOTTY.into());
        }

        let mode = fcntl_getfl(fd)? & OFlags::ACCMODE;
        let fd = if mode == OFlags::RDWR {
            file.try_into_fd()?
        } else {
            open(ttyname(fd, vec![])?)?
        };

        let stat = fstat(fd.as_fd())?;
        let lock_stdin = is_same_tty(&stat, io::stdin())?;
        let lock_stdout = is_same_tty(&stat, io::stdout())?;
        let lock_stderr = is_same_tty(&stat, io::stderr())?;

        Ok(Self {
            fd,
            lock_stdin,
            lock_stdout,
            lock_stderr,
        })
    }

    /// Lock access to this terminal.
    ///
    /// For the lifetime of the retured [TerminalLock], this object as well as
    /// any standard I/O streams that refer to the same terminal will be locked.
    pub fn lock(&mut self) -> TerminalLock<'_> {
        let fd = self.fd.as_fd();
        let _stdin = self.lock_stdin.then(|| io::stdin().lock());
        let _stdout = self.lock_stdout.then(|| io::stdout().lock());
        let _stderr = self.lock_stderr.then(|| io::stderr().lock());

        TerminalLock {
            fd,
            _stdin,
            _stdout,
            _stderr,
        }
    }
}

impl AsFd for Terminal {
    fn as_fd(&self) -> BorrowedFd<'_> {
        self.fd.as_fd()
    }
}

impl AsRawFd for Terminal {
    fn as_raw_fd(&self) -> RawFd {
        self.fd.as_raw_fd()
    }
}

impl IntoFd for Terminal {
    fn into_fd(self) -> OwnedFd {
        self.fd.into()
    }
}

impl IntoRawFd for Terminal {
    fn into_raw_fd(self) -> RawFd {
        self.fd.into_raw_fd()
    }
}

impl Read for Terminal {
    fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
        self.lock().read(buf)
    }
}

impl Write for Terminal {
    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
        self.lock().write(buf)
    }

    fn flush(&mut self) -> io::Result<()> {
        self.lock().flush()
    }
}

impl<'a> Transceiver for &'a mut Terminal {
    type Write = TerminalLock<'a>;
    type Read = TerminalLock<'a>;

    fn transceive<W, R, T>(self, write: W, read: R) -> io::Result<T>
    where
        W: FnOnce(&mut Self::Write) -> io::Result<()>,
        R: FnOnce(&mut Self::Read) -> io::Result<T>,
    {
        self.lock().transceive_impl(write, read, None)
    }

    fn transceive_timeout<W, R, T>(self, write: W, read: R, timeout: Duration) -> io::Result<T>
    where
        W: FnOnce(&mut Self::Write) -> io::Result<()>,
        R: FnOnce(&mut Self::Read) -> io::Result<T>,
    {
        self.lock().transceive_impl(write, read, Some(timeout))
    }
}

/// A scope guard that keeps a terminal in "raw mode".
struct RawMode<'a> {
    fd: BorrowedFd<'a>,
    saved: Termios,
}

impl<'a> RawMode<'a> {
    /// Enable raw mode on the given terminal.
    fn enable(fd: BorrowedFd<'a>) -> io::Result<Self> {
        // Turn off echoing (ECHO) and blocking (ICANON)
        let saved = tcgetattr(fd)?;
        let mut new = saved;
        new.c_lflag &= !(ECHO | ICANON);
        if new.c_lflag != saved.c_lflag {
            tcsetattr(fd, OptionalActions::Now, &new)?;
        }

        Ok(Self { fd, saved })
    }

    /// Disable raw mode.  Dropping this guard does the same thing, but without
    /// reporting errors.
    fn disable(self) -> io::Result<()> {
        let mut guard = ManuallyDrop::new(self);
        guard.disable_impl()
    }

    fn disable_impl(&mut self) -> io::Result<()> {
        if self.saved.c_lflag & (ECHO | ICANON) != 0 {
            tcsetattr(self.fd, OptionalActions::Now, &self.saved)?;
        }
        Ok(())
    }
}

impl Drop for RawMode<'_> {
    fn drop(&mut self) {
        let _ = self.disable_impl();
    }
}

/// A locked reference to a [Terminal] handle.
pub struct TerminalLock<'a> {
    fd: BorrowedFd<'a>,
    _stdin: Option<StdinLock<'a>>,
    _stdout: Option<StdoutLock<'a>>,
    _stderr: Option<StderrLock<'a>>,
}

impl TerminalLock<'_> {
    fn transceive_impl<W, R, T>(
        &mut self,
        write: W,
        read: R,
        timeout: Option<Duration>,
    ) -> io::Result<T>
    where
        W: FnOnce(&mut Self) -> io::Result<()>,
        R: FnOnce(&mut Self) -> io::Result<T>,
    {
        let timeout = timeout
            .map(|t| t.as_millis())
            .map(i32::try_from)
            .transpose()
            .map_err(|_| Errno::OVERFLOW)?;

        let raw = RawMode::enable(self.fd)?;

        write(self)?;

        if let Some(timeout) = timeout {
            let mut pollfds = [PollFd::from_borrowed_fd(self.fd, PollFlags::IN)];
            if poll(&mut pollfds, timeout)? < 1 || !pollfds[0].revents().contains(PollFlags::IN) {
                return Err(Errno::TIMEDOUT.into());
            }
        }

        let ret = read(self)?;
        raw.disable()?;
        Ok(ret)
    }
}

impl AsFd for TerminalLock<'_> {
    fn as_fd(&self) -> BorrowedFd<'_> {
        self.fd
    }
}

impl AsRawFd for TerminalLock<'_> {
    fn as_raw_fd(&self) -> RawFd {
        self.fd.as_raw_fd()
    }
}

impl Read for TerminalLock<'_> {
    fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
        Ok(read(self.as_fd(), buf)?)
    }
}

impl Write for TerminalLock<'_> {
    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
        Ok(write(self.as_fd(), buf)?)
    }

    fn flush(&mut self) -> io::Result<()> {
        Ok(())
    }
}

impl<'a> Transceiver for &mut TerminalLock<'a> {
    type Write = TerminalLock<'a>;
    type Read = TerminalLock<'a>;

    fn transceive<W, R, T>(self, write: W, read: R) -> io::Result<T>
    where
        W: FnOnce(&mut Self::Write) -> io::Result<()>,
        R: FnOnce(&mut Self::Read) -> io::Result<T>,
    {
        self.transceive_impl(write, read, None)
    }

    fn transceive_timeout<W, R, T>(self, write: W, read: R, timeout: Duration) -> io::Result<T>
    where
        W: FnOnce(&mut Self::Write) -> io::Result<()>,
        R: FnOnce(&mut Self::Read) -> io::Result<T>,
    {
        self.transceive_impl(write, read, Some(timeout))
    }
}