#!/bin/sh
# Copyright 2022 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

# Retrieved from https://cos.googlesource.com/cos/overlays/board-overlays/+/refs/heads/master/project-lakitu/dev-util/cos-dev-tools/files/

set -e

PKG=$1

if ! grep -q '^ID=cos$' /etc/os-release; then
    echo "Does not appear to be a COS instance"
    echo "Refusing to replace kernel"
    exit 1
fi

if [ -z "${PKG}" ]; then
    echo "Usage: $0 cos-kernel-A.B.C-arch.txz"
    exit 1
fi
VERSION=$(tar tvfJ "${PKG}" boot/ | grep /config- | sed s/.*config-//)

if [ -z "${VERSION}" ]; then
    echo "Could not find kernel version from /boot/config-*"
    echo "Check if package is valid"
    exit 1
fi

echo "Kernel version: ${VERSION}"
mount -o rw,remount / || true
if mount | grep "on / " | grep -q ro, ; then
    echo "Could not remount rootfs as read-write"
    exit 1
fi
echo "Remounted rootfs as read-write"

# We have to do this before removing the existing modules since we might need
# to load a file system driver (such as vfat) to mount the EFI partition.
ESPMOUNT=$(mktemp -d)

# Get the EFI partition of the root device. The steps are:
# 1. Get the block devices.
# 2. Filter to the block dev that has a descendant with "/" as a mount point.
# 3. Filter to the objects that have a device maj:min number.
# 4. Filter to the leaf partitions (no children).
# 5. Take the partition name from the 12th partition.
#
# This avoids several edge cases:
# 1. The /dev/disk/by-partlabel/EFI-SYSTEM symlink can be overwritten if a second disk is
#    mounted.
# 2. There may be virtual mounts for the root partition, so it might not be at the top level.
# 3. The device minor number for the 12th partition on the root device might not be 12 on
#    NVME since NVME devices share a device major number and increment the minor number for
#    each partition.
readonly ESPNAME=$(
    lsblk --json |
        jq '.blockdevices.[]' |
        jq 'select(.. | (.mountpoints? // []).[] == "/")' |
        jq '.. | select(."maj:min"? // "" != "")' |
        jq 'select(.children == null)' |
        jq -sr '.[11].name'
)
readonly ESPDEV="/dev/${ESPNAME}"
mount "${ESPDEV}" "${ESPMOUNT}"

if [ -d "/lib/modules/${VERSION}" ]; then
    rm -rf "/lib/modules/${VERSION}"
    echo "Modules for ${VERSION} removed"
fi

tar -C / -x -J -f "${PKG}"
ln -sf "vmlinuz-${VERSION}" /boot/vmlinuz
echo "rootfs updated"

# Replace the boot image.
BOOT_IMAGE="${ESPMOUNT}/$(sed 's/.*BOOT_IMAGE=\([^ ]*\) .*/\1/' /proc/cmdline)"
echo "Updating boot image ${BOOT_IMAGE}"
if [ ! -f "${BOOT_IMAGE}" ]; then
    echo "Boot image '${BOOT_IMAGE}' is not a regular file"
    umount "${ESPMOUNT}"
    exit 1
fi
rm -f "${BOOT_IMAGE}"
cp -f "/boot/vmlinuz-${VERSION}" "${BOOT_IMAGE}"
umount "${ESPMOUNT}"
rm -rf "${ESPMOUNT}"
echo "EFI system partition updated"

mount -o ro,remount / || true
echo "Remounted rootfs as read-only"
echo "Rebooting the system to activate the new kernel"

reboot