slightly better process exit handling

2 files changed, 51 insertions, 30 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)
     }
 }
diff --git a/src/pty.rs b/src/pty.rs
index 03dd7c1..450ea8c 100644
--- a/src/pty.rs
+++ b/src/pty.rs
@@ -1,8 +1,10 @@
 use crate::error::*;
 
+use std::os::unix::io::{FromRawFd as _, IntoRawFd as _};
+
 pub struct Pty {
-    master: nix::pty::PtyMaster,
-    slave: std::fs::File,
+    master: std::fs::File,
+    slave: std::path::PathBuf,
 }
 
 impl Pty {
@@ -13,20 +15,22 @@ impl Pty {
         nix::pty::grantpt(&master)?;
         nix::pty::unlockpt(&master)?;
 
-        let name = nix::pty::ptsname_r(&master)?;
-        let slave = std::fs::OpenOptions::new()
-            .read(true)
-            .write(true)
-            .open(name)?;
+        let slave = nix::pty::ptsname_r(&master)?.into();
+
+        let master_fd = master.into_raw_fd();
+        let master = unsafe { std::fs::File::from_raw_fd(master_fd) };
 
         Ok(Self { master, slave })
     }
 
-    pub fn master(&self) -> &nix::pty::PtyMaster {
+    pub fn master(&self) -> &std::fs::File {
         &self.master
     }
 
-    pub fn slave(&self) -> &std::fs::File {
-        &self.slave
+    pub fn slave(&self) -> Result<std::fs::File> {
+        Ok(std::fs::OpenOptions::new()
+            .read(true)
+            .write(true)
+            .open(&self.slave)?)
     }
 }