std: Implement CommandExt::before_exec

This is a Unix-specific function which adds the ability to register a closure to
run pre-exec to configure the child process as required (note that these
closures are run post-fork).

cc #31398

4 files changed, 143 insertions, 4 deletions

diff --git a/src/libstd/process.rs b/src/libstd/process.rs
index c64471cc729..e11bb72a35a 100644
--- a/src/libstd/process.rs
+++ b/src/libstd/process.rs
@@ -269,7 +269,7 @@ impl Command {
         self
     }
 
-    fn spawn_inner(&self, default_io: StdioImp) -> io::Result<Child> {
+    fn spawn_inner(&mut self, default_io: StdioImp) -> io::Result<Child> {
         let default_io = Stdio(default_io);
 
         // See comment on `setup_io` for what `_drop_later` is.
@@ -283,7 +283,7 @@ impl Command {
             setup_io(self.stderr.as_ref().unwrap_or(&default_io), false)
         );
 
-        match imp::Process::spawn(&self.inner, their_stdin, their_stdout,
+        match imp::Process::spawn(&mut self.inner, their_stdin, their_stdout,
                                   their_stderr) {
             Err(e) => Err(e),
             Ok(handle) => Ok(Child {
diff --git a/src/libstd/sys/unix/ext/process.rs b/src/libstd/sys/unix/ext/process.rs
index 97938b07f8b..96727ed6674 100644
--- a/src/libstd/sys/unix/ext/process.rs
+++ b/src/libstd/sys/unix/ext/process.rs
@@ -12,6 +12,9 @@
 
 #![stable(feature = "rust1", since = "1.0.0")]
 
+use prelude::v1::*;
+
+use io;
 use os::unix::io::{FromRawFd, RawFd, AsRawFd, IntoRawFd};
 use os::unix::raw::{uid_t, gid_t};
 use process;
@@ -44,6 +47,34 @@ pub trait CommandExt {
     #[unstable(feature = "process_session_leader", reason = "recently added",
                issue = "27811")]
     fn session_leader(&mut self, on: bool) -> &mut process::Command;
+
+    /// Schedules a closure to be run just before the `exec` function is
+    /// invoked.
+    ///
+    /// The closure is allowed to return an I/O error whose OS error code will
+    /// be communicated back to the parent and returned as an error from when
+    /// the spawn was requested.
+    ///
+    /// Multiple closures can be registered and they will be called in order of
+    /// their registration. If a closure returns `Err` then no further closures
+    /// will be called and the spawn operation will immediately return with a
+    /// failure.
+    ///
+    /// # Notes
+    ///
+    /// This closure will be run in the context of the child process after a
+    /// `fork`. This primarily means that any modificatons made to memory on
+    /// behalf of this closure will **not** be visible to the parent process.
+    /// This is often a very constrained environment where normal operations
+    /// like `malloc` or acquiring a mutex are not guaranteed to work (due to
+    /// other threads perhaps still running when the `fork` was run).
+    ///
+    /// When this closure is run, aspects such as the stdio file descriptors and
+    /// working directory have successfully been changed, so output to these
+    /// locations may not appear where intended.
+    #[unstable(feature = "process_exec", issue = "31398")]
+    fn before_exec<F>(&mut self, f: F) -> &mut process::Command
+        where F: FnMut() -> io::Result<()> + Send + Sync + 'static;
 }
 
 #[stable(feature = "rust1", since = "1.0.0")]
@@ -62,6 +93,13 @@ impl CommandExt for process::Command {
         self.as_inner_mut().session_leader(on);
         self
     }
+
+    fn before_exec<F>(&mut self, f: F) -> &mut process::Command
+        where F: FnMut() -> io::Result<()> + Send + Sync + 'static
+    {
+        self.as_inner_mut().before_exec(Box::new(f));
+        self
+    }
 }
 
 /// Unix-specific extensions to `std::process::ExitStatus`
diff --git a/src/libstd/sys/unix/process.rs b/src/libstd/sys/unix/process.rs
index ed512b834f8..7387e9def9f 100644
--- a/src/libstd/sys/unix/process.rs
+++ b/src/libstd/sys/unix/process.rs
@@ -58,6 +58,7 @@ pub struct Command {
     gid: Option<gid_t>,
     session_leader: bool,
     saw_nul: bool,
+    closures: Vec<Box<FnMut() -> io::Result<()> + Send + Sync>>,
 }
 
 impl Command {
@@ -75,6 +76,7 @@ impl Command {
             gid: None,
             session_leader: false,
             saw_nul: saw_nul,
+            closures: Vec::new(),
         }
     }
 
@@ -164,6 +166,11 @@ impl Command {
     pub fn session_leader(&mut self, session_leader: bool) {
         self.session_leader = session_leader;
     }
+
+    pub fn before_exec(&mut self,
+                       f: Box<FnMut() -> io::Result<()> + Send + Sync>) {
+        self.closures.push(f);
+    }
 }
 
 fn os2c(s: &OsStr, saw_nul: &mut bool) -> CString {
@@ -283,7 +290,7 @@ impl Process {
         Ok(())
     }
 
-    pub fn spawn(cfg: &Command,
+    pub fn spawn(cfg: &mut Command,
                  in_fd: Stdio,
                  out_fd: Stdio,
                  err_fd: Stdio) -> io::Result<Process> {
@@ -387,7 +394,7 @@ impl Process {
     // allocation). Instead we just close it manually. This will never
     // have the drop glue anyway because this code never returns (the
     // child will either exec() or invoke libc::exit)
-    unsafe fn exec(cfg: &Command,
+    unsafe fn exec(cfg: &mut Command,
                    in_fd: Stdio,
                    out_fd: Stdio,
                    err_fd: Stdio) -> io::Error {
@@ -497,6 +504,10 @@ impl Process {
             }
         }
 
+        for callback in cfg.closures.iter_mut() {
+            try!(callback());
+        }
+
         libc::execvp(cfg.argv[0], cfg.argv.as_ptr());
         io::Error::last_os_error()
     }
diff --git a/src/test/run-pass/command-before-exec.rs b/src/test/run-pass/command-before-exec.rs
new file mode 100644
index 00000000000..16560637b69
--- /dev/null
+++ b/src/test/run-pass/command-before-exec.rs
@@ -0,0 +1,90 @@
+// Copyright 2016 The Rust Project Developers. See the COPYRIGHT
+// file at the top-level directory of this distribution and at
+// http://rust-lang.org/COPYRIGHT.
+//
+// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
+// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
+// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
+// option. This file may not be copied, modified, or distributed
+// except according to those terms.
+
+// ignore-windows - this is a unix-specific test
+
+#![feature(process_exec, libc)]
+
+extern crate libc;
+
+use std::env;
+use std::io::Error;
+use std::os::unix::process::CommandExt;
+use std::process::Command;
+use std::sync::Arc;
+use std::sync::atomic::{AtomicUsize, Ordering};
+
+fn main() {
+    if let Some(arg) = env::args().skip(1).next() {
+        match &arg[..] {
+            "test1" => println!("hello2"),
+            "test2" => assert_eq!(env::var("FOO").unwrap(), "BAR"),
+            "test3" => assert_eq!(env::current_dir().unwrap()
+                                      .to_str().unwrap(), "/"),
+            "empty" => {}
+            _ => panic!("unknown argument: {}", arg),
+        }
+        return
+    }
+
+    let me = env::current_exe().unwrap();
+
+    let output = Command::new(&me).arg("test1").before_exec(|| {
+        println!("hello");
+        Ok(())
+    }).output().unwrap();
+    assert!(output.status.success());
+    assert!(output.stderr.is_empty());
+    assert_eq!(output.stdout, b"hello\nhello2\n");
+
+    let output = Command::new(&me).arg("test2").before_exec(|| {
+        env::set_var("FOO", "BAR");
+        Ok(())
+    }).output().unwrap();
+    assert!(output.status.success());
+    assert!(output.stderr.is_empty());
+    assert!(output.stdout.is_empty());
+
+    let output = Command::new(&me).arg("test3").before_exec(|| {
+        env::set_current_dir("/").unwrap();
+        Ok(())
+    }).output().unwrap();
+    assert!(output.status.success());
+    assert!(output.stderr.is_empty());
+    assert!(output.stdout.is_empty());
+
+    let output = Command::new(&me).arg("bad").before_exec(|| {
+        Err(Error::from_raw_os_error(102))
+    }).output().err().unwrap();
+    assert_eq!(output.raw_os_error(), Some(102));
+
+    let pid = unsafe { libc::getpid() };
+    assert!(pid >= 0);
+    let output = Command::new(&me).arg("empty").before_exec(move || {
+        let child = unsafe { libc::getpid() };
+        assert!(child >= 0);
+        assert!(pid != child);
+        Ok(())
+    }).output().unwrap();
+    assert!(output.status.success());
+    assert!(output.stderr.is_empty());
+    assert!(output.stdout.is_empty());
+
+    let mem = Arc::new(AtomicUsize::new(0));
+    let mem2 = mem.clone();
+    let output = Command::new(&me).arg("empty").before_exec(move || {
+        assert_eq!(mem2.fetch_add(1, Ordering::SeqCst), 0);
+        Ok(())
+    }).output().unwrap();
+    assert!(output.status.success());
+    assert!(output.stderr.is_empty());
+    assert!(output.stdout.is_empty());
+    assert_eq!(mem.load(Ordering::SeqCst), 0);
+}