#!/bin/bash
# SPDX-License-Identifier: GPL-3.0+
# Copyright (C) 2017 Omar Sandoval

shopt -s extglob

_warning() {
	echo "$0: $*" >&2
}

_error() {
	echo "$0: $*" >&2
	exit 1
}

_found_test() {
	local test_name="$1"
	local explicit="$2"

	unset DESCRIPTION QUICK TIMED requires device_requires test test_device

	# shellcheck disable=SC1090
	if ! . "tests/${test_name}"; then
		return 1
	fi

	if [[ -z $DESCRIPTION ]]; then
		_warning "${test_name} does not define DESCRIPTION"
		return 1
	fi

	if declare -fF test >/dev/null && declare -fF test_device >/dev/null; then
		_warning "${test_name} defines both test() and test_device()"
		return 1
	fi

	if ! declare -fF test >/dev/null && ! declare -fF test_device >/dev/null; then
		_warning "${test_name} does not define test() or test_device()"
		return 1
	fi

	if declare -fF device_requires >/dev/null && ! declare -fF test_device >/dev/null; then
		_warning "${test_name} defines device_requires() but not test_device()"
		return 1
	fi

	if (( QUICK && TIMED )); then
		_warning "${test_name} cannot be both QUICK and TIMED"
		return 1
	fi

	if (( ! explicit )); then
		if [[ $DEVICE_ONLY -ne 0 ]] && ! declare -fF test_device >/dev/null; then
			return
		fi
		if (( QUICK_RUN && !QUICK && !TIMED )); then
			return
		fi
		if [[ -n ${EXCLUDE["$test_name"]} ]]; then
			return
		fi
	fi

	printf '%s\0' "$test_name"
}

_found_group() {
	local group="$1"
	local explicit="$2"

	local test_path

	if (( ! explicit )); then
		if [[ -n ${EXCLUDE["$group"]} ]]; then
			return
		fi
	fi

	while IFS= read -r -d '' test_path; do
		_found_test "${test_path#tests/}" 0
	done < <(find "tests/$group" -type f -name '[0-9][0-9][0-9]' -print0)
}

_find_tests() {
	if [[ $# -eq 0 ]]; then
		# No filters given, run all tests except for the meta group.
		local group_path
		while IFS= read -r -d '' group_path; do
			if [[ $group_path != tests/meta ]]; then
				_found_group "${group_path#tests/}" 0
			fi
		done < <(find tests -mindepth 2 -type f -name rc -printf '%h\0')
	else
		local filter
		for filter in "$@"; do
			# A filter is bad if it:
			#     - Is an absolute path
			#     - Is a relative path containing ".."
			if [[ $filter =~ ^/|(^|/)\.\.(/|$) ]]; then
				_warning "bad filter ${filter}"
				continue
			fi
			# Remove leading and internal "." components.
			filter="${filter##+(./)}"
			filter="${filter//\/+(.\/)/\/}"
			# Remove repeated "/".
			filter="${filter//\/+(\/)/\/}"
			# Remove leading "tests/"
			filter="${filter#tests/}"

			if [[ -d tests/$filter ]]; then
				_found_group "$filter" 1
			elif [[ -f tests/$filter ]]; then
				_found_test "$filter" 1
			else
				_warning "no group or test named ${filter}"
			fi
		done
	fi
}

_check_dmesg() {
	local dmesg_marker="$1"
	local seqres="${RESULTS_DIR}/${TEST_NAME}"

	if [[ $CHECK_DMESG -eq 0 ]]; then
		return 0
	fi

	dmesg | bash -c "$DMESG_FILTER" | grep -A 9999 "$dmesg_marker" >"${seqres}.dmesg"
	grep -q -e "kernel BUG at" \
	     -e "WARNING:" \
	     -e "BUG:" \
	     -e "Oops:" \
	     -e "possible recursive locking detected" \
	     -e "Internal error" \
	     -e "INFO: suspicious RCU usage" \
	     -e "INFO: possible circular locking dependency detected" \
	     -e "general protection fault:" \
	     "${seqres}.dmesg"
	# shellcheck disable=SC2181
	if [[ $? -eq 0 ]]; then
		return 1
	else
		rm -f "${seqres}.dmesg"
		return 0
	fi
}

# Associative arrays are local by default. declare -g was added in 4.2, but it
# was apparently broken for associative arrays initially.
declare -A LAST_TEST_RUN
declare -A TEST_RUN

_read_last_test_run() {
	local seqres="${RESULTS_DIR}/${TEST_NAME}"

	LAST_TEST_RUN["date"]=""
	LAST_TEST_RUN["status"]=""
	LAST_TEST_RUN["reason"]=""
	LAST_TEST_RUN["exit_status"]=""
	LAST_TEST_RUN["runtime"]=""

	if [[ ! -e $seqres ]]; then
		return
	fi

	local key value
	while IFS=$'\t' read -r key value; do
		LAST_TEST_RUN["$key"]="$value"
	done <"$seqres"
}

_write_test_run() {
	local key value
	for key in "${!TEST_RUN[@]}"; do
		value="${TEST_RUN["$key"]}"
		printf '%s\t%s\n' "$key" "$value" >>"$seqres"
	done
}

_output_status() {
	local test="$1"
	local status="$2"

	if [[ -v DESCRIPTION ]]; then
		printf '%-60s' "$test ($DESCRIPTION)"
	else
		printf '%-60s' "$test"
	fi
	if [[ -z $status ]]; then
		echo
		return
	fi

	echo -n " ["
	if [[ -t 1 ]]; then
		case "$status" in
		passed)
			tput setaf 2
			;;
		failed)
			tput setaf 1
			;;
		"not run")
			tput setaf 3
			;;
		esac
	fi
	echo -n "$status"
	if [[ -t 1 ]]; then
		tput sgr0
	fi
	echo "]"
}

