slightly better process exit handling

1 files changed, 37 insertions, 20 deletions

diff --git a/src/command.rs b/src/command.rs
index b2d8ef7..481804f 100644
--- a/src/command.rs
+++ b/src/command.rs
@@ -1,41 +1,58 @@
 use crate::error::*;
 
 use std::os::unix::io::{AsRawFd as _, FromRawFd as _};
+use std::os::unix::process::CommandExt as _;
 
 pub struct Command {
     pty: crate::pty::Pty,
     command: std::process::Command,
+    slave_fh: Option<std::fs::File>,
 }
 
 impl Command {
-    pub fn new<S: std::convert::AsRef<std::ffi::OsStr>>(
-        program: S,
-    ) -> Result<Self> {
+    pub fn new<S: AsRef<std::ffi::OsStr>>(program: S) -> Result<Self> {
         let pty = crate::pty::Pty::new()?;
-        let fd = pty.master().as_raw_fd();
-        let stdin = unsafe { std::process::Stdio::from_raw_fd(fd) };
-        let stdout = unsafe { std::process::Stdio::from_raw_fd(fd) };
-        let stderr = unsafe { std::process::Stdio::from_raw_fd(fd) };
+        let master_fd = pty.master().as_raw_fd();
+        let slave_fh = pty.slave()?;
+        let slave_fd = slave_fh.as_raw_fd();
         let mut command = std::process::Command::new(program);
-        command.stdin(stdin).stdout(stdout).stderr(stderr);
-        Ok(Self { pty, command })
+        command
+            .stdin(unsafe { std::process::Stdio::from_raw_fd(slave_fd) })
+            .stdout(unsafe { std::process::Stdio::from_raw_fd(slave_fd) })
+            .stderr(unsafe { std::process::Stdio::from_raw_fd(slave_fd) });
+        unsafe {
+            command.pre_exec(move || {
+                // XXX unwrap
+                nix::unistd::close(master_fd)
+                    .map_err(|e| e.as_errno().unwrap())?;
+                nix::unistd::close(slave_fd)
+                    .map_err(|e| e.as_errno().unwrap())?;
+                Ok(())
+            });
+        }
+        Ok(Self {
+            pty,
+            command,
+            slave_fh: Some(slave_fh),
+        })
     }
 
     pub fn pty(&self) -> &std::fs::File {
-        self.pty.slave()
+        self.pty.master()
     }
-}
-
-impl std::ops::Deref for Command {
-    type Target = std::process::Command;
 
-    fn deref(&self) -> &Self::Target {
-        &self.command
+    pub fn args<I, S>(&mut self, args: I) -> &mut Self
+    where
+        I: IntoIterator<Item = S>,
+        S: AsRef<std::ffi::OsStr>,
+    {
+        self.command.args(args);
+        self
     }
-}
 
-impl std::ops::DerefMut for Command {
-    fn deref_mut(&mut self) -> &mut Self::Target {
-        &mut self.command
+    pub fn spawn(&mut self) -> Result<std::process::Child> {
+        let child = self.command.spawn()?;
+        self.slave_fh = None;
+        Ok(child)
     }
 }