Skip to content

Instantly share code, notes, and snippets.

@tavianator
Created August 26, 2022 22:27
Show Gist options
  • Save tavianator/d66d425399a57c51629999ae716bbd24 to your computer and use it in GitHub Desktop.
Save tavianator/d66d425399a57c51629999ae716bbd24 to your computer and use it in GitHub Desktop.
termtx
//! 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))
}
}
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment