use std::io;
use std::pin::Pin;
use std::task::{Context, Poll};
use tokio::io::{AsyncRead, AsyncWrite, ReadBuf};
use crate::backend::ChildExit;
use crate::config::SessionConfig;
use crate::error::{ExpectError, Result, SpawnError};
use crate::types::ProcessExitStatus;
pub struct PtyTransport {
reader: Box<dyn AsyncRead + Unpin + Send>,
writer: Box<dyn AsyncWrite + Unpin + Send>,
pid: Option<u32>,
}
impl PtyTransport {
pub fn new<R, W>(reader: R, writer: W) -> Self
where
R: AsyncRead + Unpin + Send + 'static,
W: AsyncWrite + Unpin + Send + 'static,
{
Self {
reader: Box::new(reader),
writer: Box::new(writer),
pid: None,
}
}
pub const fn set_pid(&mut self, pid: u32) {
self.pid = Some(pid);
}
#[must_use]
pub const fn pid(&self) -> Option<u32> {
self.pid
}
}
impl AsyncRead for PtyTransport {
fn poll_read(
mut self: Pin<&mut Self>,
cx: &mut Context<'_>,
buf: &mut ReadBuf<'_>,
) -> Poll<io::Result<()>> {
Pin::new(&mut self.reader).poll_read(cx, buf)
}
}
impl AsyncWrite for PtyTransport {
fn poll_write(
mut self: Pin<&mut Self>,
cx: &mut Context<'_>,
buf: &[u8],
) -> Poll<io::Result<usize>> {
Pin::new(&mut self.writer).poll_write(cx, buf)
}
fn poll_flush(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
Pin::new(&mut self.writer).poll_flush(cx)
}
fn poll_shutdown(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
Pin::new(&mut self.writer).poll_shutdown(cx)
}
}
#[derive(Debug, Clone)]
#[non_exhaustive]
pub struct PtyConfig {
pub dimensions: (u16, u16),
pub login_shell: bool,
pub env_mode: EnvMode,
pub env: std::collections::HashMap<String, String>,
pub working_directory: Option<std::path::PathBuf>,
}
impl Default for PtyConfig {
fn default() -> Self {
Self {
dimensions: (80, 24),
login_shell: false,
env_mode: EnvMode::Inherit,
env: std::collections::HashMap::new(),
working_directory: None,
}
}
}
impl From<&SessionConfig> for PtyConfig {
fn from(config: &SessionConfig) -> Self {
Self {
dimensions: config.dimensions,
login_shell: false,
env_mode: match (config.inherit_env, config.env.is_empty()) {
(false, _) => EnvMode::Clear,
(true, true) => EnvMode::Inherit,
(true, false) => EnvMode::Extend,
},
env: config.env.clone(),
working_directory: config.working_dir.clone(),
}
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum EnvMode {
Inherit,
Clear,
Extend,
}
pub struct PtySpawner {
config: PtyConfig,
}
impl PtySpawner {
#[must_use]
pub fn new() -> Self {
Self {
config: PtyConfig::default(),
}
}
#[must_use]
pub const fn with_config(config: PtyConfig) -> Self {
Self { config }
}
pub const fn set_dimensions(&mut self, cols: u16, rows: u16) {
self.config.dimensions = (cols, rows);
}
#[cfg(unix)]
pub async fn spawn(&self, command: &str, args: &[String]) -> Result<PtyHandle> {
use rust_pty::{PtySystem, UnixPtySystem};
if let Some(dir) = &self.config.working_directory
&& !dir.is_dir()
{
return Err(ExpectError::Spawn(SpawnError::InvalidWorkingDir {
path: dir.display().to_string(),
}));
}
let built_env: Option<std::collections::HashMap<std::ffi::OsString, std::ffi::OsString>> =
match self.config.env_mode {
EnvMode::Inherit if self.config.env.is_empty() => None,
EnvMode::Inherit | EnvMode::Extend => {
let mut m: std::collections::HashMap<_, _> = std::env::vars_os().collect();
for (k, v) in &self.config.env {
m.insert(std::ffi::OsString::from(k), std::ffi::OsString::from(v));
}
Some(m)
}
EnvMode::Clear => Some(
self.config
.env
.iter()
.map(|(k, v)| (std::ffi::OsString::from(k), std::ffi::OsString::from(v)))
.collect(),
),
};
let pty_config = rust_pty::PtyConfig {
window_size: self.config.dimensions,
env: match self.config.env_mode {
EnvMode::Clear if self.config.env.is_empty() => {
Some(std::collections::HashMap::new())
}
_ => built_env,
},
working_directory: self.config.working_directory.clone(),
..Default::default()
};
let (master, child) =
UnixPtySystem::spawn(command, args.iter().map(String::as_str), &pty_config)
.await
.map_err(|e| {
ExpectError::Spawn(SpawnError::PtyAllocation {
reason: format!("Unix PTY spawn failed: {e}"),
})
})?;
Ok(PtyHandle {
master,
child,
dimensions: self.config.dimensions,
})
}
#[cfg(windows)]
pub async fn spawn(&self, command: &str, args: &[String]) -> Result<WindowsPtyHandle> {
use rust_pty::{PtySystem, WindowsPtySystem};
let built_env: Option<std::collections::HashMap<std::ffi::OsString, std::ffi::OsString>> =
match self.config.env_mode {
EnvMode::Inherit if self.config.env.is_empty() => None,
EnvMode::Inherit | EnvMode::Extend => {
let mut m: std::collections::HashMap<_, _> = std::env::vars_os().collect();
for (k, v) in &self.config.env {
m.insert(std::ffi::OsString::from(k), std::ffi::OsString::from(v));
}
Some(m)
}
EnvMode::Clear => Some(
self.config
.env
.iter()
.map(|(k, v)| (std::ffi::OsString::from(k), std::ffi::OsString::from(v)))
.collect(),
),
};
let pty_config = rust_pty::PtyConfig {
window_size: self.config.dimensions,
env: match self.config.env_mode {
EnvMode::Clear if self.config.env.is_empty() => {
Some(std::collections::HashMap::new())
}
_ => built_env,
},
working_directory: self.config.working_directory.clone(),
..Default::default()
};
let (master, child) =
WindowsPtySystem::spawn(command, args.iter().map(|s| s.as_str()), &pty_config)
.await
.map_err(|e| {
ExpectError::Spawn(SpawnError::PtyAllocation {
reason: format!("Windows ConPTY spawn failed: {e}"),
})
})?;
Ok(WindowsPtyHandle {
master,
child,
dimensions: self.config.dimensions,
})
}
}
impl Default for PtySpawner {
fn default() -> Self {
Self::new()
}
}
#[cfg(unix)]
#[derive(Debug)]
pub struct PtyHandle {
pub(crate) master: rust_pty::UnixPtyMaster,
pub(crate) child: rust_pty::UnixPtyChild,
dimensions: (u16, u16),
}
#[cfg(windows)]
pub struct WindowsPtyHandle {
pub(crate) master: rust_pty::WindowsPtyMaster,
pub(crate) child: rust_pty::WindowsPtyChild,
dimensions: (u16, u16),
}
#[cfg(windows)]
impl std::fmt::Debug for WindowsPtyHandle {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("WindowsPtyHandle")
.field("dimensions", &self.dimensions)
.finish_non_exhaustive()
}
}
#[cfg(unix)]
impl PtyHandle {
#[must_use]
pub const fn pid(&self) -> u32 {
self.child.pid()
}
#[must_use]
pub const fn dimensions(&self) -> (u16, u16) {
self.dimensions
}
pub fn resize(&mut self, cols: u16, rows: u16) -> Result<()> {
use rust_pty::{PtyMaster, WindowSize};
self.master
.resize(WindowSize::new(cols, rows))
.map_err(|e| ExpectError::Io(io::Error::other(format!("resize failed: {e}"))))?;
self.dimensions = (cols, rows);
Ok(())
}
}
#[cfg(windows)]
impl WindowsPtyHandle {
#[must_use]
pub fn pid(&self) -> u32 {
self.child.pid()
}
#[must_use]
pub const fn dimensions(&self) -> (u16, u16) {
self.dimensions
}
pub fn resize(&mut self, cols: u16, rows: u16) -> Result<()> {
use rust_pty::{PtyMaster, WindowSize};
let size = WindowSize::new(cols, rows);
self.master
.resize(size)
.map_err(|e| ExpectError::Io(io::Error::other(format!("resize failed: {e}"))))?;
self.dimensions = (cols, rows);
Ok(())
}
#[must_use]
pub fn is_running(&self) -> bool {
self.child.is_running()
}
pub fn kill(&mut self) -> Result<()> {
self.child
.kill()
.map_err(|e| ExpectError::Io(io::Error::other(format!("kill failed: {e}"))))
}
}
#[cfg(unix)]
pub struct AsyncPty {
master: rust_pty::UnixPtyMaster,
child: rust_pty::UnixPtyChild,
pid: u32,
dimensions: (u16, u16),
}
#[cfg(unix)]
impl AsyncPty {
pub fn from_handle(handle: PtyHandle) -> io::Result<Self> {
let pid = handle.child.pid();
let dimensions = handle.dimensions;
Ok(Self {
master: handle.master,
child: handle.child,
pid,
dimensions,
})
}
pub fn try_wait(&mut self) -> Option<ProcessExitStatus> {
match self.child.try_wait() {
Ok(Some(rust_pty::ExitStatus::Exited(code))) => Some(ProcessExitStatus::Exited(code)),
Ok(Some(rust_pty::ExitStatus::Signaled(sig))) => Some(ProcessExitStatus::Signaled(sig)),
Ok(None) | Err(_) => None,
}
}
pub fn is_running(&mut self) -> bool {
self.try_wait().is_none()
}
#[must_use]
pub const fn pid(&self) -> u32 {
self.pid
}
#[must_use]
pub const fn dimensions(&self) -> (u16, u16) {
self.dimensions
}
pub fn resize(&mut self, cols: u16, rows: u16) -> Result<()> {
use rust_pty::{PtyMaster, WindowSize};
self.master
.resize(WindowSize::new(cols, rows))
.map_err(|e| ExpectError::Io(io::Error::other(format!("resize failed: {e}"))))?;
self.dimensions = (cols, rows);
Ok(())
}
#[allow(unsafe_code)]
pub fn signal(&mut self, signal: i32) -> Result<()> {
if self.try_wait().is_some() {
return Err(ExpectError::SessionClosed);
}
let result = unsafe { libc::kill(self.pid as i32, signal) };
if result == 0 {
Ok(())
} else {
let err = io::Error::last_os_error();
if err.raw_os_error() == Some(libc::ESRCH) {
Err(ExpectError::SessionClosed)
} else {
Err(ExpectError::Io(err))
}
}
}
pub fn kill(&mut self) -> Result<()> {
self.signal(libc::SIGKILL)
}
}
#[cfg(unix)]
impl AsyncRead for AsyncPty {
fn poll_read(
mut self: Pin<&mut Self>,
cx: &mut Context<'_>,
buf: &mut ReadBuf<'_>,
) -> Poll<io::Result<()>> {
Pin::new(&mut self.master).poll_read(cx, buf)
}
}
#[cfg(unix)]
impl AsyncWrite for AsyncPty {
fn poll_write(
mut self: Pin<&mut Self>,
cx: &mut Context<'_>,
buf: &[u8],
) -> Poll<io::Result<usize>> {
if matches!(self.child.try_wait(), Ok(Some(_))) {
return Poll::Ready(Err(io::ErrorKind::BrokenPipe.into()));
}
Pin::new(&mut self.master).poll_write(cx, buf)
}
fn poll_flush(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
Pin::new(&mut self.master).poll_flush(cx)
}
fn poll_shutdown(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
Pin::new(&mut self.master).poll_shutdown(cx)
}
}
#[cfg(unix)]
impl std::fmt::Debug for AsyncPty {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("AsyncPty")
.field("pid", &self.pid)
.field("dimensions", &self.dimensions)
.finish_non_exhaustive()
}
}
#[cfg(unix)]
impl ChildExit for AsyncPty {
fn try_exit_status(&mut self) -> Option<ProcessExitStatus> {
self.try_wait()
}
}
#[cfg(windows)]
pub struct WindowsAsyncPty {
master: rust_pty::WindowsPtyMaster,
child: rust_pty::WindowsPtyChild,
pid: u32,
dimensions: (u16, u16),
}
#[cfg(windows)]
impl WindowsAsyncPty {
pub fn from_handle(handle: WindowsPtyHandle) -> Self {
let pid = handle.child.pid();
let dimensions = handle.dimensions;
Self {
master: handle.master,
child: handle.child,
pid,
dimensions,
}
}
#[must_use]
pub const fn pid(&self) -> u32 {
self.pid
}
#[must_use]
pub const fn dimensions(&self) -> (u16, u16) {
self.dimensions
}
pub fn resize(&mut self, cols: u16, rows: u16) -> Result<()> {
use rust_pty::{PtyMaster, WindowSize};
let size = WindowSize::new(cols, rows);
self.master
.resize(size)
.map_err(|e| ExpectError::Io(io::Error::other(format!("resize failed: {e}"))))?;
self.dimensions = (cols, rows);
Ok(())
}
#[must_use]
pub fn is_running(&self) -> bool {
self.child.is_running()
}
pub fn kill(&mut self) -> Result<()> {
self.child
.kill()
.map_err(|e| ExpectError::Io(io::Error::other(format!("kill failed: {e}"))))
}
}
#[cfg(windows)]
impl ChildExit for WindowsAsyncPty {
fn try_exit_status(&mut self) -> Option<ProcessExitStatus> {
match self.child.try_wait() {
Ok(Some(rust_pty::ExitStatus::Exited(code))) => Some(ProcessExitStatus::Exited(code)),
Ok(Some(rust_pty::ExitStatus::Terminated(code))) => {
Some(ProcessExitStatus::Exited(code as i32))
}
Ok(None) | Err(_) => None,
}
}
}
#[cfg(windows)]
impl AsyncRead for WindowsAsyncPty {
fn poll_read(
mut self: Pin<&mut Self>,
cx: &mut Context<'_>,
buf: &mut ReadBuf<'_>,
) -> Poll<io::Result<()>> {
Pin::new(&mut self.master).poll_read(cx, buf)
}
}
#[cfg(windows)]
impl AsyncWrite for WindowsAsyncPty {
fn poll_write(
mut self: Pin<&mut Self>,
cx: &mut Context<'_>,
buf: &[u8],
) -> Poll<io::Result<usize>> {
if matches!(self.child.try_wait(), Ok(Some(_))) {
return Poll::Ready(Err(io::ErrorKind::BrokenPipe.into()));
}
Pin::new(&mut self.master).poll_write(cx, buf)
}
fn poll_flush(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
Pin::new(&mut self.master).poll_flush(cx)
}
fn poll_shutdown(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
Pin::new(&mut self.master).poll_shutdown(cx)
}
}
#[cfg(windows)]
impl std::fmt::Debug for WindowsAsyncPty {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("WindowsAsyncPty")
.field("pid", &self.pid)
.field("dimensions", &self.dimensions)
.finish_non_exhaustive()
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn pty_config_default() {
let config = PtyConfig::default();
assert_eq!(config.dimensions.0, 80);
assert_eq!(config.dimensions.1, 24);
assert_eq!(config.env_mode, EnvMode::Inherit);
}
#[test]
fn pty_config_from_session() {
let session_config = SessionConfig {
dimensions: (120, 40),
..Default::default()
};
let pty_config = PtyConfig::from(&session_config);
assert_eq!(pty_config.dimensions.0, 120);
assert_eq!(pty_config.dimensions.1, 40);
}
#[cfg(unix)]
#[tokio::test]
async fn spawn_rejects_null_byte_in_command() {
let spawner = PtySpawner::new();
let result = spawner.spawn("test\0command", &[]).await;
assert!(result.is_err());
let err = result.unwrap_err();
let err_str = err.to_string();
assert!(
err_str.contains("nul byte"),
"Expected error about a nul byte, got: {err_str}"
);
}
#[cfg(unix)]
#[tokio::test]
async fn spawn_rejects_null_byte_in_args() {
let spawner = PtySpawner::new();
let result = spawner
.spawn("/bin/echo", &["hello\0world".to_string()])
.await;
assert!(result.is_err());
let err = result.unwrap_err();
let err_str = err.to_string();
assert!(
err_str.contains("nul byte"),
"Expected error about a nul byte, got: {err_str}"
);
}
}