Process.hpp

Go to the documentation of this file.

#pragma once
 
#ifdef _WIN32
#error "Process.hpp is currently UNIX only, and does not support Windows. Feel free to open a PR to change this"
#endif
 
#include <array>
#include <atomic>
#include <cstdlib>
#include <cstring>
#include <fcntl.h>
#include <filesystem>
#include <format>
#include <functional>
#include <iostream>
#include <map>
#include <memory>
#include <mutex>
#include <optional>
#include <pty.h>
#include <sstream>
#include <string>
#include <string_view>
#include <sys/poll.h>
#include <sys/wait.h>
#include <thread>
#include <unistd.h>
#include <variant>
#include <vector>
 
// TODO: this cannot be "unix", or the build inexplicably dies ("unexpected { before numeric constant")
// It's probably a macro
namespace stc::Unix {
 
struct LowLevelWrapper {
    ssize_t writeToFd(const std::string& data, int fd) {
        if (fd < 0) {
            throw std::runtime_error("Illegal write on closed or invalid fd");
        }
        ssize_t bytes = write(fd, data.data(), data.size());
        if (bytes > 0) {
            return bytes;
        }
        return 0;
    }
 
    ssize_t readFromFd(std::array<char, 4096>& out, int fd) {
        ssize_t sum = 0;
 
        nfds_t nfds = 1;
        pollfd pdfs = {
            .fd = fd,
            .events = POLLIN,
            .revents = 0
        };
 
        // we need a small timeout here to prevent race conditions
        while (poll(&pdfs, nfds, 10)) {
            ssize_t bytes = read(
                fd,
                out.data(),
                out.size()
            );
            if (bytes > 0) {
                sum += bytes;
                break;
            }
        }
        return sum;
    }
 
    ssize_t readFromFd(std::stringstream& out, int fd) {
        std::array<char, 4096> buff;
        ssize_t sum = 0;
 
        nfds_t nfds = 1;
        pollfd pdfs = {
            .fd = fd,
            .events = POLLIN,
            .revents = 0
        };
 
        // we need a small timeout here to prevent race conditions
        while (poll(&pdfs, nfds, 10)) {
            ssize_t bytes = read(
                fd,
                buff.data(),
                buff.size()
            );
            if (bytes > 0) {
                out << std::string_view {
                    buff.begin(), buff.begin() + bytes 
                };
                sum += bytes;
            }
        }
        return sum;
    }
 
    virtual int readFd() = 0;
};
 
struct Pipe : public LowLevelWrapper {
    std::array<int, 2> fds;
    Pipe() {
        if (pipe(fds.data()) != 0) {
            throw std::runtime_error("Failed to open pipe");
        }
    }
 
    ~Pipe() {
        die();
    }
 
    void die() {
        closeRead();
        closeWrite();
    }
    void closeRead() {
        if (fds[0] != -1) {
            close(fds[0]);
        }
        fds[0] = -1;
    }
    void closeWrite() {
        if (fds[1] != -1) {
            close(fds[1]);
        }
 
        fds[1] = -1;
    }
 
    int readFd() override {
        return fds[0];
    }
    int writeFd() {
        return fds[1];
    }
 
    ssize_t readData(std::stringstream& out) {
        return readFromFd(out, readFd());
    }
};
 
struct PTY : public LowLevelWrapper {
    int master, slave;
 
    PTY() {
        // TODO: figure out if it makes sense to:
        // 1. store the name
        // 2. Allow customising whatever the last two parameters are
        if (openpty(&master, &slave, nullptr, nullptr, nullptr) == -1) {
            throw std::runtime_error("Failed to open PTY");
        }
    }
    ~PTY() {
        die();
    }
 
    void die() {
        closeMasterChannel();
        closeSlaveChannel();
    }
 
    void closeMasterChannel() {
        if (master >= 0) {
            close(master);
        }
        master = -1;
    }
    void closeSlaveChannel() {
        if (slave >= 0) {
            close(slave);
        }
        slave = -1;
    }
 
    virtual int readFd() override {
        return master;
    }
 
