#!/bin/bash
# Copyright (c) Microsoft Corporation.
# Licensed under the MIT License.

# Description: This script is designed to set up OverlayFS. It also works with
# DM-Verity environment. It is driven by kernel parameters and is invoked during
# the dracut initramfs phase.

# Kernel Parameters:
# - root: Specifies the path to the root filesystem. This script is designed to
#   support both DM-Verity protected devices and general filesystems. When a
#   DM-Verity protected device is detected (typically '/dev/mapper/root' for
#   systemd), the script performs steps specific to Verity. For non-DM-Verity
#   setups, the script will proceed with the standard OverlayFS setup, ensuring
#   versatility in its application.
# - rd.overlayfs: A comma-separated list defining the OverlayFS configuration.
#   Each entry should specify the lower, upper, work directories, optional
#   volume, and optional mount mode for an OverlayFS instance.

# Behavior:
# - Verifies the presence of the 'dracut-lib' for necessary utilities.
# - Mounts the DM-Verity root filesystem as read-only at a predefined mount
#   point or mount the root device at a predefined mount point.
# - Sets up the OverlayFS based on the provided kernel parameters. If a
#   persistent volume is specified, it's used as the upper layer for the
#   OverlayFS; otherwise, a volatile overlay is created.
# - Mounts the OverlayFS on top of the root filesystem, merging the root with
#   the writable overlay, allowing system modifications without altering the
#   base system.

parse_kernel_cmdline_args() {
    # Ensure that the 'dracut-lib' is present and loaded.
    type getarg > /dev/null 2>&1 || . /lib/dracut-lib.sh

    ROOT_VERITY_MOUNTPOINT="/mnt/root_verity_mnt"
    OVERLAY_MOUNT="/mnt/overlay_mnt"

    # Retrieve the root device.
    [ -z "$root" ] && root=$(getarg root=)
    # Check if we're in a dm-verity environment and the root variable matches
    # the expected path. The path "/dev/mapper/root" is hardcoded here because
    # it is a fixed target name generated by systemd-veritysetup. The name of
    # this dm-verity target is determined by systemd and cannot be changed,
    # hence the explicit check against this specific path.
    if [[ "$root" == *"/dev/mapper/root"* ]]; then
        is_verity=true
    else
        is_verity=false
    fi

    # Retrieve the OverlayFS parameters.
    [ -z "${overlayfs}" ] && overlayfs=$(getarg rd.overlayfs=)
}

# Modified function to mount volatile or persistent volume.
mount_volatile_persistent_volume() {
    local _volume=$1
    local _overlay_mount=$2

    mkdir -p "${_overlay_mount}"

    if [[ "${_volume}" == "volatile" ]]; then
        # Fallback to volatile overlay if no persistent volume is specified.
        echo "No overlayfs persistent volume specified. Creating a volatile overlay."
        mount -t tmpfs tmpfs "${_overlay_mount}" || \
            die "Failed to create overlay tmpfs at ${_overlay_mount}"
    else
        # Check if the specified Overlay RAID volume is present in the system.
        if mdadm --examine --scan | grep -qs "${_volume}"; then
            # If the specified Overlay RAID volume is found, attempt to assemble it.
            mdadm --assemble "${_volume}" || \
                die "Failed to assemble RAID volume."
        fi

        # Mount the specified persistent volume.
        mount "${_volume}" "${_overlay_mount}" || \
            die "Failed to mount ${_volume} at ${_overlay_mount}"
    fi
}

create_overlayfs() {
    local _lower=$1
    local _upper=$2
    local _work=$3
    local _mode=$4

    [ -d "$_lower" ] || die "Unable to create overlay as $_lower does not exist"

    mkdir -p "${_upper}" || die "Failed to create upper directory ${_upper}"
    mkdir -p "${_work}" || die "Failed to create work directory ${_work}"

    # Note for now, the mountpoint / mergedir is set to the same directory as the lowerdir.
    # This means the overlay will be mounted directly on the lower directory.
    # TODO: Add support for a customized mountpoint in future versions.
    mount -t overlay overlay -o "${_mode}",lowerdir="${_lower}",upperdir="${_upper}",workdir="${_work}" "${_lower}" || \
        die "Failed to mount overlay in ${_lower}"
}

