From 0425c4de6d1ebade3c082698a9d47b2f1b0080db Mon Sep 17 00:00:00 2001
From: Will Fancher <elvishjerricco@gmail.com>
Date: Sat, 21 Oct 2023 11:56:47 -0400
Subject: [PATCH] systemd-stage-1: bcachefs decryption

---
 nixos/modules/tasks/filesystems/bcachefs.nix | 50 ++++++++++++++++----
 nixos/tests/installer-systemd-stage-1.nix    |  2 +
 nixos/tests/installer.nix                    |  4 ++
 3 files changed, 46 insertions(+), 10 deletions(-)

diff --git a/nixos/modules/tasks/filesystems/bcachefs.nix b/nixos/modules/tasks/filesystems/bcachefs.nix
index 19ef188ce783..b90ff894624e 100644
--- a/nixos/modules/tasks/filesystems/bcachefs.nix
+++ b/nixos/modules/tasks/filesystems/bcachefs.nix
@@ -34,17 +34,43 @@ let
     }
   '';
 
-  openCommand = name: fs:
-    let
-      # we need only unlock one device manually, and cannot pass multiple at once
-      # remove this adaptation when bcachefs implements mounting by filesystem uuid
-      # also, implement automatic waiting for the constituent devices when that happens
-      # bcachefs does not support mounting devices with colons in the path, ergo we don't (see #49671)
-      firstDevice = head (splitString ":" fs.device);
-    in
-      ''
-        tryUnlock ${name} ${firstDevice}
+  # we need only unlock one device manually, and cannot pass multiple at once
+  # remove this adaptation when bcachefs implements mounting by filesystem uuid
+  # also, implement automatic waiting for the constituent devices when that happens
+  # bcachefs does not support mounting devices with colons in the path, ergo we don't (see #49671)
+  firstDevice = fs: head (splitString ":" fs.device);
+
+  openCommand = name: fs: ''
+    tryUnlock ${name} ${firstDevice fs}
+  '';
+
+  mkUnits = prefix: name: fs: let
+    mountUnit = "${utils.escapeSystemdPath (prefix + (lib.removeSuffix "/" fs.mountPoint))}.mount";
+    device = firstDevice fs;
+    deviceUnit = "${utils.escapeSystemdPath device}.device";
+  in {
+    name = "unlock-bcachefs-${utils.escapeSystemdPath fs.mountPoint}";
+    value = {
+      description = "Unlock bcachefs for ${fs.mountPoint}";
+      requiredBy = [ mountUnit ];
+      before = [ mountUnit ];
+      bindsTo = [ deviceUnit ];
+      after = [ deviceUnit ];
+      unitConfig.DefaultDependencies = false;
+      serviceConfig = {
+        Type = "oneshot";
+        ExecCondition = "${pkgs.bcachefs-tools}/bin/bcachefs unlock -c \"${device}\"";
+        Restart = "on-failure";
+        RestartMode = "direct";
+        # Ideally, this service would lock the key on stop.
+        # As is, RemainAfterExit doesn't accomplish anything.
+        RemainAfterExit = true;
+      };
+      script = ''
+        ${config.boot.initrd.systemd.package}/bin/systemd-ask-password --timeout=0 "enter passphrase for ${name}" | exec ${pkgs.bcachefs-tools}/bin/bcachefs unlock "${device}"
       '';
+    };
+  };
 
 in
 
@@ -59,6 +85,8 @@ in
 
       # use kernel package with bcachefs support until it's in mainline
       boot.kernelPackages = pkgs.linuxPackages_testing_bcachefs;
+
+      systemd.services = lib.mapAttrs' (mkUnits "") (lib.filterAttrs (n: fs: (fs.fsType == "bcachefs") && (!utils.fsNeededForBoot fs)) config.fileSystems);
     }
 
     (mkIf ((elem "bcachefs" config.boot.initrd.supportedFilesystems) || (bootFs != {})) {
@@ -79,6 +107,8 @@ in
       '';
 
       boot.initrd.postDeviceCommands = commonFunctions + concatStrings (mapAttrsToList openCommand bootFs);
+
+      boot.initrd.systemd.services = lib.mapAttrs' (mkUnits "/sysroot") bootFs;
     })
   ]);
 }
diff --git a/nixos/tests/installer-systemd-stage-1.nix b/nixos/tests/installer-systemd-stage-1.nix
index 85155a6c682b..8f10a461d1c6 100644
--- a/nixos/tests/installer-systemd-stage-1.nix
+++ b/nixos/tests/installer-systemd-stage-1.nix
@@ -8,6 +8,8 @@
   # them when fixed.
   inherit (import ./installer.nix { inherit system config pkgs; systemdStage1 = true; })
     # bcache
+    bcachefsSimple
+    bcachefsEncrypted
     btrfsSimple
     btrfsSubvolDefault
     btrfsSubvolEscape
diff --git a/nixos/tests/installer.nix b/nixos/tests/installer.nix
index 3268a16967d7..c21f56089404 100644
--- a/nixos/tests/installer.nix
+++ b/nixos/tests/installer.nix
@@ -918,6 +918,10 @@ in {
     enableOCR = true;
     preBootCommands = ''
       machine.start()
+      # Enter it wrong once
+      machine.wait_for_text("enter passphrase for ")
+      machine.send_chars("wrong\n")
+      # Then enter it right.
       machine.wait_for_text("enter passphrase for ")
       machine.send_chars("password\n")
     '';