New mirage nat interface, less mutable state

client_net.ml

@@ -101,7 +101,7 @@ let add_vif get_ts { Dao.ClientVif.domid; device_id } dns_client dns_servers ~cl
                         (Ipaddr.V4.to_string client_ip)
                         Fmt.(list ~sep:(any "@.") Pf_qubes.Parse_qubes.pp_rule) new_rules);
             (* empty NAT table if rules are updated: they might deny old connections *)
-            My_nat.remove_connections router.Router.nat router.Router.ports client_ip;
+            My_nat.remove_connections router.Router.nat client_ip;
           end);
           update new_db new_rules
         in

firewall.ml

@@ -47,7 +47,7 @@ let translate t packet =
 let add_nat_and_forward_ipv4 t packet =
   let open Router in
   let xl_host = t.uplink#my_ip in
-  My_nat.add_nat_rule_and_translate t.nat t.ports ~xl_host `NAT packet >>= function
+  match My_nat.add_nat_rule_and_translate t.nat ~xl_host `NAT packet with
   | Ok packet -> forward_ipv4 t packet
   | Error e ->
     Log.warn (fun f -> f "Failed to add NAT rewrite rule: %s (%a)" e Nat_packet.pp packet);
@@ -60,7 +60,7 @@ let nat_to t ~host ~port packet =
   | Ipaddr.V6 _ -> Log.warn (fun f -> f "Cannot NAT with IPv6"); Lwt.return_unit
   | Ipaddr.V4 target ->
     let xl_host = t.uplink#my_ip in
-    My_nat.add_nat_rule_and_translate t.nat t.ports ~xl_host (`Redirect (target, port)) packet >>= function
+    match My_nat.add_nat_rule_and_translate t.nat ~xl_host (`Redirect (target, port)) packet with
     | Ok packet -> forward_ipv4 t packet
     | Error e ->
       Log.warn (fun f -> f "Failed to add NAT redirect rule: %s (%a)" e Nat_packet.pp packet);
@@ -86,44 +86,43 @@ let apply_rules t (rules : ('a, 'b) Packet.t -> Packet.action Lwt.t) ~dst (annot
 let handle_low_memory t =
   match Memory_pressure.status () with
   | `Memory_critical -> (* TODO: should happen before copying and async *)
-      Log.warn (fun f -> f "Memory low - dropping packet and resetting NAT table");
-      My_nat.reset t.Router.nat t.Router.ports >|= fun () ->
+      Log.warn (fun f -> f "Memory low - dropping packet");
       `Memory_critical
-  | `Ok -> Lwt.return `Ok
+  | `Ok -> `Ok
 
 let ipv4_from_client resolver dns_servers t ~src packet =
-  handle_low_memory t >>= function
+  match handle_low_memory t with
   | `Memory_critical -> Lwt.return_unit
   | `Ok ->
   (* Check for existing NAT entry for this packet *)
-  translate t packet >>= function
-  | Some frame -> forward_ipv4 t frame  (* Some existing connection or redirect *)
-  | None ->
-  (* No existing NAT entry. Check the firewall rules. *)
-  let `IPv4 (ip, _transport) = packet in
-  let dst = Router.classify t (Ipaddr.V4 ip.Ipv4_packet.dst) in
-  match of_mirage_nat_packet ~src:(`Client src) ~dst packet with
-  | None -> Lwt.return_unit
-  | Some firewall_packet -> apply_rules t (Rules.from_client resolver dns_servers) ~dst firewall_packet
+    match translate t packet with
+    | Some frame -> forward_ipv4 t frame  (* Some existing connection or redirect *)
+    | None ->
+      (* No existing NAT entry. Check the firewall rules. *)
+      let `IPv4 (ip, _transport) = packet in
+      let dst = Router.classify t (Ipaddr.V4 ip.Ipv4_packet.dst) in
+      match of_mirage_nat_packet ~src:(`Client src) ~dst packet with
+      | None -> Lwt.return_unit
+      | Some firewall_packet -> apply_rules t (Rules.from_client resolver dns_servers) ~dst firewall_packet
 
 let ipv4_from_netvm t packet =
