# Appliance image: a GPT disk image with a dm-verity protected, read-only
# /nix/store partition, a volatile (tmpfs) root and a small persistent state
# partition for /var. Built with systemd-repart; the unified kernel image on the
# ESP carries the hash of the store partition, so a modified store does not boot.
#
# The system cannot be rebuilt on the device (no Nix); it is replaced as a whole
# (A/B): the store comes in two slots, systemd-sysupdate writes a new version
# into the slot not in use and adds its UKI with boot counting. A version that
# does not reach a healthy state within its tries is skipped by systemd-boot,
# which then boots the previous one again.
#
#   ESP (systemd-boot, UKIs) | store-verity_A | store_A | store-verity_B | store_B | state (/var)
#
# Importing this module turns a configuration into such an image
# (`config.system.build.image`) and its update (`config.system.build.update`).
{
  config,
  lib,
  pkgs,
  modulesPath,
  ...
}:
let
  cfg = config.router.image;
  inherit (config.image.repart.verityStore) partitionIds;
  inherit (config.system.boot.loader) ukiFile;
  efiArch = pkgs.stdenv.hostPlatform.efiArch;
  version = toString config.system.image.version;
  name = config.boot.uki.name;
  usrVerityType = config.image.repart.partitions.${partitionIds.store-verity}.repartConfig.Type;

  # Both slots have the same fixed size, so any version fits either slot.
  fixedSize = size: {
    SizeMinBytes = size;
    SizeMaxBytes = size;
  };
  emptySlot = type: size: {
    repartConfig = {
      Type = type;
      Label = "_empty";
      SplitName = "-";
    }
    // fixedSize size;
  };

  sysupdate = "${config.systemd.package}/lib/systemd/systemd-sysupdate";
  source = pattern: {
    Type = "url-file";
    Path = cfg.update.url;
    MatchPattern = pattern;
  };
