#!/bin/sh
# Copyright (c) Fathi Boudra <fathi.boudra@linaro.org>
# All rights reserved.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions
# are met:
# 1. Redistributions of source code must retain the above copyright
#    notice, this list of conditions and the following disclaimer.
# 2. Redistributions in binary form must reproduce the above copyright
#    notice, this list of conditions and the following disclaimer in the
#    documentation and/or other materials provided with the distribution.
#
# THIS SOFTWARE IS PROVIDED BY THE REGENTS AND CONTRIBUTORS ``AS IS'' AND
# ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
# ARE DISCLAIMED.  IN NO EVENT SHALL THE REGENTS OR CONTRIBUTORS BE LIABLE
# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS
# OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION)
# HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT
# LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY
# OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF
# SUCH DAMAGE.

# we must be root
[ $(whoami) = "root" ] || { echo "E: You must be root" && exit 1; }

# we must have few tools
RESIZE2FS=$(which resize2fs) || { echo "E: You must have resize2fs" && exit 1; }

# find root device
# ROOT_DEVICE=$(findmnt --noheadings --output=SOURCE / | cut -d'[' -f1)
DEVICES=$(ls /sys/fs/ext4/)

# userfs is the one partition shared across both RAUC A/B rootfs slots -
# unlike the marker this script uses to avoid re-running itself (disabling
# its own systemd unit, which lives in /etc and so is reset every time a
# fresh rootfs-a/rootfs-b image is installed), a marker recorded here
# survives OTA updates and slot switches, so userfs is only ever actually
# grown/fscked once, not on every first boot into a newly-installed slot.
USERFS_RESIZE_MARKER=/usr/local/.resize-helper-done
USERFS_ALREADY_RESIZED=0
[ -f "$USERFS_RESIZE_MARKER" ] && USERFS_ALREADY_RESIZED=1

# On the SD card build, userfs is created as a small fixed-size GPT
# partition (200M, see the .wks.in) with the rest of the card left
# genuinely unpartitioned - resize2fs/systemd-growfs only ever grow a
# filesystem to fill its *existing* partition boundary, they never touch
# the partition table itself, so without this neither one has anything
# to grow into and userfs stays stuck at 200M forever. (eMMC doesn't hit
# this: its flashlayout sizes userfs to consume the whole remaining
# capacity directly at flash time, so there's normally no free space
# after it for this to find or need to act on - a safe no-op there.)
#
# Confirmed on hardware: after writing today's SD card image and
# booting, userfs was stuck at 200M despite both the initrd-stage and
# this script's own resize attempts completing "successfully" - because
# both were only ever resizing the filesystem within an unchanged 200M
# partition. Extending the GPT partition first (guarded to only ever do
# so when this is genuinely the last partition on the disk with real
# free space after it, so it can't collide with a following partition
# or fire pointlessly on a normal, already-correctly-sized eMMC) fixed
# it: verified on both DUTs' SD cards, userfs grew from 200M to ~21GB.
extend_partition_if_last_and_free_space() {
    device="$1"
    disk="/dev/$(lsblk -lno PKNAME "$device" 2>/dev/null)"
    partnum=$(cat "/sys/class/block/$(basename "$device")/partition" 2>/dev/null)
    [ -n "$disk" ] && [ -n "$partnum" ] || return 0
    [ -f /usr/sbin/sgdisk ] || return 0

    # Only proceed if $device is the last partition on $disk and there's
    # real free space after it - otherwise this is either not our
    # situation (already fills the disk, e.g. eMMC) or would be unsafe
    # (extending anything but the true last partition risks colliding
    # with whatever comes after it).
    last_part_line=$(sgdisk -p "$disk" 2>/dev/null | awk '$1 ~ /^[0-9]+$/{line=$0} END{print line}')
    last_partnum=$(echo "$last_part_line" | awk '{print $1}')
    [ "$last_partnum" = "$partnum" ] || return 0

    end_sector=$(echo "$last_part_line" | awk '{print $3}')
    # sgdisk -e first: on an image dd'd onto a larger physical disk than
    # it was built for, the backup GPT header (and the primary header's
    # own recorded "last usable sector") still reflect the original,
    # smaller image size until this runs once - without it, the free
    # space after the last partition isn't visible to sgdisk at all, and
    # the extend below silently does nothing.
    sgdisk -e "$disk" >/dev/null 2>&1
    last_usable=$(sgdisk -p "$disk" 2>/dev/null | sed -n 's/.*last usable sector is \([0-9]*\).*/\1/p')
    [ -n "$last_usable" ] || return 0
    # Require a real, meaningful amount of free space (>16MiB) before
    # bothering - a few sectors of alignment slack isn't worth acting on.
    free_sectors=$((last_usable - end_sector))
    [ "$free_sectors" -gt 32768 ] || return 0

    echo "Extending $device ($disk partition $partnum) to use $free_sectors free sectors" > /dev/kmsg
    was_mounted=$(findmnt --noheadings --output=TARGET "$device" 2>/dev/null)
    [ -n "$was_mounted" ] && umount "$device"

    start_sector=$(sgdisk -i "$partnum" "$disk" 2>/dev/null | sed -n 's/^First sector: \([0-9]*\).*/\1/p')
    type_guid=$(sgdisk -i "$partnum" "$disk" 2>/dev/null | sed -n 's/^Partition GUID code: \([0-9A-Fa-f-]*\).*/\1/p')
    unique_guid=$(sgdisk -i "$partnum" "$disk" 2>/dev/null | sed -n "s/^Partition unique GUID: \([0-9A-Fa-f-]*\).*/\1/p")
    part_name=$(sgdisk -i "$partnum" "$disk" 2>/dev/null | sed -n "s/^Partition name: '\(.*\)'$/\1/p")

    if [ -n "$start_sector" ] && [ -n "$type_guid" ] && [ -n "$unique_guid" ]; then
        sgdisk -d "$partnum" \
               -n "${partnum}:${start_sector}:0" \
               -t "${partnum}:${type_guid}" \
               -c "${partnum}:${part_name}" \
               -u "${partnum}:${unique_guid}" \
               "$disk"
        resizepart "$disk" "$partnum" $((last_usable - start_sector + 1)) 2>/dev/null || partprobe "$disk" 2>/dev/null
    fi

    [ -n "$was_mounted" ] && mount "$device" "$was_mounted"
}

