Run zfs load-key if needed in dracut

'zfs load-key -a' will only be called if needed.  If a dataset not
needed for boot does not have its key loaded (home directories for
example) boot can still continue.

zfs:AUTO was not working via dracut, so we still need the generator
script to do its thing.

Reviewed-by: Richard Yao <ryao@gentoo.org>
Reviewed-by: Manuel Amador (Rudd-O) <rudd-o@rudd-o.com>
Reviewed-by: Brian Behlendorf <behlendorf1@llnl.gov>
Reviewed-by: loli10K <ezomori.nozomu@gmail.com>
Signed-off-by: Matthew Thode <mthode@mthode.org>
Closes #6982 
Closes #7004

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
 

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"

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.

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

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

etc/systemd/system/zfs-import-cache.service.in

@@ -14,6 +14,7 @@ Type=oneshot
 RemainAfterExit=yes
 ExecStartPre=/sbin/modprobe zfs
 ExecStart=@sbindir@/zpool import -c @sysconfdir@/zfs/zpool.cache -aN
+ExecStartPost=/bin/bash -c "/usr/bin/systemctl set-environment BOOTFS=$(@sbindir@/zpool list -H -o bootfs)"
 
 [Install]
 WantedBy=zfs-import.target

etc/systemd/system/zfs-import-scan.service.in

@@ -13,6 +13,7 @@ Type=oneshot
 RemainAfterExit=yes
 ExecStartPre=/sbin/modprobe zfs
 ExecStart=@sbindir@/zpool import -aN -o cachefile=none
+ExecStartPost=/bin/bash -c "/usr/bin/systemctl set-environment BOOTFS=$(@sbindir@/zpool list -H -o bootfs)"
 
 [Install]
 WantedBy=zfs-import.target