source/externaltask.cpp

   1 #include <cstdlib>
   2 #include <unistd.h>
   3 #include <sys/wait.h>
   4 #include <msp/fs/dir.h>
   5 #include <msp/io/console.h>
   6 #include <msp/io/file.h>
   7 #include <msp/io/print.h>
   8 #include <msp/time/timedelta.h>
   9 #include "externaltask.h"
  10
  11 using namespace std;
  12 using namespace Msp;
  13
  14 ExternalTask::ExternalTask(const Arguments &a, const FS::Path &wd):
  15         argv(a),
  16         work_dir(wd),
  17         process(0),
  18         exit_code(-1),
  19         stdin_action(PASSTHROUGH),
  20         stdout_action(PASSTHROUGH),
  21         stderr_action(PASSTHROUGH),
  22         capture_pipe(0)
  23 {
  24         if(argv.empty())
  25                 throw invalid_argument("ExternalTask::ExternalTask");
  26 }
  27
  28 ExternalTask::~ExternalTask()
  29 {
  30         delete capture_pipe;
  31 }
  32
  33 string ExternalTask::get_command() const
  34 {
  35         string cmd;
  36         for(vector<string>::const_iterator i=argv.begin(); i!=argv.end(); ++i)
  37         {
  38                 if(i!=argv.begin())
  39                         cmd += ' ';
  40
  41                 for(string::const_iterator j=i->begin(); j!=i->end(); ++j)
  42                 {
  43                         if(*j=='"' || *j=='\'' || *j==' ' || *j=='\\' || *j=='&')
  44                                 cmd += '\\';
  45                         cmd += *j;
  46                 }
  47         }
  48
  49         if(stdin_action==REDIRECT)
  50         {
  51                 cmd += " <";
  52                 cmd += stdin_file.str();
  53         }
  54
  55         if(stdout_action==REDIRECT)
  56         {
  57                 cmd += " >";
  58                 cmd += stdout_file.str();
  59         }
  60
  61         return cmd;
  62 }
  63
  64 void ExternalTask::start()
  65 {
  66         IO::File *devnull = 0;
  67         IO::File *infile = 0;
  68         IO::File *outfile = 0;
  69
  70         prepare();
  71
  72         process = new Process;
  73
  74         if(stdin_action==IGNORE || stdout_action==IGNORE || stderr_action==IGNORE)
  75         {
  76                 devnull = new IO::File("/dev/null", IO::M_RDWR);
  77                 if(stdin_action==IGNORE)
  78                         process->redirect_cin(*devnull);
  79                 if(stdout_action==IGNORE)
  80                         process->redirect_cout(*devnull);
  81                 if(stderr_action==IGNORE)
  82                         process->redirect_cerr(*devnull);
  83         }
  84
  85         if(stdout_action==REDIRECT)
  86         {
  87                 outfile = new IO::File((work_dir/stdout_file).str(), IO::M_WRITE);
  88                 process->redirect_cout(*outfile);
  89         }
  90
  91         if(stdout_action==CAPTURE || stderr_action==CAPTURE)
  92         {
  93                 capture_pipe = new IO::Pipe;
  94                 if(stdout_action==CAPTURE)
  95                         process->redirect_cout(*capture_pipe);
  96                 if(stderr_action==CAPTURE)
  97                         process->redirect_cerr(*capture_pipe);
  98         }
  99
 100         if(stdin_action==REDIRECT)
 101         {
 102                 infile = new IO::File((work_dir/stdin_file).str());
 103                 process->redirect_cin(*infile);
 104         }
 105
 106         if(!work_dir.empty())
 107                 process->set_working_directory(work_dir);
 108
 109         Process::Arguments args(argv.begin()+1, argv.end());
 110         process->execute(argv.front(), args);
 111         if(capture_pipe)
 112                 capture_pipe->set_mode(IO::M_READ);
 113
 114         delete devnull;
 115         delete infile;
 116         delete outfile;
 117 }
 118
 119 Task::Status ExternalTask::check()
 120 {
 121         return do_wait(false);
 122 }
 123
 124 Task::Status ExternalTask::wait()
 125 {
 126         return do_wait(true);
 127 }
 128
 129 Task::Status ExternalTask::do_wait(bool block)
 130 {
 131         while(process)
 132         {
 133                 if(process->wait(!capture_pipe))
 134                 {
 135                         exit_code = process->get_exit_code();
 136                         delete process;
 137                         process = 0;
 138                 }
 139
 140                 // Do this after waiting to avoid a race condition
 141                 while(capture_pipe && IO::poll(*capture_pipe, IO::P_INPUT, 10*Time::msec))
 142                 {
 143                         char buf[1024];
 144                         unsigned len = capture_pipe->read(buf, sizeof(buf));
 145                         if(len)
 146                                 output.append(buf, len);
 147                         else
 148                                 break;
 149                 }
 150
 151                 if(process)
 152                 {
 153                         if(!block)
 154                                 return RUNNING;
 155                 }
 156                 else
 157                         signal_finished.emit(!exit_code);
 158         }
 159
 160         return exit_code ? ERROR : SUCCESS;
 161 }
 162
 163 void ExternalTask::set_stdin(const FS::Path &f)
 164 {
 165         stdin_action = REDIRECT;
 166         stdin_file = f;
 167 }
 168
 169 void ExternalTask::set_stdout(StreamAction a)
 170 {
 171         if(a==REDIRECT)
 172                 throw invalid_argument("ExternalTask::set_stdout");
 173         stdout_action = a;
 174 }
 175
 176 void ExternalTask::set_stdout(const FS::Path &f)
 177 {
 178         stdout_action = REDIRECT;
 179         stdout_file = f;
 180 }
 181
 182 void ExternalTask::set_stderr(StreamAction a)
 183 {
 184         if(a==REDIRECT)
 185                 throw invalid_argument("ExternalTask::set_stdout");
 186         stderr_action = a;
 187 }
 188
 189 string ExternalTask::run_and_capture_output(const Arguments &argv, const FS::Path &wd, bool capture_stderr)
 190 {
 191         ExternalTask task(argv, wd);
 192         task.stdin_action = IGNORE;
 193         task.set_stdout(CAPTURE);
 194         task.set_stderr(capture_stderr ? CAPTURE : IGNORE);
 195         task.start();
 196         if(task.wait()!=SUCCESS)
 197                 throw runtime_error(format("%s failed", argv.front()));
 198         return task.get_output();
 199 }