#!/bin/sh
# /etc/initramfs-tools/scripts/init-bottom/overlay

PREQ=""
prereqs() {
    echo "$PREQ"
}

case $1 in
    prereqs)
        prereqs
        exit 0
        ;;
esac

. /scripts/functions

# -- User data partition --
# Change this to your user data partition if not done
USERDATA="/dev/data_partition"

VERSION="1.0.2"

log_begin_msg "[STRAWBERRY_INIT] --- Starting overlayfs mounting process (sbos-init v${VERSION}) ---"

# Check if overlay is enabled
OVERLAY="no"
for arg in $(cat /proc/cmdline)
do
    case "$arg" in
        overlay=*)
            OVERLAY=${arg#overlay=}
            ;;
        *)
            ;;
    esac
done

if [ "$OVERLAY" != "yes" ]
then
    log_warning_msg "[STRAWBERRY_INIT] Warning: Overlay not enabled, leaving the root fs alone"
    exit 0
fi

log_begin_msg "[STRAWBERRY_INIT] Remounting root fs in overlay mode"

# Load overlay kernel module
if ! modprobe -qb overlay
then
    log_failure_msg "[STRAWBERRY_INIT] Error: Cannot load overlayfs kernel module"
    sleep 0.1
    log_begin_msg "[STRAWBERRY_INIT] --- Exited overlayfs mounting process ---"
    exit 0
fi

# -- OverlayFS mount process --
# Initially we have:
#  /           initramfs
#  ${rootmnt}  readonly-mounted real root FS  (mounted from device ${ROOT})
#
# We want:
#  /                      initramfs
#  ${rootmnt}             overlayfs using ${rootmnt}/system and ${rootmnt}/user
#  ${rootmnt}/system      readonly-mounted real root FS  (mounted from device ${ROOT})
#  ${rootmnt}/user        consistent file system used for user data
#  ${rootmnt}/user/upper  overlayfs upper dir
#  ${rootmnt}/user/work   overlayfs work dir
#

# create and mount /system and /user on the initramfs, we will move them later
for dir in /system /user
do
  [ -d "${dir}" ] || mkdir -p "${dir}"
  if [ $? -ne 0 ]
  then
    log_failure_msg "[STRAWBERRY_INIT] Error: Cannot create ${dir}"
    sleep 0.1
    log_begin_msg "[STRAWBERRY_INIT] --- Exited overlayfs mounting process ---"
    exit 0
  fi
done

# set up /user
if ! mount ${USERDATA} /user
then
    log_failure_msg "[STRAWBERRY_INIT] Error: Cannot mount user data (${USERDATA}) on /user"
    sleep 0.1
    log_begin_msg "[STRAWBERRY_INIT] --- Exited overlayfs mounting process ---"
    exit 0
fi

for dir in /user/upper /user/work
do
  if ! mkdir -p ${dir}
  then
      log_failure_msg "[STRAWBERRY_INIT] Error: Cannot create ${dir}"
      sleep 0.1
      log_begin_msg "[STRAWBERRY_INIT] --- Exited overlayfs mounting process ---"
      exit 0
  fi
done

# set up /syytem
if ! mount -o move "${rootmnt}" /system
then
    log_failure_msg "[STRAWBERRY_INIT] Error: Cannot move root fs to /system"
    sleep 0.1
    log_begin_msg "[STRAWBERRY_INIT] --- Exited overlayfs mounting process ---"
    exit 0
fi

# set up an overlayfs on rootmnt
if ! mount -t overlay -o lowerdir=/system,upperdir=/user/upper,workdir=/user/work rootfs "${rootmnt}"
then
    log_failure_msg "[STRAWBERRY_INIT] can't move root fs to /system"

    # try to recover
    if ! mount -o move /system "${rootmnt}"
    then
        panic "[STRAWBERRY_INIT] Error: Recovering the old root fs failed, panicking"
        sleep 0.1
        log_begin_msg "[STRAWBERRY_INIT] --- Exited overlayfs mounting process ---"
        exit 0
    fi

    log_warning_msg "[STRAWBERRY_INIT] Error: moved regular rootfs back into place and continuing"
    sleep 0.1
    log_begin_msg "[STRAWBERRY_INIT] --- Exited overlayfs mounting process ---"
    exit 0
fi

for dir in "${rootmnt}/system" "${rootmnt}/user"
do
  [ -d "${dir}" ] || mkdir -p "${dir}"
  if [ $? -ne 0 ]
  then
    log_failure_msg "[STRAWBERRY_INIT] Error: Cannot create ${dir}"
    sleep 0.1
    log_begin_msg "[STRAWBERRY_INIT] --- Exited overlayfs mounting process ---"
    exit 0
  fi
done

# move /system and /user into place so we can access them later
if ! mount -o move /system "${rootmnt}/system"
then
    log_failure_msg "[STRAWBERRY_INIT] Error: Cannot move /system to ${rootmnt}/system"
    sleep 0.1
    log_begin_msg "[STRAWBERRY_INIT] --- Exited overlayfs mounting process ---"
    exit 0
fi

if ! mount -o move /user "${rootmnt}/user"
then
    log_failure_msg "[STRAWBERRY_INIT] Error: Cannot move /user to ${rootmnt}/user"
    sleep 0.1
    log_begin_msg "[STRAWBERRY_INIT] --- Exited overlayfs mounting process ---"
    exit 0
fi

# populate fstab on the new root

{
    cat <<EOF
# overlayfs mount is in use; changes to the root FS are
# written to tmpfs and will be discarded on reboot.
#
# to disable, pass "overlay=no" on the kernel command line
# (in /boot/cmdline.txt)
#
# the underlying root FS is mounted readonly on /system
# changes are written to /user

EOF

    {
        while read dev dir fstype fsopts freq pass
        do
            if [ "${dir}" = "${rootmnt}" ]
            then
                echo "${dev} / ${fstype} ${fsopts} ${freq} ${pass}"
            fi
        done
    } </proc/mounts

    echo "# original fstab follows"

    {
        while read dev dir fstype fsopts freq pass
        do
            if [ "${dir}" != "/" ]
            then
                echo "${dev} ${dir} ${fstype} ${fsopts} ${freq} ${pass}"
            fi
        done
    } <"${rootmnt}/system/etc/fstab"
} >"${rootmnt}/etc/fstab"

log_end_msg
exit 0