    ssize_t writeToStdin(const std::string& data) {
        return writeToFd(data, master);
    }
 
    ssize_t readData(std::stringstream& out) {
        return readFromFd(out, master);
    }
};
 
inline std::shared_ptr<Pipe> createPipe() {
    return std::make_shared<Pipe>();
}
 
struct Pipes {
    std::shared_ptr<Pipe> stdoutPipe = nullptr;
    std::shared_ptr<Pipe> stderrPipe = nullptr;
    std::shared_ptr<Pipe> stdinPipe = nullptr;
 
    void die() {
        if (stdoutPipe) stdoutPipe->die();
        if (stderrPipe) stderrPipe->die();
        if (stdinPipe) stdinPipe->die();
    }
 
    ssize_t writeToStdin(const std::string& data) {
        if (stdinPipe == nullptr) {
            throw std::runtime_error("Must open stdin to write to stdin");
        }
        return stdinPipe->writeToFd(data, stdinPipe->writeFd());
    }
 
    static Pipes separate(bool withStdin = true) {
        return Pipes {
            createPipe(),
            createPipe(),
            withStdin ? createPipe() : nullptr
        };
    }
    static Pipes shared(bool withStdin = true) {
        auto outPipe = createPipe();
        return Pipes {
            outPipe,
            outPipe,
            withStdin ? createPipe() : nullptr
        };
 
    }
};
 
 
inline std::shared_ptr<PTY> createPTY() {
    return std::make_shared<PTY>();
}
 
struct Environment {
    std::map<std::string, std::string> env = {};
    bool extendEnviron = true;
 
    std::optional<std::string> workingDirectory = std::nullopt;
 
    void validate() const {
        for (auto& [k, v] : env) {
            if (k.find('=') != std::string::npos) {
                throw std::runtime_error("Illegal key: " + k);
            }
        }
 
        if (workingDirectory.has_value() && !std::filesystem::is_directory(*workingDirectory)) {
            throw std::runtime_error(
                std::format(
                    "Working directory set to {}, which does not exist or isn't a directory",
                    *workingDirectory
                )
            );
        }
    }
};
 
struct Config {
    bool verboseUserOutput = false;
};
 
struct ReadHandler {
    virtual void read(
        LowLevelWrapper* primitive
    ) = 0;
};
 
struct InMemoryReadHandler : public ReadHandler {
    std::stringstream ss = {};
    virtual void read(
        LowLevelWrapper* primitive
    ) override {
        primitive->readFromFd(
            ss,
            primitive->readFd()
        );
    }
 
    std::string getStream(bool reset = false) {
        std::string d = ss.str();
 
        if (reset) {
            ss = {};
        }
 
        return d;
    }
};
 
struct FdRedirectInputHandler : public ReadHandler {
    int fd;
 
    FdRedirectInputHandler(int fd) : fd(fd) {}
 
    virtual void read(
        LowLevelWrapper* primitive
    ) override {
        thread_local std::array<char, 4096> buff;
        ssize_t bytes = primitive->readFromFd(
            buff,
            primitive->readFd()
        );
 
        if (bytes > 0) {
            auto written = write(fd, buff.data(), bytes);
 
            if (written <= 0) {
                std::cerr << "Writing failed: " << strerror(errno) << std::endl;
                throw std::runtime_error("Failed to write to output buffer");
            }
        }
    }
};
 
struct ReadHandlers {
    std::shared_ptr<ReadHandler> stdoutHandler, stderrHandler;
 
    static ReadHandlers inMemory(bool separateStderr = true) {
        return {
            .stdoutHandler = std::make_shared<InMemoryReadHandler>(),
            .stderrHandler = separateStderr ? std::make_shared<InMemoryReadHandler>() : nullptr,
        };
    }
 
    static ReadHandlers stdStreamRedirect(bool separateStderr = true) {
        return {
            .stdoutHandler = std::make_shared<FdRedirectInputHandler>(STDOUT_FILENO),
            .stderrHandler = separateStderr ? std::make_shared<FdRedirectInputHandler>(STDERR_FILENO) : nullptr,
        };
    }
};
 
class Process {
protected:
    std::optional<decltype(fork())> pid = std::nullopt;
 
