dotfiles/nixos/modules/microvm-extras-host.nix

{ config, lib, notnft, ... }:
let
  inherit (lib)
    mapAttrsToList
    mkOption
    hasAttr
    types
    traceVal
    flip
    mapAttrs'
    mapAttrs
    nameValuePair;

  # a = [
  #   [ (is.eq ip.protocol (f: with f; set [ tcp ])) (is.eq ip.daddr "10.80.1.2") (is.eq th.dport "22") accept ]
  # ];

  cfg = config.microvm;

  protocolEnumToNft = f: proto:
    f.${proto};

  tcpUdpServiceOptions.options = {
    hostName = mkOption {
      type = types.str;
    };

    port = mkOption {
      type = types.port;
    };

    protocol = mkOption {
      type = types.listOf (types.enum [ "tcp" "udp" ]);
    };
  };

  httpServiceOptions.options = {
    hostName = mkOption {
      type = types.str;
    };

    port = mkOption {
      type = types.port;
    };
  };

  icmpServiceOptions.options = {
    hostName = mkOption {
      type = types.str;
    };
  };

  tcpUdpConnectionOptions.options = {
    target = mkOption {
      type = types.str;
    };
  };

  icmpConnectionOptions.options = {
    target = mkOption {
      type = types.str;
    };
  };

  httpConnectionOptions.options = {
    target = mkOption {
      type = types.str;
    };
  };

  lookupService = name: type: context:
    if hasAttr name cfg.services.${type} then
      cfg.services.${type}.${name}
    else
      throw "Unknown ${type} service ${name} at ${context}";

  lookupIds = hostName: context:
    if hasAttr hostName subConfigurations then
      {
        inherit (subConfigurations.${hostName}.config.config.microvm)
          groupId
          taskId;
      }
    else
      throw "Unknown hostName ${hostName} at ${context}";

  subConfigurations = cfg.vms // config.containers;
in
{
  options.microvm = {
    services = {
      tcpUdp = mkOption {
        type = with types; types.attrsOf (submodule tcpUdpServiceOptions);
      };

      icmp = mkOption {
        type = with types; types.attrsOf (submodule icmpServiceOptions);
      };

      http = mkOption {
        type = with types; types.attrsOf (submodule httpServiceOptions);
      };
    };

    connections = {
      tcpUdp = mkOption {
        type = with types;
          listOf (submodule tcpUdpConnectionOptions);
        default = [];
      };

      icmp = mkOption {
        type = with types;
          listOf (submodule icmpConnectionOptions);
        default = [];
      };

      http = mkOption {
        type = with types;
          listOf (submodule httpConnectionOptions);
        default = [];
      };
    };
  };

  config.microvm.services.tcpUdp = flip mapAttrs' cfg.services.http
    (n: v:
      nameValuePair
        (n + "@http")
        {
          inherit (v)
            hostName
            port;
          protocol = [ "tcp" ];
        }
    );

  config.microvm.connections.tcpUdp = flip map cfg.connections.http
    (v:
      {
        target = v.target + "@http";
      }
    );

  config.networking.notnft.rules =
    with notnft.dsl; with payload; ruleset {
      bridge-t = add table { family = f: f.bridge; } {
        output-body = lib.foldl (acc: x: acc x) (add chain) ((flip mapAttrsToList subConfigurations
          (n: v:
            let
              microvmConfig = v.config.config.microvm;
              tcpUdpRules =
                flip map microvmConfig.connections.tcpUdp (connection:
                  let
                    service = lookupService connection.target "tcpUdp" n;
                    ids = lookupIds service.hostName n;
                  in
                  [
                    (is.eq meta.oifname "mvm-${microvmConfig.hostName}")
                    (is.eq ip.protocol (f: with f; set (map (protocolEnumToNft f) service.protocol)))
                    (is.eq ip.saddr "10.80.${toString microvmConfig.groupId}.${toString microvmConfig.taskId}")
                    (is.eq ip.daddr "10.80.${toString ids.groupId}.${toString ids.taskId}")
                    (is.eq th.dport service.port)
                    accept
                  ]);
              icmpRules =
                flip map microvmConfig.connections.icmp (connection:
                  let
                    service = lookupService connection.target "icmp" n;
                    ids = lookupIds service.hostName n;
                  in
                  [
                    (is.eq meta.oifname "mvm-${microvmConfig.hostName}")
                    (is.eq ip.protocol (f: with f; icmp))
                    (is.eq ip.saddr "10.80.${toString microvmConfig.groupId}.${toString microvmConfig.taskId}")
                    (is.eq ip.daddr "10.80.${toString ids.groupId}.${toString ids.taskId}")
                    accept
                  ]);
            in
              tcpUdpRules ++ icmpRules
          )) ++ (flip map cfg.connections.icmp (connection:
            let
              service = lookupService connection.target "icmp" "host";
              ids = lookupIds service.hostName "host";
            in
              [
                (is.eq meta.oifname "mvm-${service.hostName}")
                (is.eq ip.protocol (f: with f; icmp))
                (is.eq ip.saddr "10.80.${toString ids.groupId}.1")
                (is.eq ip.daddr "10.80.${toString ids.groupId}.${toString ids.taskId}")
                accept
              ]
          )) ++ (flip map cfg.connections.tcpUdp (connection:
            let
              service = lookupService connection.target "tcpUdp" "host";
              ids = lookupIds service.hostName "host";
            in
              [
                (is.eq meta.oifname "mvm-${service.hostName}")
                (is.eq ip.protocol (f: with f; set (map (protocolEnumToNft f) service.protocol)))
                (is.eq ip.saddr "10.80.${toString ids.groupId}.1")
                (is.eq ip.daddr "10.80.${toString ids.groupId}.${toString ids.taskId}")
                (is.eq th.dport service.port)
                accept
              ]
          )));
      };
    };
}