mount_overlayfs() {
    local cnt=0
    local volume_mount_with_cnt
    declare -A volume_mount_map

    mkdir -p "${ROOT_VERITY_MOUNTPOINT}"
    if [ "$is_verity" = true ]; then
        echo "Mounting DM-Verity Target"
        mount -o ro "/dev/mapper/root" "${ROOT_VERITY_MOUNTPOINT}" || \
            die "Failed to mount dm-verity root target"
    else
        echo "Mounting regular root"
        # Remove 'block:' prefix if present.
        root_device=$(expand_persistent_dev "${root#block:}")
        mount "$root_device" "${ROOT_VERITY_MOUNTPOINT}" || \
            die "Failed to mount root"
    fi

    echo "Starting to create OverlayFS"
    IFS=';' read -ra overlay_groups <<< "$overlayfs"
    for _group in "${overlay_groups[@]}"; do
        IFS=',' read -r lower upper work volume mode <<< "$_group"

        # Set mode as ro to fulfill Trident specific requirement.
        mode=${mode:-ro}
        # Resolve volume to its full device path.
        volume=$(expand_persistent_dev "$volume")

        if [[ "$volume" == "" ]]; then
            volume_mount_with_cnt="${OVERLAY_MOUNT}/${cnt}"
            mount_volatile_persistent_volume "volatile" $volume_mount_with_cnt
        else
            if [[ -n "${volume_mount_map[$volume]}" ]]; then
                # Volume already mounted, retrieve existing mount point from map.
                volume_mount_with_cnt=${volume_mount_map[$volume]}
            else
                # Not in map, so mount and update the map.
                volume_mount_with_cnt="${OVERLAY_MOUNT}/${cnt}"
                mount_volatile_persistent_volume $volume $volume_mount_with_cnt
                volume_mount_map[$volume]=$volume_mount_with_cnt
            fi
        fi

        echo "Creating OverlayFS with lower: $lower, upper: ${volume_mount_with_cnt}/${upper}, work: ${volume_mount_with_cnt}/${work}, mode: ${mode}"
        create_overlayfs "${ROOT_VERITY_MOUNTPOINT}/${lower}" "${volume_mount_with_cnt}/${upper}" "${volume_mount_with_cnt}/${work}" "${mode}"

        cnt=$((cnt + 1))
    done

    echo "Done Verity Root Mounting and OverlayFS Mounting"
    # Re-mount the verity mount along with overlayfs to the sysroot.
    mount --rbind "${ROOT_VERITY_MOUNTPOINT}" "${NEWROOT}"
}

# Keep a copy of this function here from verity-read-only-root package.
expand_persistent_dev() {
    local _dev=$1

    case "$_dev" in
        LABEL=*)
            _dev="/dev/disk/by-label/${_dev#LABEL=}"
            ;;
        UUID=*)
            _dev="${_dev#UUID=}"
            _dev="${_dev,,}"
            _dev="/dev/disk/by-uuid/${_dev}"
            ;;
        PARTUUID=*)
            _dev="${_dev#PARTUUID=}"
            _dev="${_dev,,}"
            _dev="/dev/disk/by-partuuid/${_dev}"
            ;;
        PARTLABEL=*)
            _dev="/dev/disk/by-partlabel/${_dev#PARTLABEL=}"
            ;;
    esac
    printf "%s" "$_dev"
}

# Parse kernel command line arguments to set environment variables.
# This function populates variables based on the kernel command line, such as overlayfs.
parse_kernel_cmdline_args

# Check if the overlayfs variable is set, indicating that overlay filesystem parameters were found.
# If not set, the process to enable and mount the overlay filesystem will be skipped.
if [ -n "${overlayfs}" ]; then
    mount_overlayfs
else
    echo "OverlayFS parameter not found in kernel cmdline, skipping mount_overlayfs."
fi
