use anyhow::{Context, Result};
use nix::fcntl::{fcntl, FcntlArg, OFlag};
use nix::pty::{openpty, Winsize};
use nix::unistd::dup;
use std::ffi::CStr;
use std::fs::{self, File};
use std::io::{self, Read, Write};
use std::os::unix::fs::symlink;
use std::os::unix::io::{AsRawFd, FromRawFd};
use std::path::{Path, PathBuf};
use std::pin::Pin;
use std::task::{Context as TaskContext, Poll};
use tokio::fs::remove_file;
use tokio::io::unix::AsyncFd;
use tokio::io::{AsyncRead, AsyncWrite, ReadBuf};
use tracing::info;
pub struct PtyEndpoint {
pub async_io: AsyncPty,
pub link_path: PathBuf,
}
pub struct AsyncPty {
inner: AsyncFd<File>,
}
impl AsyncPty {
pub fn new(master_fd: std::os::unix::io::RawFd) -> Result<Self> {
fcntl(master_fd, FcntlArg::F_SETFL(OFlag::O_NONBLOCK))
.context("failed to set PTY master non-blocking")?;
let file = unsafe { File::from_raw_fd(master_fd) };
let inner = AsyncFd::new(file).context("failed to wrap PTY master in AsyncFd")?;
Ok(Self { inner })
}
}
impl AsyncRead for AsyncPty {
fn poll_read(
self: Pin<&mut Self>,
cx: &mut TaskContext<'_>,
buf: &mut ReadBuf<'_>,
) -> Poll<io::Result<()>> {
loop {
let mut guard = match self.inner.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| inner.get_ref().read(unfilled)) {
Ok(Ok(0)) => continue,
Ok(Ok(n)) => {
buf.advance(n);
return Poll::Ready(Ok(()));
}
Ok(Err(e)) => return Poll::Ready(Err(e)),
Err(_) => continue,
}
}
}
}
impl AsyncWrite for AsyncPty {
fn poll_write(
self: Pin<&mut Self>,
cx: &mut TaskContext<'_>,
buf: &[u8],
) -> Poll<io::Result<usize>> {
loop {
let mut guard = match self.inner.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| inner.get_ref().write(buf)) {
Ok(result) => return Poll::Ready(result),
Err(_) => continue,
}
}
}
fn poll_flush(self: Pin<&mut Self>, _cx: &mut TaskContext<'_>) -> Poll<io::Result<()>> {
Poll::Ready(Ok(()))
}
fn poll_shutdown(self: Pin<&mut Self>, _cx: &mut TaskContext<'_>) -> Poll<io::Result<()>> {
Poll::Ready(Ok(()))
}
}
pub fn create_pty_endpoint(link_path: &Path) -> Result<PtyEndpoint> {
if link_path.exists() {
anyhow::bail!(
"link path {} already exists; remove it or choose another path",
link_path.display()
);
}
let winsize = Winsize {
ws_row: 24,
ws_col: 80,
ws_xpixel: 0,
ws_ypixel: 0,
};
let pty = openpty(Some(&winsize), None).context("failed to open PTY pair")?;
let slave_path = slave_path_for_master(pty.master.as_raw_fd())?;
if let Some(parent) = link_path.parent() {
fs::create_dir_all(parent)
.with_context(|| format!("failed to create parent dir {}", parent.display()))?;
}
symlink(&slave_path, link_path).with_context(|| {
format!(
"failed to symlink {} -> {}",
link_path.display(),
slave_path
)
})?;
info!(
link = %link_path.display(),
slave = %slave_path,
"PTY ready"
);
let owned_fd = dup(pty.master.as_raw_fd()).context("failed to dup PTY master fd")?;
Ok(PtyEndpoint {
async_io: AsyncPty::new(owned_fd)?,
link_path: link_path.to_path_buf(),
})
}
pub async fn cleanup_link(link_path: &Path) -> Result<()> {
if link_path.exists() || link_path.is_symlink() {
remove_file(link_path)
.await
.with_context(|| format!("failed to remove {}", link_path.display()))?;
info!(link = %link_path.display(), "removed PTY symlink");
}
Ok(())
}
fn slave_path_for_master(master_fd: std::os::unix::io::RawFd) -> Result<String> {
let name_ptr = unsafe { libc::ptsname(master_fd) };
if name_ptr.is_null() {
return Err(anyhow::anyhow!("ptsname failed for PTY master"));
}
let name = unsafe { CStr::from_ptr(name_ptr) };
Ok(name.to_string_lossy().into_owned())
}