1use std::{
10 io,
11 sync::{
12 Arc,
13 atomic::{AtomicBool, Ordering},
14 },
15 time::Duration,
16};
17
18#[derive(Clone, Debug)]
21pub struct Waker(Arc<WakeInner>);
22
23#[derive(Debug)]
25struct WakeInner {
26 woken: AtomicBool,
28 #[cfg(any(unix, windows))]
30 sys: sys::Wake,
31}
32
33impl Waker {
34 pub fn new() -> io::Result<Self> {
39 Ok(Self(Arc::new(WakeInner {
40 woken: AtomicBool::new(false),
41 #[cfg(any(unix, windows))]
42 sys: sys::Wake::new()?,
43 })))
44 }
45
46 pub fn wake(&self) {
48 self.0.woken.store(true, Ordering::Release);
49 #[cfg(any(unix, windows))]
50 self.0.sys.signal();
51 }
52
53 pub(crate) fn take(&self) -> bool {
55 self.0.woken.swap(false, Ordering::AcqRel)
56 }
57}
58
59#[cfg(any(unix, windows))]
61static OWNED: AtomicBool = AtomicBool::new(false);
62
63#[cfg(any(unix, windows))]
65#[derive(Debug)]
66struct Ownership;
67
68#[cfg(any(unix, windows))]
69impl Ownership {
70 fn acquire() -> io::Result<Self> {
71 OWNED
72 .compare_exchange(false, true, Ordering::AcqRel, Ordering::Acquire)
73 .map(|_| Self)
74 .map_err(|_| {
75 io::Error::new(io::ErrorKind::AlreadyExists, "the process tty is already in use")
76 })
77 }
78}
79
80#[cfg(any(unix, windows))]
81impl Drop for Ownership {
82 fn drop(&mut self) {
83 OWNED.store(false, Ordering::Release);
84 }
85}
86
87pub trait Terminal {
89 fn write(&mut self, bytes: &[u8]) -> io::Result<()>;
94
95 fn read(&mut self, buf: &mut [u8], timeout: Option<Duration>) -> io::Result<Option<usize>>;
102
103 fn restore(&mut self) -> io::Result<()> {
109 Ok(())
110 }
111
112 fn modes(&mut self, _paste: bool, _kitty: bool) {}
115
116 fn waker(&self) -> Option<Waker> {
119 None
120 }
121}
122
123#[cfg(any(unix, windows))]
125pub fn is_tty() -> bool {
126 sys::is_tty()
127}
128
129#[cfg(any(unix, windows))]
131pub fn stdout_is_tty() -> bool {
132 sys::stdout_is_tty()
133}
134
135#[cfg(any(unix, windows))]
137pub fn columns() -> Option<u16> {
138 sys::columns()
139}
140
141#[cfg(any(unix, windows))]
144#[derive(Debug)]
145pub struct Tty {
146 inner: sys::Raw,
148 restored: bool,
150 ownership: Option<Ownership>,
152 waker: Waker,
154}
155
156#[cfg(any(unix, windows))]
157impl Tty {
158 pub fn open() -> io::Result<Self> {
164 let ownership = Ownership::acquire()?;
165 let waker = Waker::new()?;
166 Ok(Self { inner: sys::Raw::enter()?, restored: false, ownership: Some(ownership), waker })
167 }
168}
169
170#[cfg(any(unix, windows))]
171impl Terminal for Tty {
172 fn write(&mut self, bytes: &[u8]) -> io::Result<()> {
173 use io::Write as _;
174 let mut out = io::stdout().lock();
175 out.write_all(bytes)?;
176 out.flush()
177 }
178
179 fn read(&mut self, buf: &mut [u8], timeout: Option<Duration>) -> io::Result<Option<usize>> {
180 self.inner.read(buf, timeout, &self.waker.0.sys)
181 }
182
183 fn waker(&self) -> Option<Waker> {
184 Some(self.waker.clone())
185 }
186
187 fn restore(&mut self) -> io::Result<()> {
188 if self.restored {
189 return Ok(());
190 }
191 self.restored = true;
192 let result = self.inner.leave();
193 self.ownership.take();
194 result
195 }
196
197 fn modes(&mut self, paste: bool, kitty: bool) {
198 sys::set_modes(paste, kitty);
199 }
200}
201
202#[cfg(any(unix, windows))]
203impl Drop for Tty {
204 fn drop(&mut self) {
205 if let Err(err) = self.restore() {
206 tracing::warn!(%err, "could not restore the terminal");
207 }
208 }
209}
210
211#[cfg(all(test, any(unix, windows)))]
212mod tests {
213 use super::Ownership;
214
215 #[test]
216 fn tty_ownership_is_exclusive_and_released_by_drop() {
217 let first = Ownership::acquire().expect("first owner");
218 assert!(Ownership::acquire().is_err());
219 drop(first);
220 assert!(Ownership::acquire().is_ok());
221 }
222}
223
224#[cfg(any(windows, test))]
229fn utf16_to_utf8(high: &mut Option<u16>, units: &[u16], buf: &mut [u8]) -> usize {
230 let mut units = high
231 .take()
232 .into_iter()
233 .chain(units.iter().copied())
234 .peekable();
235 let mut n = 0;
236 while let Some(unit) = units.next() {
237 let ch = if (0xd800..0xdc00).contains(&unit) {
238 match units.peek().copied() {
239 None => {
240 *high = Some(unit);
241 break;
242 },
243 Some(low) if (0xdc00..0xe000).contains(&low) => {
244 units.next();
245 char::decode_utf16([unit, low])
246 .next()
247 .and_then(Result::ok)
248 .unwrap_or(char::REPLACEMENT_CHARACTER)
249 },
250 Some(_) => char::REPLACEMENT_CHARACTER,
251 }
252 } else {
253 char::from_u32(u32::from(unit)).unwrap_or(char::REPLACEMENT_CHARACTER)
254 };
255 n += ch.encode_utf8(&mut buf[n..]).len();
256 }
257 n
258}
259
260#[cfg(test)]
261mod utf16_tests {
262 use super::utf16_to_utf8;
263
264 fn read(pieces: &[&[u16]]) -> String {
266 let mut high = None;
267 let mut out = Vec::new();
268 for units in pieces {
269 let mut buf = [0u8; 64];
270 let n = utf16_to_utf8(&mut high, units, &mut buf);
271 out.extend_from_slice(&buf[..n]);
272 }
273 String::from_utf8(out).expect("UTF-8")
274 }
275
276 #[test]
277 fn console_text_reaches_the_program_as_utf8() {
278 let units = |s: &str| s.encode_utf16().collect::<Vec<_>>();
279 assert_eq!(read(&[&units("中")]), "中");
281 assert_eq!(read(&[&units("\x1b[Aa中文")]), "\x1b[Aa中文");
282 let emoji = units("😀");
284 assert_eq!(read(&[&emoji[..1], &emoji[1..]]), "😀");
285 assert_eq!(read(&[&[0xdc00], &[0xd800, 0x61]]), "\u{fffd}\u{fffd}a");
286 }
287}
288
289pub(crate) const fn mode_reset(paste: bool, kitty: bool) -> &'static [u8] {
291 match (paste, kitty) {
292 (true, true) => b"\x1b[<u\x1b[?2004l",
293 (false, true) => b"\x1b[<u",
294 (true, false) => b"\x1b[?2004l",
295 (false, false) => b"",
296 }
297}
298
299#[cfg(unix)]
300mod sys {
301 use std::{
302 cell::UnsafeCell,
303 io,
304 mem::MaybeUninit,
305 sync::atomic::{AtomicBool, AtomicU8, Ordering},
306 time::Duration,
307 };
308
309 struct Saved(UnsafeCell<MaybeUninit<libc::termios>>);
311
312 unsafe impl Sync for Saved {}
315
316 static SAVED: Saved = Saved(UnsafeCell::new(MaybeUninit::uninit()));
318 static ARMED: AtomicBool = AtomicBool::new(false);
320 static MODES: AtomicU8 = AtomicU8::new(0);
322
323 const SIGNALS: [libc::c_int; 4] = [libc::SIGTERM, libc::SIGHUP, libc::SIGINT, libc::SIGQUIT];
325
326 #[derive(Debug)]
328 pub(super) struct Wake {
329 read: libc::c_int,
331 write: libc::c_int,
333 }
334
335 impl Wake {
336 pub(super) fn new() -> io::Result<Self> {
337 let mut fds = [0; 2];
338 if unsafe { libc::pipe(fds.as_mut_ptr()) } != 0 {
340 return Err(io::Error::last_os_error());
341 }
342 let wake = Self { read: fds[0], write: fds[1] };
343 for fd in fds {
344 let ok = unsafe {
346 let flags = libc::fcntl(fd, libc::F_GETFL);
347 flags >= 0
348 && libc::fcntl(fd, libc::F_SETFL, flags | libc::O_NONBLOCK) == 0
349 && libc::fcntl(fd, libc::F_SETFD, libc::FD_CLOEXEC) == 0
350 };
351 if !ok {
352 return Err(io::Error::last_os_error());
353 }
354 }
355 Ok(wake)
356 }
357
358 pub(super) fn signal(&self) {
360 unsafe { libc::write(self.write, [1u8].as_ptr().cast(), 1) };
362 }
363
364 fn drain(&self) {
366 let mut buf = [0u8; 64];
367 while unsafe { libc::read(self.read, buf.as_mut_ptr().cast(), buf.len()) } > 0 {}
370 }
371 }
372
373 impl Drop for Wake {
374 fn drop(&mut self) {
375 unsafe {
377 libc::close(self.read);
378 libc::close(self.write);
379 }
380 }
381 }
382
383 #[derive(Debug)]
385 pub(super) struct Raw {
386 saved: libc::termios,
388 previous: Vec<(libc::c_int, libc::sigaction)>,
390 }
391
392 pub(super) fn is_tty() -> bool {
393 unsafe { libc::isatty(0) == 1 && libc::isatty(1) == 1 }
395 }
396
397 pub(super) fn stdout_is_tty() -> bool {
398 unsafe { libc::isatty(1) == 1 }
400 }
401
402 pub(super) fn columns() -> Option<u16> {
403 let mut size = MaybeUninit::<libc::winsize>::zeroed();
404 let ok = unsafe { libc::ioctl(1, libc::TIOCGWINSZ, size.as_mut_ptr()) } == 0;
406 let size = unsafe { size.assume_init() };
408 (ok && size.ws_col > 0).then_some(size.ws_col)
409 }
410
411 pub(super) fn set_modes(paste: bool, kitty: bool) {
412 MODES.store(u8::from(paste) | (u8::from(kitty) << 1), Ordering::Release);
413 }
414
415 extern "C" fn on_signal(sig: libc::c_int) {
417 if ARMED.load(Ordering::Acquire) {
418 let modes = MODES.load(Ordering::Acquire);
419 let reset = super::mode_reset(modes & 1 != 0, modes & 2 != 0);
420 unsafe {
423 libc::write(1, reset.as_ptr().cast(), reset.len());
424 libc::tcsetattr(0, libc::TCSANOW, (*SAVED.0.get()).as_ptr());
425 }
426 }
427 unsafe {
429 libc::signal(sig, libc::SIG_DFL);
430 libc::raise(sig);
431 }
432 }
433
434 impl Raw {
435 pub(super) fn enter() -> io::Result<Self> {
436 let mut termios = MaybeUninit::<libc::termios>::uninit();
437 if unsafe { libc::tcgetattr(0, termios.as_mut_ptr()) } != 0 {
439 return Err(io::Error::last_os_error());
440 }
441 let saved = unsafe { termios.assume_init() };
443 let mut raw = saved;
444 raw.c_iflag &= !(libc::IGNBRK
445 | libc::BRKINT
446 | libc::PARMRK
447 | libc::ISTRIP
448 | libc::INLCR
449 | libc::IGNCR
450 | libc::ICRNL
451 | libc::IXON);
452 raw.c_lflag &= !(libc::ECHO | libc::ECHONL | libc::ICANON | libc::ISIG | libc::IEXTEN);
453 raw.c_cflag &= !(libc::CSIZE | libc::PARENB);
454 raw.c_cflag |= libc::CS8;
455 raw.c_cc[libc::VMIN] = 1;
456 raw.c_cc[libc::VTIME] = 0;
457 unsafe { (*SAVED.0.get()).write(saved) };
459 ARMED.store(true, Ordering::Release);
460 let mut previous = Vec::with_capacity(SIGNALS.len());
461 for sig in SIGNALS {
462 unsafe {
464 let mut action: libc::sigaction = std::mem::zeroed();
465 action.sa_sigaction = on_signal as extern "C" fn(libc::c_int) as usize;
466 libc::sigemptyset(&mut action.sa_mask);
467 let mut old: libc::sigaction = std::mem::zeroed();
468 if libc::sigaction(sig, &action, &mut old) != 0 {
469 let err = io::Error::last_os_error();
470 let mut state = Self { saved, previous };
471 let _ = state.leave();
472 return Err(err);
473 }
474 previous.push((sig, old));
475 }
476 }
477 let mut state = Self { saved, previous };
478 if unsafe { libc::tcsetattr(0, libc::TCSANOW, &raw) } != 0 {
480 let err = io::Error::last_os_error();
481 let _ = state.leave();
482 return Err(err);
483 }
484 Ok(state)
485 }
486
487 #[allow(clippy::unused_self, reason = "the Windows version reads through its handle")]
488 pub(super) fn read(
489 &self,
490 buf: &mut [u8],
491 timeout: Option<Duration>,
492 wake: &Wake,
493 ) -> io::Result<Option<usize>> {
494 let ms = timeout
495 .map_or(-1, |t| libc::c_int::try_from(t.as_millis()).unwrap_or(libc::c_int::MAX));
496 let mut fds = [libc::pollfd { fd: 0, events: libc::POLLIN, revents: 0 }, libc::pollfd {
497 fd: wake.read,
498 events: libc::POLLIN,
499 revents: 0,
500 }];
501 let n = unsafe { libc::poll(fds.as_mut_ptr(), 2, ms) };
503 if n < 0 {
504 let err = io::Error::last_os_error();
505 return if err.kind() == io::ErrorKind::Interrupted {
506 Ok(None)
507 } else {
508 Err(err)
509 };
510 }
511 if fds[1].revents != 0 {
512 wake.drain();
513 }
514 if n == 0 || fds[0].revents == 0 {
515 return Ok(None);
516 }
517 let got = unsafe { libc::read(0, buf.as_mut_ptr().cast(), buf.len()) };
519 if got < 0 {
520 let err = io::Error::last_os_error();
521 return match err.kind() {
522 io::ErrorKind::Interrupted | io::ErrorKind::WouldBlock => Ok(None),
523 _ => Err(err),
524 };
525 }
526 Ok(Some(got as usize))
527 }
528
529 pub(super) fn leave(&mut self) -> io::Result<()> {
530 let result = if unsafe { libc::tcsetattr(0, libc::TCSANOW, &self.saved) } == 0 {
532 Ok(())
533 } else {
534 Err(io::Error::last_os_error())
535 };
536 ARMED.store(false, Ordering::Release);
537 MODES.store(0, Ordering::Release);
538 for (sig, old) in self.previous.drain(..) {
539 unsafe { libc::sigaction(sig, &old, std::ptr::null_mut()) };
541 }
542 result
543 }
544 }
545}
546
547#[cfg(windows)]
548mod sys {
549 use std::{io, time::Duration};
550
551 use windows_sys::Win32::{
552 Foundation::{CloseHandle, FALSE, HANDLE, WAIT_OBJECT_0},
553 System::{
554 Console::{
555 CONSOLE_SCREEN_BUFFER_INFO, ENABLE_ECHO_INPUT, ENABLE_LINE_INPUT,
556 ENABLE_PROCESSED_INPUT, ENABLE_PROCESSED_OUTPUT, ENABLE_VIRTUAL_TERMINAL_INPUT,
557 ENABLE_VIRTUAL_TERMINAL_PROCESSING, GetConsoleMode, GetConsoleScreenBufferInfo,
558 GetStdHandle, INPUT_RECORD, KEY_EVENT, PeekConsoleInputW, ReadConsoleInputW,
559 ReadConsoleW, STD_INPUT_HANDLE, STD_OUTPUT_HANDLE, SetConsoleMode,
560 },
561 Threading::{CreateEventW, INFINITE, SetEvent, WaitForMultipleObjects},
562 },
563 };
564
565 #[derive(Debug)]
567 pub(super) struct Wake {
568 event: HANDLE,
570 }
571
572 unsafe impl Send for Wake {}
574 unsafe impl Sync for Wake {}
576
577 impl Wake {
578 pub(super) fn new() -> io::Result<Self> {
579 let event = unsafe { CreateEventW(std::ptr::null(), FALSE, FALSE, std::ptr::null()) };
581 if event.is_null() {
582 return Err(io::Error::last_os_error());
583 }
584 Ok(Self { event })
585 }
586
587 pub(super) fn signal(&self) {
589 unsafe { SetEvent(self.event) };
591 }
592 }
593
594 impl Drop for Wake {
595 fn drop(&mut self) {
596 unsafe { CloseHandle(self.event) };
598 }
599 }
600
601 const WIDE: usize = 2048;
603
604 #[derive(Debug)]
606 pub(super) struct Raw {
607 input: HANDLE,
609 output: HANDLE,
611 in_mode: u32,
613 out_mode: u32,
615 high: Option<u16>,
617 }
618
619 fn mode(handle: HANDLE) -> Option<u32> {
621 let mut mode = 0;
622 (unsafe { GetConsoleMode(handle, &mut mode) } != 0).then_some(mode)
624 }
625
626 pub(super) fn is_tty() -> bool {
627 let (input, output) =
629 unsafe { (GetStdHandle(STD_INPUT_HANDLE), GetStdHandle(STD_OUTPUT_HANDLE)) };
630 mode(input).is_some() && mode(output).is_some()
631 }
632
633 pub(super) fn stdout_is_tty() -> bool {
634 mode(unsafe { GetStdHandle(STD_OUTPUT_HANDLE) }).is_some()
636 }
637
638 pub(super) fn columns() -> Option<u16> {
639 unsafe {
641 let mut info: CONSOLE_SCREEN_BUFFER_INFO = std::mem::zeroed();
642 if GetConsoleScreenBufferInfo(GetStdHandle(STD_OUTPUT_HANDLE), &mut info) == 0 {
643 return None;
644 }
645 u16::try_from(info.srWindow.Right - info.srWindow.Left + 1).ok()
646 }
647 }
648
649 pub(super) const fn set_modes(_paste: bool, _kitty: bool) {}
650
651 impl Raw {
652 pub(super) fn enter() -> io::Result<Self> {
653 let (input, output) =
655 unsafe { (GetStdHandle(STD_INPUT_HANDLE), GetStdHandle(STD_OUTPUT_HANDLE)) };
656 let in_mode = mode(input).ok_or_else(io::Error::last_os_error)?;
657 let out_mode = mode(output).ok_or_else(io::Error::last_os_error)?;
658 let raw_in = (in_mode & !(ENABLE_LINE_INPUT | ENABLE_ECHO_INPUT | ENABLE_PROCESSED_INPUT))
659 | ENABLE_VIRTUAL_TERMINAL_INPUT;
660 let raw_out = out_mode | ENABLE_PROCESSED_OUTPUT | ENABLE_VIRTUAL_TERMINAL_PROCESSING;
661 unsafe {
663 if SetConsoleMode(input, raw_in) == 0 {
664 return Err(io::Error::last_os_error());
665 }
666 if SetConsoleMode(output, raw_out) == 0 {
667 let err = io::Error::last_os_error();
668 SetConsoleMode(input, in_mode);
669 return Err(err);
670 }
671 }
672 Ok(Self { input, output, in_mode, out_mode, high: None })
673 }
674
675 fn key_ready(&self) -> io::Result<bool> {
678 loop {
679 let mut record = INPUT_RECORD::default();
680 let mut n = 0;
681 if unsafe { PeekConsoleInputW(self.input, &mut record, 1, &mut n) } == 0 {
683 return Err(io::Error::last_os_error());
684 }
685 if n == 0 {
686 return Ok(false);
687 }
688 let key_down = u32::from(record.EventType) == KEY_EVENT
690 && unsafe {
691 let key = record.Event.KeyEvent;
692 key.bKeyDown != 0
693 && (key.uChar.UnicodeChar != 0
694 || matches!(key.wVirtualKeyCode,
695 0x08 | 0x09 | 0x0d | 0x1b | 0x20..=0x28 | 0x2d | 0x2e | 0x70..=0x87))
696 };
697 if key_down {
698 return Ok(true);
699 }
700 if unsafe { ReadConsoleInputW(self.input, &mut record, 1, &mut n) } == 0 {
702 return Err(io::Error::last_os_error());
703 }
704 }
705 }
706
707 pub(super) fn read(
708 &mut self,
709 buf: &mut [u8],
710 timeout: Option<Duration>,
711 wake: &Wake,
712 ) -> io::Result<Option<usize>> {
713 let ms =
714 timeout.map_or(INFINITE, |t| u32::try_from(t.as_millis()).unwrap_or(INFINITE - 1));
715 let handles = [self.input, wake.event];
716 if unsafe { WaitForMultipleObjects(2, handles.as_ptr(), FALSE, ms) } != WAIT_OBJECT_0 {
718 return Ok(None);
719 }
720 if !self.key_ready()? {
721 return Ok(None);
722 }
723 let room = (buf.len() / 3).saturating_sub(1).clamp(1, WIDE);
731 let mut wide = [0u16; WIDE];
732 let mut got = 0;
733 if unsafe {
735 ReadConsoleW(
736 self.input,
737 wide.as_mut_ptr().cast(),
738 u32::try_from(room).unwrap_or(u32::MAX),
739 &mut got,
740 std::ptr::null(),
741 )
742 } == 0
743 {
744 return Err(io::Error::last_os_error());
745 }
746 let n = super::utf16_to_utf8(&mut self.high, &wide[..(got as usize).min(room)], buf);
747 Ok((n > 0).then_some(n))
750 }
751
752 pub(super) fn leave(&mut self) -> io::Result<()> {
753 unsafe {
755 let input = SetConsoleMode(self.input, self.in_mode);
756 let input_error = (input == 0).then(io::Error::last_os_error);
757 let output = SetConsoleMode(self.output, self.out_mode);
758 if let Some(err) = input_error {
759 return Err(err);
760 }
761 if output == 0 {
762 return Err(io::Error::last_os_error());
763 }
764 }
765 Ok(())
766 }
767 }
768}