graduate IOBuf out of folly/experimental
[folly.git] / folly / Subprocess.cpp
1 /*
2  * Copyright 2013 Facebook, Inc.
3  *
4  * Licensed under the Apache License, Version 2.0 (the "License");
5  * you may not use this file except in compliance with the License.
6  * You may obtain a copy of the License at
7  *
8  *   http://www.apache.org/licenses/LICENSE-2.0
9  *
10  * Unless required by applicable law or agreed to in writing, software
11  * distributed under the License is distributed on an "AS IS" BASIS,
12  * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13  * See the License for the specific language governing permissions and
14  * limitations under the License.
15  */
16
17 #include "folly/Subprocess.h"
18
19 #include <fcntl.h>
20 #include <poll.h>
21 #include <unistd.h>
22 #include <wait.h>
23
24 #include <array>
25 #include <algorithm>
26 #include <system_error>
27
28 #include <boost/container/flat_set.hpp>
29 #include <boost/range/adaptors.hpp>
30
31 #include <glog/logging.h>
32
33 #include "folly/Conv.h"
34 #include "folly/ScopeGuard.h"
35 #include "folly/String.h"
36 #include "folly/io/Cursor.h"
37
38 extern char** environ;
39
40 namespace folly {
41
42 ProcessReturnCode::State ProcessReturnCode::state() const {
43   if (rawStatus_ == RV_NOT_STARTED) return NOT_STARTED;
44   if (rawStatus_ == RV_RUNNING) return RUNNING;
45   if (WIFEXITED(rawStatus_)) return EXITED;
46   if (WIFSIGNALED(rawStatus_)) return KILLED;
47   throw std::runtime_error(to<std::string>(
48       "Invalid ProcessReturnCode: ", rawStatus_));
49 }
50
51 void ProcessReturnCode::enforce(State s) const {
52   if (state() != s) {
53     throw std::logic_error(to<std::string>("Invalid state ", s));
54   }
55 }
56
57 int ProcessReturnCode::exitStatus() const {
58   enforce(EXITED);
59   return WEXITSTATUS(rawStatus_);
60 }
61
62 int ProcessReturnCode::killSignal() const {
63   enforce(KILLED);
64   return WTERMSIG(rawStatus_);
65 }
66
67 bool ProcessReturnCode::coreDumped() const {
68   enforce(KILLED);
69   return WCOREDUMP(rawStatus_);
70 }
71
72 std::string ProcessReturnCode::str() const {
73   switch (state()) {
74   case NOT_STARTED:
75     return "not started";
76   case RUNNING:
77     return "running";
78   case EXITED:
79     return to<std::string>("exited with status ", exitStatus());
80   case KILLED:
81     return to<std::string>("killed by signal ", killSignal(),
82                            (coreDumped() ? " (core dumped)" : ""));
83   }
84   CHECK(false);  // unreached
85 }
86
87 CalledProcessError::CalledProcessError(ProcessReturnCode rc)
88   : returnCode_(rc),
89     what_(returnCode_.str()) {
90 }
91
92 namespace {
93
94 // Copy pointers to the given strings in a format suitable for posix_spawn
95 std::unique_ptr<const char*[]> cloneStrings(const std::vector<std::string>& s) {
96   std::unique_ptr<const char*[]> d(new const char*[s.size() + 1]);
97   for (int i = 0; i < s.size(); i++) {
98     d[i] = s[i].c_str();
99   }
100   d[s.size()] = nullptr;
101   return d;
102 }
103
104 // Helper to throw std::system_error
105 void throwSystemError(int err, const char* msg) __attribute__((noreturn));
106 void throwSystemError(int err, const char* msg) {
107   throw std::system_error(err, std::system_category(), msg);
108 }
109
110 // Helper to throw std::system_error from errno
111 void throwSystemError(const char* msg) __attribute__((noreturn));
112 void throwSystemError(const char* msg) {
113   throwSystemError(errno, msg);
114 }
115
116 // Check a Posix return code (0 on success, error number on error), throw
117 // on error.
118 void checkPosixError(int err, const char* msg) {
119   if (err != 0) {
120     throwSystemError(err, msg);
121   }
122 }
123
124 // Check a traditional Uinx return code (-1 and sets errno on error), throw
125 // on error.
126 void checkUnixError(ssize_t ret, const char* msg) {
127   if (ret == -1) {
128     throwSystemError(msg);
129   }
130 }
131 void checkUnixError(ssize_t ret, int savedErrno, const char* msg) {
132   if (ret == -1) {
133     throwSystemError(savedErrno, msg);
134   }
135 }
136
137 // Check a wait() status, throw on non-successful
138 void checkStatus(ProcessReturnCode returnCode) {
139   if (returnCode.state() != ProcessReturnCode::EXITED ||
140       returnCode.exitStatus() != 0) {
141     throw CalledProcessError(returnCode);
142   }
143 }
144
145 }  // namespace
146
147 Subprocess::Options& Subprocess::Options::fd(int fd, int action) {
148   if (action == Subprocess::PIPE) {
149     if (fd == 0) {
150       action = Subprocess::PIPE_IN;
151     } else if (fd == 1 || fd == 2) {
152       action = Subprocess::PIPE_OUT;
153     } else {
154       throw std::invalid_argument(
155           to<std::string>("Only fds 0, 1, 2 are valid for action=PIPE: ", fd));
156     }
157   }
158   fdActions_[fd] = action;
159   return *this;
160 }
161
162 Subprocess::Subprocess(
163     const std::vector<std::string>& argv,
164     const Options& options,
165     const char* executable,
166     const std::vector<std::string>* env)
167   : pid_(-1),
168     returnCode_(RV_NOT_STARTED) {
169   if (argv.empty()) {
170     throw std::invalid_argument("argv must not be empty");
171   }
172   if (!executable) executable = argv[0].c_str();
173   spawn(cloneStrings(argv), executable, options, env);
174 }
175
176 Subprocess::Subprocess(
177     const std::string& cmd,
178     const Options& options,
179     const std::vector<std::string>* env)
180   : pid_(-1),
181     returnCode_(RV_NOT_STARTED) {
182   if (options.usePath_) {
183     throw std::invalid_argument("usePath() not allowed when running in shell");
184   }
185   const char* shell = getenv("SHELL");
186   if (!shell) {
187     shell = "/bin/sh";
188   }
189
190   std::unique_ptr<const char*[]> argv(new const char*[4]);
191   argv[0] = shell;
192   argv[1] = "-c";
193   argv[2] = cmd.c_str();
194   argv[3] = nullptr;
195   spawn(std::move(argv), shell, options, env);
196 }
197
198 Subprocess::~Subprocess() {
199   CHECK_NE(returnCode_.state(), ProcessReturnCode::RUNNING)
200     << "Subprocess destroyed without reaping child";
201 }
202
203 namespace {
204 void closeChecked(int fd) {
205   checkUnixError(::close(fd), "close");
206 }
207 }  // namespace
208
209 void Subprocess::closeAll() {
210   for (auto& p : pipes_) {
211     closeChecked(p.parentFd);
212   }
213   pipes_.clear();
214 }
215
216 void Subprocess::setAllNonBlocking() {
217   for (auto& p : pipes_) {
218     int fd = p.parentFd;
219     int flags = ::fcntl(fd, F_GETFL);
220     checkUnixError(flags, "fcntl");
221     int r = ::fcntl(fd, F_SETFL, flags | O_NONBLOCK);
222     checkUnixError(r, "fcntl");
223   }
224 }
225
226 void Subprocess::spawn(
227     std::unique_ptr<const char*[]> argv,
228     const char* executable,
229     const Options& optionsIn,
230     const std::vector<std::string>* env) {
231   if (optionsIn.usePath_ && env) {
232     throw std::invalid_argument(
233         "usePath() not allowed when overriding environment");
234   }
235
236   // Make a copy, we'll mutate options
237   Options options(optionsIn);
238
239   // Parent work, pre-fork: create pipes
240   std::vector<int> childFds;
241   for (auto& p : options.fdActions_) {
242     if (p.second == PIPE_IN || p.second == PIPE_OUT) {
243       int fds[2];
244       int r = ::pipe(fds);
245       checkUnixError(r, "pipe");
246       PipeInfo pinfo;
247       pinfo.direction = p.second;
248       int cfd;
249       if (p.second == PIPE_IN) {
250         // Child gets reading end
251         pinfo.parentFd = fds[1];
252         cfd = fds[0];
253       } else {
254         pinfo.parentFd = fds[0];
255         cfd = fds[1];
256       }
257       p.second = cfd;  // ensure it gets dup2()ed
258       pinfo.childFd = p.first;
259       childFds.push_back(cfd);
260       pipes_.push_back(pinfo);
261     }
262   }
263
264   // This should already be sorted, as options.fdActions_ is
265   DCHECK(std::is_sorted(pipes_.begin(), pipes_.end()));
266
267   // Note that the const casts below are legit, per
268   // http://pubs.opengroup.org/onlinepubs/009695399/functions/exec.html
269
270   char** argVec = const_cast<char**>(argv.get());
271
272   // Set up environment
273   std::unique_ptr<const char*[]> envHolder;
274   char** envVec;
275   if (env) {
276     envHolder = cloneStrings(*env);
277     envVec = const_cast<char**>(envHolder.get());
278   } else {
279     envVec = environ;
280   }
281
282   // Block all signals around vfork; see http://ewontfix.com/7/.
283   //
284   // As the child may run in the same address space as the parent until
285   // the actual execve() system call, any (custom) signal handlers that
286   // the parent has might alter parent's memory if invoked in the child,
287   // with undefined results.  So we block all signals in the parent before
288   // vfork(), which will cause them to be blocked in the child as well (we
289   // rely on the fact that Linux, just like all sane implementations, only
290   // clones the calling thread).  Then, in the child, we reset all signals
291   // to their default dispositions (while still blocked), and unblock them
292   // (so the exec()ed process inherits the parent's signal mask)
293   //
294   // The parent also unblocks all signals as soon as vfork() returns.
295   sigset_t allBlocked;
296   int r = ::sigfillset(&allBlocked);
297   checkUnixError(r, "sigfillset");
298   sigset_t oldSignals;
299   r = pthread_sigmask(SIG_SETMASK, &allBlocked, &oldSignals);
300   checkPosixError(r, "pthread_sigmask");
301
302   pid_t pid = vfork();
303   if (pid == 0) {
304     // While all signals are blocked, we must reset their
305     // dispositions to default.
306     for (int sig = 1; sig < NSIG; ++sig) {
307       ::signal(sig, SIG_DFL);
308     }
309     // Unblock signals; restore signal mask.
310     int r = pthread_sigmask(SIG_SETMASK, &oldSignals, nullptr);
311     if (r != 0) abort();
312
313     runChild(executable, argVec, envVec, options);
314     // This should never return, but there's nothing else we can do here.
315     abort();
316   }
317   // In parent.  We want to restore the signal mask even if vfork fails,
318   // so we'll save errno here, restore the signal mask, and only then
319   // throw.
320   int savedErrno = errno;
321
322   // Restore signal mask; do this even if vfork fails!
323   // We only check for errors from pthread_sigmask after we recorded state
324   // that the child is alive, so we know to reap it.
325   r = pthread_sigmask(SIG_SETMASK, &oldSignals, nullptr);
326   checkUnixError(pid, savedErrno, "vfork");
327
328   // Child is alive
329   pid_ = pid;
330   returnCode_ = ProcessReturnCode(RV_RUNNING);
331
332   // Parent work, post-fork: close child's ends of pipes
333   for (int f : childFds) {
334     closeChecked(f);
335   }
336
337   checkPosixError(r, "pthread_sigmask");
338 }
339
340 namespace {
341
342 // Checked version of close() to use in the child: abort() on error
343 void childClose(int fd) {
344   int r = ::close(fd);
345   if (r == -1) abort();
346 }
347
348 // Checked version of dup2() to use in the child: abort() on error
349 void childDup2(int oldfd, int newfd) {
350   int r = ::dup2(oldfd, newfd);
351   if (r == -1) abort();
352 }
353
354 }  // namespace
355
356 void Subprocess::runChild(const char* executable,
357                           char** argv, char** env,
358                           const Options& options) const {
359   // Close parent's ends of all pipes
360   for (auto& p : pipes_) {
361     childClose(p.parentFd);
362   }
363
364   // Close all fds that we're supposed to close.
365   // Note that we're ignoring errors here, in case some of these
366   // fds were set to close on exec.
367   for (auto& p : options.fdActions_) {
368     if (p.second == CLOSE) {
369       ::close(p.first);
370     } else {
371       childDup2(p.second, p.first);
372     }
373   }
374
375   // If requested, close all other file descriptors.  Don't close
376   // any fds in options.fdActions_, and don't touch stdin, stdout, stderr.
377   // Ignore errors.
378   if (options.closeOtherFds_) {
379     for (int fd = getdtablesize() - 1; fd >= 3; --fd) {
380       if (options.fdActions_.count(fd) == 0) {
381         ::close(fd);
382       }
383     }
384   }
385
386   // Now, finally, exec.
387   int r;
388   if (options.usePath_) {
389     ::execvp(executable, argv);
390   } else {
391     ::execve(executable, argv, env);
392   }
393
394   // If we're here, something's wrong.
395   abort();
396 }
397
398 ProcessReturnCode Subprocess::poll() {
399   returnCode_.enforce(ProcessReturnCode::RUNNING);
400   DCHECK_GT(pid_, 0);
401   int status;
402   pid_t found = ::waitpid(pid_, &status, WNOHANG);
403   checkUnixError(found, "waitpid");
404   if (found != 0) {
405     returnCode_ = ProcessReturnCode(status);
406     pid_ = -1;
407   }
408   return returnCode_;
409 }
410
411 bool Subprocess::pollChecked() {
412   if (poll().state() == ProcessReturnCode::RUNNING) {
413     return false;
414   }
415   checkStatus(returnCode_);
416   return true;
417 }
418
419 ProcessReturnCode Subprocess::wait() {
420   returnCode_.enforce(ProcessReturnCode::RUNNING);
421   DCHECK_GT(pid_, 0);
422   int status;
423   pid_t found;
424   do {
425     found = ::waitpid(pid_, &status, 0);
426   } while (found == -1 && errno == EINTR);
427   checkUnixError(found, "waitpid");
428   DCHECK_EQ(found, pid_);
429   returnCode_ = ProcessReturnCode(status);
430   return returnCode_;
431 }
432
433 void Subprocess::waitChecked() {
434   wait();
435   checkStatus(returnCode_);
436 }
437
438 void Subprocess::sendSignal(int signal) {
439   returnCode_.enforce(ProcessReturnCode::RUNNING);
440   int r = ::kill(pid_, signal);
441   checkUnixError(r, "kill");
442 }
443
444 namespace {
445 void setNonBlocking(int fd) {
446   int flags = ::fcntl(fd, F_GETFL);
447   checkUnixError(flags, "fcntl");
448   int r = ::fcntl(fd, F_SETFL, flags | O_NONBLOCK);
449   checkUnixError(r, "fcntl");
450 }
451
452 std::pair<const uint8_t*, size_t> queueFront(const IOBufQueue& queue) {
453   auto* p = queue.front();
454   if (!p) return std::make_pair(nullptr, 0);
455   return io::Cursor(p).peek();
456 }
457
458 // fd write
459 bool handleWrite(int fd, IOBufQueue& queue) {
460   for (;;) {
461     auto p = queueFront(queue);
462     if (p.second == 0) {
463       return true;  // EOF
464     }
465
466     ssize_t n;
467     do {
468       n = ::write(fd, p.first, p.second);
469     } while (n == -1 && errno == EINTR);
470     if (n == -1 && errno == EAGAIN) {
471       return false;
472     }
473     checkUnixError(n, "write");
474     queue.trimStart(n);
475   }
476 }
477
478 // fd read
479 bool handleRead(int fd, IOBufQueue& queue) {
480   for (;;) {
481     auto p = queue.preallocate(100, 65000);
482     ssize_t n;
483     do {
484       n = ::read(fd, p.first, p.second);
485     } while (n == -1 && errno == EINTR);
486     if (n == -1 && errno == EAGAIN) {
487       return false;
488     }
489     checkUnixError(n, "read");
490     if (n == 0) {
491       return true;
492     }
493     queue.postallocate(n);
494   }
495 }
496
497 bool discardRead(int fd) {
498   static const size_t bufSize = 65000;
499   // Thread unsafe, but it doesn't matter.
500   static std::unique_ptr<char[]> buf(new char[bufSize]);
501
502   for (;;) {
503     ssize_t n;
504     do {
505       n = ::read(fd, buf.get(), bufSize);
506     } while (n == -1 && errno == EINTR);
507     if (n == -1 && errno == EAGAIN) {
508       return false;
509     }
510     checkUnixError(n, "read");
511     if (n == 0) {
512       return true;
513     }
514   }
515 }
516
517 }  // namespace
518
519 std::pair<std::string, std::string> Subprocess::communicate(
520     const CommunicateFlags& flags,
521     StringPiece data) {
522   IOBufQueue dataQueue;
523   dataQueue.wrapBuffer(data.data(), data.size());
524
525   auto outQueues = communicateIOBuf(flags, std::move(dataQueue));
526   auto outBufs = std::make_pair(outQueues.first.move(),
527                                 outQueues.second.move());
528   std::pair<std::string, std::string> out;
529   if (outBufs.first) {
530     outBufs.first->coalesce();
531     out.first.assign(reinterpret_cast<const char*>(outBufs.first->data()),
532                      outBufs.first->length());
533   }
534   if (outBufs.second) {
535     outBufs.second->coalesce();
536     out.second.assign(reinterpret_cast<const char*>(outBufs.second->data()),
537                      outBufs.second->length());
538   }
539   return out;
540 }
541
542 std::pair<IOBufQueue, IOBufQueue> Subprocess::communicateIOBuf(
543     const CommunicateFlags& flags,
544     IOBufQueue data) {
545   std::pair<IOBufQueue, IOBufQueue> out;
546
547   auto readCallback = [&] (int pfd, int cfd) {
548     if (cfd == 1 && flags.readStdout_) {
549       return handleRead(pfd, out.first);
550     } else if (cfd == 2 && flags.readStderr_) {
551       return handleRead(pfd, out.second);
552     } else {
553       // Don't close the file descriptor, the child might not like SIGPIPE,
554       // just read and throw the data away.
555       return discardRead(pfd);
556     }
557   };
558
559   auto writeCallback = [&] (int pfd, int cfd) {
560     if (cfd == 0 && flags.writeStdin_) {
561       return handleWrite(pfd, data);
562     } else {
563       // If we don't want to write to this fd, just close it.
564       return false;
565     }
566   };
567
568   communicate(std::move(readCallback), std::move(writeCallback));
569
570   return out;
571 }
572
573 void Subprocess::communicate(FdCallback readCallback,
574                              FdCallback writeCallback) {
575   returnCode_.enforce(ProcessReturnCode::RUNNING);
576   setAllNonBlocking();
577
578   std::vector<pollfd> fds;
579   fds.reserve(pipes_.size());
580   std::vector<int> toClose;
581   toClose.reserve(pipes_.size());
582
583   while (!pipes_.empty()) {
584     fds.clear();
585     toClose.clear();
586
587     for (auto& p : pipes_) {
588       pollfd pfd;
589       pfd.fd = p.parentFd;
590       // Yes, backwards, PIPE_IN / PIPE_OUT are defined from the
591       // child's point of view.
592       pfd.events = (p.direction == PIPE_IN ?  POLLOUT : POLLIN);
593       fds.push_back(pfd);
594     }
595
596     int r;
597     do {
598       r = ::poll(fds.data(), fds.size(), -1);
599     } while (r == -1 && errno == EINTR);
600     checkUnixError(r, "poll");
601
602     for (int i = 0; i < pipes_.size(); ++i) {
603       auto& p = pipes_[i];
604       DCHECK_EQ(fds[i].fd, p.parentFd);
605       short events = fds[i].revents;
606
607       bool closed = false;
608       if (events & POLLOUT) {
609         DCHECK(!(events & POLLIN));
610         if (writeCallback(p.parentFd, p.childFd)) {
611           toClose.push_back(i);
612           closed = true;
613         }
614       }
615
616       if (events & POLLIN) {
617         DCHECK(!(events & POLLOUT));
618         if (readCallback(p.parentFd, p.childFd)) {
619           toClose.push_back(i);
620           closed = true;
621         }
622       }
623
624       if ((events & (POLLHUP | POLLERR)) && !closed) {
625         toClose.push_back(i);
626         closed = true;
627       }
628     }
629
630     // Close the fds in reverse order so the indexes hold after erase()
631     for (int idx : boost::adaptors::reverse(toClose)) {
632       auto pos = pipes_.begin() + idx;
633       closeChecked(pos->parentFd);
634       pipes_.erase(pos);
635     }
636   }
637 }
638
639 int Subprocess::findByChildFd(int childFd) const {
640   auto pos = std::lower_bound(
641       pipes_.begin(), pipes_.end(), childFd,
642       [] (const PipeInfo& info, int fd) { return info.childFd < fd; });
643   if (pos == pipes_.end() || pos->childFd != childFd) {
644     throw std::invalid_argument(folly::to<std::string>(
645         "child fd not found ", childFd));
646   }
647   return pos - pipes_.begin();
648 }
649
650 void Subprocess::closeParentFd(int childFd) {
651   int idx = findByChildFd(childFd);
652   closeChecked(pipes_[idx].parentFd);
653   pipes_.erase(pipes_.begin() + idx);
654 }
655
656 namespace {
657
658 class Initializer {
659  public:
660   Initializer() {
661     // We like EPIPE, thanks.
662     ::signal(SIGPIPE, SIG_IGN);
663   }
664 };
665
666 Initializer initializer;
667
668 }  // namespace
669
670 }  // namespace folly
671