| 1 | /*
 | 
| 2 |  * Souffle - A Datalog Compiler
 | 
| 3 |  * Copyright (c) 2021, Oracle and/or its affiliates. All rights reserved
 | 
| 4 |  * Licensed under the Universal Permissive License v 1.0 as shown at:
 | 
| 5 |  * - https://opensource.org/licenses/UPL
 | 
| 6 |  * - <souffle root>/licenses/SOUFFLE-UPL.txt
 | 
| 7 |  */
 | 
| 8 | 
 | 
| 9 | /************************************************************************
 | 
| 10 |  *
 | 
| 11 |  * @file SubProcess.h
 | 
| 12 |  *
 | 
| 13 |  * Wrapper for launching subprocesses.
 | 
| 14 |  *
 | 
| 15 |  ***********************************************************************/
 | 
| 16 | 
 | 
| 17 | #pragma once
 | 
| 18 | 
 | 
| 19 | #include "souffle/utility/Types.h"
 | 
| 20 | #include "souffle/utility/span.h"
 | 
| 21 | #include <algorithm>
 | 
| 22 | #include <cassert>
 | 
| 23 | #include <cstdlib>
 | 
| 24 | #include <optional>
 | 
| 25 | #include <type_traits>
 | 
| 26 | 
 | 
| 27 | #ifdef _MSC_VER
 | 
| 28 | #define NOMINMAX
 | 
| 29 | #include <windows.h>
 | 
| 30 | #else
 | 
| 31 | #include <sys/wait.h>
 | 
| 32 | #include <unistd.h>
 | 
| 33 | #endif
 | 
| 34 | 
 | 