-  handle_low_memory t >>= function
+  match handle_low_memory t with
   | `Memory_critical -> Lwt.return_unit
   | `Ok ->
-  let `IPv4 (ip, _transport) = packet in
-  let src = Router.classify t (Ipaddr.V4 ip.Ipv4_packet.src) in
-  let dst = Router.classify t (Ipaddr.V4 ip.Ipv4_packet.dst) in
-  match Packet.of_mirage_nat_packet ~src ~dst packet with
-  | None -> Lwt.return_unit
-  | Some _ ->
-  match src with
-  | `Client _ | `Firewall ->
-    Log.warn (fun f -> f "Frame from NetVM has internal source IP address! %a" Nat_packet.pp packet);
-    Lwt.return_unit
-  | `External _ | `NetVM as src ->
-  translate t packet >>= function
-  | Some frame -> forward_ipv4 t frame
-  | None ->
+    let `IPv4 (ip, _transport) = packet in
+    let src = Router.classify t (Ipaddr.V4 ip.Ipv4_packet.src) in
+    let dst = Router.classify t (Ipaddr.V4 ip.Ipv4_packet.dst) in
     match Packet.of_mirage_nat_packet ~src ~dst packet with
     | None -> Lwt.return_unit
-    | Some packet -> apply_rules t Rules.from_netvm ~dst packet
+    | Some _ ->
+      match src with
+      | `Client _ | `Firewall ->
+        Log.warn (fun f -> f "Frame from NetVM has internal source IP address! %a" Nat_packet.pp packet);
+        Lwt.return_unit
+      | `External _ | `NetVM as src ->
+        match translate t packet with
+        | Some frame -> forward_ipv4 t frame
+        | None ->
+          match Packet.of_mirage_nat_packet ~src ~dst packet with
+          | None -> Lwt.return_unit
+          | Some packet -> apply_rules t Rules.from_netvm ~dst packet

my_dns.ml

