fix: ignore stdin redirection to ensure always accessing the real tty (#2425)

This commit is contained in:
三咲雅 · Misaki Masa 2025-03-02 06:09:28 +08:00 committed by GitHub
parent 66f015de95
commit 541d849157
No known key found for this signature in database
GPG key ID: B5690EEEBB952194

View file

@ -1,10 +1,27 @@
use std::{io::{Error, ErrorKind}, time::{Duration, Instant}}; use std::{io::{Error, ErrorKind}, ops::Deref, time::{Duration, Instant}};
pub struct AsyncStdin { pub struct AsyncStdin {
fd: Fd,
#[cfg(unix)] #[cfg(unix)]
fds: libc::fd_set, fds: libc::fd_set,
} }
impl Default for AsyncStdin {
fn default() -> Self {
let fd = Fd::new().expect("failed to open stdin");
#[cfg(unix)]
{
let mut me = Self { fd, fds: unsafe { std::mem::MaybeUninit::zeroed().assume_init() } };
me.reset();
me
}
#[cfg(windows)]
{
Self { fd }
}
}
}
impl AsyncStdin { impl AsyncStdin {
pub fn read_until<P>(&mut self, timeout: Duration, predicate: P) -> (Vec<u8>, std::io::Result<()>) pub fn read_until<P>(&mut self, timeout: Duration, predicate: P) -> (Vec<u8>, std::io::Result<()>)
where where
@ -21,7 +38,7 @@ impl AsyncStdin {
continue; continue;
} }
let b = Self::read_u8()?; let b = self.read_u8()?;
buf.push(b); buf.push(b);
if predicate(b, &buf) { if predicate(b, &buf) {
@ -36,15 +53,7 @@ impl AsyncStdin {
} }
} }
#[cfg(unix)] // --- Unix
impl Default for AsyncStdin {
fn default() -> Self {
let mut me = Self { fds: unsafe { std::mem::MaybeUninit::zeroed().assume_init() } };
me.reset();
me
}
}
#[cfg(unix)] #[cfg(unix)]
impl AsyncStdin { impl AsyncStdin {
pub fn poll(&mut self, timeout: Duration) -> std::io::Result<bool> { pub fn poll(&mut self, timeout: Duration) -> std::io::Result<bool> {
@ -54,13 +63,7 @@ impl AsyncStdin {
}; };
let result = unsafe { let result = unsafe {
libc::select( libc::select(*self.fd + 1, &mut self.fds, std::ptr::null_mut(), std::ptr::null_mut(), &mut tv)
libc::STDIN_FILENO + 1,
&mut self.fds,
std::ptr::null_mut(),
std::ptr::null_mut(),
&mut tv,
)
}; };
match result { match result {
@ -73,9 +76,9 @@ impl AsyncStdin {
} }
} }
pub fn read_u8() -> std::io::Result<u8> { pub fn read_u8(&mut self) -> std::io::Result<u8> {
let mut b = 0; let mut b = 0;
match unsafe { libc::read(libc::STDIN_FILENO, &mut b as *mut _ as *mut _, 1) } { match unsafe { libc::read(*self.fd, &mut b as *mut _ as *mut _, 1) } {
-1 => Err(Error::last_os_error()), -1 => Err(Error::last_os_error()),
0 => Err(Error::new(ErrorKind::UnexpectedEof, "unexpected EOF")), 0 => Err(Error::new(ErrorKind::UnexpectedEof, "unexpected EOF")),
_ => Ok(b), _ => Ok(b),
@ -85,42 +88,31 @@ impl AsyncStdin {
fn reset(&mut self) { fn reset(&mut self) {
unsafe { unsafe {
libc::FD_ZERO(&mut self.fds); libc::FD_ZERO(&mut self.fds);
libc::FD_SET(libc::STDIN_FILENO, &mut self.fds); libc::FD_SET(*self.fd, &mut self.fds);
} }
} }
} }
#[cfg(windows)] // --- Windows
impl Default for AsyncStdin {
fn default() -> Self { Self {} }
}
#[cfg(windows)] #[cfg(windows)]
impl AsyncStdin { impl AsyncStdin {
pub fn poll(&mut self, timeout: Duration) -> std::io::Result<bool> { pub fn poll(&mut self, timeout: Duration) -> std::io::Result<bool> {
use std::os::windows::io::AsRawHandle;
use windows_sys::Win32::{Foundation::{WAIT_FAILED, WAIT_OBJECT_0}, System::Threading::WaitForSingleObject}; use windows_sys::Win32::{Foundation::{WAIT_FAILED, WAIT_OBJECT_0}, System::Threading::WaitForSingleObject};
let handle = std::io::stdin().as_raw_handle();
let millis = timeout.as_millis(); let millis = timeout.as_millis();
match unsafe { WaitForSingleObject(handle, millis as u32) } { match unsafe { WaitForSingleObject(*self.fd, millis as u32) } {
WAIT_FAILED => Err(Error::last_os_error()), WAIT_FAILED => Err(Error::last_os_error()),
WAIT_OBJECT_0 => Ok(true), WAIT_OBJECT_0 => Ok(true),
_ => Ok(false), _ => Ok(false),
} }
} }
pub fn read_u8() -> std::io::Result<u8> { pub fn read_u8(&mut self) -> std::io::Result<u8> {
use std::os::windows::io::AsRawHandle;
use windows_sys::Win32::Storage::FileSystem::ReadFile; use windows_sys::Win32::Storage::FileSystem::ReadFile;
let mut buf = 0; let mut buf = 0;
let mut bytes = 0; let mut bytes = 0;
let success = unsafe { let success = unsafe { ReadFile(*self.fd, &mut buf, 1, &mut bytes, std::ptr::null_mut()) };
ReadFile(std::io::stdin().as_raw_handle(), &mut buf, 1, &mut bytes, std::ptr::null_mut())
};
if success == 0 { if success == 0 {
return Err(Error::last_os_error()); return Err(Error::last_os_error());
@ -130,3 +122,74 @@ impl AsyncStdin {
Ok(buf) Ok(buf)
} }
} }
// --- Fd
struct Fd {
#[cfg(unix)]
inner: std::os::fd::RawFd,
#[cfg(windows)]
inner: std::os::windows::io::RawHandle,
close: bool,
}
impl Deref for Fd {
#[cfg(unix)]
type Target = std::os::fd::RawFd;
#[cfg(windows)]
type Target = std::os::windows::io::RawHandle;
fn deref(&self) -> &Self::Target { &self.inner }
}
impl Drop for Fd {
fn drop(&mut self) {
#[cfg(unix)]
if self.close {
unsafe { libc::close(self.inner) };
}
#[cfg(windows)]
if self.close {
unsafe { windows_sys::Win32::Foundation::CloseHandle(self.inner) };
}
}
}
impl Fd {
#[cfg(unix)]
fn new() -> std::io::Result<Self> {
use std::{fs::OpenOptions, os::fd::IntoRawFd};
Ok(if unsafe { libc::isatty(libc::STDIN_FILENO) } == 1 {
Self { inner: libc::STDIN_FILENO, close: false }
} else {
Self {
inner: OpenOptions::new().read(true).write(true).open("/dev/tty")?.into_raw_fd(),
close: true,
}
})
}
#[cfg(windows)]
fn new() -> std::io::Result<Self> {
use windows_sys::Win32::{Foundation::{GENERIC_READ, GENERIC_WRITE, INVALID_HANDLE_VALUE}, Storage::FileSystem::{CreateFileW, FILE_SHARE_READ, FILE_SHARE_WRITE, OPEN_EXISTING}};
let name: Vec<u16> = "CONIN$\0".encode_utf16().collect();
let result = unsafe {
CreateFileW(
name.as_ptr(),
GENERIC_READ | GENERIC_WRITE,
FILE_SHARE_READ | FILE_SHARE_WRITE,
std::ptr::null_mut(),
OPEN_EXISTING,
0,
std::ptr::null_mut(),
)
};
if result == INVALID_HANDLE_VALUE {
Err(std::io::Error::last_os_error())
} else {
Ok(Self { inner: result, close: true })
}
}
}