#!/bin/bash
# SPDX-License-Identifier: GPL-2.0-only

shopt -s extglob

use_systemd=0
unshare=0
keepresolvconf=0

#!/hint/bash
# SPDX-License-Identifier: GPL-2.0-only
# shellcheck disable=SC2059 # $1 and $2 can contain the printf modifiers

out() { printf "$1 $2\n" "${@:3}"; }
error() { out "==> ERROR:" "$@"; } >&2
warning() { out "==> WARNING:" "$@"; } >&2
msg() { out "==>" "$@"; }
die() { error "$@"; exit 1; }

ignore_error() {
  "$@" 2>/dev/null
  return 0
}

chroot_add_mount() {
  mount "$@" && CHROOT_ACTIVE_MOUNTS=("$2" "${CHROOT_ACTIVE_MOUNTS[@]}")
}

chroot_maybe_add_mount() {
  local cond=$1; shift
  if eval "$cond"; then
    chroot_add_mount "$@"
  fi
}

chroot_init() {
  CHROOT_ACTIVE_MOUNTS=()
  CHROOT_ACTIVE_LAZY=()
  CHROOT_ACTIVE_FILES=()
  [[ $(trap -p EXIT) ]] && die '(BUG): attempting to overwrite existing EXIT trap'
  trap 'chroot_teardown' EXIT
}

