# Copyright (C) 2025 Dave Jones <dave.jones@canonical.com>

# This program is free software; you can redistribute it and/or
# modify it under the terms of the GNU General Public License
# as published by the Free Software Foundation; either version 2
# of the License, or (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program; if not, write to the Free Software
# Foundation, Inc., 51 Franklin St, Fifth Floor, Boston, MA 02110-1301,
# USA.

PROC_DT_BOOTLOADER="${FK_PROC_DT_BOOTLOADER:-/proc/device-tree/chosen/bootloader}"
FK_ETC_VALIDATE="${FK_ETC_VALIDATE:-/etc/flash-kernel/piboot-validate}"

parse_find_opts() {
	# A parser for options common to all the find_* functions below. Simply
	# arranges the --suffix argument to be echo'd first, and all other
	# positional parameters to follow

	local sources
	local suffix

	eval set -- $(getopt -o "s:" -l "suffix:" -- "$@")
	while true; do
		case "$1" in
			-s|--suffix)
				suffix="$2"
				shift 2
			;;
			--)
				shift
				break
			;;
			*)
				printf "internal error %s" "$0"
				exit 1
			;;
		esac
	done

	echo "suffix=$suffix" "$@"
}

unique_filenames() {
	# Given a list of filepaths on stdin, filters the list to include only
	# unique file *names* (with the first encountered winning). The filtered
	# list is printed on stdout.
	#
	# For example, given input of "/foo/quux" "/foo/bar" "/bar/quux" the
	# function will echo "/foo/quux" "/foo/bar", eliminating "/bar/quux" as
	# there was already a "quux" file encountered.

	awk '
	{
		file=$0;
		sub("^.*/", "", file);
		if (!(file in files)) files[file]=$0;
	}
	END {
		for (file in files) print files[file];
	}'
}

find_bootloader_assets() {
	# Searches for the Raspberry Pi bootloader files on the paths specified
	# as positional parameters. If --suffix is specified, searches for files
	# with the specified suffix appended (e.g. ".bak").

	local suffix

	set -- $(parse_find_opts "$@")
	suffix="${1#suffix=}"
	shift

	find "$@" -maxdepth 1 \
		-name "bootcode.bin${suffix}" -o \
		-name "start*.elf${suffix}" -o \
		-name "fixup*.dat${suffix}" | unique_filenames
}

find_device_trees() {
	# Searches for base device-trees (*.dtb) on the paths specified as
	# positional parameters. If --suffix is specified, searches for files with
	# the specified suffix appended (e.g. ".bak").

	local suffix

	set -- $(parse_find_opts "$@")
	suffix="${1#suffix=}"
	shift

	find "$@" -maxdepth 1 \
		-name "*.dtb${suffix}" | unique_filenames
}

find_device_tree_overlays() {
	# Searches for device-tree overlays on the paths specified as positional
	# parameters. If --suffix is specified, searches for files with the
	# specified suffix appended (e.g. ".bak").
	#
	# NOTE: the paths specified are assumed to be base directories. The
	# "overlays/" folder should NOT be specified

	local suffix

	set -- $(parse_find_opts "$@")
	suffix="${1#suffix=}"
	shift

	find "$@" -maxdepth 2 \
		-path "*/overlays/*.dtbo${suffix}" -o \
		-path "*/overlays/hat_map.dtb${suffix}" -o \
		-path "*/overlays/overlay_map.dtb${suffix}" -o \
		-path "*/overlays/README${suffix}" | unique_filenames
}

find_kernel_assets() {
	# Searches for Linux kernel assets (vmlinuz and the initramfs) on the paths
	# specified as positional parameters. If --suffix is specified, searches
	# for files with the specified suffix appended (e.g. ".bak").

	local suffix

	set -- $(parse_find_opts "$@")
	suffix="${1#suffix=}"
	shift

	find "$@" -maxdepth 1 \
		-name "vmlinuz${suffix}" -o \
		-name "initrd.img${suffix}" -o \
		-name "cmdline.txt${suffix}" | unique_filenames
}

find_all_assets() {
	# Searches for all boot assets relevant to the Raspberry Pi on the paths
	# specified as positional parameters. If --suffix is specified, searches
	# for files with the specified suffix appended (e.g. ".bak").

	find_bootloader_assets "$@"
	find_kernel_assets "$@"
	find_device_trees "$@"
	find_device_tree_overlays "$@"
}

