image/repart: add version and compression options

The version option is needed if you want to implement partition &
systemd-boot based A/B booting where the version information is encoded
in the files on the ESP. See systemd-sysupate docs for more details on
this:
https://www.freedesktop.org/software/systemd/man/latest/sysupdate.d.html

Note, however, that this is not *only* useful for systemd-sysupdate but
also for other similar updating tools/mechanisms.

2 files changed, 107 insertions, 5 deletions

diff --git a/nixos/modules/image/repart-image.nix b/nixos/modules/image/repart-image.nix
index b4a1dfe51ff3a..a12b4fb14fb16 100644
--- a/nixos/modules/image/repart-image.nix
+++ b/nixos/modules/image/repart-image.nix
@@ -10,6 +10,8 @@
 , systemd
 , fakeroot
 , util-linux
+
+  # filesystem tools
 , dosfstools
 , mtools
 , e2fsprogs
@@ -18,8 +20,13 @@
 , btrfs-progs
 , xfsprogs
 
+  # compression tools
+, zstd
+, xz
+
   # arguments
-, name
+, imageFileBasename
+, compression
 , fileSystems
 , partitions
 , split
@@ -52,14 +59,25 @@ let
   };
 
   fileSystemTools = builtins.concatMap (f: fileSystemToolMapping."${f}") fileSystems;
+
+  compressionPkg = {
+    "zstd" = zstd;
+    "xz" = xz;
+  }."${compression.algorithm}";
+
+  compressionCommand = {
+    "zstd" = "zstd --no-progress --threads=0 -${toString compression.level}";
+    "xz" = "xz --keep --verbose --threads=0 -${toString compression.level}";
+  }."${compression.algorithm}";
 in
 
-runCommand name
+runCommand imageFileBasename
 {
   nativeBuildInputs = [
     systemd
     fakeroot
     util-linux
+    compressionPkg
   ] ++ fileSystemTools;
 } ''
   amendedRepartDefinitions=$(${amendRepartDefinitions} ${partitions} ${definitionsDirectory})
@@ -67,6 +85,7 @@ runCommand name
   mkdir -p $out
   cd $out
 
+  echo "Building image with systemd-repart..."
   unshare --map-root-user fakeroot systemd-repart \
     --dry-run=no \
     --empty=create \
@@ -75,6 +94,17 @@ runCommand name
     --definitions="$amendedRepartDefinitions" \
     --split="${lib.boolToString split}" \
     --json=pretty \
-    image.raw \
+    ${imageFileBasename}.raw \
     | tee repart-output.json
+
+  # Compression is implemented in the same derivation as opposed to in a
+  # separate derivation to allow users to save disk space. Disk images are
+  # already very space intensive so we want to allow users to mitigate this.
+  if ${lib.boolToString compression.enable}; then
+    for f in ${imageFileBasename}*; do
+      echo "Compressing $f with ${compression.algorithm}..."
+      # Keep the original file when compressing and only delete it afterwards
+      ${compressionCommand} $f && rm $f
+    done
+  fi
 ''
diff --git a/nixos/modules/image/repart.nix b/nixos/modules/image/repart.nix
index da4f45d9a6392..ed584d9bf997b 100644
--- a/nixos/modules/image/repart.nix
+++ b/nixos/modules/image/repart.nix
@@ -66,7 +66,53 @@ in
 
     name = lib.mkOption {
       type = lib.types.str;
-      description = lib.mdDoc "The name of the image.";
+      description = lib.mdDoc ''
+        Name of the image.
+
+        If this option is unset but config.system.image.id is set,
+        config.system.image.id is used as the default value.
+      '';
+    };
+
+    version = lib.mkOption {
+      type = lib.types.nullOr lib.types.str;
+      default = config.system.image.version;
+      defaultText = lib.literalExpression "config.system.image.version";
+      description = lib.mdDoc "Version of the image";
+    };
+
+    imageFileBasename = lib.mkOption {
+      type = lib.types.str;
+      readOnly = true;
+      description = lib.mdDoc ''
+        Basename of the image filename without any extension (e.g. `image_1`).
+      '';
+    };
+
+    imageFile = lib.mkOption {
+      type = lib.types.str;
+      readOnly = true;
+      description = lib.mdDoc ''
+        Filename of the image including all extensions (e.g `image_1.raw` or
+        `image_1.raw.zst`).
+      '';
+    };
+
+    compression = {
+      enable = lib.mkEnableOption (lib.mdDoc "Image compression");
+
+      algorithm = lib.mkOption {
+        type = lib.types.enum [ "zstd" "xz" ];
+        default = "zstd";
+        description = lib.mdDoc "Compression algorithm";
+      };
+
+      level = lib.mkOption {
+        type = lib.types.int;
+        description = lib.mdDoc ''
+          Compression level. The available range depends on the used algorithm.
+        '';
+      };
     };
 
     seed = lib.mkOption {
@@ -131,6 +177,32 @@ in
 
   config = {
 
+    image.repart =
+      let
+        version = config.image.repart.version;
+        versionInfix = if version != null then "_${version}" else "";
+        compressionSuffix = lib.optionalString cfg.compression.enable
+          {
+            "zstd" = ".zst";
+            "xz" = ".xz";
+          }."${cfg.compression.algorithm}";
+      in
+      {
+        name = lib.mkIf (config.system.image.id != null) (lib.mkOptionDefault config.system.image.id);
+        imageFileBasename = cfg.name + versionInfix;
+        imageFile = cfg.imageFileBasename + ".raw" + compressionSuffix;
+
+        compression = {
+          # Generally default to slightly faster than default compression
+          # levels under the assumption that most of the building will be done
+          # for development and release builds will be customized.
+          level = lib.mkOptionDefault {
+            "zstd" = 3;
+            "xz" = 3;
+          }."${cfg.compression.algorithm}";
+        };
+      };
+
     system.build.image =
       let
         fileSystems = lib.filter
@@ -160,7 +232,7 @@ in
       in
       pkgs.callPackage ./repart-image.nix {
         systemd = cfg.package;
-        inherit (cfg) name split seed;
+        inherit (cfg) imageFileBasename compression split seed;
         inherit fileSystems definitionsDirectory partitions;
       };