| 35 | namespace souffle {
 | 
| 36 | 
 | 
| 37 | namespace detail {
 | 
| 38 | [[noreturn]] inline void perrorExit(char const* msg) {
 | 
| 39 |     ::perror(msg);
 | 
| 40 |     std::exit(EXIT_FAILURE);
 | 
| 41 | }
 | 
| 42 | 
 | 
| 43 | // These are used by bash and are a defacto standard on Linux.
 | 
| 44 | // This list is incomplete.
 | 
| 45 | enum LinuxExitCode : int {
 | 
| 46 |     cannot_execute = 126,
 | 
| 47 |     command_not_found = 127,
 | 
| 48 | };
 | 
| 49 | 
 | 
| 50 | using LinuxWaitStatus = int;
 | 
| 51 | }  // namespace detail
 | 
| 52 | 
 | 
| 53 | /**
 | 
| 54 |  * Executes a process w/ the given `argv` arguments and `envp` overriden env-vars.
 | 
| 55 |  *
 | 
| 56 |  * @param   argv  The arguments to the process.
 | 
| 57 |  *                Do not include the 'program invoked as' argument 0. This is implicitly done for you.
 | 
| 58 |  * @param   envp  Collection of env vars to override.
 | 
| 59 |  *                Any env not specified in `envp` is inherited from this process' environment.
 | 
| 60 |  * @return  `None` IFF unable to launch `program`, otherwise `program`'s `wait` status.
 | 
| 61 |  *           NB: This is not the exit code, though the exit code can be obtained from it.
 | 
| 62 |  *               However, you can do `execute(...) == 0` if you only care that it succeeded.
 | 
| 63 |  */
 | 
| 64 | template <typename Envp = span<std::pair<char const*, char const*>>,
 | 
| 65 |         typename = std::enable_if_t<is_iterable_of<Envp, std::pair<char const*, char const*> const>>>
 | 
| 66 | std::optional<detail::LinuxWaitStatus> execute(
 | 
| 67 |         std::string const& program, span<char const* const> argv = {}, Envp&& envp = {}) {
 | 
| 68 | #ifndef _MSC_VER
 | 
| 69 |     using EC = detail::LinuxExitCode;
 | 
| 70 | 
 | 
| 71 |     auto pid = ::fork();
 | 
| 72 |     switch (pid) {
 | 
| 73 |         case -1: return std::nullopt;  // unable to fork. likely hit a resource limit of some kind.
 | 
| 74 | 
 | 
| 75 |         case 0: {  // child
 | 
| 76 |             // thankfully we're a fork. we can trash this proc's `::environ` w/o reprocussions
 | 
| 77 |             for (auto&& [k, v] : envp) {
 | 
| 78 |                 if (::setenv(k, v, 1)) detail::perrorExit("setenv");
 | 
| 79 |             }
 | 
| 80 | 
 | 
| 81 |             char* argv_temp[argv.size() + 2];
 | 
| 82 |             argv_temp[0] = const_cast<char*>(program.c_str());
 | 
| 83 |             std::copy_n(argv.data(), argv.size(), const_cast<char const**>(argv_temp) + 1);
 | 
| 84 |             argv_temp[argv.size() + 1] = nullptr;
 | 
| 85 | 
 | 
| 86 |             ::execvp(program.c_str(), argv_temp);
 | 
| 87 |             std::exit(EC::cannot_execute);
 | 
| 88 |         }
 | 
| 89 | 
 | 
| 90 |         default: {  // parent
 | 
| 91 |             detail::LinuxWaitStatus status;
 | 
| 92 |             if (::waitpid(pid, &status, 0) == -1) {
 | 
| 93 |                 // not recoverable / should never happen.
 | 
| 94 |                 detail::perrorExit("`waitpid` failed");
 | 
| 95 |             }
 | 
| 96 | 
 | 
| 97 |             // check it exited or signaled (didn't specify `WNOHANG` or `WUNTRACED`)
 | 
| 98 |             assert(WIFSIGNALED(status) || WIFEXITED(status));
 | 
| 99 | 
 | 
| 100 |             // check that the fork child successfully `exec`'d
 | 
| 101 |             if (WIFEXITED(status)) {
 | 
| 102 |                 switch (WEXITSTATUS(status)) {
 | 
| 103 |                     default: return WEXITSTATUS(status);
 | 
| 104 | 
 | 
| 105 |                     case EC::cannot_execute:                          // FALL THRU: command_not_found
 | 
| 106 |                     case EC::command_not_found: return std::nullopt;  // fork couldn't execute the program
 | 
| 107 |                 }
 | 
| 108 |             }
 | 
| 109 |             // what should be returned on signal? Treat as error
 | 
| 110 |             return EXIT_FAILURE;
 | 
| 111 |         }
 | 
| 112 |     }
 | 
| 113 | #else
 | 
| 114 |     STARTUPINFOW si;
 | 
| 115 |     PROCESS_INFORMATION pi;
 | 
| 116 |     DWORD exit_code = 0;
 | 
| 117 | 
 | 
| 118 |     memset(&si, 0, sizeof(si));
 | 
| 119 |     si.cb = sizeof(si);
 | 
| 120 |     memset(&pi, 0, sizeof(pi));
 | 
| 121 | 
 | 
| 122 |     std::size_t l;
 | 
| 123 |     std::wstring program_w(program.length() + 1, L' ');
 | 
| 124 |     ::mbstowcs_s(&l, program_w.data(), program_w.size(), program.data(), program.size());
 | 
| 125 |     program_w.resize(l - 1);
 | 
| 126 | 
 | 
| 127 |     WCHAR FoundPath[PATH_MAX];
 | 
| 128 |     int64_t Found = (int64_t)FindExecutableW(program_w.c_str(), nullptr, FoundPath);
 | 
| 129 |     if (Found <= 32) {
 | 
| 130 |         std::cerr << "Cannot find executable '" << program << "'.\n";
 | 
| 131 |         return std::nullopt;
 | 
| 132 |     }
 | 
| 133 | 
 | 
| 134 |     std::wstringstream args_w;
 | 
| 135 |     args_w << program_w;
 | 
| 136 |     for (const auto& arg : argv) {
 | 
| 137 |         std::string arg_s(arg);
 | 
| 138 |         std::wstring arg_w(arg_s.length() + 1, L' ');
 | 
| 139 |         ::mbstowcs_s(&l, arg_w.data(), arg_w.size(), arg_s.data(), arg_s.size());
 | 
| 140 |         arg_w.resize(l - 1);
 | 
| 141 |         args_w << L' ' << arg_w;
 | 
| 142 |     }
 | 
| 143 | 
 | 
| 144 |     std::string envir;
 | 
| 145 |     for (const auto& couple : envp) {
 | 
| 146 |         envir += couple.first;
 | 
| 147 |         envir += '=';
 | 
| 148 |         envir += couple.second;
 | 
| 149 |         envir += '\0';
 | 
| 150 |     }
 | 
| 151 |     envir += '\0';
 | 
| 152 | 
 | 
| 153 |     if (!CreateProcessW(FoundPath, args_w.str().data(), NULL, NULL, FALSE, 0, /*envir.data()*/ nullptr, NULL,
 | 
| 154 |                 &si, &pi)) {
 | 
| 155 |         return std::nullopt;
 | 
| 156 |     }
 | 
| 157 | 
 | 
| 158 |     WaitForSingleObject(pi.hProcess, INFINITE);
 | 
| 159 | 
 | 
| 160 |     if (!GetExitCodeProcess(pi.hProcess, &exit_code)) {
 | 
| 161 |         return std::nullopt;
 | 
| 162 |     }
 | 
| 163 | 
 | 
| 164 |     CloseHandle(pi.hProcess);
 | 
| 165 |     CloseHandle(pi.hThread);
 | 
| 166 | 
 | 
| 167 |     return static_cast<int>(exit_code);
 | 
| 168 | 
 | 
| 169 | #endif
 | 
| 170 | }
 | 
| 171 | 
 | 
| 172 | /**
 | 
| 173 |  * Executes a process w/ the given `argv` arguments and `envp` overriden env-vars.
 | 
| 174 |  *
 | 
| 175 |  * @param   argv  The arguments to the process.
 | 
| 176 |  *                Do not include the 'program invoked as' argument 0. This is implicitly done for you.
 | 
| 177 |  * @param   envp  Collection of env vars to override.
 | 
| 178 |  *                Any env not specified in `envp` is inherited from this process' environment.
 | 
| 179 |  * @return  `None` IFF unable to launch `program`, otherwise `program`'s `wait` status.
 | 
| 180 |  *           NB: This is not the exit code, though the exit code can be obtained from it.
 | 
| 181 |  *               However, you can do `execute(...) == 0` if you only care that it succeeded.
 | 
| 182 |  */
 | 
| 183 | template <typename Envp = span<std::pair<char const*, std::string>>,
 | 
| 184 |         typename = std::enable_if_t<is_iterable_of<Envp, std::pair<char const*, std::string> const>>>
 | 
| 185 | std::optional<detail::LinuxWaitStatus> execute(
 | 
| 186 |         std::string const& program, span<std::string const> argv, Envp&& envp = {}) {
 | 
| 187 |     auto go = [](auto* dst, auto&& src, auto&& f) {
 | 
| 188 |         size_t i = 0;
 | 
| 189 |         for (auto&& x : src)
 | 
| 190 |             dst[i++] = f(x);
 | 
| 191 |         return span<std::remove_pointer_t<decltype(dst)>>{dst, dst + src.size()};
 | 
| 192 |     };
 | 
| 193 | 
 | 
| 194 |     std::unique_ptr<char const*[]> argv_temp = std::make_unique<char const*[]>(argv.size());
 | 
| 195 |     std::unique_ptr<std::pair<char const*, char const*>[]> envp_temp =
 | 
| 196 |             std::make_unique<std::pair<char const*, char const*>[]>(envp.size());
 | 
| 197 |     auto argv_ptr = go(argv_temp.get(), argv, [](auto&& x) { return x.c_str(); });
 | 
| 198 |     auto envp_ptr = go(envp_temp.get(), envp, [](auto&& kv) {
 | 
| 199 |         return std::pair{kv.first, kv.second.c_str()};
 | 
| 200 |     });
 | 
| 201 |     return souffle::execute(program, argv_ptr, envp_ptr);
 | 
| 202 | }
 | 
| 203 | 
 | 
| 204 | }  // namespace souffle
 |