From 936d5fa20ea275d1a573d030aca0bd5e8cf2bf92 Mon Sep 17 00:00:00 2001 From: Jesse Luehrs Date: Mon, 3 Jan 2022 04:32:02 -0500 Subject: fix pipeline handling for builtins --- src/builtins.rs | 219 ++++++++++++++++++++++++++++++++++++++++++++++++-------- 1 file changed, 188 insertions(+), 31 deletions(-) (limited to 'src/builtins.rs') 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 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 { + self.fds + .get(&0.as_raw_fd()) + .copied() + .map(|fd| unsafe { async_std::fs::File::from_raw_fd(fd) }) + } + + fn set_stdin(&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 { + self.fds + .get(&1.as_raw_fd()) + .copied() + .map(|fd| unsafe { async_std::fs::File::from_raw_fd(fd) }) + } + + fn set_stdout(&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 { + self.fds + .get(&2.as_raw_fd()) + .copied() + .map(|fd| unsafe { async_std::fs::File::from_raw_fd(fd) }) + } + + fn set_stderr(&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(&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 { + 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 + Sync + Send); @@ -24,9 +171,7 @@ static BUILTINS: once_cell::sync::Lazy< pub struct Command { exe: crate::parse::Exe, f: Builtin, - stdin: Box, - stdout: Box, - stderr: Box, + 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(&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 { - (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 { 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 { 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 { 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 { 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 { 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)?)), -- cgit v1.2.3-54-g00ecf