Use LRU cache to prevent out-of-memory errors

README.md

@@ -31,8 +31,8 @@ This took about 10 minutes on my laptop (it will be much quicker if you run it a
 
 2. Install mirage, pinning a few unreleased features we need:
 
-        opam pin add -n -y tcpip 'https://github.com/talex5/mirage-tcpip.git#fix-length-checks'
-        opam pin add -y mirage-nat 'https://github.com/talex5/mirage-nat.git#cleanup'
+        opam pin add -n -y tcpip.3.0.0 'https://github.com/talex5/mirage-tcpip.git#fix-length-checks'
+        opam pin add -y mirage-nat 'https://github.com/talex5/mirage-nat.git#lru'
         opam install mirage
 
 3. Build mirage-firewall:

config.ml

@@ -5,8 +5,18 @@
 
 open Mirage
 
+let table_size =
+  let open Functoria_key in
+  let info = Arg.info
+      ~doc:"The number of NAT entries to allocate."
+      ~docv:"ENTRIES" ["nat-table-size"]
+  in
+  let key = Arg.opt ~stage:`Both Arg.int 5_000 info in
+  create "nat_table_size" key
+
 let main =
   foreign
+    ~keys:[Functoria_key.abstract table_size]
     ~packages:[
       package "vchan";
       package "cstruct";

my_nat.ml

@@ -18,8 +18,10 @@ type t = {
   get_time : unit -> Mirage_nat.time;
 }
 
-let create ~get_time =
-  Nat.empty () >|= fun table ->
+let create ~get_time ~max_entries =
+  let tcp_size = 7 * max_entries / 8 in
+  let udp_size = max_entries - tcp_size in
+  Nat.empty ~tcp_size ~udp_size ~icmp_size:100 >|= fun table ->
   { get_time; table }
 
 let translate t packet =
@@ -53,26 +55,27 @@ let add_nat_rule_and_translate t ~xl_host action packet =
     let xl_port = random_user_port () in
     apply_action xl_port >>= function
     | Error `Out_of_memory ->
-        (* 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...");
-        Nat.reset t.table >>= fun () ->
-        aux ~retries:(retries - 1)
+      (* 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...");
+      Nat.reset t.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");
-          Nat.reset t.table >>= fun () ->
-          aux ~retries:(retries - 1)
-        ) else (
-          aux ~retries:(retries - 1)
-        )
+      if retries = 0 then (
+        Log.warn (fun f -> f "Failed to find a free port; resetting NAT table");
+        Nat.reset t.table >>= fun () ->
+        aux ~retries:(retries - 1)
+      ) else (
+        aux ~retries:(retries - 1)
+      )
     | Error `Cannot_NAT ->
-        Lwt.return (Error "Cannot NAT this packet")
+      Lwt.return (Error "Cannot NAT this packet")
     | Ok () ->
-        translate t packet >|= function
-        | None -> Error "No NAT entry, even after adding one!"
-        | Some packet ->
-          Ok packet
+      Log.debug (fun f -> f "Updated NAT table: %a" Nat.pp_summary t.table);
+      translate t packet >|= function
+      | None -> Error "No NAT entry, even after adding one!"
+      | Some packet ->
+        Ok packet
   in
   aux ~retries:100

my_nat.mli

@@ -10,7 +10,7 @@ type action = [
   | `Redirect of Mirage_nat.endpoint
 ]
 
-val create : get_time:(unit -> Mirage_nat.time) -> t Lwt.t
+val create : get_time:(unit -> Mirage_nat.time) -> max_entries:int -> t Lwt.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.V4.t ->

unikernel.ml

@@ -72,7 +72,8 @@ module Main (Clock : Mirage_clock_lwt.MCLOCK) = struct
       return () in
     (* Set up networking *)
     let get_time () = Clock.elapsed_ns clock in
-    My_nat.create ~get_time >>= fun nat ->
+    let max_entries = Key_gen.nat_table_size () in
+    My_nat.create ~get_time ~max_entries >>= fun nat ->
     let net_listener = network ~clock nat qubesDB in
     (* Report memory usage to XenStore *)
     Memory_pressure.init ();