subprocess.hpp

Go to the documentation of this file.

 
#pragma once
 
#include <cstring>
#include <functional>
#include <sys/wait.h>
#include <csignal>
#include <vector>
#include <string>
#include <unistd.h>
#include <sys/types.h>
#include <sys/prctl.h>
#include <sys/stat.h>
#include <fcntl.h>
#include <sys/mman.h>
#include <ranges>
#include <filesystem>
#include <utility>
#include <memory>
 
#include "log.hpp"
#include "../macro.hpp"
#include "../std/vector.hpp"
#include "subprocess/pipe.hpp"
#include "subprocess/child.hpp"

namespace ns_subprocess
{

enum class Stream
{
  Inherit,  // Child inherits parent's stdin/stdout/stderr
  Pipe,     // Redirect to pipes with callbacks
  Null,     // Redirect to /dev/null (silent)
};

struct ArgsCallbackChild
{
  pid_t parent_pid;  // Parent process PID
  int stdin_fd;      // Standard input file descriptor (usually 0)
  int stdout_fd;     // Standard output file descriptor (usually 1)
  int stderr_fd;     // Standard error file descriptor (usually 2)
};

struct ArgsCallbackParent
{
  pid_t child_pid;   // Child process PID
};

namespace stream
{

inline std::fstream& null()
{
  static std::fstream null("/dev/null", std::ios::in | std::ios::out);
  return null;
}
 
}

inline std::unique_ptr<const char*[]> to_carray(std::vector<std::string> const& vec)
{
  auto arr = std::make_unique<const char*[]>(vec.size() + 1);
  std::ranges::transform(vec, arr.get(), [](auto const& e) { return e.c_str(); });
  arr[vec.size()] = nullptr;
  return arr;
}
 
class Subprocess
{
  private:
    std::filesystem::path m_program;
    std::vector<std::string> m_args;
    std::vector<std::string> m_env;
    std::reference_wrapper<std::istream> m_stdin;
    std::reference_wrapper<std::ostream> m_stdout;
    std::reference_wrapper<std::ostream> m_stderr;
    Stream m_stream_mode;
    std::optional<pid_t> m_die_on_pid;
    std::filesystem::path m_path_file_log;
    ns_log::Level m_log_level;
    std::optional<std::function<void(ArgsCallbackChild)>> m_callback_child;
    std::optional<std::function<void(ArgsCallbackParent)>> m_callback_parent;
    bool m_daemon_mode;
 
    void die_on_pid(pid_t pid);
    void to_dev_null();
    std::vector<std::jthread> setup_pipes(pid_t child_pid
      , int pipestdin[2]
      , int pipestdout[2]
      , int pipestderr[2]
      , std::filesystem::path const& path_file_log);
    [[noreturn]] void exec_child();
 
  public:
    template<ns_concept::StringRepresentable T>
    [[nodiscard]] Subprocess(T&& t);
    ~Subprocess();
 
    Subprocess(Subprocess const&) = delete;
    Subprocess& operator=(Subprocess const&) = delete;
 
    Subprocess(Subprocess&&) = delete;
    Subprocess& operator=(Subprocess&&) = delete;
 
 
    [[maybe_unused]] [[nodiscard]] Subprocess& env_clear();
 
    template<ns_concept::StringRepresentable K, ns_concept::StringRepresentable V>
    [[maybe_unused]] [[nodiscard]] Subprocess& with_var(K&& k, V&& v);
 
    template<ns_concept::StringRepresentable K>
    [[maybe_unused]] [[nodiscard]] Subprocess& rm_var(K&& k);
 
    template<typename... Args>
    [[maybe_unused]] [[nodiscard]] Subprocess& with_args(Args&&... args);
 
    template<typename... Args>
    [[maybe_unused]] [[nodiscard]] Subprocess& with_env(Args&&... args);
 
    [[maybe_unused]] [[nodiscard]] Subprocess& with_die_on_pid(pid_t pid);
 
    [[maybe_unused]] [[nodiscard]] Subprocess& with_stdio(Stream mode);
 
    [[maybe_unused]] [[nodiscard]] Subprocess& with_log_file(std::filesystem::path const& path);
 
    [[maybe_unused]] [[nodiscard]] Subprocess& with_log_level(ns_log::Level const& level);
 
    [[maybe_unused]] [[nodiscard]] Subprocess& with_streams(std::istream& stdin_stream, std::ostream& stdout_stream, std::ostream& stderr_stream);
 
    template<typename F>
    [[maybe_unused]] [[nodiscard]] Subprocess& with_callback_child(F&& f);
 
