#!/bin/bash
# fsck output of the upper file system (rw) is
# written to /dev/kmsg

# Kernel arguments for this script:
# This script allows one to specify the root read-only device with
# kernel argument:
#     root=
# What follows can be a PARTUUID, PARTLABEL, or some other mechanism
# understood by findfs, including the plain device name.
#
# The rw file system overlayed on top of root defaults to:
#     rootrw='PARTLABEL=user_data'
# This is the first GPT partition found with the name user_data.
#
# Another file system can be specified by PARTUUID, or the
# plain device name, or anything understood by findfs.
# For instance:
#     rootrw=/dev/mmcblk0p10
# to use the SD card.
#
# fsck is not done to the root file system.  It is assumed it cannot be
# corrupted.
# By default, fsck -p [rootrw device] is done.
#
# There are two kernel options governing fsck:
#
# fsck.repair=true|false|force
#    false skips fsck, force adds the -f option.
#    true is the default (fsck with no -f option)
# fsck.mode=preen|yes|no
#    These represent -p, -y, and -n options.
#    Preen (-p) is the default
#
# Additional options exist to specify the init program
# mount options, file system type.  rootrwreset="yes"
# will cause the read/write file system to be erased.
#
# shintramfs starts a shell.  To continue boot from
# shell:
#
# bash-4.4# export ROOT_MOUNT=/mnt
# bash-4.4# export INIT=/sbin/init
# bash-4.4# exec switch_root $ROOT_MOUNT $INIT
# 
#
# U-Boot environment example to add an init overlay parameter:
#
# => printenv args_mmc
# args_mmc=run finduuid;setenv bootargs console=${console} ${optargs} root=PARTUUID=${uuid} rw rootfstype=${mmcrootfstype}
#
# => setenv args_mmc 'run finduuid;setenv bootargs console=${console} ${optargs} root=PARTUUID=${uuid} rw rootfstype=${mmcrootfstype} rootrwreset=yes'


# Enable strict shell mode
set -euo pipefail

PATH=/sbin:/bin:/usr/sbin:/usr/bin

MOUNT="/bin/mount"
UMOUNT="/bin/umount"

INIT="/sbin/init"
ROOT_ROINIT="/sbin/init"

ROOT_MOUNT="/mnt"
ROOT_RODEVICE=""
ROOT_RWDEVICE=""
ROOT_ROMOUNT="/run/media/rfs/ro"
ROOT_RWMOUNT="/run/media/rfs/rw"
ROOT_RWRESET="no"
DO_ERASE_USERDATA_FILE="$ROOT_RWMOUNT/.persistent/mts_do_erase_userdata"

ROOT_ROFSTYPE=""
ROOT_ROMOUNTOPTIONS="bind"
ROOT_ROMOUNTOPTIONS_DEVICE="noatime,nodiratime,discard,ro"

ROOT_RWFSTYPE=""
ROOT_RWMOUNTOPTIONS="rw,noatime,mode=755 tmpfs"
ROOT_RWMOUNTOPTIONS_DEVICE="rw,noatime,nodiratime,discard"

# Arithmetic assignments of 0 are false, and then
# script uses set -e, so use let instead.
DO_FSCK=1
DO_STOP=0
FSCKOPT="-p"


finddevice() {
    DEVICE="$1"
    if dev=$(findfs "${DEVICE}") && ((${#dev})) ; then
        # Replace the rootfs string with the findfs device result
        log "Searched for ${DEVICE} and found $dev"
        echo "$dev"
        return 0
    fi
    log "Could not find $DEVICE"
    return 1
}

early_setup() {
	mkdir -p /proc
	mkdir -p /sys
	$MOUNT -t proc proc /proc
	$MOUNT -t sysfs sysfs /sys
	grep -w "/dev" /proc/mounts >/dev/null || $MOUNT -t devtmpfs none /dev
}

read_args() {
        # Default is to do fsck with -p option on user_data partition.
	[ -z "${CMDLINE+x}" ] && CMDLINE=`cat /proc/cmdline`
	for arg in $CMDLINE; do
		# Set optarg to option parameter, and '' if no parameter was
		# given
		optarg=`expr "x$arg" : 'x[^=]*=\(.*\)' || echo ''`
			case $arg in
			shinitramfs)
                                ((DO_STOP=1)) ;;
			root=*)
				ROOT_RODEVICE=$optarg
                                if rootdev=$(finddevice "${ROOT_RODEVICE}") ; then
                                        # Replace the rootfs string with the findfs device result
                                        ROOT_RODEVICE="${rootdev}"
                                        log "Actual root device to be used: $ROOT_RODEVICE"
                                fi
				;;
			rootfstype=*)
				ROOT_ROFSTYPE="$optarg" ;;
			rootinit=*)
				ROOT_ROINIT=$optarg ;;
			rootoptions=*)
				ROOT_ROMOUNTOPTIONS_DEVICE="$optarg" ;;
			rootrw=*)
				ROOT_RWDEVICE=$optarg
				if rootrwdev="$(finddevice ${ROOT_RWDEVICE})" ; then
                                        # Replace the rootfs string with the findfs device result
                                        ROOT_RWDEVICE="${rootrwdev}"
                                fi
                                ;;
			rootrwfstype=*)
				ROOT_RWFSTYPE="$optarg" ;;
			rootrwreset=*)
				ROOT_RWRESET=$optarg ;;
			rootrwoptions=*)
				ROOT_RWMOUNTOPTIONS_DEVICE="$optarg" ;;
			init=*)
                            INIT=$optarg ;;
                        fsck.mode=*)
                                fsck_mode=$optarg
                                case $fsck_mode in
                                    preen)
                                        FSCKOPT="-p" ;;
                                    yes)
                                        FSCKOPT="-y" ;;
                                    no)
                                        FSCKOPT="-n" ;;
                                esac
                                ;;
                        fsck.repair=*)
                                fsck_repair=$optarg
                                case $fsck_repair in
                                    force)
                                        if ((${#FSCKOPT})) ; then
                                            FSCKOPT+="f"
                                        else
                                            FSCKOPT="-f"
                                        fi
                                        DO_FSCK=1
                                        ;;
                                    yes)
                                        DO_FSCK=1
                                        ;;
                                    no)
                                        DO_FSCK=0
                                        ;;
                                esac

                                ;;
		esac
	done
}

