From 3f1f20b2a7708b7f77f6252a00f8038608e1fdf1 Mon Sep 17 00:00:00 2001
From: WilliButz <willibutz@posteo.de>
Date: Mon, 30 Sep 2024 16:31:56 +0200
Subject: [PATCH 1/2] nixos/repart-image: pass partition attrs to builder
 instead of JSON file

Having access to the original Nix partition definitions in the builder
should make it a bit easier to manipulate them and still provide access
to the manipulated results.
---
 nixos/modules/image/repart-image.nix        |  8 +++++---
 nixos/modules/image/repart-verity-store.nix | 18 ++++++++----------
 nixos/modules/image/repart.nix              |  6 ++----
 3 files changed, 15 insertions(+), 17 deletions(-)

diff --git a/nixos/modules/image/repart-image.nix b/nixos/modules/image/repart-image.nix
index 41f68a0282ac..cc4c2211e3d3 100644
--- a/nixos/modules/image/repart-image.nix
+++ b/nixos/modules/image/repart-image.nix
@@ -30,7 +30,7 @@
 , imageFileBasename
 , compression
 , fileSystems
-, partitionsJSON
+, finalPartitions
 , split
 , seed
 , definitionsDirectory
@@ -110,7 +110,9 @@ in
 
   env = mkfsEnv;
 
-  inherit partitionsJSON definitionsDirectory;
+  inherit finalPartitions definitionsDirectory;
+
+  partitionsJSON = builtins.toJSON finalAttrs.finalPartitions;
 
   # relative path to the repart definitions that are read by systemd-repart
   finalRepartDefinitions = "repart.d";
@@ -136,7 +138,7 @@ in
   patchPhase = ''
     runHook prePatch
 
-    amendedRepartDefinitionsDir=$(${amendRepartDefinitions} $partitionsJSON $definitionsDirectory)
+    amendedRepartDefinitionsDir=$(${amendRepartDefinitions} <(echo "$partitionsJSON") $definitionsDirectory)
     ln -vs $amendedRepartDefinitionsDir $finalRepartDefinitions
 
     runHook postPatch
diff --git a/nixos/modules/image/repart-verity-store.nix b/nixos/modules/image/repart-verity-store.nix
index 3f341ca421f2..e576c2d649b5 100644
--- a/nixos/modules/image/repart-verity-store.nix
+++ b/nixos/modules/image/repart-verity-store.nix
@@ -163,21 +163,19 @@ in
           createEmpty = false;
         }).overrideAttrs
           (
-            finalAttrs: previousAttrs:
-            let
-              copyUki = "CopyFiles=${config.system.build.uki}/${config.system.boot.loader.ukiFile}:${cfg.ukiPath}";
-            in
-            {
+            finalAttrs: previousAttrs: {
+              # add entry to inject UKI into ESP
+              finalPartitions = lib.recursiveUpdate previousAttrs.finalPartitions {
+                ${cfg.partitionIds.esp}.contents = {
+                  "${cfg.ukiPath}".source = "${config.system.build.uki}/${config.system.boot.loader.ukiFile}";
+                };
+              };
+
               nativeBuildInputs = previousAttrs.nativeBuildInputs ++ [
                 pkgs.systemdUkify
                 verityHashCheck
               ];
 
-              postPatch = ''
-                # add entry to inject UKI into ESP
-                echo '${copyUki}' >> $finalRepartDefinitions/${cfg.partitionIds.esp}.conf
-              '';
-
               preBuild = ''
                 # check that we build the final image with the same intermediate image for
                 # which the injected UKI was built by comparing the UKI cmdline with the repart output
diff --git a/nixos/modules/image/repart.nix b/nixos/modules/image/repart.nix
index b2f03d96ad35..544779cd6e9b 100644
--- a/nixos/modules/image/repart.nix
+++ b/nixos/modules/image/repart.nix
@@ -318,14 +318,12 @@ in
           format
           (lib.mapAttrs (_n: v: { Partition = v.repartConfig; }) cfg.finalPartitions);
 
-        partitionsJSON = pkgs.writeText "partitions.json" (builtins.toJSON cfg.finalPartitions);
-
         mkfsEnv = mkfsOptionsToEnv cfg.mkfsOptions;
       in
       pkgs.callPackage ./repart-image.nix {
         systemd = cfg.package;
-        inherit (cfg) name version imageFileBasename compression split seed sectorSize;
-        inherit fileSystems definitionsDirectory partitionsJSON mkfsEnv;
+        inherit (cfg) name version imageFileBasename compression split seed sectorSize finalPartitions;
+        inherit fileSystems definitionsDirectory mkfsEnv;
       };
 
     meta.maintainers = with lib.maintainers; [ nikstur willibutz ];

From 93c61c1e58dd74f0f864eb98820149a5cfe9dc3f Mon Sep 17 00:00:00 2001
From: WilliButz <willibutz@posteo.de>
Date: Mon, 30 Sep 2024 22:45:05 +0200
Subject: [PATCH 2/2] nixos/repart-verity-store: include original roothashes in
 repart-output.json

---
 nixos/modules/image/repart-verity-store.nix | 19 +++++++++++++++++++
 1 file changed, 19 insertions(+)

diff --git a/nixos/modules/image/repart-verity-store.nix b/nixos/modules/image/repart-verity-store.nix
index e576c2d649b5..e8706a5f79d0 100644
--- a/nixos/modules/image/repart-verity-store.nix
+++ b/nixos/modules/image/repart-verity-store.nix
@@ -174,6 +174,7 @@ in
               nativeBuildInputs = previousAttrs.nativeBuildInputs ++ [
                 pkgs.systemdUkify
                 verityHashCheck
+                pkgs.jq
               ];
 
               preBuild = ''
@@ -192,6 +193,24 @@ in
                 chmod +w ${config.image.repart.imageFileBasename}.raw
               '';
 
+              # replace "TBD" with the original roothash values
+              preInstall = ''
+                mv -v repart-output{.json,_orig.json}
+
+                jq --slurp --indent -1 \
+                  '.[0] as $intermediate | .[1] as $final
+                    | $intermediate | map(select(.roothash != null) | { "uuid":.uuid,"roothash":.roothash }) as $uuids
+                    | $final + $uuids
+                    | group_by(.uuid)
+                    | map(add)
+                    | sort_by(.offset)' \
+                      ${config.system.build.intermediateImage}/repart-output.json \
+                      repart-output_orig.json \
+                  > repart-output.json
+
+                rm -v repart-output_orig.json
+              '';
+
               # the image will be self-contained so we can drop references
               # to the closure that was used to build it
               unsafeDiscardReferences.out = true;