fix pipeline handling for builtins

1 files changed, 188 insertions, 31 deletions

diff --git a/src/builtins.rs b/src/builtins.rs
index 18eb1aa..8ff1222 100644
--- a/src/builtins.rs
+++ b/src/builtins.rs
@@ -1,9 +1,156 @@
-use async_std::io::WriteExt as _;
+use async_std::io::{ReadExt as _, WriteExt as _};
+use std::os::unix::io::{AsRawFd as _, FromRawFd as _, IntoRawFd as _};
 use std::os::unix::process::ExitStatusExt as _;
 
+struct Io {
+    fds: std::collections::HashMap<
+        std::os::unix::io::RawFd,
+        std::os::unix::io::RawFd,
+    >,
+    pre_exec: Option<
+        Box<dyn 'static + FnMut() -> std::io::Result<()> + Send + Sync>,
+    >,
+}
+
+impl Io {
+    fn new() -> Self {
+        let mut fds = std::collections::HashMap::new();
+        fds.insert(0.as_raw_fd(), 0.as_raw_fd());
+        fds.insert(1.as_raw_fd(), 1.as_raw_fd());
+        fds.insert(2.as_raw_fd(), 2.as_raw_fd());
+        Self {
+            fds,
+            pre_exec: None,
+        }
+    }
+
+    fn stdin(&self) -> Option<async_std::fs::File> {
+        self.fds
+            .get(&0.as_raw_fd())
+            .copied()
+            .map(|fd| unsafe { async_std::fs::File::from_raw_fd(fd) })
+    }
+
+    fn set_stdin<T: std::os::unix::io::IntoRawFd>(&mut self, stdin: T) {
+        if let Some(fd) = self.fds.get(&0.as_raw_fd()) {
+            if *fd > 2 {
+                drop(unsafe { async_std::fs::File::from_raw_fd(*fd) });
+            }
+        }
+        self.fds.insert(0.as_raw_fd(), stdin.into_raw_fd());
+    }
+
+    fn stdout(&self) -> Option<async_std::fs::File> {
+        self.fds
+            .get(&1.as_raw_fd())
+            .copied()
+            .map(|fd| unsafe { async_std::fs::File::from_raw_fd(fd) })
+    }
+
+    fn set_stdout<T: std::os::unix::io::IntoRawFd>(&mut self, stdout: T) {
+        if let Some(fd) = self.fds.get(&1.as_raw_fd()) {
+            if *fd > 2 {
+                drop(unsafe { async_std::fs::File::from_raw_fd(*fd) });
+            }
+        }
+        self.fds.insert(1.as_raw_fd(), stdout.into_raw_fd());
+    }
+
+    fn stderr(&self) -> Option<async_std::fs::File> {
+        self.fds
+            .get(&2.as_raw_fd())
+            .copied()
+            .map(|fd| unsafe { async_std::fs::File::from_raw_fd(fd) })
+    }
+
+    fn set_stderr<T: std::os::unix::io::IntoRawFd>(&mut self, stderr: T) {
+        if let Some(fd) = self.fds.get(&2.as_raw_fd()) {
+            if *fd > 2 {
+                drop(unsafe { async_std::fs::File::from_raw_fd(*fd) });
+            }
+        }
+        self.fds.insert(2.as_raw_fd(), stderr.into_raw_fd());
+    }
+
+    pub unsafe fn pre_exec<F>(&mut self, f: F)
+    where
+        F: 'static + FnMut() -> std::io::Result<()> + Send + Sync,
+    {
+        self.pre_exec = Some(Box::new(f));
+    }
+
+    async fn read_stdin(&self, buf: &mut [u8]) -> anyhow::Result<usize> {
+        if let Some(mut fh) = self.stdin() {
+            let res = fh.read(buf).await;
+            let _ = fh.into_raw_fd();
+            Ok(res?)
+        } else {
+            Ok(0)
+        }
+    }
+
+    async fn write_stdout(&self, buf: &[u8]) -> anyhow::Result<()> {
+        if let Some(mut fh) = self.stdout() {
+            let res = fh.write_all(buf).await;
+            let _ = fh.into_raw_fd();
+            Ok(res.map(|_| ())?)
+        } else {
+            Ok(())
+        }
+    }
+
+    async fn write_stderr(&self, buf: &[u8]) -> anyhow::Result<()> {
+        if let Some(mut fh) = self.stderr() {
+            let res = fh.write_all(buf).await;
+            let _ = fh.into_raw_fd();
+            Ok(res.map(|_| ())?)
+        } else {
+            Ok(())
+        }
+    }
+
+    fn setup_command(mut self, cmd: &mut crate::command::Command) {
+        if let Some(stdin) = self.stdin() {
+            let stdin = stdin.into_raw_fd();
+            if stdin != 0 {
+                cmd.stdin(unsafe { std::fs::File::from_raw_fd(stdin) });
+                self.fds.remove(&0.as_raw_fd());
+            }
+        }
+        if let Some(stdout) = self.stdout() {
+            let stdout = stdout.into_raw_fd();
+            if stdout != 1 {
+                cmd.stdout(unsafe { std::fs::File::from_raw_fd(stdout) });
+                self.fds.remove(&1.as_raw_fd());
+            }
+        }
+        if let Some(stderr) = self.stderr() {
+            let stderr = stderr.into_raw_fd();
+            if stderr != 2 {
+                cmd.stderr(unsafe { std::fs::File::from_raw_fd(stderr) });
+                self.fds.remove(&2.as_raw_fd());
+            }
+        }
+        if let Some(pre_exec) = self.pre_exec.take() {
+            unsafe { cmd.pre_exec(pre_exec) };
+        }
+    }
+}
+
+impl Drop for Io {
+    fn drop(&mut self) {
+        for fd in self.fds.values() {
+            if *fd > 2 {
+                drop(unsafe { std::fs::File::from_raw_fd(*fd) });
+            }
+        }
+    }
+}
+
 type Builtin = &'static (dyn Fn(
     &crate::parse::Exe,
     &crate::command::Env,
