From 1638ec77df7886c04f6c15a5f8d4b435906c0140 Mon Sep 17 00:00:00 2001 From: Ivan Petkov Date: Sat, 27 Sep 2014 17:15:50 -0700 Subject: [PATCH] Implement inheriting more than 3 file descriptors for UNIX Commands * This allows creating sub-processes through the std::io::process::Command struct with file descriptors beyond stdin/stdout/stderr * Support is still missing under Windows --- src/libnative/io/process.rs | 150 +++++++++++++++++++++++++++++------- src/libstd/io/process.rs | 147 ++++++++++++++++++++++++++++++++++- 2 files changed, 269 insertions(+), 28 deletions(-) diff --git a/src/libnative/io/process.rs b/src/libnative/io/process.rs index cb392e1675feb..2890b3327a07b 100644 --- a/src/libnative/io/process.rs +++ b/src/libnative/io/process.rs @@ -68,11 +68,6 @@ impl Process { pub fn spawn(cfg: ProcessConfig) -> IoResult<(Process, Vec>)> { - // right now we only handle stdin/stdout/stderr. - if cfg.extra_io.len() > 0 { - return Err(super::unimpl()); - } - fn get_io(io: rtio::StdioContainer, ret: &mut Vec>) -> IoResult> @@ -97,12 +92,16 @@ impl Process { } let mut ret_io = Vec::new(); - let res = spawn_process_os(cfg, - try!(get_io(cfg.stdin, &mut ret_io)), - try!(get_io(cfg.stdout, &mut ret_io)), - try!(get_io(cfg.stderr, &mut ret_io))); + let stdin = try!(get_io(cfg.stdin, &mut ret_io)); + let stdout = try!(get_io(cfg.stdout, &mut ret_io)); + let stderr = try!(get_io(cfg.stderr, &mut ret_io)); + + let mut child_extra_io = Vec::new(); + for io in cfg.extra_io.iter() { + child_extra_io.push(try!(get_io(*io, &mut ret_io))); + } - match res { + match spawn_process_os(cfg, stdin, stdout, stderr, child_extra_io) { Ok(res) => { let p = Process { pid: res.pid, @@ -277,7 +276,8 @@ struct SpawnProcessResult { fn spawn_process_os(cfg: ProcessConfig, in_fd: Option, out_fd: Option, - err_fd: Option) + err_fd: Option, + extra_fds: Vec>) -> IoResult { use libc::types::os::arch::extra::{DWORD, HANDLE, STARTUPINFO}; use libc::consts::os::extra::{ @@ -298,6 +298,11 @@ fn spawn_process_os(cfg: ProcessConfig, use std::iter::Iterator; use std::str::StrSlice; + // right now we only handle stdin/stdout/stderr. + if extra_fds.len() > 0 { + return Err(super::unimpl()); + } + if cfg.gid.is_some() || cfg.uid.is_some() { return Err(IoError { code: libc::ERROR_CALL_NOT_IMPLEMENTED as uint, @@ -529,12 +534,15 @@ fn make_command_line(prog: &CString, args: &[CString]) -> String { fn spawn_process_os(cfg: ProcessConfig, in_fd: Option, out_fd: Option, - err_fd: Option) + err_fd: Option, + extra_fds: Vec>) -> IoResult { use libc::funcs::posix88::unistd::{fork, dup2, close, chdir, execvp}; use libc::funcs::bsd44::getdtablesize; use io::c; + use std::collections::HashSet; + use std::iter::FromIterator; mod rustrt { extern { @@ -565,16 +573,31 @@ fn spawn_process_os(cfg: ProcessConfig, mem::transmute::>(cfg) }; + // Pre-calculate file descriptors to save/close + // before forking to avoid allocation issues + let mut extra_io: Vec> = extra_fds.iter().map(|opt| match opt { + &Some(ref fd) => Some(fd.fd()), + &None => None, + }).collect(); + + let inherit_fd_lookup: HashSet = FromIterator::from_iter( + extra_fds.iter() + .filter(|opt| opt.is_some()) + .map(|opt| opt.as_ref().unwrap().fd()) + ); + + let fds_to_close: Vec = range(3, unsafe { getdtablesize() }) + .filter(|fd| !inherit_fd_lookup.contains(fd)) + .collect(); + with_envp(cfg.env, proc(envp) { with_argv(cfg.program, cfg.args, proc(argv) unsafe { - let (mut input, mut output) = try!(pipe()); + let (mut input, output) = try!(pipe()); // We may use this in the child, so perform allocations before the // fork let devnull = "/dev/null".to_c_str(); - set_cloexec(output.fd()); - let pid = fork(); if pid < 0 { return Err(super::last_error()) @@ -634,7 +657,15 @@ fn spawn_process_os(cfg: ProcessConfig, // child will either exec() or invoke libc::exit) let _ = libc::close(input.fd()); - fn fail(output: &mut file::FileDesc) -> ! { + // If the child process needs to inherit a large number of + // file descriptors, we might have to overwrite the fd of + // the output pipe back to the parent. Thus we'll duplicate + // its fd around as necessary, but we'll use it directly, + // since we cannot allocate a new FileDesc wrapper at this time. + let mut output_fd = output.fd(); + mem::forget(output); + + fn fail(output_fd: c_int) -> ! { let errno = os::errno(); let bytes = [ (errno << 24) as u8, @@ -642,7 +673,12 @@ fn spawn_process_os(cfg: ProcessConfig, (errno << 8) as u8, (errno << 0) as u8, ]; - assert!(output.inner_write(bytes).is_ok()); + let num = unsafe{ + libc::write(output_fd, + bytes.as_ptr() as *const libc::c_void, + bytes.len() as libc::size_t) + }; + assert_eq!(num, bytes.len() as libc::ssize_t); unsafe { libc::_exit(1) } } @@ -675,21 +711,80 @@ fn spawn_process_os(cfg: ProcessConfig, src != -1 && retry(|| dup2(src, dst)) != -1 }; - if !setup(in_fd, libc::STDIN_FILENO) { fail(&mut output) } - if !setup(out_fd, libc::STDOUT_FILENO) { fail(&mut output) } - if !setup(err_fd, libc::STDERR_FILENO) { fail(&mut output) } + if !setup(in_fd, libc::STDIN_FILENO) { fail(output_fd) } + if !setup(out_fd, libc::STDOUT_FILENO) { fail(output_fd) } + if !setup(err_fd, libc::STDERR_FILENO) { fail(output_fd) } - // close all other fds - for fd in range(3, getdtablesize()).rev() { - if fd != output.fd() { + // close all fds we don't care about + for fd in fds_to_close.into_iter().rev() { + if fd != output_fd { let _ = close(fd as c_int); } } + // Now we must juggle any inherited file descriptors and position + // them in the order requested. + for idx in range(0, extra_io.len()) { + let child_fd = 3 + idx as c_int; // stdio fds (0, 1, 2) already handled + + let backup = match retry(|| libc::dup(child_fd)) { + // Current fd is not used, no need to restore later + -1 if os::errno() as c_int == libc::EBADF => None, + -1 => fail(output_fd), // Unforseen error, fail + fd => { // Current fd is used, "restore" it later + if child_fd == output_fd { + output_fd = fd; + } + + Some(fd) + }, + }; + + match extra_io[idx] { + None => (), + Some(cur_fd) => { + // Copy the fd in the child's next available slot + if retry(|| dup2(cur_fd, child_fd)) == -1 { + fail(output_fd); + } + + // Close the inherited fd unless it's to be duplicated later + let searchable = extra_io.slice(idx, extra_io.len()); + if !searchable.iter().any(|&fd| fd == Some(cur_fd)) { + if retry(|| close(cur_fd)) == -1 { + fail(output_fd); + } + } + }, + } + + // If the current child fd we overwrote was to be inherited, + // replace any references to it with the duplicate fd. + match backup { + None => (), + Some(fd) => { + let mut will_be_inherited = false; + for i in range(idx, extra_io.len()) { + if extra_io[idx] == Some(child_fd) { + *extra_io.get_mut(idx) = Some(fd); + will_be_inherited = true; + } + } + + // Close the saved fd if we won't need it later + if !will_be_inherited { + if retry(|| close(fd)) == -1 { + fail(output_fd); + } + } + }, + } + } + match cfg.gid { Some(u) => { if libc::setgid(u as libc::gid_t) != 0 { - fail(&mut output); + fail(output_fd); } } None => {} @@ -710,7 +805,7 @@ fn spawn_process_os(cfg: ProcessConfig, let _ = setgroups(0, 0 as *const libc::c_void); if libc::setuid(u as libc::uid_t) != 0 { - fail(&mut output); + fail(output_fd); } } None => {} @@ -722,13 +817,14 @@ fn spawn_process_os(cfg: ProcessConfig, let _ = libc::setsid(); } if !dirp.is_null() && chdir(dirp) == -1 { - fail(&mut output); + fail(output_fd); } if !envp.is_null() { set_environ(envp); } + set_cloexec(output_fd); let _ = execvp(*argv, argv as *mut _); - fail(&mut output); + fail(output_fd); }) }) } diff --git a/src/libstd/io/process.rs b/src/libstd/io/process.rs index 83890d2b1275e..f9b14b6f261a9 100644 --- a/src/libstd/io/process.rs +++ b/src/libstd/io/process.rs @@ -1,4 +1,4 @@ -// Copyright 2013 The Rust Project Developers. See the COPYRIGHT +// Copyright 2013-2014 The Rust Project Developers. See the COPYRIGHT // file at the top-level directory of this distribution and at // http://rust-lang.org/COPYRIGHT. // @@ -1097,6 +1097,151 @@ mod tests { assert!(fdes.inner_write("extra write\n".as_bytes()).is_ok()); }) + #[cfg(unix)] + iotest!(fn test_extra_io() { + let (shell, cflag) = if cfg!(windows) { + ("cmd", "/c") + } else { + ("bash", "-c") + }; + + let mut cmd = Command::new(shell); + cmd.arg(cflag) + .arg("echo hello world 1>&3") + .stdin(Ignored) + .stdout(Ignored) + .stderr(Ignored) + .extra_io(CreatePipe(false, true)); + + let mut process = cmd.spawn().ok().expect("command failed to spawn"); + let mut pipe = process.extra_io.pop() + .expect("no pipe entry found in process.extra_io") + .expect("a pipe was not opened to child"); + + let read = pipe.read_to_string().unwrap(); + assert_eq!(read, "hello world\n".to_string()); + }) + + #[cfg(unix)] + iotest!(fn test_extra_io_with_pipe() { + use native::io::process; + + let (shell, cflag) = if cfg!(windows) { + ("cmd", "/c") + } else { + ("bash", "-c") + }; + + let (mut reader, writer) = match process::pipe() { + Ok((r, w)) => (r, w), + Err(_) => fail!("unable to create an OS pipe"), + }; + + let mut cmd = Command::new(shell); + let _ = cmd.arg(cflag) + .arg("echo hello world 1>&3") + .stdin(Ignored) + .stdout(Ignored) + .stderr(Ignored) + .extra_io(InheritFd(writer.fd())); + + let process = cmd.spawn().ok().expect("command failed to spawn"); + assert_eq!(process.extra_io.len(), 1); + assert!(process.extra_io[0].is_none()); + + drop(writer); + let mut buf = [0u8, ..32]; + let output = match reader.inner_read(buf) { + Ok(num_read) => { + let mut string = String::from_utf8(buf.into_vec()).unwrap(); + string.truncate(num_read); + string + }, + Err(..) => fail!("unable to read from pipe"), + }; + + assert_eq!(output.as_slice(), "hello world\n"); + }) + + #[cfg(unix)] + iotest!(fn test_extra_io_with_hole() { + let (shell, cflag) = if cfg!(windows) { + ("cmd", "/c") + } else { + ("bash", "-c") + }; + + let mut cmd = Command::new(shell); + let _ = cmd.arg(cflag) + .arg("echo hello world 1>&4") + .stdin(Ignored) + .stdout(Ignored) + .stderr(Ignored) + .extra_io(Ignored) + .extra_io(CreatePipe(false, true)); + + let mut process = cmd.spawn().ok().expect("command failed to spawn"); + assert_eq!(process.extra_io.len(), 2); + assert!(process.extra_io[0].is_none()); + assert!(process.extra_io[1].is_some()); + + let mut pipe = process.extra_io.pop() + .expect("no pipe entry found in process.extra_io") + .expect("a pipe was not opened to child"); + + let read = pipe.read_to_string().unwrap(); + assert_eq!(read, "hello world\n".to_string()); + }) + + #[cfg(unix)] + iotest!(fn test_extra_io_reorder_fds() { + use native::io::file; + use native::io::process; + + let (shell, cflag) = if cfg!(windows) { + ("cmd", "/c") + } else { + ("bash", "-c") + }; + + let (reader2, writer2) = process::pipe().ok().expect("unable to create OS pipe"); + let (reader3, writer3) = process::pipe().ok().expect("unable to create OS pipe"); + let (reader1, writer1) = process::pipe().ok().expect("unable to create OS pipe"); + + let mut cmd = Command::new(shell); + let _ = cmd.arg(cflag) + .arg("echo one 1>&3; echo two 1>&4; echo three 1>&5") + .stdin(Ignored) + .stdout(Ignored) + .stderr(Ignored) + .extra_io(InheritFd(writer1.fd())) + .extra_io(InheritFd(writer2.fd())) + .extra_io(InheritFd(writer3.fd())); + + let process = cmd.spawn().ok().expect("command failed to spawn"); + assert_eq!(process.extra_io.len(), 3); + + drop(writer1); + drop(writer2); + drop(writer3); + + let read_pipe = |mut reader: file::FileDesc| -> String { + let mut buf = [0, ..32]; + match reader.inner_read(buf) { + Ok(num_read) => { + let mut string = String::from_utf8(buf.into_vec()).unwrap(); + string.truncate(num_read); + string + }, + Err(..) => fail!("unable to read from pipe"), + } + }; + + assert_eq!(read_pipe(reader1).as_slice(), "one\n"); + assert_eq!(read_pipe(reader2).as_slice(), "two\n"); + assert_eq!(read_pipe(reader3).as_slice(), "three\n"); + }) + #[test] #[cfg(windows)] fn env_map_keys_ci() {