atomic_cp() {
	# An "atomic" cp operation that will copy $1 to $2. This is only atomic in
	# the sense that, afterward, the destination $2 is guaranteed to be either
	# its original content (if an error occurred) or the new content from $1
	# (if no error occurred), and not some in-between state. It is NOT atomic
	# in the sense of preventing simultaneous operations on the target.

	local src="$1"
	local dest="$2"
	if [ -d "$dest" ]; then
		dest="${dest%/}"/"$(basename "$src")"
	fi
	local tmp="$(mktemp "$dest".XXX)"

	if ! cp "$src" "$tmp"; then
		rm "$tmp"
		return 1
	fi
	if ! mv "$tmp" "$dest"; then
		rm "$tmp"
		return 1
	fi
	return 0 # true
}

needs_migrate() {
	# Determines if the boot files under $1 are laid out in the legacy manner
	# (everything in the root of the boot partition).

	local boot_mnt="$1"

	if  [ -f "$boot_mnt"/current/state ] || \
		grep -q -e "^os_prefix=" "$boot_mnt"/config.txt 2>/dev/null
	then
		return 1 # false
	else
		return 0 # true
	fi
}

write_autoboot() {
	local boot_mnt="$1"
	local boot_dev
	local first_dev
	local mnt

	boot_dev="$(findmnt -n -o SOURCE --target "$boot_mnt")"
	first_dev="$(echo "$boot_dev" | sed -e 's/[0-9][0-9]*$//')1"
	(
		if [ "$first_dev" != "$boot_dev" ]; then
			echo "Mounting recovery partition" >&2
			mnt=$(mktemp -d)
			mount "$first_dev" "$mnt"
			cleanup() {
				echo "Unmounting recovery partition" >&2
				umount $1
				rmdir $1
			}
			trap "cleanup $mnt" EXIT
		else
			mnt="$boot_mnt"
		fi
		"${FK_CHECKOUT:-$FK_DIR}"/migrate-autoboot "$mnt"/autoboot.txt
	)
}

migrate() {
	# Migrates boot files under $1 from the legacy layout (everything under
	# a single directory) to the new layout (all assets stored under current/
	# with a config.txt pointing at this).

	local boot_mnt="$1"
	local src
	local dest

	echo "Removing backup boot files to make space for migration" >&2
	find_all_assets --suffix .bak "$boot_mnt" | {
		while read src; do
			rm -f "$src"
		done
	}

	if [ -f "$boot_mnt"/boot.scr ]; then
		echo "Removing legacy u-boot assets" >&2
		rm -f "$boot_mnt"/uboot*.bin \
			"$boot_mnt"/uboot*.bin.bak \
			"$boot_mnt"/boot.scr \
			"$boot_mnt"/boot.scr.bak
	fi

	echo "Copying current boot files to current/ directory" >&2
	mkdir -p "$boot_mnt"/current
	find_all_assets "$boot_mnt" | {
		while read src; do
			# We deliberately avoid basename here as we want to strip *only*
			# $boot_mnt, leaving sub-directories like overlays/ intact
			dest="${src#"$boot_mnt"}"
			dest="${dest#/}"
			dest="$boot_mnt"/current/"$dest"
			mkdir -p "$(dirname "$dest")"
			cp "$src" "$dest"
		done
	}
	set_state "$boot_mnt" current good
	sync

	echo "Adjusting config.txt to point to current/ directory" >&2
	local tmp="$(mktemp "$boot_mnt"/config.XXX)"
	"${FK_CHECKOUT:-$FK_DIR}"/migrate-config "$boot_mnt"/config.txt > "$tmp"
	mv "$tmp" "$boot_mnt"/config.txt

	echo "Writing autoboot.txt" >&2
	write_autoboot "$boot_mnt"
	sync

	echo "Removing unneeded files" >&2
	{
		find_kernel_assets "$boot_mnt"
		find_device_trees "$boot_mnt"
		find_device_tree_overlays "$boot_mnt"
	} | {
		while read src; do
			rm -f "$src"
		done
		rmdir "$boot_mnt"/overlays
	}

	echo "Migration completed successfully" >&2
}

stable_1() {
	# Step 1 of the transition to stable state. Atomically exchanges the "new/"
	# and "current/" folders. This guarantees current/ exists, and points to
	# "known good" boot assets at all times in the transition. $1 is the boot
	# partition mount point.

	local boot_mnt="$1"

	# NOTE: mv --exchange requires coreutils 9.5 or above (plucky onwards)
	# NOTE: rust-coreutils does not (currently) support --exchange
	gnumv -T --no-copy --exchange "$boot_mnt"/new "$boot_mnt"/current
}

