Give exact types for Packet.src

Before, the packet passed to rules.ml could have any host as its src.
Now, `from_client` knows that `src` must be a `Client`, and `from_netvm`
knows that `src` is `External` or `NetVM`.

client_net.ml

@@ -56,7 +56,7 @@ let input_arp ~fixed_arp ~iface request =
       iface#writev `ARP (fun b -> Arp_packet.encode_into response b; Arp_packet.size)
 
 (** Handle an IPv4 packet from the client. *)
-let input_ipv4 ~client_ip ~router packet =
+let input_ipv4 ~iface ~router packet =
   match Nat_packet.of_ipv4_packet packet with
   | Error e ->
     Log.warn (fun f -> f "Ignored unknown IPv4 message: %a" Nat_packet.pp_error e);
@@ -64,10 +64,10 @@ let input_ipv4 ~client_ip ~router packet =
   | Ok packet ->
     let `IPv4 (ip, _) = packet in
     let src = ip.Ipv4_packet.src in
-    if src = client_ip then Firewall.ipv4_from_client router packet
+    if src = iface#other_ip then Firewall.ipv4_from_client router ~src:iface packet
     else (
       Log.warn (fun f -> f "Incorrect source IP %a in IP packet from %a (dropping)"
-                   Ipaddr.V4.pp src Ipaddr.V4.pp client_ip);
+                   Ipaddr.V4.pp src Ipaddr.V4.pp iface#other_ip);
       return ()
     )
 
@@ -94,7 +94,7 @@ let add_vif { Dao.ClientVif.domid; device_id } ~client_ip ~router ~cleanup_tasks
     | Ok (eth, payload) ->
         match eth.Ethernet_packet.ethertype with
         | `ARP -> input_arp ~fixed_arp ~iface payload
