#!/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 -n"
UMOUNT="/bin/umount"
MKFS_EXT4="/sbin/mkfs.ext4"

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

ROOT_MOUNT="/mnt"
ROOT_RODEVICE=""
ROOT_RWDEVICE="ubi0:user-data"
ROOT_ROMOUNT="/run/media/rfs/ro"
# ROOT_RWMOUNT is also known as MNT_USER in some scripts
ROOT_RWMOUNT="/run/media/rfs/rw"
VAR_PERSISTENT="/var/persistent"

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

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

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


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
}

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

function logpipe {
    OIFS="${IFS}"
    IFS=$'\n'
    while read a
    do
        IFS="${OIFS}"
        log "$a"
        IFS=$'\n'
	done
	IFS="${OIFS}"
}

# select files for deletion
do_clear_old() {
    # when "clear user data" is requested - mark all non-hidden files for deletion
    log "Clearing user data. Persistent data will be left "

    for FILE_PATH in "$ROOT_RWMOUNT"/*; do
        FILE_NAME=$(basename "$FILE_PATH")
        # rename all files and folders that exist in /mnt/user
        mv "$FILE_PATH" "$ROOT_RWMOUNT/$FILE_NAME.old" 2>&1 | logpipe
    done
    do_remove_old
    fw_setenv default_reset_f 0
}

do_remove_old() {
    shopt -s dotglob
    rm -rf $ROOT_RWMOUNT/*.old
    shopt -u dotglob
}


read_args() {
	log "read_args: ENTER"
	# Default is to do fsck with -p option on user_data partition.
	[ -z "${CMDLINE+x}" ] && CMDLINE=`cat /proc/cmdline`
	log "CMDLINE: ${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
			    ;;
			rootflags=*)
				ROOT_ROMOUNTOPTIONS_DEVICE="$optarg"
				log "ROOT_ROMOUNTOPTIONS_DEVICE: $optarg"
				# Make this Read-only if not set by u-boot
				if ! echo "$ROOT_ROMOUNTOPTIONS_DEVICE" | grep -qE '(^ro,)|(,ro$)|(^ro$)' ; then
				    if [ -n "$ROOT_ROMOUNTOPTIONS_DEVICE" ] ; then
				        ROOT_ROMOUNTOPTIONS_DEVICE="${ROOT_ROMOUNTOPTIONS_DEVICE},ro"
					else
					    ROOT_ROMOUNTOPTIONS_DEVICE=",ro"
					fi
				fi
				;;
			rootfstype=*)
				ROOT_ROFSTYPE="$optarg"
				log "ROOT_ROFSTYPE: $ROOT_ROFSTYPE"
				;;
			rootinit=*)
				ROOT_ROINIT=$optarg ;;
			rootoptions=*)
				ROOT_ROMOUNTOPTIONS_DEVICE="$optarg" ;;
			rootrw=*)
				ROOT_RWDEVICE=$optarg ;;
			rootrwfstype=*)
				ROOT_RWFSTYPE="$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() {
	logger "$1"
	logger ""
	exec bash
}

early_setup

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


read_args

echo "ROOT_ROFSTYPE is $ROOT_ROFSTYPE"
echo "ROOT_RODEVICE is $ROOT_RODEVICE"

# user data file system is cleared if needed.

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
    log "ROMOUNT: $MOUNT $ROOT_ROMOUNTPARAMS $ROOT_ROMOUNT"
	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

	DO_ERASE_PERSISTENT=$(fw_printenv erase_persistent_f | cut -d '"' -f2 | grep -Eo '[0-9]+$' || true)
	log "DO_ERASE_PERSISTENT0: $DO_ERASE_PERSISTENT"
	if ((DO_ERASE_PERSISTENT == 1)) ; then
		log "Erase user-data partition $DEV"
		# Would be nice if ubiupdatevol took volume names instead of numbers
		ubinfo -d 0 -N user-data | logpipe
		[[ $(ubinfo -d 0 -N user-data | grep 'Volume ID') =~ .*[[:space:]]+([0-9]+) ]]
		DEV=/dev/ubi0_"${BASH_REMATCH[1]}"
		ubiupdatevol ${DEV} -t 2>&1 | logpipe
		fw_setenv erase_persistent_f 0
	fi

	# 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 ! [[ $ROOT_RWDEVICE =~ ^ubi ]] ; then
		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
	fi

	log "user_data: $MOUNT $ROOT_RWMOUNTPARAMS $ROOT_RWMOUNT"

	$MOUNT $ROOT_RWMOUNTPARAMS $ROOT_RWMOUNT
	mkdir $ROOT_RWMOUNT/upper
	mkdir $ROOT_RWMOUNT/work
	log "user_data:upper"
	ls $ROOT_RWMOUNT/upper
	log "user_data:work"
	ls $ROOT_RWMOUNT/work

	# Determine which unification file system to use
	union_fs_type=""
	if grep -w "overlay" /proc/filesystems >/dev/null; then
		union_fs_type="overlay"
	else
		fatal "No overlay file system for this kernel, aborting"
	fi

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

	log "$(cat /proc/mounts)"


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

	# If PERSISTENT variables are set too soon in this script, they do not get set.
	DO_CLEAR_PERSISTENT=$(fw_printenv default_reset_f | cut -d '"' -f2 | grep -Eo '[0-9]+$' || true)

	if ((DO_ERASE_PERSISTENT != 1)) ; then
		if ((DO_CLEAR_PERSISTENT == 1)) ; then
			do_clear_old
		else
			do_remove_old
		fi
	fi
	mkdir -p ${ROOT_MOUNT}/proc ${ROOT_MOUNT}/sys /dev ${ROOT_MOUNT}/dev
	$MOUNT --move /proc ${ROOT_MOUNT}/proc
	$MOUNT --move /sys ${ROOT_MOUNT}/sys
	$MOUNT --move /dev ${ROOT_MOUNT}/dev

	log "INIT is $INIT"
	log $(ls /sbin)

	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