in
{
  imports = [
    (modulesPath + "/image/repart.nix")
    (modulesPath + "/profiles/image-based-appliance.nix")
    (modulesPath + "/profiles/perlless.nix")
  ];

  options.router.image = {
    espSize = lib.mkOption {
      type = lib.types.str;
      default = "256M";
      description = "Size of the EFI system partition (systemd-boot and the UKIs of both slots).";
    };
    storeSize = lib.mkOption {
      type = lib.types.str;
      default = "2G";
      description = ''
        Size of each store slot. Every future version must fit; it cannot be
        changed on an installed device.
      '';
    };
    storeVeritySize = lib.mkOption {
      type = lib.types.str;
      default = "128M";
      description = "Size of each slot's dm-verity hash partition.";
    };
    stateSize = lib.mkOption {
      type = lib.types.str;
      default = "512M";
      description = "Minimum size of the persistent state partition mounted on /var.";
    };
    tries = lib.mkOption {
      type = lib.types.ints.positive;
      default = 3;
      description = ''
        Boots a freshly installed version gets to become healthy before
        systemd-boot falls back to the previous one.
      '';
    };
    healthChecks = lib.mkOption {
      type = lib.types.listOf lib.types.str;
      description = ''
        Units that must be active for a new version to count as good. The
        router's own services are always included.
      '';
    };
    update = {
      url = lib.mkOption {
        type = lib.types.nullOr lib.types.str;
        default = null;
        example = "https://updates.example.com/router/";
        description = ''
          Where `router-update` looks for new versions: a directory with the
          files of `system.build.update`, a SHA256SUMS and, when
          `publicKey` is set, its detached signature SHA256SUMS.gpg.
        '';
      };
      publicKey = lib.mkOption {
        type = lib.types.nullOr lib.types.path;
        default = null;
        description = ''
          OpenPGP public key (armored) that signs SHA256SUMS. Without it updates
          are not verified.
        '';
      };
    };
  };

  config = {
    assertions = [
      {
        assertion = config.system.image.version != null;
        message = "router.image: set system.image.version; updates are told apart by it.";
      }
    ];
    warnings = lib.optional (cfg.update.url != null && cfg.update.publicKey == null) ''
      router.image.update.publicKey is not set: updates are installed without verifying them.
    '';

    system.image.id = lib.mkDefault "nixos-router";

    # No Nix on the device: the rollback helpers need it, and boot counting
    # takes over their job in images.
    router.rollback.enable = false;

    boot.loader.grub.enable = false;

    fileSystems = {
      "/" = {
        fsType = "tmpfs";
        options = [
          "mode=0755"
          "size=50%"
        ];
      };
      # The only persistent, writable place: DHCP leases, SSH host keys,
      # journal, secrets (/var/lib/router).
      "/var" = {
        device = "/dev/disk/by-partlabel/state";
        fsType = "ext4";
        options = [
          "nosuid"
          "nodev"
        ];
      };
      # UKIs are installed and blessed here.
      "/boot" = {
        device = "/dev/disk/by-partlabel/esp";
        fsType = "vfat";
        options = [ "umask=0077" ];
      };
    };

    # /etc is a read-only overlay; users are created by userborn, which keeps
    # its password files on the state partition.
    system.etc.overlay.mutable = false;
    services.userborn.passwordFilesLocation = "/var/lib/nixos";

    # Host keys live on the state partition as well.
    services.openssh.hostKeys = [
      {
        path = "/var/lib/ssh/ssh_host_ed25519_key";
        type = "ed25519";
      }
    ];

    # /usr is the read-only verity mount; there is nothing to create there.
    system.activationScripts.usrbinenv = lib.mkForce "";

    image.repart = {
      enable = true;
      split = true;

      # systemd-repart passes the image's sector size (512) to mkfs.erofs as the
      # block size; libblkid cannot verify the superblock checksum of an erofs
      # with blocks <= 1 KiB, so udev never reports the store's file system and
      # systemd waits forever for /dev/mapper/usr. Appended options win.
      mkfsOptions.erofs = [ "-b4096" ];

      verityStore.enable = true;

      partitions = {
        ${partitionIds.esp} = {
          contents = {
            "/EFI/BOOT/BOOT${lib.toUpper efiArch}.EFI".source =
              "${config.systemd.package}/lib/systemd/boot/efi/systemd-boot${efiArch}.efi";
            "/loader/loader.conf".source = pkgs.writeText "loader.conf" ''
              timeout 0
            '';
          };
          repartConfig = {
            Type = "esp";
            Format = "vfat";
            SplitName = "-";
          }
          // fixedSize cfg.espSize;
        };
        # Slot A holds this version. The labels carry the version, the UUIDs
        # come from the verity root hash (that is how the UKI finds its slot),
        # and the split files named after both are the update.
        ${partitionIds.store-verity}.repartConfig = {
          Label = "store-verity_${version}";
          Minimize = "off";
          SplitName = "store-verity-%U";
        }
        // fixedSize cfg.storeVeritySize;
        ${partitionIds.store}.repartConfig = {
          Label = "store_${version}";
          Minimize = "off";
          Compression = "lz4hc";
          SplitName = "store-%U";
        }
        // fixedSize cfg.storeSize;
        # Slot B starts empty; systemd-sysupdate fills it.
        "30-store-verity-empty" = emptySlot "usr-verity" cfg.storeVeritySize;
        "40-store-empty" = emptySlot "usr" cfg.storeSize;
        "50-state".repartConfig = {
          Type = "var";
          Format = "ext4";
          Label = "state";
          SplitName = "-";
          SizeMinBytes = cfg.stateSize;
        };
      };
    };

    # Upstream reads the usrhash from every usr-verity partition; the empty
    # slot B has none (its roothash is null), so pick the one that has one.
    system.build.uki = lib.mkOverride 98 (
      pkgs.runCommand ukiFile
        {
          nativeBuildInputs = [
            pkgs.buildPackages.jq
            pkgs.buildPackages.systemdUkify
          ];
        }
        ''
          mkdir -p $out
          usrhash=$(jq -r \
            '.[] | select(.type=="${usrVerityType}" and .roothash != null) | .roothash' \
            ${config.system.build.intermediateImage}/repart-output.json
          )
          ukify build \
              --config=${config.boot.uki.configFile} \
              --cmdline="init=${config.system.build.toplevel}/init ${toString config.boot.kernelParams} usrhash=$usrhash" \
              --output="$out/${ukiFile}"
        ''
    );

    # What a device needs to move to this version: both store partitions
    # (compressed; mostly empty space) and the UKI.
    system.build.update =
      pkgs.runCommand "${config.image.baseName}-update" { nativeBuildInputs = [ pkgs.zstd ]; }
        ''
          mkdir -p $out
          for f in ${config.system.build.intermediateImage}/${config.image.baseName}.store*.raw; do
            zstd -q -T0 -o $out/$(basename $f).zst $f
          done
          cp ${config.system.build.uki}/${ukiFile} $out/
          (cd $out && sha256sum * > SHA256SUMS)
        '';

    router.image.healthChecks = [
      "systemd-networkd.service"
      "nftables.service"
    ]
    ++ lib.optional config.services.unbound.enable "unbound.service"
    ++ lib.optional config.services.kea.dhcp4.enable "kea-dhcp4-server.service";

    # Boot assessment: only boots that are being counted pull in
    # boot-complete.target (via systemd-bless-boot). If the router does not
    # come up, reboot; once the tries are used up systemd-boot falls back to
    # the previous version.
    systemd.services.router-healthy = {
      description = "Check that the router came up";
      requiredBy = [ "boot-complete.target" ];
      before = [ "boot-complete.target" ];
      after = cfg.healthChecks;
      unitConfig.FailureAction = "reboot";
      serviceConfig = {
        Type = "oneshot";
        RemainAfterExit = true;
        # One per unit: `systemctl is-active a b` succeeds if any one is active.
        ExecStart = map (
          unit: "${lib.getExe' config.systemd.package "systemctl"} is-active --quiet ${unit}"
        ) cfg.healthChecks;
      };
    };

    systemd.sysupdate = lib.mkIf (cfg.update.url != null) {
      enable = true;
      transfers =
        let
          common = {
            Transfer = {
              ProtectVersion = "%A";
              Verify = cfg.update.publicKey != null;
            };
          };
          partition = label: type: split: {
            Source = source "${name}_@v.${split}-@u.raw.zst";
            Target = {
              Type = "partition";
              Path = "auto";
              MatchPattern = "${label}_@v";
              MatchPartitionType = type;
              ReadOnly = true;
            };
          };
        in
        {
          # Applied in this order: the UKI last, so a half-written update can
          # never be booted.
          "10-store-verity" = common // partition "store-verity" "usr-verity" "store-verity";
          "20-store" = common // partition "store" "usr" "store";
          "90-uki" = common // {
            Source = source "${name}_@v.efi";
            Target = {
              Type = "regular-file";
              Path = "/EFI/Linux";
              PathRelativeTo = "boot";
              MatchPattern = [
                "${name}_@v+@l-@d.efi"
                "${name}_@v+@l.efi"
                "${name}_@v.efi"
              ];
              Mode = "0444";
              TriesLeft = cfg.tries;
              TriesDone = 0;
              InstancesMax = 2;
            };
          };
        };
    };

    environment.etc."systemd/import-pubring.gpg" = lib.mkIf (cfg.update.publicKey != null) {
      source = cfg.update.publicKey;
    };

    # `router-update list`, `router-update update`, then reboot.
    environment.systemPackages = [
      # systemd-pull runs gpg to check SHA256SUMS.gpg.
      (pkgs.writeShellScriptBin "router-update" ''
        PATH=${lib.makeBinPath [ pkgs.gnupgMinimal ]}:$PATH exec ${sysupdate} "$@"
      '')
    ];
  };
}