use std::io;
use std::io::{Read, Write};
use std::os::fd::{FromRawFd, OwnedFd};
use std::os::unix::io::AsRawFd;
use std::os::unix::process::CommandExt;
use std::pin::Pin;
use std::process::Stdio;
use std::task::{Context, Poll};
use tokio::io::unix::AsyncFd;
use tokio::io::{AsyncRead, AsyncWrite, ReadBuf};
use tokio::process::{Child, Command};
use crate::sys::SpawnOptions;
use crate::sys::pid_gate::PidGate;
use super::{EofOnEio, PtyExitStatus, PtyReader, PtySpawn, PtyWriter};
pub(crate) struct PtyChild {
child: Child,
resize_fd: OwnedFd,
}
impl PtyChild {
pub(crate) fn id(&self) -> Option<u32> {
self.child.id()
}
pub(crate) fn resize(&self, cols: u16, rows: u16) -> io::Result<()> {
let winsize = libc::winsize {
ws_row: rows,
ws_col: cols,
ws_xpixel: 0,
ws_ypixel: 0,
};
let rc = unsafe {
libc::ioctl(
self.resize_fd.as_raw_fd(),
libc::TIOCSWINSZ as _,
&raw const winsize,
)
};
if rc != 0 {
return Err(io::Error::last_os_error());
}
Ok(())
}
pub(crate) async fn reap(&mut self, gate: &PidGate) -> io::Result<PtyExitStatus> {
use std::future::Future;
let mut wait = std::pin::pin!(self.child.wait());
let status = std::future::poll_fn(|cx| {
let mut out = std::task::Poll::Pending;
gate.reap_under_lock(|| match wait.as_mut().poll(cx) {
std::task::Poll::Ready(res) => {
out = std::task::Poll::Ready(res);
true
}
std::task::Poll::Pending => false,
});
out
})
.await?;
Ok(PtyExitStatus::from_std(status))
}
pub(crate) async fn wait(&mut self) -> io::Result<PtyExitStatus> {
self.child.wait().await.map(PtyExitStatus::from_std)
}
pub(crate) fn try_wait(&mut self) -> io::Result<Option<PtyExitStatus>> {
Ok(self.child.try_wait()?.map(PtyExitStatus::from_std))
}
pub(crate) fn start_kill(&mut self) -> io::Result<()> {
self.child.start_kill()
}
}
fn disable_echo(fd: &OwnedFd) -> io::Result<()> {
let raw = fd.as_raw_fd();
let mut termios: libc::termios = unsafe { std::mem::zeroed() };
if unsafe { libc::tcgetattr(raw, &mut termios) } != 0 {
return Err(io::Error::last_os_error());
}
termios.c_lflag &= !(libc::ECHO | libc::ECHOE | libc::ECHOK | libc::ECHONL);
if unsafe { libc::tcsetattr(raw, libc::TCSANOW, &termios) } != 0 {
return Err(io::Error::last_os_error());
}
Ok(())
}
fn open_pty(cols: u16, rows: u16) -> io::Result<(OwnedFd, OwnedFd)> {
let mut master: libc::c_int = -1;
let mut slave: libc::c_int = -1;
let mut winsize = libc::winsize {
ws_row: rows,
ws_col: cols,
ws_xpixel: 0,
ws_ypixel: 0,
};
#[allow(clippy::unnecessary_mut_passed)]
let rc = unsafe {
libc::openpty(
&mut master,
&mut slave,
std::ptr::null_mut(),
std::ptr::null_mut::<libc::termios>(),
&mut winsize,
)
};
if rc != 0 {
return Err(io::Error::last_os_error());
}
let master = unsafe { OwnedFd::from_raw_fd(master) };
let slave = unsafe { OwnedFd::from_raw_fd(slave) };
Ok((master, slave))
}
fn set_nonblocking(fd: &OwnedFd) -> io::Result<()> {
let raw = fd.as_raw_fd();
let flags = unsafe { libc::fcntl(raw, libc::F_GETFL) };
if flags < 0 {
return Err(io::Error::last_os_error());
}
if unsafe { libc::fcntl(raw, libc::F_SETFL, flags | libc::O_NONBLOCK) } < 0 {
return Err(io::Error::last_os_error());
}
Ok(())
}
#[derive(Debug)]
struct AsyncPtyMaster {
master: AsyncFd<std::fs::File>,
}
impl AsyncPtyMaster {
fn new(fd: OwnedFd) -> io::Result<Self> {
set_nonblocking(&fd)?;
Ok(Self {
master: AsyncFd::new(std::fs::File::from(fd))?,
})
}
}
impl AsyncRead for AsyncPtyMaster {
fn poll_read(
self: Pin<&mut Self>,
cx: &mut Context<'_>,
buf: &mut ReadBuf<'_>,
) -> Poll<io::Result<()>> {
let this = self.get_mut();
loop {
let mut guard = match this.master.poll_read_ready(cx) {
Poll::Ready(Ok(guard)) => guard,
Poll::Ready(Err(e)) => return Poll::Ready(Err(e)),
Poll::Pending => return Poll::Pending,
};
let unfilled = buf.initialize_unfilled();
match guard.try_io(|inner| {
let mut file: &std::fs::File = inner.get_ref();
file.read(unfilled)
}) {
Ok(Ok(n)) => {
buf.advance(n);
return Poll::Ready(Ok(()));
}
Ok(Err(e)) => return Poll::Ready(Err(e)),
Err(_would_block) => continue,
}
}
}
}
impl AsyncWrite for AsyncPtyMaster {
fn poll_write(
self: Pin<&mut Self>,
cx: &mut Context<'_>,
data: &[u8],
) -> Poll<io::Result<usize>> {
let this = self.get_mut();
loop {
let mut guard = match this.master.poll_write_ready(cx) {
Poll::Ready(Ok(guard)) => guard,
Poll::Ready(Err(e)) => return Poll::Ready(Err(e)),
Poll::Pending => return Poll::Pending,
};
match guard.try_io(|inner| {
let mut file: &std::fs::File = inner.get_ref();
file.write(data)
}) {
Ok(result) => return Poll::Ready(result),
Err(_would_block) => continue,
}
}
}
fn poll_flush(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<io::Result<()>> {
Poll::Ready(Ok(()))
}
fn poll_shutdown(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<io::Result<()>> {
Poll::Ready(Ok(()))
}
}
pub(crate) fn spawn_pty<F>(cmd: &mut Command, opts: &SpawnOptions, spawn: F) -> io::Result<PtySpawn>
where
F: FnOnce(&mut Command, &SpawnOptions) -> io::Result<Child>,
{
let (cols, rows) = opts.pty_size.unwrap_or(super::DEFAULT_PTY_SIZE);
let (master, slave) = open_pty(cols, rows)?;
disable_echo(&slave)?;
let slave_out = slave.try_clone()?;
let slave_err = slave.try_clone()?;
cmd.stdin(Stdio::from(slave));
cmd.stdout(Stdio::from(slave_out));
cmd.stderr(Stdio::from(slave_err));
let mut pty_opts = *opts;
if !pty_opts.setsid {
unsafe {
cmd.as_std_mut().pre_exec(|| {
if libc::setsid() == -1 {
Err(io::Error::last_os_error())
} else {
Ok(())
}
});
}
}
pty_opts.setsid = true;
unsafe {
cmd.as_std_mut().pre_exec(|| {
if libc::ioctl(0, libc::TIOCSCTTY as _, 0) == -1 {
Err(io::Error::last_os_error())
} else {
Ok(())
}
});
}
let child = spawn(cmd, &pty_opts)?;
let pid = child.id();
let master_w = master.try_clone()?;
let master_resize = master.try_clone()?;
let reader: PtyReader = Box::new(EofOnEio(AsyncPtyMaster::new(master)?));
let writer: PtyWriter = Box::new(AsyncPtyMaster::new(master_w)?);
Ok(PtySpawn {
child: PtyChild {
child,
resize_fd: master_resize,
},
reader,
writer,
pid,
})
}