Pigeon/vm.nix

{
  pkgs,
  pigeon-toolchain,
  pigeon-vcpu-count ? 1,
  pigeon-memory-size ? 512,
  pigeon-disc-size ? "1G",
  pigeon-kernel-params ? [
    "64BIT" "TTY" "SERIAL_8250" "SERIAL_8250_CONSOLE" "PRINTK"

    "HYPERVISOR_GUEST" "PARAVIRT" "KVM_GUEST"

    "BLOCK" "VIRTIO_MENU" "VIRTIO" "VIRTIO_PCI" "VIRTIO_MMIO" "VIRTIO_BLK"
    "EXT4_FS" "SHMEM" "TMPFS" "PROC_FS" "PROC_SYSCTL" "SYSFS" "DEVTMPFS"

    "ACPI" "PCI" "PCI_MSI" "MAGIC_SYSRQ"

    "BINFMT_ELF" "BINFMT_SCRIPT" "MULTIUSER" "FUTEX" "EPOLL" "EVENTFD"
    "SIGNALFD" "TIMERFD" "POSIX_TIMERS" "FILE_LOCKING" "ADVISE_SYSCALLS"
    "NET" "UNIX"
  ]
}:

let
  kernel = pkgs.stdenv.mkDerivation {
    inherit (pkgs.linux_6_1) src version;
    pname = "pigeon-kernel";

    nativeBuildInputs = with pkgs; [ bc bison elfutils flex perl ];
    hardeningDisable = [ "all" ];
    enableParallelBuilding = true;

    configurePhase = ''
      runHook preConfigure
      patchShebangs scripts tools
      make tinyconfig
      for OPTION in ${toString pigeon-kernel-params}; do
        scripts/config --enable $OPTION
      done
      make olddefconfig
      for OPTION in ${toString pigeon-kernel-config}; do
        grep -qx "CONFIG_$OPTION=y" .config || exit 1
      done
      runHook postConfigure
    '';

    buildFlags = [ "vmlinux" ];

    installPhase = ''
      runHook preInstall
      cp vmlinux "$out"
      runHook postInstall
    '';
  };

  toolchain = pkgs.buildEnv { name = "pigeon-toolchain"; paths = pigeon-toolchain; };
  closure = pkgs.closureInfo { rootPaths = [ toolchain ]; };
  rootfs = pkgs.runCommand "pigeon-rootfs" { nativeBuildInputs = [ pkgs.e2fsprogs ]; } ''
    mkdir -p rootfs && cd rootfs
    mkdir -p dev proc sys tmp usr nix/store root/student root/grader root/results
    chmod 700 root

    for FILE in $(cat ${closure}/store-paths); do
      cp -a "$FILE" nix/store
    done
    ln -s ${toolchain}/bin bin
    ln -s ${toolchain}/bin usr/bin
    install -m 755 ${./init.sh} init
    install -m 755 ${./run.sh} run

    cd ..
    mkfs.ext4 -q -F -d rootfs $out ${pigeon-disc-size}
  '';
  
  vmconfig = pkgs.writeText "pigeon-vm.json" (builtins.toJSON {
    boot-source = {
      kernel_image_path = "${kernel}";
      boot_args = "reboot=k panic=1 quiet init=/init nomodule 8250.nr_uarts=0 i8042.noaux i8042.nomux i8042.dumbkbd swiotlb=noforce";
    };
    drives = [
      {
        drive_id = "rootfs";
        path_on_host = "${rootfs}";
        is_root_device = true;
        is_read_only = true;
      }
      {
        drive_id = "student";
        path_on_host = "";
        is_root_device = false;
        is_read_only = true;
      }
      {
        drive_id = "grader";
        path_on_host = "";
        is_root_device = false;
        is_read_only = true;
      }
      {
        drive_id = "results";
        path_on_host = "";
        is_root_device = false;
        is_read_only = false;
      }
    ];
    machine-config = {
      vcpu_count = pigeon-vcpu-count;
      mem_size_mib = pigeon-memory-size;
    };
  });
in
  pkgs.writeShellApplication {
    name = "pigeon-vm";
    runtimeInputs = with pkgs; [ coreutils firecracker jq ];
    text = ''
      [ "$#" -ne 3 ] && exit 1

      CONFIG="$(mktemp)"
      trap 'rm -rf "$CONFIG"' EXIT

      jq --arg s "$1" --arg g "$2" --arg r "$3" '
        (.drives[] | select(.drive_id == "student").path_on_host) = $s |
        (.drives[] | select(.drive_id == "grader").path_on_host) = $g |
        (.drives[] | select(.drive_id == "result").path_on_host) = $r
      ' ${vmconfig} >"$CONFIG"

      firecracker --no-api --config-file "$CONFIG"
    '';
  }