tvl-depot/users/grfn/xanthous/server/src/pty.rs

use std::io::{self};
use std::os::unix::prelude::{AsRawFd, CommandExt, FromRawFd};
use std::pin::Pin;
use std::process::{abort, Command};
use std::task::{Context, Poll};

use eyre::{bail, Result};
use futures::Future;
use nix::pty::forkpty;
use nix::pty::Winsize;
use nix::sys::termios::Termios;
use nix::sys::wait::{waitpid, WaitPidFlag, WaitStatus};
use nix::unistd::{ForkResult, Pid};
use tokio::fs::File;
use tokio::io::{AsyncRead, AsyncWrite};
use tokio::signal::unix::{signal, Signal, SignalKind};
use tokio::task::spawn_blocking;

mod ioctl {
    use super::Winsize;
    use libc::TIOCSWINSZ;
    use nix::ioctl_write_ptr_bad;

    ioctl_write_ptr_bad!(tiocswinsz, TIOCSWINSZ, Winsize);
}

async fn asyncify<F, T>(f: F) -> Result<T>
where
    F: FnOnce() -> Result<T> + Send + 'static,
    T: Send + 'static,
{
    match spawn_blocking(f).await {
        Ok(res) => res,
        Err(_) => bail!("background task failed",),
    }
}

pub struct Child {
    pub tty: File,
    pub pid: Pid,
}

pub struct ChildHandle {
    pub tty: File,
}

pub struct WaitPid {
    pid: Pid,
    signal: Signal,
}

impl WaitPid {
    pub fn new(pid: Pid) -> Self {
        Self {
            pid,
            signal: signal(SignalKind::child()).unwrap(),
        }
    }
}

impl Future for WaitPid {
    type Output = nix::Result<WaitStatus>;

    fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
        let _ = self.signal.poll_recv(cx);
        match waitpid(self.pid, Some(WaitPidFlag::WNOHANG)) {
            Ok(WaitStatus::StillAlive) => Poll::Pending,
            result => Poll::Ready(result),
        }
    }
}

impl Child {
    pub async fn handle(&self) -> io::Result<ChildHandle> {
        Ok(ChildHandle {
            tty: self.tty.try_clone().await?,
        })
    }
}

impl ChildHandle {
    pub async fn resize_window(&mut self, winsize: Winsize) -> Result<()> {
        let fd = self.tty.as_raw_fd();
        asyncify(move || unsafe {
            ioctl::tiocswinsz(fd, &winsize as *const Winsize)?;
            Ok(())
        })
        .await
    }
}

pub async fn spawn(
    mut cmd: Command,
    winsize: Option<Winsize>,
    termios: Option<Termios>,
) -> Result<Child> {
    asyncify(move || unsafe {
        let res = forkpty(winsize.as_ref(), termios.as_ref())?;
        match res.fork_result {
            ForkResult::Parent { child } => Ok(Child {
                pid: child,
                tty: File::from_raw_fd(res.master),
            }),
            ForkResult::Child => {
                cmd.exec();
                abort();
            }
        }
    })
    .await
}

impl AsyncRead for Child {
    fn poll_read(
        mut self: Pin<&mut Self>,
        cx: &mut Context<'_>,
        buf: &mut tokio::io::ReadBuf<'_>,
    ) -> Poll<io::Result<()>> {
        Pin::new(&mut self.tty).poll_read(cx, buf)
    }
}

impl AsyncWrite for Child {
    fn poll_write(
        mut self: Pin<&mut Self>,
        cx: &mut Context<'_>,
        buf: &[u8],
    ) -> Poll<Result<usize, io::Error>> {
        Pin::new(&mut self.tty).poll_write(cx, buf)
    }

    fn poll_flush(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), io::Error>> {
        Pin::new(&mut self.tty).poll_flush(cx)
    }

    fn poll_shutdown(
        mut self: Pin<&mut Self>,
        cx: &mut Context<'_>,
    ) -> Poll<Result<(), io::Error>> {
        Pin::new(&mut self.tty).poll_shutdown(cx)
    }
}

impl AsyncRead for ChildHandle {
    fn poll_read(
        mut self: Pin<&mut Self>,
        cx: &mut Context<'_>,
        buf: &mut tokio::io::ReadBuf<'_>,
    ) -> Poll<io::Result<()>> {
        Pin::new(&mut self.tty).poll_read(cx, buf)
    }
}

impl AsyncWrite for ChildHandle {
    fn poll_write(
        mut self: Pin<&mut Self>,
        cx: &mut Context<'_>,
        buf: &[u8],
    ) -> Poll<Result<usize, io::Error>> {
        Pin::new(&mut self.tty).poll_write(cx, buf)
    }

    fn poll_flush(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), io::Error>> {
        Pin::new(&mut self.tty).poll_flush(cx)
    }

    fn poll_shutdown(
        mut self: Pin<&mut Self>,
        cx: &mut Context<'_>,
    ) -> Poll<Result<(), io::Error>> {
        Pin::new(&mut self.tty).poll_shutdown(cx)
    }
}