From 570c40b434b4f6f41f89053c8fa2b5878c0f72e9 Mon Sep 17 00:00:00 2001
From: Jeltz <jeltz@auro.re>
Date: Tue, 30 Aug 2022 01:58:52 +0200
Subject: [PATCH] Add an implementation for zone_of_json

---
 config.ml   | 13 +++++++++++--
 nftables.ml | 26 +++++++++++++++++++++++---
 2 files changed, 34 insertions(+), 5 deletions(-)

diff --git a/config.ml b/config.ml
index 2d0cf3b..08dc90b 100644
--- a/config.ml
+++ b/config.ml
@@ -17,8 +17,17 @@ type l4_rule =
 type rule = { src : addrs list; dest : addrs list; l4 : l4_rule }
 type config = { zones : (string * zone) list; rules : rule list }
 
-(* TODO *)
-let zone_of_json _ = Zone "xxx"
+let rec zone_of_json json =
+  match json with
+  | `String s -> (
+      let zone = Zone s in
+      let zone =
+        match V6.Prefix.of_string s with Ok p -> ZoneIpv6 p | Error _ -> zone
+      in
+      match V4.Prefix.of_string s with Ok p -> ZoneIpv4 p | Error _ -> zone)
+  | `Assoc [ ("exclude", e) ] -> ZoneExclude (zone_of_json e)
+  | `List l -> ZoneList (List.map zone_of_json l)
+  | _ -> assert false
 
 let addrs_of_json json =
   let open Yojson.Basic.Util in
diff --git a/nftables.ml b/nftables.ml
index abfbf07..b45ce83 100644
--- a/nftables.ml
+++ b/nftables.ml
@@ -3,7 +3,14 @@ open Ipaddr
 
 type _ udp = UdpDport : int udp | UdpSport : int udp
 type _ tcp = TcpDport : int tcp | TcpSport : int tcp
-type _ payload = Udp : 'a udp -> 'a payload | Tcp : 'a tcp -> 'a payload
+type _ ipv4 = Ipv4Saddr : V4.Prefix.t ipv4 | Ipv4Daddr : V4.Prefix.t ipv4
+type _ ipv6 = Ipv6Saddr : V6.Prefix.t ipv6 | Ipv6Daddr : V6.Prefix.t ipv6
+
+type _ payload =
+  | Udp : 'a udp -> 'a payload
+  | Tcp : 'a tcp -> 'a payload
+  | Ipv4 : 'a ipv4 -> 'a payload
+  | Ipv6 : 'a ipv6 -> 'a payload
 
 type _ expr =
   | String : string -> string expr
@@ -39,7 +46,7 @@ type _ stmt =
   | NoTrack : unit stmt
   | Log : { prefix : string option; group : int option } -> unit stmt
 
-type family = Ip6 | Ip4 | Inet
+type family = FamilyIpv6 | FamilyIpv4 | FamilyInet
 type table = { family : family; table_name : string }
 type chain = { family : family; table : string; chain_name : string }
 
@@ -66,11 +73,21 @@ let string_of_tcp : type a. a tcp -> string = function
   | TcpSport -> "sport"
   | TcpDport -> "dport"
 
+let string_of_ipv4 : type a. a ipv4 -> string = function
+  | Ipv4Saddr -> "saddr"
+  | Ipv4Daddr -> "daddr"
+
+let string_of_ipv6 : type a. a ipv6 -> string = function
+  | Ipv6Saddr -> "saddr"
+  | Ipv6Daddr -> "daddr"
+
 let assoc_one key value = `Assoc [ (key, value) ]
 
 let json_of_payload (type a) (payload : a payload) =
   let protocol, field =
     match payload with
+    | Ipv4 ipv4 -> ("ip", string_of_ipv4 ipv4)
+    | Ipv6 ipv6 -> ("ip6", string_of_ipv6 ipv6)
     | Udp udp -> ("udp", string_of_udp udp)
     | Tcp tcp -> ("tcp", string_of_tcp tcp)
   in
@@ -139,7 +156,10 @@ let json_of_stmt : type a. a stmt -> Yojson.Basic.t = function
       in
       assoc_one "log" (`Assoc (deoptionalise elems))
 
-let string_of_family = function Ip6 -> "ip6" | Ip4 -> "ip4" | Inet -> "inet"
+let string_of_family = function
+  | FamilyIpv4 -> "ip"
+  | FamilyIpv6 -> "ip6"
+  | FamilyInet -> "inet"
 
 let json_of_table { family; table_name } =
   assoc_one "table"