From 425ba262863e8b4077160004a36305dcc33e90ee Mon Sep 17 00:00:00 2001
From: Thomas Leonard <thomas.leonard@unikernel.com>
Date: Sat, 2 Jan 2016 15:50:05 +0000
Subject: [PATCH] Reset NAT table if memory gets low

---
 firewall.ml | 18 +++++++++++-------
 router.ml   |  5 ++++-
 router.mli  |  5 ++++-
 3 files changed, 19 insertions(+), 9 deletions(-)

diff --git a/firewall.ml b/firewall.ml
index ad18ac4..f90b5c0 100644
--- a/firewall.ml
+++ b/firewall.ml
@@ -145,11 +145,17 @@ let apply_rules t rules info =
       Log.info "Dropped packet (%s) %a" (fun f -> f reason pp_packet info);
       return ()
 
-let ipv4_from_client t frame =
+let handle_low_memory t =
   match Memory_pressure.status () with
   | `Memory_critical -> (* TODO: should happen before copying and async *)
-      Log.warn "Memory low - dropping packet" Logs.unit;
-      return ()
+      Log.warn "Memory low - dropping packet and resetting NAT table" Logs.unit;
+      Router.reset t;
+      `Memory_critical
+  | `Ok -> `Ok
+
+let ipv4_from_client t frame =
+  match handle_low_memory t with
+  | `Memory_critical -> return ()
   | `Ok ->
   (* Check for existing NAT entry for this packet *)
   match translate t frame with
@@ -161,10 +167,8 @@ let ipv4_from_client t frame =
   | Some info -> apply_rules t Rules.from_client info
 
 let ipv4_from_netvm t frame =
-  match Memory_pressure.status () with
-  | `Memory_critical -> (* TODO: should happen before copying and async *)
-      Log.warn "Memory low - dropping packet" Logs.unit;
-      return ()
+  match handle_low_memory t with
+  | `Memory_critical -> return ()
   | `Ok ->
   match classify t frame with
   | None -> return ()
diff --git a/router.ml b/router.ml
index 3e6dd8b..ba1a2c3 100644
--- a/router.ml
+++ b/router.ml
@@ -10,7 +10,7 @@ module Log = (val Logs.src_log src : Logs.LOG)
 
 type t = {
   client_eth : Client_eth.t;
-  nat : Nat_lookup.t;
+  mutable nat : Nat_lookup.t;
   uplink : interface;
 }
 
@@ -42,3 +42,6 @@ let resolve t = function
   | `Firewall_uplink -> Ipaddr.V4 t.uplink#my_ip
   | `NetVM -> Ipaddr.V4 t.uplink#other_ip
   | #Client_eth.host as host -> Client_eth.resolve t.client_eth host
+
+let reset t =
+  t.nat <- Nat_lookup.empty ()
diff --git a/router.mli b/router.mli
index f5cd8bd..8743b57 100644
--- a/router.mli
+++ b/router.mli
@@ -7,7 +7,7 @@ open Utils
 
 type t = private {
   client_eth : Client_eth.t;
-  nat : Nat_lookup.t;
+  mutable nat : Nat_lookup.t;
   uplink : interface;
 }
 (** A routing table. *)
@@ -30,3 +30,6 @@ val remove_client : t -> client_link -> unit
 
 val classify : t -> Ipaddr.t -> Packet.host
 val resolve : t -> Packet.host -> Ipaddr.t
+
+val reset : t -> unit
+(** Clear the NAT table (to free memory). *)