#!/bin/sh

# To use a different OCI-compliant container runtime,
# update both CONTAINER_RUNTIME and CONTAINER_RUNTIME_PATH:
#
# - CONTAINER_RUNTIME should be set to the name of the runtime binary (e.g., 'crun', 'runc').
# - CONTAINER_RUNTIME_PATH should point to the directory where the binary is installed
#   (e.g., '/usr/bin' for a system-installed runtime).

CONTAINER_RUNTIME_PATH=$PWD
CONTAINER_RUNTIME=crun
CONTAINER_STATUS_PATH=/run/$CONTAINER_RUNTIME
CONTAINER_CMD="$CONTAINER_RUNTIME_PATH/$CONTAINER_RUNTIME --root=$CONTAINER_STATUS_PATH"
CONTAINER_CMD_ADDITIONAL_OPTIONS=""
CONTAINER_LOGS="/var/log/dante_container.log"

IMAGES_PATH=$PWD/dante_data/images
ROOTFS_MOUNTPOINT=$PWD/bundle/rootfs
ACTIVE_IMAGE_ID_PATH=$IMAGES_PATH/active
DANTE_JSON=$PWD/dante_data/capability/dante.json

# Check we can actually start/stop containers
# NOTE: on some systems 'id' might not be available, hence we check manually
if [ "$(grep -E '^Uid:' /proc/self/status | awk '{print $2}')" -ne 0 ]; then
	echo "This script must be executed with root privileges."
	echo ""
	exit 1
fi

# This function assumes that:
# - the JSON file is well-formed
# - key and value appear on the same line
# - strings are double-quoted and don’t contain escaped quotes
# - assumes the key exists exactly once per line
get_json_field()
{
	json_file=$1
	field_name=$2
	default="__NOT_FOUND__"

	if [ ! -f "$json_file" ]; then
		echo "error: file '$json_file' not found" >&2
		exit 1
	fi

	# explaining each sed:
	# - 's/^[^:]*://' removes everything up to and including the first colon
	# - 's/ //' removes the first space character after the colon, if present
	# - 's/^"//' and 's/"$//' removes leading and trailing double quotes from the value
	# - 's/,[[:space:]]*$//' removes a trailing comma and any following whitespace (e.g. to handle lists)
	# - 's/[[:space:]]*$//' trims any remaining trailing whitespace from the value
	value=$(grep "\"$field_name\"" "$json_file" | \
		sed -e 's/^[^:]*://' -e 's/ //' | \
		sed -e 's/^"//' -e 's/"$//' | \
		sed -e 's/,[[:space:]]*$//' | \
		sed -e 's/[[:space:]]*$//' | head -n 1)

	if [ -z "$value" ]; then
		echo "$default"
	else
		echo "$value"
	fi
}

check_cgroup_mounts()
{
	cgroup_version=$(get_json_field "$DANTE_JSON" "cgroupVersion")

	if [ "$cgroup_version" = "__NOT_FOUND__" ]; then
		return
	fi

	check_mount() {
		path=$1
		expectedType=$2

		while IFS= read -r line; do
			# get the second field (mount point) and the third field (type)
			mountPoint=$(echo "$line" | awk '{print $2}')
			mountType=$(echo "$line" | awk '{print $3}')

			# if the mountPoint doesn't start with /sys/fs/cgroup, skip it
			case "$mountPoint" in
				/sys/fs/cgroup*) ;;
				*) continue ;;
			esac

			# if mount point and type exactly match the expected values, we're good
			if [ "$mountPoint" = "$path" ] && [ "$mountType" = "$expectedType" ]; then
				echo "mount OK: $path ($expectedType)"
				return
			fi

			# There is a chance multiple controllers are mounted on the same path,
			# for instance we might be looking for /sys/fs/cgroup/cpu and have
			#
			#   /sys/fs/cgroup/cpu,cpuacct cgroup etc..
			#
			# mounted instead.

			# because we skip entries that do not start with /sys/fs/cgroup at
			# the beginning of the loop, we know getting the substring after
			# /sys/fs/cgroup at this point will yield an empty string at worst
			cgroupSubstring=${mountPoint#/sys/fs/cgroup/}

			# do the same with $path
			cgroupPathSubstring=${path#/sys/fs/cgroup/}

			# check if cgroupPathSubstring is part of cgroupSubstring
			# eg this would successfully match 'cpu' against both 'cpuacct,cpu' and 'cpu,cpuacct'
			if echo "$cgroupSubstring" | grep -qw "$cgroupPathSubstring"; then
				if [ "$mountType" = "$expectedType" ]; then
					echo "mount OK: $path ($expectedType)"
					return
				fi
			fi
		done < /proc/mounts

		echo "warning: missing or incorrect mountpoint: $path (expected type: $expectedType)"
	}

	if [ "$cgroup_version" = "1" ]; then
		echo "cgroup version set to v1 in $DANTE_JSON"
		echo "checking mounts..."
		check_mount "/sys/fs/cgroup" "tmpfs"
		check_mount "/sys/fs/cgroup/cpuset" "cgroup"
		check_mount "/sys/fs/cgroup/cpu" "cgroup"
		check_mount "/sys/fs/cgroup/memory" "cgroup"
		check_mount "/sys/fs/cgroup/devices" "cgroup"
	elif [ "$cgroup_version" = "2" ]; then
		echo "cgroup version set to v2 in $DANTE_JSON"
		echo "checking mounts..."
		check_mount "/sys/fs/cgroup" "cgroup2"
	else
		echo "error: unsupported cgroupVersion value ($cgroup_version) in $DANTE_JSON"
		exit 1
	fi
}