echo "RESIZE-HELPER START" > /dev/kmsg

# disable service to be sure it will be executed only one time
if [ -f /bin/systemctl ]; then
    /bin/systemctl --no-reload disable resize-helper.service
else
    rm /etc/rc[35].d/S22resize-helper.sh &> /dev/null
fi
sync

if `grep -q nfsroot /proc/cmdline` ; then
    exit 0;
fi

if [ -f /lib/systemd/systemd-growfs ]; then
    echo "RESIZE-HELPER: Using systemd-growfs" > /dev/kmsg
    # force to mount partitions
    if [ -z $(findmnt --noheadings --output=SOURCE /boot|cut  -d'[' -f1) ];
    then
        [ -f /sbin/mount-partitions.sh ] && /sbin/mount-partitions.sh stop
    fi
    if [ "$USERFS_ALREADY_RESIZED" = "0" ]; then
        # Resolve via the root filesystem's own disk, not via /usr/local's
        # current mount state - mount-partitions.sh stop, just above, may
        # have already unmounted /usr/local by this point (it's in
        # MOUNT_PARTITIONS_LIST), which would make a findmnt-on-/usr/local
        # lookup here return nothing despite the device still existing.
        root_device=$(findmnt --noheadings --output=SOURCE / 2>/dev/null | cut -d'[' -f1)
        mmc_disk=""
        case "$root_device" in
            /dev/mmcblk0*) mmc_disk=/dev/mmcblk0 ;;
            /dev/mmcblk1*) mmc_disk=/dev/mmcblk1 ;;
            /dev/mmcblk2*) mmc_disk=/dev/mmcblk2 ;;
        esac
        userfs_device=""
        if [ -n "$mmc_disk" ]; then
            userfs_part=$(lsblk -lno NAME,PARTLABEL "$mmc_disk" 2>/dev/null | awk '$2=="userfs"{print $1}')
            [ -n "$userfs_part" ] && userfs_device="/dev/${userfs_part}"
        fi
        [ -n "$userfs_device" ] && extend_partition_if_last_and_free_space "$userfs_device"
        /lib/systemd/systemd-growfs /usr/local/
        touch "$USERFS_RESIZE_MARKER"
    fi
else
    echo "RESIZE-HELPER: Using directly resize2fs" > /dev/kmsg
    # umount partition before to resize it
    [ -f /sbin/mount-partitions.sh ] && /sbin/mount-partitions.sh stop
    for device in ${DEVICES}; do
        PART_TABLE_TYPE=$(udevadm info --query=property --name=/dev/${device} | grep '^ID_PART_TABLE_TYPE=' | cut -d'=' -f2)
        if [ -e /dev/${device} ]; then
            if [ "$PART_TABLE_TYPE" = "gpt" ]; then
                PART_LABEL=$(udevadm info --query=property --name=/dev/${device} | grep '^ID_PART_ENTRY_NAME=' | cut -d'=' -f2)
                if [ "$PART_LABEL" = "userfs" ] && [ "$USERFS_ALREADY_RESIZED" = "0" ]; then
                    echo "Resize /dev/${device}"
                    extend_partition_if_last_and_free_space "/dev/${device}"
                    e2fsck -f -y -c -C 0 "/dev/${device}"
                    ${RESIZE2FS} "/dev/${device}"
                    mkdir -p /tmp/userfs-mark
                    mount "/dev/${device}" /tmp/userfs-mark
                    touch "/tmp/userfs-mark/$(basename "$USERFS_RESIZE_MARKER")"
                    umount /tmp/userfs-mark
                    rmdir /tmp/userfs-mark
                fi
            fi
        fi
    done
fi

df -h > /dev/kmsg

#echo "RESIZE-HELPER: For integrity of file-system after a resize2fs, the system need to be rebooted" > /dev/kmsg
#echo "RESIZE-HELPER REBOOTING of system" > /dev/kmsg
echo "RESIZE-HELPER FINISH" > /dev/kmsg