@@ -35,12 +35,14 @@ module Transport (R : Mirage_random.S) (C : Mirage_clock.MCLOCK) (Time : Mirage_
     let open My_nat in
     let dst, dst_port = ctx.nameserver in
     let router, send_udp, answer = ctx.stack in
-    let src_port = Ports.pick_free_port ~consult:router.ports.nat_udp router.ports.dns_udp in
+    let src_port, evict =
+      My_nat.free_udp_port router.nat ~src:router.uplink#my_ip ~dst ~dst_port:53
+    in
     with_timeout ctx.timeout_ns
       ((send_udp ~src_port ~dst ~dst_port buf >|= Rresult.R.open_error_msg) >>= function
         | Ok () -> (Lwt_mvar.take answer >|= fun (_, dns_response) -> Ok dns_response)
         | Error _ as e -> Lwt.return e) >|= fun result ->
-    router.ports.dns_udp := Ports.remove src_port !(router.ports.dns_udp);
+    evict ();
     result
 
   let close _ = Lwt.return_unit

my_nat.ml

@@ -11,34 +11,62 @@ type action = [
   | `Redirect of Mirage_nat.endpoint
 ]
 
-type ports = {
-  nat_tcp : Ports.t ref;
-  nat_udp : Ports.t ref;
-  nat_icmp : Ports.t ref;
-  dns_udp : Ports.t ref;
-}
-
-let empty_ports () =
-  let nat_tcp = ref Ports.empty in
-  let nat_udp = ref Ports.empty in
-  let nat_icmp = ref Ports.empty in
-  let dns_udp = ref Ports.empty in
-  { nat_tcp ; nat_udp ; nat_icmp ; dns_udp }
-
 module Nat = Mirage_nat_lru
 
+module S =
+  Set.Make(struct type t = int let compare (a : int) (b : int) = compare a b end)
+
 type t = {
   table : Nat.t;
+  mutable udp_dns : S.t;
+  last_resort_port : int
 }
 
+let pick_port () =
+  1024 + Random.int (0xffff - 1024)
+
 let create ~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 ->
-  { table }
+  let table = Nat.empty ~tcp_size ~udp_size ~icmp_size:100 in
+  let last_resort_port = pick_port () in
+  { table ; udp_dns = S.empty ; last_resort_port }
+
+let pick_free_port t proto =
+  let rec go retries =
+    if retries = 0 then
+      None
+    else
+      let p = 1024 + Random.int (0xffff - 1024) in
+      match proto with
+      | `Udp when S.mem p t.udp_dns || p = t.last_resort_port ->
+        go (retries - 1)
+      | _ -> Some p
+  in
+  go 10
+
+let free_udp_port t ~src ~dst ~dst_port =
+  let rec go () =
+    let src_port =
+      Option.value ~default:t.last_resort_port (pick_free_port t `Udp)
+    in
+    if Nat.is_port_free t.table `Udp ~src ~dst ~src_port ~dst_port then begin
+      let remove =
+        if src_port <> t.last_resort_port then begin
+          t.udp_dns <- S.add src_port t.udp_dns;
+          (fun () -> t.udp_dns <- S.remove src_port t.udp_dns)
+        end else Fun.id
+      in
+      src_port, remove
+    end else
+      go ()
+  in
+  go ()
+
+let dns_port t port = S.mem port t.udp_dns || port = t.last_resort_port
 
 let translate t packet =
-  Nat.translate t.table packet >|= function
+  match Nat.translate t.table packet with
   | Error (`Untranslated | `TTL_exceeded as e) ->
     Log.debug (fun f -> f "Failed to NAT %a: %a"
                   Nat_packet.pp packet
@@ -47,63 +75,19 @@ let translate t packet =
     None
   | Ok packet -> Some packet
 
-let pick_free_port ~nat_ports ~dns_ports =
-  Ports.pick_free_port ~consult:dns_ports nat_ports
-
-(* just clears the nat ports, dns ports stay as is *)
-let reset t ports =
-  ports.nat_tcp := Ports.empty;
-  ports.nat_udp := Ports.empty;
-  ports.nat_icmp := Ports.empty;
-  Nat.reset t.table
+let remove_connections t ip =
+  ignore (Nat.remove_connections t.table ip)
 
-let remove_connections t ports ip =
-  let freed_ports = Nat.remove_connections t.table ip in
-  ports.nat_tcp := Ports.diff !(ports.nat_tcp) (Ports.of_list freed_ports.Mirage_nat.tcp);
-  ports.nat_udp := Ports.diff !(ports.nat_udp) (Ports.of_list freed_ports.Mirage_nat.udp);
-  ports.nat_icmp := Ports.diff !(ports.nat_icmp) (Ports.of_list freed_ports.Mirage_nat.icmp)
-
-let add_nat_rule_and_translate t ports ~xl_host action packet =
-  let apply_action xl_port =
-    Lwt.catch (fun () ->
-        Nat.add t.table packet (xl_host, xl_port) action
-      )
-      (function
-        | Out_of_memory -> Lwt.return (Error `Out_of_memory)
-        | x -> Lwt.fail x
-      )
-  in
-  let rec aux ~retries =
-    let nat_ports, dns_ports =
-      match packet with
-      | `IPv4 (_, `TCP _) -> ports.nat_tcp, ref Ports.empty
-      | `IPv4 (_, `UDP _) -> ports.nat_udp, ports.dns_udp
-      | `IPv4 (_, `ICMP _) -> ports.nat_icmp, ref Ports.empty
-    in
-    let xl_port = pick_free_port ~nat_ports ~dns_ports 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...");
-      reset t ports >>= 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 t ports >>= fun () ->
-        aux ~retries:(retries - 1)
-      ) else (
-        aux ~retries:(retries - 1)
-      )
-    | Error `Cannot_NAT ->
-      Lwt.return (Error "Cannot NAT this packet")
-    | Ok () ->
-      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
+let add_nat_rule_and_translate t ~xl_host action packet =
+  let proto = match packet with
+    | `IPv4 (_, `TCP _) -> `Tcp
+    | `IPv4 (_, `UDP _) -> `Udp
+    | `IPv4 (_, `ICMP _) -> `Icmp
   in
-  aux ~retries:100
+  match Nat.add t.table packet xl_host (fun () -> pick_free_port t proto) action with
+  | Error `Overlap -> Error "Too many retries"
+  | Error `Cannot_NAT -> Error "Cannot NAT this packet"
+  | Ok () ->
+    Log.debug (fun f -> f "Updated NAT table: %a" Nat.pp_summary t.table);
+    Option.to_result ~none:"No NAT entry, even after adding one!"
+      (translate t packet)

my_nat.mli

@@ -3,25 +3,18 @@
 
 (* Abstract over NAT interface (todo: remove this) *)
 
-type ports = private {
-  nat_tcp : Ports.t ref;
-  nat_udp : Ports.t ref;
-  nat_icmp : Ports.t ref;
-  dns_udp : Ports.t ref;
-}
-
-val empty_ports : unit -> ports
-
 type t
 
 type action = [
   | `NAT
   | `Redirect of Mirage_nat.endpoint
 ]
 
-val create : max_entries:int -> t Lwt.t
-val reset : t -> ports -> unit Lwt.t
-val remove_connections : t -> ports -> Ipaddr.V4.t -> unit
-val translate : t -> Nat_packet.t -> Nat_packet.t option Lwt.t
-val add_nat_rule_and_translate : t -> ports ->
-  xl_host:Ipaddr.V4.t -> action -> Nat_packet.t -> (Nat_packet.t, string) result Lwt.t
+val free_udp_port : t -> src:Ipaddr.V4.t -> dst:Ipaddr.V4.t -> dst_port:int ->
+  int * (unit -> unit)
+val dns_port : t -> int -> bool
+val create : max_entries:int -> t
+val remove_connections : t -> Ipaddr.V4.t -> unit
+val translate : t -> Nat_packet.t -> Nat_packet.t option
+val add_nat_rule_and_translate : t ->
+  xl_host:Ipaddr.V4.t -> action -> Nat_packet.t -> (Nat_packet.t, string) result

router.ml

@@ -9,13 +9,10 @@ type t = {
   client_eth : Client_eth.t;
   nat : My_nat.t;
   uplink : interface;
-  (* NOTE: do not try to make this pure, it relies on mvars / side effects *)
-  ports : My_nat.ports;
 }
 
 let create ~client_eth ~uplink ~nat =
-  let ports = My_nat.empty_ports () in
-  { client_eth; nat; uplink; ports }
+  { client_eth; nat; uplink }
 
 let target t buf =
   let dst_ip = buf.Ipv4_packet.dst in

router.mli

@@ -9,7 +9,6 @@ type t = private {
   client_eth : Client_eth.t;
   nat : My_nat.t;
   uplink : interface;
-  ports : My_nat.ports;
 }
 
 val create :

unikernel.ml

@@ -45,7 +45,7 @@ module Main (R : Mirage_random.S)(Clock : Mirage_clock.MCLOCK)(Time : Mirage_tim
       Lwt.return_unit in
     (* Set up networking *)
     let max_entries = Key_gen.nat_table_size () in
-    My_nat.create ~max_entries >>= fun nat ->
+    let nat = My_nat.create ~max_entries in
 
     (* Read network configuration from QubesDB *)
     Dao.read_network_config qubesDB >>= fun config ->

uplink.ml

@@ -47,7 +47,7 @@ end
 
       Log.debug (fun f -> f "received ipv4 packet from %a on uplink" Ipaddr.V4.pp ip_header.Ipv4_packet.src);
       match ip_packet with
-      | `UDP (header, packet) when Ports.mem header.dst_port !(router.Router.ports.My_nat.dns_udp) ->
+      | `UDP (header, packet) when My_nat.dns_port router.Router.nat header.dst_port ->
         Log.debug (fun f -> f "found a DNS packet whose dst_port (%d) was in the list of dns_client ports" header.dst_port);
         Lwt_mvar.put dns_responses (header, packet)
       | _ ->