#!/usr/bin/ash
# -*- mode: sh; indent-tabs-mode: nil; sh-basic-offset: 4; -*-
# vim: et sts=4 sw=4

#  SPDX-License-Identifier: LGPL-2.1+
#
#  Copyright © 2022-2023 Collabora Ltd.
#  Copyright © 2022-2023 Valve Corporation.
#
#  This file is part of steamos-customizations.
#
#  steamos-customizations is free software; you can redistribute it and/or modify
#  it under the terms of the GNU Lesser General Public License as published
#  by the Free Software Foundation; either version 2.1 of the License, or
#  (at your option) any later version.

NEWROOT=/new_root

# these are set by steamos_parse_cmdline based on the kernel cmdline:
steamos_efi=
steamos_factory_reset=0

emergency_shell() {
    launch_interactive_shell --exec
}

ismounted() {
    findmnt "$1" > /dev/null 2>&1
}

vinfo() {
    local x
    while read -r x; do
        msg "$x"
    done
}

to_integer() {
    if [ -n "$1" ] && [ "$1" -ne 0 ]; then
        echo 1
    else
        echo 0
    fi
}

# under dracut we'd just call getarg whenever we needed a command line value
# but initcpio does things a little differently so we cluster all the fetches
# here in one callback:
steamos_parse_cmdline() {
    local key=$1 value=$2

    case $key in
        rd.steamos.factory-reset|steamos.factory-reset)
            steamos_factory_reset=$(to_integer "$value")
            ;;
        rd.steamos.efi|steamos.efi)
            steamos_efi="$value"
            msg "EFI partition from kernel cmdline is \"$steamos_efi\""
            ;;
    esac
}

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

downcase() {
    echo -n "$@" | tr '[:upper:]' '[:lower:]'
}

expand_dev() {
    local dev

    case "$1" in
    LABEL=*)
        dev="/dev/disk/by-label/${1#LABEL=}"
        ;;
    UUID=*)
        dev="${1#UUID=}"
        dev="/dev/disk/by-uuid/$(downcase "$dev")"
        ;;
    PARTUUID=*)
        dev="${1#PARTUUID=}"
        dev="/dev/disk/by-partuuid/$(downcase "$dev")"
        ;;
    PARTLABEL=*)
        dev="/dev/disk/by-partlabel/${1#PARTLABEL=}"
        ;;
    *)
        dev="$1"
        ;;
    esac

    echo "$dev"
}

udev_rules() {
    local partset="${1##*/}"

    # ignore the "other" partitions, so they don't know up in dolphin
    local udisks
    [ "$partset" = "other" ] && udisks=1 || udisks=0

    while read -r name partuuid; do
        cat <<EOF
ENV{ID_PART_ENTRY_SCHEME}=="gpt", ENV{ID_PART_ENTRY_UUID}=="$partuuid", SYMLINK+="disk/by-partsets/$partset/$name", ENV{UDISKS_IGNORE}="$udisks"
EOF
    done < "$1"
}

steamos_generate_partsets() {
    local dev=$1

    msg "Mounting $dev on /mnt"

    mkdir -p /mnt
    mount -o ro "$dev" /mnt 2>&1 | vinfo
    if ! ismounted /mnt; then
        err "Mounting $dev failed"
        emergency_shell
    fi

    mkdir -p /run/udev/rules.d
    for partset in /mnt/SteamOS/partsets/*; do
        [ -e "$partset" ] || continue
        msg "Generating udev rules from $partset"

        # must be after 60-persistent-storage.rules
        udev_rules "$partset" > "/run/udev/rules.d/90-steamos-partsets-${partset##*/}.rules"
    done
    umount /mnt

    udevadm control --reload-rules
    udevadm trigger
    # Explicitly wait for _everything_ to settle. We do not want trigger
    # --settle here since that may lead to deadlock or other issues. See the
    # manual for details how the two differ.
    udevadm settle
}

steamos_setup_partsets() {
    local efi_dev

    msg "Scanning for EFI partition"

    if [ -z "$1" ]; then
        err "EFI partition not found"
        emergency_shell
    fi

    efi_dev=$(expand_dev "$1")

    msg "Waiting for $efi_dev"
    poll_device "$efi_dev"

    steamos_generate_partsets "$efi_dev"
}

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

