#!/bin/sh

. /lib/dracut-zfs-lib.sh

ZFS_DATASET=""
ZFS_POOL=""

case "${root}" in
	zfs:*) ;;
	*) return ;;
esac

GENERATOR_FILE=/run/systemd/generator/sysroot.mount
GENERATOR_EXTENSION=/run/systemd/generator/sysroot.mount.d/zfs-enhancement.conf

if [ -e "$GENERATOR_FILE" -a -e "$GENERATOR_EXTENSION" ] ; then
	# If the ZFS sysroot.mount flag exists, the initial RAM disk configured
	# it to mount ZFS on root.  In that case, we bail early.  This flag
	# file gets created by the zfs-generator program upon successful run.
	info "ZFS: There is a sysroot.mount and zfs-generator has extended it."
	info "ZFS: Delegating root mount to sysroot.mount."
	# Let us tell the initrd to run on shutdown.
	# We have a shutdown hook to run
	# because we imported the pool.
	# We now prevent Dracut from running this thing again.
	for zfsmounthook in "$hookdir"/mount/*zfs* ; do
		if [ -f "$zfsmounthook" ] ; then
			rm -f "$zfsmounthook"
		fi
	done
	return
fi
info "ZFS: No sysroot.mount exists or zfs-generator did not extend it."
info "ZFS: Mounting root with the traditional mount-zfs.sh instead."

# Delay until all required block devices are present.
udevadm settle

if [ "${root}" = "zfs:AUTO" ] ; then
	ZFS_DATASET="$(find_bootfs)"
	if [ $? -ne 0 ] ; then
		zpool import -N -a ${ZPOOL_IMPORT_OPTS}
		ZFS_DATASET="$(find_bootfs)"
		if [ $? -ne 0 ] ; then
			warn "ZFS: No bootfs attribute found in importable pools."
			export_all -F

			rootok=0
			return 1
		fi
	fi
	info "ZFS: Using ${ZFS_DATASET} as root."
fi

ZFS_DATASET="${ZFS_DATASET:-${root#zfs:}}"
ZFS_POOL="${ZFS_DATASET%%/*}"

if import_pool "${ZFS_POOL}" ; then
	# Load keys if we can or if we need to
	if [ $(zpool list -H -o feature@encryption $(echo "${ZFS_POOL}" | awk -F\/ '{print $1}')) == 'active' ]; then
		# if the root dataset has encryption enabled
		if $(zfs list -H -o encryption "${ZFS_DATASET}" | grep -q -v off); then
			# figure out where the root dataset has its key, the keylocation should not be none
			while true; do
				if [[ $(zfs list -H -o keylocation "${ZFS_DATASET}") == 'none' ]]; then
					ZFS_DATASET=$(echo -n "${ZFS_DATASET}" | awk 'BEGIN{FS=OFS="/"}{NF--; print}')
					if [[ "${ZFS_DATASET}" == '' ]]; then
						rootok=0
						break
					fi
				else
					rootok=1
					break
				fi
			done
			[[ "${rootok}" -eq 0 ]]&& return 1
			# decrypt them
			TRY_COUNT=5
			while [ $TRY_COUNT != 0 ]; do
				zfs load-key "${ZFS_DATASET}"
				[ $? == 0 ] && break
				((TRY_COUNT-=1))
			done
		fi
	fi
	# Let us tell the initrd to run on shutdown.
	# We have a shutdown hook to run
	# because we imported the pool.
	info "ZFS: Mounting dataset ${ZFS_DATASET}..."
	if mount_dataset "${ZFS_DATASET}" ; then
		ROOTFS_MOUNTED=yes
		return 0
	fi
fi

rootok=0