fatal() {
	echo "rorootfs-overlay: $1" >$CONSOLE
	echo >$CONSOLE
	exec bash
}

log() {
	echo "rorootfs-overlay: ${BASH_LINENO[*]}: ${*-""}" >$CONSOLE
}

early_setup

[ -z "${CONSOLE+x}" ] && CONSOLE="/dev/kmsg"

log "Kernel args are:"
log "$(cat /proc/cmdline)"
read_args


mount_and_boot() {
	mkdir -p $ROOT_MOUNT $ROOT_ROMOUNT $ROOT_RWMOUNT

	# Build mount options for read only root file system.
	# If no read-only device was specified via kernel command line, use
	# current root file system via bind mount.
	ROOT_ROMOUNTPARAMS_BIND="-o ${ROOT_ROMOUNTOPTIONS} /"
	if [ -n "${ROOT_RODEVICE}" ]; then
		ROOT_ROMOUNTPARAMS="-o ${ROOT_ROMOUNTOPTIONS_DEVICE} $ROOT_RODEVICE"
		if [ -n "${ROOT_ROFSTYPE}" ]; then
			ROOT_ROMOUNTPARAMS="-t $ROOT_ROFSTYPE $ROOT_ROMOUNTPARAMS"
		fi
	else
                log "ROOT_RODEVICE: empty: ${ROOT_RODEVICE}"
		ROOT_ROMOUNTPARAMS="$ROOT_ROMOUNTPARAMS_BIND"
	fi

	mountlog=$($MOUNT $ROOT_ROMOUNTPARAMS "$ROOT_ROMOUNT" 2>&1)
	mountresult=$?
	if ((mountresult != 0)) ; then
            log "Failed root mount (result $mountresult):  $mountlog"
        else
            log "Mounted root: $MOUNT $ROOT_ROMOUNTPARAMS "$ROOT_ROMOUNT": $mountlog"
            log "Mounted ro root: $(grep "$ROOT_ROMOUNT" /proc/mounts)"
	fi

	# If future init is the same as current file, use $ROOT_ROINIT
	# Tries to avoid loop to infinity if init is set to current file via
	# kernel command line
	if cmp -s "$0" "$INIT"; then
		INIT="$ROOT_ROINIT"
	fi

	# find user_data
        if [[ -z ${ROOT_RWDEVICE} ]] ; then
            userdata=$(findfs PARTLABEL=user_data)
            blkid_out=$(blkid $userdata)
            if [[ -n $userdata ]] ; then
                if [[ $blkid_out =~ TYPE=\"([^\"]+)\" ]]; then
                    log "Detected user_data partition with ${BASH_REMATCH[1]} file system"
                else
                    log "Create new user_data partition ext4 file system"
                    mkfs.ext4 -L user_data -O 64bit $userdata
                fi
                ROOT_RWDEVICE="$userdata"
            fi # Found user_data GPT partition
	fi # Empty ROOT_RWDEVICE

	# Build mount options for read write root file system.
	# If a read-write device was specified via kernel command line
	# or partition table, use it, otherwise default to tmpfs.
	if [ -n "${ROOT_RWDEVICE}" ]; then
		ROOT_RWMOUNTPARAMS="-o $ROOT_RWMOUNTOPTIONS_DEVICE $ROOT_RWDEVICE"
		if [ -n "${ROOT_RWFSTYPE}" ]; then
			ROOT_RWMOUNTPARAMS="-t $ROOT_RWFSTYPE $ROOT_RWMOUNTPARAMS"
		fi
	else
		ROOT_RWMOUNTPARAMS="-t tmpfs -o $ROOT_RWMOUNTOPTIONS"
	fi

        if ((DO_FSCK == 1)) ; then
            if log_result=$(fsck $FSCKOPT $ROOT_RWDEVICE 2>&1) ; then
                log "OK fsck user_data: fsck $FSCKOPT $ROOT_RWDEVICE:"
                log "fsck user_data: $log_result"
            else
                log "ERR fsck user_data: fsck $FSCKOPT $ROOT_RWDEVICE:"
                log "$log_result"
            fi
	fi

	log "user_data: $MOUNT $ROOT_RWMOUNTPARAMS $ROOT_RWMOUNT"
	# Mount read-write file system into initram root file system
	if ! $MOUNT $ROOT_RWMOUNTPARAMS $ROOT_RWMOUNT ; then
		if [ -n "${ROOT_RWDEVICE}" ]; then
			# fsck didn't help. Reformat the partition
			log "Could not mount user_data partition. Erasing ${ROOT_RWDEVICE}."
			mkfs.ext4 -L user_data -O 64bit ${ROOT_RWDEVICE}
			log "mounting user_data: $MOUNT $ROOT_RWMOUNTPARAMS $ROOT_RWMOUNT"
			if ! $MOUNT $ROOT_RWMOUNTPARAMS $ROOT_RWMOUNT ; then
				fatal "Could not mount read-write rootfs."
			fi
			log "${ROOT_RWDEVICE} partition has been erased and mounted successfully."
		else
			fatal "Could not mount user_data partition."
		fi
	fi

	# Reset read-write file system if specified
	log "ROOT_RWRESET=$ROOT_RWRESET ROOT_RWMOUNT=$ROOT_RWMOUNT"
	if [ "yes" == "$ROOT_RWRESET" -a -n "${ROOT_RWMOUNT}" ]; then  # JAK
		log "Removing user_data files"
		rm -rf $ROOT_RWMOUNT/*
		if [ -f $DO_ERASE_USERDATA_FILE ]; then
			log "Removing persistent user_data files"
			shopt -s dotglob
			rm -rf $ROOT_RWMOUNT/*
			shopt -u dotglob
		fi
		log "Completed removal of user_data files"
	fi

	# Determine which unification file system to use
	union_fs_type=""
	if grep -w "overlay" /proc/filesystems >/dev/null; then
		union_fs_type="overlay"
	elif grep -w "aufs" /proc/filesystems >/dev/null; then
		union_fs_type="aufs"
	else
		union_fs_type=""
	fi

	# Create/Mount overlay root file system
	case $union_fs_type in
		"overlay")
			mkdir -p $ROOT_RWMOUNT/upperdir $ROOT_RWMOUNT/work
			if ! $MOUNT -t overlay overlay \
				-o "$(printf "%s%s%s" \
					"lowerdir=$ROOT_ROMOUNT," \
					"upperdir=$ROOT_RWMOUNT/upperdir," \
					"workdir=$ROOT_RWMOUNT/work")" \
				$ROOT_MOUNT ; then
			    log "$MOUNT -t overlay overlay \
				-o \"$(printf \"%s%s%s\" \
					\"lowerdir=$ROOT_ROMOUNT,\" \
					\"upperdir=$ROOT_RWMOUNT/upperdir,\" \
					\"workdir=$ROOT_RWMOUNT/work\")\" \
				$ROOT_MOUNT"
                            fatal "Cound not mount overlay"
			fi
			;;
		"aufs")
			$MOUNT -t aufs i\
				-o "dirs=$ROOT_RWMOUNT=rw:$ROOT_ROMOUNT=ro" \
				aufs $ROOT_MOUNT
			;;
		"")
			fatal "No overlay filesystem type available"
			;;
	esac

	# Move read-only and read-write root file system into the overlay
	# file system
	mkdir -p $ROOT_MOUNT/$ROOT_ROMOUNT $ROOT_MOUNT/$ROOT_RWMOUNT

	$MOUNT -n --move $ROOT_ROMOUNT ${ROOT_MOUNT}/$ROOT_ROMOUNT
	$MOUNT -n --move $ROOT_RWMOUNT ${ROOT_MOUNT}/$ROOT_RWMOUNT


	# Create persistent storage directory and bind-mount it as /var/persistent.
	# The directory name starts with "." so "rm -rf" will skip it on read-write file system reset.
	mkdir -p ${ROOT_MOUNT}/$ROOT_RWMOUNT/.persistent
	mkdir -p ${ROOT_MOUNT}/var/persistent
	$MOUNT -n --bind ${ROOT_MOUNT}/$ROOT_RWMOUNT/.persistent ${ROOT_MOUNT}/var/persistent

	if ((DO_STOP)) ; then
            fatal "Initramfs Shell mode selected -- bash shell"
        fi

	$MOUNT -n --move /proc ${ROOT_MOUNT}/proc
	$MOUNT -n --move /sys ${ROOT_MOUNT}/sys
	$MOUNT -n --move /dev ${ROOT_MOUNT}/dev
 
        CONSOLE="${ROOT_MOUNT}/dev/kmsg"
	log "Mounted filesystems:"
        log "$(cat ${ROOT_MOUNT}/proc/mounts)"
	cd $ROOT_MOUNT

	# switch to actual init in the overlay root file system
	exec switch_root $ROOT_MOUNT $INIT ||
		fatal "Couldn't chroot, dropping to shell"
}

mount_and_boot