# WireGuard: road warrior, site-to-site and a provider ("Freifunk-style")
# tunnel, all on one router.
#
#   laptop -----+                                 +--- wg-rw  (dual-stack, zone vpn)
#   siteb ------+-- vlan 1 --  router  -- vlan 2 -- main LAN (zone trusted)
#   isp/provider+                                 +--- wg-s2s (no transfer net, zone remote)
#                                                 +--- wg-ff  (dual-stack, zone freifunk)
#
#   rules: vpn -> trusted, trusted -> remote, vpn -> remote, trusted -> freifunk
#
# The keys below are throwaway test keys, generated for this test only.
{ self, pkgs }:
let
  inherit (pkgs) lib;
  common = import ./common.nix;

  key = {
    router-rw = {
      priv = "wHjLTos8CeZxnCPWPKlZjlzDuer2uyhhNbXR6BYYkXY=";
      pub = "7BBr4sUxdMI9OlHcEUY+sg5eh8fbeo2nZzSXEI7IJQQ=";
    };
    router-s2s = {
      priv = "CM1gl4vlKo7iULJ07wykWxrHke4YTahZq+LCakFmT0g=";
      pub = "8/XG5fihWpnHQWQ/cKEjCBZvtk+DmBK/05gR26/Wfzo=";
    };
    router-ff = {
      priv = "2FFVahDzQTlzIpHmX95WeliWpSFKGVlfvJ6iC9iqA1Q=";
      pub = "lgTbki7gCrRFQa1WmV2IVlocWDMh5Fta0S2+DfVF93k=";
    };
    laptop = {
      priv = "uEWQDRJkpi7JFThKSdLASqd+Z8+2buzzEdcehNXU9Uc=";
      pub = "ru65s6+UROJ6CDfU15MS83Wlq1jYOJ5Leus1+3FkWmo=";
    };
    siteb = {
      priv = "WPqvH1WTD2nqq/+Ja474J79PAEUNomSVNvUvuQIAw1Q=";
      pub = "VPHx2uwfYY8BVeZRMUHRBAyFJ+CUnulnlkvzibBGSXo=";
    };
    isp = {
      priv = "sNot2jWhW+MJKBVScvm56FiA0Ekt/DfrFfKYXpK5RUQ=";
      pub = "7MkFXOxlCsmXm98uv0txvYzx5Ce80mkg2mDsswY4FyI=";
    };
  };

  # Plain (non-router) hosts: static address on the shared segment plus one
  # WireGuard device configured straight through networkd.
  segmentHost =
    {
      address,
      wgName,
      wgKey,
      listenPort ? null,
      wgAddresses ? [ ],
      peer,
    }:
    { lib, ... }:
    {
      virtualisation.vlans = [ 1 ];
      networking.useNetworkd = true;
      networking.useDHCP = false;
      networking.firewall.enable = false;
      networking.interfaces.eth1.ipv4.addresses = lib.mkForce [ ];

      environment.etc."wg/${wgName}.key" = {
        text = wgKey.priv;
        mode = "0400";
        # These hosts use plain networkd, which reads the key itself.
        user = "systemd-network";
        group = "systemd-network";
      };
      systemd.network.networks."10-seg" = {
        matchConfig.Name = "eth1";
        address = [ "${address}/24" ];
      };
      systemd.network.netdevs."30-${wgName}" = {
        netdevConfig = {
          Kind = "wireguard";
          Name = wgName;
        };
        wireguardConfig = {
          PrivateKeyFile = "/etc/wg/${wgName}.key";
        }
        // lib.optionalAttrs (listenPort != null) { ListenPort = listenPort; };
        wireguardPeers = [ peer ];
      };
      systemd.network.networks."30-${wgName}" = {
        matchConfig.Name = wgName;
        address = wgAddresses;
        routes = map (ip: { Destination = ip; }) peer.AllowedIPs;
      };
    };

  dummy = name: addresses: {
    systemd.network.netdevs."20-${name}".netdevConfig = {
      Kind = "dummy";
      Name = name;
    };
    systemd.network.networks."20-${name}" = {
      matchConfig.Name = name;
      address = addresses;
    };
  };

  echo = {
    services.nginx = {
      enable = true;
      virtualHosts.default = {
        default = true;
        locations."/".return = ''200 "$remote_addr"'';
      };
    };
  };