stable_2() {
	# Step 2 of the transition to stable state. Atomically rename the "new/"
	# folder (previously "current/" before step 1) to "old/". $1 is the boot
	# partition mount point.

	local boot_mnt="$1"

	mv "$boot_mnt"/new "$boot_mnt"/old
}

stable_3() {
	# Step 3 (final) of the transition to stable state. Atomically re-write the
	# "current/" state to "good". This must be performed last so that we can
	# detect partially executed transitions and recover from them. $1 is the
	# boot partition mount point.

	local boot_mnt="$1"

	set_state "$boot_mnt" current good
}

set_state() {
	# Atomically replace the "state" file under the specified directory. $1 is
	# the boot partition mount point. $2 is the boot assets sub-directory
	# (typically "new" or "current"). $3 is the new state to write (e.g.
	# "good", "bad", "unknown" or "trying").
	#
	# NOTE: This has the same limitations as atomic_cp above.

	local boot_mnt="$1"
	local dir="$2"
	local state="$3"
	local tmp="$(mktemp "$boot_mnt"/"$dir"/state.XXX)"

	if ! echo "$state" > "$tmp"; then
		rm -f "$tmp"
		return 1
	fi
	if ! mv "$tmp" "$boot_mnt"/"$dir"/state; then
		rm -f "$tmp"
		return 1
	fi
	return 0 # true
}

boot_service_key() {
	# Calculates the "key" for the boot-service. This is essentially the
	# amalgamation (actually simple concatenation) of the tryboot status, and
	# the old, current, and new boot asset states. This "key" dictates which
	# transition (if any) a given service should be performing.
	#
	# The state graph for the whole system, along with which service carries
	# out each transition is as follows:
	#
	# ┌────────┐         p-t-validate         ┌────────┐
	# │ stable │◂─────────────────────────────┤ trying ├─┐
	# └───┬────┘                              └────────┘ │
	#     │               ┌─────┐                  ▴     │
	#     │               ▾ f-k │                  │     │
	#     │           ┌─────────┴┐    p-t-reboot   │     │ p-t-reboot /
	#     └──────────▸│ untested ├─────────────────┘     │ p-t-validate
	#        f-k      └──────────┘                       │
	#                      ▴                             │
	#                      │         ┌────────┐          │
	#                      └─────────┤ failed │◂─────────┘
	#                          f-k   └────────┘
	#
	#   f-k = an execution of flash-kernel
	#   p-t-reboot = the (early) piboot-try-reboot service
	#   p-t-validate = the (late) piboot-try-validate service
	#
	# Refer to comments within the case statements of the boot_service_*
	# functions to see which transition a given key corresponds to.

	local boot_mnt="$1"
	local tryboot_status
	local current_state
	local old_state="-"
	local new_state="-"

	tryboot_status=$(
		od -A n -t u4 --endian big "$PROC_DT_BOOTLOADER"/tryboot |
		tr -d '[:space:]')
	current_state=$(cat "$boot_mnt"/current/state)
	if [ -f "$boot_mnt"/new/state ]; then
		new_state=$(cat "$boot_mnt"/new/state)
	fi
	if [ -f "$boot_mnt"/old/state ]; then
		old_state=$(cat "$boot_mnt"/old/state)
	fi

	echo "${tryboot_status}/${old_state}/${current_state}/${new_state}"
}

boot_service_reboot() {
	# Implements piboot-try --reboot. $1 is the boot partition mount point

	local boot_mnt="$1"
	local key="$(boot_service_key "$boot_mnt")"
	local boot_partition

	boot_partition=$(
		od -A n -t u4 --endian big "$PROC_DT_BOOTLOADER"/partition |
		tr -d '[:space:]')

	case "$key" in
		0/-/good/-|[01]/good/good/-|[01]/-/good/good)
			# stable states with / without older / restored kernels
		;;
		0/-/good/bad)
			# failed state, fallen back to "known good" kernel
		;;
		[01]/-/good/unknown)
			# untested->trying transition
			set_state "$boot_mnt" new trying
			# No need to sync as the reboot will do this anyway
			echo "Rebooting to test new boot assets" >&2
			reboot "$boot_partition tryboot"
		;;
		1/-/good/trying)
			# trying state, awaiting validation by later service
		;;
		0/-/good/trying)
			# trying->failed transition due to reboot
			# This MUST be handled by piboot-try-reboot otherwise a race exists
			# with the piboot-try-validate service which may cause the latter
			# to mark boot assets about to be tested as "bad" before the
			# tryboot reboot can occur
			set_state "$boot_mnt" new bad
			echo "Marked new boot assets bad" >&2
		;;
		0/-/trying/good|0/good/trying/-)
			# trying->stable transition interrupted; leave this for
			# piboot-try-validate to correct
		;;
		*)
			unexpected_state "$key"
		;;
	esac

	return 0
}

