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;
in {
  options.microvm = {
    services = {
      tcpUdp = mkOption {
        type = with types; types.attrsOf (submodule tcpUdpServiceOptions);
        default = {};
      };

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

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

    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
            ]
          )));
      };
    };
}