Add A/B updates with boot counting to the appliance image

The image gets a second, empty store slot. systemd-sysupdate (wrapped as
router-update) installs a new version into the slot not in use: both
store partitions, then the UKI with boot counting, verified against a
signed SHA256SUMS. system.build.update holds a version's files.

Counted boots must pass router-healthy (the router's units plus
router.image.healthChecks, each checked on its own) to reach
boot-complete.target; otherwise the box reboots, and once the tries are
used up systemd-boot falls back to the previous version.

checks.image now installs v2 into the empty slot, boots and blesses it,
refuses a tampered update, and rolls a broken v3 back to v2.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

examples/image/default.nix

@@ -9,6 +9,8 @@
   imports = [ ../../modules/image ];
 
   networking.hostName = "router";
+  # Bump for every build that is meant to be installed as an update.
+  system.image.version = "1";
   system.stateVersion = "26.11";
 
   # Generic UEFI PC: storage, USB keyboard and NIC drivers in the initrd.

modules/image/default.nix

@@ -3,9 +3,16 @@
 # 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.
+# 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`).
+# (`config.system.build.image`) and its update (`config.system.build.update`).
 {
   config,
   lib,
@@ -16,7 +23,32 @@
 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 = [
@@ -27,18 +59,76 @@ in
 
   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 the EFI system partition (holds the unified kernel image).";
+      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
@@ -65,6 +155,12 @@ in
           "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
@@ -85,31 +181,180 @@ in
 
     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;
+      verityStore.enable = true;
-        # The firmware's default boot path: no NVRAM entry needed.
-        ukiPath = "/EFI/BOOT/BOOT${lib.toUpper efiArch}.EFI";
-      };
 
       partitions = {
-        ${partitionIds.esp}.repartConfig = {
+        ${partitionIds.esp} = {
-          Type = "esp";
+          contents = {
-          Format = "vfat";
+            "/EFI/BOOT/BOOT${lib.toUpper efiArch}.EFI".source =
-          SizeMinBytes = cfg.espSize;
+              "${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;
         };
-        "30-state".repartConfig = {
+        # 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} "$@"
+      '')
+    ];
   };
 }

tests/image.nix

@@ -1,14 +1,16 @@
 # 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.
+# 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;
-in
+  keyring = import "${pkgs.path}/nixos/tests/common/gpg-keyring.nix" { inherit pkgs; };
-pkgs.testers.runNixOSTest {
-  name = "router-image";
 
-  nodes.machine =
+  # The router as an image; v1 boots it, v2 and v3 only provide their updates.
+  router =
     { lib, ... }:
     {
       imports = [
@@ -29,7 +31,16 @@ pkgs.testers.runNixOSTest {
       networking.interfaces.eth1.ipv4.addresses = lib.mkForce [ ];
       networking.interfaces.eth2.ipv4.addresses = lib.mkForce [ ];
 
-      system.image.version = "1";
+      # 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;
@@ -46,13 +57,83 @@ pkgs.testers.runNixOSTest {
       };
     };
 
+  # 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",
@@ -62,12 +143,32 @@ pkgs.testers.runNixOSTest {
       ], 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"):
+      with subtest("the disk has the appliance layout, slot B is empty"):
-          parts = machine.succeed("lsblk -n -o PARTLABEL /dev/vda")
+          parts = labels()
-          for label in ("esp", "store-verity", "store", "state"):
+          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 /")
@@ -78,7 +179,7 @@ pkgs.testers.runNixOSTest {
           machine.fail("touch /etc/should-not-work")
 
       with subtest("the image identifies itself"):
-          machine.succeed("grep -qE '^IMAGE_VERSION=.?1.?$' /etc/os-release")
+          assert booted_version() == "1"
           machine.succeed("grep -qE '^IMAGE_ID=.?nixos-router.?$' /etc/os-release")
 
       with subtest("the router comes up from the image"):
@@ -94,11 +195,57 @@ pkgs.testers.runNixOSTest {
           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")
-          machine.shutdown()
+          reboot()
-          machine.start()
-          machine.wait_for_unit("multi-user.target")
           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")
     '';
 }