From 8ceff1b73139d650578600649d70f43c6d90e2ff Mon Sep 17 00:00:00 2001
From: matthewcroughan <matt@croughan.sh>
Date: Tue, 12 Dec 2023 17:53:56 +0000
Subject: [PATCH] make-disk-image: use nix path-info to automatically set the
 imageSize

---
 lib/make-disk-image.nix | 14 +++++++++++++-
 lib/types/disk.nix      | 10 +++++-----
 2 files changed, 18 insertions(+), 6 deletions(-)

diff --git a/lib/make-disk-image.nix b/lib/make-disk-image.nix
index 01d70b4..43ce08b 100644
--- a/lib/make-disk-image.nix
+++ b/lib/make-disk-image.nix
@@ -30,7 +30,19 @@ let
     util-linux
   ] ++ nixosConfig.config.disko.extraDependencies;
   preVM = ''
-    ${lib.concatMapStringsSep "\n" (disk: "truncate -s ${disk.imageSize} ${disk.name}.raw") (lib.attrValues nixosConfig.config.disko.devices.disk)}
+    export PATH=$PATH:${lib.makeBinPath (with pkgs; [ nix jq ])}
+    export NIX_STATE_DIR=$TMPDIR/state
+    export NIX_CONFIG="experimental-features = nix-command"
+    nix-store --load-db < ${pkgs.closureInfo {
+      rootPaths = [ systemToInstall.config.system.build.toplevel ];
+    }}/registration
+
+    ${lib.concatMapStringsSep "\n" (disk: ''
+      closureSize=$(nix path-info --closure-size --json ${systemToInstall.config.system.build.toplevel} | jq .[].closureSize)
+      size=$((closureSize + ${toString disk.imageSizeExtraBytes}))
+      echo "Truncating ${disk.name}.raw to $size bytes..."
+      truncate -s "$size" ${disk.name}.raw
+    '') (lib.attrValues nixosConfig.config.disko.devices.disk)}
   '';
   postVM = ''
     # shellcheck disable=SC2154
diff --git a/lib/types/disk.nix b/lib/types/disk.nix
index 60f31b5..6dc76e1 100644
--- a/lib/types/disk.nix
+++ b/lib/types/disk.nix
@@ -16,13 +16,13 @@
       type = diskoLib.optionTypes.absolute-pathname; # TODO check if subpath of /dev ? - No! eg: /.swapfile
       description = "Device path";
     };
-    imageSize = lib.mkOption {
-      type = lib.types.strMatching "[0-9]+[KMGTP]?";
+    imageSizeExtraBytes = lib.mkOption {
+      type = lib.types.int;
       description = ''
-        size of the image if the makeDiskImages function from diksoLib is used.
-        is used as an argument to truncate -s
+        extra size in bytes to add to the end of the disk image if the makeDiskImages function from diskoLib is used.
+        This is used as an argument to truncate -s. Defaults to 512MiB.
       '';
-      default = "2G";
+      default = 536870912;
     };
     content = diskoLib.deviceType { parent = config; device = config.device; };
     _meta = lib.mkOption {