# Updating an installed router to a newer version of nixos-router.
#
# A release is a store path (see `lib.mkRelease`): the new source, its
# nixpkgs, the configuration flake that points at them and, for offline
# routers, the reference system with everything the rebuild needs. It travels
# as a binary cache signed with the update key: a bundle file (uploaded in the
# web UI or copied over) or a URL the router pulls from. `router-update`
# accepts it only when every path in its closure carries a signature of
# `routerUpdate.publicKeys`, then
#
#   1. imports it into the store,
#   2. points the configuration flake at it,
#   3. migrates settings.json to the release's schema (migrations/N.jq),
#   4. rebuilds with `router-apply` (commit-confirm: unless the update is
#      confirmed in time, the previous system comes back).
#
# The configuration directory as it was is kept in config.pre-update until the
# update is confirmed; a rollback (or a failed update) puts it back.
{
  config,
  lib,
  pkgs,
  ...
}:
let
  cfg = config.routerSettings;
  keys = config.routerUpdate.publicKeys;
  settings = lib.importJSON cfg.file;
  snapshot = "${cfg.stateDirectory}/config.pre-update";
  updateDir = "${cfg.stateDirectory}/update";
  webui = pkgs.callPackage ../../pkgs/router-webui/package.nix { };
  migrate = pkgs.callPackage ../../pkgs/router-migrate.nix { };

  restore = pkgs.writeShellApplication {
    name = "router-update-restore";
    runtimeInputs = [
      pkgs.coreutils
      pkgs.findutils
    ];
    text = ''
      [ -d ${snapshot} ] || exit 0
      find ${cfg.directory} -mindepth 1 -delete
      cp -a ${snapshot}/. ${cfg.directory}/
      rm -rf ${snapshot}
      echo "router: restored the configuration from before the update"
    '';
  };

  update = pkgs.writeShellApplication {
    name = "router-update";
    runtimeInputs = [
      config.nix.package
      config.systemd.package
      webui
      migrate
      pkgs.coreutils
      pkgs.curl
      pkgs.jq
      pkgs.util-linux
    ];
    text = ''
      usage() {
        echo "usage: router-update [--force] (--request | --bundle FILE | --url URL) [timeout]" >&2
        echo "  --bundle FILE  a bundle made by router-release" >&2
        echo "  --url URL      a server with the release's latest and cache/ (HTTP or HTTPS)" >&2
        echo "  --request      what the web UI asked for (${updateDir}/request.json)" >&2
        echo "  --force        also install a release that is not newer" >&2
        exit 2
      }
      die() { echo "router-update: $*" >&2; exit 1; }

      force=0 mode="" arg="" timeout=""
      while [ $# -gt 0 ]; do
        case $1 in
          --force) force=1 ;;
          --bundle) shift; [ $# -gt 0 ] || usage; mode=bundle arg=$1 ;;
          --url) shift; [ $# -gt 0 ] || usage; mode=url arg=$1 ;;
          --request)
            req=${updateDir}/request.json
            [ -r "$req" ] || die "no request"
            mode=$(jq -er .mode "$req") arg=$(jq -r '.url // ""' "$req")
            force=$(jq -r 'if .force then 1 else 0 end' "$req")
            rm -f "$req" ;;
          -*) usage ;;
          *) timeout=$1 ;;
        esac
        shift
      done
      [ -n "$mode" ] || usage
      [ "$(id -u)" -eq 0 ] || die "must run as root"
      exec 9>/run/router-update.lock
      flock -n 9 || die "another update is running"

      trusted=$(tr '\n' ' ' < /etc/router-update-keys)
      [ -n "''${trusted// /}" ] || die "this router has no update key"

      staged=${updateDir}/staged
      case $mode in
        bundle)
          rm -rf "$staged"
          router-webui unbundle "$arg" "$staged"
          mode=staged ;;
      esac
      case $mode in
        staged)
          store="file://$staged/cache"
          latest=$(cat "$staged/latest") ;;
        url)
          base=''${arg%/}
          store="$base/cache"
          latest=$(curl -fsS --max-time 30 "$base/latest") ;;
        *) usage ;;
      esac

      release=$(jq -er .release <<<"$latest") || die "the update has no release"
      [[ $release =~ ^/nix/store/[0-9a-z]{32}-[A-Za-z0-9+._-]+$ ]] || die "not a store path: $release"

      # Nothing is trusted until every path of the closure is signed by our key.
      echo "router-update: verifying $release"
      nix --option trusted-public-keys "$trusted" store verify \
        --store "$store" --no-contents --sigs-needed 1 -r "$release" \
        || die "the update is not signed with this router's update key"
      echo "router-update: importing it"
      nix --option trusted-public-keys "$trusted" copy --from "$store" "$release"
      rm -rf "$staged"

      new_version=$(cat "$release/version")
      new_serial=$(cat "$release/serial")
      cur_version=$(jq -r '.version // "unknown"' /etc/router-release)
      cur_serial=$(jq -r '.serial // 0' /etc/router-release)
      echo "router-update: installed $cur_version, update $new_version"
      if [ "$new_serial" -le "$cur_serial" ] && [ "$force" -ne 1 ]; then
        die "this release is not newer than the installed one (--force installs it anyway)"
      fi

      config=${cfg.directory}
      [ ! -d ${snapshot} ] || die "an earlier update is not confirmed yet"
      cp -a "$config" ${snapshot}
      fail() {
        echo "router-update: $*; keeping $cur_version" >&2
        ${restore}/bin/router-update-restore
        exit 1
      }

      install -m 0644 "$release/flake.nix" "$config/flake.nix"
      rm -f "$config/flake.lock"
      # A nixpkgs that router-upgrade chose stays, rather than the release's.
      if [ -r "$config/nixpkgs-pin.json" ]; then
        nix flake lock --override-input nixos-router/nixpkgs "$(jq -r .ref "$config/nixpkgs-pin.json")" "path:$config" \
          || fail "pinning the upgraded nixpkgs failed"
      fi

      router-migrate-settings "$release" "$config/settings.json" || fail "migrating the settings failed"

      echo "router-update: building and activating $new_version"
      if ! router-apply ''${timeout:+"$timeout"}; then
        if systemctl is-active --quiet router-rollback.timer; then
          # Activation failed half way: the rollback restores the configuration.
          systemctl start router-rollback.service
        else
          ${restore}/bin/router-update-restore
        fi
        die "applying $new_version failed; $cur_version is back"
      fi
      if systemctl is-active --quiet router-rollback.timer; then
        echo "router-update: $new_version is active; confirm it (router-confirm, or in the web UI) or $cur_version comes back"
      else
        rm -rf ${snapshot} # nothing changed: nothing to confirm
      fi
    '';
  };
