use std::{
io,
sync::{
Arc,
atomic::{AtomicBool, Ordering},
},
time::Duration,
};
#[derive(Clone, Debug)]
pub struct Waker(Arc<WakeInner>);
#[derive(Debug)]
struct WakeInner {
woken: AtomicBool,
#[cfg(any(unix, windows))]
sys: sys::Wake,
}
impl Waker {
pub fn new() -> io::Result<Self> {
Ok(Self(Arc::new(WakeInner {
woken: AtomicBool::new(false),
#[cfg(any(unix, windows))]
sys: sys::Wake::new()?,
})))
}
pub fn wake(&self) {
self.0.woken.store(true, Ordering::Release);
#[cfg(any(unix, windows))]
self.0.sys.signal();
}
pub(crate) fn take(&self) -> bool {
self.0.woken.swap(false, Ordering::AcqRel)
}
}
#[cfg(any(unix, windows))]
static OWNED: AtomicBool = AtomicBool::new(false);
#[cfg(any(unix, windows))]
#[derive(Debug)]
struct Ownership;
#[cfg(any(unix, windows))]
impl Ownership {
fn acquire() -> io::Result<Self> {
OWNED
.compare_exchange(false, true, Ordering::AcqRel, Ordering::Acquire)
.map(|_| Self)
.map_err(|_| {
io::Error::new(io::ErrorKind::AlreadyExists, "the process tty is already in use")
})
}
}
#[cfg(any(unix, windows))]
impl Drop for Ownership {
fn drop(&mut self) {
OWNED.store(false, Ordering::Release);
}
}
pub trait Terminal {
fn write(&mut self, bytes: &[u8]) -> io::Result<()>;
fn read(&mut self, buf: &mut [u8], timeout: Option<Duration>) -> io::Result<Option<usize>>;
fn restore(&mut self) -> io::Result<()> {
Ok(())
}
fn modes(&mut self, _paste: bool, _kitty: bool) {}
fn waker(&self) -> Option<Waker> {
None
}
}
#[cfg(any(unix, windows))]
pub fn is_tty() -> bool {
sys::is_tty()
}
#[cfg(any(unix, windows))]
pub fn stdout_is_tty() -> bool {
sys::stdout_is_tty()
}
#[cfg(any(unix, windows))]
pub fn columns() -> Option<u16> {
sys::columns()
}
#[cfg(any(unix, windows))]
#[derive(Debug)]
pub struct Tty {
inner: sys::Raw,
restored: bool,
ownership: Option<Ownership>,
waker: Waker,
}
#[cfg(any(unix, windows))]
impl Tty {
pub fn open() -> io::Result<Self> {
let ownership = Ownership::acquire()?;
let waker = Waker::new()?;
Ok(Self { inner: sys::Raw::enter()?, restored: false, ownership: Some(ownership), waker })
}
}
#[cfg(any(unix, windows))]
impl Terminal for Tty {
fn write(&mut self, bytes: &[u8]) -> io::Result<()> {
use io::Write as _;
let mut out = io::stdout().lock();
out.write_all(bytes)?;
out.flush()
}
fn read(&mut self, buf: &mut [u8], timeout: Option<Duration>) -> io::Result<Option<usize>> {
self.inner.read(buf, timeout, &self.waker.0.sys)
}
fn waker(&self) -> Option<Waker> {
Some(self.waker.clone())
}
fn restore(&mut self) -> io::Result<()> {
if self.restored {
return Ok(());
}
self.restored = true;
let result = self.inner.leave();
self.ownership.take();
result
}
fn modes(&mut self, paste: bool, kitty: bool) {
sys::set_modes(paste, kitty);
}
}
#[cfg(any(unix, windows))]
impl Drop for Tty {
fn drop(&mut self) {
if let Err(err) = self.restore() {
tracing::warn!(%err, "could not restore the terminal");
}
}
}
#[cfg(all(test, any(unix, windows)))]
mod tests {
use super::Ownership;
#[test]
fn tty_ownership_is_exclusive_and_released_by_drop() {
let first = Ownership::acquire().expect("first owner");
assert!(Ownership::acquire().is_err());
drop(first);
assert!(Ownership::acquire().is_ok());
}
}
#[cfg(any(windows, test))]
fn utf16_to_utf8(high: &mut Option<u16>, units: &[u16], buf: &mut [u8]) -> usize {
let mut units = high
.take()
.into_iter()
.chain(units.iter().copied())
.peekable();
let mut n = 0;
while let Some(unit) = units.next() {
let ch = if (0xd800..0xdc00).contains(&unit) {
match units.peek().copied() {
None => {
*high = Some(unit);
break;
},
Some(low) if (0xdc00..0xe000).contains(&low) => {
units.next();
char::decode_utf16([unit, low])
.next()
.and_then(Result::ok)
.unwrap_or(char::REPLACEMENT_CHARACTER)
},
Some(_) => char::REPLACEMENT_CHARACTER,
}
} else {
char::from_u32(u32::from(unit)).unwrap_or(char::REPLACEMENT_CHARACTER)
};
n += ch.encode_utf8(&mut buf[n..]).len();
}
n
}
#[cfg(test)]
mod utf16_tests {
use super::utf16_to_utf8;
fn read(pieces: &[&[u16]]) -> String {
let mut high = None;
let mut out = Vec::new();
for units in pieces {
let mut buf = [0u8; 64];
let n = utf16_to_utf8(&mut high, units, &mut buf);
out.extend_from_slice(&buf[..n]);
}
String::from_utf8(out).expect("UTF-8")
}
#[test]
fn console_text_reaches_the_program_as_utf8() {
let units = |s: &str| s.encode_utf16().collect::<Vec<_>>();
assert_eq!(read(&[&units("ä¸")]), "ä¸");
assert_eq!(read(&[&units("\x1b[Aa䏿–‡")]), "\x1b[Aa䏿–‡");
let emoji = units("😀");
assert_eq!(read(&[&emoji[..1], &emoji[1..]]), "😀");
assert_eq!(read(&[&[0xdc00], &[0xd800, 0x61]]), "\u{fffd}\u{fffd}a");
}
}
pub(crate) const fn mode_reset(paste: bool, kitty: bool) -> &'static [u8] {
match (paste, kitty) {
(true, true) => b"\x1b[<u\x1b[?2004l",
(false, true) => b"\x1b[<u",
(true, false) => b"\x1b[?2004l",
(false, false) => b"",
}
}
#[cfg(unix)]
mod sys {
use std::{
cell::UnsafeCell,
io,
mem::MaybeUninit,
sync::atomic::{AtomicBool, AtomicU8, Ordering},
time::Duration,
};
struct Saved(UnsafeCell<MaybeUninit<libc::termios>>);
unsafe impl Sync for Saved {}
static SAVED: Saved = Saved(UnsafeCell::new(MaybeUninit::uninit()));
static ARMED: AtomicBool = AtomicBool::new(false);
static MODES: AtomicU8 = AtomicU8::new(0);
const SIGNALS: [libc::c_int; 4] = [libc::SIGTERM, libc::SIGHUP, libc::SIGINT, libc::SIGQUIT];
#[derive(Debug)]
pub(super) struct Wake {
read: libc::c_int,
write: libc::c_int,
}
impl Wake {
pub(super) fn new() -> io::Result<Self> {
let mut fds = [0; 2];
if unsafe { libc::pipe(fds.as_mut_ptr()) } != 0 {
return Err(io::Error::last_os_error());
}
let wake = Self { read: fds[0], write: fds[1] };
for fd in fds {
let ok = unsafe {
let flags = libc::fcntl(fd, libc::F_GETFL);
flags >= 0
&& libc::fcntl(fd, libc::F_SETFL, flags | libc::O_NONBLOCK) == 0
&& libc::fcntl(fd, libc::F_SETFD, libc::FD_CLOEXEC) == 0
};
if !ok {
return Err(io::Error::last_os_error());
}
}
Ok(wake)
}
pub(super) fn signal(&self) {
unsafe { libc::write(self.write, [1u8].as_ptr().cast(), 1) };
}
fn drain(&self) {
let mut buf = [0u8; 64];
while unsafe { libc::read(self.read, buf.as_mut_ptr().cast(), buf.len()) } > 0 {}
}
}
impl Drop for Wake {
fn drop(&mut self) {
unsafe {
libc::close(self.read);
libc::close(self.write);
}
}
}
#[derive(Debug)]
pub(super) struct Raw {
saved: libc::termios,
previous: Vec<(libc::c_int, libc::sigaction)>,
}
pub(super) fn is_tty() -> bool {
unsafe { libc::isatty(0) == 1 && libc::isatty(1) == 1 }
}
pub(super) fn stdout_is_tty() -> bool {
unsafe { libc::isatty(1) == 1 }
}
pub(super) fn columns() -> Option<u16> {
let mut size = MaybeUninit::<libc::winsize>::zeroed();
let ok = unsafe { libc::ioctl(1, libc::TIOCGWINSZ, size.as_mut_ptr()) } == 0;
let size = unsafe { size.assume_init() };
(ok && size.ws_col > 0).then_some(size.ws_col)
}
pub(super) fn set_modes(paste: bool, kitty: bool) {
MODES.store(u8::from(paste) | (u8::from(kitty) << 1), Ordering::Release);
}
extern "C" fn on_signal(sig: libc::c_int) {
if ARMED.load(Ordering::Acquire) {
let modes = MODES.load(Ordering::Acquire);
let reset = super::mode_reset(modes & 1 != 0, modes & 2 != 0);
unsafe {
libc::write(1, reset.as_ptr().cast(), reset.len());
libc::tcsetattr(0, libc::TCSANOW, (*SAVED.0.get()).as_ptr());
}
}
unsafe {
libc::signal(sig, libc::SIG_DFL);
libc::raise(sig);
}
}
impl Raw {
pub(super) fn enter() -> io::Result<Self> {
let mut termios = MaybeUninit::<libc::termios>::uninit();
if unsafe { libc::tcgetattr(0, termios.as_mut_ptr()) } != 0 {
return Err(io::Error::last_os_error());
}
let saved = unsafe { termios.assume_init() };
let mut raw = saved;
raw.c_iflag &= !(libc::IGNBRK
| libc::BRKINT
| libc::PARMRK
| libc::ISTRIP
| libc::INLCR
| libc::IGNCR
| libc::ICRNL
| libc::IXON);
raw.c_lflag &= !(libc::ECHO | libc::ECHONL | libc::ICANON | libc::ISIG | libc::IEXTEN);
raw.c_cflag &= !(libc::CSIZE | libc::PARENB);
raw.c_cflag |= libc::CS8;
raw.c_cc[libc::VMIN] = 1;
raw.c_cc[libc::VTIME] = 0;
unsafe { (*SAVED.0.get()).write(saved) };
ARMED.store(true, Ordering::Release);
let mut previous = Vec::with_capacity(SIGNALS.len());
for sig in SIGNALS {
unsafe {
let mut action: libc::sigaction = std::mem::zeroed();
action.sa_sigaction = on_signal as extern "C" fn(libc::c_int) as usize;
libc::sigemptyset(&mut action.sa_mask);
let mut old: libc::sigaction = std::mem::zeroed();
if libc::sigaction(sig, &action, &mut old) != 0 {
let err = io::Error::last_os_error();
let mut state = Self { saved, previous };
let _ = state.leave();
return Err(err);
}
previous.push((sig, old));
}
}
let mut state = Self { saved, previous };
if unsafe { libc::tcsetattr(0, libc::TCSANOW, &raw) } != 0 {
let err = io::Error::last_os_error();
let _ = state.leave();
return Err(err);
}
Ok(state)
}
#[allow(clippy::unused_self, reason = "the Windows version reads through its handle")]
pub(super) fn read(
&self,
buf: &mut [u8],
timeout: Option<Duration>,
wake: &Wake,
) -> io::Result<Option<usize>> {
let ms = timeout
.map_or(-1, |t| libc::c_int::try_from(t.as_millis()).unwrap_or(libc::c_int::MAX));
let mut fds = [libc::pollfd { fd: 0, events: libc::POLLIN, revents: 0 }, libc::pollfd {
fd: wake.read,
events: libc::POLLIN,
revents: 0,
}];
let n = unsafe { libc::poll(fds.as_mut_ptr(), 2, ms) };
if n < 0 {
let err = io::Error::last_os_error();
return if err.kind() == io::ErrorKind::Interrupted {
Ok(None)
} else {
Err(err)
};
}
if fds[1].revents != 0 {
wake.drain();
}
if n == 0 || fds[0].revents == 0 {
return Ok(None);
}
let got = unsafe { libc::read(0, buf.as_mut_ptr().cast(), buf.len()) };
if got < 0 {
let err = io::Error::last_os_error();
return match err.kind() {
io::ErrorKind::Interrupted | io::ErrorKind::WouldBlock => Ok(None),
_ => Err(err),
};
}
Ok(Some(got as usize))
}
pub(super) fn leave(&mut self) -> io::Result<()> {
let result = if unsafe { libc::tcsetattr(0, libc::TCSANOW, &self.saved) } == 0 {
Ok(())
} else {
Err(io::Error::last_os_error())
};
ARMED.store(false, Ordering::Release);
MODES.store(0, Ordering::Release);
for (sig, old) in self.previous.drain(..) {
unsafe { libc::sigaction(sig, &old, std::ptr::null_mut()) };
}
result
}
}
}
#[cfg(windows)]
mod sys {
use std::{io, time::Duration};
use windows_sys::Win32::{
Foundation::{CloseHandle, FALSE, HANDLE, WAIT_OBJECT_0},
System::{
Console::{
CONSOLE_SCREEN_BUFFER_INFO, ENABLE_ECHO_INPUT, ENABLE_LINE_INPUT,
ENABLE_PROCESSED_INPUT, ENABLE_PROCESSED_OUTPUT, ENABLE_VIRTUAL_TERMINAL_INPUT,
ENABLE_VIRTUAL_TERMINAL_PROCESSING, GetConsoleMode, GetConsoleScreenBufferInfo,
GetStdHandle, INPUT_RECORD, KEY_EVENT, PeekConsoleInputW, ReadConsoleInputW,
ReadConsoleW, STD_INPUT_HANDLE, STD_OUTPUT_HANDLE, SetConsoleMode,
},
Threading::{CreateEventW, INFINITE, SetEvent, WaitForMultipleObjects},
},
};
#[derive(Debug)]
pub(super) struct Wake {
event: HANDLE,
}
unsafe impl Send for Wake {}
unsafe impl Sync for Wake {}
impl Wake {
pub(super) fn new() -> io::Result<Self> {
let event = unsafe { CreateEventW(std::ptr::null(), FALSE, FALSE, std::ptr::null()) };
if event.is_null() {
return Err(io::Error::last_os_error());
}
Ok(Self { event })
}
pub(super) fn signal(&self) {
unsafe { SetEvent(self.event) };
}
}
impl Drop for Wake {
fn drop(&mut self) {
unsafe { CloseHandle(self.event) };
}
}
const WIDE: usize = 2048;
#[derive(Debug)]
pub(super) struct Raw {
input: HANDLE,
output: HANDLE,
in_mode: u32,
out_mode: u32,
high: Option<u16>,
}
fn mode(handle: HANDLE) -> Option<u32> {
let mut mode = 0;
(unsafe { GetConsoleMode(handle, &mut mode) } != 0).then_some(mode)
}
pub(super) fn is_tty() -> bool {
let (input, output) =
unsafe { (GetStdHandle(STD_INPUT_HANDLE), GetStdHandle(STD_OUTPUT_HANDLE)) };
mode(input).is_some() && mode(output).is_some()
}
pub(super) fn stdout_is_tty() -> bool {
mode(unsafe { GetStdHandle(STD_OUTPUT_HANDLE) }).is_some()
}
pub(super) fn columns() -> Option<u16> {
unsafe {
let mut info: CONSOLE_SCREEN_BUFFER_INFO = std::mem::zeroed();
if GetConsoleScreenBufferInfo(GetStdHandle(STD_OUTPUT_HANDLE), &mut info) == 0 {
return None;
}
u16::try_from(info.srWindow.Right - info.srWindow.Left + 1).ok()
}
}
pub(super) const fn set_modes(_paste: bool, _kitty: bool) {}
impl Raw {
pub(super) fn enter() -> io::Result<Self> {
let (input, output) =
unsafe { (GetStdHandle(STD_INPUT_HANDLE), GetStdHandle(STD_OUTPUT_HANDLE)) };
let in_mode = mode(input).ok_or_else(io::Error::last_os_error)?;
let out_mode = mode(output).ok_or_else(io::Error::last_os_error)?;
let raw_in = (in_mode & !(ENABLE_LINE_INPUT | ENABLE_ECHO_INPUT | ENABLE_PROCESSED_INPUT))
| ENABLE_VIRTUAL_TERMINAL_INPUT;
let raw_out = out_mode | ENABLE_PROCESSED_OUTPUT | ENABLE_VIRTUAL_TERMINAL_PROCESSING;
unsafe {
if SetConsoleMode(input, raw_in) == 0 {
return Err(io::Error::last_os_error());
}
if SetConsoleMode(output, raw_out) == 0 {
let err = io::Error::last_os_error();
SetConsoleMode(input, in_mode);
return Err(err);
}
}
Ok(Self { input, output, in_mode, out_mode, high: None })
}
fn key_ready(&self) -> io::Result<bool> {
loop {
let mut record = INPUT_RECORD::default();
let mut n = 0;
if unsafe { PeekConsoleInputW(self.input, &mut record, 1, &mut n) } == 0 {
return Err(io::Error::last_os_error());
}
if n == 0 {
return Ok(false);
}
let key_down = u32::from(record.EventType) == KEY_EVENT
&& unsafe {
let key = record.Event.KeyEvent;
key.bKeyDown != 0
&& (key.uChar.UnicodeChar != 0
|| matches!(key.wVirtualKeyCode,
0x08 | 0x09 | 0x0d | 0x1b | 0x20..=0x28 | 0x2d | 0x2e | 0x70..=0x87))
};
if key_down {
return Ok(true);
}
if unsafe { ReadConsoleInputW(self.input, &mut record, 1, &mut n) } == 0 {
return Err(io::Error::last_os_error());
}
}
}
pub(super) fn read(
&mut self,
buf: &mut [u8],
timeout: Option<Duration>,
wake: &Wake,
) -> io::Result<Option<usize>> {
let ms =
timeout.map_or(INFINITE, |t| u32::try_from(t.as_millis()).unwrap_or(INFINITE - 1));
let handles = [self.input, wake.event];
if unsafe { WaitForMultipleObjects(2, handles.as_ptr(), FALSE, ms) } != WAIT_OBJECT_0 {
return Ok(None);
}
if !self.key_ready()? {
return Ok(None);
}
let room = (buf.len() / 3).saturating_sub(1).clamp(1, WIDE);
let mut wide = [0u16; WIDE];
let mut got = 0;
if unsafe {
ReadConsoleW(
self.input,
wide.as_mut_ptr().cast(),
u32::try_from(room).unwrap_or(u32::MAX),
&mut got,
std::ptr::null(),
)
} == 0
{
return Err(io::Error::last_os_error());
}
let n = super::utf16_to_utf8(&mut self.high, &wide[..(got as usize).min(room)], buf);
Ok((n > 0).then_some(n))
}
pub(super) fn leave(&mut self) -> io::Result<()> {
unsafe {
let input = SetConsoleMode(self.input, self.in_mode);
let input_error = (input == 0).then(io::Error::last_os_error);
let output = SetConsoleMode(self.output, self.out_mode);
if let Some(err) = input_error {
return Err(err);
}
if output == 0 {
return Err(io::Error::last_os_error());
}
}
Ok(())
}
}
}