kdeploy

#!/bin/bash
# SPDX-License-Identifier: LGPL-2.1-or-later

set -euo pipefail

################################################################
## Following functions must be self-contained since they may be
## sent to a remote and executed there. Global params won't work

local_sync_to_dest() {
	local tmpdir="$1"
	shift
	local sync_files_args=("$@")

	if [[ -z "$tmpdir" ]]; then
		return
	fi

	rsync -tcKrv "${sync_files_args[@]}" --progress "$tmpdir/." /
}

initrd() {
	local module_only=$1
	shift
	local k=$1

	if [[ "$module_only" != "no" ]]; then
		return
	fi

	# shellcheck disable=SC1091
	. /etc/os-release

	if [[ "$ID_LIKE" == "arch" ]] || [[ "$ID" == "arch" ]]; then
		mkinitcpio -k "/boot/vmlinuz-$k" -g "/boot/initramfs-$k.img"
	elif [[ "$ID" == "ubuntu" ]]; then
		update-initramfs -c -k "$k"
	fi
}

bootloader() {
	local module_only=$1

	if [[ "$module_only" != "no" ]]; then
		return
	fi

	# shellcheck disable=SC1091
	. /etc/os-release

	update-grub
}

reboot() {
	local should_reboot="$1"

	if [[ "$should_reboot" == "yes" ]]; then
		command reboot
	fi
}

