nixos/nat: Leverage firewall module

nixos/modules/services/networking/nat.nix

@@ -12,6 +12,41 @@ let
 
   dest = if cfg.externalIP == null then "-j MASQUERADE" else "-j SNAT --to-source ${cfg.externalIP}";
 
+  flushNat = ''
+    iptables -w -t nat -F PREROUTING
+    iptables -w -t nat -F POSTROUTING
+    iptables -w -t nat -X
+  '';
+
+  setupNat = ''
+    # We can't match on incoming interface in POSTROUTING, so
+    # mark packets coming from the external interfaces.
+    ${concatMapStrings (iface: ''
+      iptables -w -t nat -A PREROUTING \
+        -i '${iface}' -j MARK --set-mark 1
+    '') cfg.internalInterfaces}
+
+    # NAT the marked packets.
+    ${optionalString (cfg.internalInterfaces != []) ''
+      iptables -w -t nat -A POSTROUTING -m mark --mark 1 \
+        -o ${cfg.externalInterface} ${dest}
+    ''}
+
+    # NAT packets coming from the internal IPs.
+    ${concatMapStrings (range: ''
+      iptables -w -t nat -A POSTROUTING \
+        -s '${range}' -o ${cfg.externalInterface} ${dest}
+    '') cfg.internalIPs}
+
+    # NAT from external ports to internal ports.
+    ${concatMapStrings (fwd: ''
+      iptables -w -t nat -A PREROUTING \
+        -i ${cfg.externalInterface} -p tcp \
+        --dport ${builtins.toString fwd.sourcePort} \
+        -j DNAT --to-destination ${fwd.destination}
+    '') cfg.forwardPorts}
+  '';
+
 in
 
 {
@@ -109,57 +144,34 @@ in
 
     environment.systemPackages = [ pkgs.iptables ];
 
-    boot.kernelModules = [ "nf_nat_ftp" ];
-
-    jobs.nat =
-      { description = "Network Address Translation";
-
-        startOn = "started network-interfaces";
-
-        path = [ pkgs.iptables ];
-
-        preStart =
-          ''
-            iptables -w -t nat -F PREROUTING
-            iptables -w -t nat -F POSTROUTING
-            iptables -w -t nat -X
-
-            # We can't match on incoming interface in POSTROUTING, so
-            # mark packets coming from the external interfaces.
-            ${concatMapStrings (iface: ''
-              iptables -w -t nat -A PREROUTING \
-                -i '${iface}' -j MARK --set-mark 1
-            '') cfg.internalInterfaces}
-
-            # NAT the marked packets.
-            ${optionalString (cfg.internalInterfaces != []) ''
-              iptables -w -t nat -A POSTROUTING -m mark --mark 1 \
-                -o ${cfg.externalInterface} ${dest}
-            ''}
-
-            # NAT packets coming from the internal IPs.
-            ${concatMapStrings (range: ''
-              iptables -w -t nat -A POSTROUTING \
-                -s '${range}' -o ${cfg.externalInterface} ${dest}
-            '') cfg.internalIPs}
-
-            # NAT from external ports to internal ports.
-            ${concatMapStrings (fwd: ''
-              iptables -w -t nat -A PREROUTING \
-                -i ${cfg.externalInterface} -p tcp \
-                --dport ${builtins.toString fwd.sourcePort} \
-                -j DNAT --to-destination ${fwd.destination}
-            '') cfg.forwardPorts}
-
-            echo 1 > /proc/sys/net/ipv4/ip_forward
-          '';
-
-        postStop =
-          ''
-            iptables -w -t nat -F PREROUTING
-            iptables -w -t nat -F POSTROUTING
-            iptables -w -t nat -X
-          '';
+    boot = {
+      kernelModules = [ "nf_nat_ftp" ];
+      kernel.sysctl = mkOverride 99 {
+        "net.ipv4.conf.all.forwarding" = true;
+        "net.ipv4.conf.default.forwarding" = true;
       };
+    };
+
+    networking.firewall = mkIf config.networking.firewall.enable {
+      extraCommands = mkMerge [ (mkBefore flushNat) setupNat ];
+      extraStopCommands = flushNat;
+    };
+
+    systemd.services =  mkIf (!config.networking.firewall.enable) { nat = {
+      description = "Network Address Translation";
+      wantedBy = [ "network.target" ];
+      after = [ "network-interfaces.target" "systemd-modules-load.service" ];
+      path = [ pkgs.iptables ];
+      unitConfig.ConditionCapability = "CAP_NET_ADMIN";
+
+      serviceConfig = {
+        Type = "oneshot";
+        RemainAfterExit = true;
+      };
+
+      script = flushNat + setupNat;
+
+      postStop = flushNat;
+    }; };
   };
 }