Refactor ShellManager process spawning

1 files changed, 133 insertions, 100 deletions

diff --git a/src/shell_manager.cc b/src/shell_manager.cc
index e1d03ffd..d2d6789a 100644
--- a/src/shell_manager.cc
+++ b/src/shell_manager.cc
@@ -13,8 +13,6 @@
 namespace Kakoune
 {
 
-static const Regex env_var_regex(R"(\bkak_(\w+)\b)");
-
 ShellManager::ShellManager()
 {
     const char* path = getenv("PATH");
@@ -22,117 +20,152 @@ ShellManager::ShellManager()
     setenv("PATH", new_path.c_str(), 1);
 }
 
+namespace
+{
+
+struct Pipe
+{
+    Pipe() { pipe(m_fd); }
+    ~Pipe() { close_read_fd(); close_write_fd(); }
+
+    int read_fd() const { return m_fd[0]; }
+    int write_fd() const { return m_fd[1]; }
+
+    void close_read_fd() { close_fd(m_fd[0]); }
+    void close_write_fd() { close_fd(m_fd[1]); }
+
+private:
+    void close_fd(int& fd) { if (fd != -1) { close(fd); fd = -1; } }
+    int m_fd[2];
+};
+
+pid_t spawn_process(StringView cmdline, ConstArrayView<String> params, ConstArrayView<String> kak_env,
+                    const Pipe& child_stdout, const Pipe& child_stdin, const Pipe& child_stderr)
+{
+    Vector<const char*> envptrs;
+    for (char** envp = environ; *envp; ++envp)
+        envptrs.push_back(*envp);
+    for (auto& env : kak_env)
+        envptrs.push_back(env.c_str());
+    envptrs.push_back(nullptr);
+
+    const char* shell = "/bin/sh";
+    auto cmdlinezstr = cmdline.zstr();
+    Vector<const char*> execparams = { shell, "-c", cmdlinezstr };
+    if (not params.empty())
+        execparams.push_back(shell);
+    for (auto& param : params)
+        execparams.push_back(param.c_str());
+    execparams.push_back(nullptr);
+
+    if (pid_t pid = fork())
+        return pid;
+
+    auto move = [](int oldfd, int newfd) { dup2(oldfd, newfd); close(oldfd); };
+
+    close(child_stdout.write_fd());
+    move(child_stdout.read_fd(), 0);
+
+    close(child_stdin.read_fd());
+    move(child_stdin.write_fd(), 1);
+
+    close(child_stderr.read_fd());
+    move(child_stderr.write_fd(), 2);
+
+    execve(shell, (char* const*)execparams.data(), (char* const*)envptrs.data());
+    exit(-1);
+    return -1;
+}
+
+}
+
 std::pair<String, int> ShellManager::eval(
     StringView cmdline, const Context& context, StringView input,
     Flags flags, ConstArrayView<String> params, const EnvVarMap& env_vars)
 {
-    int write_pipe[2]; // child stdin
-    int read_pipe[2];  // child stdout
-    int error_pipe[2]; // child stderr
+    static const Regex re(R"(\bkak_(\w+)\b)");
 
-    ::pipe(write_pipe);
-    ::pipe(read_pipe);
-    ::pipe(error_pipe);
-
-    if (pid_t pid = fork())
+    Vector<String> kak_env;
+    for (RegexIterator<const char*> it{cmdline.begin(), cmdline.end(), re}, end;
+         it != end; ++it)
     {
-        close(write_pipe[0]);
-        close(read_pipe[1]);
-        close(error_pipe[1]);
+        StringView name{(*it)[1].first, (*it)[1].second};
 
-        write(write_pipe[1], input.data(), (int)input.length());
-        close(write_pipe[1]);
+        auto match_name = [&](const String& s) {
+            return s.length() > name.length()  and
+                   prefix_match(s, name) and s[name.length()] == '=';
+        };
+        if (find_if(kak_env, match_name) != kak_env.end())
+            continue;
 
-        String child_stdout, child_stderr;
-        int status = 0;
+        auto var_it = env_vars.find(name);
+        try
         {
-            auto pipe_reader = [](String& output) {
-                return [&output](FDWatcher& watcher, EventMode) {
-                    char buffer[1024];
-                    size_t size = read(watcher.fd(), buffer, 1024);
-                    if (size <= 0)
-                        watcher.close_fd();
-                    output += StringView{buffer, buffer+size};
-                };
-            };
-
-            FDWatcher stdout_watcher{read_pipe[0], pipe_reader(child_stdout)};
-            FDWatcher stderr_watcher{error_pipe[0], pipe_reader(child_stderr)};
-
-            // block SIGCHLD to make sure we wont receive it before
-            // our call to pselect, that will end up blocking indefinitly.
-            sigset_t mask, orig_mask;
-            sigemptyset(&mask);
-            sigaddset(&mask, SIGCHLD);
-            sigprocmask(SIG_BLOCK, &mask, &orig_mask);
-
-            // check for termination now that SIGCHLD is blocked
-            bool terminated = waitpid(pid, &status, WNOHANG);
-
-            while (not terminated or
-                   ((flags & Flags::WaitForStdout) and
-                    (not stdout_watcher.closed() or
-                     not stderr_watcher.closed())))
-            {
-                EventManager::instance().handle_next_events(EventMode::Urgent, &orig_mask);
-                if (not terminated)
-                    terminated = waitpid(pid, &status, WNOHANG);
-            }
-
-            stdout_watcher.close_fd();
-            stderr_watcher.close_fd();
-
-            sigprocmask(SIG_SETMASK, &orig_mask, nullptr);
-        }
-
-        if (not child_stderr.empty())
-            write_to_debug_buffer(format("shell stderr: <<<\n{}>>>", child_stderr));
-
-        return { child_stdout, WIFEXITED(status) ? WEXITSTATUS(status) : - 1 };
+            const String& value = var_it != env_vars.end() ?
+                var_it->value : get_val(name, context);
+
+            kak_env.push_back(format("kak_{}={}", name, value));
+        } catch (runtime_error&) {}
     }
-    else try
-    {
-        close(write_pipe[1]);
-        close(read_pipe[0]);
-        close(error_pipe[0]);
 
-        dup2(read_pipe[1], 1); close(read_pipe[1]);
-        dup2(error_pipe[1], 2); close(error_pipe[1]);
-        dup2(write_pipe[0], 0); close(write_pipe[0]);
+    Pipe child_stdin, child_stdout, child_stderr;
+    pid_t pid = spawn_process(cmdline, params, kak_env,
+                              child_stdin, child_stdout, child_stderr);
 
-        using RegexIt = RegexIterator<StringView::iterator>;
-        for (RegexIt it{cmdline.begin(), cmdline.end(), env_var_regex}, end;
-             it != end; ++it)
-        {
-            auto& match = *it;
-
-            StringView name{match[1].first, match[1].second};
-            kak_assert(name.length() > 0);
-
-            auto local_var = env_vars.find(name);
-            if (local_var != env_vars.end())
-                setenv(("kak_" + name).c_str(), local_var->value.c_str(), 1);
-            else try
-            {
-                String value = get_val(name, context);
-                setenv(format("kak_{}", name).c_str(), value.c_str(), 1);
-            }
-            catch (runtime_error&) {}
-        }
-        const char* shell = "/bin/sh";
-        auto cmdlinezstr = cmdline.zstr();
-        Vector<const char*> execparams = { shell, "-c", cmdlinezstr };
-        if (not params.empty())
-            execparams.push_back(shell);
-        for (auto& param : params)
-            execparams.push_back(param.c_str());
-        execparams.push_back(nullptr);
-
-        execvp(shell, (char* const*)execparams.data());
-        exit(-1);
+    child_stdin.close_read_fd();
+    child_stdout.close_write_fd();
+    child_stderr.close_write_fd();
+
+    write(child_stdin.write_fd(), input.data(), (int)input.length());
+    child_stdin.close_write_fd();
+
+    struct PipeReader : FDWatcher
+    {
+        PipeReader(Pipe& pipe, String& contents)
+            : FDWatcher(pipe.read_fd(),
+                        [&contents, &pipe](FDWatcher& watcher, EventMode) {
+                            char buffer[1024];
+                            size_t size = read(pipe.read_fd(), buffer, 1024);
+                            if (size <= 0)
+                            {
+                                pipe.close_read_fd();
+                                watcher.disable();
+                                return;
+                            }
+                            contents += StringView{buffer, buffer+size};
+                        })
+        {}
+    };
+
+    String stdout_contents, stderr_contents;
+    PipeReader stdout_reader{child_stdout, stdout_contents};
+    PipeReader stderr_reader{child_stderr, stderr_contents};
+
+    // block SIGCHLD to make sure we wont receive it before
+    // our call to pselect, that will end up blocking indefinitly.
+    sigset_t mask, orig_mask;
+    sigemptyset(&mask);
+    sigaddset(&mask, SIGCHLD);
+    sigprocmask(SIG_BLOCK, &mask, &orig_mask);
+    auto restore_mask = on_scope_end([&] { sigprocmask(SIG_SETMASK, &orig_mask, nullptr); });
+
+    int status = 0;
+    // check for termination now that SIGCHLD is blocked
+    bool terminated = waitpid(pid, &status, WNOHANG);
+
+    while (not terminated or
+           ((flags & Flags::WaitForStdout) and
+            (child_stdout.read_fd() != -1 or child_stderr.read_fd() != -1)))
+    {
+        EventManager::instance().handle_next_events(EventMode::Urgent, &orig_mask);
+        if (not terminated)
+            terminated = waitpid(pid, &status, WNOHANG);
     }
-    catch (...) { exit(-1); }
-    return {};
+
+    if (not stderr_contents.empty())
+        write_to_debug_buffer(format("shell stderr: <<<\n{}>>>", stderr_contents));
+
+    return { stdout_contents, WIFEXITED(status) ? WEXITSTATUS(status) : -1 };
 }
 
 void ShellManager::register_env_var(StringView str, bool prefix,