start()
{
	# A poorly-timed stop() could leave the container mounted while
	# the processes inside the container were successfully shut down.
	# Instead of relying on whether the container is there or not when 
	# deciding to start DEP, check whether dep_manager is actually running.

	# shellcheck disable=SC2009
	# (SC2009 recommends using pgrep, but it is not always available)
	if ps -e >/dev/null 2>&1; then
		PS_CMD="ps -e"
	else
		PS_CMD="ps" # assume heavily stripped-down BusyBox
	fi

	if $PS_CMD | grep -q "[d]ep_manager"; then
		echo "DEP is already running"
		exit 0
	fi

	# Some basic checks before proceeding
	if [ ! -f "$ACTIVE_IMAGE_ID_PATH" ]; then
		echo "error: $ACTIVE_IMAGE_ID_PATH not found, can't select active rootfs"
		exit 1
	fi

	active_image_id=$(cat "$ACTIVE_IMAGE_ID_PATH")
	rootfs="$IMAGES_PATH/$active_image_id/rootfs_squash"

	if [ ! -f "$rootfs" ]; then
		echo "error: $rootfs not found"
		exit 1
	fi

	check_cgroup_mounts

	mkdir -p /var/run/dante
	mkdir -p ${CONTAINER_STATUS_PATH}

	# Make sure /etc/resolv.conf is there when later on we
	# try to bind mount it from the container.
	if [ ! -f "/etc/resolv.conf" ]; then
		touch /etc/resolv.conf
	fi

	if ! grep -q " $ROOTFS_MOUNTPOINT " /proc/mounts; then
		if ! mount "$rootfs" "$ROOTFS_MOUNTPOINT" >/dev/null 2>&1; then
			echo "error: could not mount $rootfs"
			exit 1
		fi
	fi

	# At this point, it's safe to always forcefully delete the container.
	#
	# This may be necessary in scenarios where the DEP processes did not actually
	# start after running ./dep.sh start — for example, due to invalid configuration
	# in dante.json. In such cases, a user would typically inspect the logs,
	# fix the underlying issue, and then retry with ./dep.sh start.
	#
	# However, if the dante container remains mounted, the container runtime's 'run'
	# command will fail, forcing the user to manually delete the container - either
	# by using ./dep.sh stop (which is not intuitive) or manually.
	#
	# To avoid these issues and make the recovery easier to execute, unconditionally
	# remove the dante container before attempting to run it again.
	#
	# NOTE: while we could check whether the container exists before removing it,
	# not all systems provide the necessary cgroup status layers to reliably list
	# configured containers.
	${CONTAINER_CMD} delete --force dante

	# rootfs (only mount with no parent mount) cannot be pivot_root()ed. The check hereafter
	# relies on the fact that rootfs will be either a ramfs or tmpfs. This is a bit more restrictive
	# than necessary, as the container could in practice be started from a ramfs or tmpfs (as long as
	# it is not the rootfs).
	# WARNING: crun falls back to chroot when --no-pivot is enabled, and a process running in the container
	# can in practice access the tree outside of the chroot.
	ROOT_FSTYPE=$(mount|grep 'on / type'|awk '{print $5}')
	if [ "$ROOT_FSTYPE" = "rootfs" ] || [ "$ROOT_FSTYPE" = "ramfs" ] || [ "$ROOT_FSTYPE" = "tmpfs" ]; then
		CONTAINER_CMD_ADDITIONAL_OPTIONS="$CONTAINER_CMD_ADDITIONAL_OPTIONS --no-pivot"
	fi

	if ! ${CONTAINER_CMD} run ${CONTAINER_CMD_ADDITIONAL_OPTIONS} --detach --bundle ./bundle dante > "$CONTAINER_LOGS" 2>&1; then
		echo "error: failed to start dante container, more details available in $CONTAINER_LOGS"
		exit 1
	else
		echo "DEP started"
	fi
}

stop()
{
	# in some cases we might have the mountpoint but no container running:
	# check if that's the case before proceeding
	if ${CONTAINER_CMD} list | grep dante >/dev/null 2>&1; then
		# stop the init process (dep_manager) by sending a SIGTERM signal
		echo "stopping DEP..."
		${CONTAINER_CMD} kill dante TERM

		for _ in $(seq 1 10); do
			sleep 1
			DEP_PROCS=$(${CONTAINER_CMD} ps dante | grep -v PID -c)
			if [ "$DEP_PROCS" -eq 0 ]; then
				break
			fi
		done

		DEP_PROCS=$(${CONTAINER_CMD} ps dante | grep -v PID -c)
		if [ "$DEP_PROCS" -ne 0 ]; then
			echo "DEP still running, sending SIGKILL"
			${CONTAINER_CMD} kill -a dante KILL
			sleep 1
		fi

		echo "removing container..."
		${CONTAINER_CMD} delete --force dante
	fi

	if grep -q " $ROOTFS_MOUNTPOINT " /proc/mounts; then
		echo "umount rootfs..."
		umount "$PWD"/bundle/rootfs
	fi

	echo "done"
}

USAGE_MESSAGE="Usage: dep.sh <start|stop>"

if [ "$#" -eq 0 ]; then
	echo "$USAGE_MESSAGE"
	exit 1
fi

case $1 in
	"start" ) start "$2" ;;
	"stop" ) stop ;;
	* )
		echo "$USAGE_MESSAGE"
		exit 1
		;;
esac
