Create modules in nftables.ml

nftables.ml

@@ -1,146 +1,144 @@
 open Utils
 open Ipaddr
 
-type _ udp = UdpDport : int udp | UdpSport : int udp
-type _ tcp = TcpDport : int tcp | TcpSport : int tcp
-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
-  | Number : int -> int expr
-  | Boolean : bool -> int expr
-  | Ipv4 : V4.Prefix.t -> V4.Prefix.t expr
-  | Ipv6 : V6.Prefix.t -> V6.Prefix.t expr
-  | List : 'a expr list -> 'a expr
-  | Set : 'a expr list -> 'a expr
-  | Range : 'a expr * 'a expr -> 'a expr
-  | Payload : 'a payload -> 'a expr
-
-type (_, _) match_op = Eq : ('a, 'a) match_op | NotEq : ('a, 'a) match_op
-type counter = Named of string | Anon of { packets : int; bytes : int }
-
-type verdict =
-  | Accept
-  | Drop
-  | Continue
-  | Return
-  | Jump of string
-  | Goto of string
-
-type _ stmt =
-  | Match : (('a, 'b) match_op * 'a expr * 'b expr) -> unit stmt
-  | Counter : counter -> unit stmt
-  | Verdict : verdict -> unit stmt
-  | NoTrack : unit stmt
-  | Log : { prefix : string option; group : int option } -> unit stmt
-
-type family = FamilyIpv6 | FamilyIpv4 | FamilyInet
-type table = { family : family; table_name : string }
-type chain = { family : family; table : string; chain_name : string }
-
-type rule = {
-  family : family;
-  table : string;
-  chain : string;
-  expr : unit stmt list;
-}
-
-type command =
-  | FlushRuleset
-  | AddTable of table
-  | FlushTable of table
-  | AddChain of chain
-  | FlushChain of chain
-  | AddRule of rule
-
-let string_of_udp : type a. a udp -> string = function
-  | UdpSport -> "sport"
-  | UdpDport -> "dport"
-
-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) =
+module Payload = struct
+  module Udp = struct
+    type _ t = Sport : int t | Dport : int t
+
+    let to_string : type a. a t -> string = function
+      | Sport -> "sport"
+      | Dport -> "dport"
+  end
+
+  module Tcp = struct
+    type _ t = Sport : int t | Dport : int t
+
+    let to_string : type a. a t -> string = function
+      | Sport -> "sport"
+      | Dport -> "dport"
+  end
+
+  module Ipv4 = struct
+    type _ t = Saddr : V4.Prefix.t t | Daddr : V4.Prefix.t t
+
+    let to_string : type a. a t -> string = function
+      | Saddr -> "saddr"
+      | Daddr -> "saddr"
+  end
+
+  module Ipv6 = struct
+    type _ t = Saddr : V6.Prefix.t t | Daddr : V6.Prefix.t t
+
+    let to_string : type a. a t -> string = function
+      | Saddr -> "saddr"
+      | Daddr -> "saddr"
+  end
+
+  type _ t =
+    | Udp : 'a Udp.t -> 'a t
+    | Tcp : 'a Tcp.t -> 'a t
+    | Ipv4 : 'a Ipv4.t -> 'a t
+    | Ipv6 : 'a Ipv6.t -> 'a t
+
+  let to_json (type a) (payload : a t) =
     let protocol, field =
       match payload with
