6 files changed, 90 insertions, 11 deletions
diff --git a/contrib/dracut/90zfs/Makefile.am b/contrib/dracut/90zfs/Makefile.am
index a8f08bed7..afcd51022 100644
--- a/contrib/dracut/90zfs/Makefile.am
+++ b/contrib/dracut/90zfs/Makefile.am
@@ -5,6 +5,7 @@ pkgdracut_SCRIPTS = \
 	mount-zfs.sh \
 	parse-zfs.sh \
 	zfs-generator.sh \
+	zfs-load-key.sh \
 	zfs-needshutdown.sh \
 	zfs-lib.sh
 
@@ -14,6 +15,7 @@ EXTRA_DIST = \
 	$(top_srcdir)/contrib/dracut/90zfs/mount-zfs.sh.in \
 	$(top_srcdir)/contrib/dracut/90zfs/parse-zfs.sh.in \
 	$(top_srcdir)/contrib/dracut/90zfs/zfs-generator.sh.in \
+	$(top_srcdir)/contrib/dracut/90zfs/zfs-load-key.sh.in \
 	$(top_srcdir)/contrib/dracut/90zfs/zfs-needshutdown.sh.in \
 	$(top_srcdir)/contrib/dracut/90zfs/zfs-lib.sh.in
 
diff --git a/contrib/dracut/90zfs/module-setup.sh.in b/contrib/dracut/90zfs/module-setup.sh.in
index 1d41c265f..5ae5ad285 100755
--- a/contrib/dracut/90zfs/module-setup.sh.in
+++ b/contrib/dracut/90zfs/module-setup.sh.in
@@ -63,6 +63,7 @@ install() {
 	if [ -n "$systemdutildir" ] ; then
 		inst_script "${moddir}/zfs-generator.sh" "$systemdutildir"/system-generators/dracut-zfs-generator
 	fi
+	inst_hook pre-mount 90 "${moddir}/zfs-load-key.sh"
 	inst_hook mount 98 "${moddir}/mount-zfs.sh"
 	inst_hook cleanup 99 "${moddir}/zfs-needshutdown.sh"
 	inst_hook shutdown 20 "${moddir}/export-zfs.sh"
diff --git a/contrib/dracut/90zfs/mount-zfs.sh.in b/contrib/dracut/90zfs/mount-zfs.sh.in
index e7f217736..36f07d667 100755
--- a/contrib/dracut/90zfs/mount-zfs.sh.in
+++ b/contrib/dracut/90zfs/mount-zfs.sh.in
@@ -56,6 +56,33 @@ 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.
diff --git a/contrib/dracut/90zfs/zfs-generator.sh.in b/contrib/dracut/90zfs/zfs-generator.sh.in
index c6384f583..8cc85a3d3 100755
--- a/contrib/dracut/90zfs/zfs-generator.sh.in
+++ b/contrib/dracut/90zfs/zfs-generator.sh.in
@@ -23,13 +23,6 @@ type getarg >/dev/null 2>&1 || {
 # If root is not ZFS= or zfs: or rootfstype is not zfs
 # then we are not supposed to handle it.
 [ "${root##zfs:}" = "${root}" -a "${root##ZFS=}" = "${root}" -a "$rootfstype" != "zfs" ] && exit 0
-# If root is set to zfs:AUTO, then we are also not
-# supposed to handle it, and it should be handled
-# by the traditional Dracut mount hook.
-# See https://github.com/zfsonlinux/zfs/pull/4558#discussion_r61118952
-if [ "${root}" = "zfs:AUTO" ] ; then
-  exit 0
-fi
 
 rootfstype=zfs
 if echo "${rootflags}" | grep -Eq '^zfsutil$|^zfsutil,|,zfsutil$|,zfsutil,' ; then
@@ -40,9 +33,6 @@ else
     rootflags=zfsutil
 fi
 
-root="${root##zfs:}"
-root="${root##ZFS=}"
-
 echo "zfs-generator: writing extension for sysroot.mount to $GENERATOR_DIR"/sysroot.mount.d/zfs-enhancement.conf >> /dev/kmsg
 
 [ -d "$GENERATOR_DIR" ] || mkdir "$GENERATOR_DIR"
@@ -54,7 +44,14 @@ echo "zfs-generator: writing extension for sysroot.mount to $GENERATOR_DIR"/sysr
     echo "After=zfs-import-scan.service"
     echo "After=zfs-import-cache.service"
     echo "[Mount]"
-    echo "What=${root}"
+    if [ "${root}" = "zfs:AUTO" ] ; then
+      echo "PassEnvironment=BOOTFS"
+      echo 'What=${BOOTFS}'
+    else
+      root="${root##zfs:}"
+      root="${root##ZFS=}"
+      echo "What=${root}"
+    fi
     echo "Type=${rootfstype}"
     echo "Options=${rootflags}"
 } > "$GENERATOR_DIR"/sysroot.mount.d/zfs-enhancement.conf
diff --git a/contrib/dracut/90zfs/zfs-load-key.sh.in b/contrib/dracut/90zfs/zfs-load-key.sh.in
new file mode 100755
index 000000000..d86763fcc
--- /dev/null
+++ b/contrib/dracut/90zfs/zfs-load-key.sh.in
@@ -0,0 +1,52 @@
+#!/bin/bash
+
+# This script only gets executed on systemd systems, see mount-zfs.sh for non-systemd systems
+
+# import the libs now that we know the pool imported
+[ -f /lib/dracut-lib.sh ] && dracutlib=/lib/dracut-lib.sh
+[ -f /usr/lib/dracut/modules.d/99base/dracut-lib.sh ] && dracutlib=/usr/lib/dracut/modules.d/99base/dracut-lib.sh
+. "$dracutlib"
+
+# load the kernel command line vars
+[ -z "$root" ] && root=$(getarg root=)
+# If root is not ZFS= or zfs: or rootfstype is not zfs then we are not supposed to handle it.
+[ "${root##zfs:}" = "${root}" -a "${root##ZFS=}" = "${root}" -a "$rootfstype" != "zfs" ] && exit 0
+
+# There is a race between the zpool import and the pre-mount hooks, so we wait for a pool to be imported
+while true; do
+    zpool list -H | grep -q -v '^$' && break
+    [[ $(systemctl is-failed zfs-import-cache.service) == 'failed' ]] && exit 1
+    [[ $(systemctl is-failed zfs-import-scan.service) == 'failed' ]] && exit 1
+    sleep 0.1s
+done
+
+# run this after import as zfs-import-cache/scan service is confirmed good
+if [[ "${root}" = "zfs:AUTO" ]] ; then
+    root=$(zpool list -H -o bootfs | awk '$1 != "-" {print; exit}')
+else
+    root="${root##zfs:}"
+    root="${root##ZFS=}"
+fi
+
+# if pool encryption is active and the zfs command understands '-o encryption'
+if [[ $(zpool list -H -o feature@encryption $(echo "${root}" | awk -F\/ '{print $1}')) == 'active' ]]; then
+    # check if root dataset has encryption enabled
+    if $(zfs list -H -o encryption "${root}" | 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 "${root}") == 'none' ]]; then
+                root=$(echo -n "${root}" | awk 'BEGIN{FS=OFS="/"}{NF--; print}')
+                [[ "${root}" == '' ]] && exit 1
+            else
+                break
+            fi
+        done
+        # decrypt them
+        TRY_COUNT=5
+        while [ $TRY_COUNT != 0 ]; do
+            zfs load-key "$root" <<< $(systemd-ask-password "Encrypted ZFS password for ${root}: ")
+            [[ $? == 0 ]] && break
+            ((TRY_COUNT-=1))
+        done
+    fi
+fi
diff --git a/contrib/dracut/90zfs/zfs-needshutdown.sh.in b/contrib/dracut/90zfs/zfs-needshutdown.sh.in
index e3d1b59cc..e3d1b59cc 100644..100755
--- a/contrib/dracut/90zfs/zfs-needshutdown.sh.in
+++ b/contrib/dracut/90zfs/zfs-needshutdown.sh.in