_output_notrun() {
	_output_status "$1" "not run"
	echo "    $SKIP_REASON"
}

_output_last_test_run() {
	if [[ -v TEST_DEV ]]; then
		_output_status "$TEST_NAME => $(basename "$TEST_DEV")" ""
	else
		_output_status "$TEST_NAME" ""
	fi

	(
	local key value
	while IFS= read -r key; do
		if [[ $key =~ ^date|status|reason|exit_status$ ]]; then
			continue
		fi
		value="${LAST_TEST_RUN["$key"]}"
		printf '    %s\t%s\t...\n' "${key}" "${value}"
	done < <(printf '%s\n' "${!LAST_TEST_RUN[@]}" | sort)
	) | column -t -s $'\t'
}

_output_test_run() {
	if [[ -t 1 ]]; then
		# -4 for date, status, reason, and exit_status, which we don't
		# output, +1 for the test name.
		tput cuu $((${#LAST_TEST_RUN[@]} - 3))
	fi

	if [[ -v TEST_DEV ]]; then
		_output_status "$TEST_NAME => $(basename "$TEST_DEV")" "${TEST_RUN["status"]}ed"
	else
		_output_status "$TEST_NAME" "${TEST_RUN["status"]}ed"
	fi

	(
	local key last_value value
	while IFS= read -r key; do
		if [[ $key =~ ^date|status|reason|exit_status$ ]]; then
			continue
		fi
		last_value="${LAST_TEST_RUN["$key"]}"
		value="${TEST_RUN["$key"]}"
		printf '    %s\t%s\t...\t%s\n' "$key" "$last_value" "$value"
	done < <(printf '%s\n' "${!LAST_TEST_RUN[@]}" | sort)

	while IFS= read -r key; do
		# [[ -v array[key] ]] was added in Bash 4.3. Do it this ugly
		# way to support older versions.
		if [[ -n ${LAST_TEST_RUN["$key"]} || ${LAST_TEST_RUN["$key"]-unset} != unset ]]; then
			continue
		fi
		value="${TEST_RUN["$key"]}"
		printf '    %s\t\t...\t%s\n' "$key" "$value"
	done < <(printf '%s\n' "${!TEST_RUN[@]}" | sort)
	) | column -t -s $'\t'
}

_cleanup() {
	if [[ -v TMPDIR ]]; then
		rm -rf "$TMPDIR"
		unset TMPDIR
	fi

	local key value
	for key in "${!TEST_DEV_QUEUE_SAVED[@]}"; do
		value="${TEST_DEV_QUEUE_SAVED["$key"]}"
		echo "$value" >"${TEST_DEV_SYSFS}/queue/${key}"
		unset TEST_DEV_QUEUE_SAVED["$key"]
	done

	if [[ -v RESTORE_CPUS_ONLINE ]]; then
		local cpu
		for cpu in "${!CPUS_ONLINE_SAVED[@]}"; do
			echo "${CPUS_ONLINE_SAVED["$cpu"]}" >"/sys/devices/system/cpu/cpu$cpu/online"
		done
		unset RESTORE_CPUS_ONLINE
	fi

	_exit_cgroup2
}

_call_test() {
	local test_func="$1"
	local seqres="${RESULTS_DIR}/${TEST_NAME}"
	# shellcheck disable=SC2034
	FULL="${seqres}.full"
	declare -A TEST_DEV_QUEUE_SAVED

	_read_last_test_run
	_output_last_test_run

	TEST_RUN["date"]="$(date "+%F %T")"

	mkdir -p "$(dirname "$seqres")"
	# Remove leftovers from last time.
	rm -f "${seqres}" "${seqres}."*

	if [[ -w /dev/kmsg ]]; then
		local dmesg_marker="run blktests $TEST_NAME at ${TEST_RUN["date"]}"
		echo "$dmesg_marker" >> /dev/kmsg
	else
		local dmesg_marker=""
		CHECK_DMESG=0
	fi
	$LOGGER_PROG "run blktests $TEST_NAME"

	trap _cleanup EXIT
	if ! TMPDIR="$(mktemp --tmpdir -p "$OUTPUT" -d "tmpdir.${TEST_NAME//\//.}.XXX")"; then
		return
	fi

	TIMEFORMAT="%Rs"
	pushd . >/dev/null || return
	{ time "$test_func" >"${seqres}.out" 2>&1; } 2>"${seqres}.runtime"
	TEST_RUN["exit_status"]=$?
	popd >/dev/null || return
	TEST_RUN["runtime"]="$(cat "${seqres}.runtime")"
	rm -f "${seqres}.runtime"

	_cleanup

	if ! diff "tests/${TEST_NAME}.out" "${seqres}.out" >/dev/null; then
		mv "${seqres}.out" "${seqres}.out.bad"
		TEST_RUN["status"]=fail
		TEST_RUN["reason"]=output
	elif [[ ${TEST_RUN["exit_status"]} -ne 0 ]]; then
		TEST_RUN["status"]=fail
		TEST_RUN["reason"]="exit"
	elif ! _check_dmesg "$dmesg_marker"; then
		TEST_RUN["status"]=fail
		TEST_RUN["reason"]=dmesg
	else
		TEST_RUN["status"]=pass
	fi
	rm -f "${seqres}.out"

	_write_test_run
	_output_test_run

	if [[ ${TEST_RUN["status"]} = pass ]]; then
		return 0
	else
		case "${TEST_RUN["reason"]}" in
		output)
			diff -u "tests/${TEST_NAME}.out" "${seqres}.out.bad" | awk "
			{
				if (NR > 10) {
					print \"    ...\"
					print \"    (Run 'diff -u tests/${TEST_NAME}.out ${seqres}.out.bad' to see the entire diff)\"
					exit
				}
				print \"    \" \$0
			}"
			;;
		exit)
			echo "    exited with status ${TEST_RUN["exit_status"]}"
			;;
		dmesg)
			echo "    something found in dmesg:"
			awk "
			{
				if (NR > 10) {
					print \"    ...\"
					print \"    (See '${seqres}.dmesg' for the entire message)\"
					exit
				}
				print \"    \" \$0
			}" "${seqres}.dmesg"
			;;
		esac
		return 1
	fi
}

_run_test() {
	TEST_NAME="$1"
	CHECK_DMESG=1
	DMESG_FILTER="cat"

	# shellcheck disable=SC1090
	. "tests/${TEST_NAME}"

	if declare -fF test >/dev/null; then
		if declare -fF requires >/dev/null && ! requires; then
			_output_notrun "$TEST_NAME"
			return 0
		fi

		RESULTS_DIR="$OUTPUT/nodev"
		_call_test test
	else
		if [[ ${#TEST_DEVS[@]} -eq 0 ]]; then
			return 0
		fi

		if declare -fF requires >/dev/null && ! requires; then
			_output_notrun "$TEST_NAME"
			return 0
		fi

		local ret=0
		for TEST_DEV in "${TEST_DEVS[@]}"; do
			TEST_DEV_SYSFS="${TEST_DEV_SYSFS_DIRS["$TEST_DEV"]}"
			if declare -fF device_requires >/dev/null && ! device_requires; then
				_output_notrun "$TEST_NAME => $(basename "$TEST_DEV")"
				continue
			fi
			RESULTS_DIR="$OUTPUT/$(basename "$TEST_DEV")"
			if ! _call_test test_device; then
				ret=1
			fi
		done
		return $ret
	fi
}

_run_group() {
	local tests=("$@")
	local group="${tests["0"]%/*}"

	# shellcheck disable=SC1090
	. "tests/${group}/rc"

	if declare -fF group_requires >/dev/null && ! group_requires; then
		_output_notrun "${group}/***"
		return 0
	fi

	if declare -fF group_device_requires >/dev/null; then
		local i
		for i in "${!TEST_DEVS[@]}"; do
			TEST_DEV="${TEST_DEVS[$i]}"
			TEST_DEV_SYSFS="${TEST_DEV_SYSFS_DIRS["$TEST_DEV"]}"
			if ! group_device_requires; then
				_output_notrun "${group}/*** => $(basename "$TEST_DEV")"
				unset TEST_DEVS["$i"]
			fi
		done
		# Fix the array indices.
		TEST_DEVS=("${TEST_DEVS[@]}")
		unset TEST_DEV
	fi

	local ret=0
	local test_name
	for test_name in "${tests[@]}"; do
		if ! ( _run_test "$test_name" ); then
			ret=1
		fi
	done
	return $ret
}

_find_sysfs_dir() {
	local test_dev="$1"
	local major=$((0x$(stat -L -c '%t' "$test_dev")))
	local minor=$((0x$(stat -L -c '%T' "$test_dev")))
	local dev="$major:$minor"

	local block_dir part_dir
	for block_dir in /sys/block/*; do
		if [[ $(cat "${block_dir}/dev") = "$dev" ]]; then
			echo "$block_dir"
			return
		fi
		for part_dir in "$block_dir"/*; do
			if [[ -r ${part_dir}/dev && $(cat "${part_dir}/dev") = "$dev" ]]; then
				echo "$block_dir"
				return
			fi
		done
	done

	return 1
}

declare -A TEST_DEV_SYSFS_DIRS
_check() {
	# shellcheck disable=SC2034
	SRCDIR="$(realpath src)"

	local test_dev
	for test_dev in "${TEST_DEVS[@]}"; do
		if [[ ! -e $test_dev ]]; then
			_error "${test_dev} does not exist"
		elif [[ ! -b $test_dev ]]; then
			_error "${test_dev} is not a block device"
		fi

		local sysfs_dir
		if ! sysfs_dir="$(_find_sysfs_dir "$test_dev")"; then
			_error "could not find sysfs directory for ${test_dev}"
		fi
		TEST_DEV_SYSFS_DIRS["$test_dev"]="$sysfs_dir"
	done

	local test_name group prev_group
	local tests=()
	local ret=0
	while IFS= read -r -d '' test_name; do
		group="${test_name%/*}"
		if [[ $group != "$prev_group" ]]; then
			prev_group="$group"
			if [[ ${#tests[@]} -gt 0 ]]; then
				if ! ( _run_group "${tests[@]}" ); then
					ret=1
				fi
				tests=()
			fi
		fi
		tests+=("$test_name")
	done < <(_find_tests "$@" | sort -zu)

	if [[ ${#tests[@]} -gt 0 ]]; then
		if ! ( _run_group "${tests[@]}" ); then
			ret=1
		fi
	fi

	return $ret
}

usage () {
	USAGE_STRING="\
usage: $0 [options] [group-or-test...]

Run blktests

Test runs:
  -d, --device-only	 only run tests which use a test device from the
			 TEST_DEVS config setting

  -o, --output=DIR	 output results to the given directory (the default is
			 ./results)

  -q, --quick=SECONDS	 do a quick run (only run quick tests and limit the
			 runtime of longer tests to the given timeout,
			 defaulting to 30 seconds)

  -x, --exclude=TEST	 exclude a test (or test group) from the list of
			 tests to run

Miscellaneous:
  -h, --help             display this help message and exit"

	case "$1" in
		out)
			echo "$USAGE_STRING"
			exit 0
			;;
		err)
			echo "$USAGE_STRING" >&2
			exit 1
			;;
	esac
}

if ! TEMP=$(getopt -o 'do:q::x:h' --long 'quick::,exclude:,output:,help' -n "$0" -- "$@"); then
	exit 1
fi

eval set -- "$TEMP"
unset TEMP

LOGGER_PROG="$(type -P logger)" || LOGGER_PROG=true

if [[ -r config ]]; then
	# shellcheck disable=SC1091
	. config
fi

# Default configuration.
: "${DEVICE_ONLY:=0}"
: "${QUICK_RUN:=0}"
: "${OUTPUT:=results}"
if [[ -v EXCLUDE ]] && ! declare -p EXCLUDE | grep -q '^declare -a'; then
	# If EXCLUDE was not defined as an array, convert it to one.
	# shellcheck disable=SC2190,SC2206
	EXCLUDE=($EXCLUDE)
elif [[ ! -v EXCLUDE ]]; then
	EXCLUDE=()
fi
if [[ -v TEST_DEVS ]] && ! declare -p TEST_DEVS | grep -q '^declare -a'; then
	# If TEST_DEVS was not defined as an array, convert it to one.
	# shellcheck disable=SC2206
	TEST_DEVS=($TEST_DEVS)
elif [[ ! -v TEST_DEVS ]]; then
	TEST_DEVS=()
fi

while true; do
	case "$1" in
		'-d'|'--device-only')
			DEVICE_ONLY=1
			shift
			;;
		'-o'|'--output')
			OUTPUT="$2"
			shift 2
			;;
		'-q'|'--quick')
			QUICK_RUN=1
			# Use the timeout specified on the command line, from
			# the config, or the default.
			TIMEOUT="${2:-${TIMEOUT:-30}}"
			shift 2
			;;
		'-x'|'--exclude')
			# shellcheck disable=SC2190
			EXCLUDE+=("$2")
			shift 2
			;;
		'-h'|'--help')
			usage out
			;;
		'--')
			shift
			break
			;;
		*)
			usage err
			;;
	esac
done

if [[ $QUICK_RUN -ne 0 && ! -v TIMEOUT ]]; then
	_error "QUICK_RUN specified without TIMEOUT"
fi

if [[ $DEVICE_ONLY -ne 0 && ${#TEST_DEVS[@]} -eq 0 ]]; then
	_error "DEVICE_ONLY specified without TEST_DEVS"
fi

# Convert the exclude list to an associative array.
TEMP_EXCLUDE=("${EXCLUDE[@]}")
unset EXCLUDE
declare -A EXCLUDE
for filter in "${TEMP_EXCLUDE[@]}"; do
	EXCLUDE["$filter"]=1
done
unset TEMP_EXCLUDE

mkdir -p "$OUTPUT"
OUTPUT="$(realpath "$OUTPUT")"

_check "$@"