-    | Ipv4 ipv4 -> ("ip", string_of_ipv4 ipv4)
+      | Udp udp -> ("udp", Udp.to_string udp)
-    | Ipv6 ipv6 -> ("ip6", string_of_ipv6 ipv6)
+      | Tcp tcp -> ("tcp", Tcp.to_string tcp)
-    | Udp udp -> ("udp", string_of_udp udp)
+      | Ipv4 ipv4 -> ("ip", Ipv4.to_string ipv4)
-    | Tcp tcp -> ("tcp", string_of_tcp tcp)
+      | Ipv6 ipv6 -> ("ip6", Ipv6.to_string ipv6)
     in
     assoc_one "payload"
       (`Assoc [ ("protocol", `String protocol); ("field", `String field) ])
+end
 
-let rec json_of_expr : type a. a expr -> Yojson.Basic.t = function
+module Expr = struct
-  | String s -> `String s
+  type _ t =
-  | Number n -> `Int n
+    | String : string -> string t
-  | Boolean b -> `Bool b
+    | Number : int -> int t
-  | Ipv4 i ->
+    | Boolean : bool -> int t
-      assoc_one "prefix"
+    | Ipv4 : V4.Prefix.t -> V4.Prefix.t t
-        (`Assoc
+    | Ipv6 : V6.Prefix.t -> V6.Prefix.t t
-          [
+    | List : 'a t list -> 'a t
-            ("addr", `String V4.(to_string (Prefix.network i)));
+    | Set : 'a t list -> 'a t
-            ("len", `Int (V4.Prefix.bits i));
+    | Range : 'a t * 'a t -> 'a t
-          ])
+    | Payload : 'a Payload.t -> 'a t
-  | Ipv6 i ->
-      assoc_one "prefix"
-        (`Assoc
-          [
-            ("addr", `String V6.(to_string (Prefix.network i)));
-            ("len", `Int (V6.Prefix.bits i));
-          ])
-  | List l -> `List (List.map json_of_expr l)
-  | Set s -> assoc_one "set" (`List (List.map json_of_expr s))
-  | Range (a, b) -> assoc_one "range" (`List [ json_of_expr a; json_of_expr b ])
-  | Payload p -> json_of_payload p
 
-let string_of_match_op : type a b. (a, b) match_op -> string = function
+  let rec to_json : type a. a t -> Yojson.Basic.t = function
+    | String str -> `String str
+    | Number num -> `Int num
+    | Boolean bool -> `Bool bool
+    | Ipv4 ipv4 ->
+        let network = V4.(to_string (Prefix.network ipv4)) in
+        let len = V4.Prefix.bits ipv4 in
+        assoc_one "prefix"
+          (`Assoc [ ("addr", `String network); ("len", `Int len) ])
+    | Ipv6 ipv6 ->
+        let network = V6.(to_string (Prefix.network ipv6)) in
+        let len = V6.Prefix.bits ipv6 in
+        assoc_one "prefix"
+          (`Assoc [ ("addr", `String network); ("len", `Int len) ])
+    | List list -> `List (List.map to_json list)
+    | Set set -> assoc_one "set" (`List (List.map to_json set))
+    | Range (a, b) -> assoc_one "range" (`List [ to_json a; to_json b ])
+    | Payload payload -> Payload.to_json payload
+end
+
+module Match = struct
+  type (_, _) t = Eq : ('a, 'a) t | NotEq : ('a, 'a) t
+
+  let to_string : type a b. (a, b) t -> string = function
     | Eq -> "=="
     | NotEq -> "!="
 
-let json_of_counter = function
+  let to_json (type a b) (op : (a, b) t) = `String (to_string op)
-  | Named n -> assoc_one "counter" (`String n)
+end
-  | Anon { packets; bytes } ->
+
+module Counter = struct
+  type t =
+    | NamedCounter of string
+    | AnonCounter of { packets : int; bytes : int }
+
+  let to_json = function
+    | NamedCounter n -> assoc_one "counter" (`String n)
+    | AnonCounter { packets; bytes } ->
         assoc_one "counter"
           (`Assoc [ ("packets", `Int packets); ("bytes", `Int bytes) ])
+end
 
-let json_of_verdict = function
+module Verdict = struct
+  type t = Accept | Drop | Continue | Return | Jump of string | Goto of string
+
+  let to_json = function
     | Accept -> assoc_one "accept" `Null
     | Drop -> assoc_one "drop" `Null
     | Continue -> assoc_one "continue" `Null
     | Return -> assoc_one "return" `Null
     | Jump s -> assoc_one "jump" (assoc_one "target" (`String s))
     | Goto s -> assoc_one "goto" (assoc_one "target" (`String s))
+end
 
-let json_of_stmt : type a. a stmt -> Yojson.Basic.t = function
+module Stmt = struct
+  type _ t =
+    | Match : (('a, 'b) Match.t * 'a Expr.t * 'b Expr.t) -> unit t
+    | Counter : Counter.t -> unit t
+    | Verdict : Verdict.t -> unit t
+    | NoTrack : unit t
+    | Log : { prefix : string option; group : int option } -> unit t
+
+  let to_json : type a. a t -> Yojson.Basic.t = function
     | Match (op, left, right) ->
         assoc_one "match"
           (`Assoc
             [
-            ("left", json_of_expr left);
+              ("left", Expr.to_json left);
-            ("right", json_of_expr right);
+              ("right", Expr.to_json right);
-            ("op", `String (string_of_match_op op));
+              ("op", Match.to_json op);
             ])
-  | Counter c -> json_of_counter c
+    | Counter counter -> Counter.to_json counter
-  | Verdict v -> json_of_verdict v
+    | Verdict verdict -> Verdict.to_json verdict
     | NoTrack -> assoc_one "notrack" `Null
     | Log { prefix; group } ->
         let elems =
@@ -150,46 +148,72 @@ let json_of_stmt : type a. a stmt -> Yojson.Basic.t = function
           ]
         in
         assoc_one "log" (`Assoc (deoptionalise elems))
+end
 
-let string_of_family = function
+module Family = struct
-  | FamilyIpv4 -> "ip"
+  type t = Ipv4 | Ipv6 | Inet
-  | FamilyIpv6 -> "ip6"
-  | FamilyInet -> "inet"
 
-let json_of_table { family; table_name } =
+  let to_string = function Ipv4 -> "ip" | Ipv6 -> "ip6" | Inet -> "inet"
+  let to_json f = `String (to_string f)
+end
+
+module Table = struct
+  type t = { family : Family.t; name : string }
+
+  let to_json { family; name } =
     assoc_one "table"
-    (`Assoc
+      (`Assoc [ ("family", Family.to_json family); ("name", `String name) ])
-      [
+end
-        ("family", `String (string_of_family family));
-        ("name", `String table_name);
-      ])
 
-let json_of_chain { family; table; chain_name } =
+module Chain = struct
+  type t = { family : Family.t; table : string; name : string }
+
+  let to_json { family; table; name } =
     assoc_one "chain"
       (`Assoc
         [
-        ("family", `String (string_of_family family));
+          ("family", Family.to_json family);
           ("table", `String table);
-        ("name", `String chain_name);
+          ("name", `String name);
         ])
+end
 
-let json_of_rule { family; table; chain; expr } =
+module Rule = struct
+  type t = {
+    family : Family.t;
+    table : string;
+    chain : string;
+    expr : unit Stmt.t list;
+  }
+
+  let to_json { family; table; chain; expr } =
     assoc_one "rule"
       (`Assoc
         [
-        ("family", `String (string_of_family family));
+          ("family", Family.to_json family);
           ("table", `String table);
           ("chain", `String chain);
-        ("expr", `List (List.map json_of_stmt expr));
+          ("expr", `List (List.map Stmt.to_json expr));
         ])
+end
 
-let json_of_command = function
+module Command = struct
+  type t =
+    | FlushRuleset
+    | AddTable of Table.t
+    | FlushTable of Table.t
+    | AddChain of Chain.t
+    | FlushChain of Chain.t
+    | AddRule of Rule.t
+
+  let to_json = function
     | FlushRuleset -> assoc_one "flush" (assoc_one "ruleset" `Null)
-  | AddTable t -> assoc_one "add" (json_of_table t)
+    | AddTable table -> assoc_one "add" (Table.to_json table)
-  | FlushTable t -> assoc_one "flush" (json_of_table t)
+    | FlushTable table -> assoc_one "flush" (Table.to_json table)
-  | AddChain c -> assoc_one "add" (json_of_chain c)
+    | AddChain chain -> assoc_one "add" (Chain.to_json chain)
-  | FlushChain c -> assoc_one "flush" (json_of_chain c)
+    | FlushChain chain -> assoc_one "flush" (Chain.to_json chain)
-  | AddRule r -> assoc_one "add" (json_of_rule r)
+    | AddRule rule -> assoc_one "add" (Rule.to_json rule)
+end
 
-let json_of_nftables n =
+let to_json commands =
-  assoc_one "nftables" (`List (List.map json_of_command n))
+  assoc_one "nftables" (`List (List.map Command.to_json commands))