run_hook() {
    # This hook runs _before_ the real rootfs is mounted
    local rdev

    parse_cmdline steamos_parse_cmdline </proc/cmdline

    ######################################################################
    # reuse dracut : pre-mount : 90 : steamos-udev-rules.sh
    steamos_setup_partsets "$steamos_efi"

    ######################################################################
    # from dracut : mount : 90 : steamos-root.sh
    rdev=/dev/disk/by-partsets/self/rootfs
    msg "Waiting for root device $rdev"
    poll_device "$rdev"

    msg "Setting root target to $rdev"
    # shellcheck disable=SC2034 # used by init_functions:default_mount_handler()
    root=$rdev

    # HERE BE DRAGONS:
    #
    # The whole RO rootfs setup and situation is quite tricky and should
    # probably be reworked. Here's some historical information and current state
    # of affairs:
    #
    # Currently we mount the rootfs (with dracut) without an explicit RO/RW
    # flag. This was required for the ext234 rootfs to work properly. In
    # particular during steamos-readonly we toggle the tune2fs RO flag, which is
    # the default state used by mount, whenever there is no explicit RO/RW
    # passed.
    #
    # At the same time, that caused systemd to try and remount the rootfs (just
    # after the pivot), resulting in failures. That was masked out via a config
    # file.
    #
    # Seemingly the explicit RO mount was also causing issues with btrfs rootfs.
    # Since the `btrfs property set` was unable to toggle the RO flag.
    #
    # Now enter mkinitcpio - it explicitly passes the RO/RW flag at mount with
    # RO being the default. If we keep it as-is, the tune2fs RO flag is ignored
    # and the ext234 fs will always be RO, so one has to manually remount on
    # every boot ... even if they use steamos-readonly disable.
    #
    # On the btrfs front, the situation is similar and/or worse. One needs to
    # remount RW manually even if they want to toggle back to steamos-readonly
    # enable.
    #
    # As a sane-ish default we opt for RW here, so while it means that for
    # ext234 fs the RO tune2fs state will be invalid, yet both fs will just work
    # correctly.
    #
    # In addition/parallel to this behaviour, the usual expectation is to have
    # /etc as non-overlay with RW rootfs, which cannot (trivially) happen. So
    # all in all, we might need something like Android, whereby when the
    # bootloader is unlocked the device is rebooted and some data (/etc overlay
    # in our case) is wiped. Otherwise, any meaningful changes to the rootfs
    # /etc will be masked by the overlay. Ideally steamos-readonly will toggle a
    # file/flag where bootloader (or initrd) will parse it and do the right
    # thing.
    #
    # But all that will happen another day.
    #
    # HERE BE DRAGONS:
    # shellcheck disable=SC2034 # used by init_functions:default_mount_handler()
    rwopt=rw

    # NOTE: under dracut we explicitly mounted the root target. mkinitcpio will
    # implicitly attempt to mount whatever 'root' is set to after 'run_hook'
}

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

# Perform a factory reset
# >>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>> X <<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<
# NOTE: This is never a complete nuke-the-site-from-orbit reset
# a) we boot from /esp, so we can't really fix it as we
#    have nothing to fix it _from_
# b) likewise we don't reset the root fses as we haven't got a source
#    for a vanilla rootfs image
#
# The recovery tool/reset tool _does_ have vanilla data sources for those,
# so it is able to reset them.
# >>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>> X <<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<
#
# The current design suffers a few additional issues, namely:
#
# Currently if the formatting fails, we simply log it and continue booting.
# Which means that if the user has exxpected to have the data wiped, they would
# be surprised. Then on their next (re)boot the formatting will trigger again,
# causing even further confusion/annoyance.
#
# One way to resolve that is to have the system reboot on format failure.
# Although that in itself could lead to the device boot-looping if the failure
# is consistent.
#
# So overall, we should do some reasearch how Android, Apple, others are doing
# it and redesign, or at least reconsider our approach.
#

reset_device_ext4() {
    local device=$1
    local label=$2
    local casefold=$3
    local noreserved=$4

    # not considering it an error if a device we were meant to wipe does not exist
    if ! [ -b "$device" ]; then
        return 0
    fi

    [ "$casefold" -eq 1 ] && fmt_case="-O casefold"
    [ "$noreserved" -eq 1 ] && fmt_nores="-m 0"

    msg "Making ext4 filesystem on device $device"
    msg "Options: with casefold: $casefold, w/o res blocks: $noreserved"
    # shellcheck disable=SC2086 # do not quote these optional arguments
    mkfs.ext4 -qF -L "$label" $fmt_case $fmt_nores "$device"
}

