#!/bin/bash

show_help() {
	echo "usage: cleanup-state <pre-invariant|post-invariant>"
	echo
	echo "Runs cleanup actions of the selected class:"
	echo "   pre-invariant:"
	echo "      Clean state that is currently expected to leak from"
	echo "      any test, due to imperfections in the code."
	echo "   post-invariant:"
	echo "      Clean state that is not expected to be restored by"
	echo "      each test, that should not leak across tests."
}

main() {
	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
				;;
			pre-invariant)
				action=pre-invariant
				shift
				;;
			post-invariant)
				action=post-invariant
				shift
				;;
			-*)
				echo "cleanup-state: unsupported argument $1" >&2
				exit 1
				;;
			*)
				echo "cleanup-state: unknown action $1" >&2
				exit 1
				;;
		esac
	done

	case "$action" in
		pre-invariant)
			# we try to remove snaps with user services, nevertheless some user
			# services may have been started for root or $SPREAD_SYSTEM_USER and
			# attempts to stop them during snap removal are best effort
			find /sys/fs/cgroup/ -type d -path '*/user.slice/*/snap.*service' -prune | while read -r svc; do
				# do not worry about putting a time bound/iteration limit,
				# worst case we hit a test timeout which will allow to identify
				# a problem
				while (("$(wc -l < "$svc/cgroup.procs" || echo 0)" > 0)); do
					# for compatibility with v1/hybrid and old kernels not having cgroup.kill
					if [ -f "$svc/cgroup.kill" ]; then
						echo 1 > "$svc/cgroup.kill"
					else
						# shellcheck disable=SC2002
						cat "$svc/cgroup.procs" | while read -r killpid; do
							kill -9 "$killpid" || true
						done
					fi
					sleep 1
				done
				rmdir "$svc" || true
			done

			# If using cgroups v1, the current snapd code does not remove
			# freezer or device cgroups, just clean them up here
			if [ -d /sys/fs/cgroup/freezer ]; then
				find /sys/fs/cgroup/freezer/ -type d -name 'snap.*' -prune -ls -exec rmdir \{\} \;
			fi

			if [ -d /sys/fs/cgroup/devices ]; then
				find /sys/fs/cgroup/devices/ -type d -name 'snap.*' -prune -ls -exec rmdir \{\} \;
			fi

			# Remove any cgroup pinned map files might be left behind
			if [ -d /sys/fs/bpf/snap ]; then
				find /sys/fs/bpf/snap -type f -name "snap_*" -ls -exec rm \{\} \;
			fi

			# If the root user has a systemd --user instance then ask it to reload.
			# This prevents tests from leaking user-session services that stay in
			# memory but are not present on disk, or have been modified on disk, as is
			# common with tests that use snaps with user services _or_ with tests that
			# cause installation of the snapd.session-agent.service unit via re-exec
			# machinery.
			#
			# This is done AHEAD of the invariant checks as it is very widespread
			# and fixing it in each test is not a priority right now.
			#
			# Note that similar treatment is not required for the "test" user as
			# correct usage of tests.session ensures that the session and all the
			# processes of the "test" user are terminated.
			if pgrep -u root --full "systemd --user"; then
				systemctl --user daemon-reload
				# Following that, check if there's a snapd.session-agent.socket and
				# if one exists stop it and then start it, ignoring errors.  If the
				# unit was removed, stopping it clears it from memory. This is
				# different from restarting the unit, which doesn't do anything if
				# the unit on disk is gone.
				if systemctl --user is-active snapd.session-agent.socket; then
					systemctl --user stop snapd.session-agent.socket
				fi
				# XXX: if there's a way to check if an unit exists but is stopped,
				# use it here to avoid starting a non-existing unit.
				systemctl --user start snapd.session-agent.socket || true

				# Do the same for root's D-Bus session bus.
				# This will stop the bus and any clients that
				# may be connected to it.
				if systemctl --user is-active dbus.socket; then
					systemctl --user stop dbus.socket
				fi
				systemctl --user start dbus.socket || true
			fi
			;;
		post-invariant)
			true
			;;
		*)
			echo "cleanup-state: unknown action $action" >&2
			exit 1
			;;
	esac
}

main "$@"