    std::optional<
        std::variant<Pipes, std::shared_ptr<PTY>>
    > interface;
    // std::stringstream stdoutBuff, stderrBuff;
    ReadHandlers readHandlers;
    std::mutex lock;
 
    std::thread inputCollector;
    std::atomic<int> statusCode = -1;
    std::atomic<std::optional<bool>> exitedNormally;
    bool running = true;
 
    Config config;
    
    bool waitPid(int opts = 0) {
        int wstatus;
        if (!pid.has_value()) {
            std::cerr << "waitPid called, but pid has no value. Something has gone very wrong" << std::endl;
            exit(70);
        }
        if (waitpid(*pid, &wstatus, opts) > 0) {
            if (WIFEXITED(wstatus)) {
                statusCode = WEXITSTATUS(wstatus);
                exitedNormally = true;
            } else if (WIFSIGNALED(wstatus)) {
                statusCode = WTERMSIG(wstatus);
                exitedNormally = false;
            } else if (WIFSTOPPED(wstatus)) {
                statusCode = WSTOPSIG(wstatus);
                exitedNormally = false;
            } else {
                std::cerr
                    << "WARNING: stc::Unix::Process got an unknown status: " << wstatus 
                    << std::endl;
            }
            this->running = false;
            return true;
        }
        return false;
    }
 
    char* const* createEnviron(
        const std::optional<Environment>& env
    ) {
        if (env == std::nullopt) {
            return environ;
        }
 
        size_t size = 0;
        for (char **env = environ; *env != nullptr; env++) {
            ++size;
        }
 
        // This is disgusting, but it beats fucking around with the real raw types. This will technically leak a vector,
        // but it does not matter because it's disappeared once exec is called:
        // https://stackoverflow.com/a/3617385
        // Would prefer to do this better, but I just don't want to
        std::vector<char*>* data = new std::vector<char*>;
        if (env->extendEnviron && size > 0) {
            data->assign(
                environ, environ + size
            );
        }
 
        data->reserve(
            // Existing data or 0
            data->size()
            // nullptr
            + 1
            // Extra envs
            + env->env.size()
        );
        for (const auto& [k, v] : env->env) {
            if (env->extendEnviron) {
                // If we're extending environ, keys here can conflict with environ. They won't conflict internally,
                // because env is a non-multimap, and we enforce keys not containing '=', so no weird injection shit
                // resulting in identical strings.
                data->erase(
                    std::remove_if(
                        data->begin(),
                        data->end(), 
                        [&](const auto& v) -> bool {
                            return strncmp(v, k.data(), k.size()) == 0;
                        }
                    ), data->end()
                );
            }
 
            std::string combined = std::format(
                "{}={}", k, v
            );
            auto* newStr = strdup(combined.c_str());
            if (newStr == nullptr) {
                std::cerr << "Failed to copy string to env" << std::endl;
                exit(69);
            }
            data->push_back(newStr);
        }
        data->push_back(nullptr);
 
        return data->data();
    }
 