chroot_teardown() {
  if (( ${#CHROOT_ACTIVE_MOUNTS[@]} )); then
    umount "${CHROOT_ACTIVE_MOUNTS[@]}"
  fi
  unset CHROOT_ACTIVE_MOUNTS

  if (( ${#CHROOT_ACTIVE_LAZY[@]} )); then
    umount --lazy "${CHROOT_ACTIVE_LAZY[@]}"
  fi
  unset CHROOT_ACTIVE_LAZY

  if (( ${#CHROOT_ACTIVE_FILES[@]} )); then
    rm "${CHROOT_ACTIVE_FILES[@]}"
  fi
  unset CHROOT_ACTIVE_FILES
}

chroot_setup() {
  chroot_add_mount proc "$1/proc" -t proc -o nosuid,noexec,nodev &&
  chroot_add_mount sys "$1/sys" -t sysfs -o nosuid,noexec,nodev,ro &&
  ignore_error chroot_maybe_add_mount "[[ -d '$1/sys/firmware/efi/efivars' ]]" \
      efivarfs "$1/sys/firmware/efi/efivars" -t efivarfs -o nosuid,noexec,nodev &&
  chroot_add_mount udev "$1/dev" -t devtmpfs -o mode=0755,nosuid &&
  chroot_add_mount devpts "$1/dev/pts" -t devpts -o mode=0620,gid=5,nosuid,noexec &&
  chroot_add_mount shm "$1/dev/shm" -t tmpfs -o mode=1777,nosuid,nodev &&
  chroot_add_mount /run "$1/run" --bind --make-private &&
  chroot_add_mount tmp "$1/tmp" -t tmpfs -o mode=1777,strictatime,nodev,nosuid
}

chroot_add_mount_lazy() {
  mount "$@" && CHROOT_ACTIVE_LAZY=("$2" "${CHROOT_ACTIVE_LAZY[@]}")
}

chroot_bind_device() {
  touch "$2" && CHROOT_ACTIVE_FILES=("$2" "${CHROOT_ACTIVE_FILES[@]}")
  chroot_add_mount "$1" "$2" --bind
}

chroot_add_link() {
  ln -sf "$1" "$2" && CHROOT_ACTIVE_FILES=("$2" "${CHROOT_ACTIVE_FILES[@]}")
}

unshare_setup() {
  chroot_add_mount_lazy "$1" "$1" --bind &&
  chroot_add_mount proc "$1/proc" -t proc -o nosuid,noexec,nodev &&
  chroot_add_mount_lazy /sys "$1/sys" --rbind &&
  chroot_add_link /proc/self/fd "$1/dev/fd" &&
  chroot_add_link /proc/self/fd/0 "$1/dev/stdin" &&
  chroot_add_link /proc/self/fd/1 "$1/dev/stdout" &&
  chroot_add_link /proc/self/fd/2 "$1/dev/stderr" &&
  chroot_bind_device /dev/full "$1/dev/full" &&
  chroot_bind_device /dev/null "$1/dev/null" &&
  chroot_bind_device /dev/random "$1/dev/random" &&
  chroot_bind_device /dev/tty "$1/dev/tty" &&
  chroot_bind_device /dev/urandom "$1/dev/urandom" &&
  chroot_bind_device /dev/zero "$1/dev/zero" &&
  chroot_add_mount run "$1/run" -t tmpfs -o nosuid,nodev,mode=0755 &&
  chroot_add_mount tmp "$1/tmp" -t tmpfs -o mode=1777,strictatime,nodev,nosuid
}

pid_unshare="unshare --fork --pid"
mount_unshare="$pid_unshare --mount --map-auto --map-root-user --setuid 0 --setgid 0"

# This outputs code for declaring all variables to stdout. For example, if
# FOO=BAR, then running
#     declare -p FOO
# will result in the output
#     declare -- FOO="bar"
# This function may be used to re-declare all currently used variables and
# functions in a new shell.
declare_all() {
  # Remove read-only variables to avoid warnings. Unfortunately, declare +r -p
  # doesn't work like it looks like it should (declaring only read-write
  # variables). However, declare -rp will print out read-only variables, which
  # we can then use to remove those definitions.
  declare -p | grep -Fvf <(declare -rp)
  # Then declare functions
  declare -pf
}


usage() {
  cat <<EOF
usage: ${0##*/} chroot-dir [command] [arguments...]

    -h                  Print this help message
    -S                  Perform chroot using systemd-run
    -N                  Run in unshare mode as a regular user
    -u <user>[:group]   Specify non-root user and optional group to use
    -r                  Do not change the resolv.conf within the chroot

If 'command' is unspecified, ${0##*/} will launch /bin/bash.

Note that when using arch-chroot, the target chroot directory *should* be a
mountpoint. This ensures that tools such as pacman(8) or findmnt(8) have an
accurate hierarchy of the mounted filesystems within the chroot.

If your chroot target is not a mountpoint, you can bind mount the directory on
itself to make it a mountpoint, i.e. 'mount --bind /your/chroot /your/chroot'.

EOF
}

chroot_add_resolv_conf() {
  local dest="$chrootdir/etc/resolv.conf"

  # If we don't have a source resolv.conf file, there's nothing useful we can do.
  [[ -e /etc/resolv.conf ]] || return 0

  if [[ ! -e "$dest" && ! -L "$dest" ]]; then
    # There may be no resolv.conf in the chroot. In this case, we'll just exit.
    # The chroot environment must not be concerned with DNS resolution.
    return 0
  fi

  chroot_add_mount /etc/resolv.conf "$dest" --bind -o X-mount.nocanonicalize=target
}

systemd_sanity_check() {
  local systemd_ver

  command -v systemd-run &>/dev/null || die 'systemd-run is not available'

  IFS=' .-~^' read -r _ systemd_ver _ < <(SYSTEMD_COLORS=0 systemd-run --version)
  (( systemd_ver >= 257 )) || die 'Unsupported systemd version: %s (requires >= 257)' "$systemd_ver"

  [[ -d /run/systemd/system ]] || die 'System is not booted with systemd'
}

systemd_mode() {
  local user=root group="" sd_args

  systemd_sanity_check

  # NB: we always have systemd spawn the service main process as root
  # and call into setpriv later. The reason is two-fold:
  # systemd resolves users using the host's passwd/groups database,
  # while we want to resolve inside the chroot. Plus, setpriv
  # does us a favor of clearing extraneous envvars from systemd,
  # notably SYSTEMD_EXEC_PID which might confuse some programs.

  if [[ $userspec ]]; then
    IFS=':' read -r user group <<< "$userspec"
    [[ $user ]] || die 'Invalid -u argument: %s' "$userspec"
  fi

  sd_args=(systemd-run --unit="arch-chroot-$$_$(systemd-escape --path "$chrootdir")"
                       --description="[arch-chroot] $chrootdir (as $user)"
                       --service-type=exec --collect
                       --pty --pipe --send-sighup
                       --expand-environment=no
                       -p RootDirectory="$chrootdir"
                       -p MountAPIVFS=yes
                       -p ProtectControlGroupsEx=private
                       -p BindLogSockets=no
                       -p BindReadOnlyPaths=/run/udev # https://github.com/archlinux/arch-install-scripts/pull/52
                       -p TemporaryFileSystem=/tmp
                       -p IgnoreSIGPIPE=no)

  if ! [[ -t 0 && -t 1 ]]; then
    sd_args+=(--quiet)
  fi

  sd_args+=(setpriv --reset-env --reuid="$user" --regid="${group:-"$user"}" --init-groups)

  if [[ ! -v args ]]; then
    args=(/bin/bash -i)
  fi

  "${sd_args[@]}" -- "${args[@]}"
}

offline_mode() {
  local chroot_args=()

  $setup "$chrootdir" || die "failed to setup chroot %s" "$chrootdir"

  [[ $userspec ]] && chroot_args+=(--userspec "$userspec")

  SHELL=/bin/bash SYSTEMD_IN_CHROOT=1 $pid_unshare chroot "${chroot_args[@]}" -- "$chrootdir" "${args[@]}"
}

arch-chroot() {
  (( EUID == 0 )) || die 'This script must be run with root privileges'

  [[ -d $chrootdir ]] || die "Can't create chroot on non-directory %s" "$chrootdir"

  if ! mountpoint -q "$chrootdir"; then
    warning "$chrootdir is not a mountpoint. This may have undesirable side effects."
  fi

  chroot_init
  if (( ! keepresolvconf )); then
    chroot_add_resolv_conf || die "failed to setup resolv.conf"
  fi

  if (( use_systemd )); then
    systemd_mode
  else
    offline_mode
  fi
}

while getopts ':hSNu:r' flag; do
  case $flag in
    h)
      usage
      exit 0
      ;;
    S)
      use_systemd=1
      ;;
    N)
      unshare=1
      ;;
    u)
      userspec=$OPTARG
      ;;
    r)
      keepresolvconf=1
      ;;
    :)
      die '%s: option requires an argument -- '\''%s'\' "${0##*/}" "$OPTARG"
      ;;
    ?)
      die '%s: invalid option -- '\''%s'\' "${0##*/}" "$OPTARG"
      ;;
  esac
done
shift $(( OPTIND - 1 ))

(( $# )) || die 'No chroot directory specified'
chrootdir=$1
shift

args=("$@")
if (( unshare )); then
  if (( use_systemd )); then
    die 'systemd mode (-S) cannot be used in conjunction with unshare (-N)'
  fi

  setup=unshare_setup
  $mount_unshare bash -c "$(declare_all); arch-chroot"
else
  setup=chroot_setup
  arch-chroot
fi