do_reset() {
    local wait_pids

    msg "A factory reset has been requested."

    # Make sure we bail out if the reset fails at any stage
    # we do this to make sure the reset will be re-attempted
    # or resumed if it does not complete here (possibly because
    # the user got bored and leaned on the power button)

    # There is a small chance of a reset loop occurring if the reset cannot
    # complete for fundamental reasons (unable to format filesystem and so
    # forth) BUT
    #
    # a) the device is probably hosed anyway if this happens
    #
    # b) we care more (for now) about doing a genuine reset to stop
    #    leaking private data / things worth actual €£$¥ to the next owner
    #    than we do about [hopefully] unlikely reset loops

    # We want to reset each filesystem in parallel _but_ we must wait for
    # them to finish as we have to release all fds before the pivot to the
    # real sysroot happens:
    # NOTE: the rootfs would need to be reset _before_ the EFI fs if we were
    # handling it, as its fs uuid must be known for the EFI reset - but since
    # we're not touching it everything is parallelisable:

    for cfg in "/esp/efi/steamos/factory-reset"/*.cfg; do
        [ -r "$cfg" ] || continue

        # shellcheck disable=SC2094 # yes, the read/rm is perfectly safe
        while read -r type instance dev casefold noreserved; do
            msg "Processing manifest file $cfg (async)"
            name="${instance##*/}"
            case $type in
                EFI)
                    msg "Reset of efi partition ($instance, $dev) is obsolete, ignoring"
                    rm -f "$cfg"
                    ;;
                VAR|HOME)
                    # these are slow so we want them done in parallel and async
                    # BUT we need to wait until they're all done before proceeding
                    msg "Formatting data partition $dev ($instance)"
                    (
                        if reset_device_ext4 "$dev" "$name" "$casefold" "$noreserved"; then
                            rm -f "$cfg"
                            msg "Reset of $dev ($instance) complete"
                        else
                            err "Reset of $dev ($instance) failed, factory reset incomplete"
                        fi
                    ) &
                    wait_pids="$wait_pids $!"
                    ;;
                *)
                    err "Unexpected SteamOS reset type $type ($instance, $dev)"
                    rm -f "$cfg"
                    ;;
            esac
        done < "$cfg"
    done

    if [ -n "$wait_pids" ]; then
        msg "Waiting for $wait_pids"
        # shellcheck disable=SC2086 # a space separated list
        wait $wait_pids
        msg "Formatting complete"
    fi
}

