#!/bin/bash
#
# Smart chroot script. All processes version, checks for any processes in
# the chroot instead of just shells.
#
# This will set up read-only bind mounts for /dev /proc /sys and /dev/pts
# if they don't already exist before chrooting. It will automatically tear
# down the mount points when the chroot terminates, if there are no other
# processes chrooted (it leaves them up if there are other chroot sessions
# on the image).
#
# You can set an environment variable called IMAGE_ROOT
# with a location to look for chroot images in, so you can do something
# like "chroot example" from anywhere and have it find the proper chroot.
#
# Will also check if /bin/sh in the image is (or points to) a 32-bit
# executeable, and run chroot under "linux32" if it does.
#
# This can also override the system chroot command without getting into
# any infinite loops.
#
# By default this script sets up bind mounts for /proc, /sys, /dev and
# /dev/pts. Extra directories can be set up by defining a space separated
# BIND_MOUNTS environment variable. This contains a semicolon separated set
# of host source directory and image mountpoint.
#
# Author: Patrick McLean <chutzpah@gentoo.org>
#

# default system bind mounts
SYSTEM_BIND_MOUNTS=(proc sys dev)
TMPFS_MOUNTS=(tmp /var/tmp/portage)

# list of shells to search
SHELL_LIST=(
	"${SHELL}"
	/bin/{bash,sh,bb}
)

RUNDIRS=(lock mount)

# control groups to use
CGROUPS=(cpu blkio cpuacct memory)
NAMESPACES=(mount uts ipc)

# name to put cgroups under
CGROUP_NAME="chroot-wrapper"

# maximum memory usage for a particular chroot
CHROOT_MAX_MEM="8192"

# define some useful functions
has() { [[ " ${*:2} " == *" $1 "* ]]; }

has_chrooted_processes() {
	# find if there are any shells chrooted into the given chroot
	local root="${1}"

	# search /proc/*/root entries for one that points to our chroot
	[[ -z "$(find /proc/*/root -nowarn -lname "${root}" 2>/dev/null)" ]] && return 1
	return 0
}

