# ZFS boot stub for initramfs-tools.
#
# In the initramfs environment, the /init script sources this stub to
# override the default functions in the /scripts/local script.
#
# Enable this by passing boot=zfs on the kernel command line.
#

# Source the common init script
. /etc/zfs/zfs-functions

# Paths to what we need - in the initrd, these paths are hardcoded,
# so override the defines in zfs-functions.
ZFS="/sbin/zfs"
ZPOOL="/sbin/zpool"
ZPOOL_CACHE="/etc/zfs/zpool.cache"
export ZFS ZPOOL ZPOOL_CACHE

# Needs to be set for rest of the code to work.
[ -z "${ZFS_INITRD_PRE_MOUNTROOT_SLEEP}" ] && ZFS_INITRD_PRE_MOUNTROOT_SLEEP=0
[ -z "${ZFS_INITRD_POST_MODPROBE_SLEEP}" ] && ZFS_INITRD_POST_MODPROBE_SLEEP=0

# This runs any scripts that should run before we start importing
# pools and mounting any filesystems.
pre_mountroot()
{
	if type run_scripts > /dev/null 2>&1 && \
	    [ -f "/scripts/local-top" -o -d "/scripts/local-top" ]
	then
		check_boolean "${quiet}" || \
		    zfs_log_begin_msg "Running /scripts/local-top"
		run_scripts /scripts/local-top
		check_boolean "${quiet}" || zfs_log_end_msg
	fi

	if type run_scripts > /dev/null 2>&1 && \
	    [ -f "/scripts/local-premount" -o -d "/scripts/local-premount" ]
	then
		check_boolean "${quiet}" || \
		    zfs_log_begin_msg "Running /scripts/local-premount"
		run_scripts /scripts/local-premount
		check_boolean "${quiet}" || zfs_log_end_msg
	fi
}

# Load ZFS modules
# Loading a module in a initrd require a slightly different approach,
# with more logging etc.
load_module_initrd()
{
	if [ "${ZFS_INITRD_PRE_MOUNTROOT_SLEEP}" -gt 0 ]
	then
		check_boolean "${quiet}" || \
			zfs_log_begin_msg "Sleeping for" \
				"${ZFS_INITRD_PRE_MOUNTROOT_SLEEP} seconds..."
		sleep "${ZFS_INITRD_PRE_MOUNTROOT_SLEEP}"
		check_boolean "${quiet}" || zfs_log_end_msg
	fi

	# Wait for all of the /dev/{hd,sd}[a-z] device nodes to appear.
	if type wait_for_udev > /dev/null 2>&1 ; then
		wait_for_udev 10
	elif type wait_for_dev > /dev/null 2>&1 ; then
		wait_for_dev
	fi

	# zpool import refuse to import without a valid mtab
	[ ! -f /proc/mounts ] && mount proc /proc
	[ ! -f /etc/mtab ] && cat /proc/mounts > /etc/mtab

	# Load the module
	load_module "zfs" || return 1

	if [ "${ZFS_INITRD_POST_MODPROBE_SLEEP}" -gt 0 ]
	then
		check_boolean "${quiet}" || \
			zfs_log_begin_msg "Sleeping for" \
				"${ZFS_INITRD_POST_MODPROBE_SLEEP} seconds..."
		sleep "${ZFS_INITRD_POST_MODPROBE_SLEEP}"
		check_boolean "${quiet}" || zfs_log_end_msg
	fi

	return 0
}

# ================================================================

