// we need this so sched.h exports unshare and CLONE_*
#define _GNU_SOURCE
#include <sched.h>
#include <errno.h>
#include <err.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <sys/types.h>
#include <sys/wait.h>
#include <unistd.h>
#include <signal.h>

static
pid_t pid_child;

static
void drop_root(void)
{
  /// Drop root privileges
  // First group then user because we might not
  // be able to drop group once we dropped user
  gid_t gid = getgid();
  if (setresgid(-1,gid,gid) == -1)
    err(errno, "Failed to drop root privileges with setresgid");

  uid_t uid = getuid();
  if (setresuid(-1,uid,uid) == -1)
    err(errno, "Failed to drop root privileges with setresuid");

  // sanity check
  if (seteuid(0) != -1)
    errx(1, "Sanity check failed. Able to regain root");
}

static
void forward_signal(int sig)
{
  if (kill(pid_child, sig) == -1)
    {
      if (sig == SIGTERM)
        exit(1);
    }
}

int main(int argc, char* const* argv)
{
  struct sigaction forward_signal_descriptor;
  forward_signal_descriptor.sa_flags = SA_RESTART;
  forward_signal_descriptor.sa_handler = &forward_signal;

  if (argc == 1)
    {
      fprintf(stderr,"Usage: %s PROGRAM ARGUMENTS...\n"
              "Run command within its own pid namespace. Integrated init process.\n",
              argv[0]);
      return 0;
    }

  // next fork shall be in a new pid namespace
  if (unshare(CLONE_NEWPID) != 0)
    {
      err(errno, "Failed to unshare pid namespace");
    }

  // Drop root privileges, we only needed those for the unshare call.
  drop_root();

  pid_t pid = fork();

  if (pid == -1)
    {
      err(errno, "Failed to fork");
    }

  if (pid != 0)
    {
      /// Head process
      // Wait for the init process in the PID namespace to terminate and forward its exit code.
      // Also forward SIGTERM signals towards that init process.

      // Setup signal handler to forward SIGTERM
      pid_child = pid;
      if (sigaction(SIGTERM, &forward_signal_descriptor, NULL) == -1)
        {
          int saved_errno = errno;
          // Have to kill child here, otherwise that gets orphaned and runs anyway.
          // Use SIGKILL here because it might forward SIGTERM to its child and that
          // decides not to stop.
          kill(pid_child, SIGKILL);
          // Restore errno as it might've been overwritten by kill
          errno = saved_errno;
          err(errno, "Unable to set up signal handler in head process");
        }
      // parent waits for child then exits
      // Could be interrupt due to a signal. Retry in that case.
      int status;
      if (waitpid(pid, &status, 0) == -1)
        {
          err(errno, "Failed to wait for init process");
        }

      return WEXITSTATUS(status);
    }
  else
    {
      // Child should be in new pid namespace and
      // functions as the the init process
      // it needs to fork again then wait for any child.
      // if the forked child exits then exit.

      pid = fork();
      if (pid == -1)
        {
          err(errno, "Failed to fork in init process");
        }

      if (pid != 0)
        {
          /// Init process
          // This part of the program runs as first process in the pid namespace
          // When this terminates then Linux terminates all remaining processes
          // in the PID namespace. As we want this to happen when our first child
          // terminates, we wait for our first child to terminate before terminating
          // ourselves.
          // As first process in the PID namespace, this also functions as adopting parent
          // for orphaned processes in the PID namespace and therefore has to wait for
          // any child process and then check if the a child process that has terminated
          // is the one we were waiting for.

          pid_t first_child = pid;
          pid_t exited_child;
          int child_status;

          // Setup forward for SIGTERM
          pid_child = first_child;
          if (sigaction(SIGTERM, &forward_signal_descriptor, NULL) == -1)
            {
              err(1, "Unable to setup signal forward in init");
            }

          int wait_errno;
          // wait could be interrupt due to a signal. In that case just call wait again.
          do {
            exited_child = wait(&child_status);
            wait_errno = errno;
          } while (!(exited_child == first_child || (exited_child == -1 && wait_errno == ECHILD)));

          if (exited_child == -1)
            {
              err(wait_errno, "Error while waiting for subprocess");
            }
          else
            {
              return WEXITSTATUS(child_status);
            }
        }
      else
        {
          // First child of init process. do exec here
          // use cli arguments for subprocess. skip 0 as it's our programs name.
          argv++;

          if (execvp(argv[0], argv) == -1)
            {
              err(errno, "Failed to exec");
            }
        }
    }
}