    template<typename F>
    [[maybe_unused]] [[nodiscard]] Subprocess& with_callback_parent(F&& f);
 
    [[maybe_unused]] [[nodiscard]] Subprocess& with_daemon();
 
    [[maybe_unused]] [[nodiscard]] std::unique_ptr<Child> spawn();
};

template<ns_concept::StringRepresentable T>
Subprocess::Subprocess(T&& t)
  : m_program(ns_string::to_string(t))
  , m_args()
  , m_env()
  , m_stdin(stream::null())
  , m_stdout(stream::null())
  , m_stderr(stream::null())
  , m_stream_mode(Stream::Inherit)
  , m_die_on_pid(std::nullopt)
  , m_path_file_log("/dev/null")
  , m_log_level(ns_log::get_level())
  , m_callback_child(std::nullopt)
  , m_callback_parent(std::nullopt)
  , m_daemon_mode(false)
{
  // argv0 is program name
  m_args.push_back(m_program);
  // Copy environment
  for(char** i = environ; *i != nullptr; ++i)
  {
    m_env.push_back(*i);
  } // for
}

inline Subprocess::~Subprocess()
{
}

inline Subprocess& Subprocess::env_clear()
{
  m_env.clear();
  return *this;
}

template<ns_concept::StringRepresentable K, ns_concept::StringRepresentable V>
Subprocess& Subprocess::with_var(K&& k, V&& v)
{
  std::ignore = rm_var(k);
  m_env.push_back(std::format("{}={}", k, v));
  return *this;
}

template<ns_concept::StringRepresentable K>
Subprocess& Subprocess::rm_var(K&& k)
{
  // Find variable
  auto it = std::ranges::find_if(m_env, [&](std::string const& e)
  {
    auto vec = ns_vector::from_string(e, '=');
    return_if(vec.empty(), false);
    return vec.front() == k;
  });
 
  // Erase if found
  if ( it != std::ranges::end(m_env) )
  {
    logger("D::Erased var entry: {}", *it);
    m_env.erase(it);
  } // if
 
  return *this;
}

template<typename... Args>
Subprocess& Subprocess::with_args(Args&&... args)
{
  // Helper lambda to add a single argument
  auto add_arg = [this]<typename T>(T&& arg) -> void
  {
    if constexpr ( ns_concept::Uniform<T, std::string> )
    {
      this->m_args.push_back(std::forward<T>(arg));
    }
    else if constexpr ( ns_concept::Container<T> )
    {
      std::copy(arg.begin(), arg.end(), std::back_inserter(m_args));
    }
    else if constexpr ( ns_concept::StringRepresentable<T> )
    {
      this->m_args.push_back(ns_string::to_string(std::forward<T>(arg)));
    }
    else
    {
      static_assert(false, "Could not determine argument type");
    }
  };
 
  // Process each argument
  (add_arg(std::forward<Args>(args)), ...);
 
  return *this;
}

template<typename... Args>
Subprocess& Subprocess::with_env(Args&&... args)
{
  // Helper to erase existing entries by key
  auto f_erase_existing = [this](auto&& entries)
  {
    for (auto&& entry : entries)
    {
      auto parts = ns_vector::from_string(entry, '=');
      continue_if(parts.size() < 2, "E::Entry '{}' is not valid", entry);
      std::string key = parts.front();
      std::ignore = this->rm_var(key);
    }
  };
 
  // Helper lambda to add a single env entry or container
  auto add_env = [this, &f_erase_existing]<typename T>(T&& arg) -> void
  {
    if constexpr ( ns_concept::Uniform<std::remove_cvref_t<T>, std::string> )
    {
      f_erase_existing(std::vector<std::string>{arg});
      this->m_env.push_back(std::forward<T>(arg));
    }
    else if constexpr ( ns_concept::Iterable<T> )
    {
      f_erase_existing(arg);
      std::ranges::copy(arg, std::back_inserter(m_env));
    }
    else if constexpr ( ns_concept::StringRepresentable<T> )
    {
      auto entry = ns_string::to_string(std::forward<T>(arg));
      f_erase_existing(std::vector<std::string>{entry});
      this->m_env.push_back(entry);
    }
    else
    {
      static_assert(false, "Could not determine argument type");
    }
  };
 
  // Process each argument
  (add_env(std::forward<Args>(args)), ...);
 
  return *this;
}

inline Subprocess& Subprocess::with_die_on_pid(pid_t pid)
{
  m_die_on_pid = pid;
  return *this;
}

inline Subprocess& Subprocess::with_stdio(Stream mode)
{
  m_stream_mode = mode;
  return *this;
}

inline Subprocess& Subprocess::with_log_file(std::filesystem::path const& path)
{
  m_path_file_log = path;
  // Set up pipe mode - pipe readers will configure logger sink to the specified path
  return this->with_stdio(Stream::Pipe);
}

inline Subprocess& Subprocess::with_log_level(ns_log::Level const& level)
{
  m_log_level = level;
  return *this;
}

inline void Subprocess::die_on_pid(pid_t pid)
{
  // Set death signal when pid dies
  return_if(prctl(PR_SET_PDEATHSIG, SIGKILL) < 0,,"E::Failed to set PR_SET_PDEATHSIG: {}", strerror(errno));
  // Abort if pid is not running
  if (::kill(pid, 0) < 0)
  {
    logger("E::Parent died, prctl will not have effect: {}", strerror(errno));
    _exit(1);
  } // if
  // Log pid and current pid
  logger("D::{} dies with {}", getpid(), pid);
}

inline void Subprocess::to_dev_null()
{
  int fd = open("/dev/null", O_WRONLY);
  return_if(fd < 0,,"E::Failed to open /dev/null: {}", strerror(errno));
  return_if(dup2(fd, STDIN_FILENO) < 0,,"E::Failed to redirect stdin: {}", strerror(errno));
  return_if(dup2(fd, STDOUT_FILENO) < 0,,"E::Failed to redirect stdout: {}", strerror(errno));
  return_if(dup2(fd, STDERR_FILENO) < 0,,"E::Failed to redirect stderr: {}", strerror(errno));
  close(fd);
}

inline std::vector<std::jthread> Subprocess::setup_pipes(pid_t child_pid
  , int pipestdin[2]
  , int pipestdout[2]
  , int pipestderr[2]
  , std::filesystem::path const& log)
{
  // Setup pipes for parent or child and return threads
  return ns_pipe::setup(child_pid
    , pipestdin
    , pipestdout
    , pipestderr
    , m_stdin.get()
    , m_stdout.get()
    , m_stderr.get()
    , log
  );
}

[[noreturn]] inline void Subprocess::exec_child()
{
  // Create null-terminated arrays for execve
  auto argv_custom = to_carray(m_args);
  auto envp_custom = to_carray(m_env);
 
  // Perform execve
  execve(m_program.c_str(), const_cast<char**>(argv_custom.get()), const_cast<char**>(envp_custom.get()));
 
  // Log error (will go to log_file if redirected)
  logger("E::execve() failed: {}", strerror(errno));
 
  // Child should stop here
  _exit(1);
}

inline Subprocess& Subprocess::with_streams(std::istream& stdin_stream, std::ostream& stdout_stream, std::ostream& stderr_stream)
{
  this->m_stdin = stdin_stream;
  this->m_stdout = stdout_stream;
  this->m_stderr = stderr_stream;
  return *this;
}

template<typename F>
Subprocess& Subprocess::with_callback_child(F&& f)
{
  this->m_callback_child = std::forward<F>(f);
  return *this;
}

template<typename F>
inline Subprocess& Subprocess::with_callback_parent(F&& f)
{
  m_callback_parent = std::forward<F>(f);
  return *this;
}

inline Subprocess& Subprocess::with_daemon()
{
  m_daemon_mode = true;
  return *this;
}

inline std::unique_ptr<Child> Subprocess::spawn()
{
  // Ignore on empty vec_argv
  return_if(m_args.empty(), Child::create(-1, m_program), "E::No arguments to spawn subprocess");
  logger("D::Spawn command: {}", m_args);
 
  // Create pipes BEFORE fork (if needed for Stream::Pipe)
  int pipestdin[2];
  int pipestdout[2];
  int pipestderr[2];
 
  if ( m_stream_mode == Stream::Pipe )
  {
    return_if(pipe(pipestdin), Child::create(-1, m_program), "E::{}", strerror(errno));
    return_if(pipe(pipestdout), Child::create(-1, m_program), "E::{}", strerror(errno));
    return_if(pipe(pipestderr), Child::create(-1, m_program), "E::{}", strerror(errno));
  }
 
  // Create shared memory for daemon grandchild PID (if daemon mode)
  pid_t* grandchild_pid_ptr = nullptr;
  if ( m_daemon_mode )
  {
    grandchild_pid_ptr = static_cast<pid_t*>(mmap(
      nullptr,
      sizeof(pid_t),
      PROT_READ | PROT_WRITE,
      MAP_SHARED | MAP_ANONYMOUS,
      -1,
      0
    ));
    return_if(grandchild_pid_ptr == MAP_FAILED, nullptr, "E::mmap failed: {}", strerror(errno));
    *grandchild_pid_ptr = -1;  // Initialize
  }
 
  // Create child
  pid_t pid = fork();
 
  // Failed to fork
  if (pid < 0)
  {
    if (grandchild_pid_ptr != nullptr)
    {
      munmap(grandchild_pid_ptr, sizeof(pid_t));
    }
    logger("E::Failed to fork");
    return Child::create(-1, m_program);
  }
 
  // Parent returns here, child continues to execve
  if ( pid > 0 )
  {
    std::vector<std::jthread> pipe_threads;
 
    // If daemon mode, wait for intermediate child to exit and read grandchild PID
    if ( m_daemon_mode )
    {
      int status;
      log_if(waitpid(pid, &status, 0) < 0, "E::Waitpid failed: {}", strerror(errno));
      logger("D::Daemon mode: intermediate process exited");
 
      // Read grandchild PID from shared memory
      pid_t pid_grandchild = *grandchild_pid_ptr;
 
      // Cleanup shared memory
      munmap(grandchild_pid_ptr, sizeof(pid_t));
 
      // Return Child handle with grandchild PID
      logger("D::Daemon grandchild PID: {}", pid_grandchild);
      return Child::create(pid_grandchild, m_program);
    }
    // Setup pipes for parent and capture threads
    else if ( m_stream_mode == Stream::Pipe)
    {
      pipe_threads = this->setup_pipes(pid, pipestdin, pipestdout, pipestderr, m_path_file_log);
      logger("D::Parent pipes configured");
    }
 
    // Execute parent callback if provided
    if (m_callback_parent)
    {
      ArgsCallbackParent args{.child_pid = pid};
      m_callback_parent.value()(args);
    }
 
    // Return Child handle with process and transfer pipe threads ownership
    return Child::create(pid, m_program, std::move(pipe_threads));
  }
 
  // Child process continues here (intermediate child in daemon mode, final child otherwise)
 
  // Daemon mode: perform second fork to create grandchild
  if ( m_daemon_mode )
  {
    // Create session leader to detach from controlling terminal
    if ( setsid() < 0 )
    {
      logger("E::setsid() failed: {}", strerror(errno));
      _exit(1);
    }
 
    // Second fork - creates grandchild
    pid = fork();
 
    if ( pid < 0 )
    {
      logger("E::Second fork failed: {}", strerror(errno));
      _exit(1);
    }
 
    if ( pid > 0 )
    {
      // Intermediate child: write grandchild PID to shared memory
      *grandchild_pid_ptr = pid;
 
      // Intermediate child exits immediately
      // This causes grandchild to be orphaned and adopted by init
      _exit(0);
    }
 
    // Grandchild continues here - now a daemon
    // Cleanup shared memory in grandchild (no longer needed)
    munmap(grandchild_pid_ptr, sizeof(pid_t));
 
    // Change working directory to root to avoid keeping directories busy
    if ( chdir("/") < 0 )
    {
      logger("W::chdir() to / failed: {}", strerror(errno));
    }
 
    // Reset file mode creation mask
    umask(0);
  }
  // Setup pipes for child
  else if ( m_stream_mode == Stream::Pipe )
  {
    this->setup_pipes(pid, pipestdin, pipestdout, pipestderr, m_path_file_log);
    logger("D::child pipes configured", m_daemon_mode);
  }
 
  ns_log::set_level(m_log_level);
 
  // Handle stdio redirection based on mode
  // Daemon should not inherit the parent's IO
  if (m_stream_mode == Stream::Null or m_daemon_mode)
  {
    this->to_dev_null();
  }
  // Stream::Pipe: child's stdout/stderr redirected by pipes_child() via setup_pipes()
  // Stream::Inherit: do nothing, child keeps parent's stdio
 
  // Die with pid
  if(m_die_on_pid)
  {
    this->die_on_pid(m_die_on_pid.value());
  }
 
  // Execute child callback if provided
  if (m_callback_child)
  {
    ArgsCallbackChild args
    {
      .parent_pid = getppid(),
      .stdin_fd = STDIN_FILENO,
      .stdout_fd = STDOUT_FILENO,
      .stderr_fd = STDERR_FILENO
    };
    m_callback_child.value()(args);
  }
 
  // Execute child process (never returns)
  this->exec_child();
}
 
} // namespace ns_subprocess
 
/* vim: set expandtab fdm=marker ts=2 sw=2 tw=100 et :*/