Refactoring + working compiler for ip6? (s|d)addr

compile.ml

@@ -12,6 +12,16 @@ module Prefix = struct
     | Not a, Not b -> -compare a b
     | Not _, _ -> -1
     | _, Not _ -> 1
+
+  let rec to_ipv4_list negate = function
+    | Ipv4 ipv4 -> if negate then [] else [ ipv4 ]
+    | Ipv6 _ -> []
+    | Not prefix -> to_ipv4_list (not negate) prefix
+
+  let rec to_ipv6_list negate = function
+    | Ipv4 _ -> []
+    | Ipv6 ipv6 -> if negate then [] else [ ipv6 ]
+    | Not prefix -> to_ipv6_list (not negate) prefix
 end
 
 module PrefixSet = struct
@@ -31,19 +41,19 @@ module PrefixSet = struct
     List.fold_left (fun acc addrs -> union (of_addrs zone addrs) acc) empty
 end
 
-let rec zone_deps =
+module Zones = struct
-  let open Config.Zone in
+  open Config.Zone
-  function
+
+  let dependencies zone =
+    let rec aux = function
       | Ipv4 _ | Ipv6 _ -> []
       | Name name -> [ name ]
-  | List list -> List.flatten (List.map zone_deps list)
+      | List list -> List.flatten (List.map aux list)
-  | Not not -> zone_deps not
+      | Not not -> aux not
+    in
+    List.map (fun (k, v) -> (k, aux v)) zone
 
-let zones_deps zone = List.map (fun (k, v) -> (k, zone_deps v)) zone
+  let rec compile_zone assoc = function
-
-let rec compile_zone assoc =
-  let open Config.Zone in
-  function
     | Ipv4 ipv4 -> PrefixSet.singleton (Prefix.Ipv4 ipv4)
     | Ipv6 ipv6 -> PrefixSet.singleton (Prefix.Ipv6 ipv6)
     | Name name -> List.assoc name assoc
@@ -51,10 +61,11 @@ let rec compile_zone assoc =
         List.fold_left
           (fun acc zone -> PrefixSet.union (compile_zone assoc zone) acc)
           PrefixSet.empty list
-  | Not zone -> PrefixSet.map (fun p -> Prefix.Not p) (compile_zone assoc zone)
+    | Not zone ->
+        PrefixSet.map (fun p -> Prefix.Not p) (compile_zone assoc zone)
 
-let compile_zones zones =
+  let compile zones =
-  match Tsort.sort (zones_deps zones) with
+    match Tsort.sort (dependencies zones) with
     | Tsort.Sorted sorted ->
         List.fold_right
           (fun name acc ->
@@ -63,43 +74,56 @@ let compile_zones zones =
             (name, compiled) :: acc)
           sorted []
     | _ -> failwith "cyclic dependency in zones definitions"
+end
 
-let find_ipv4 zones negate addrs_list =
+module Rules = struct
-  let open Prefix in
+  open Nftables
-  let prefixes = PrefixSet.of_addrs_list zones addrs_list in
+  open Config.Rule
-  let rec filter_prefix negate prefix acc =
-    match prefix with
-    | Ipv4 ipv4 -> if negate then acc else ipv4 :: acc
-    | Ipv6 _ -> acc
-    | Not prefix -> filter_prefix (not negate) prefix acc
-  in
-  PrefixSet.fold (filter_prefix negate) prefixes []
 
-let compile_addrs_ipv4 zones field addrs =
+  (* Bon, ce module n'est vraiment pas très joli… *)
-  let open Nftables in
+
-  let equal = find_ipv4 zones false addrs |> List.map (fun p -> Expr.Ipv4 p) in
+  let compile_addrs_list getter expr negate zones addrs_list =
-  let not_equal =
+    Expr.Set
-    find_ipv4 zones true addrs |> List.map (fun p -> Expr.Ipv4 p)
+      (PrefixSet.fold
-  in
+         (fun prefix acc -> getter negate prefix @ acc)
-  (* TODO: handle empty sets *)
+         (PrefixSet.of_addrs_list zones addrs_list)
+         []
+      |> List.map expr)
+
+  let compile_match_addrs getter expr field zones addrs_list =
     [
-    Stmt.Match (Match.Eq, Expr.Payload (Payload.Ipv4 field), Expr.Set equal);
       Stmt.Match
-      (Match.NotEq, Expr.Payload (Payload.Ipv4 field), Expr.Set not_equal);
+        ( Match.Eq,
+          Expr.Payload field,
+          compile_addrs_list getter expr false zones addrs_list );
+      Stmt.Match
+        ( Match.NotEq,
+          Expr.Payload field,
+          compile_addrs_list getter expr true zones addrs_list );
     ]
 
-let compile_rule zones rule =
+  let compile_match_ipv4 field =
-  let open Config.Rule in
+    compile_match_addrs Prefix.to_ipv4_list Expr.ipv4 (Payload.Ipv4 field)
-  let open Nftables in
+
-  let ipv4_src = compile_addrs_ipv4 zones Payload.Ipv4.Saddr rule.src in
+  let compile_match_ipv6 field =
-  let ipv4_dest = compile_addrs_ipv4 zones Payload.Ipv4.Daddr rule.dest in
+    compile_match_addrs Prefix.to_ipv6_list Expr.ipv6 (Payload.Ipv6 field)
-  List.flatten [ ipv4_src; ipv4_dest ]
+
+  let compile_rule zones { src; dest; _ } =
+    let ipv4_src = compile_match_ipv4 Payload.Ipv4.Saddr zones src in
+    let ipv4_dest = compile_match_ipv4 Payload.Ipv4.Daddr zones dest in
+    let ipv6_src = compile_match_ipv6 Payload.Ipv6.Saddr zones src in
+    let ipv6_dest = compile_match_ipv6 Payload.Ipv6.Daddr zones dest in
+    let verdict = [ Stmt.Verdict Verdict.Accept ] in
+    [ ipv4_src @ ipv4_dest @ verdict; ipv6_src @ ipv6_dest @ verdict ]
+
+  let compile zones rules = List.flatten (List.map (compile_rule zones) rules)
+end
 
 let compile config =
   let open Nftables in
   let open Config in
-  let zones = compile_zones config.zones in
+  let zones = Zones.compile config.zones in
-  let exprs = List.map (compile_rule zones) config.rules in
+  let exprs = Rules.compile zones config.rules in
   let family = Family.Inet in
   let table = "filter" in
   let chain = "forward" in