factory_reset() {
    local want_reset=0
    local cleanup_esp=0

    ########################################################################
    # mount /esp if it isn't mounted, the reset config is located in there
    if [ ! -d /esp/efi ]; then
        local dev="/dev/disk/by-partsets/all/esp"
        msg "Checking ESP partition $dev"
        if ! ismounted "$dev"; then
            msg "Mounting $dev at /esp"
            mkdir -p /esp
            mount "$dev" /esp 2>&1 | vinfo
            cleanup_esp=1
        fi
    fi

    ########################################################################
    # if reset config exists, we want a reset:
    if [ -d "/esp/efi/steamos/factory-reset" ]; then
        for cfg in "/esp/efi/steamos/factory-reset"/*.cfg; do
            [ -r "$cfg" ] || continue

            msg "Factory reset request found in /esp/efi/steamos/factory-reset"
            want_reset=1
            break
        done
    fi

    ########################################################################
    # if we don't already have a reset config then check to see if
    # the bootloader asked us to generate one:
    if [ "$want_reset" -eq 0 ]; then
        want_reset=$1
        if [ "$want_reset" -ne 0 ]; then
            steamos-factory-reset-config
        fi
    fi

    ########################################################################
    # perform the actual reset operations
    if [ "$want_reset" -eq 1 ]; then
        do_reset
    fi

    ########################################################################
    # unmount /esp if we mounted it
    if [ "$cleanup_esp" -eq 1 ]; then
        msg "Unmounting /esp"
        umount /esp
    fi
}

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

mount_var() {
    msg "Mounting /dev/disk/by-partsets/self/var"
    mount "/dev/disk/by-partsets/self/var" "$NEWROOT/var" 2>&1 | vinfo
    if ismounted "$NEWROOT/var"; then
        return
    fi

    err "Mounting /dev/disk/by-partsets/self/var failed! Fallback using tmpfs!"
    mount -t tmpfs -o size=512m tmpfs "$NEWROOT/var" 2>&1 | vinfo
    if ismounted "$NEWROOT/var"; then
        return
    fi

    err "Mounting /dev/disk/by-partsets/self/var failed! Compile the kernel with CONFIG_TMPFS!"
    err "*** Dropping you to a shell; the system will continue"
    err "*** when you leave the shell."
    emergency_shell
}

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

# Mount the etc overlay for steamos atomic
#
# Mount the /etc overlay now (that is, in the initrd), which is needed for
# systemd to successfully commit a generated machine-id on the very first
# boot, and to find this existing machine-id on subsequent boots.
#
# Since the overlay is defined in `/etc/fstab`, one could try to add the
# option `x-initrd.mount`, and let things happen automatically, right?
#
# Well, no, for two reasons. One is that we need to make sure that the
# various directories used for the overlay (upper and work dir) exist,
# which is not the case for a first boot scenario.
#
# Another reason is that, even though systemd-fstab-generator is smart enough
# to prefix the mountpoints found in fstab with `/sysroot`, I don't think
# it's smart enough to do the same with the overlay options `upperdir=`,
# `workdir=` and so on.
#
# For these reasons, we do the job manually here.


setup_etc_overlay() {
    local lowerdir="$NEWROOT/etc"
    local upperdir="$NEWROOT/var/lib/overlays/etc/upper"
    local workdir="$NEWROOT/var/lib/overlays/etc/work"

    # upper dir contains persistent data, create it only if it doesn't exist
    msg "Preparing /etc overlay"
    if [ ! -d "$upperdir" ]; then
        msg "Creating overlay upper directory '$upperdir'"
        rm -fr "$upperdir"
        mkdir -p "$upperdir"
    fi

    # work dir must exist and be empty
    msg "Clearing overlay work directory $workdir"
    rm -fr "$workdir"
    mkdir -p "$workdir"

    # Mount the /etc overlay
    msg "Mounting overlay $upperdir on $lowerdir ($workdir)"
    mount -v \
        -t overlay \
        -o "lowerdir=$lowerdir,upperdir=$upperdir,workdir=$workdir" \
        overlay \
        "$lowerdir" 2>&1 | vinfo

    if ismounted "$lowerdir"; then
        return
    fi

    err "Mounting $upperdir failed: Compile the kernel with CONFIG_OVERLAY_FS!"
    err "*** Dropping you to a shell; the system will continue"
    err "*** when you leave the shell."
    emergency_shell
}

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

# Initialize /var in case of first boot
#
# /var is empty on first boot or after a factory reset. A 'usual' way of
# initializing is to use systemd-tmpfiles, and indeed it's how it's done.
# However it doesn't run early enough to initialize /var/lib/modules in time.
#
# For the record, SteamOS has a custom layout for kernel modules things,
# due to the requirement to have DKMS working with a read-only rootfs,
# so we end up having kernel module things in /var rather than /usr.
# The kernel modules must be available super early at boot, earlier than
# what systemd-tmpfiles provides out of the box.
#
# So systemd-tmpfiles won't work for us. I thought about using it from the
# initrd instead, and using the `--root=/sysroot` option, but that doesn't
# really work either, see <https://github.com/systemd/systemd/issues/12467>.
#
# So, let's do the job manually here.

initialize_var_lib_modules() {

    # Create /var/lib/modules from factory

    local moddir="var/lib/modules"
    local orig="$NEWROOT/usr/share/factory/$moddir"
    local dest="$NEWROOT/$moddir"

    msg "Checking modules source and destination ($orig; $dest)"
    [ -e "$dest" ] && return
    [ -d "$orig" ] || return

    msg "Creating module directory '$dest'"
    mkdir -p "$(dirname "$dest")"

    # purge any half copied content or leftovers
    rm -rf "$dest".new

    msg "Copying $orig to $dest.new"
    if cp -a "$orig" "$dest.new"; then
        msg "Copy successful, installing to $dest"
        mv "$dest.new" "$dest"
        return
    fi

    err "Could not install kernel modules to $dest, system may need rescue"
    emergency_shell
}

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

run_latehook() {
    # This is run after the rootfs is mounted, but before
    # switch_root is invoked to pivot to it:

    # re-parse the command line (settings from run_hook will have
    # disappeared by now):
    parse_cmdline steamos_parse_cmdline </proc/cmdline

    # reuse dracut : pre-pivot : 88 : steamos-factory-reset.sh
    factory_reset "$steamos_factory_reset"

    # reuse dracut : pre-pivot : 89 : steamos-var.sh
    mount_var

    # reuse dracut : pre-pivot : 90 : steamos-etc-overlay.sh
    setup_etc_overlay

    # reuse dracut : pre-pivot : 90 : steamos-var-lib-modules.sh
    initialize_var_lib_modules
}
