use std::sync::atomic::{AtomicBool, AtomicU64, Ordering};
use std::sync::{Arc, Mutex, OnceLock};
use std::time::{Duration, Instant};
use std::{io, thread};
use windows::Win32::Foundation::{CloseHandle, ERROR_NOT_FOUND, HANDLE, WAIT_EVENT, WAIT_OBJECT_0};
use windows::Win32::Globalization::CP_UTF8;
use windows::Win32::Storage::FileSystem::{ReadFile, WriteFile};
use windows::Win32::System::Console::{
CONSOLE_MODE, CONSOLE_SCREEN_BUFFER_INFO, CTRL_BREAK_EVENT, CTRL_C_EVENT, ENABLE_ECHO_INPUT,
ENABLE_LINE_INPUT, ENABLE_PROCESSED_INPUT, ENABLE_PROCESSED_OUTPUT,
ENABLE_VIRTUAL_TERMINAL_INPUT, ENABLE_VIRTUAL_TERMINAL_PROCESSING, ENABLE_WINDOW_INPUT,
ENABLE_WRAP_AT_EOL_OUTPUT, GetConsoleCP, GetConsoleMode, GetConsoleOutputCP,
GetConsoleScreenBufferInfo, GetStdHandle, INPUT_RECORD, KEY_EVENT, PeekConsoleInputW,
ReadConsoleInputW, STD_HANDLE, STD_INPUT_HANDLE, STD_OUTPUT_HANDLE, SetConsoleCP,
SetConsoleCtrlHandler, SetConsoleMode, SetConsoleOutputCP, WINDOW_BUFFER_SIZE_EVENT,
};
use windows::Win32::System::IO::CancelIoEx;
use windows::Win32::System::Threading::{CreateEventW, INFINITE, SetEvent, WaitForMultipleObjects};
use windows::core::{BOOL, HRESULT};
use crate::pal::error::{PalError, PalErrorKind};
use crate::pal::ids::RelayLeaseId;
use crate::pal::local_console::{ConsoleInput, LocalConsole};
use crate::pal::pseudoconsole::WindowSize;
use crate::pal::raw_handle::RawHandle;
#[derive(Debug, Default)]
pub(crate) struct BuildTargetConsole;
const INPUT_READ_BUF: usize = 4096;
const PEEK_INPUT_RECORDS: usize = 16;
const WAIT_INPUT: WAIT_EVENT = WAIT_EVENT(WAIT_OBJECT_0.0 + 1);
const CANCEL_READ_HANDOFF_TIMEOUT: Duration = Duration::from_secs(1);
#[derive(Debug, Default)]
struct TakenConsole {
in_mode: Option<CONSOLE_MODE>,
out_mode: Option<CONSOLE_MODE>,
code_pages: Option<(u32, u32)>,
ctrl_handler_installed: bool,
cancel_event: Option<Arc<CancelEvent>>,
}
#[derive(Debug)]
struct CancelEvent {
handle: RawHandle,
cancelled: AtomicBool,
read_active: AtomicBool,
}
impl CancelEvent {
fn new() -> Result<Arc<Self>, PalError> {
let handle = unsafe { CreateEventW(None, true, false, None) }
.map_err(|error| PalError::with_source(PalErrorKind::Other, error))?;
Ok(Arc::new(Self {
handle: RawHandle::from_handle(handle),
cancelled: AtomicBool::new(false),
read_active: AtomicBool::new(false),
}))
}
fn as_handle(&self) -> HANDLE {
self.handle.as_handle()
}
fn begin_read(&self) -> bool {
self.read_active.store(true, Ordering::SeqCst);
if self.cancelled.load(Ordering::SeqCst) {
self.read_active.store(false, Ordering::SeqCst);
return false;
}
true
}
fn end_read(&self) {
self.read_active.store(false, Ordering::SeqCst);
}
fn await_read_end(&self, deadline: Instant) -> Result<(), PalError> {
while self.read_active.load(Ordering::SeqCst) {
if Instant::now() >= deadline {
return Err(PalError::new(PalErrorKind::Other));
}
thread::yield_now();
}
Ok(())
}
fn cancel(&self, input: HANDLE) -> Result<(), PalError> {
self.cancelled.store(true, Ordering::SeqCst);
unsafe { SetEvent(self.as_handle()) }
.map_err(|error| PalError::with_source(PalErrorKind::Other, error))?;
let deadline = Instant::now()
.checked_add(CANCEL_READ_HANDOFF_TIMEOUT)
.expect("the cancellation handoff timeout fits in Instant");
while self.read_active.load(Ordering::SeqCst) {
match unsafe { CancelIoEx(input, None) } {
Ok(()) => {
return self.await_read_end(deadline);
}
Err(error) if error.code() == HRESULT::from_win32(ERROR_NOT_FOUND.0) => {
if Instant::now() >= deadline {
return Err(PalError::with_source(PalErrorKind::Other, error));
}
thread::yield_now();
}
Err(error) => {
return Err(PalError::with_source(PalErrorKind::Other, error));
}
}
}
Ok(())
}
}
impl Drop for CancelEvent {
fn drop(&mut self) {
_ = unsafe { CloseHandle(self.as_handle()) };
}
}
fn relay_lease() -> &'static Mutex<Option<(RelayLeaseId, TakenConsole)>> {
static LEASE: OnceLock<Mutex<Option<(RelayLeaseId, TakenConsole)>>> = OnceLock::new();
LEASE.get_or_init(|| Mutex::new(None))
}
fn active_cancel_event() -> Result<Arc<CancelEvent>, PalError> {
relay_lease()
.lock()
.expect("the relay lease is only replaced, never held across a panic")
.as_ref()
.and_then(|(_id, taken)| taken.cancel_event.as_ref())
.cloned()
.ok_or_else(|| PalError::new(PalErrorKind::Other))
}
fn next_lease_id() -> RelayLeaseId {
static NEXT: AtomicU64 = AtomicU64::new(1);
RelayLeaseId(NEXT.fetch_add(1, Ordering::Relaxed))
}
#[cfg_attr(coverage_nightly, coverage(off))]
unsafe extern "system" fn relay_ctrl_handler(ctrl_type: u32) -> BOOL {
BOOL::from(ctrl_type == CTRL_C_EVENT || ctrl_type == CTRL_BREAK_EVENT)
}
fn std_handle(kind: STD_HANDLE) -> Result<HANDLE, PalError> {
let handle = unsafe { GetStdHandle(kind) }
.map_err(|error| PalError::with_source(PalErrorKind::Other, error))?;
if handle.is_invalid() {
return Err(PalError::new(PalErrorKind::Other));
}
Ok(handle)
}
fn console_mode(handle: HANDLE) -> Result<CONSOLE_MODE, PalError> {
let mut mode = CONSOLE_MODE(0);
unsafe { GetConsoleMode(handle, &raw mut mode) }
.map_err(|error| PalError::with_source(PalErrorKind::Other, error))?;
Ok(mode)
}
fn restore_mode(kind: STD_HANDLE, mode: CONSOLE_MODE) -> Result<(), PalError> {
let handle = std_handle(kind)?;
unsafe { SetConsoleMode(handle, mode) }
.map_err(|error| PalError::with_source(PalErrorKind::Other, error))
}
fn set_code_pages(input: u32, output: u32) -> Result<(), PalError> {
let input = unsafe { SetConsoleCP(input) };
let output = unsafe { SetConsoleOutputCP(output) };
input
.and(output)
.map_err(|error| PalError::with_source(PalErrorKind::Other, error))
}
fn read_window_size(output: HANDLE) -> Result<WindowSize, PalError> {
let mut info = CONSOLE_SCREEN_BUFFER_INFO::default();
unsafe { GetConsoleScreenBufferInfo(output, &raw mut info) }
.map_err(|error| PalError::with_source(PalErrorKind::Other, error))?;
let width = info
.srWindow
.Right
.checked_sub(info.srWindow.Left)
.and_then(|delta| delta.checked_add(1))
.unwrap_or(1);
let height = info
.srWindow
.Bottom
.checked_sub(info.srWindow.Top)
.and_then(|delta| delta.checked_add(1))
.unwrap_or(1);
WindowSize::new(
u16::try_from(width.max(1)).unwrap_or(1),
u16::try_from(height.max(1)).unwrap_or(1),
)
.ok_or_else(|| PalError::new(PalErrorKind::Other))
}
fn event_kind(record: &INPUT_RECORD) -> u32 {
u32::from(record.EventType)
}
fn peek_input(handle: HANDLE) -> Result<([INPUT_RECORD; PEEK_INPUT_RECORDS], usize), PalError> {
let mut peek = [INPUT_RECORD::default(); PEEK_INPUT_RECORDS];
let mut count = 0_u32;
unsafe { PeekConsoleInputW(handle, &mut peek, &raw mut count) }
.map_err(|error| PalError::with_source(PalErrorKind::Other, error))?;
Ok((peek, count as usize))
}
fn consume_records(handle: HANDLE, count: usize) -> Result<(), PalError> {
if count == 0 {
return Ok(());
}
let mut discarded = [INPUT_RECORD::default(); PEEK_INPUT_RECORDS];
let discarded = discarded
.get_mut(..count)
.ok_or_else(|| PalError::new(PalErrorKind::Other))?;
let mut read = 0_u32;
unsafe { ReadConsoleInputW(handle, discarded, &raw mut read) }
.map_err(|error| PalError::with_source(PalErrorKind::Other, error))?;
Ok(())
}
fn take_leading_resize(handle: HANDLE) -> Result<Option<WindowSize>, PalError> {
let (peek, count) = peek_input(handle)?;
let leading_resizes = peek
.iter()
.take(count)
.take_while(|record| event_kind(record) == WINDOW_BUFFER_SIZE_EVENT)
.count();
if leading_resizes == 0 {
return Ok(None);
}
consume_records(handle, leading_resizes)?;
let output = std_handle(STD_OUTPUT_HANDLE)?;
read_window_size(output).map(Some)
}
fn discard_leading_noise(handle: HANDLE) -> Result<bool, PalError> {
let (peek, count) = peek_input(handle)?;
let leading_noise = peek
.iter()
.take(count)
.take_while(|record| {
let kind = event_kind(record);
kind != WINDOW_BUFFER_SIZE_EVENT && kind != KEY_EVENT
})
.count();
if leading_noise == 0 {
return Ok(false);
}
consume_records(handle, leading_noise)?;
Ok(true)
}
fn leading_record_is_key(handle: HANDLE) -> Result<Option<bool>, PalError> {
let (peek, count) = peek_input(handle)?;
Ok(peek
.first()
.filter(|_record| count != 0)
.map(|record| event_kind(record) == KEY_EVENT))
}
#[cfg_attr(coverage_nightly, coverage(off))]
fn take_over_console(
taken: &mut TakenConsole,
input: HANDLE,
output: HANDLE,
in_mode: CONSOLE_MODE,
out_mode: CONSOLE_MODE,
) -> Result<(), PalError> {
let raw_in = CONSOLE_MODE(
(in_mode.0 & !(ENABLE_ECHO_INPUT.0 | ENABLE_LINE_INPUT.0 | ENABLE_PROCESSED_INPUT.0))
| ENABLE_VIRTUAL_TERMINAL_INPUT.0
| ENABLE_WINDOW_INPUT.0,
);
let raw_out = CONSOLE_MODE(
out_mode.0
| ENABLE_VIRTUAL_TERMINAL_PROCESSING.0
| ENABLE_PROCESSED_OUTPUT.0
| ENABLE_WRAP_AT_EOL_OUTPUT.0,
);
unsafe { SetConsoleMode(input, raw_in) }
.map_err(|error| PalError::with_source(PalErrorKind::Other, error))?;
taken.in_mode = Some(in_mode);
unsafe { SetConsoleMode(output, raw_out) }
.map_err(|error| PalError::with_source(PalErrorKind::Other, error))?;
taken.out_mode = Some(out_mode);
Ok(())
}
#[cfg_attr(coverage_nightly, coverage(off))]
fn hand_back_console(taken: &TakenConsole) -> Result<(), PalError> {
let input = taken
.in_mode
.map_or(Ok(()), |mode| restore_mode(STD_INPUT_HANDLE, mode));
let output = taken
.out_mode
.map_or(Ok(()), |mode| restore_mode(STD_OUTPUT_HANDLE, mode));
let code_pages = taken.code_pages.map_or(Ok(()), |(in_page, out_page)| {
set_code_pages(in_page, out_page)
});
let ctrl_handler = if taken.ctrl_handler_installed {
unsafe { SetConsoleCtrlHandler(Some(relay_ctrl_handler), false) }
.map_err(|error| PalError::with_source(PalErrorKind::Other, error))
} else {
Ok(())
};
input.and(output).and(code_pages).and(ctrl_handler)
}
#[cfg_attr(coverage_nightly, coverage(off))]
#[cfg_attr(test, mutants::skip)]
impl LocalConsole for BuildTargetConsole {
fn has_console(&self) -> bool {
[STD_INPUT_HANDLE, STD_OUTPUT_HANDLE]
.into_iter()
.all(|kind| {
std_handle(kind)
.ok()
.and_then(|handle| console_mode(handle).ok())
.is_some()
})
}
fn stdin_is_terminal(&self) -> bool {
std_handle(STD_INPUT_HANDLE)
.ok()
.and_then(|handle| console_mode(handle).ok())
.is_some()
}
fn begin_raw_relay(&self) -> Result<RelayLeaseId, PalError> {
let mut lease = relay_lease()
.lock()
.expect("the relay lease is only replaced, never held across a panic");
if lease.is_some() {
return Err(PalError::new(PalErrorKind::Other));
}
let input = std_handle(STD_INPUT_HANDLE)?;
let output = std_handle(STD_OUTPUT_HANDLE)?;
let in_mode = console_mode(input)?;
let out_mode = console_mode(output)?;
let in_code_page = unsafe { GetConsoleCP() };
let out_code_page = unsafe { GetConsoleOutputCP() };
let mut taken = TakenConsole {
cancel_event: Some(CancelEvent::new()?),
..TakenConsole::default()
};
let result = take_over_console(&mut taken, input, output, in_mode, out_mode)
.and_then(|()| {
taken.code_pages = Some((in_code_page, out_code_page));
set_code_pages(CP_UTF8, CP_UTF8)
})
.and_then(|()| {
unsafe { SetConsoleCtrlHandler(Some(relay_ctrl_handler), true) }
.map_err(|error| PalError::with_source(PalErrorKind::Other, error))
})
.inspect(|()| taken.ctrl_handler_installed = true);
if let Err(error) = result {
_ = hand_back_console(&taken);
return Err(error);
}
let id = next_lease_id();
*lease = Some((id, taken));
Ok(id)
}
fn end_raw_relay(&self, lease: RelayLeaseId) -> Result<(), PalError> {
let taken = {
let mut held = relay_lease()
.lock()
.expect("the relay lease is only replaced, never held across a panic");
match held.as_ref() {
Some((id, _taken)) if *id == lease => held
.take()
.map(|(_id, taken)| taken)
.ok_or_else(|| PalError::new(PalErrorKind::Other))?,
_ => return Err(PalError::new(PalErrorKind::Other)),
}
};
hand_back_console(&taken)
}
fn window_size(&self) -> Result<WindowSize, PalError> {
read_window_size(std_handle(STD_OUTPUT_HANDLE)?)
}
fn read_input(&self) -> Result<ConsoleInput, PalError> {
let handle = std_handle(STD_INPUT_HANDLE)?;
let cancel_event = active_cancel_event()?;
loop {
let wait = unsafe {
WaitForMultipleObjects(&[cancel_event.as_handle(), handle], false, INFINITE)
};
if wait == WAIT_OBJECT_0 {
return Err(PalError::new(PalErrorKind::Disconnected));
}
if wait != WAIT_INPUT {
return Err(PalError::new(PalErrorKind::Other));
}
if let Some(size) = take_leading_resize(handle)? {
return Ok(ConsoleInput::Resize(size));
}
if discard_leading_noise(handle)? {
continue;
}
match leading_record_is_key(handle)? {
Some(true) | None => {}
Some(false) => continue,
}
if !cancel_event.begin_read() {
return Err(PalError::new(PalErrorKind::Disconnected));
}
let mut buf = vec![0_u8; INPUT_READ_BUF];
let mut transferred = 0_u32;
let read = unsafe {
ReadFile(
handle,
Some(buf.as_mut_slice()),
Some(&raw mut transferred),
None,
)
};
cancel_event.end_read();
read.map_err(|error| PalError::with_source(PalErrorKind::Disconnected, error))?;
if transferred == 0 {
return Err(PalError::new(PalErrorKind::Disconnected));
}
buf.truncate(transferred as usize);
return Ok(ConsoleInput::Bytes(buf));
}
}
fn cancel_input(&self) -> Result<(), PalError> {
let input = std_handle(STD_INPUT_HANDLE)?;
active_cancel_event()?.cancel(input)
}
fn write_output(&self, data: &[u8]) -> Result<(), PalError> {
let handle = std_handle(STD_OUTPUT_HANDLE)?;
let mut remaining = data;
while !remaining.is_empty() {
let mut transferred = 0_u32;
unsafe { WriteFile(handle, Some(remaining), Some(&raw mut transferred), None) }
.map_err(|error| PalError::with_source(PalErrorKind::Other, error))?;
if transferred == 0 {
return Err(PalError::new(PalErrorKind::Other));
}
remaining = remaining
.get(transferred as usize..)
.ok_or_else(|| PalError::new(PalErrorKind::Other))?;
}
Ok(())
}
fn read_prompt_line(&self) -> Result<String, PalError> {
let mut line = String::new();
io::stdin()
.read_line(&mut line)
.map_err(PalError::from_io)?;
Ok(line.trim_end_matches(['\r', '\n']).to_string())
}
}