unikernel/duniverse/mirage-tcpip/test/vnetif_common.ml

@@ -0,0 +1,135 @@
+(*
+ * Copyright (c) 2015 Magnus Skjegstad <magnus@skjegstad.com>
+ *
+ * Permission to use, copy, modify, and distribute this software for any
+ * purpose with or without fee is hereby granted, provided that the above
+ * copyright notice and this permission notice appear in all copies.
+ *
+ * THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES
+ * WITH REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF
+ * MERCHANTABILITY AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR
+ * ANY SPECIAL, DIRECT, INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES
+ * WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS, WHETHER IN AN
+ * ACTION OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF
+ * OR IN CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE.
+ *)
+
+open Common
+open Lwt.Infix
+
+module type VNETIF_STACK =
+sig
+  type backend
+  module Stack : Tcpip.Stack.V4V6
+
+  (** Create a new backend *)
+  val create_backend : unit -> backend
+
+  (** Create a new stack connected to an existing backend *)
+  val create_stack : ?mtu:int -> cidr:Ipaddr.V4.Prefix.t ->
+    ?gateway:Ipaddr.V4.t -> ?cidr6:Ipaddr.V6.Prefix.t ->
+    ?gateway6:Ipaddr.V6.t -> backend -> Stack.t Lwt.t
+
+  val create_backend_listener : backend -> (Cstruct.t -> unit Lwt.t) -> int
+
+  (** Disable a listener function *)
+  val disable_backend_listener : backend -> int -> unit Lwt.t
+
+  (** Records pcap data from the backend while running the specified
+      function. Disables the pcap recorder when the function exits. *)
+  val record_pcap : backend -> string -> (unit -> unit Lwt.t) -> unit Lwt.t
+end
+
+module VNETIF_STACK (B: Vnetif_backends.Backend): sig
+  include VNETIF_STACK with type backend = B.t
+
+  module T : sig
+    val num_open_channels : Stack.TCP.t -> int
+  end
+end
+= struct
+  type backend = B.t
+  module V = Vnetif.Make(B)
+  module E = Ethernet.Make(V)
+
+  module A = Arp.Make(E)
+  module Ip4 = Static_ipv4.Make(E)(A)
+  module Icmp4 = Icmpv4.Make(Ip4)
+  module Ip6 = Ipv6.Make(V)(E)
+  module Ip46 = Tcpip_stack_direct.IPV4V6(Ip4)(Ip6)
+  module U = Udp.Make(Ip46)
+  module T = Tcp.Flow.Make(Ip46)
+
+  module Stack =
+    Tcpip_stack_direct.MakeV4V6(V)(E)(A)(Ip46)(Icmp4)(U)(T)
+
+  let create_backend () =
+    B.create ()
+
+  let create_stack ?mtu ~cidr ?gateway ?cidr6 ?gateway6 backend =
+    let size_limit = match mtu with None -> None | Some x -> Some x in
+    V.connect ?size_limit backend >>= fun netif ->
+    E.connect netif >>= fun ethif ->
+    A.connect ethif >>= fun arpv4 ->
+    Ip4.connect ~cidr ?gateway ethif arpv4 >>= fun ipv4 ->
+    Icmp4.connect ipv4 >>= fun icmpv4 ->
+    Ip6.connect ?cidr:cidr6 ?gateway:gateway6 netif ethif >>= fun ipv6 ->
+    Ip46.connect ~ipv4_only:false ~ipv6_only:false ipv4 ipv6 >>= fun ip ->
+    U.connect ip >>= fun udp ->
+    T.connect ip >>= fun tcp ->
+    Stack.connect netif ethif arpv4 ip icmpv4 udp tcp
+
+  let create_backend_listener backend listenf =
+    match (B.register backend) with
+    | Error _ -> failf "Error occurred while registering to backend"
+    | Ok id -> (B.set_listen_fn backend id listenf); id
+
+  let disable_backend_listener backend id =
+    B.unregister_and_flush backend id
+
+  let create_pcap_recorder backend channel =
+    let header_buf = Cstruct.create Pcap.sizeof_pcap_header in
+    Pcap.LE.set_pcap_header_magic_number header_buf Pcap.magic_number;
+    Pcap.LE.set_pcap_header_network header_buf Pcap.Network.(to_int32 Ethernet);
+    Pcap.LE.set_pcap_header_sigfigs header_buf 0l;
+    Pcap.LE.set_pcap_header_snaplen header_buf 0xffffl;
+    Pcap.LE.set_pcap_header_thiszone header_buf 0l;
+    Pcap.LE.set_pcap_header_version_major header_buf Pcap.major_version;
+    Pcap.LE.set_pcap_header_version_minor header_buf Pcap.minor_version;
+    Lwt_io.write channel (Cstruct.to_string header_buf) >>= fun () ->
+    let pcap_record channel buffer =
+      let pcap_buf = Cstruct.create Pcap.sizeof_pcap_packet in
+      let time = Unix.gettimeofday () in
+      Pcap.LE.set_pcap_packet_incl_len pcap_buf (Int32.of_int (Cstruct.length buffer));
+      Pcap.LE.set_pcap_packet_orig_len pcap_buf (Int32.of_int (Cstruct.length buffer));
+      Pcap.LE.set_pcap_packet_ts_sec pcap_buf (Int32.of_float time);
+      let frac = (time -. (float_of_int (truncate time))) *. 1000000.0 in
+      Pcap.LE.set_pcap_packet_ts_usec pcap_buf (Int32.of_float frac);
+      (try
+          Lwt_io.write channel ((Cstruct.to_string pcap_buf) ^ (Cstruct.to_string buffer))
+      with
+        Lwt_io.Channel_closed msg -> Printf.printf "Warning: Pcap output channel already closed: %s.\n" msg; Lwt.return_unit
+      )
+      >>= fun () ->
+      Lwt.return_unit
+    in
+    let recorder_id = create_backend_listener backend (pcap_record channel) in
+    Lwt.return recorder_id
+
+  let record_pcap backend pcap_file fn =
+    Lwt.catch
+      (fun _ ->
+        Lwt_io.with_file ~mode:Lwt_io.output pcap_file (fun oc ->
+        create_pcap_recorder backend oc >>= fun recorder_id ->
+        fn () >>= fun () ->
+        disable_backend_listener backend recorder_id >>= fun () ->
+        Lwt.return_unit
+        )
+      )
+      (function
+        | Unix.Unix_error _ ->
+          Printf.printf "Could not create pcap file %s - something along the way doesn't exist.\n" pcap_file;
+          fn ()
+        | e -> Lwt.fail e
+      )
+end