boot_service_validate() {
	# Implements piboot-try --validate. $1 is the boot partition mount point

	local boot_mnt="$1"
	local key="$(boot_service_key "$boot_mnt")"
	local boot_partition

	boot_partition=$(
		od -A n -t u4 --endian big "$PROC_DT_BOOTLOADER"/partition |
		tr -d '[:space:]')

	case "$key" in
		0/-/good/-|[01]/good/good/-|[01]/-/good/good)
			# stable states with / without older / restored kernels
		;;
		0/-/good/bad)
			# failed state, fallen back to "known good" kernel
		;;
		[01]/-/good/unknown)
			# untested->trying transition; piboot-try-reboot about to reboot
		;;
		0/-/good/trying)
			# trying->failed transition due to reboot; piboot-try-reboot
			# handles this case
		;;
		1/-/good/trying)
			if ${FK_ETC_VALIDATE}; then
				# trying->stable transition
				stable_1 "$boot_mnt"
				stable_2 "$boot_mnt"
				stable_3 "$boot_mnt"
				echo "Switched new boot assets to current successfully" >&2
			else
				# trying->failed transition due to validate script failure
				echo "Validation failed" >&2
				set_state "$boot_mnt" new bad
				echo "Marked new boot assets bad" >&2
				echo "Rebooting to known good boot assets" >&2
				reboot "$boot_partition"
			fi
		;;
		0/-/trying/good)
			# trying->stable transition interrupted after stable_1; continue
			stable_2 "$boot_mnt"
			stable_3 "$boot_mnt"
			echo "Switched new boot assets to current successfully" >&2
		;;
		0/good/trying/-)
			# trying->stable transition interrupted after stable_2; continue
			stable_3 "$boot_mnt"
			echo "Switched new boot assets to current successfully" >&2
		;;
		*)
			unexpected_state "$key"
		;;
	esac

	return 0
}

boot_service_test() {
	# Implements piboot-try --test. $1 is the boot partition mount point

	local boot_mnt="$1"
	local key="$(boot_service_key "$boot_mnt")"

	case "$key" in
		[01]/-/good/unknown)
			# untested->trying transition; we will double boot next time
			return 0 # true
		;;
		0/-/good/-|[01]/good/good/-|[01]/-/good/good)
			# stable states with / without older / restored kernels
			return 1 # false
		;;
		0/-/good/bad)
			# failed state, fallen back to "known good" kernel
			return 1 # false
		;;
		*)
			unexpected_state "$key"
		;;
	esac
}

boot_service_status() {
	# Implements piboot-try --status. $1 is the boot partition mount
	# point

	local boot_mnt="$1"
	local key="$(boot_service_key "$boot_mnt")"

	case "$key" in
		[01]/-/good/unknown)
			# untested->trying transition; we will double boot next time
			echo
			echo piboot-try:
			fmt << EOF
Untested boot assets present in $boot_mnt/new. Next boot will double-boot; use
piboot-try --reboot to reboot and test new assets immediately without
double-booting
EOF
		;;
		0/-/good/bad)
			# failed state, fallen back to "known good" kernel
			echo
			echo piboot-try:
			fmt << EOF
New boot assets in $boot_mnt/new failed. Fallen back to known good state. Use
piboot-try --reset-new to re-test boot assets. If you are able to determine
the boot failure is due to a bug, please file an issue: ubuntu-bug flash-kernel
EOF
		;;
		0/-/good/-|[01]/good/good/-|[01]/-/good/good)
			# stable states with / without older / restored kernels
		;;
		*)
			unexpected_state "$key"
		;;
	esac
}

unexpected_state() {
	local key="$1"

	echo "piboot is in an invalid or unexpected state: ${key}" >&2
	echo "Please file an issue: ubuntu-bug flash-kernel" >&2
	return 2
}