+    Io,
 ) -> anyhow::Result<Child>
               + Sync
               + Send);
@@ -24,9 +171,7 @@ static BUILTINS: once_cell::sync::Lazy<
 pub struct Command {
     exe: crate::parse::Exe,
     f: Builtin,
-    stdin: Box<dyn async_std::io::Read>,
-    stdout: Box<dyn async_std::io::Write>,
-    stderr: Box<dyn async_std::io::Write>,
+    io: Io,
 }
 
 impl Command {
@@ -34,26 +179,32 @@ impl Command {
         BUILTINS.get(exe.exe()).map(|f| Self {
             exe: exe.clone(),
             f,
-            stdin: Box::new(async_std::io::stdin()),
-            stdout: Box::new(async_std::io::stdout()),
-            stderr: Box::new(async_std::io::stderr()),
+            io: Io::new(),
         })
     }
 
     pub fn stdin(&mut self, fh: std::fs::File) {
-        self.stdin = Box::new(async_std::fs::File::from(fh));
+        self.io.set_stdin(fh);
     }
 
     pub fn stdout(&mut self, fh: std::fs::File) {
-        self.stdout = Box::new(async_std::fs::File::from(fh));
+        self.io.set_stdout(fh);
     }
 
     pub fn stderr(&mut self, fh: std::fs::File) {
-        self.stderr = Box::new(async_std::fs::File::from(fh));
+        self.io.set_stderr(fh);
+    }
+
+    pub unsafe fn pre_exec<F>(&mut self, f: F)
+    where
+        F: 'static + FnMut() -> std::io::Result<()> + Send + Sync,
+    {
+        self.io.pre_exec(f);
     }
 
     pub fn spawn(self, env: &crate::command::Env) -> anyhow::Result<Child> {
-        (self.f)(&self.exe, env)
+        let Self { f, exe, io } = self;
+        (f)(&exe, env, io)
     }
 }
 
