/* 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 struct {
  char *name;
  int ns;
  char *fs;
} special_mount_id[] = {
  [SPECIAL_PROC] = { "proc", CLONE_NEWPID, "proc"  },
  [SPECIAL_MAX]  = { NULL, 0, NULL }
};

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;
  const char **filter;
  DIR *dir = NULL;
  char *path = NULL;
  struct mount_info *mt;
  bool success = false;

  memset(special_mounts, '\0', sizeof special_mounts);

  if ((dir = opendir("/")) == NULL) {
    perror("opening root directory");
    return false;
  }

  for (num_mounts = 0; (de = readdir(dir)); num_mounts++);
  rewinddir(dir);
  mts = calloc(num_mounts + 1, sizeof *mts);

  for (mt = mts; mt - mts < num_mounts && (de = readdir(dir)); mt++) {
    if (asprintf(&mt->from, "/%s", de->d_name) == -1 ||
        asprintf(&mt->to, "%s/%s", new_root, de->d_name) == -1)
      goto fail;

    if (de->d_type == DT_LNK) {
      int rc;
      path = malloc(PATH_MAX);
      if (path == NULL)
        goto fail;
      rc = readlinkat(dirfd(dir), de->d_name, path, PATH_MAX - 1);
      if (rc == -1 || rc >= PATH_MAX - 1)
        goto fail;
      path[rc]='\0';
      errno = 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_id[special].name); special++);

    if (is_debug()) {
      fprintf(stderr, "creating mount point %s in 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 (opt.new_ns & special_mount_id[special].ns) {
      free(mt->from);
      mt->from = strdup(special_mount_id[special].name);
      mt->rc = special_mount(mt->to, special_mount_id[special].fs, mt->from, NULL);
    } else {
      mt->rc = mount(mt->from, mt->to, NULL, MS_BIND | MS_REC | MS_SLAVE, 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;
  }

  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);
  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_SLAVE, "");
  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;

  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_REC | MS_SLAVE, 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);
  old_dir = -1;

  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");

finish:
  if (old_dir != -1)
    close(old_dir);

  return success;
}