From 4df3c4c17b3622c2f3f2dc8bb877ac8470efc1c6 Mon Sep 17 00:00:00 2001
From: misuzu <bakalolka@gmail.com>
Date: Wed, 12 Jun 2024 20:06:38 +0300
Subject: [PATCH] nixos/clevis: add support for parent encrypted zfs datasets

---
 nixos/modules/system/boot/clevis.nix      |  2 +-
 nixos/modules/tasks/filesystems/zfs.nix   |  3 +--
 nixos/tests/installer-systemd-stage-1.nix |  2 ++
 nixos/tests/installer.nix                 | 16 ++++++++++++++--
 4 files changed, 18 insertions(+), 5 deletions(-)

diff --git a/nixos/modules/system/boot/clevis.nix b/nixos/modules/system/boot/clevis.nix
index d9390f5bc15f..ac881e953576 100644
--- a/nixos/modules/system/boot/clevis.nix
+++ b/nixos/modules/system/boot/clevis.nix
@@ -48,7 +48,7 @@ in
 
     assertions = (attrValues (mapAttrs
       (device: _: {
-        assertion = (any (fs: fs.device == device && (elem fs.fsType supportedFs)) config.system.build.fileSystems) || (hasAttr device config.boot.initrd.luks.devices);
+        assertion = (any (fs: fs.device == device && (elem fs.fsType supportedFs) || (fs.fsType == "zfs" && hasPrefix "${device}/" fs.device)) config.system.build.fileSystems) || (hasAttr device config.boot.initrd.luks.devices);
         message = ''
           No filesystem or LUKS device with the name ${device} is declared in your configuration.'';
       })
diff --git a/nixos/modules/tasks/filesystems/zfs.nix b/nixos/modules/tasks/filesystems/zfs.nix
index b75817a011cb..7861a4b946b4 100644
--- a/nixos/modules/tasks/filesystems/zfs.nix
+++ b/nixos/modules/tasks/filesystems/zfs.nix
@@ -17,8 +17,7 @@ let
   cfgZED = config.services.zfs.zed;
 
   selectModulePackage = package: config.boot.kernelPackages.${package.kernelModuleAttribute};
-  clevisDatasets = map (e: e.device) (filter (e: e.device != null && (hasAttr e.device config.boot.initrd.clevis.devices) && e.fsType == "zfs" && (fsNeededForBoot e)) config.system.build.fileSystems);
-
+  clevisDatasets = attrNames (filterAttrs (device: _: any (e: e.fsType == "zfs" && (fsNeededForBoot e) && (e.device == device || hasPrefix "${device}/" e.device)) config.system.build.fileSystems) config.boot.initrd.clevis.devices);
 
   inInitrd = config.boot.initrd.supportedFilesystems.zfs or false;
   inSystem = config.boot.supportedFilesystems.zfs or false;
diff --git a/nixos/tests/installer-systemd-stage-1.nix b/nixos/tests/installer-systemd-stage-1.nix
index 00205f941771..3b5e0ed8e7bb 100644
--- a/nixos/tests/installer-systemd-stage-1.nix
+++ b/nixos/tests/installer-systemd-stage-1.nix
@@ -37,6 +37,8 @@
     clevisLuksFallback
     clevisZfs
     clevisZfsFallback
+    clevisZfsParentDataset
+    clevisZfsParentDatasetFallback
     gptAutoRoot
     clevisBcachefs
     clevisBcachefsFallback
diff --git a/nixos/tests/installer.nix b/nixos/tests/installer.nix
index 3f57a64333dd..bb6ad79615fa 100644
--- a/nixos/tests/installer.nix
+++ b/nixos/tests/installer.nix
@@ -714,7 +714,7 @@ let
     '';
   };
 
-  mkClevisZfsTest = { fallback ? false }: makeInstallerTest "clevis-zfs${optionalString fallback "-fallback"}" {
+  mkClevisZfsTest = { fallback ? false, parentDataset ? false }: makeInstallerTest "clevis-zfs${optionalString parentDataset "-parent-dataset"}${optionalString fallback "-fallback"}" {
     clevisTest = true;
     clevisFallbackTest = fallback;
     enableOCR = fallback;
@@ -731,17 +731,27 @@ let
         "udevadm settle",
         "mkswap /dev/vda2 -L swap",
         "swapon -L swap",
+    '' + optionalString (!parentDataset) ''
         "zpool create -O mountpoint=legacy rpool /dev/vda3",
         "echo -n password | zfs create"
         + " -o encryption=aes-256-gcm -o keyformat=passphrase rpool/root",
+    '' + optionalString (parentDataset) ''
+        "echo -n password | zpool create -O mountpoint=none -O encryption=on -O keyformat=passphrase rpool /dev/vda3",
+        "zfs create -o mountpoint=legacy rpool/root",
+    '' +
+    ''
         "mount -t zfs rpool/root /mnt",
         "mkfs.ext3 -L boot /dev/vda1",
         "mkdir -p /mnt/boot",
         "mount LABEL=boot /mnt/boot",
         "udevadm settle")
     '';
-    extraConfig = ''
+    extraConfig = optionalString (!parentDataset) ''
       boot.initrd.clevis.devices."rpool/root".secretFile = "/etc/nixos/clevis-secret.jwe";
+    '' + optionalString (parentDataset) ''
+      boot.initrd.clevis.devices."rpool".secretFile = "/etc/nixos/clevis-secret.jwe";
+    '' +
+    ''
       boot.zfs.requestEncryptionCredentials = true;
 
 
@@ -1359,6 +1369,8 @@ in {
   clevisLuksFallback = mkClevisLuksTest { fallback = true; };
   clevisZfs = mkClevisZfsTest { };
   clevisZfsFallback = mkClevisZfsTest { fallback = true; };
+  clevisZfsParentDataset = mkClevisZfsTest { parentDataset = true; };
+  clevisZfsParentDatasetFallback = mkClevisZfsTest { parentDataset = true; fallback = true; };
 } // optionalAttrs systemdStage1 {
   stratisRoot = makeInstallerTest "stratisRoot" {
     createPartitions = ''