# The appliance image: boot the raw disk image under UEFI and check what makes
# it an appliance (read-only verity store, volatile root, persistent /var) while
# the router itself runs. Then update it A/B: version 2 goes into the empty
# slot, boots and is blessed; a broken version 3 replaces version 1, fails its
# health check on every try and systemd-boot falls back to version 2.
{ self, pkgs }:
let
  inherit (pkgs) lib;
  common = import ./common.nix;
  keyring = import "${pkgs.path}/nixos/tests/common/gpg-keyring.nix" { inherit pkgs; };

  # The router as an image; v1 boots it, v2 and v3 only provide their updates.
  router =
    { lib, ... }:
    {
      imports = [
        self.nixosModules.router
        self.nixosModules.image
        common.tools
      ];

      # The VM boots the image's own disk, not the usual test plumbing.
      virtualisation.fileSystems = lib.mkVMOverride { };
      virtualisation.directBoot.enable = false;
      virtualisation.mountHostNixStore = false;
      virtualisation.useEFIBoot = true;
      virtualisation.vlans = [
        1
        2
      ];
      networking.interfaces.eth1.ipv4.addresses = lib.mkForce [ ];
      networking.interfaces.eth2.ipv4.addresses = lib.mkForce [ ];

      # Keep the disk (and what the test writes to it) small.
      router.image = {
        storeSize = "1G";
        storeVeritySize = "64M";
        tries = 2;
        update = {
          url = "http://192.168.10.2/";
          publicKey = "${keyring}/pubkey.gpg";
        };
      };

      router = {
        enable = true;
        wan.interface = "eth1";
        lans.main = {
          interface = "eth2";
          internet = true;
          ipv4.address = "192.168.10.1";
          dhcp.pool = {
            start = "192.168.10.100";
            end = "192.168.10.199";
          };
        };
      };
    };

  # An update server's directory: the updates plus signed checksums.
  site =
    name: updates:
    pkgs.runCommand "updates-${name}" { nativeBuildInputs = [ pkgs.gnupg ]; } ''
      mkdir -p $out
      cd $out
      for u in ${toString updates}; do cp $u/*.zst $u/*.efi .; done
      sha256sum *.zst *.efi > SHA256SUMS
      export GNUPGHOME=$(mktemp -d)
      cp -R ${keyring}/* $GNUPGHOME
      gpg --batch --sign --detach-sign --output SHA256SUMS.gpg SHA256SUMS
    '';
in
pkgs.testers.runNixOSTest {
  name = "router-image";

  nodes = {
    machine = {
      imports = [ router ];
      system.image.version = "1";
    };
    v2 = {
      imports = [ router ];
      system.image.version = "2";
    };
    # Never gets healthy.
    v3 = {
      imports = [ router ];
      system.image.version = "3";
      systemd.services.broken = {
        wantedBy = [ "multi-user.target" ];
        serviceConfig.Type = "oneshot";
        script = "exit 1";
      };
      router.image.healthChecks = [ "broken.service" ];
    };

    server =
      { lib, ... }:
      {
        virtualisation.vlans = [ 2 ];
        networking.interfaces.eth1.ipv4.addresses = lib.mkForce [
          {
            address = "192.168.10.2";
            prefixLength = 24;
          }
        ];
        networking.firewall.allowedTCPPorts = [ 80 ];
        services.nginx = {
          enable = true;
          virtualHosts.updates = {
            default = true;
            root = "/srv/updates";
          };
        };
      };
  };

  testScript =
    { nodes, ... }:
    let
      v2 = site "v2" [ nodes.v2.system.build.update ];
      v3 = site "v3" [
        nodes.v2.system.build.update
        nodes.v3.system.build.update
      ];
    in
    ''
      import datetime as dt
      import os
      import re
      import subprocess
      import tempfile

      # Before NIX_DISK_IMAGE is set: it applies to every VM started later.
      server.start()

      disk = tempfile.NamedTemporaryFile()
      subprocess.run([
        "${nodes.machine.virtualisation.qemu.package}/bin/qemu-img", "create",
        "-f", "qcow2",
        "-b", "${nodes.machine.system.build.image}/${nodes.machine.image.filePath}",
        "-F", "raw", disk.name,
      ], check=True)
      os.environ["NIX_DISK_IMAGE"] = disk.name

      def labels():
          return machine.succeed("lsblk -n -o PARTLABEL /dev/vda").split()

      def booted_version():
          return machine.succeed(". /etc/os-release; echo $IMAGE_VERSION").strip()

      def usr_slot():
          """The label of the partition the verity /usr is read from."""
          devs = re.findall(r"\((\w+)\)", machine.succeed("dmsetup deps -o devname usr"))
          parts = [machine.succeed(f"lsblk -n -o PARTLABEL /dev/{d}").strip() for d in devs]
          [data] = [p for p in parts if p.startswith("store_")]
          return data

      def reboot():
          machine.shutdown()
          machine.start()
          machine.wait_for_unit("multi-user.target")

      machine.wait_for_unit("multi-user.target")

      with subtest("the disk has the appliance layout, slot B is empty"):
          parts = labels()
          for label in ("esp", "store-verity_1", "store_1", "state"):
              assert label in parts, parts
          assert parts.count("_empty") == 2, parts
          assert usr_slot() == "store_1"

      with subtest("root is volatile and /nix/store is a dm-verity protected read-only partition"):
          machine.succeed("findmnt --kernel --type tmpfs /")
          assert "ACTIVE" in machine.succeed("dmsetup info --target verity usr")
          backing = machine.succeed("df --output=source /nix/store | tail -n1").strip()
          assert backing == "/dev/mapper/usr", backing
          machine.fail("touch /nix/store/should-not-work")
          machine.fail("touch /etc/should-not-work")

      with subtest("the image identifies itself"):
          assert booted_version() == "1"
          machine.succeed("grep -qE '^IMAGE_ID=.?nixos-router.?$' /etc/os-release")

      with subtest("the router comes up from the image"):
          machine.wait_for_unit("nftables.service")
          machine.wait_for_unit("unbound.service")
          machine.wait_for_unit("kea-dhcp4-server.service")
          machine.wait_until_succeeds("ip -4 addr show dev eth2 | grep -q 192.168.10.1")

      with subtest("there is no Nix on the device"):
          machine.fail("command -v nix-env")

      with subtest("/var is on the state partition and survives a reboot, the root does not"):
          src = machine.succeed("findmnt -n -o SOURCE /var").strip()
          assert "vda" in src, src
          machine.succeed("echo kept > /var/lib/marker; echo lost > /root/marker; sync")
          reboot()
          assert machine.succeed("cat /var/lib/marker").strip() == "kept"
          machine.fail("test -e /root/marker")
          machine.wait_for_unit("nftables.service")

      server.wait_for_unit("nginx.service")

      with subtest("an update whose checksums do not match the signature is refused"):
          server.succeed("mkdir -p /srv/updates && cp ${v2}/* /srv/updates/")
          server.succeed("echo tampered >> /srv/updates/SHA256SUMS")
          machine.wait_until_succeeds("curl -sf http://192.168.10.2/SHA256SUMS")
          out = machine.fail("router-update update 2>&1")
          assert "Signature verification failed" in out, out
          assert labels().count("_empty") == 2

      with subtest("version 2 is installed into the empty slot"):
          server.succeed("rm -r /srv/updates && ln -s ${v2} /srv/updates")
          machine.succeed("router-update update >&2")
          parts = labels()
          assert "store_2" in parts and "store-verity_2" in parts, parts
          assert "store_1" in parts, parts
          machine.succeed("test -e /boot/EFI/Linux/nixos-router_2+2-0.efi")

      with subtest("version 2 boots from slot B, is healthy and gets blessed"):
          reboot()
          assert booted_version() == "2"
          assert usr_slot() == "store_2"
          machine.wait_for_unit("router-healthy.service")
          machine.wait_for_unit("systemd-bless-boot.service")
          machine.succeed("test -e /boot/EFI/Linux/nixos-router_2.efi")

      with subtest("version 3 replaces version 1, not the running version"):
          server.succeed("ln -sfn ${v3} /srv/updates")
          machine.succeed("router-update update >&2")
          parts = labels()
          assert "store_3" in parts and "store_2" in parts and "store_1" not in parts, parts
          machine.succeed("test -e /boot/EFI/Linux/nixos-router_3+2-0.efi")
          machine.fail("test -e /boot/EFI/Linux/nixos-router_1.efi")

      with subtest("version 3 never gets healthy and version 2 boots again"):
          machine.shutdown()
          # Each try of version 3 fails its health check and reboots, which
          # ends the VM (-no-reboot); start it again like a reset would.
          for _ in range(2):
              machine.start()
              machine.wait_for_shutdown(timeout=dt.timedelta(minutes=3))
          machine.start()
          machine.wait_for_unit("multi-user.target")
          assert booted_version() == "2"
          assert usr_slot() == "store_2"
          machine.succeed("test -e /boot/EFI/Linux/nixos-router_3+0-2.efi")
          machine.wait_for_unit("nftables.service")
    '';
}