deno.land / x / deno@v1.28.2 / test_util / src / pty.rs
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457use std::collections::HashMap;use std::io::Read;use std::path::Path;
pub trait Pty: Read { fn write_text(&mut self, text: &str);
fn write_line(&mut self, text: &str) { self.write_text(&format!("{}\n", text)); }
/// Reads the output to the EOF. fn read_all_output(&mut self) -> String { let mut text = String::new(); self.read_to_string(&mut text).unwrap(); text }}
#[cfg(unix)]pub fn create_pty( program: impl AsRef<Path>, args: &[&str], cwd: impl AsRef<Path>, env_vars: Option<HashMap<String, String>>,) -> Box<dyn Pty> { let fork = pty::fork::Fork::from_ptmx().unwrap(); if fork.is_parent().is_ok() { Box::new(unix::UnixPty { fork }) } else { std::process::Command::new(program.as_ref()) .current_dir(cwd) .args(args) .envs(env_vars.unwrap_or_default()) .spawn() .unwrap() .wait() .unwrap(); std::process::exit(0); }}
#[cfg(unix)]mod unix { use std::io::Read; use std::io::Write;
use super::Pty;
pub struct UnixPty { pub fork: pty::fork::Fork, }
impl Drop for UnixPty { fn drop(&mut self) { self.fork.wait().unwrap(); } }
impl Pty for UnixPty { fn write_text(&mut self, text: &str) { let mut master = self.fork.is_parent().unwrap(); master.write_all(text.as_bytes()).unwrap(); } }
impl Read for UnixPty { fn read(&mut self, buf: &mut [u8]) -> std::io::Result<usize> { let mut master = self.fork.is_parent().unwrap(); master.read(buf) } }}
#[cfg(target_os = "windows")]pub fn create_pty( program: impl AsRef<Path>, args: &[&str], cwd: impl AsRef<Path>, env_vars: Option<HashMap<String, String>>,) -> Box<dyn Pty> { let pty = windows::WinPseudoConsole::new( program, args, &cwd.as_ref().to_string_lossy(), env_vars, ); Box::new(pty)}
#[cfg(target_os = "windows")]mod windows { use std::collections::HashMap; use std::io::Read; use std::io::Write; use std::path::Path; use std::ptr; use std::time::Duration;
use winapi::shared::minwindef::FALSE; use winapi::shared::minwindef::LPVOID; use winapi::shared::minwindef::TRUE; use winapi::shared::winerror::S_OK; use winapi::um::consoleapi::ClosePseudoConsole; use winapi::um::consoleapi::CreatePseudoConsole; use winapi::um::fileapi::ReadFile; use winapi::um::fileapi::WriteFile; use winapi::um::handleapi::DuplicateHandle; use winapi::um::handleapi::INVALID_HANDLE_VALUE; use winapi::um::namedpipeapi::CreatePipe; use winapi::um::processthreadsapi::CreateProcessW; use winapi::um::processthreadsapi::DeleteProcThreadAttributeList; use winapi::um::processthreadsapi::GetCurrentProcess; use winapi::um::processthreadsapi::InitializeProcThreadAttributeList; use winapi::um::processthreadsapi::UpdateProcThreadAttribute; use winapi::um::processthreadsapi::LPPROC_THREAD_ATTRIBUTE_LIST; use winapi::um::processthreadsapi::PROCESS_INFORMATION; use winapi::um::synchapi::WaitForSingleObject; use winapi::um::winbase::CREATE_UNICODE_ENVIRONMENT; use winapi::um::winbase::EXTENDED_STARTUPINFO_PRESENT; use winapi::um::winbase::INFINITE; use winapi::um::winbase::STARTUPINFOEXW; use winapi::um::wincontypes::COORD; use winapi::um::wincontypes::HPCON; use winapi::um::winnt::DUPLICATE_SAME_ACCESS; use winapi::um::winnt::HANDLE;
use super::Pty;
macro_rules! assert_win_success { ($expression:expr) => { let success = $expression; if success != TRUE { panic!("{}", std::io::Error::last_os_error().to_string()) } }; }
pub struct WinPseudoConsole { stdin_write_handle: WinHandle, stdout_read_handle: WinHandle, // keep these alive for the duration of the pseudo console _process_handle: WinHandle, _thread_handle: WinHandle, _attribute_list: ProcThreadAttributeList, }
impl WinPseudoConsole { pub fn new( program: impl AsRef<Path>, args: &[&str], cwd: &str, maybe_env_vars: Option<HashMap<String, String>>, ) -> Self { // https://docs.microsoft.com/en-us/windows/console/creating-a-pseudoconsole-session // SAFETY: // Generous use of winapi to create a PTY (thus large unsafe block). unsafe { let mut size: COORD = std::mem::zeroed(); size.X = 800; size.Y = 500; let mut console_handle = std::ptr::null_mut(); let (stdin_read_handle, stdin_write_handle) = create_pipe(); let (stdout_read_handle, stdout_write_handle) = create_pipe();
let result = CreatePseudoConsole( size, stdin_read_handle.as_raw_handle(), stdout_write_handle.as_raw_handle(), 0, &mut console_handle, ); assert_eq!(result, S_OK);
let mut environment_vars = maybe_env_vars.map(get_env_vars); let mut attribute_list = ProcThreadAttributeList::new(console_handle); let mut startup_info: STARTUPINFOEXW = std::mem::zeroed(); startup_info.StartupInfo.cb = std::mem::size_of::<STARTUPINFOEXW>() as u32; startup_info.lpAttributeList = attribute_list.as_mut_ptr();
let mut proc_info: PROCESS_INFORMATION = std::mem::zeroed(); let command = format!( "\"{}\" {}", program.as_ref().to_string_lossy(), args.join(" ") ) .trim() .to_string(); let mut application_str = to_windows_str(&program.as_ref().to_string_lossy()); let mut command_str = to_windows_str(&command); let mut cwd = to_windows_str(cwd);
assert_win_success!(CreateProcessW( application_str.as_mut_ptr(), command_str.as_mut_ptr(), ptr::null_mut(), ptr::null_mut(), FALSE, EXTENDED_STARTUPINFO_PRESENT | CREATE_UNICODE_ENVIRONMENT, environment_vars .as_mut() .map(|v| v.as_mut_ptr() as LPVOID) .unwrap_or(ptr::null_mut()), cwd.as_mut_ptr(), &mut startup_info.StartupInfo, &mut proc_info, ));
// close the handles that the pseudoconsole now has drop(stdin_read_handle); drop(stdout_write_handle);
// start a thread that will close the pseudoconsole on process exit let thread_handle = WinHandle::new(proc_info.hThread); std::thread::spawn({ let thread_handle = thread_handle.duplicate(); let console_handle = WinHandle::new(console_handle); move || { WaitForSingleObject(thread_handle.as_raw_handle(), INFINITE); // wait for the reading thread to catch up std::thread::sleep(Duration::from_millis(200)); // close the console handle which will close the // stdout pipe for the reader ClosePseudoConsole(console_handle.into_raw_handle()); } });
Self { stdin_write_handle, stdout_read_handle, _process_handle: WinHandle::new(proc_info.hProcess), _thread_handle: thread_handle, _attribute_list: attribute_list, } } } }
impl Read for WinPseudoConsole { fn read(&mut self, buf: &mut [u8]) -> std::io::Result<usize> { loop { let mut bytes_read = 0; // SAFETY: // winapi call let success = unsafe { ReadFile( self.stdout_read_handle.as_raw_handle(), buf.as_mut_ptr() as _, buf.len() as u32, &mut bytes_read, ptr::null_mut(), ) };
// ignore zero-byte writes let is_zero_byte_write = bytes_read == 0 && success == TRUE; if !is_zero_byte_write { return Ok(bytes_read as usize); } } } }
impl Pty for WinPseudoConsole { fn write_text(&mut self, text: &str) { // windows pseudo console requires a \r\n to do a newline let newline_re = regex::Regex::new("\r?\n").unwrap(); self .write_all(newline_re.replace_all(text, "\r\n").as_bytes()) .unwrap(); } }
impl std::io::Write for WinPseudoConsole { fn write(&mut self, buffer: &[u8]) -> std::io::Result<usize> { let mut bytes_written = 0; // SAFETY: // winapi call assert_win_success!(unsafe { WriteFile( self.stdin_write_handle.as_raw_handle(), buffer.as_ptr() as *const _, buffer.len() as u32, &mut bytes_written, ptr::null_mut(), ) }); Ok(bytes_written as usize) }
fn flush(&mut self) -> std::io::Result<()> { Ok(()) } }
struct WinHandle { inner: HANDLE, }
impl WinHandle { pub fn new(handle: HANDLE) -> Self { WinHandle { inner: handle } }
pub fn duplicate(&self) -> WinHandle { // SAFETY: // winapi call let process_handle = unsafe { GetCurrentProcess() }; let mut duplicate_handle = ptr::null_mut(); // SAFETY: // winapi call assert_win_success!(unsafe { DuplicateHandle( process_handle, self.inner, process_handle, &mut duplicate_handle, 0, 0, DUPLICATE_SAME_ACCESS, ) });
WinHandle::new(duplicate_handle) }
pub fn as_raw_handle(&self) -> HANDLE { self.inner }
pub fn into_raw_handle(self) -> HANDLE { let handle = self.inner; // skip the drop implementation in order to not close the handle std::mem::forget(self); handle } }
// SAFETY: These handles are ok to send across threads. unsafe impl Send for WinHandle {} // SAFETY: These handles are ok to send across threads. unsafe impl Sync for WinHandle {}
impl Drop for WinHandle { fn drop(&mut self) { if !self.inner.is_null() && self.inner != INVALID_HANDLE_VALUE { // SAFETY: winapi call unsafe { winapi::um::handleapi::CloseHandle(self.inner); } } } }
struct ProcThreadAttributeList { buffer: Vec<u8>, }
impl ProcThreadAttributeList { pub fn new(console_handle: HPCON) -> Self { // SAFETY: // Generous use of unsafe winapi calls to create a ProcThreadAttributeList. unsafe { // discover size required for the list let mut size = 0; let attribute_count = 1; assert_eq!( InitializeProcThreadAttributeList( ptr::null_mut(), attribute_count, 0, &mut size ), FALSE );
let mut buffer = vec![0u8; size]; let attribute_list_ptr = buffer.as_mut_ptr() as _;
assert_win_success!(InitializeProcThreadAttributeList( attribute_list_ptr, attribute_count, 0, &mut size, ));
const PROC_THREAD_ATTRIBUTE_PSEUDOCONSOLE: usize = 0x00020016; assert_win_success!(UpdateProcThreadAttribute( attribute_list_ptr, 0, PROC_THREAD_ATTRIBUTE_PSEUDOCONSOLE, console_handle, std::mem::size_of::<HPCON>(), ptr::null_mut(), ptr::null_mut(), ));
ProcThreadAttributeList { buffer } } }
pub fn as_mut_ptr(&mut self) -> LPPROC_THREAD_ATTRIBUTE_LIST { self.buffer.as_mut_slice().as_mut_ptr() as *mut _ } }
impl Drop for ProcThreadAttributeList { fn drop(&mut self) { // SAFETY: // winapi call unsafe { DeleteProcThreadAttributeList(self.as_mut_ptr()) }; } }
fn create_pipe() -> (WinHandle, WinHandle) { let mut read_handle = std::ptr::null_mut(); let mut write_handle = std::ptr::null_mut();
// SAFETY: // Creating an anonymous pipe with winapi. assert_win_success!(unsafe { CreatePipe(&mut read_handle, &mut write_handle, ptr::null_mut(), 0) });
(WinHandle::new(read_handle), WinHandle::new(write_handle)) }
fn to_windows_str(str: &str) -> Vec<u16> { use std::os::windows::prelude::OsStrExt; std::ffi::OsStr::new(str) .encode_wide() .chain(Some(0)) .collect() }
fn get_env_vars(env_vars: HashMap<String, String>) -> Vec<u16> { // See lpEnvironment: https://docs.microsoft.com/en-us/windows/win32/api/processthreadsapi/nf-processthreadsapi-createprocessw let mut parts = env_vars .into_iter() // each environment variable is in the form `name=value\0` .map(|(key, value)| format!("{}={}\0", key, value)) .collect::<Vec<_>>();
// all strings in an environment block must be case insensitively // sorted alphabetically by name // https://docs.microsoft.com/en-us/windows/win32/procthread/changing-environment-variables parts.sort_by_key(|part| part.to_lowercase());
// the entire block is terminated by NULL (\0) format!("{}\0", parts.join("")) .encode_utf16() .collect::<Vec<_>>() }}
deno
Version Info
a year ago
