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

DEFINE_ERROR(command_error);
DEFINE_CHILD_ERROR(command_strerror, command_error);

namespace command_util {
  /**
   * Wrapper used to execute command in a subprocess.
   * In-/output streams are opened to enable ipc.
   *
   * Example usage:
   *
   * @code cpp
   *   auto cmd = command_util::make_command("cat /etc/rc.local");
   *   cmd->exec();
   *   cmd->tail([](string s) { std::cout << s << std::endl; });
   * @endcode
   *
   * @code cpp
   *   auto cmd = command_util::make_command(
   *    "/bin/sh\n-c\n while read -r line; do echo data from parent process: $line; done");
   *   cmd->exec(false);
   *   cmd->writeline("Test");
   *   cout << cmd->readline();
   *   cmd->wait();
   * @endcode
   *
   * @code cpp
   *   vector<string> exec{{"/bin/sh"}, {"-c"}, {"for i in 1 2 3; do echo $i; done"}};
   *   auto cmd = command_util::make_command(exec);
   *   cmd->exec();
   *   cout << cmd->readline(); // 1
   *   cout << cmd->readline() << cmd->readline(); // 23
   * @endcode
   */
  class command {
   public:
    explicit command(const logger& logger, vector<string> cmd)
        : command(logger, string_util::join(cmd, "\n")) {}

    explicit command(const logger& logger, string cmd) : m_log(logger), m_cmd(cmd) {
      if (pipe(m_stdin) != 0)
        throw command_strerror("Failed to allocate input stream");
      if (pipe(m_stdout) != 0)
        throw command_strerror("Failed to allocate output stream");
    }

    ~command() {
      if (is_running())
        terminate();

      if (m_stdin[PIPE_READ] > 0)
        close(m_stdin[PIPE_READ]);
      if (m_stdin[PIPE_WRITE] > 0)
        close(m_stdin[PIPE_WRITE]);
      if (m_stdout[PIPE_READ] > 0)
        close(m_stdout[PIPE_READ]);
      if (m_stdout[PIPE_WRITE] > 0)
        close(m_stdout[PIPE_WRITE]);
    }

    /**
     * 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 command_strerror("Failed to redirect stdin in child process");
        if (dup2(m_stdout[PIPE_WRITE], STDOUT_FILENO) == -1)
          throw command_strerror("Failed to redirect stdout in child process");
        if (dup2(m_stdout[PIPE_WRITE], STDERR_FILENO) == -1)
          throw command_strerror("Failed to redirect stderr in child process");

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

        // Make sure SIGTERM is raised
        process_util::unblock_signal(SIGTERM);

        setpgid(m_forkpid, 0);
        process_util::exec(m_cmd);

        throw command_error("Exec failed");
      } else {
        // Close file descriptors that won't be used by the parent
        if ((m_stdin[PIPE_READ] = close(m_stdin[PIPE_READ])) == -1)
          throw command_strerror("Failed to close fd");
        if ((m_stdout[PIPE_WRITE] = close(m_stdout[PIPE_WRITE])) == -1)
          throw command_strerror("Failed to close fd");

        if (wait_for_completion) {
          auto status = wait();
          m_forkpid = -1;
          return status;
        }
      }

      return EXIT_SUCCESS;
    }

    void terminate() {
      try {
        if (is_running()) {
          m_log.trace("command: Sending SIGTERM to running child process (%d)", m_forkpid);
          killpg(m_forkpid, SIGTERM);
          wait();
        }
      } catch (const command_error& err) {
        m_log.warn("%s", err.what());
      }

      m_forkpid = -1;
    }

    /**
     * Wait for the child processs to finish
     */
    int wait() {
      auto waitflags = WCONTINUED | WUNTRACED;

      do {
        process_util::wait_for_completion(m_forkpid, &m_forkstatus, waitflags);

        if (WIFEXITED(m_forkstatus) && m_forkstatus > 0)
          m_log.warn("command: Exited with failed status %d", WEXITSTATUS(m_forkstatus));
        else if (WIFEXITED(m_forkstatus))
          m_log.trace("command: Exited with status %d", WEXITSTATUS(m_forkstatus));
        else if (WIFSIGNALED(m_forkstatus))
          m_log.trace("command: killed by signal %d", WTERMSIG(m_forkstatus));
        else if (WIFSTOPPED(m_forkstatus))
          m_log.trace("command: Stopped by signal %d", 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);
    }

    /**
     * Read a line from the commands output stream
     */
    string readline() {
      std::lock_guard<threading_util::spin_lock> lck(m_pipelock);
      return io_util::readline(m_stdout[PIPE_READ]);
    }

    /**
     * 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() {
      if (m_forkpid > 0)
        return process_util::wait_for_completion_nohang(m_forkpid, &m_forkstatus) > -1;
      return false;
    }

    /**
     * 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;
  };

  using command_t = unique_ptr<command>;

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

LEMONBUDDY_NS_END