#!/bin/bash

show_help() {
	echo "usage: tests.invariant check [INVARIANT...]"
	echo
	echo "Supported invariants:"
	echo "    root-files-in-home: most of /home/* does not contain root-owned files"
	echo "    crashed-snap-confine: $TESTSTMP/snap.rootfs_* does not exist"
	echo "    lxcfs-mounted: /var/lib/lxcfs is a mount point"
	echo "    stray-dbus-daemon: at most one dbus-daemon is running"
	echo "    leftover-defer-sh: defer.sh must not be left over by tests"
	echo "    broken-snaps: snaps must not be left around that are in a broken state"
	echo "    segmentation-violations: snapd must not have segmentation-violation errors in journal logs"
}

if [ $# -eq 0 ]; then
	show_help
	exit 1
fi

action=
while [ $# -gt 0 ]; do
	case "$1" in
		-h|--help)
			show_help
			exit 0
			;;
		--)
			shift
			break
			;;
		check)
			action=check
			shift
			;;
		-*)
			echo "tests.invariant: unsupported argument $1" >&2
			exit 1
			;;
		*)
			break
			;;
	esac
done

check_root_files_in_home() {
	n="$1" # invariant name
	# This find expression looks for root owned file in /home with the following exceptions:
	# - The /home/ubuntu directory is root owned but is not used by the tests so we ignore it
	# - The /home/gopath and everything inside it comes from spread.yaml and is too inconvenient to change
	# As a note, the working theory for the origin of /home/ubuntu is cloud-init and data from GCE.
	find /home -mindepth 1 -user root -a ! -path "/home/ubuntu" -a ! -path "$GOHOME" -a ! -path "$GOHOME/*" 2>/dev/null >"$TESTSTMP/tests.invariant.$n"
	if [ -s "$TESTSTMP/tests.invariant.$n" ]; then
		echo "tests.invariant: the following files should not be owned by root" >&2
		cat "$TESTSTMP/tests.invariant.$n" >&2
		return 1
	fi
}

check_crashed_snap_confine() {
	n="$1" # invariant name
	find "$TESTSTMP" -name 'snap.rootfs_*' > "$TESTSTMP/tests.invariant.$n"
	# NOTE: it may be a mount point but we are not checking that here.
	if [ -s "$TESTSTMP/tests.invariant.$n" ]; then
		echo "tests.invariant: it seems snap-confine has crashed" >&2
		cat "$TESTSTMP/tests.invariant.$n" >&2
		return 1
	fi
}

check_lxcfs_mounted() {
	n="$1" # invariant name
	"$TESTSTOOLS"/mountinfo.query /var/lib/lxcfs > "$TESTSTMP/tests.invariant.$n"
	if [ -s "$TESTSTMP/tests.invariant.$n" ]; then
		echo "tests.invariant: it seems lxcfs is mounted" >&2
		cat "$TESTSTMP/tests.invariant.$n" >&2
		return 1
	fi
}

check_stray_dbus_daemon() {
	n="$1" # invariant name
	failed=()
	skipped_system=0
	skipped_root_session=0
	for pid in $(pgrep -x dbus-daemon); do
		cmdline="$(tr '\0' ' ' < "/proc/$pid/cmdline")"
		# Ignore one dbus-daemon responsible for the system bus.
		if echo "$cmdline" | grep -q 'dbus-daemon --system' && [ "$skipped_system" -eq 0 ]; then
			skipped_system=1
			continue
		fi
		if echo "$cmdline" | grep -q 'dbus-daemon --session' && [ "$(stat -c %u "/proc/$pid")" -eq 0 ] && [ "$skipped_root_session" -eq 0 ]; then
			skipped_root_session=1
			continue
		fi
		# Ignore dbus-daemon running the session of the "external" user.
		# This may happen when testing a core device using the ad-hoc
		# backend. This user is never used by tests explicitly.
		user="$(stat -c %U "/proc/$pid")"
		if [ "$user" = "external" ]; then
			continue
		fi
		# On qemu backend, sudo from user "ubuntu" is used. So the session
		# for user "ubuntu" is expected to exist.
		if [ "$user" = "ubuntu" ]; then
			continue
		fi
		# Report stray dbus-daemon.
		failed+=("pid:$pid user:$user cmdline:$cmdline")
	done

	if [ "${#failed[@]}" -ne 0 ]; then
		(for line in "${failed[@]}"; do
			echo "${line}"
		done) >"$TESTSTMP/tests.invariant.$n"
		cat "$TESTSTMP/tests.invariant.$n" >&2
		return 1
	fi
}

check_leftover_defer_sh() {
	n="$1" # invariant name
	(
		find "$PROJECT_PATH" -name defer.sh > "$TESTSTMP/tests.invariant.$n"
	) > "$TESTSTMP/tests.invariant.$n"
	if [ -s "$TESTSTMP/tests.invariant.$n" ]; then
		echo "tests.invariant: leftover defer.sh script" >&2
		cat "$TESTSTMP/tests.invariant.$n" >&2
		return 1
	fi
}

