Process::send_input_to_fd

3 files changed, 125 insertions, 1 deletions

diff --git a/paludis/util/process.cc b/paludis/util/process.cc
index a45e1b3..93c84de 100644
--- a/paludis/util/process.cc
+++ b/paludis/util/process.cc
@@ -39,6 +39,7 @@
 #include <grp.h>
 #include <pwd.h>
 #include <signal.h>
+#include <fcntl.h>
 #include <sys/types.h>
 #include <sys/wait.h>
 #include <sys/select.h>
@@ -171,6 +172,9 @@ namespace paludis
         std::ostream * capture_output_to_fd;
         std::unique_ptr<Pipe> capture_output_to_fd_pipe;
 
+        std::istream * send_input_to_fd;
+        std::unique_ptr<Pipe> send_input_to_fd_pipe;
+
         ProcessPipeCommandFunction pipe_command_handler;
         std::unique_ptr<Pipe> pipe_command_handler_command_pipe;
         std::unique_ptr<Pipe> pipe_command_handler_response_pipe;
@@ -183,7 +187,8 @@ namespace paludis
             ctl_pipe(true),
             capture_stdout(0),
             capture_stderr(0),
-            capture_output_to_fd(0)
+            capture_output_to_fd(0),
+            send_input_to_fd(0)
         {
         }
 
