polybar/include/utils/command.hpp

#pragma once

#include <functional>
#include <memory>
#include <mutex>
#include <stdexcept>
#include <string>

#include "common.hpp"
#include "components/logger.hpp"
#include "utils/io.hpp"
#include "utils/process.hpp"
#include "utils/threading.hpp"

LEMONBUDDY_NS

namespace command_util {
  /**
   * Example usage:
   * @code cpp
   *   auto cmd = make_unique<Command>("cat /etc/rc.local");
   *   cmd->exec();
   *   cmd->tail(callback); //---> the contents of /etc/rc.local is sent to callback()
   *
   *   auto cmd = make_unique<Command>(
   *    "/bin/sh\n-c\n while read -r line; do echo data from parent process: $line; done");
   *   cmd->exec();
   *   cmd->writeline("Test");
   *   cout << cmd->readline(); //---> data from parent process: Test
   *
   *   auto cmd = make_unique<Command>("/bin/sh\n-c\nfor i in 1 2 3; do echo $i; done");
   *   cout << cmd->readline(); //---> 1
   *   cout << cmd->readline() << cmd->readline(); //---> 23
   * @encode
   */
  class command {
   public:
    explicit command(const logger& logger, string cmd, int out[2] = nullptr, int in[2] = nullptr)
        : m_log(logger), m_cmd(cmd) {
      if (in != nullptr) {
        m_stdin[PIPE_READ] = in[PIPE_READ];
        m_stdin[PIPE_WRITE] = in[PIPE_WRITE];
      } else if (pipe(m_stdin) != 0) {
        throw system_error("Failed to allocate pipe");
      }

      if (out != nullptr) {
        m_stdout[PIPE_READ] = out[PIPE_READ];
        m_stdout[PIPE_WRITE] = out[PIPE_WRITE];
      } else if (pipe(m_stdout) != 0) {
        close(m_stdin[PIPE_READ]);
        close(m_stdin[PIPE_WRITE]);
        throw system_error("Failed to allocate pipe");
      }
    }

    ~command() {
      if (is_running()) {
        m_log.warn("command: Sending SIGKILL to forked process (%d)", m_forkpid);
        kill(m_forkpid, SIGKILL);
      }
      if (m_stdin[PIPE_READ] > 0 && (close(m_stdin[PIPE_READ]) == -1))
        m_log.err("command: Failed to close fd: %s (%d)", strerror(errno), errno);
      if (m_stdin[PIPE_WRITE] > 0 && (close(m_stdin[PIPE_WRITE]) == -1))
        m_log.err("command: Failed to close fd: %s (%d)", strerror(errno), errno);
      if (m_stdout[PIPE_READ] > 0 && (close(m_stdout[PIPE_READ]) == -1))
        m_log.err("command: Failed to close fd: %s (%d)", strerror(errno), errno);
      if (m_stdout[PIPE_WRITE] > 0 && (close(m_stdout[PIPE_WRITE]) == -1))
        m_log.err("command: Failed to close fd: %s (%d)", strerror(errno), errno);
    }

    /**
     * Execute the command
     */
    int exec(bool wait_for_completion = true) {
      if ((m_forkpid = fork()) == -1)
        throw system_error("Failed to fork process");

      if (process_util::in_forked_process(m_forkpid)) {
        if (dup2(m_stdin[PIPE_READ], STDIN_FILENO) == -1)
          throw system_error("Failed to redirect stdin in child process");
        if (dup2(m_stdout[PIPE_WRITE], STDOUT_FILENO) == -1)
          throw system_error("Failed to redirect stdout in child process");
        if (dup2(m_stdout[PIPE_WRITE], STDERR_FILENO) == -1)
          throw system_error("Failed to redirect stderr in child process");

        // Close file descriptors that won't be used by forked process
        if ((m_stdin[PIPE_READ] = close(m_stdin[PIPE_READ])) == -1)
          throw system_error("Failed to close fd");
        if ((m_stdin[PIPE_WRITE] = close(m_stdin[PIPE_WRITE])) == -1)
          throw system_error("Failed to close fd");
        if ((m_stdout[PIPE_READ] = close(m_stdout[PIPE_READ])) == -1)
          throw system_error("Failed to close fd");
        if ((m_stdout[PIPE_WRITE] = close(m_stdout[PIPE_WRITE])) == -1)
          throw system_error("Failed to close fd");

        // Replace the forked process with the given command
        process_util::exec(m_cmd);
        std::exit(0);
      } else {
        // Close file descriptors that won't be used by parent process
        if ((m_stdin[PIPE_READ] = close(m_stdin[PIPE_READ])) == -1)
          throw system_error("Failed to close fd");
        if ((m_stdout[PIPE_WRITE] = close(m_stdout[PIPE_WRITE])) == -1)
          throw system_error("Failed to close fd");

        if (wait_for_completion)
          return wait();
      }

      return EXIT_SUCCESS;
    }

    /**
     * Wait for the child processs to finish
     */
    int wait() {
      do {
        pid_t pid;

        if ((pid = process_util::wait_for_completion(
                 m_forkpid, &m_forkstatus, WCONTINUED | WUNTRACED)) == -1) {
          throw system_error(
              "Process did not finish successfully (" + to_string(m_forkstatus) + ")");
        }

        if (WIFEXITED(m_forkstatus))
          m_log.trace("command: Exited with status %d", WEXITSTATUS(m_forkstatus));
        else if (WIFSIGNALED(m_forkstatus))
          m_log.trace("command: Got killed by signal %d (%s)", WTERMSIG(m_forkstatus),
              strerror(WTERMSIG(m_forkstatus)));
        else if (WIFSTOPPED(m_forkstatus))
          m_log.trace("command: Stopped by signal %d (%s)", WSTOPSIG(m_forkstatus),
              strerror(WSTOPSIG(m_forkstatus)));
        else if (WIFCONTINUED(m_forkstatus) == true)
          m_log.trace("command: Continued");
      } while (!WIFEXITED(m_forkstatus) && !WIFSIGNALED(m_forkstatus));

      return m_forkstatus;
    }

    /**
     * Tail command output
     */
    void tail(function<void(string)> callback) {
      io_util::tail(m_stdout[PIPE_READ], callback);
    }

    /**
     * Write line to command input channel
     */
    int writeline(string data) {
      std::lock_guard<threading_util::spin_lock> lck(m_pipelock);
      return io_util::writeline(m_stdin[PIPE_WRITE], data);
    }

    /**
     * Get command output channel
     */
    int get_stdout(int c) {
      return m_stdout[c];
    }

    /**
     * Get command input channel
     */
    int get_stdin(int c) {
      return m_stdin[c];
    }

    /**
     * Get command pid
     */
    pid_t get_pid() {
      return m_forkpid;
    }

    /**
     * Check if command is running
     */
    bool is_running() {
      return process_util::wait_for_completion_nohang(m_forkpid, &m_forkstatus) > -1;
    }

    /**
     * Get command exit status
     */
    int get_exit_status() {
      return m_forkstatus;
    }

   protected:
    const logger& m_log;

    string m_cmd;

    int m_stdout[2];
    int m_stdin[2];

    pid_t m_forkpid;
    int m_forkstatus;

    threading_util::spin_lock m_pipelock;
  };

  template <typename... Args>
  auto make_command(Args&&... args) {
    return make_unique<command>(
        logger::configure().create<const logger&>(), forward<Args>(args)...);
  }
}

LEMONBUDDY_NS_END