    void doSpawnCommand(
        const std::vector<std::string>& command,
        const std::function<void()>& readImpl,
        const std::function<void()>& prepDuping,
        const std::optional<Environment>& env
    ) {
        if (command.size() == 0) {
            throw std::runtime_error("Cannot run null command");
        }
        std::vector<const char*> convertedCommand;
 
        if (env) {
            env->validate();
        }
 
        if (config.verboseUserOutput) {
            std::cout << "Exec: ";
        }
        convertedCommand.reserve(command.size() + 1);
        for (auto& str : command) {
            if (config.verboseUserOutput) {
                std::cout << std::quoted(str);
                // This isn't strictly speaking necessary, but avoids a situation where the tests need a .starts_with(),
                // or it'll require the full string to contain a load-bearing linebreak AND a load-bearing trailing
                // space.
                if (convertedCommand.size() != command.size() - 1) {
                    std::cout << " ";
                }
            }
            convertedCommand.push_back(str.c_str());
        }
        if (config.verboseUserOutput) {
            std::cout << "\n";
        }
        convertedCommand.push_back(nullptr);
 
        pid = fork();
        if (pid < 0) {
            throw std::runtime_error("Failed to fork");
        } else if (pid == 0) {
            // Child process
            if (prepDuping != nullptr) {
                prepDuping();
            }
 
            // Close handles if they're opened
            if (interface) {
                std::visit([](auto& resolved) {
                    using T = std::decay_t<decltype(resolved)>;
                    if constexpr (std::is_same_v<T, std::shared_ptr<PTY>>) {
                        resolved->die();
                    } else {
                        resolved.die();
                    }
                }, interface.value());
            }
 
            if (env.has_value()) {
                if (env->workingDirectory.has_value()) {
                    // C++17 <3
                    // Avoids chdir() from <unistd.h> so there's one less thing to do if this file can be made portable.
                    std::filesystem::current_path(
                        env->workingDirectory.value()
                    );
                }
            }
 
            execve(
                convertedCommand.at(0),
                (char**) convertedCommand.data(),
                createEnviron(env)
            );
        } else {
            // Parent process
            if (readImpl != nullptr) {
                this->inputCollector = std::thread(
                    std::bind(&Process::run, this, readImpl)
                );
            }
        }
    }
 
    void run(const std::function<void()>& readImpl) {
        // TODO: readImpl should bake in some timeout here, but is that enough? Is this thread going to be too busy?
        do {
            readImpl();
        } while (!waitPid(WNOHANG));
        // Read anything left in the buffer at exit time
        readImpl();
    }
public:
    [[nodiscard("Discarding immediately terminates the process. You probably don't want this")]]
    Process(
        const std::vector<std::string>& command,
        const std::optional<Environment>& env = std::nullopt,
        const Config& config = {}
    ): config(config) {
        doSpawnCommand(command, nullptr, nullptr, env);
    }
 
    [[nodiscard("Discarding immediately terminates the process. You probably don't want this")]]
    Process(
        const std::vector<std::string>& command,
        const Pipes& pipes,
        const std::optional<Environment>& env = std::nullopt,
        const Config& config = {},
        const ReadHandlers& readHandlers = ReadHandlers::inMemory()
    ): readHandlers(readHandlers), config(config) {
        interface = pipes;
 
        doSpawnCommand(command, [this]() {
            auto& pipes = std::get<Pipes>(this->interface.value());
            if (pipes.stdoutPipe != nullptr && this->readHandlers.stdoutHandler != nullptr) {
                std::lock_guard l(lock);
                this->readHandlers.stdoutHandler->read(pipes.stdoutPipe.get());
            }
            if (pipes.stderrPipe != nullptr && this->readHandlers.stderrHandler != nullptr) {
                std::lock_guard l(lock);
                this->readHandlers.stderrHandler->read(pipes.stderrPipe.get());
            }
        }, [&]() {
            if (pipes.stdinPipe != nullptr) {
                dup2(pipes.stdinPipe->readFd(), STDIN_FILENO);
            }
            if (pipes.stdoutPipe != nullptr) {
                dup2(pipes.stdoutPipe->writeFd(), STDOUT_FILENO);
            }
            if (pipes.stderrPipe != nullptr) {
                dup2(pipes.stderrPipe->writeFd(), STDERR_FILENO);
            }
            std::get<Pipes>(*interface).die();
        }, env);
    }
 
