Systemd tpm fixes (#343307)

nixos/modules/module-list.nix

@@ -1641,6 +1641,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

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"
       ];
 
     };

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"

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"

nixos/modules/system/boot/systemd/tpm2.nix

@@ -0,0 +1,80 @@
+{
+  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"
+          "systemd-tpm2-setup-early.service"
+          "systemd-tpm2-setup.service"
+        ];
+      }
+    )
+
+    # Stage 1
+    (
+      let
+        cfg = config.boot.initrd.systemd;
+      in
+      lib.mkIf cfg.tpm2.enable {
+        boot.initrd.systemd.additionalUpstreamUnits = [
+          "tpm2.target"
+          "systemd-tpm2-setup-early.service"
+        ];
+
+        boot.initrd.availableKernelModules =
+          [ "tpm-tis" ]
+          ++ lib.optional (
+            !(pkgs.stdenv.hostPlatform.isRiscV64 || pkgs.stdenv.hostPlatform.isArmv7)
+          ) "tpm-crb";
+        boot.initrd.systemd.storePaths = [
+          pkgs.tpm2-tss
+          "${cfg.package}/lib/systemd/systemd-tpm2-setup"
+          "${cfg.package}/lib/systemd/system-generators/systemd-tpm2-generator"
+        ];
+      }
+    )
+  ];
+}