udm-init

#!/bin/sh
# Init script for booting into a custom kernel on the UDM (Pro).
#
# Copyright (C) 2021 Fabian Mastenbroek.
#
# 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.

# Perform initial setup of the system by mounting the procfs and sysfs and
# installing busybox.
setup_system() {
    # Create all the symlinks to /bin/busybox
    /bin/busybox --install -s

    mount -t proc proc /proc
    mount -t sysfs sysfs /sys
    mount -t devtmpfs none /dev
}

# Setup watchdog to prevent the device from resetting during the initramfs boot
# procedure. If we do not enable the watchdog, the system will reset in 5 seconds.
setup_watchdog() {
    if [ -e /dev/watchdog ]; then
        echo "udm-init: starting watchdog..."
        watchdog -t 2 /dev/watchdog
    fi
}

# Procedure to wait for the root device to appear. This may not be the case
# immediately when the device boots up. Retry for five times or fail.
setup_wait_root() {
    n=0
    while [ "$n" -lt 25 ] && [ ! -e "$1" ]; do
        n=$(( n + 1 ))
        sleep 2
    done

    [ -e "$1" ]
}

# Setup the modprobe configuration to force load older versions of proprietary
# kernel modules (e.g., ubnthal and ubnt_common).
setup_modprobe_config() {
    mkdir -p "$OVERLAY_UPPER"/etc/modprobe.d
    touch "$OVERLAY_UPPER"/etc/modprobe.d/udm-kernel-tools.conf
    KERNEL_MODDIR=$(find "$OVERLAY_ROOT"/lib/modules/ -mindepth 1 -maxdepth 1 -type d | head -n1)
    KERNEL_RELEASE=$(basename "$KERNEL_MODDIR")
    if [ -d "$KERNEL_MODDIR/extra" ]; then
        find "$KERNEL_MODDIR/extra" -name '*.ko' -not -name 'wireguard.ko' \
            | sed "s/.*\///; s/\.ko//" \
            | awk "{print \"install \" \$0 \" /sbin/modprobe --ignore-install -f -S $KERNEL_RELEASE \" \$0}" \
            > "$OVERLAY_UPPER"/etc/modprobe.d/udm-kernel-tools.conf
    else
        echo "udm-init: no extra modules found"
    fi
}

# Disable the RAM overlay that UniFi OS tries to mount. We already do it.
setup_disable_overlay() {
    mkdir -p "$OVERLAY_UPPER"/sbin
    echo '#!/bin/sh' > "$OVERLAY_UPPER"/sbin/mount_ram_overlay.sh
    echo 'echo RAM overlay already exists' >> "$OVERLAY_UPPER"/sbin/mount_ram_overlay.sh
    chmod +x "$OVERLAY_UPPER"/sbin/mount_ram_overlay.sh
}

# Setup kmod to manage the kernel modules. We need kmod to force insert the
# proprietary kernel modules from Ubiquiti.
setup_kmod() {
    mkdir -p "$OVERLAY_UPPER"/sbin "$OVERLAY_UPPER"/bin
    cp /bin/kmod "$OVERLAY_UPPER"/bin/kmod
    for path in bin/lsmod sbin/depmod sbin/insmod sbin/lsmod sbin/modinfo sbin/modprobe sbin/rmmod; do
        ln -sf /bin/kmod "$OVERLAY_UPPER"/$path
    done
}

# Setup the kernel modules. This copies the kernel modules from the kernel
# into the root filesystem.
setup_modules() {
    mkdir -p "$OVERLAY_UPPER"/lib/modules
    cp -r /lib/modules/* "$OVERLAY_UPPER"/lib/modules
}

# Create a mark to indicate that the system booted via udm-kernel-tools.
setup_mark() {
    echo "Kernel booted via udm-kernel-tools" > "$OVERLAY_UPPER"/.udm-kernel-tools
}

# Function for parsing command line options with "=" in them
get_opt() {
    echo "$@" | cut -d "=" -f 2
}

# Perform initial system setup
setup_system
setup_watchdog

# Defaults
init="/sbin/init"
root="/dev/sdb3"

# Process command line options
for i in $(cat /proc/cmdline); do
    case $i in
    root\=*)
        root=$(get_opt $i)
        ;;
    init\=*)
        init=$(get_opt $i)
        ;;
    esac
done

echo "udm-init: running initial setup [root=$root, init=$init]"

# Wait for root device to appear. This may take a few seconds.
echo "udm-init: waiting for root device..."
if ! setup_wait_root "$root"; then
     echo "udm-init: root device not found, aborting boot..."
     exit 1
fi

OVERLAY=/overlay
OVERLAY_ROOT=$OVERLAY/root_ro
OVERLAY_RW=$OVERLAY/rw
OVERLAY_UPPER=$OVERLAY_RW/upper
OVERLAY_OVERRIDE=$OVERLAY_RW/override
OVERLAY_WORK=$OVERLAY_RW/work

# Mount the root device
mkdir -p $OVERLAY_ROOT
mount "$root" $OVERLAY_ROOT

# Prepare read-write overlay layer
mkdir -p $OVERLAY_RW
mount -t tmpfs none $OVERLAY_RW
mkdir -p $OVERLAY_UPPER $OVERLAY_OVERRIDE $OVERLAY_WORK

# Mount data drive and prepare read-only override layer
ln -s $OVERLAY_ROOT/etc/fstab /etc/fstab
mkdir -p /mnt/data
mount /mnt/data
mount --move /mnt/data $OVERLAY_ROOT/mnt/data
if [ -d "$OVERLAY_ROOT/mnt/data/udm-kernel-tools/root" ]; then
    mount --bind "$OVERLAY_ROOT/mnt/data/udm-kernel-tools/root" $OVERLAY_OVERRIDE
fi

# Mount overlay
mount -t overlay none -olowerdir=$OVERLAY_OVERRIDE:$OVERLAY_ROOT,upperdir=$OVERLAY_UPPER,workdir=$OVERLAY_WORK /root
mount --rbind $OVERLAY /root/overlay

echo "udm-init: initializing module workaround..."
setup_modprobe_config
setup_kmod
setup_disable_overlay

setup_modules
setup_mark

echo "udm-init: booting into UniFi OS..."
# Check if $init exists and is executable
if [ -x "/root$init" ]; then
    # Umount proc and sysfs, but move devtmpfs as we cannot umount it during
    # the initialization phase.
    umount /proc
    umount /sys
    mount --move /dev /root/dev

    # Switch to the new root and execute init
    exec switch_root /root "$init"
fi

# This will only be run if the exec above failed
echo "udm-init: failed to switch_root, aborting boot..."
exit 1