NixOSConfig/lib/modules/custom-build-vm.nix

{ config, lib, options, pkgs, ... }:
let
  packageName = "custom-nixos-vm";

  # Determine `system.build` configuration without this file's influence
  mergedBuildOption =
    with options.system;
    lib.mergeDefinitions
      build.loc
      build.type
      (lib.lists.forEach
        (
          builtins.filter
            (item:
              !(lib.path.hasPrefix ./. (/. + item.file)))
            build.definitionsWithLocations)
        (item: { inherit (item) file value; }));

  # Get vanilla `config.system.build.vm`
  vanillaVM = mergedBuildOption.mergedValue.vm;
in {
  options =
  let
    # Add new options to `config.virtualisation.vmVariant` and `config.virtualisation.vmVariantWithBootLoader`
    vmVariantOptions = {
      virtualisation = {
        runAsRoot = lib.mkOption {
          type = lib.types.bool;
          description = "Whether to launch the VM as root.";
          default = false;
        };

        sharedHostKeys = lib.mkOption {
          type = lib.types.bool;
          description = "Whether to share the local host SSH keys with the VM.";
          default = false;
        };

        usb-redirect = lib.mkOption {
          type = lib.types.bool;
          description = lib.mdDoc "Whether to enable USB redirection to the VM.";
          default = false;
        };

        virt-viewer = lib.mkOption {
          type = lib.types.bool;
          description = "Whether to use `remote-viewer` for displaying the VM.";
          default = false;
        };

        qemu = {
          runInBackground = lib.mkOption {
            type = lib.types.bool;
            description = "Whether to run the QEMU command in a background job";
            default = false;
          };

          spice = {
            enable = lib.mkEnableOption "spice";

            bindAddress = lib.mkOption {
              type = lib.types.str;
              description = "The IP address for listening to incoming SPICE connections.";
              default = "127.0.0.1";
            };

            port = lib.mkOption {
              type = lib.types.port;
              description = "The port for listening to incoming SPICE connections.";
              default = 5900;
            };
          };
        };
      };
    };
  in {
    virtualisation = {
      vmVariant = vmVariantOptions;
      vmVariantWithBootLoader = vmVariantOptions;
    };
  };

  config = {
    virtualisation =
      let
        extendVMConfig =
          vmVariant: {
            # Prevent GRUB2 errors in `nixos-rebuild build-vm-with-bootloader`
            boot.loader.efi.efiSysMountPoint = lib.mkVMOverride "/boot";

            virtualisation = {
              # Enable root permissions to get access to the `/etc/ssh` directory
              runAsRoot = lib.mkIf
                (vmVariant.virtualisation.sharedHostKeys || vmVariant.virtualisation.usb-redirect)
                true;

              # Enable spice and run QEMU in background to let `remote-viewer` take over
              qemu = {
                spice.enable = lib.mkIf vmVariant.virtualisation.virt-viewer true;
                runInBackground = lib.mkIf vmVariant.virtualisation.virt-viewer true;

                options =
                  with {
                    inherit (vmVariant.virtualisation.qemu) spice;
                    inherit (vmVariant.virtualisation) usb-redirect;
                  };
                  (
                    lib.optionals usb-redirect (
                      [
                        "-usb"
                        "-device qemu-xhci"
                      ] ++ (builtins.concatMap
                        (index:
                          let
                            devName = "usbredirchardev${toString index}";
                          in [
                            "-chardev spicevmc,name=usbredir,id=${devName}"
                            "-device usb-redir,chardev=${devName},id=usbredirdev${toString index}"
                          ])
                        (lib.lists.range 1 3)))) ++
                  (
                    lib.optional (spice.enable)
                      ("-spice " + (
                        lib.concatStringsSep "," [
                          "addr=${lib.escapeShellArg spice.bindAddress}"
                          "port=${toString spice.port}"
                          "disable-ticketing=on"
                        ])));
              };

              # Map SSH keys into the vm if necessary
              sharedDirectories = lib.optionalAttrs (vmVariant.virtualisation.sharedHostKeys) {
                hostKeys =
                let
                  path = "/etc/ssh";
                in {
                  source = path;
                  target = path;
                };
              };
            };
          };

        inherit (config.virtualisation)
          vmVariant
          vmVariantWithBootLoader
        ;
      in {
        vmVariant = extendVMConfig vmVariant;
        vmVariantWithBootLoader = extendVMConfig vmVariantWithBootLoader;
      };

    system.build =
      {
        vm = lib.mkForce (
          let
            vm = vanillaVM;
          in
            if (vm.name == packageName)
            then
              vm
            else
              let
                originalCommand = "${vm}/bin/run-${config.system.name}-vm";

                # Have the command run in background if requested
                suffix =
                  lib.concatStringsSep " " (
                    lib.optional config.virtualisation.qemu.runInBackground "&");

                vmRunner = pkgs.writeShellApplication {
                  name = "run-${config.system.name}-vm";
                  runtimeInputs = lib.optional config.virtualisation.usb-redirect pkgs.spice-gtk;

                  text = lib.strings.concatLines (
                    [
                      "${originalCommand} ${suffix}"
                    ] ++ (
                      # Run `remote-viewer` as normal user to limit access
                      (
                        lib.optionals
                          config.virtualisation.virt-viewer (
                            let
                              spice = config.virtualisation.qemu.spice;
                              remoteAddress = "spice://${lib.escapeShellArg spice.bindAddress}:${toString spice.port}";
                            in
                              [
                                "${pkgs.virt-viewer}/bin/remote-viewer ${remoteAddress}"
                                # Kill QEMU after `remote-viewer` finished running
                                "kill %1"
                              ]))));
                };

                # Run VM as root if requested
                wrapped =
                  if !config.virtualisation.runAsRoot
                  then
                    vmRunner
                  else
                    pkgs.writeShellApplication {
                      inherit (vmRunner) name;
                      
                      text = ''
                        sudo -E "${vmRunner}/bin/${vmRunner.name}"
                      '';
                    };
              in
                pkgs.symlinkJoin {
                  name = packageName;
                  paths = [ wrapped ];
                });
      };
  };
}