help() {
	cat << EOF
Usage: $0 [OPTION]

Manages boot assets on the boot partition (under $FK_BOOT_MOUNT).

Options:
  -h, --help     Display this help
  --test         If new, untested boot assets are present, return 0 (indicating
                 that the next boot will be a double-boot). Otherwise, return 1
  --reboot       If new, untested boot assets are present on the boot partition
                 (under $FK_BOOT_MOUNT/new) then immediately reboot to test
                 them. This is used by the piboot-try-reboot service
  --validate     If new boot assets are actively being tested, mark them as
                 "good" or "bad" depending on whether we have fallen back from
                 tryboot mode. This is used by the piboot-try-validate service
  --reset-new    If new boot assets have been marked "bad", reset their state
                 to "unknown"
  --restore-old  Restore "old" boot assets to "current", and "current" to
                 "new"

This may be used by those wishing to avoid the double-boot by initiating it
manually with piboot-try --reboot (in the event that piboot-try --test
indicates that new, untested boot assets are present). See piboot-try(8) for
further information.
EOF
	return 0
}

piboot_try() {
	# A parser for the piboot-try executable

	local op
	local ops=0

	eval set -- $(getopt -o "h" -l "help,reboot,validate,test,reset-new,restore-old,status" -- "$@")
	while true; do
		case "$1" in
			-h|--help)
				help
				return 0
			;;
			--restore-old)
				ops=$((ops | 0x01))
			;;
			--reset-new)
				ops=$((ops | 0x02))
			;;
			--test)
				ops=$((ops | 0x04))
			;;
			--validate)
				ops=$((ops | 0x08))
			;;
			--reboot)
				ops=$((ops | 0x10))
			;;
			--status)
				ops=$((ops | 0x20))
			;;
			--)
				shift
				break
			;;
			*)
				printf "internal error %s" "$0"
				return 1
			;;
		esac
		shift
	done

	# Test for invalid option combinations
	if [ "$ops" -eq 0 ]; then
		echo "No command specified; showing --help" >&2
		help
		return 1
	fi
	if [ "$((ops & 0x08 && ops != 0x08))" -ne 0 ]; then
		echo "Cannot specify --validate with other options" >&2
		return 1
	fi
	if [ "$((ops & 0x04 && ops != 0x04))" -ne 0 ]; then
		echo "Cannot specify --test with other options" >&2
		return 1
	fi
	if [ "$((ops & 0x11 == 0x11))" -ne 0 ]; then
		echo "--restore-old --reboot makes no sense without --reset-new" >&2
		return 1
	fi

	# If the method defined by flash-kernel is not pi-try (if the user has
	# decided to stay on the "pi" method, for instance), bail out
	local method="$(get_machine; get_machine_field "$machine" "Method")"
	if [ "$method" != "pi-try" ]; then
		if [ "$((ops & 0x1F))" -ne 0 ]; then
			echo "flash-kernel method is not pi-try" >&2
			return 1
		else
			return 0
		fi
	fi

	# Carry out options in a specific order; reboot must always be last,
	# restore-old must always be first, etc.
	if [ "$((ops & 0x01))" -ne 0 ]; then
		if ! [ -d "$FK_BOOT_MOUNT"/old ]; then
			echo "No old boot assets to restore" >&2
			return 1
		else
			gnumv -T --no-copy --exchange \
				"$FK_BOOT_MOUNT"/old "$FK_BOOT_MOUNT"/current
			mv "$FK_BOOT_MOUNT"/old "$FK_BOOT_MOUNT"/new
			echo "Restored old boot assets to current, and current to new" >&2
			echo "Former current assets will not be booted unless --reset-new is used" >&2
		fi
	fi
	if [ "$((ops & 0x02))" -ne 0 ]; then
		if ! [ -d "$FK_BOOT_MOUNT"/new ]; then
			echo "No new boot assets to reset" >&2
			return 1
		else
			set_state "$FK_BOOT_MOUNT" new unknown
			echo "Marked new boot assets unknown" >&2
			echo "Please be aware next reboot will boot twice" >&2
		fi
	fi
	if [ "$((ops & 0x04))" -ne 0 ]; then
		boot_service_test "$FK_BOOT_MOUNT"
	fi
	if [ "$((ops & 0x08))" -ne 0 ]; then
		boot_service_validate "$FK_BOOT_MOUNT"
	fi
	if [ "$((ops & 0x10))" -ne 0 ]; then
		boot_service_reboot "$FK_BOOT_MOUNT"
	fi
	if [ "$((ops & 0x20))" -ne 0 ]; then
		boot_service_status "$FK_BOOT_MOUNT"
	fi
}