check_broken_snaps() {
	n="$1" # invariant name
	(
		SNAP_MOUNT_DIR="$(os.paths snap-mount-dir)"
		# first column is the snap name, revision is 3rd
		snap list --all | awk '/,?broken,?/ {print $1,$3}' | while read -r name rev; do
			echo "snap $name ($rev) is broken:"
			systemctl status "$(systemd-escape -p "$SNAP_MOUNT_DIR/${name}/${rev}).mount")" || true
			ls -l "/${SNAP_MOUNT_DIR}/${name}/" || true
			echo "---"
		done
	) > "$TESTSTMP/tests.invariant.$n"
	if [ -s "$TESTSTMP/tests.invariant.$n" ]; then
		echo "tests.invariant: broken snaps" >&2
		cat "$TESTSTMP/tests.invariant.$n" >&2
		return 1
	fi
}

check_cgroup_scopes() {
	# xenial has leftover scope files possibly due to an old systemd bug
	# https://bugs.launchpad.net/ubuntu/+source/systemd/+bug/1934147
	if [[ "$SPREAD_SYSTEM" == ubuntu-core-16-* ]] || [[ "$SPREAD_SYSTEM" == ubuntu-16.04-* ]]; then
		return 0
	fi

	wait_for_system_scopes=5
	while [ "${wait_for_system_scopes}" -gt 0 ]; do
		has_active_scope=0
		while read -r -d '' scope; do
			if systemctl is-active "$(basename "${scope}")"; then
				has_active_scope=1
				break
			fi
		done < <(find /sys/fs/cgroup/system.slice -name 'snap.*.scope' -print0)
		if [ "${has_active_scope}" -eq 0 ]; then
			break
		fi
		: $((wait_for_system_scopes--))
		sleep 1
	done

	n="$1" # invariant name
	find /sys/fs/cgroup -name 'snap.*.scope' > "$TESTSTMP/tests.invariant.$n"
	if [ -s "$TESTSTMP/tests.invariant.$n" ]; then
		echo "tests.invariant: leftover cgroup scope files" >&2
		cat "$TESTSTMP/tests.invariant.$n" >&2
		return 1
	fi
}

check_cgroups() {
	n="$1" # invariant name
	# ignore mounts which we expect, especially on Ubuntu Core
	#
	# TODO: cleanup isn't perfect, it is possible that we're leaking socket and
	# units in user sessions
	find /sys/fs/cgroup -type d -name 'snap.*' \
		! -name '*.mount' \
		! -name '*.socket' -ls > "$TESTSTMP/tests.invariant.$n"
	if [ -s "$TESTSTMP/tests.invariant.$n" ]; then
		echo "tests.invariant: leftover cgroups" >&2
		cat "$TESTSTMP/tests.invariant.$n" >&2
		return 1
	fi
}

check_segmentation_violations() {
	n="$1" # invariant name

	"$TESTSTOOLS"/journal-state get-log -u snapd | grep "segmentation violation" > "$TESTSTMP/tests.invariant.$n"
	if [ -s "$TESTSTMP/tests.invariant.$n" ]; then
		echo "tests.invariant: segmentation violation found" >&2
		cat "$TESTSTMP/tests.invariant.$n" >&2
		return 1
	fi
}

check_fakestore_cleaned() {
    # Check if fakestore was properly cleaned to avoid leaking into other tests.
    if [ -f "/etc/systemd/system/snapd.service.d/store.conf" ]; then
        echo "/etc/systemd/system/snapd.service.d/store.conf was not cleaned properly"
        exit 1
    fi
}

check_invariant() {
	case "$1" in
		root-files-in-home)
			check_root_files_in_home "$1"
			;;
		crashed-snap-confine)
			check_crashed_snap_confine "$1"
			;;
		lxcfs-mounted)
			check_lxcfs_mounted "$1"
			;;
		stray-dbus-daemon)
			check_stray_dbus_daemon "$1"
			;;
		leftover-defer-sh)
			check_leftover_defer_sh "$1"
			;;
		broken-snaps)
			check_broken_snaps "$1"
			;;
		cgroup-scopes)
			check_cgroup_scopes "$1"
			;;
		cgroups)
			check_cgroups "$1"
			;;
		segmentation-violations)
			check_segmentation_violations "$1"
			;;
		check-fakestore-cleaned)
			check_fakestore_cleaned
			;;
		*)
			echo "tests.invariant: unknown invariant $1" >&2
			exit 1
			;;
	esac
}

main() {
	ALL_INVARIANTS="
	root-files-in-home
	crashed-snap-confine
	lxcfs-mounted
	stray-dbus-daemon
	leftover-defer-sh
	broken-snaps
	cgroup-scopes
	cgroups
	segmentation-violations
	check-fakestore-cleaned
	"

	case "$action" in
		check)
			ok=1
			if [ $# -gt 0 ]; then
				INV_LIST="$*"
			else
				INV_LIST="$ALL_INVARIANTS"
			fi
			for inv in $INV_LIST; do
				if check_invariant "$inv"; then
					echo "tests.invariant: $inv ok"
				else
					echo "tests.invariant: $inv not-ok" >&2
					ok=0
				fi
			done
			if [ $ok -eq 0 ]; then
				echo "tests.invariant: system is corrupted" >&2
				exit 1
			fi
			;;
		*)
			echo "tests.invariant: unknown action $action" >&2
			exit 1
			;;
	esac
}

main "$@"