in
{
  options = {
    routerSettings.release = lib.mkOption {
      type = lib.types.nullOr (lib.types.attrsOf lib.types.anything);
      default = null;
      description = "Version, serial, revision and settings schema of this nixos-router (`/etc/router-release`).";
    };
    routerUpdate.publicKeys = lib.mkOption {
      type = lib.types.listOf lib.types.str;
      default = [ ];
      example = [ "my-update-key-1:AAAA…=" ];
      description = ''
        Public keys (`nix key convert-secret-to-public`) of the keys whose
        signatures make an update acceptable.
      '';
    };
  };

  config = {
    environment.etc."router-release".text = builtins.toJSON (
      if config.routerSettings.release == null then { } else config.routerSettings.release
    );
    environment.etc."router-update-keys".text = lib.concatLines keys;
    environment.systemPackages = [
      update
      restore
    ];

    assertions = lib.optional (config.routerSettings.release != null) {
      assertion = (settings.version or 0) == config.routerSettings.release.schema;
      message = ''
        settings.json is at version ${
          toString (settings.version or 0)
        }, but this release of nixos-router expects
        version ${toString config.routerSettings.release.schema}. `router-update` migrates the settings; if you edited
        them by hand, set "version" accordingly.
      '';
    };

    # A rolled back or failed update puts the configuration directory back;
    # a confirmed one keeps the new one.
    router.rollback.onRollback = "${restore}/bin/router-update-restore";
    router.rollback.onConfirm = "rm -rf ${snapshot}";

    systemd.services.router-update = {
      description = "Install a router update";
      restartIfChanged = false;
      stopIfChanged = false;
      path = [ "/run/current-system/sw" ];
      serviceConfig = {
        Type = "oneshot";
        ExecStart = "${lib.getExe update} --request";
        TimeoutStartSec = "2h";
      };
    };

    # Old generations (and the updates' leftovers) would fill the disk.
    nix.gc = {
      automatic = true;
      dates = "weekly";
      options = "--delete-older-than 30d";
    };
  };
}