@@ -197,6 +202,8 @@ void
 RunningProcessThread::thread_func()
 {
     bool prefix_stdout_buffer_has_newline(false), prefix_stderr_buffer_has_newline(false);
+    std::string input_stream_pending;
+
     bool done(false);
     while (! done)
     {
@@ -226,6 +233,12 @@ RunningProcessThread::thread_func()
             max_fd = std::max(max_fd, capture_output_to_fd_pipe->read_fd());
         }
 
+        if (send_input_to_fd)
+        {
+            FD_SET(send_input_to_fd_pipe->write_fd(), &write_fds);
+            max_fd = std::max(max_fd, send_input_to_fd_pipe->write_fd());
+        }
+
         if (pipe_command_handler)
         {
             FD_SET(pipe_command_handler_command_pipe->read_fd(), &read_fds);
@@ -289,6 +302,38 @@ RunningProcessThread::thread_func()
             done_anything = true;
         }
 
+        if (send_input_to_fd && FD_ISSET(send_input_to_fd_pipe->write_fd(), &write_fds))
+        {
+            while ((! input_stream_pending.empty()) || send_input_to_fd->good())
+            {
+                if (input_stream_pending.empty() && send_input_to_fd->good())
+                {
+                    send_input_to_fd->read(buf, sizeof(buf));
+                    input_stream_pending.assign(buf, send_input_to_fd->gcount());
+                }
+
+                int w(::write(send_input_to_fd_pipe->write_fd(), input_stream_pending.data(),
+                            input_stream_pending.length()));
+
+                if (0 == w || (-1 == w && (errno == EAGAIN || errno == EWOULDBLOCK)))
+                    break;
+                else if (-1 == w)
+                    throw ProcessError("write() send_input_to_fd_pipe write_fd failed");
+                else
+                    input_stream_pending.erase(0, w);
+            }
+
+            if (input_stream_pending.empty() && ! send_input_to_fd->good())
+            {
+                if (0 != ::close(send_input_to_fd_pipe->write_fd()))
+                    throw ProcessError("close() send_input_to_fd_pipe write_fd failed");
+                send_input_to_fd_pipe->clear_write_fd();
+                send_input_to_fd = 0;
+            }
+
+            done_anything = true;
+        }
+
         if (pipe_command_handler && FD_ISSET(pipe_command_handler_command_pipe->read_fd(), &read_fds))
         {
             int n(::read(pipe_command_handler_command_pipe->read_fd(), &buf, sizeof(buf)));
@@ -409,6 +454,9 @@ namespace paludis
         std::ostream * capture_output_to_fd_stream;
         int capture_output_to_fd_fd;
         std::string capture_output_to_fd_env_var;
+        std::istream * send_input_to_fd_stream;
+        int send_input_to_fd_fd;
+        std::string send_input_to_fd_env_var;
         ProcessPipeCommandFunction pipe_command_handler;
         std::string pipe_command_handler_env_var;
         int set_stdin_fd;
@@ -431,6 +479,8 @@ namespace paludis
             capture_stderr(0),
             capture_output_to_fd_stream(0),
             capture_output_to_fd_fd(-1),
+            send_input_to_fd_stream(0),
+            send_input_to_fd_fd(-1),
             set_stdin_fd(-1),
             clearenv(false),
             setuid(getuid()),
@@ -505,6 +555,19 @@ Process::run()
             thread->capture_output_to_fd_pipe.reset(new Pipe(true));
         }
 
+        if (_imp->send_input_to_fd_stream)
+        {
+            thread->send_input_to_fd = _imp->send_input_to_fd_stream;
+            thread->send_input_to_fd_pipe.reset(new Pipe(true));
+
+            int arg(::fcntl(thread->send_input_to_fd_pipe->write_fd(), F_GETFL, NULL));
+            if (-1 == arg)
+                throw ProcessError("fcntl(F_GETFL) failed");
+            arg |= O_NONBLOCK;
+            if (-1 == ::fcntl(thread->send_input_to_fd_pipe->write_fd(), F_SETFL, arg))
+                throw ProcessError("fcntl(F_SETFL) failed");
+        }
+
         if (_imp->pipe_command_handler)
         {
             thread->pipe_command_handler = _imp->pipe_command_handler;
@@ -584,6 +647,20 @@ Process::run()
                     ::setenv(_imp->capture_output_to_fd_env_var.c_str(), stringify(fd).c_str(), 1);
             }
 
+            if (thread && thread->send_input_to_fd_pipe)
+            {
+                int fd(_imp->send_input_to_fd_fd);
+                if (-1 == fd)
+                    fd = ::dup(thread->send_input_to_fd_pipe->read_fd());
+                else
+                    fd = ::dup2(thread->send_input_to_fd_pipe->read_fd(), fd);
+
+                if (-1 == fd)
+                    throw ProcessError("dup failed");
+                else if (! _imp->send_input_to_fd_env_var.empty())
+                    ::setenv(_imp->send_input_to_fd_env_var.c_str(), stringify(fd).c_str(), 1);
+            }
+
             if (thread && thread->pipe_command_handler)
             {
                 int read_fd(::dup(thread->pipe_command_handler_response_pipe->read_fd()));
@@ -684,6 +761,16 @@ Process::capture_output_to_fd(std::ostream & s, int fd_or_minus_one, const std::
 }
 
 Process &
+Process::send_input_to_fd(std::istream & s, int fd_or_minus_one, const std::string & env_var_with_fd)
+{
+    _imp->need_thread = true;
+    _imp->send_input_to_fd_stream = &s;
+    _imp->send_input_to_fd_fd = fd_or_minus_one;
+    _imp->send_input_to_fd_env_var = env_var_with_fd;
+    return *this;
+}
+
+Process &
 Process::set_stdin_fd(int f)
 {
     _imp->set_stdin_fd = f;
diff --git a/paludis/util/process.hh b/paludis/util/process.hh
index a1193f1..21e28f7 100644
--- a/paludis/util/process.hh
+++ b/paludis/util/process.hh
@@ -88,6 +88,7 @@ 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 & send_input_to_fd(std::istream &, int fd_or_minus_one, const std::string & env_var_with_fd);
             Process & set_stdin_fd(int);
             Process & pipe_command_handler(const std::string &, const ProcessPipeCommandFunction &);
 
diff --git a/paludis/util/process_TEST.cc b/paludis/util/process_TEST.cc
index be63e49..89ca85a 100644
--- a/paludis/util/process_TEST.cc
+++ b/paludis/util/process_TEST.cc
@@ -413,5 +413,41 @@ namespace test_cases
             TEST_CHECK_EQUAL(stdout_stream.str(), "");
         }
     } test_clearenv;
+
+    struct SendFDTest : TestCase
+    {
+        SendFDTest() : TestCase("send fd") { }
+
+        void run()
+        {
+            std::stringstream stdout_stream, in_stream;
+            in_stream << "monkey" << std::endl;
+
+            Process cat_process(ProcessCommand({"sh", "-c", "cat <&$MAGIC_FD"}));
+            cat_process.send_input_to_fd(in_stream, -1, "MAGIC_FD");
+            cat_process.capture_stdout(stdout_stream);
+
+            TEST_CHECK_EQUAL(cat_process.run().wait(), 0);
+            TEST_CHECK_EQUAL(stdout_stream.str(), "monkey\n");
+        }
+    } test_send_fd;
+
+    struct SendFDFixedTest : TestCase
+    {
+        SendFDFixedTest() : TestCase("send fd fixed") { }
+
+        void run()
+        {
+            std::stringstream stdout_stream, in_stream;
+            in_stream << "monkey" << std::endl;
+
+            Process cat_process(ProcessCommand({"sh", "-c", "cat <&5"}));
+            cat_process.send_input_to_fd(in_stream, 5, "");
+            cat_process.capture_stdout(stdout_stream);
+
+            TEST_CHECK_EQUAL(cat_process.run().wait(), 0);
+            TEST_CHECK_EQUAL(stdout_stream.str(), "monkey\n");
+        }
+    } test_send_fd_fixed;
 }