nixos/modules/wireguard.nix

{
  config,
  lib,
  pkgs,
  ...
}: {
  options.elss.wireguard = with lib; {
    enable = mkEnableOption "wireguard overlay network";

    interfaces = mkOption {
      default = {};
      type = types.attrsOf (types.submodule {
        options = {
          servers = mkOption {
            type = types.attrsOf (types.submodule {
              options = {
                localIp = mkOption {
                  type = types.str;
                  description = "local IP part for the interfaces";
                };

                extraIps = mkOption {
                  type = types.listOf types.str;
                  default = [];
                  description = "extra IPs to add to allowedIPs";
                };

                listenPort = mkOption {
                  type = types.port;
                  description = "Port to listen on";
                  default = 51820;
                };

                publicKey = mkOption {
                  type = types.str;
                  description = "Wireguard public key for this peer";
                };

                endpoint = mkOption {
                  type = types.str;
                  description = "Wireguard endpoint for this peer.";
                };
              };
            });
          };

          peers = mkOption {
            type = types.attrsOf (types.submodule {
              options = {
                localIp = mkOption {
                  type = types.str;
                  description = "local IP part for the interfaces";
                };

                listenPort = mkOption {
                  type = types.port;
                  description = "Port to listen on";
                  default = 51820;
                };

                publicKey = mkOption {
                  type = types.str;
                  description = "Wireguard public key for this peer";
                };

                additionalAllowedIps = mkOption {
                  type = types.listOf types.str;
                  description = "Additional IPs to add to allowedIPs ";
                  default = [];
                };
              };
            });
          };

          prefixes = {
            ipv4 = mkOption {
              type = types.listOf types.str;
              description = "IPv4 prefixes to use for wireguard addressing";
            };

            ipv6 = {
              ula = mkOption {
                type = types.listOf types.str;
                description = "IPv6 prefixes to use for ULA wireguard addressing";
              };

              gua = mkOption {
                type = types.listOf types.str;
                description = "IPv6 prefixes to use for GUA wireguard addressing";
              };
            };

            serial = mkOption {
              type = types.str;
              description = "serial for the generated DNS zone";
            };
          };
        };
      });
    };
  };

  config = let
    cfg = config.elss;
    hostName = config.system.name;
    secretsFile =
      ../machines
      + builtins.toPath "/${hostName}/secrets/wireguard.yaml";
    takeNonEmpty = lib.filter (interface: interface != "");
    testInterface = predicate:
      lib.mapAttrsToList
      (interface: value:
        if (predicate interface value)
        then interface
        else "")
      cfg.wireguard.interfaces;
    onlyInterfaces = predicate: takeNonEmpty (testInterface predicate);
    peerInterfaces =
      onlyInterfaces (interface: value: builtins.hasAttr hostName value.peers);
    serverInterfaces =
      onlyInterfaces
      (interface: value: builtins.hasAttr hostName value.servers);
    interfaces = serverInterfaces ++ peerInterfaces;

    mkAddresses = prefixes: localIp:
      (map (prefix: "${prefix}.${localIp}/32") prefixes.ipv4)
      ++ (map (prefix: "${prefix}::${localIp}/128") prefixes.ipv6.ula)
      ++ (map (prefix: "${prefix}::${localIp}/128") prefixes.ipv6.gua);

    mkServerAddresses = prefixes: serverIp:
      (map (prefix: "${prefix}.${serverIp}") prefixes.ipv4)
      ++ (map (prefix: "${prefix}::${serverIp}") prefixes.ipv6.ula)
      ++ (map (prefix: "${prefix}::${serverIp}") prefixes.ipv6.gua);

    mkInterfaceName = interface: "wg-${interface}";

    mkServerPeer = prefixes: peer: {
      allowedIPs = mkAddresses prefixes peer.localIp;
      inherit (peer) publicKey;
    };

    mkPeerPeer = prefixes: peers: peer: {
      allowedIPs =
        (mkAddresses prefixes peer.localIp)
        ++ (lib.concatMap (mkAddresses prefixes) peer.extraIps)
        ++ (
          if lib.hasAttr hostName peers
          then peers.${hostName}.additionalAllowedIps
          else []
        );
      persistentKeepalive = 25;
      inherit (peer) publicKey endpoint;
    };

    mkPostSetup = name: prefixes: servers: let
      ifName = mkInterfaceName name;
      serverIps = name: server: mkServerAddresses prefixes server.localIp;
      dnsServers = lib.concatLists (lib.mapAttrsToList serverIps servers);
    in
      lib.concatStrings ([
          ''
            ${pkgs.systemd}/bin/resolvectl domain ${ifName} ${name}.${config.elss.dns.wgZone}
            ${pkgs.systemd}/bin/resolvectl default-route ${ifName} true
          ''
        ]
        ++ (map
          (ip: ''
            ${pkgs.systemd}/bin/resolvectl dns ${ifName} ${ip}
          '')
          dnsServers));

    mkInterfaceConfig = hostName: interface: value: let
      isServer = builtins.hasAttr hostName value.servers;
      isPeer = builtins.hasAttr hostName value.peers;
      myConfig =
        if isServer
        then value.servers."${hostName}"
        else value.peers."${hostName}";
    in
      assert lib.asserts.assertMsg
      ((isServer || isPeer) && !(isServer && isPeer))
      "host must be either server or peer";
        lib.nameValuePair (mkInterfaceName interface) ({
            privateKeyFile = config.sops.secrets."wireguard-${interface}".path;
            ips = mkAddresses value.prefixes myConfig.localIp;
            inherit (myConfig) listenPort;
          }
          // (
            if isServer
            then {
              peers = lib.mapAttrsToList (_: mkServerPeer value.prefixes) value.peers;
            }
            else if isPeer
            then {
              peers = lib.mapAttrsToList (_: mkPeerPeer value.prefixes value.peers) value.servers;
              postSetup = mkPostSetup interface value.prefixes value.servers;
            }
            else {}
          ));

    mkInterfaceSecret = interface: {
      "wireguard-${interface}" = {sopsFile = secretsFile;};
    };

    mkListenPorts = hostName: interface: value:
      if builtins.hasAttr hostName value.servers
      then value.servers."${hostName}".listenPort
      else if builtins.hasAttr hostName value.peers
      then value.peers."${hostName}".listenPort
      else -1;

    mkSysctl = hostName: interface: [
      {
        name = "net.ipv4.conf.${mkInterfaceName interface}.forwarding";
        value = "1";
      }
      {
        name = "net.ipv6.conf.${mkInterfaceName interface}.forwarding";
        value = "1";
      }
      {
        name = "net.ipv6.conf.all.forwarding";
        value = "1";
      }
    ];
  in
    lib.mkIf cfg.wireguard.enable {
      networking = {
        wireguard.interfaces =
          lib.mapAttrs' (mkInterfaceConfig hostName) cfg.wireguard.interfaces;
        firewall = {
          # allowedUDPPorts = lib.filter (port: port > 0)
          #   (lib.mapAttrsToList (mkListenPorts hostName) cfg.wireguard.interfaces);
          allowedUDPPorts = lib.filter (port: port > 0) (map
            (interface:
              lib.attrByPath [interface "servers" hostName "listenPort"] (-1)
              cfg.wireguard.interfaces)
            serverInterfaces);
          trustedInterfaces = map mkInterfaceName interfaces;
        };
        interfaces = lib.listToAttrs (map
          (interface: {
            name = mkInterfaceName interface;
            value = {mtu = 1300;};
          })
          interfaces);
      };

      services.unbound.settings.server.interface = map mkInterfaceName serverInterfaces;
      services.resolved.enable = lib.mkDefault true;
      systemd.services = lib.listToAttrs (map
        (interface: {
          name = "wireguard-${mkInterfaceName interface}";
          value = {serviceConfig.Restart = "on-failure";};
        })
        interfaces);

      boot.kernel.sysctl =
        builtins.listToAttrs (lib.concatMap (mkSysctl hostName) serverInterfaces);

      sops.secrets = lib.mkMerge (map mkInterfaceSecret interfaces);
    };
}