-        | `IPv4 -> input_ipv4 ~client_ip ~router payload
+        | `IPv4 -> input_ipv4 ~iface ~router payload
         | `IPv6 -> return () (* TODO: oh no! *)
   )
   >|= or_raise "Listen on client interface" Netback.pp_error

firewall.ml

@@ -48,8 +48,21 @@ let forward_ipv4 t packet =
 
 (* Packet classification *)
 
-let classify t packet =
-  let `IPv4 (ip, transport) = packet in
+let parse_ips ips = List.map (fun (ip_str, id) -> (Ipaddr.of_string_exn ip_str, id)) ips
+
+let clients = parse_ips Rules.clients
+let externals = parse_ips Rules.externals
+
+let resolve_client client =
+  `Client (try List.assoc (Ipaddr.V4 client#other_ip) clients with Not_found -> `Unknown)
+
+let resolve_host = function
+  | `Client c -> resolve_client c
+  | `External ip -> `External (try List.assoc ip externals with Not_found -> `Unknown)
+  | (`Client_gateway | `Firewall_uplink | `NetVM) as x -> x
+
+let classify ~src ~dst packet =
+  let `IPv4 (_ip, transport) = packet in
   let proto =
     match transport with
     | `TCP ({Tcp.Tcp_packet.src_port; dst_port; _}, _) -> `TCP {sport = src_port; dport = dst_port}
@@ -58,8 +71,8 @@ let classify t packet =
   in
   Some {
     packet;
-    src = Router.classify t (Ipaddr.V4 ip.Ipv4_packet.src);
-    dst = Router.classify t (Ipaddr.V4 ip.Ipv4_packet.dst);
+    src;
+    dst;
     proto;
   }
 
@@ -80,7 +93,10 @@ let pp_proto fmt = function
   | `ICMP -> Format.pp_print_string fmt "ICMP"
   | `Unknown -> Format.pp_print_string fmt "UnknownProtocol"
 
-let pp_packet fmt {src; dst; proto; packet = _} =
+let pp_packet t fmt {src = _; dst = _; proto; packet} =
+  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
   Format.fprintf fmt "[src=%a dst=%a proto=%a]"
     pp_host src
     pp_host dst
@@ -125,30 +141,18 @@ let nat_to t ~host ~port packet =
 
 (* Handle incoming packets *)
 
-let parse_ips ips = List.map (fun (ip_str, id) -> (Ipaddr.of_string_exn ip_str, id)) ips
-
-let clients = parse_ips Rules.clients
-let externals = parse_ips Rules.externals
-
-let resolve_host = function
-  | `Client c -> `Client (try List.assoc (Ipaddr.V4 c#other_ip) clients with Not_found -> `Unknown)
-  | `External ip -> `External (try List.assoc ip externals with Not_found -> `Unknown)
-  | (`Client_gateway | `Firewall_uplink | `NetVM) as x -> x
-
-let apply_rules t rules info =
+let apply_rules t rules ~dst info =
   let packet = info.packet in
-  let resolved_info = { info with src = resolve_host info.src;
-                                  dst = resolve_host info.dst } in
-  match rules resolved_info, info.dst with
+  match rules info, dst with
   | `Accept, `Client client_link -> transmit_ipv4 packet client_link
   | `Accept, (`External _ | `NetVM) -> transmit_ipv4 packet t.Router.uplink
   | `Accept, (`Firewall_uplink | `Client_gateway) ->
-      Log.warn (fun f -> f "Bad rule: firewall can't accept packets %a" pp_packet info);
+      Log.warn (fun f -> f "Bad rule: firewall can't accept packets %a" (pp_packet t) info);
       return ()
   | `NAT, _ -> add_nat_and_forward_ipv4 t packet
   | `NAT_to (host, port), _ -> nat_to t packet ~host ~port
   | `Drop reason, _ ->
-      Log.info (fun f -> f "Dropped packet (%s) %a" reason pp_packet info);
+      Log.info (fun f -> f "Dropped packet (%s) %a" reason (pp_packet t) info);
       return ()
 
 let handle_low_memory t =
@@ -159,7 +163,7 @@ let handle_low_memory t =
       `Memory_critical
   | `Ok -> Lwt.return `Ok
 
-let ipv4_from_client t packet =
+let ipv4_from_client t ~src packet =
   handle_low_memory t >>= function
   | `Memory_critical -> return ()
   | `Ok ->
@@ -168,23 +172,28 @@ let ipv4_from_client t packet =
   | Some frame -> forward_ipv4 t frame  (* Some existing connection or redirect *)
   | None ->
   (* No existing NAT entry. Check the firewall rules. *)
-  match classify t packet with
+  let `IPv4 (ip, _transport) = packet in
+  let dst = Router.classify t (Ipaddr.V4 ip.Ipv4_packet.dst) in
+  match classify ~src:(resolve_client src) ~dst:(resolve_host dst) packet with
   | None -> return ()
-  | Some info -> apply_rules t Rules.from_client info
+  | Some info -> apply_rules t Rules.from_client ~dst info
 
 let ipv4_from_netvm t packet =
   handle_low_memory t >>= function
   | `Memory_critical -> return ()
   | `Ok ->
-  match classify t packet with
+  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 classify ~src ~dst:(resolve_host dst) packet with
   | None -> return ()
   | Some info ->
-  match info.src with
+  match src with
   | `Client _ | `Firewall_uplink | `Client_gateway ->
-    Log.warn (fun f -> f "Frame from NetVM has internal source IP address! %a" pp_packet info);
+    Log.warn (fun f -> f "Frame from NetVM has internal source IP address! %a" (pp_packet t) info);
     return ()
-  | `External _ | `NetVM ->
+  | `External _ | `NetVM as src ->
   translate t packet >>= function
   | Some frame -> forward_ipv4 t frame
   | None ->
-  apply_rules t Rules.from_netvm info
+    apply_rules t Rules.from_netvm ~dst { info with src }

firewall.mli

@@ -6,6 +6,6 @@
 val ipv4_from_netvm : Router.t -> Nat_packet.t -> unit Lwt.t
 (** Handle a packet from the outside world (this module will validate the source IP). *)
 
-val ipv4_from_client : Router.t -> Nat_packet.t -> unit Lwt.t
+val ipv4_from_client : Router.t -> src:Fw_utils.client_link -> Nat_packet.t -> unit Lwt.t
 (** Handle a packet from a client. Caller must check the source IP matches the client's
     before calling this. *)

packet.ml

@@ -13,11 +13,10 @@ type ports = {
 type host = 
   [ `Client of client_link | `Client_gateway | `Firewall_uplink | `NetVM | `External of Ipaddr.t ]
 
-(* Note: 'a is either [host], or the result of applying [Rules.clients] and [Rules.externals] to a host. *)
-type 'a info = {
+type ('src, 'dst) info = {
   packet : Nat_packet.t;
-  src : 'a;
-  dst : 'a;
+  src : 'src;
+  dst : 'dst;
   proto : [ `UDP of ports | `TCP of ports | `ICMP | `Unknown ];
 }
 

rules.ml

@@ -1,12 +1,9 @@
 (* Copyright (C) 2015, Thomas Leonard <thomas.leonard@unikernel.com>
    See the README file for details. *)
 
-(** Put your firewall rules here. *)
+(** Put your firewall rules in this file. *)
 
-open Packet
-
-(* OCaml normally warns if you don't match all fields, but that's OK here. *)
-[@@@ocaml.warning "-9"]
+open Packet   (* Allow us to use definitions in packet.ml *)
 
 (* List your AppVM IP addresses here if you want to match on them in the rules below.
    Any client not listed here will appear as [`Client `Unknown]. *)
@@ -25,11 +22,29 @@ let externals = [
   *)
 ]
 
-(** Decide what to do with a packet from a client VM.
+(* OCaml normally warns if you don't match all fields, but that's OK here. *)
+[@@@ocaml.warning "-9"]
+
+(** This function decides what to do with a packet from a client VM.
+
+    It takes as input an argument [info] (of type [Packet.info]) describing the
+    packet, and returns an action (of type [Packet.action]) to perform.
+
+    See packet.ml for the definitions of [info] and [action].
+
     Note: If the packet matched an existing NAT rule then this isn't called. *)
-let from_client (info : _ info) : action =
+let from_client (info : ([`Client of _], _) Packet.info) : Packet.action =
   match info with
-  (* Examples (add your own rules here): *)
+  (* Examples (add your own rules here):
+
+     1. Allows Dev to send SSH packets to Untrusted.
+        Note: responses are not covered by this!
+     2. Allows clients to continue existing TCP connections with other clients.
+        This allows responses to SSH packets from the previous rule.
+     3. Blocks an external site.
+
+     In all cases, make sure you've added the VM name to [clients] or [externals] above, or it won't
+     match anything! *)
   (*
   | { src = `Client `Dev; dst = `Client `Untrusted; proto = `TCP { dport = 22 } } -> `Accept
   | { src = `Client _; dst = `Client _; proto = `TCP _; packet }
@@ -43,6 +58,6 @@ let from_client (info : _ info) : action =
 
 (** Decide what to do with a packet received from the outside world.
     Note: If the packet matched an existing NAT rule then this isn't called. *)
-let from_netvm (info : _ info) : action =
+let from_netvm (info : ([`NetVM | `External of _], _) Packet.info) : Packet.action =
   match info with
   | _ -> `Drop "drop by default"