Use new Nat.reset function to clear the table

my_nat.ml

@@ -11,58 +11,37 @@ type action = [
   | `Redirect of Ipaddr.t * int
 ]
 
-(* To avoid needing to allocate a new NAT table when we've run out of
-   memory, pre-allocate the new one ahead of time. *)
-type 'a with_standby = {
-  mutable current :'a;
-  mutable next : 'a;
-}
+type t = Nat : (module Mirage_nat.S with type t = 't) * 't -> t
 
-type t = Nat : (module Mirage_nat.S with type t = 't and type config = 'c) * 'c * 't with_standby -> t
+let create (type t) (nat:(module Mirage_nat.S with type t = t)) (table:t) =
+  let (module Nat : Mirage_nat.S with type t = t) = nat in
+  Nat (nat, table)
 
-let create (type c t) (nat:(module Mirage_nat.S with type config = c and type t = t)) (c:c) =
-  let (module Nat : Mirage_nat.S with type config = c and type t = t) = nat in
-  Nat.empty c >>= fun current ->
-  Nat.empty c >>= fun next ->
-  let table = { current; next } in
-  Lwt.return (Nat (nat, c, table))
-
-let translate (Nat ((module Nat), _, table)) packet =
-  Nat.translate table.current packet >|= function
+let translate (Nat ((module Nat), table)) packet =
+  Nat.translate table packet >|= function
   | Error `Untranslated -> None
   | Ok packet -> Some packet
 
 let random_user_port () =
   1024 + Random.int (0xffff - 1024)
 
-let reset (Nat ((module Nat), c, table)) =
-  table.current <- table.next;
-  (* (at this point, the big old NAT table can be GC'd, so allocating
-     a new one should be OK) *)
-  Nat.empty c >|= fun next ->
-  table.next <- next
+let reset (Nat ((module Nat), table)) =
+  Nat.reset table
 
-let add_nat_rule_and_translate ((Nat ((module Nat), c, table)) as t) ~xl_host action packet =
+let add_nat_rule_and_translate ((Nat ((module Nat), table)) as t) ~xl_host action packet =
   let apply_action xl_port =
     Lwt.catch (fun () ->
         match action with
         | `Rewrite ->
-          Nat.add_nat table.current packet (xl_host, xl_port)
+          Nat.add_nat table packet (xl_host, xl_port)
         | `Redirect target ->
-          Nat.add_redirect table.current packet (xl_host, xl_port) target
+          Nat.add_redirect table packet (xl_host, xl_port) target
       )
       (function
         | Out_of_memory -> Lwt.return (Error `Out_of_memory)
         | x -> Lwt.fail x
       )
   in
-  let reset () =
-    table.current <- table.next;
-    (* (at this point, the big old NAT table can be GC'd, so allocating
-       a new one should be OK) *)
-    Nat.empty c >|= fun next ->
-    table.next <- next
-  in
   let rec aux ~retries =
     let xl_port = random_user_port () in
     apply_action xl_port >>= function
@@ -70,13 +49,13 @@ let add_nat_rule_and_translate ((Nat ((module Nat), c, table)) as t) ~xl_host ac
         (* Because hash tables resize in big steps, this can happen even if we have a fair
            chunk of free memory. *)
         Log.warn (fun f -> f "Out_of_memory adding NAT rule. Dropping NAT table...");
-        reset () >>= fun () ->
+        Nat.reset table >>= fun () ->
         aux ~retries:(retries - 1)
     | Error `Overlap when retries < 0 -> Lwt.return (Error "Too many retries")
     | Error `Overlap ->
         if retries = 0 then (
           Log.warn (fun f -> f "Failed to find a free port; resetting NAT table");
-          reset () >>= fun () ->
+          Nat.reset table >>= fun () ->
           aux ~retries:(retries - 1)
         ) else (
           aux ~retries:(retries - 1)
@@ -87,22 +66,6 @@ let add_nat_rule_and_translate ((Nat ((module Nat), c, table)) as t) ~xl_host ac
         translate t packet >|= function
         | None -> Error "No NAT entry, even after adding one!"
         | Some packet ->
-(*
-          Log.debug (fun f ->
-              match action with
-              | `Rewrite ->
-                let (ip, trans) = packet in
-                let src, dst = Nat_rewrite.addresses_of_ip ip in
-                let sport, dport = Nat_rewrite.ports_of_transport transport in
-                f "added NAT entry: %s:%d -> firewall:%d -> %d:%s" (Ipaddr.to_string src) sport xl_port dport (Ipaddr.to_string dst)
-              | `Redirect ->
-                let (ip, transport) = packet in
-                let src, _dst = Nat_rewrite.addresses_of_ip ip in
-                let sport, dport = Nat_rewrite.ports_of_transport transport in
-                f "added NAT redirect %s:%d -> %d:firewall:%d -> %d:%a"
-                  (Ipaddr.to_string src) sport dport xl_port port pp_host host
-            );
-*)
           Ok packet
   in
   aux ~retries:100

my_nat.mli

@@ -10,7 +10,7 @@ type action = [
   | `Redirect of Ipaddr.t * int
 ]
 
-val create : (module Mirage_nat.S with type t = 'a and type config = 'c) -> 'c -> t Lwt.t
+val create : (module Mirage_nat.S with type t = 'a) -> 'a -> t
 val reset : t -> unit Lwt.t
 val translate : t -> Nat_packet.t -> Nat_packet.t option Lwt.t
 val add_nat_rule_and_translate : t -> xl_host:Ipaddr.t ->

unikernel.ml

@@ -72,7 +72,7 @@ module Main (Clock : Mirage_clock_lwt.MCLOCK) = struct
       OS.Lifecycle.await_shutdown_request () >>= fun (`Poweroff | `Reboot) ->
       return () in
     (* Set up networking *)
-    My_nat.create (module Nat) clock >>= fun nat ->
+    Nat.empty clock >|= My_nat.create (module Nat) >>= fun nat ->
     let net_listener = network ~clock nat qubesDB in
     (* Report memory usage to XenStore *)
     Memory_pressure.init ();