sudo-function() {
	(($#)) || { echo "Usage: sudo-function FUNC [ARGS...]" >&2; return 1; }
	sudo bash -c "$(declare -f "$1");$(printf ' %q' "$@")"
}

################################################################

SCRIPT_FILENAME=$(basename "${BASH_SOURCE[0]}")

usage() {
	cat - <<-EOF
	$SCRIPT_FILENAME - Deploy kernel builds
	USAGE:
	    $SCRIPT_FILENAME [OPTIONS]

	Deploy Linux kernel build from the current directory to a local/remote system.
	This command is expected to be used after building a kernel checkout and in the
	directory where that build happened.

	For remote deployment, both ssh and rsync are used - it's recommended to have
	ssh key configured to avoid multiple authentication requests and the remote user
	should be configured to be able to use sudo without password (TODO: remove this requirement)

	OPTIONS:
	    -h, --help			Show this help
	    -i, --include[=] GLOB	Include only specific module glob
	    -C DIR			Change to DIR before running
	    -m, --module-only		Only deploy modules, skip kernel/bootloader changes
	    -r, --remote-host[=] HOST	Host to be used for remote deploy via ssh/rsync
	    -u, --remote-user[=] USER	User to be used for remote deploy via ssh/rsync
	    -p, --remote-port[=] PORT	Port to be used for remote deploy via ssh/rsync
	    --remote-reboot		Reboot remote host after complete [ default: no ]
	    --wait-remote-reboot	Like --remote-reboot (implied), but also wait it
	                                come back [ default: no ]
	    -t, --tmpdir[=] DIR		Directory to use for temporary installation
	    --remote-tmpdir[=] DIR	Set the remote directory to use as temporary location
	                                for installation. An empty string disables it, with
					the same behavior as --no-remote-tmpdir.
					[ default: /tmp/kdeploy ]
	    --no-remote-tmpdir		Disable use of tmpdir on the remote and it's the same
					as passing --remote-tmpdir="". It requires remote
	                                user to have sudo permission without password,
	                                i.e. '<user> ALL=(ALL) NOPASSWD: ALL' to be configured
	                                in sudoers.

	EXAMPLES:

	Install on the same host:

	    $ sudo $SCRIPT_FILENAME

	Deploy kernel to a remote system:

	    $ $SCRIPT_FILENAME -r 192.168.0.2

	Deploy kernel to a VM running on this host with port 4022 forwarded and with a
	different user name (requirement: remote user must have sudo power):

	    $ $SCRIPT_FILENAME -p 4022 -r localhost -u test-user

	Deploy only one module to a remote system:

	    $ $SCRIPT_FILENAME -r remote-system -m -i 'xe.ko*'

	Deploy 2 modules to a remote system:

	    $ $SCRIPT_FILENAME -r remote-system -m -i 'xe.ko*' -i 'i915.ko*'
EOF

}

argerr() {
	echo "invalid argument: $*" >&2
	echo >&2
	usage >&2
	exit 1
}

parse_arg_val() {
	local -n ret=$1
	shift
	local key val n=0

	if [[ "$1" =~ ^([^=]+)(=(.*)) ]]; then
		key=${BASH_REMATCH[1]}
		val=${BASH_REMATCH[3]}
		n=0
	else
		if [[ $# -lt 2 ]]; then
			argerr "Invalid option $1"
		fi
		key="$1"
		val="$2"
		n=1
	fi

	# shellcheck disable=SC2034
	ret=("$n" "$key" "$val")
}

OPT_MODULE_INCLUDE=()
OPT_MODULE_ONLY=no
OPT_REMOTE_USER=
OPT_REMOTE_PORT=
OPT_REMOTE_HOST=
OPT_REMOTE_REBOOT=no
OPT_REMOTE_TMPDIR=/tmp/kdeploy
OPT_WAIT_REMOTE_REBOOT=no
OPT_TMPDIR=
OPT_CHDIR=

SSH_REMOTE_ARGS=()
SYNC_FILES_ARGS=()

parse_args() {
	local args

	while [[ $# -gt 0 ]]; do
		case  "$1" in
		-h|--help)
			usage
			exit 0
			;;
		-C)
			parse_arg_val args "$@"
			shift "${args[0]}"
			OPT_CHDIR="${args[2]}"
			;;
		-i|--include|--include=*)
			parse_arg_val args "$@"
			shift "${args[0]}"
			OPT_MODULE_INCLUDE+=("${args[2]}")
			;;
		-m|--module-only)
			OPT_MODULE_ONLY=yes
			;;
		-r|--remote-host|--remote-host=*)
			parse_arg_val args "$@"
			shift "${args[0]}"
			OPT_REMOTE_HOST="${args[2]}"
			;;
		-u|--remote-user|--remote-user=*)
			parse_arg_val args "$@"
			shift "${args[0]}"
			OPT_REMOTE_USER="${args[2]}"
			;;
		-p|--remote-port|--remote-port=*)
			parse_arg_val args "$@"
			shift "${args[0]}"
			OPT_REMOTE_PORT="${args[2]}"
			;;
		--remote-tmpdir|--remote-tmpdir=*)
			parse_arg_val args "$@"
			shift "${args[0]}"
			OPT_REMOTE_TMPDIR="${args[2]}"
			;;
		--no-remote-tmpdir)
			OPT_REMOTE_TMPDIR=
			;;
		-t|--tmpdir|--tmpdir=*)
			parse_arg_val args "$@"
			shift "${args[0]}"
			OPT_TMPDIR="$(realpath -m "${args[2]}")"
			;;
		--remote-reboot)
			OPT_REMOTE_REBOOT="yes"
			;;
		--wait-remote-reboot)
			OPT_REMOTE_REBOOT="yes"
			OPT_WAIT_REMOTE_REBOOT="yes"
			;;
		*)
			argerr "Unknown option $1"
			;;
		esac
		shift
	done

	if [[ -z "$OPT_REMOTE_HOST" ]]; then
		if [[ -n "$OPT_REMOTE_USER" ]] ||
			[[ -n "$OPT_REMOTE_PORT" ]] ||
			[[ "$OPT_REMOTE_REBOOT" != no ]] ||
			[[ "$OPT_WAIT_REMOTE_REBOOT" != no ]]; then
			argerr "remote options require a remote specified with -r"
		fi
	fi

	if [[ -z "$OPT_TMPDIR" ]]; then
		OPT_TMPDIR=$(mktemp --tmpdir -d tmp.install-kernel.XXXX)
		# shellcheck disable=SC2064
		trap "rm -rf $OPT_TMPDIR" EXIT
	fi

	if [[ -n "$OPT_REMOTE_PORT" ]]; then
		SSH_REMOTE_ARGS+=("-p" "$OPT_REMOTE_PORT")
	fi

	if [[ -n "$OPT_REMOTE_USER" ]]; then
		SSH_REMOTE_ARGS+=("-l" "$OPT_REMOTE_USER")
	fi

	if [[ ${#OPT_MODULE_INCLUDE[@]} -gt 0 ]]; then
		# include specific files given in the command line
		readarray -t SYNC_FILES_ARGS< <(printf -- "--include=%s\n" "${OPT_MODULE_INCLUDE[@]}")

		# include our script to finalize installation
		SYNC_FILES_ARGS+=("--include=kdeploy-remote.sh")

		# recurse subdirs, but exclude anything not included above
		SYNC_FILES_ARGS+=("--include=*/")
		SYNC_FILES_ARGS+=("--exclude=*")
	fi
}

install_kernel() {
	local krel="$1"

	if [[ "$OPT_MODULE_ONLY" != "no" ]]; then
		return
	fi

	mkdir -p "$OPT_TMPDIR/boot/"

	# TODO: figure out arch
	cp arch/x86/boot/bzImage "$OPT_TMPDIR/boot/vmlinuz-$krel"
}

install_modules() {
	make "-j$(nproc)" INSTALL_MOD_PATH="$OPT_TMPDIR" modules_install
}

sync_to_dest() {
	local rsync_path=() rtmp

	if [[ -n "$OPT_REMOTE_HOST" ]]; then
		# send functions to the remote that will finish the installation later
		mkdir -p "$OPT_TMPDIR/tmp/"
		typeset -f initrd bootloader reboot local_sync_to_dest \
			> "$OPT_TMPDIR/tmp/kdeploy-remote.sh"

		if [[ -n "$OPT_REMOTE_TMPDIR" ]]; then
			rtmp=$OPT_REMOTE_TMPDIR
		else
			rsync_path+=("--rsync-path=sudo -u root rsync")
			rtmp="/"
		fi

		rsync -tcKrv \
			-e "ssh ${SSH_REMOTE_ARGS[*]}" \
			"${rsync_path[@]}" \
			"${SYNC_FILES_ARGS[@]}" \
			--progress "$OPT_TMPDIR/." "$OPT_REMOTE_HOST:$rtmp"
	else
		sudo-function local_sync_to_dest "$OPT_TMPDIR" "${SYNC_FILES_ARGS[@]}"
	fi
}

post_kernel_install() {
	local ret rtmp

	if [[ -n "$OPT_REMOTE_HOST" ]]; then
		rtmp=${OPT_REMOTE_TMPDIR:-}
		set +e
		# shellcheck disable=SC2029
		ssh -t "${SSH_REMOTE_ARGS[@]}" "$OPT_REMOTE_HOST" \
			"sudo -s bash -c 'source $rtmp/tmp/kdeploy-remote.sh; 		\
			 local_sync_to_dest $OPT_REMOTE_TMPDIR ${SYNC_FILES_ARGS[*]} && \
			 initrd $OPT_MODULE_ONLY $krel &&	\
			 bootloader $OPT_MODULE_ONLY && 	\
			 reboot $OPT_REMOTE_REBOOT'"
		ret=$?
		set -e

		if [[ $ret != 0 ]] && [[ "$OPT_REMOTE_REBOOT" != "yes" ]]; then
			exit $ret
		fi
	else
		sudo-function initrd "$OPT_MODULE_ONLY" "$krel" && sudo-function bootloader "$OPT_MODULE_ONLY"
	fi
}

wait_reboot() {
	local retries=20 i

	if [[ -z "$OPT_REMOTE_HOST" ]] || [[ "$OPT_WAIT_REMOTE_REBOOT" == "no" ]]; then
		return
	fi

	for (( i=1; i <= retries; i++ )); do
		echo "Trying to re-connect [$i/$retries]"
		if ssh -o BatchMode=yes \
			"${SSH_REMOTE_ARGS[@]}"
			"$OPT_REMOTE_HOST" true; then
			break
		fi
		sleep 3
	done

	if [[ $retries -eq 0 ]]; then
		echo "Max timeout (60s) to reconnect reached" >&2
		exit 1
	fi
}

parse_args "$@"

if [[ -n "$OPT_CHDIR" ]]; then
	cd "$OPT_CHDIR"
fi

krel=$(make kernelrelease)

install_kernel "$krel"
install_modules

sync_to_dest

post_kernel_install "$krel"

wait_reboot