make-disk-image: convert into NixOS module

As `makeDiskImages` always requires a NixOS configuration, we can
simplify the code by convering it into a NixOS module. Then we can make
it responsible for populating `system.build.diskoImages` and
`system.build.diskoImagesScript`.

lib/default.nix

@@ -9,13 +9,6 @@ with builtins;
 let
   outputs = import ../default.nix { inherit lib diskoLib; };
   diskoLib = {
-
-    # like make-disk-image.nix from nixpkgs, but with disko config
-    makeDiskImages = args: (import ./make-disk-image.nix ({ inherit diskoLib; } // args)).pure;
-
-    # a version of makeDiskImage which runs outside of the store
-    makeDiskImagesScript = args: (import ./make-disk-image.nix ({ inherit diskoLib; } // args)).impure;
-
     testLib = import ./tests.nix { inherit lib makeTest eval-config; };
     # like lib.types.oneOf but instead of a list takes an attrset
     # uses the field "type" to find the correct type in the attrset

lib/interactive-vm.nix

@@ -1,6 +1,4 @@
-# We need to specify extendModules here to ensure that it is available
-# in args for makeDiskImages
-{ diskoLib, modulesPath, config, pkgs, lib, extendModules, ... }@args:
+{ diskoLib, modulesPath, config, pkgs, lib, ... }:
 
 let
   vm_disko = (diskoLib.testLib.prepareDiskoConfig config diskoLib.testLib.devices).disko;
@@ -20,15 +18,6 @@ let
     };
   }).config;
   disks = lib.attrValues cfg_.disko.devices.disk;
-  diskoImages = diskoLib.makeDiskImages {
-    nixosConfig = args;
-    copyNixStore = false;
-    extraConfig = {
-      disko.devices = cfg_.disko.devices;
-    };
-    testMode = true;
-    imageFormat = "qcow2";
-  };
   rootDisk = {
     name = "root";
     file = ''"$tmp"/${(builtins.head disks).name}.qcow2'';
@@ -58,6 +47,14 @@ in
     diskoBasedConfiguration
   ];
 
+  disko.testMode = true;
+
+  disko.imageBuilder.copyNixStore = false;
+  disko.imageBuilder.extraConfig = {
+    disko.devices = cfg_.disko.devices;
+  };
+  disko.imageBuilder.imageFormat = "qcow2";
+
   virtualisation.useEFIBoot = config.disko.tests.efi;
   virtualisation.memorySize = config.disko.memSize;
   virtualisation.useDefaultFilesystems = false;
@@ -72,7 +69,7 @@ in
     trap 'rm -rf "$tmp"' EXIT
     ${lib.concatMapStringsSep "\n" (disk: ''
       ${pkgs.qemu}/bin/qemu-img create -f qcow2 \
-      -b ${diskoImages}/${disk.name}.qcow2 \
+      -b ${config.system.build.diskoImages}/${disk.name}.qcow2 \
       -F qcow2 "$tmp"/${disk.name}.qcow2
     '') disks}
     set +f

lib/make-disk-image.nix

@@ -1,30 +1,29 @@
-{ nixosConfig
+{ config
 , diskoLib
-, pkgs ? nixosConfig.config.disko.imageBuilderPkgs
-, lib ? pkgs.lib
-, name ? "${nixosConfig.config.networking.hostName}-disko-images"
-, extraPostVM ? nixosConfig.config.disko.extraPostVM
-, checked ? false
-, copyNixStore ? true
-, testMode ? false
-, extraConfig ? { }
-, imageFormat ? "raw"
+, lib
+, extendModules
+, ...
 }:
 let
-  configSupportsZfs = nixosConfig.config.boot.supportedFilesystems.zfs or false;
+  diskoCfg = config.disko;
+  cfg = diskoCfg.imageBuilder;
+  inherit (cfg) pkgs imageFormat;
+  checked = diskoCfg.checkScripts;
+
+  configSupportsZfs = config.boot.supportedFilesystems.zfs or false;
   vmTools = pkgs.vmTools.override {
     rootModules = [ "9p" "9pnet_virtio" "virtio_pci" "virtio_blk" ]
       ++ (lib.optional configSupportsZfs "zfs")
-      ++ nixosConfig.config.disko.extraRootModules;
-    customQemu = nixosConfig.config.disko.imageBuilderQemu;
+      ++ cfg.extraRootModules;
+    customQemu = cfg.qemu;
     kernel = pkgs.aggregateModules
-      (with nixosConfig.config.disko.imageBuilderKernelPackages; [ kernel ]
-        ++ lib.optional (lib.elem "zfs" nixosConfig.config.disko.extraRootModules || configSupportsZfs) zfs);
+      (with cfg.kernelPackages; [ kernel ]
+        ++ lib.optional (lib.elem "zfs" cfg.extraRootModules || configSupportsZfs) zfs);
   };
-  cleanedConfig = diskoLib.testLib.prepareDiskoConfig nixosConfig.config diskoLib.testLib.devices;
-  systemToInstall = nixosConfig.extendModules {
+  cleanedConfig = diskoLib.testLib.prepareDiskoConfig config diskoLib.testLib.devices;
+  systemToInstall = extendModules {
     modules = [
-      extraConfig
+      cfg.extraConfig
       {
         disko.testMode = true;
         disko.devices = lib.mkForce cleanedConfig.disko.devices;
@@ -41,9 +40,9 @@ let
     nix
     util-linux
     findutils
-  ] ++ nixosConfig.config.disko.extraDependencies;
+  ] ++ cfg.extraDependencies;
   preVM = ''
-    ${lib.concatMapStringsSep "\n" (disk: "${pkgs.qemu}/bin/qemu-img create -f ${imageFormat} ${disk.name}.${imageFormat} ${disk.imageSize}") (lib.attrValues nixosConfig.config.disko.devices.disk)}
+    ${lib.concatMapStringsSep "\n" (disk: "${pkgs.qemu}/bin/qemu-img create -f ${imageFormat} ${disk.name}.${imageFormat} ${disk.imageSize}") (lib.attrValues diskoCfg.devices.disk)}
     # This makes disko work, when canTouchEfiVariables is set to true.
     # Technically these boot entries will no be persisted this way, but
     # in most cases this is OK, because we can rely on the standard location for UEFI executables.
@@ -52,8 +51,8 @@ let
   postVM = ''
     # shellcheck disable=SC2154
     mkdir -p "$out"
-    ${lib.concatMapStringsSep "\n" (disk: "mv ${disk.name}.${imageFormat} \"$out\"/${disk.name}.${imageFormat}") (lib.attrValues nixosConfig.config.disko.devices.disk)}
-    ${extraPostVM}
+    ${lib.concatMapStringsSep "\n" (disk: "mv ${disk.name}.${imageFormat} \"$out\"/${disk.name}.${imageFormat}") (lib.attrValues diskoCfg.devices.disk)}
+    ${cfg.extraPostVM}
   '';
 
   closureInfo = pkgs.closureInfo {
@@ -76,13 +75,13 @@ let
     udevadm trigger --action=add
     udevadm settle
 
-    ${lib.optionalString testMode ''
+    ${lib.optionalString diskoCfg.testMode ''
       export IN_DISKO_TEST=1
     ''}
     ${systemToInstall.config.system.build.diskoScript}
   '';
 
-  installer = lib.optionalString copyNixStore ''
+  installer = lib.optionalString cfg.copyNixStore ''
     # populate nix db, so nixos-install doesn't complain
     export NIX_STATE_DIR=${systemToInstall.config.disko.rootMountPoint}/nix/var/nix
     nix-store --load-db < "${closureInfo}/registration"
@@ -102,17 +101,18 @@ let
     (disk:
       "-drive file=${disk.name}.${imageFormat},if=virtio,cache=unsafe,werror=report,format=${imageFormat}"
     )
-    (lib.attrValues nixosConfig.config.disko.devices.disk));
+    (lib.attrValues diskoCfg.devices.disk));
 in
 {
-  pure = vmTools.runInLinuxVM (pkgs.runCommand name
+  system.build.diskoImages = vmTools.runInLinuxVM (pkgs.runCommand cfg.name
     {
       buildInputs = dependencies;
       inherit preVM postVM QEMU_OPTS;
-      memSize = nixosConfig.config.disko.memSize;
+      inherit (diskoCfg) memSize;
     }
     (partitioner + installer));
-  impure = diskoLib.writeCheckedBash { inherit checked pkgs; } name ''
+
+  system.build.diskoImagesScript = diskoLib.writeCheckedBash { inherit checked pkgs; } cfg.name ''
     set -efu
     export PATH=${lib.makeBinPath dependencies}
     showUsage() {

lib/types/disk.nix

@@ -19,7 +19,7 @@
     imageSize = lib.mkOption {
       type = lib.types.strMatching "[0-9]+[KMGTP]?";
       description = ''
-        size of the image if the makeDiskImages function from diksoLib is used.
+        size of the image when disko images are created
         is used as an argument to "qemu-img create ..."
       '';
       default = "2G";

module.nix

@@ -1,24 +1,20 @@
-{ config, lib, pkgs, extendModules, ... }@args:
+{ config, lib, pkgs, extendModules, diskoLib, ... }:
 let
-  diskoLib = import ./lib {
-    inherit lib;
-    rootMountPoint = config.disko.rootMountPoint;
-    makeTest = import (pkgs.path + "/nixos/tests/make-test-python.nix");
-    eval-config = import (pkgs.path + "/nixos/lib/eval-config.nix");
-  };
   cfg = config.disko;
 
   vmVariantWithDisko = extendModules {
     modules = [
       ./lib/interactive-vm.nix
-      { _module.args = { inherit diskoLib; }; }
       config.disko.tests.extraConfig
     ];
   };
 in
 {
+  imports = [ ./lib/make-disk-image.nix ];
+
   options.disko = {
-    imageBuilderQemu = lib.mkOption {
+    imageBuilder = {
+      qemu = lib.mkOption {
         type = lib.types.nullOr lib.types.str;
         description = ''
           the qemu emulator string used when building disk images via make-disk-image.nix.
@@ -28,7 +24,8 @@ in
         default = null;
         example = lib.literalExpression "\${pkgs.qemu_kvm}/bin/qemu-system-aarch64";
       };
-    imageBuilderPkgs = lib.mkOption {
+
+      pkgs = lib.mkOption {
         type = lib.types.attrs;
         description = ''
           the pkgs instance used when building disk images via make-disk-image.nix.
@@ -38,7 +35,8 @@ in
         defaultText = lib.literalExpression "pkgs";
         example = lib.literalExpression "pkgs";
       };
-    imageBuilderKernelPackages = lib.mkOption {
+
+      kernelPackages = lib.mkOption {
         type = lib.types.attrs;
         description = ''
           the kernel used when building disk images via make-disk-image.nix.
@@ -48,6 +46,7 @@ in
         defaultText = lib.literalExpression "config.boot.kernelPackages";
         example = lib.literalExpression "pkgs.linuxPackages_testing";
       };
+
       extraRootModules = lib.mkOption {
         type = lib.types.listOf lib.types.str;
         description = ''
@@ -56,6 +55,7 @@ in
         default = [ ];
         example = [ "bcachefs" ];
       };
+
       extraPostVM = lib.mkOption {
         type = lib.types.str;
         description = ''
@@ -67,18 +67,7 @@ in
           rm $out/*raw
         '';
       };
-    memSize = lib.mkOption {
-      type = lib.types.int;
-      description = ''
-        size of the memory passed to runInLinuxVM, in megabytes
-      '';
-      default = 1024;
-    };
-    devices = lib.mkOption {
-      type = diskoLib.toplevel;
-      default = { };
-      description = "The devices to set up";
-    };
+
       extraDependencies = lib.mkOption {
         type = lib.types.listOf lib.types.package;
         description = ''
@@ -86,11 +75,54 @@ in
         '';
         default = [ ];
       };
+
+      name = lib.mkOption {
+        type = lib.types.str;
+        description = "name for the disk images";
+        default = "${config.networking.hostName}-disko-images";
+      };
+
+      copyNixStore = lib.mkOption {
+        type = lib.types.bool;
+        description = "whether to copy the nix store into the disk images we just created";
+        default = true;
+      };
+
+      extraConfig = lib.mkOption {
+        description = ''
+          Extra NixOS config for your test. Can be used to specify a different luks key for tests.
+          A dummy key is in /tmp/secret.key
+        '';
+        default = { };
+      };
+
+      imageFormat = lib.mkOption {
+        type = lib.types.enum [ "raw" "qcow2" ];
+        description = "QEMU image format to use for the disk images";
+        default = "raw";
+      };
+    };
+
+    memSize = lib.mkOption {
+      type = lib.types.int;
+      description = ''
+        size of the memory passed to runInLinuxVM, in megabytes
+      '';
+      default = 1024;
+    };
+
+    devices = lib.mkOption {
+      type = diskoLib.toplevel;
+      default = { };
+      description = "The devices to set up";
+    };
+
     rootMountPoint = lib.mkOption {
       type = lib.types.str;
       default = "/mnt";
       description = "Where the device tree should be mounted by the mountScript";
     };
+
     enableConfig = lib.mkOption {
       description = ''
         configure nixos with the specified devices
@@ -100,6 +132,7 @@ in
       type = lib.types.bool;
       default = true;
     };
+
     checkScripts = lib.mkOption {
       description = ''
         Whether to run shellcheck on script outputs
@@ -107,6 +140,7 @@ in
       type = lib.types.bool;
       default = false;
     };
+
     testMode = lib.mkOption {
       internal = true;
       description = ''
@@ -116,6 +150,7 @@ in
       type = lib.types.bool;
       default = false;
     };
+
     tests = {
       efi = lib.mkOption {
         description = ''
@@ -126,6 +161,7 @@ in
         defaultText = "config.boot.loader.systemd-boot.enable || config.boot.loader.grub.efiSupport";
         default = config.boot.loader.systemd-boot.enable || config.boot.loader.grub.efiSupport;
       };
+
       extraChecks = lib.mkOption {
         description = ''
           extra checks to run in the `system.build.installTest`.
@@ -136,6 +172,7 @@ in
           machine.succeed("test -e /var/secrets/my.secret")
         '';
       };
+
       extraConfig = lib.mkOption {
         description = ''
           Extra NixOS config for your test. Can be used to specify a different luks key for tests.
@@ -155,20 +192,14 @@ in
     visible = "shallow";
   };
 
-  config = lib.mkIf (cfg.devices.disk != { }) {
+  config = lib.mkMerge [
+    (lib.mkIf (cfg.devices.disk != { }) {
       system.build = (cfg.devices._scripts { inherit pkgs; checked = cfg.checkScripts; }) // {
 
         # we keep these old outputs for compatibility
         disko = builtins.trace "the .disko output is deprecated, please use .diskoScript instead" (cfg.devices._scripts { inherit pkgs; }).diskoScript;
         diskoNoDeps = builtins.trace "the .diskoNoDeps output is deprecated, please use .diskoScriptNoDeps instead" (cfg.devices._scripts { inherit pkgs; }).diskoScriptNoDeps;
 
-      diskoImages = diskoLib.makeDiskImages {
-        nixosConfig = args;
-      };
-      diskoImagesScript = diskoLib.makeDiskImagesScript {
-        nixosConfig = args;
-      };
-
         installTest = diskoLib.testLib.makeDiskoTest {
           inherit extendModules pkgs;
           name = "${config.networking.hostName}-disko";
@@ -188,5 +219,14 @@ in
       fileSystems = lib.mkIf cfg.enableConfig cfg.devices._config.fileSystems or { };
       boot = lib.mkIf cfg.enableConfig cfg.devices._config.boot or { };
       swapDevices = lib.mkIf cfg.enableConfig cfg.devices._config.swapDevices or [ ];
+    })
+    {
+      _module.args.diskoLib = import ./lib {
+        inherit lib;
+        rootMountPoint = config.disko.rootMountPoint;
+        makeTest = import (pkgs.path + "/nixos/tests/make-test-python.nix");
+        eval-config = import (pkgs.path + "/nixos/lib/eval-config.nix");
       };
     }
+  ];
+}

tests/make-disk-image-impure.nix

@@ -1,15 +1,13 @@
 { pkgs ? import <nixpkgs> { }
-, diskoLib ? pkgs.callPackage ../lib { }
+, ...
 }:
-diskoLib.makeDiskImagesScript {
-  nixosConfig = pkgs.nixos [
+
+(pkgs.nixos [
   ../module.nix
   ../example/simple-efi.nix
   ({ config, ... }: {
     documentation.enable = false;
     system.stateVersion = config.system.nixos.version;
+    disko.checkScripts = true;
   })
-  ];
-  checked = true;
-}
-
+]).config.system.build.diskoImagesScript

tests/make-disk-image.nix

@@ -1,14 +1,14 @@
 { pkgs ? import <nixpkgs> { }
-, diskoLib ? pkgs.callPackage ../lib { }
+, ...
 }:
-diskoLib.makeDiskImages {
-  nixosConfig = pkgs.nixos [
+
+(pkgs.nixos [
   ../module.nix
   ../example/simple-efi.nix
   ({ config, ... }: {
     documentation.enable = false;
     system.stateVersion = config.system.nixos.version;
     disko.memSize = 2048;
+    disko.checkScripts = true;
   })
-  ];
-}
+]).config.system.build.diskoImages