# This is the main function.
mountroot()
{
	local snaporig snapsub destfs pool rfs POOLS

	# ----------------------------------------------------------------
	# I N I T I A L   S E T U P

	# ------------
	# Run the pre-mount scripts from /scripts/local-top.
	pre_mountroot

	# ------------
	# Source the default setup variables.
	[ -r '/etc/default/zfs' ] && . /etc/default/zfs

	# ------------
	# Support debug option
	check_zfs_debug && enable_debug

	# ------------
	# Load ZFS module etc.
	if ! load_module_initrd; then
		disable_plymouth
		echo ""
		echo "Failed to load ZFS modules."
		echo "Manually load the modules and exit."
		emergency_shell
	fi

	# ------------
	# Look for the cache file (if any).
	[ ! -f "${ZPOOL_CACHE}" ] && unset ZPOOL_CACHE

	# ------------
	# Compatibility: 'ROOT' is for Debian GNU/Linux (etc),
	#		 'root' is for Redhat/Fedora (etc),
	#		 'REAL_ROOT' is for Gentoo
	if [ -z "${ROOT}" ]
	then
		[ -n "${root}" ] && ROOT="${root}"

		[ -n "${REAL_ROOT}" ] && ROOT="${REAL_ROOT}"
	fi

	# ------------
	# Where to mount the root fs in the initrd - set outside this script
	# Compatibility: 'rootmnt' is for Debian GNU/Linux (etc),
	#		 'NEWROOT' is for RedHat/Fedora (etc),
	#		 'NEW_ROOT' is for Gentoo
	if [ -z "${rootmnt}" ]
	then
		[ -n "${NEWROOT}" ] && rootmnt="${NEWROOT}"

		[ -n "${NEW_ROOT}" ] && rootmnt="${NEW_ROOT}"
	fi

	# ------------
	# No longer set in the defaults file, but it could have been set in
	# get_pools() in some circumstances. If it's something, but not 'yes',
	# it's no good to us.
	if [ -n "${USE_DISK_BY_ID}" ] && ! check_boolean "${USE_DISK_BY_ID}"; then
		unset USE_DISK_BY_ID
	fi

	# ----------------------------------------------------------------
	# P A R S E   C O M M A N D   L I N E   O P T I O N S

	# This part is the really ugly part - there's so many options and permutations
	# 'out there', and if we should make this the 'primary' source for ZFS initrd
	# scripting, we need/should support them all.
	#
	# Supports the following kernel command line argument combinations
	# (in this order - first match win):
	#
	#	rpool=<pool>			(tries to finds bootfs automatically)
	#	bootfs=<pool>/<dataset>		(uses this for rpool - first part)
	#	rpool=<pool> bootfs=<pool>/<dataset>
	#	-B zfs-bootfs=<pool>/<fs>	(uses this for rpool - first part)
	#	rpool=rpool			(default if none of the above is used)
	#	root=<pool>/<dataset>		(uses this for rpool - first part)
	#	root=ZFS=<pool>/<dataset>	(uses this for rpool - first part, without 'ZFS=')
	#	root=zfs:AUTO			(tries to detect both pool and rootfs
	#	root=zfs:<pool>/<dataset>	(uses this for rpool - first part, without 'zfs:')
	#
	# Option <dataset> could also be <snapshot>

	# ------------
	# Support force option
	# In addition, setting one of zfs_force, zfs.force or zfsforce to
	# 'yes', 'on' or '1' will make sure we force import the pool.
	# This should (almost) never be needed, but it's here for
	# completeness.
	ZPOOL_FORCE=""
	if grep -qiE '(^|[^\\](\\\\)* )(zfs_force|zfs\.force|zfsforce)=(on|yes|1)( |$)' /proc/cmdline
	then
		ZPOOL_FORCE="-f"
	fi

	# ------------
	# Look for 'rpool' and 'bootfs' parameter
	[ -n "${rpool}" ] && ZFS_RPOOL="${rpool#rpool=}"
	[ -n "${bootfs}" ] && ZFS_BOOTFS="${bootfs#bootfs=}"

	# ------------
	# If we have 'ROOT' (see above), but not 'ZFS_BOOTFS', then use
	# 'ROOT'
	[ -n "${ROOT}" -a -z "${ZFS_BOOTFS}" ] && ZFS_BOOTFS="${ROOT}"

	# ------------
	# Check for the `-B zfs-bootfs=%s/%u,...` kind of parameter.
	# NOTE: Only use the pool name and dataset. The rest is not
	#       supported by ZoL (whatever it's for).
	if [ -z "${ZFS_RPOOL}" ]
	then
		# The ${zfs-bootfs} variable is set at the kernel commmand
		# line, usually by GRUB, but it cannot be referenced here
		# directly because bourne variable names cannot contain a
		# hyphen.
		#
		# Reassign the variable by dumping the environment and
		# stripping the zfs-bootfs= prefix.  Let the shell handle
		# quoting through the eval command.
		eval ZFS_RPOOL="$(set | sed -n -e 's,^zfs-bootfs=,,p')"
	fi

	# ------------
	# No root fs or pool specified - do auto detect.
	if [ -z "${ZFS_RPOOL}" -a -z "${ZFS_BOOTFS}" ]
	then
		# Do auto detect. Do this by 'cheating' - set 'root=zfs:AUTO'
		# which will be caught later
		ROOT=zfs:AUTO
	fi

	# ----------------------------------------------------------------
	# F I N D   A N D   I M P O R T   C O R R E C T   P O O L

	# ------------
	if [ "${ROOT}" = "zfs:AUTO" ]
	then
		# Try to detect both pool and root fs.

		check_boolean "${quiet}" || \
		    zfs_log_begin_msg "Attempting to import additional pools."

		# Get a list of pools available for import
		if [ -n "${ZFS_RPOOL}" ]
		then
			# We've specified a pool - check only that
			POOLS="${ZFS_RPOOL}"
		else
			POOLS="$(get_pools)"
		fi

		OLD_IFS="${IFS}" ; IFS=";"
		for pool in ${POOLS}
		do
			[ -z "${pool}" ] && continue

			# Only (re)set what fs to use for boot if it's not
			# already set. Kernel command line is supposed to
			# correct!
			import_pool "${pool}" && \
				[ -z "${ZFS_BOOTFS}" ] && \
					rfs="$(find_rootfs "${pool}")" && \
						ZFS_BOOTFS="${rfs}"
			[ -n "${ZFS_BOOTFS}" ] && break
		done
		IFS="${OLD_IFS}"

		check_boolean "${quiet}" || zfs_log_end_msg "${ZFS_ERROR}"
	else
		# No auto - use value from the command line option.

		# Strip 'zfs:' and 'ZFS='.
		ZFS_BOOTFS="${ROOT#*[:=]}"

		# Stip everything after the first slash.
		ZFS_RPOOL="${ZFS_BOOTFS%%/*}"
	fi

	# Import the pool (if not already done so in the AUTO check above).
	if [ -n "${ZFS_RPOOL}" -a -z "${POOL_IMPORTED}" ]
	then
		# Only (re)set what fs to use for boot if it's not already
		# set. Kernel command line is supposed to correct!
		import_pool "${ZFS_RPOOL}" && \
			[ -z "${ZFS_BOOTFS}" ] && \
				rfs="$(find_rootfs "${ZFS_RPOOL}")" && \
					ZFS_BOOTFS="${rfs}"
	fi

	if [ -z "${POOL_IMPORTED}" ]
	then
		# No pool imported, this is serious!
		disable_plymouth
		echo ""
		echo "Command: ${ZFS_CMD}"
		echo "Message: ${ZFS_STDERR}"
		echo "Error: ${ZFS_ERROR}"
		echo ""
		echo "No pool imported. Manually import the root pool"
		echo "at the command prompt and then exit."
		echo "Hint: Try:  zpool import -R ${rootmnt} -N ${ZFS_RPOOL}"
		emergency_shell
	fi

	# ----------------------------------------------------------------
	# P R E P A R E   R O O T   F I L E S Y S T E M

	if [ -n "${ZFS_BOOTFS}" ]
	then
		# Booting from a snapshot?
		# Will overwrite the ZFS_BOOTFS variable like so:
		#   rpool/ROOT/debian@snap2 => rpool/ROOT/debian_snap2
		echo "${ZFS_BOOTFS}" | grep -q '@' && setup_snapshot_booting "${ZFS_BOOTFS}"
	fi

	if [ -z "${ZFS_BOOTFS}" ]
	then
		# Still nothing! Let the user sort this out.
		disable_plymouth
		echo ""
		echo "Error: Unknown root filesystem - no 'bootfs' pool property and"
		echo "       not specified on the kernel command line."
		echo ""
		echo "Manually mount the root filesystem on ${rootmnt} and then exit."
		echo "Hint: Try:  mount -o zfsutil -t zfs ${ZFS_RPOOL-rpool}/ROOT/system ${rootmnt}"
		emergency_shell
	fi

	# ----------------------------------------------------------------
	# M O U N T   F I L E S Y S T E M S

	# * Ideally, the root filesystem would be mounted like this:
	#
	#     zpool import -R "${rootmnt}" -N "${ZFS_RPOOL}"
	#     zfs mount -o mountpoint=/ "${ZFS_BOOTFS}"
	#
	#   but the MOUNTPOINT prefix is preserved on descendent filesystem
	#   after the pivot into the regular root, which later breaks things
	#   like `zfs mount -a` and the /etc/mtab refresh.
	#
	# * Mount additional filesystems required
	#   Such as /usr, /var, /usr/local etc.
	#   NOTE: Mounted in the order specified in the
	#         ZFS_INITRD_ADDITIONAL_DATASETS variable so take care!
	recursive_mount_filesystems "Using ${ZFS_BOOTFS} as root." "${ZFS_BOOTFS}"
	if [ "${ZFS_ERROR}" != 0 ]
	then
		disable_plymouth
		echo ""
		echo "Command: ${ZFS_CMD}"
		echo "Message: ${ZFS_STDERR}"
		echo "Error: ${ZFS_ERROR}"
		echo ""
		echo "Error: Failed to mount root filesystem '${ZFS_BOOTFS}'."
		echo ""
		echo "Manually mount the filesystem and exit."
		echo "Hint: Try:  mount -o zfsutil -t zfs ${ZFS_BOOTFS} ${rootmnt}"
		emergency_shell
	fi

	# ------------
	# Debugging information
	if check_zfs_debug
	then
		echo

		echo "DEBUG: imported pools:"
		"${ZPOOL}" list -H
		echo

		echo "DEBUG: mounted ZFS filesystems:"
		mount | grep zfs
		echo

		echo "=> waiting for ENTER before continuing because of 'zfsdebug=1'. "
		echo -n "   'c' for shell, 'r' for reboot, any other to continue. "
		read -n1 b

		[ "${b}" = "c" ] && emergency_shell
		[ "${b}" = "r" ] && reboot -f

		set +x
	fi

	# ------------
	# Run local bottom script
	if type run_scripts > /dev/null 2>&1 && \
	    [ -f "/scripts/local-bottom" -o -d "/scripts/local-bottom" ]
	then
		check_boolean "${quiet}" || \
		    zfs_log_begin_msg "Running /scripts/local-bottom"
		run_scripts /scripts/local-bottom
		check_boolean "${quiet}" || zfs_log_end_msg
	fi
}