make_rprivate() {
	# mount --make-rprivate the mount containing the given path
	local src=${1}
	while ! mountpoint "${src}" >/dev/null; do
		src=${src%/*}
	done
	mount --make-rprivate "${src}"
}

setup_bind_mounts() {
	local image_root="${1}" src dest mtab_opts mount_output rundir sshdir tmpfile tmpdir

	has mount ${CHROOT_NAMESPACES} && mtab_opts="--no-mtab"

	# MS_PRIVATE is the kernel mount default, but systemd explicitly
	# sets it to MS_SHARED, which prevents the earlier unshare call
	# from working as intended. Therefore, we have to explicitly
	# make relevant mounts private.
	make_rprivate "${image_root}"

	# mount the system bind mounts defined in SYSTEM_BIND_MOUNTS at the top of this script
	for fs in "${SYSTEM_BIND_MOUNTS[@]}"; do
		if ! findmnt -kcunr "${image_root}/${fs}" > /dev/null; then
			make_rprivate "/${fs}"
			mount_output="$(mount --rbind ${mtab_opts}  "/${fs}" "${image_root}/${fs}" 2>&1)"

			if [[ "${?}" != "0" ]]; then
				echo "${mount_output}" > /dev/stderr
				echo "ERROR: Could not bind mount \"/${fs}\" on \"${image_root}/${fs}\"" > /dev/stderr
				return 1
			fi
		fi
	done

	# mount /dev/pts in the chroot image
	if ! findmnt -kcunr "${image_root}/dev/pts" > /dev/null; then
		mount_output="$(mount ${mtab_opts} -t devpts none "${image_root}/dev/pts" 2>&1)"

		if [[ "${?}" != "0" ]]; then
			echo "${mount_output}" > /dev/stderr
			echo "ERROR: Could not mount \"/dev/pts\" in the chroot image." > /dev/stderr
			return 1
		fi
	fi

	[[ "${DEBUG_MOUNTS}" ]] && { echo "BIND_MOUNTS=\"${BIND_MOUNTS}\""; unset DEBUG_MOUNTS; }
	# set up bind mounts from the BIND_MOUNTS environment variable
	for mnt in ${BIND_MOUNTS}; do
		src="${mnt%%;*}"
		src="${src%/}"

		dest="${mnt##*;}"
		dest="${dest%/}"

		if [[ ! -d "${src}" ]]; then
			echo "ERROR: bind mount source \"${src}\" does not exist on host system." > /dev/stderr
			return 1
		fi

		if [[ ! -d "${image_root}/${dest}" ]]; then
			echo "ERROR: bind mount point \"${dest}\" does not exist in chroot image." > /dev/stderr
			return 1
		fi

		if ! findmnt -kcunr "${image_root}${dest}" > /dev/null; then
			make_rprivate "${src}"
			mount_output="$(mount --bind ${mtab_opts} "${src}" "${image_root}${dest}" 2>&1)"

			if [[ "${?}" != "0" ]]; then
				echo "${mount_output}" > /dev/stderr
				echo "ERROR: Could not bind mount \"${src}\" on \"${image_root}/${dest}\"" > /dev/stderr
				return 1
			fi
		fi
	done

	# mount tmpfs locations
	for mnt in "${TMPFS_MOUNTS[@]}"; do
		dest="${mnt##*;}"
		dest="${dest%/}"

		if [[ -e "${image_root}/${dest}" ]] && ! findmnt -kcunr "${image_root}${dest}" > /dev/null; then
			mount_output="$(mount -t tmpfs ${mtab_opts} none "${image_root}/${dest}")"

			if [[ "${?}" != "0" ]]; then
				echo "${mount_output}" > /dev/stderr
				echo "ERROR: Could not bind mount \"${mnt}\" on \"${image_root}/${dest}\"" > /dev/stderr
				return 1
			fi
		fi
	done

	# create dirs in /run if necessary
	if [[ -d "${image_root}/run" ]]; then
		for rundir in "${RUNDIRS[@]}"; do
			[[ -d "${image_root}/run/${rundir}" ]] || mkdir -p "${image_root}/run/${rundir}"
		done
	fi

	# forward the ssh agent
	if [[ ${SSH_AUTH_SOCK} && -e ${SSH_AUTH_SOCK} && ${SSH_AUTH_SOCK#/tmp} != ${SSH_AUTH_SOCK} ]]; then
		sshdir="$(dirname ${SSH_AUTH_SOCK})"
		[[ -d ${image_root}/${sshdir} ]] || mkdir -p "${image_root}/${sshdir}"
			mount_output="$(mount --bind ${mtab_opts} "${sshdir}" "${image_root}/${sshdir}" 2>&1)"

			if [[ "${?}" != "0" ]]; then
				echo "${mount_output}" > /dev/stderr
				echo "ERROR: Could not bind mount \"${src}\" on \"${image_root}/${dest}\"" > /dev/stderr
				return 1
			fi
	fi
}

trap_signals() {
	local arg="${1}" signals

	if [[ -z "${arg}" ]]; then
		arg="-"
		signals="$(trap -p | egrep -o 'SIG[A-Z0-9+-]+')"
	else
		signals="$(trap -l | egrep -o 'SIG[A-Z0-9+-]+')"
	fi

	[[ ${signals} ]] && trap "${arg}" ${signals}
}

cleanup_bind_mounts() {
	local image_root="${1:-${image_root}}" dest sshdir cleanup

	# make sure this only gets called once
	[[ ${cleanup_bind_mounts_called} ]] && return
	cleanup_bind_mounts_called=1

	# reset signals
	trap_signals

	# if CHROOT_KEEP_MOUNTS is defined then don't umount
	[[ "${CHROOT_KEEP_MOUNTS}" ]] && return 0

	# if there are no other processes using the chroot, unmount it
	if has mount ${CHROOT_NAMESPACES} || ! has_chrooted_processes "${image_root}"; then
		[[ ${SSH_AUTH_SOCK} ]] && sshdir="$(dirname ${SSH_AUTH_SOCK})"

		for fs in "${sshdir}" "${TMPFS_MOUNTS[@]}" dev/pts "${SYSTEM_BIND_MOUNTS[@]}"; do
			# don't unmount proc yet
			[[ -z ${fs} || ${fs} =~ ^[/]*proc[/]*$ ]] && continue

			findmnt -kcunro TARGET "${image_root}/${fs}" > /dev/null && umount -l ${mtab_opts} "${image_root}/${fs}"
		done

		for cleanup in "${sshdir}"; do
			[[ ${cleanup} && -d ${image_root}${cleanup} ]] && rmdir "${image_root}${cleanup}"
		done

		for mnt in $(echo ${BIND_MOUNTS} | tac -s ' '); do
			dest="${mnt##*;}"
			dest="${dest%/}"

			# don't unmount proc yet
			[[ -z ${dest} || ${dest} =~ ^[/]*proc[/]*$ ]] && continue

			findmnt -kcunro TARGET "${image_root}${dest}" > /dev/null && umount -l ${mtab_opts} "${image_root}${dest}"
		done
	fi

	# make sure /proc is unmounted
	findmnt -kcunro TARGET "${image_root}/proc" > /dev/null && umount -l ${mtab_opts} "${image_root}/proc"
}

setup_cgroups() {
	local image_name="$(basename "${1}")"

	[[ ! -d "/sys/fs/cgroup" ]] && return 0

	for cgroup in "${CGROUPS[@]}"; do
		if [[ -d "/sys/fs/cgroup/${cgroup}" ]]; then
			[[ -d "/sys/fs/cgroup/${cgroup}/${CGROUP_NAME}" ]] || mkdir "/sys/fs/cgroup/${cgroup}/${CGROUP_NAME}"
			[[ -d "/sys/fs/cgroup/${cgroup}/${CGROUP_NAME}/${image_name}" ]] || mkdir "/sys/fs/cgroup/${cgroup}/${CGROUP_NAME}/${image_name}"

			[[ -w "/sys/fs/cgroup/${cgroup}/${CGROUP_NAME}/${image_name}/tasks" ]] && echo $$ > "/sys/fs/cgroup/${cgroup}/${CGROUP_NAME}/${image_name}/tasks"

			if [[ "${cgroup}" == "memory" && -w "/sys/fs/cgroup/${cgroup}/${CGROUP_NAME}/${image_name}/memory.limit_in_bytes" ]]; then
				((max_mem=CHROOT_MAX_MEM*1048576))
				echo ${max_mem} > "/sys/fs/cgroup/${cgroup}/${CGROUP_NAME}/${image_name}/memory.limit_in_bytes"
			fi
		fi
	done
}

setup_namespaces() {
	local ns namespaces p

	# setup namespaces if we haven't already
	if [[ -z "${CHROOT_NAMESPACES}" ]] && command -v unshare &> /dev/null; then
		namespaces=""
		for ns in "${NAMESPACES[@]}"; do
			unshare --${ns} true &> /dev/null && namespaces+=" ${ns}"
		done

		CHROOT_NAMESPACES="${namespaces% }" exec unshare $(for p in ${namespaces}; do echo --${p}; done) -- /bin/bash "${0}" "${@}"
	fi
}

cleanup_cgroups() {
	local image_name="$(basename "${1}")"

	[[ ! -d "/sys/fs/cgroup" ]] && return 0
	if ! has_chrooted_processes "${image_root}"; then
		for cgroup in "${CGROUPS[@]}"; do
			[[ -f "/sys/fs/cgroup/${cgroup}/tasks" ]] && echo $$ > "/sys/fs/cgroup/${cgroup}/tasks"

			if [[ -d "/sys/fs/cgroup/${cgroup}/${CGROUP_NAME}/${image_name}" ]]; then
				[[ -w "/sys/fs/cgroup/${cgroup}/tasks" ]] && echo $$ > "/sys/fs/cgroup/${cgroup}/tasks"

				tasks_file="/sys/fs/cgroup/${cgroup}/${CGROUP_NAME}/${image_name}/tasks"
				if [[ -e "${tasks_file}" && -z "$(< "${tasks_file}")" ]]; then
					rmdir "/sys/fs/cgroup/${cgroup}/${CGROUP_NAME}/${image_name}"
				fi

				if [[ -d "/sys/fs/cgroup/${cgroup}/${CGROUP_NAME}" && -z "$(find -P "/sys/fs/cgroup/${cgroup}/${CGROUP_NAME}" -maxdepth 1 -mindepth 1 -and -type d -print0)" ]]; then
					rmdir "/sys/fs/cgroup/${cgroup}/${CGROUP_NAME}"
				fi
			fi
		done
	fi
}


# use namespaces if we haven't already
setup_namespaces "${@}"

trap_signals cleanup_bind_mounts

if [[ -z "${1}" ]]; then
	echo "Syntax: $(basename $0) <chroot directory> [command]"
	exit 1
fi

if command -v realpath &>/dev/null; then
	IMAGE="$(realpath "${1}")"
else
	IMAGE="${1}"
fi
shift

# find the system "chroot" command in case this is overriding it
realpath="$(type -p realpath)"
[[ -z "${realpath}" ]] && realpath="echo"
chroot_cmd="$(type -Paf chroot | grep -v "^$(${realpath} "${0}")" | head -n1)"

# look in IMAGE_ROOT for the image if the given image is not in the current
# directory is not a absolute path
if [[ "${IMAGE#/}" != "${IMAGE}" ]]; then
	image_root="$(realpath "${IMAGE%/}")"
	[[ "${IMAGE_NAME}" ]] || IMAGE_NAME="$(basename ${IMAGE})"

elif [[ -d "${PWD}/${IMAGE}" ]]; then
	image_root="$(realpath "${PWD}/${IMAGE%/}")"
	[[ "${IMAGE_NAME}" ]] || IMAGE_NAME="${IMAGE%/}"

elif [[ -n "${IMAGE_ROOT}" && -d "${IMAGE_ROOT}/${IMAGE}" ]]; then
	image_root="$(realpath "${IMAGE_ROOT%/}/${IMAGE%/}")"
	[[ "${IMAGE_NAME}" ]] || IMAGE_NAME="${IMAGE%/}"

else
	echo "Could not find chroot image \"${IMAGE}\"" > /dev/stderr
	exit 1
fi

# make sure the image exists
if [[ ! -d "${image_root}" ]]; then
	echo "Could not find chroot directory \"${IMAGE}\"" > /dev/stderr
	exit 1
fi

if ! setup_bind_mounts "${image_root}"; then
	cleanup_bind_mounts "${image_root}"
	exit 1
fi

# find a shell to run
for shell in "${SHELL_LIST[@]}"; do
	if [[ -x ${image_root}${shell} ]]; then
		image_shell="${shell}"
		break
	fi
done

if [[ ! ${image_shell} ]]; then
	echo "Could not find a valid shell" >> /dev/stderr
	exit 1
fi

# run a shell if the command is missing
[[ -z "${1}" ]] && set -- "${image_shell}"

# always use login shells
if grep -q "^${1}" /etc/shells && ! has -l "${@}" && ! has --login "${@}"; then
	CMD="${1}"
	shift
	set -- "${CMD}" -l "${@}"
else
	CMD="${@}"
	shift
	set -- "${image_shell}" -l -c "${CMD}"
fi

# try to figure out if we need to use linux32
if [[ "$(type -p file)" ]]; then
	[[ "$(uname -m)" == "x86_64" && -e "${image_root}/bin/sh" ]] && file -bL "${image_root}/bin/sh" | grep -q '32-bit' && chroot_cmd="$(type -p linux32) ${chroot_cmd}"
elif [[ -x "${image_root}/usr/bin/file" ]] && "${image_root}/usr/bin/file" -m "${image_root}"/usr/share/misc/magic -l > /dev/null 2>&1 ; then
	[[ "$(uname -m)" == "x86_64" && -e "${image_root}/bin/sh" ]] && "${image_root}/usr/bin/file" -bL "${image_root}/bin/sh" | grep -q '32-bit' && chroot_cmd="$(type -p linux32) ${chroot_cmd}"
fi

[[ -z "${CHROOT_KEEP_ENV}" ]] && env_opts="-i"

setup_cgroups "${image_root}"

# set the hostname if we have uts namespaces
has uts ${CHROOT_NAMESPACES} && hostname "${IMAGE_NAME}"

env ${env_opts} CHROOT="${IMAGE_NAME}" HOME="${HOME}" TERM="${TERM}" SSH_AUTH_SOCK="${SSH_AUTH_SOCK}" ${chroot_cmd} "${image_root}" "${@}"
chroot_return="${?}"

# clean up after ourselves
cleanup_bind_mounts "${image_root}"

cleanup_cgroups "${image_root}"

exit "${chroot_return}"

# vim:ts=4:sw=4:noet:sts=4:ai: