/* SPDX-License-Identifier: MIT */
/* SPDX-FileCopyrightText: (c) Copyright 2024 Andrew Bower <andrew@bower.uk> */

/* xchpst: eXtended Change Process State
 * A tool that is backwards compatible with chpst(8) from runit(8),
 * offering additional options to harden process with namespace isolation
 * and more. */

#include <libgen.h>
#include <time.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>
#include <sched.h>
#include <sys/file.h>
#include <sys/dir.h>
#include <sys/mount.h>
#include <sys/stat.h>
#include <sys/syscall.h>
#include <sys/time.h>

#include "xchpst.h"
#include "rootfs.h"
#include "mount.h"

/* Missing in glibc - see pivot_root(2) */
static int pivot_root(const char *new_root, const char *put_old) {
  return syscall(SYS_pivot_root, new_root, put_old);
}

static const char *special_mount_names[] = {
  [SPECIAL_PROC] = "proc",
};

static const char *exclude_root_dirs[] = {
  ".",
  "..",
  "lost+found",
  NULL,
};

struct mount_info *special_mounts[SPECIAL_MAX];
static struct mount_info *mts = NULL;
static ssize_t num_mounts = 0;

static bool bind_root_dirs(const char *new_root) {
  enum special_mount special;
  const struct dirent *de;
  struct stat statbuf;
  const char *entity;
  const char **filter;
  DIR *dir = NULL;
  char *path = NULL;
  struct mount_info *mt;
  bool success = false;
  int dir1 = -1;
  int dir2 = -1;

  mts = NULL;
  memset(special_mounts, '\0', sizeof special_mounts);

  if ((dir1 = open(entity = "/", O_RDONLY | O_DIRECTORY)) == -1) {
    perror("opening root directory");
    return false;
  }
  if ((dir2 = openat(dir1, "./", O_RDONLY | O_DIRECTORY)) == -1)
    goto fail;
  if (fstat(dir1, &statbuf) == -1) {
    perror("stat");
    goto fail;
  }
  if ((dir = fdopendir(dir1)) == NULL) {
    perror("opening root directory");
    close(dir1);
    return false;
  }

  mts = calloc(statbuf.st_nlink, sizeof *mts);
  if (mts == NULL) {
    perror("calloc");
    goto fail;
  }
  mt = mts;
  for (errno = 0; (de = readdir(dir));) {
    if (asprintf(&mt->from, "/%s", de->d_name) == -1 ||
        asprintf(&mt->to, "%s/%s", new_root, de->d_name) == -1)
      mt->to = NULL;

    if (de->d_type == DT_LNK) {
      int rc;
      path = malloc(PATH_MAX);
      if (path == NULL)
        goto fail;
      rc = readlinkat(dir2, de->d_name, path, PATH_MAX - 1);
      if (rc == -1 || rc >= PATH_MAX - 1)
        goto fail;
      path[rc]='\0';
      if (symlink(path, mt->to) == -1 || is_verbose())
        fprintf(stderr, "  symlink(%s,%s)=%s\n", path, mt->to, strerror(errno));
      free(path);
      path = NULL;
    }

    if (de->d_type != DT_DIR)
      continue;

    for (filter = exclude_root_dirs; *filter && strcmp(de->d_name, *filter); filter++);
    if (*filter)
      continue;

    for (special = 0; special < SPECIAL_MAX && strcmp(de->d_name, special_mount_names[special]); special++);

    if (is_debug()) {
      fprintf(stderr, "binding %s into new rootfs\n", de->d_name);
    }
    mkdir(mt->to, 0700);
    if (special < SPECIAL_MAX) {
      if (is_debug())
        fprintf(stderr, "  found special mount %s\n", de->d_name);
      special_mounts[special] = mt;
    }
    if (special != SPECIAL_PROC || (opt.new_ns & CLONE_NEWPID) == 0) {
      mt->rc = mount(mt->from, mt->to, NULL, MS_BIND | MS_REC | MS_PRIVATE, NULL);
      if (is_debug() || mt->rc != 0)
        fprintf(stderr, "  mount(%s,%s)=%s\n", mt->from, mt->to, strerror(mt->rc == 0 ? 0 : errno));
      if (mt->rc == 0)
        mt->mounted = true;
    }

    mt++;
  }
  num_mounts = mt - mts;

  success = true;
fail:
  if (!success)
    fprintf(stderr, "failed to bind directories into new root, %s\n", strerror(errno));

  if (mts && !success) {
    unmount_temp_rootfs();
    free_rootfs_data();
  }
  free(path);
  closedir(dir);
  if (dir2 != -1)
    close(dir2);
  return success;
}

void unmount_temp_rootfs(void) {
  struct mount_info *mt;

  for (mt = mts; mt - mts < num_mounts; mt++) {
    if (mt->to && mt->mounted)
      umount2(mt->to, MNT_DETACH);
  }
}

void free_rootfs_data(void) {
  struct mount_info *mt;

  if (!mts)
    return;

  for (mt = mts; mt - mts < num_mounts; mt++) {
    free(mt->from);
    free(mt->to);
  }
  free(mts);
  mts = NULL;
}

void close_if_open(int *fd) {
  if (*fd != -1) {
    close(*fd);
    *fd = -1;
  }
}

void pivot_tidy(struct pivot_state *rt) {
  close_if_open(&rt->new_root_fd);
  close_if_open(&rt->old_root_fd);
  close_if_open(&rt->roots_dir_fd);
  if (rt->path) {
    free(rt->path);
    rt->path = NULL;
    rt->leaf_dir = NULL;
  }
}

bool create_new_root(const char *executable,
                     struct pivot_state *rt) {
  struct timeval t = { 0 };
  bool success = false;
  int rc;
  int run_dir_fd;

  run_dir_fd = get_run_dir();
  if (run_dir_fd == -1)
    return false;

  if (ensure_dir(run_dir_fd, "roots", &rt->roots_dir_fd, 0700) == -1)
    return false;

  gettimeofday(&t, NULL);
  /* basename(3) promises that with _GNU_SOURCE defined, its argument is
     unmodified. */
  rc = asprintf(&rt->path, "%s/roots/%lld-%d-%s",
                run_dir,
                (long long) t.tv_sec, getpid(), basename((char *)executable));
  if (rc == -1) {
    perror("formatting new root");
    goto finish;
  }
  rt->leaf_dir = basename(rt->path);
  rc = private_mount(rt->path);
  if (rc == -1) {
    perror("special_mount: root");
    goto finish;
  }
  rc = mount("", rt->path, "", MS_REMOUNT | MS_BIND | MS_REC | MS_PRIVATE, "");
  if (rc == -1) {
    perror("adjusting new root fs mount options");
    goto finish;
  }
  success = bind_root_dirs(rt->path);

finish:
  if (rt->path && !success) {
    umount2(rt->path, MNT_DETACH);
  }
  return success;
}

bool pivot_to_new_root(struct pivot_state *rt) {
  bool success = false;
  int old_dir = -1;
  int rc;
  char *pwd;

  pwd = get_current_dir_name();
  if (pwd == NULL)
    perror("warning: could not save location of starting directory");

  if (chdir(rt->path) == -1)
    perror("chdir to new root");
  else
    success = true;

  old_dir = open("/", O_DIRECTORY | O_RDONLY | O_CLOEXEC);
  if (old_dir == -1) {
    perror("could not open old root for pivot");
    goto finish;
  }

  rc = pivot_root(".", ".");
  if (rc == -1) {
    fprintf(stderr, "could not pivot to new root %s, %s\n",
            rt->path,
            strerror(errno));
    goto finish;
  } else if (is_verbose()) {
    fprintf(stderr, "pivoted new root to %s\n", rt->path);
  }

   if (fchdir(old_dir) == -1) {
    perror("could not open old root to tidy up");
  }

  if (mount("", ".", "", MS_SLAVE | MS_REC, NULL) == -1)
    perror("could not disable old root mount propagation");

  if (umount2(".", MNT_DETACH) == -1)
    perror("umounting old root");

  if (chdir("/") == -1)
    perror("chdir to new root");
  else
    success = true;

  close(old_dir);

  rc = umount2(rt->path, MNT_DETACH);
  if (rc == -1)
    perror("umount2");

  rc = unlinkat(rt->roots_dir_fd, rt->leaf_dir, AT_REMOVEDIR);
  if (rc == -1)
    perror("unlinkat");

  if (pwd != NULL && chdir(pwd) == -1)
    perror("warning: could not change to original directory relative to new root");

finish:
  return success;
}