in
pkgs.testers.runNixOSTest {
  name = "router-wireguard";

  defaults.imports = [ common.tools ];
  interactive.defaults = common.interactiveDefaults;

  nodes = {
    # ISP, DHCP server for the router, and the VPN provider: the tunnel leads
    # to 198.18.0.1 / fd00:ff::1, which are also reachable via the WAN on this
    # very host (so a leak past the kill switch would succeed).
    isp = {
      imports = [
        (segmentHost {
          address = "203.0.113.1";
          wgName = "wg-ff";
          wgKey = key.isp;
          listenPort = 51822;
          wgAddresses = [
            "10.200.0.1/24"
            "fd42:200::1/64"
          ];
          peer = {
            PublicKey = key.router-ff.pub;
            AllowedIPs = [
              "10.200.0.2/32"
              "fd42:200::2/128"
            ];
          };
        })
        (dummy "ff0" [
          "198.18.0.1/24"
          "fd00:ff::1/128"
        ])
        echo
      ];
      services.dnsmasq = {
        enable = true;
        settings = {
          interface = "eth1";
          bind-interfaces = true;
          dhcp-range = "203.0.113.100,203.0.113.100,255.255.255.0,1h";
          dhcp-option = [ "option:router,203.0.113.1" ];
          address = "/example.test/203.0.113.1";
        };
      };
    };

    router =
      { lib, ... }:
      {
        imports = [ self.nixosModules.router ];
        virtualisation.vlans = [
          1
          2
        ];
        networking.interfaces.eth1.ipv4.addresses = lib.mkForce [ ];
        networking.interfaces.eth2.ipv4.addresses = lib.mkForce [ ];

        # Names for peer endpoints, so the test can move them.
        networking.hosts = {
          "203.0.113.60" = [ "siteb.test" ];
          "203.0.113.70" = [ "ghost.test" ];
        };

        environment.etc = lib.listToAttrs (
          map
            (n: {
              name = "wg/${n}.key";
              value = {
                text = key."router-${n}".priv;
                mode = "0600";
              };
            })
            [
              "rw"
              "s2s"
              "ff"
            ]
        )
        // {
          # A tunnel to a peer that does not exist: it never has a handshake.
          "wg/ghost.key" = {
            text = key.router-rw.priv;
            mode = "0600";
          };
        };

        router = {
          enable = true;
          wan.interface = "eth1";
          lans.main = {
            interface = "eth2";
            zone = "trusted";
            internet = true;
            ipv4.address = "192.168.10.1";
            dhcp.pool = {
              start = "192.168.10.100";
              end = "192.168.10.199";
            };
          };
          dns.forwarders = [ "203.0.113.1" ];

          wireguard = {
            rw = {
              privateKeyFile = "/etc/wg/rw.key";
              listenPort = 51820;
              addresses = [
                "10.100.0.1/24"
                "fd42:100::1/64"
              ];
              zone = "vpn";
              allowManagement = true;
              dns = true;
              peers.laptop = {
                publicKey = key.laptop.pub;
                allowedIPs = [
                  "10.100.0.2/32"
                  "fd42:100::2/128"
                ];
              };
            };
            # Site-to-site: no transfer network, routing by allowedIPs alone.
            s2s = {
              privateKeyFile = "/etc/wg/s2s.key";
              listenPort = 51821;
              zone = "remote";
              # No address on the tunnel: the site queries the LAN address and
              # is allowed because it routes 192.168.50.0/24 through the tunnel.
              dns = true;
              # Router-originated traffic to the site uses the LAN address.
              sourceAddresses = [ "192.168.10.1" ];
              peers.siteb = {
                publicKey = key.siteb.pub;
                endpoint = "siteb.test:51821";
                allowedIPs = [ "192.168.50.0/24" ];
                persistentKeepalive = 5;
              };
            };
            ghost = {
              privateKeyFile = "/etc/wg/ghost.key";
              zone = "ghost";
              peers.nobody = {
                publicKey = key.laptop.pub;
                endpoint = "ghost.test:51830";
                allowedIPs = [ "192.168.77.0/24" ];
              };
            };
            # Provider tunnel: only the listed destinations go through it.
            ff = {
              privateKeyFile = "/etc/wg/ff.key";
              zone = "freifunk";
              masquerade = true;
              addresses = [
                "10.200.0.2/32"
                "fd42:200::2/128"
              ];
              peers.provider = {
                publicKey = key.isp.pub;
                endpoint = "203.0.113.1:51822";
                allowedIPs = [
                  # The provider's own transfer addresses, plus what it serves.
                  "10.200.0.1/32"
                  "fd42:200::1/128"
                  "198.18.0.0/24"
                  "fd00:ff::/32"
                ];
                persistentKeepalive = 5;
              };
            };
          };

          firewall.rules = [
            {
              from = "vpn";
              to = "trusted";
            }
            {
              from = "trusted";
              to = "remote";
            }
            {
              from = "vpn";
              to = "remote";
            }
            {
              from = "trusted";
              to = "freifunk";
            }
          ];
        };

        # See basic.nix: relax Unbound for the made-up test domain only.
        services.unbound.settings.server = {
          local-zone = [ "test. transparent" ];
          domain-insecure = [ "example.test" ];
        };
      };

    laptop = segmentHost {
      address = "203.0.113.50";
      wgName = "wg-rw";
      wgKey = key.laptop;
      wgAddresses = [
        "10.100.0.2/24"
        "fd42:100::2/64"
      ];
      peer = {
        PublicKey = key.router-rw.pub;
        Endpoint = "203.0.113.100:51820";
        PersistentKeepalive = 5;
        AllowedIPs = [
          "10.100.0.0/24"
          "192.168.10.0/24"
          "192.168.50.0/24"
          "fd42:100::/64"
        ];
      };
    };

    siteb = {
      imports = [
        (segmentHost {
          address = "203.0.113.60";
          wgName = "wg-s2s";
          wgKey = key.siteb;
          listenPort = 51821;
          peer = {
            PublicKey = key.router-s2s.pub;
            Endpoint = "203.0.113.100:51821";
            PersistentKeepalive = 5;
            AllowedIPs = [
              "192.168.10.0/24"
              "10.100.0.0/24"
            ];
          };
        })
        (dummy "lan0" [ "192.168.50.1/24" ])
        echo
      ];
    };

    main =
      { lib, ... }:
      {
        virtualisation.vlans = [ 2 ];
        networking.useNetworkd = true;
        networking.useDHCP = false;
        networking.firewall.enable = false;
        networking.interfaces.eth1.ipv4.addresses = lib.mkForce [ ];
        systemd.network.networks."20-lan" = {
          matchConfig.Name = "eth1";
          networkConfig.DHCP = "ipv4";
        };
      };
  };

  testScript = ''
    start_all()

    isp.wait_for_unit("dnsmasq.service")
    isp.wait_for_unit("nginx.service")
    siteb.wait_for_unit("nginx.service")

    router.wait_for_unit("network-online.target")
    router.wait_until_succeeds("ip -4 addr show dev eth1 | grep -q 203.0.113.100")
    router.wait_for_unit("nftables.service")
    router.wait_for_unit("kea-dhcp4-server.service")
    router.wait_for_unit("unbound.service")
    main.wait_until_succeeds("ip -4 addr show dev eth1 | grep -q 192.168.10.100")

    with subtest("all three tunnels complete a handshake"):
        for dev in ("wg-rw", "wg-s2s", "wg-ff"):
            router.wait_until_succeeds(f"wg show {dev} latest-handshakes | awk '$2 > 0' | grep -q .")

    with subtest("peer endpoint names are looked up again, but only for peers without a handshake"):
        router.succeed("systemctl is-enabled router-wireguard-reresolve.timer")
        router.succeed("wg show wg-ghost endpoints | grep -q 203.0.113.70:51830")
        # The names now have other addresses.
        router.succeed("rm /etc/hosts; printf '127.0.0.1 localhost\\n203.0.113.71 ghost.test\\n203.0.113.99 siteb.test\\n' > /etc/hosts")
        router.succeed("systemctl start router-wireguard-reresolve.service")
        router.succeed("wg show wg-ghost endpoints | grep -q 203.0.113.71:51830")
        router.succeed("journalctl -u router-wireguard-reresolve --no-pager | grep -q 'endpoint 203.0.113.70:51830 -> 203.0.113.71:51830'")
        # The working tunnel keeps its endpoint (it has a fresh handshake).
        router.succeed("wg show wg-s2s endpoints | grep -q 203.0.113.60:51821")
        # A name that does not resolve leaves the endpoint as it was.
        router.succeed("printf '127.0.0.1 localhost\\n' > /etc/hosts")
        router.succeed("systemctl start router-wireguard-reresolve.service")
        router.succeed("wg show wg-ghost endpoints | grep -q 203.0.113.71:51830")
        router.succeed("journalctl -u router-wireguard-reresolve --no-pager | grep -q 'ghost.test:51830'")
        router.succeed("rm /etc/hosts; ln -s /etc/static/hosts /etc/hosts")

    with subtest("road warrior: dual-stack tunnel to the router, DNS and SSH over it"):
        laptop.succeed("ping -c 1 -W 2 10.100.0.1")
        laptop.succeed("ping -6 -c 1 -W 2 fd42:100::1")
        laptop.succeed("dig +short +time=2 +tries=1 @10.100.0.1 example.test | grep -qx 203.0.113.1")
        laptop.succeed("dig +short +time=2 +tries=1 @fd42:100::1 example.test | grep -qx 203.0.113.1")
        laptop.succeed("nc -z -w 2 10.100.0.1 22")

    with subtest("road warrior reaches the LAN, but not the other way round"):
        laptop.succeed("ping -c 1 -W 2 192.168.10.100")
        # Positive control: the LAN host can reach the router's own LAN address.
        main.succeed("ping -c 1 -W 2 192.168.10.1")
        main.fail("ping -c 1 -W 2 10.100.0.2")

    with subtest("site-to-site: routed without NAT and without a transfer network"):
        assert router.succeed("ip -4 addr show dev wg-s2s | grep -c inet || true").strip() == "0"
        main.succeed("ping -c 1 -W 2 192.168.50.1")
        out = main.succeed("curl -sf --max-time 5 http://192.168.50.1/")
        assert out.strip() == "192.168.10.100", f"expected the LAN address to be preserved, got {out!r}"
        out = laptop.succeed("curl -sf --max-time 5 http://192.168.50.1/")
        assert out.strip() == "10.100.0.2", f"expected the tunnel address to be preserved, got {out!r}"

    with subtest("site-to-site: the remote site cannot initiate towards the LAN"):
        # Positive control: it can reach the router itself over the tunnel.
        siteb.succeed("ping -c 1 -W 2 -I 192.168.50.1 192.168.10.1")
        siteb.fail("ping -c 1 -W 2 -I 192.168.50.1 192.168.10.100")

    with subtest("site-to-site: router-originated traffic uses the configured source"):
        # Without it the kernel would pick an address by itself (the tunnel has none).
        router.succeed("ip route get 192.168.50.1 | grep -q 'src 192.168.10.1'")
        router.succeed("ping -c 1 -W 2 192.168.50.1")

    with subtest("site-to-site: the remote site can use the router's resolver"):
        siteb.succeed("dig +short +time=2 +tries=1 -b 192.168.50.1 @192.168.10.1 example.test | grep -qx 203.0.113.1")

    with subtest("provider tunnel: only the listed destinations, masqueraded"):
        out = main.succeed("curl -sf --max-time 5 http://198.18.0.1/")
        assert out.strip() == "10.200.0.2", f"expected masquerade to 10.200.0.2, got {out!r}"
        # Everything else still leaves via the WAN.
        out = main.succeed("curl -sf --max-time 5 http://203.0.113.1/")
        assert out.strip() == "203.0.113.100", f"expected the WAN address, got {out!r}"
        # IPv6 over the dual-stack transfer net (router-originated).
        router.succeed("ping -6 -c 1 -W 2 fd42:200::1")
        out = router.succeed("curl -sf -g --max-time 5 'http://[fd00:ff::1]/'")
        assert out.strip() == "fd42:200::2", f"expected the IPv6 transfer address, got {out!r}"

    with subtest("zones apply to tunnels: the road warrior may not use the provider tunnel"):
        laptop.fail("curl -sf --max-time 3 http://198.18.0.1/")

    with subtest("kill switch: no leak to the WAN when the provider tunnel is down"):
        router.succeed("ip link set dev wg-ff down")
        # The route is gone, so the destination would now follow the default route.
        router.succeed("ip route get 198.18.0.1 | grep -q 'dev eth1'")
        main.fail("curl -sf --max-time 3 http://198.18.0.1/")
        # Ordinary destinations are not affected.
        main.succeed("curl -sf --max-time 5 http://203.0.113.1/")
  '';
}