From 6c8ec28a8dd13cdb76f2a367bd54ca1302894bec Mon Sep 17 00:00:00 2001 From: Ciaran McCreesh Date: Fri, 20 Aug 2010 14:56:23 +0100 Subject: Process::capture_output_to_fd --- paludis/util/process.cc | 85 ++++++++++++++++++++++++++++++++++++++++++-- paludis/util/process.hh | 4 +++ paludis/util/process_TEST.cc | 30 ++++++++++++++++ 3 files changed, 117 insertions(+), 2 deletions(-) diff --git a/paludis/util/process.cc b/paludis/util/process.cc index ced2c2b89..abcef4690 100644 --- a/paludis/util/process.cc +++ b/paludis/util/process.cc @@ -96,6 +96,20 @@ ProcessCommand::exec() throw ProcessError("execvp failed"); } +void +ProcessCommand::echo_command_to(std::ostream & s) +{ + for (auto v_begin(_imp->args.begin()), v(v_begin), v_end(_imp->args.end()) ; + v != v_end ; ++v) + { + if (v != v_begin) + s << " "; + s << *v; + } + + s << std::endl; +} + namespace paludis { struct RunningProcessThread @@ -108,13 +122,17 @@ namespace paludis std::ostream * capture_stderr; std::unique_ptr capture_stderr_pipe; + std::ostream * capture_output_to_fd; + std::unique_ptr capture_output_to_fd_pipe; + /* must be last, so the thread gets join()ed before its FDs vanish */ std::unique_ptr thread; RunningProcessThread() : ctl_pipe(true), capture_stdout(0), - capture_stderr(0) + capture_stderr(0), + capture_output_to_fd(0) { } @@ -150,6 +168,12 @@ RunningProcessThread::thread_func() max_fd = std::max(max_fd, capture_stderr_pipe->read_fd()); } + if (capture_output_to_fd) + { + FD_SET(capture_output_to_fd_pipe->read_fd(), &read_fds); + max_fd = std::max(max_fd, capture_output_to_fd_pipe->read_fd()); + } + int retval(::pselect(max_fd + 1, &read_fds, &write_fds, 0, 0, 0)); if (-1 == retval) throw ProcessError("pselect() failed"); @@ -177,6 +201,16 @@ RunningProcessThread::thread_func() done_anything = true; } + if (capture_output_to_fd_pipe && FD_ISSET(capture_output_to_fd_pipe->read_fd(), &read_fds)) + { + int n(::read(capture_output_to_fd_pipe->read_fd(), &buf, sizeof(buf))); + if (-1 == n) + throw ProcessError("read() capture_output_to_fd_pipe read_fd failed"); + else if (0 != n) + capture_output_to_fd->write(buf, n); + done_anything = true; + } + if (done_anything) continue; @@ -210,11 +244,15 @@ namespace paludis bool use_ptys; std::ostream * capture_stdout; std::ostream * capture_stderr; + std::ostream * capture_output_to_fd_stream; + int capture_output_to_fd_fd; + std::string capture_output_to_fd_env_var; std::map setenvs; std::string chdir; uid_t setuid; gid_t setgid; + std::ostream * echo_command_to; Imp(ProcessCommand && c) : command(std::move(c)), @@ -222,8 +260,11 @@ namespace paludis use_ptys(false), capture_stdout(0), capture_stderr(0), + capture_output_to_fd_stream(0), + capture_output_to_fd_fd(-1), setuid(getuid()), - setgid(getgid()) + setgid(getgid()), + echo_command_to(0) { } }; @@ -241,6 +282,9 @@ Process::~Process() RunningProcessHandle Process::run() { + if (_imp->echo_command_to) + _imp->command.echo_command_to(*_imp->echo_command_to); + std::unique_ptr thread; if (_imp->need_thread) { @@ -263,6 +307,12 @@ Process::run() else thread->capture_stderr_pipe.reset(new Pipe(true)); } + + if (_imp->capture_output_to_fd_stream) + { + thread->capture_output_to_fd = _imp->capture_output_to_fd_stream; + thread->capture_output_to_fd_pipe.reset(new Pipe(true)); + } } pid_t child(fork()); @@ -284,6 +334,20 @@ Process::run() throw ProcessError("dup2() failed"); } + if (thread && thread->capture_output_to_fd_pipe) + { + int fd(_imp->capture_output_to_fd_fd); + if (-1 == fd) + fd = ::dup(thread->capture_output_to_fd_pipe->write_fd()); + else + fd = ::dup2(thread->capture_output_to_fd_pipe->write_fd(), fd); + + if (-1 == fd) + throw ProcessError("dup failed"); + else if (! _imp->capture_output_to_fd_env_var.empty()) + ::setenv(_imp->capture_output_to_fd_env_var.c_str(), stringify(fd).c_str(), 1); + } + for (auto m(_imp->setenvs.begin()), m_end(_imp->setenvs.end()) ; m != m_end ; ++m) ::setenv(m->first.c_str(), m->second.c_str(), 1); @@ -346,6 +410,16 @@ Process::capture_stderr(std::ostream & s) return *this; } +Process & +Process::capture_output_to_fd(std::ostream & s, int fd_or_minus_one, const std::string & env_var_with_fd) +{ + _imp->need_thread = true; + _imp->capture_output_to_fd_stream = &s; + _imp->capture_output_to_fd_fd = fd_or_minus_one; + _imp->capture_output_to_fd_env_var = env_var_with_fd; + return *this; +} + Process & Process::setenv(const std::string & a, const std::string & b) { @@ -375,6 +449,13 @@ Process::setuid_setgid(uid_t u, gid_t g) return *this; } +Process & +Process::echo_command_to(std::ostream & s) +{ + _imp->echo_command_to = &s; + return *this; +} + namespace paludis { template <> diff --git a/paludis/util/process.hh b/paludis/util/process.hh index 0cc62e82c..3865911a5 100644 --- a/paludis/util/process.hh +++ b/paludis/util/process.hh @@ -58,6 +58,8 @@ namespace paludis ProcessCommand(const ProcessCommand &) = delete; ProcessCommand & operator= (const ProcessCommand &) = delete; + void echo_command_to(std::ostream &); + void exec() PALUDIS_ATTRIBUTE((noreturn)); }; @@ -75,10 +77,12 @@ namespace paludis Process & capture_stdout(std::ostream &); Process & capture_stderr(std::ostream &); + Process & capture_output_to_fd(std::ostream &, int fd_or_minus_one, const std::string & env_var_with_fd); Process & setenv(const std::string &, const std::string &); Process & chdir(const FSEntry &); Process & use_ptys(); Process & setuid_setgid(uid_t, gid_t); + Process & echo_command_to(std::ostream &); }; class PALUDIS_VISIBLE RunningProcessHandle : diff --git a/paludis/util/process_TEST.cc b/paludis/util/process_TEST.cc index a7361f3c0..fc57c197a 100644 --- a/paludis/util/process_TEST.cc +++ b/paludis/util/process_TEST.cc @@ -232,5 +232,35 @@ namespace test_cases } } } test_setuid; + + struct GrabFDTest : TestCase + { + GrabFDTest() : TestCase("grab fd") { } + + void run() + { + std::stringstream fd_stream; + Process echo_process(ProcessCommand({"sh", "-c", "echo monkey 1>&$MAGIC_FD"})); + echo_process.capture_output_to_fd(fd_stream, -1, "MAGIC_FD"); + + TEST_CHECK_EQUAL(echo_process.run().wait(), 0); + TEST_CHECK_EQUAL(fd_stream.str(), "monkey\n"); + } + } test_grab_fd; + + struct GrabFDFixedTest : TestCase + { + GrabFDFixedTest() : TestCase("grab fd fixed") { } + + void run() + { + std::stringstream fd_stream; + Process echo_process(ProcessCommand({"sh", "-c", "echo monkey 1>&5"})); + echo_process.capture_output_to_fd(fd_stream, 5, ""); + + TEST_CHECK_EQUAL(echo_process.run().wait(), 0); + TEST_CHECK_EQUAL(fd_stream.str(), "monkey\n"); + } + } test_grab_fd_fixed; } -- cgit v1.2.1