# `vmTests.<system>.wg-hive`: two hives peered over `wg-hive`, with real traffic across the tunnel.
# Unlike ../module-eval, which only reads the rendered `networking.wireguard`
# options, this boots both hosts and sends packets.
#
# Keys are minted by a throwaway derivation at build time and nothing is
# committed. The module takes peer public keys as eval-time strings
# (`swarm.hives.<name>.wireguardPublicKey`), so they are read back from
# the derivation's output.
{
  pkgs,
  self,
}:
let
  mkKey =
    name:
    pkgs.runCommand "wg-hive-test-key-${name}" { nativeBuildInputs = [ pkgs.wireguard-tools ]; } ''
      mkdir $out
      wg genkey > $out/private
      wg pubkey < $out/private > $out/public
    '';
  keys = {
    h1 = mkKey "h1";
    h2 = mkKey "h2";
    # Never installed on any node; stands in for a peer key that does not
    # match the real peer's private key.
    wrong = mkKey "wrong";
  };
  pub = name: pkgs.lib.removeSuffix "\n" (builtins.readFile "${keys.${name}}/public");

  # The test driver numbers nodes on the test vlan in alphabetical order.
  lanAddr = {
    h1 = "192.168.1.1";
    h2 = "192.168.1.2";
  };
  meshAddr = {
    h1 = "10.100.0.1";
    h2 = "10.100.0.2";
  };

  swarmHives = pkgs.lib.mapAttrs (name: _: {
    domain = "${name}.t.local";
    wireguardPublicKey = pub name;
    wireguardAddress = "${meshAddr.${name}}/32";
    wireguardEndpoint = "${lanAddr.${name}}:51820";
  }) lanAddr;

  node = name: {
    imports = [ self.nixosModules.default ];
    services.hyperhive = {
      swarm.domain = "t.local";
      swarm.hives = swarmHives;
      deploy.hive-controller.hiveName = name;
      deploy.wireguard = {
        enable = true;
        address = "${meshAddr.${name}}/24";
        privateKeyFile = "${keys.${name}}/private";
        persistentKeepalive = 1;
      };
    };
    # The listener the TCP check connects to, open on the mesh only.
    networking.firewall.interfaces.wg-hive.allowedTCPPorts = [ 9000 ];
    environment.systemPackages = [
      pkgs.wireguard-tools
      pkgs.python3
      pkgs.curl
    ];
  };
in
pkgs.testers.runNixOSTest {
  name = "wg-hive";

  nodes = {
    h1 = node "h1";
    h2 = node "h2";
  };

  testScript = ''
    def wg_state(m):
        return m.execute("wg show wg-hive; ip route; systemctl status 'wireguard-wg-hive*' --no-pager")[1]

    start_all()
    for m, peer_key in ((h1, "${pub "h2"}"), (h2, "${pub "h1"}")):
        m.wait_for_unit("wireguard-wg-hive.service")
        m.succeed("ip link show wg-hive")
        # The interface unit is active before its peers are installed: each
        # peer has its own unit, which runs after network-online.target. Until
        # then the interface has no peer and a send fails with ENOKEY.
        m.wait_until_succeeds(f"wg show wg-hive allowed-ips | grep -F {peer_key}")

    # Traffic makes the tunnel handshake.
    for src, addr in ((h1, "${meshAddr.h2}"), (h2, "${meshAddr.h1}")):
        status, out = src.execute(f"ping -c 3 -W 5 {addr}")
        assert status == 0, f"ping {addr} failed:\n{out}\n{wg_state(src)}"
    for m in (h1, h2):
        stamps = m.succeed("wg show wg-hive latest-handshakes").split()
        assert len(stamps) == 2 and int(stamps[1]) > 0, f"no handshake: {stamps}\n{wg_state(m)}"

    # The peer's mesh address is routed through the tunnel, its plain
    # address is not.
    via_mesh = h1.succeed("ip route get ${meshAddr.h2}")
    assert "dev wg-hive" in via_mesh, via_mesh
    via_lan = h1.succeed("ip route get ${lanAddr.h2}")
    assert "dev wg-hive" not in via_lan, via_lan

    # A TCP connect to a listener on the peer's mesh address, and the
    # tunnel's own counters moved.
    h2.succeed(
        "systemd-run --unit=wg-listener ${pkgs.python3}/bin/python3"
        " -m http.server 9000 --bind ${meshAddr.h2}"
    )
    h2.wait_for_open_port(9000, "${meshAddr.h2}")
    h1.succeed("${pkgs.curl}/bin/curl -sf --max-time 5 http://${meshAddr.h2}:9000/")
    transfer = h1.succeed("wg show wg-hive transfer").split()
    assert int(transfer[1]) > 0 and int(transfer[2]) > 0, transfer

    # Negative control: h2 swaps h1's public key for one that does not match
    # h1's private key. Same endpoint, same address, so the only difference
    # from the working peer above is the key, and no handshake may happen.
    h2.succeed("wg set wg-hive peer ${pub "h1"} remove")
    h2.succeed(
        "wg set wg-hive peer ${pub "wrong"} allowed-ips ${meshAddr.h1}/32"
        " endpoint ${lanAddr.h1}:51820 persistent-keepalive 1"
    )
    h2.fail("ping -c 3 -W 2 ${meshAddr.h1}")
    stamps = h2.succeed("wg show wg-hive latest-handshakes").split()
    assert stamps == ["${pub "wrong"}", "0"], f"handshake with wrong key: {stamps}"
  '';
}