@@ -90,10 +241,12 @@ impl Child {
 fn cd(
     exe: &crate::parse::Exe,
     env: &crate::command::Env,
+    io: Io,
 ) -> anyhow::Result<Child> {
     async fn async_cd(
         exe: &crate::parse::Exe,
         _env: &crate::command::Env,
+        io: Io,
     ) -> std::process::ExitStatus {
         let dir = exe
             .args()
@@ -113,10 +266,7 @@ fn cd(
                     home().join(path.strip_prefix(prefix).unwrap())
                 } else {
                     // TODO
-                    async_std::io::stderr()
-                        .write(b"unimplemented\n")
-                        .await
-                        .unwrap();
+                    io.write_stderr(b"unimplemented\n").await.unwrap();
                     return async_std::process::ExitStatus::from_raw(1 << 8);
                 }
             } else {
@@ -128,18 +278,17 @@ fn cd(
         let code = match std::env::set_current_dir(&dir) {
             Ok(()) => 0,
             Err(e) => {
-                async_std::io::stderr()
-                    .write(
-                        format!(
-                            "{}: {}: {}\n",
-                            exe.exe(),
-                            crate::format::io_error(&e),
-                            dir.display()
-                        )
-                        .as_bytes(),
+                io.write_stderr(
+                    format!(
+                        "{}: {}: {}\n",
+                        exe.exe(),
+                        crate::format::io_error(&e),
+                        dir.display()
                     )
-                    .await
-                    .unwrap();
+                    .as_bytes(),
+                )
+                .await
+                .unwrap();
                 1
             }
         };
@@ -149,7 +298,7 @@ fn cd(
     let exe = exe.clone();
     let env = env.clone();
     Ok(Child {
-        fut: Box::pin(async move { async_cd(&exe, &env).await }),
+        fut: Box::pin(async move { async_cd(&exe, &env, io).await }),
         wrapped_child: None,
     })
 }
@@ -157,10 +306,12 @@ fn cd(
 fn and(
     exe: &crate::parse::Exe,
     env: &crate::command::Env,
+    io: Io,
 ) -> anyhow::Result<Child> {
     let exe = exe.shift();
     if env.latest_status().success() {
-        let cmd = crate::command::Command::new(&exe);
+        let mut cmd = crate::command::Command::new(&exe);
+        io.setup_command(&mut cmd);
         Ok(Child {
             fut: Box::pin(async move { unreachable!() }),
             wrapped_child: Some(Box::new(cmd.spawn(env)?)),
@@ -177,6 +328,7 @@ fn and(
 fn or(
     exe: &crate::parse::Exe,
     env: &crate::command::Env,
+    io: Io,
 ) -> anyhow::Result<Child> {
     let exe = exe.shift();
     if env.latest_status().success() {
@@ -186,7 +338,8 @@ fn or(
             wrapped_child: None,
         })
     } else {
-        let cmd = crate::command::Command::new(&exe);
+        let mut cmd = crate::command::Command::new(&exe);
+        io.setup_command(&mut cmd);
         Ok(Child {
             fut: Box::pin(async move { unreachable!() }),
             wrapped_child: Some(Box::new(cmd.spawn(env)?)),
@@ -197,9 +350,11 @@ fn or(
 fn command(
     exe: &crate::parse::Exe,
     env: &crate::command::Env,
+    io: Io,
 ) -> anyhow::Result<Child> {
     let exe = exe.shift();
-    let cmd = crate::command::Command::new_binary(&exe);
+    let mut cmd = crate::command::Command::new_binary(&exe);
+    io.setup_command(&mut cmd);
     Ok(Child {
         fut: Box::pin(async move { unreachable!() }),
         wrapped_child: Some(Box::new(cmd.spawn(env)?)),
@@ -209,9 +364,11 @@ fn command(
 fn builtin(
     exe: &crate::parse::Exe,
     env: &crate::command::Env,
+    io: Io,
 ) -> anyhow::Result<Child> {
     let exe = exe.shift();
-    let cmd = crate::command::Command::new_builtin(&exe);
+    let mut cmd = crate::command::Command::new_builtin(&exe);
+    io.setup_command(&mut cmd);
     Ok(Child {
         fut: Box::pin(async move { unreachable!() }),
         wrapped_child: Some(Box::new(cmd.spawn(env)?)),