    [[nodiscard("Discarding immediately terminates the process. You probably don't want this")]]
    Process(
        const std::vector<std::string>& command,
        const std::shared_ptr<PTY>& pty,
        const std::optional<Environment>& env = std::nullopt,
        const Config& config = {},
        const ReadHandlers& readHandlers = ReadHandlers::inMemory()
    ): readHandlers(readHandlers), config(config) {
        if (pty == nullptr) {
            throw std::runtime_error(
                "pty cannot be null. If you don't want to attach anything, use the non-pipe/non-PTY constructor instead"
            );
        }
        interface = pty;
        doSpawnCommand(command, [this]() {
            auto pty = std::get<std::shared_ptr<PTY>>(this->interface.value());
 
            {
                std::lock_guard l(lock);
                if (this->readHandlers.stdoutHandler != nullptr) {
                    this->readHandlers.stdoutHandler->read(
                        pty.get()
                    );
                }
            }
            // TODO: a random python-related question I stumbled into suggested using two PTYs so the output and input
            // can be handled separately. This was in relation to closing stdin. In theory, three separate PTYs could be
            // used to achieve the same system as pipes. I imagine this is what some terminals use to highlight error
            // output separately from standard output?
        }, [&]() {
            dup2(pty->slave, STDIN_FILENO);
            dup2(pty->slave, STDOUT_FILENO);
            dup2(pty->slave, STDERR_FILENO);
            std::get<std::shared_ptr<PTY>>(*interface)->die();
        }, env);
    }
 
    virtual ~Process() {
        this->sigkill();
        this->block();
    }
 
    std::string getStdoutBuffer(bool reset = false) {
        std::lock_guard g(lock);
        if (auto ptr = std::dynamic_pointer_cast<InMemoryReadHandler>(this->readHandlers.stdoutHandler);
            ptr != nullptr
        ) {
            return ptr->getStream(reset);
        } else {
            return "";
        }
    }
 
    std::string getStderrBuffer(bool reset = false) {
        std::lock_guard g(lock);
        if (auto ptr = std::dynamic_pointer_cast<InMemoryReadHandler>(this->readHandlers.stderrHandler);
            ptr != nullptr
        ) {
            return ptr->getStream(reset);
        } else {
            return "";
        }
    }
 
    void resetBuffers() {
        std::lock_guard g(lock);
        if (auto ptr = std::dynamic_pointer_cast<InMemoryReadHandler>(this->readHandlers.stdoutHandler);
            ptr != nullptr
        ) {
            ptr->ss = {};
        } else {
            throw std::runtime_error("stdout handler is null or otherwise not an InMemoryReadHandler");
        }
        if (auto ptr = std::dynamic_pointer_cast<InMemoryReadHandler>(this->readHandlers.stderrHandler);
            ptr != nullptr
        ) {
            ptr->ss = {};
        } else {
            throw std::runtime_error("stderr handler is null or otherwise not an InMemoryReadHandler");
        }
    }
 
    ssize_t writeToStdin(const std::string& data) {
        if (!this->running) {
            throw std::runtime_error("Cannot write to closed proc");
        }
        if (interface) {
            return std::visit([&](auto& resolved) -> ssize_t {
                using T = std::decay_t<decltype(resolved)>;
                if constexpr (std::is_same_v<T, std::shared_ptr<PTY>>) {
                    return resolved->writeToStdin(data);
                } else {
                    return resolved.writeToStdin(data);
                }
            }, interface.value());
        } else {
            throw std::runtime_error("Must use pty or pipe mode to write to stdin");
        }
    }
 
    int block() {
        if (this->interface) {
            if (inputCollector.joinable()) {
                inputCollector.join();
            }
            return statusCode;
        } else {
            waitPid();
            return statusCode;
        }
    }
 
    void signal(int sig) {
        if (statusCode == -1) {
            if (pid.has_value() && *pid > 0) {
                kill(*pid, sig);
            }
        }
    }
 
    void stop() {
        signal(SIGTERM);
    }
 
    void sigkill() {
        signal(SIGKILL);
    }
 
    void closeStdin() {
        if (!this->interface.has_value()) {
            throw std::runtime_error("Must use pipe or pty mode to use this function");
        }
 
        if (std::holds_alternative<Pipes>(*this->interface)) {
            auto& ptr = std::get<Pipes>(*this->interface).stdinPipe;
            if (ptr) {
                ptr->closeWrite();
            }
        }
    }
 
    std::optional<bool> hasExitedNormally() {
        return exitedNormally;
    }
 
    std::optional<int> getExitCode() {
        if (statusCode != -1) {
            return statusCode;
        }
        return std::nullopt;
    }
 
};
 
}