From 5034450095548401f47f1a55c44cb2ecdae6b77f Mon Sep 17 00:00:00 2001
From: Will Fancher <elvishjerricco@gmail.com>
Date: Thu, 19 Sep 2024 04:10:56 -0400
Subject: [PATCH] nixos/systemd: Factor out tpm2 support into separate module

---
 nixos/modules/module-list.nix                |  1 +
 nixos/modules/system/boot/luksroot.nix       |  2 +
 nixos/modules/system/boot/systemd.nix        |  2 -
 nixos/modules/system/boot/systemd/initrd.nix | 18 +----
 nixos/modules/system/boot/systemd/tpm2.nix   | 75 ++++++++++++++++++++
 5 files changed, 79 insertions(+), 19 deletions(-)
 create mode 100644 nixos/modules/system/boot/systemd/tpm2.nix

diff --git a/nixos/modules/module-list.nix b/nixos/modules/module-list.nix
index 006cdeedcaf6..40c4ed5d4a66 100644
--- a/nixos/modules/module-list.nix
+++ b/nixos/modules/module-list.nix
@@ -1636,6 +1636,7 @@
   ./system/boot/systemd/sysupdate.nix
   ./system/boot/systemd/sysusers.nix
   ./system/boot/systemd/tmpfiles.nix
+  ./system/boot/systemd/tpm2.nix
   ./system/boot/systemd/user.nix
   ./system/boot/systemd/userdbd.nix
   ./system/boot/systemd/homed.nix
diff --git a/nixos/modules/system/boot/luksroot.nix b/nixos/modules/system/boot/luksroot.nix
index 65190e65d9b9..70b455871b4b 100644
--- a/nixos/modules/system/boot/luksroot.nix
+++ b/nixos/modules/system/boot/luksroot.nix
@@ -1088,6 +1088,8 @@ in
       storePaths = [
         "${config.boot.initrd.systemd.package}/bin/systemd-cryptsetup"
         "${config.boot.initrd.systemd.package}/lib/systemd/system-generators/systemd-cryptsetup-generator"
+      ] ++ lib.optionals config.boot.initrd.systemd.tpm2.enable [
+        "${config.boot.initrd.systemd.package}/lib/cryptsetup/libcryptsetup-token-systemd-tpm2.so"
       ];
 
     };
diff --git a/nixos/modules/system/boot/systemd.nix b/nixos/modules/system/boot/systemd.nix
index bb6fcc1d38ce..797db1048648 100644
--- a/nixos/modules/system/boot/systemd.nix
+++ b/nixos/modules/system/boot/systemd.nix
@@ -37,8 +37,6 @@ let
       "cryptsetup.target"
       "cryptsetup-pre.target"
       "remote-cryptsetup.target"
-    ] ++ optionals cfg.package.withTpm2Tss [
-      "tpm2.target"
     ] ++ [
       "sigpwr.target"
       "timers.target"
diff --git a/nixos/modules/system/boot/systemd/initrd.nix b/nixos/modules/system/boot/systemd/initrd.nix
index 2ccc964820fe..d0295e78fcaf 100644
--- a/nixos/modules/system/boot/systemd/initrd.nix
+++ b/nixos/modules/system/boot/systemd/initrd.nix
@@ -68,7 +68,6 @@ let
     "systemd-reboot.service"
     "systemd-sysctl.service"
     "timers.target"
-    "tpm2.target"
     "umount.target"
     "systemd-bsod.service"
   ] ++ cfg.additionalUpstreamUnits;
@@ -349,15 +348,6 @@ in {
       visible = "shallow";
       description = "Definition of slice configurations.";
     };
-
-    enableTpm2 = mkOption {
-      default = cfg.package.withTpm2Tss;
-      defaultText = "boot.initrd.systemd.package.withTpm2Tss";
-      type = types.bool;
-      description = ''
-        Whether to enable TPM2 support in the initrd.
-      '';
-    };
   };
 
   config = mkIf (config.boot.initrd.enable && cfg.enable) {
@@ -394,9 +384,7 @@ in {
       # systemd needs this for some features
       "autofs"
       # systemd-cryptenroll
-    ] ++ lib.optional cfg.enableTpm2 "tpm-tis"
-    ++ lib.optional (cfg.enableTpm2 && !(pkgs.stdenv.hostPlatform.isRiscV64 || pkgs.stdenv.hostPlatform.isArmv7)) "tpm-crb"
-    ++ lib.optional cfg.package.withEfi "efivarfs";
+    ] ++ lib.optional cfg.package.withEfi "efivarfs";
 
     boot.kernelParams = [
       "root=${config.boot.initrd.systemd.root}"
@@ -495,10 +483,6 @@ in {
 
         # so NSS can look up usernames
         "${pkgs.glibc}/lib/libnss_files.so.2"
-      ] ++ optionals (cfg.package.withCryptsetup && cfg.enableTpm2) [
-        # tpm2 support
-        "${cfg.package}/lib/cryptsetup/libcryptsetup-token-systemd-tpm2.so"
-        pkgs.tpm2-tss
       ] ++ optionals cfg.package.withCryptsetup [
         # fido2 support
         "${cfg.package}/lib/cryptsetup/libcryptsetup-token-systemd-fido2.so"
diff --git a/nixos/modules/system/boot/systemd/tpm2.nix b/nixos/modules/system/boot/systemd/tpm2.nix
new file mode 100644
index 000000000000..31c63f75816e
--- /dev/null
+++ b/nixos/modules/system/boot/systemd/tpm2.nix
@@ -0,0 +1,75 @@
+{
+  lib,
+  config,
+  pkgs,
+  ...
+}:
+{
+  meta.maintainers = [ lib.maintainers.elvishjerricco ];
+
+  imports = [
+    (lib.mkRenamedOptionModule
+      [
+        "boot"
+        "initrd"
+        "systemd"
+        "enableTpm2"
+      ]
+      [
+        "boot"
+        "initrd"
+        "systemd"
+        "tpm2"
+        "enable"
+      ]
+    )
+  ];
+
+  options = {
+    systemd.tpm2.enable = lib.mkEnableOption "systemd TPM2 support" // {
+      default = config.systemd.package.withTpm2Tss;
+      defaultText = "systemd.package.withTpm2Tss";
+    };
+
+    boot.initrd.systemd.tpm2.enable = lib.mkEnableOption "systemd initrd TPM2 support" // {
+      default = config.boot.initrd.systemd.package.withTpm2Tss;
+      defaultText = "boot.initrd.systemd.package.withTpm2Tss";
+    };
+  };
+
+  # TODO: pcrphase, pcrextend, pcrfs, pcrmachine
+  config = lib.mkMerge [
+    # Stage 2
+    (
+      let
+        cfg = config.systemd;
+      in
+      lib.mkIf cfg.tpm2.enable {
+        systemd.additionalUpstreamSystemUnits = [
+          "tpm2.target"
+        ];
+      }
+    )
+
+    # Stage 1
+    (
+      let
+        cfg = config.boot.initrd.systemd;
+      in
+      lib.mkIf cfg.tpm2.enable {
+        boot.initrd.systemd.additionalUpstreamUnits = [
+          "tpm2.target"
+        ];
+
+        boot.initrd.availableKernelModules =
+          [ "tpm-tis" ]
+          ++ lib.optional (
+            !(pkgs.stdenv.hostPlatform.isRiscV64 || pkgs.stdenv.hostPlatform.isArmv7)
+          ) "tpm-crb";
+        boot.initrd.systemd.storePaths = [
+          pkgs.tpm2-tss
+        ];
+      }
+    )
+  ];
+}