diff options
author | Aspen Smith <grfn@gws.fyi> | 2024-02-12T03·00-0500 |
---|---|---|
committer | clbot <clbot@tvl.fyi> | 2024-02-14T19·37+0000 |
commit | 82ecd61f5c699cf3af6c4eadf47a1c52b1d696c6 (patch) | |
tree | 429c5e078528000591742ec3211bc768ae913a78 /users/aspen/xanthous/server/src/pty.rs | |
parent | 0ba476a4266015f278f18d74094299de74a5a111 (diff) |
chore(users): grfn -> aspen r/7511
Change-Id: I6c6847fac56f0a9a1a2209792e00a3aec5e672b9 Reviewed-on: https://cl.tvl.fyi/c/depot/+/10809 Autosubmit: aspen <root@gws.fyi> Reviewed-by: sterni <sternenseemann@systemli.org> Tested-by: BuildkiteCI Reviewed-by: lukegb <lukegb@tvl.fyi>
Diffstat (limited to 'users/aspen/xanthous/server/src/pty.rs')
-rw-r--r-- | users/aspen/xanthous/server/src/pty.rs | 172 |
1 files changed, 172 insertions, 0 deletions
diff --git a/users/aspen/xanthous/server/src/pty.rs b/users/aspen/xanthous/server/src/pty.rs new file mode 100644 index 000000000000..234ecd8f2336 --- /dev/null +++ b/users/aspen/xanthous/server/src/pty.rs @@ -0,0 +1,172 @@ +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, 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) + } +} |