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