From b1898db0f10f9641c7616e93499348d4fe743ddd Mon Sep 17 00:00:00 2001 From: Alex Crichton Date: Wed, 3 Feb 2016 16:55:59 -0800 Subject: [PATCH] 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 --- src/libstd/process.rs | 4 +- src/libstd/sys/unix/ext/process.rs | 38 ++++++++++ src/libstd/sys/unix/process.rs | 15 +++- src/test/run-pass/command-before-exec.rs | 90 ++++++++++++++++++++++++ 4 files changed, 143 insertions(+), 4 deletions(-) create mode 100644 src/test/run-pass/command-before-exec.rs diff --git a/src/libstd/process.rs b/src/libstd/process.rs index c64471cc729af..e11bb72a35a58 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 { + fn spawn_inner(&mut self, default_io: StdioImp) -> io::Result { 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 97938b07f8b95..96727ed66745a 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(&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(&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 ed512b834f83b..7387e9def9f04 100644 --- a/src/libstd/sys/unix/process.rs +++ b/src/libstd/sys/unix/process.rs @@ -58,6 +58,7 @@ pub struct Command { gid: Option, session_leader: bool, saw_nul: bool, + closures: Vec 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 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 { @@ -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 0000000000000..16560637b6926 --- /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 or the MIT license +// , 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); +}