unikernel/duniverse/mirage-flow/combinators/mirage_flow_combinators.ml

@@ -0,0 +1,430 @@
+(*
+ * Copyright (c) 2011-present Anil Madhavapeddy <anil@recoil.org>
+ * Copyright (c) 2013-present Thomas Gazagnaire <thomas@gazagnaire.org>
+ * Copyright (C) 2016-present David Scott <dave.scott@docker.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 Lwt.Infix
+
+let src = Logs.Src.create "mirage-flow-combinators"
+module Log = (val Logs.src_log src : Logs.LOG)
+
+type stats = {
+  read_bytes: int64;
+  read_ops: int64;
+  write_bytes: int64;
+  write_ops: int64;
+  duration: int64;
+}
+
+let kib = 1024L
+let ( ** ) = Int64.mul
+let mib = kib ** 1024L
+let gib = mib ** 1024L
+let tib = gib ** 1024L
+
+let suffix = [
+  kib, "KiB";
+  mib, "MiB";
+  gib, "GiB";
+  tib, "TiB";
+]
+
+let add_suffix x =
+  List.fold_left (fun acc (y, label) ->
+      if Int64.div x y > 0L
+      then Printf.sprintf "%.1f %s" Int64.((to_float x) /. (to_float y)) label
+      else acc
+    ) (Printf.sprintf "%Ld bytes" x) suffix
+
+let pp_stats ppf s =
+  Fmt.pf ppf "%s bytes at %s/nanosec and %Lu IOPS/nanosec"
+    (add_suffix s.read_bytes)
+    (add_suffix Int64.(div s.read_bytes s.duration))
+    (Int64.div s.read_ops s.duration)
+
+module type CONCRETE =  Mirage_flow.S
+  with type error = [ `Msg of string ]
+   and type write_error = [ Mirage_flow.write_error | `Msg of string ]
+
+module Concrete (S: Mirage_flow.S) = struct
+  type error = [`Msg of string]
+  type write_error = [ Mirage_flow.write_error | `Msg of string]
+  type flow = S.flow
+
+  let pp_error ppf = function
+    | `Msg s -> Fmt.string ppf s
+
+  let pp_write_error ppf = function
+    | #error as e -> pp_error ppf e
+    | `Closed     -> Mirage_flow.pp_write_error ppf `Closed
+
+  let lift_read = function
+    | Ok x    -> Ok x
+    | Error e -> Error (`Msg (Fmt.str "%a" S.pp_error e))
+
+  let lift_write = function
+    | Ok ()         -> Ok ()
+    | Error `Closed -> Error `Closed
+    | Error e       -> Error (`Msg (Fmt.str "%a" S.pp_write_error e))
+
+  let read t = S.read t >|= lift_read
+  let write t b = S.write t b >|= lift_write
+  let writev t bs = S.writev t bs >|= lift_write
+  let shutdown t m = S.shutdown t m
+  let close t = S.close t
+end
+
+type time = int64
+
+type 'a stats_lwt = {
+  read_bytes: int64 ref;
+  read_ops: int64 ref;
+  write_bytes: int64 ref;
+  write_ops: int64 ref;
+  finish: time option ref;
+  start: time;
+  time: unit -> time;
+  t: (unit, 'a) result Lwt.t;
+}
+
+let stats_lwt t =
+  let duration : int64 = match !(t.finish) with
+    | None -> Int64.sub (t.time ()) t.start
+    | Some x -> Int64.sub x t.start
+  in {
+    read_bytes  = !(t.read_bytes);
+    read_ops    = !(t.read_ops);
+    write_bytes = !(t.write_bytes);
+    write_ops   = !(t.write_ops);
+    duration;
+  }
+
+module Copy (A: Mirage_flow.S) (B: Mirage_flow.S) =
+struct
+
+  type error = [`A of A.error | `B of B.write_error]
+
+  let pp_error ppf = function
+    | `A e -> A.pp_error ppf e
+    | `B e -> B.pp_write_error ppf e
+
+  let start (a: A.flow) (b: B.flow) =
+    let read_bytes = ref 0L in
+    let read_ops = ref 0L in
+    let write_bytes = ref 0L in
+    let write_ops = ref 0L in
+    let finish = ref None in
+    let start = Mirage_mtime.elapsed_ns () in
+    let rec loop () =
+      A.read a >>= function
+      | Error e ->
+        finish := Some (Mirage_mtime.elapsed_ns ());
+        Lwt.return (Error (`A e))
+      | Ok `Eof ->
+        finish := Some (Mirage_mtime.elapsed_ns ());
+        Lwt.return (Ok ())
+      | Ok (`Data buffer) ->
+        read_ops := Int64.succ !read_ops;
+        read_bytes := Int64.(add !read_bytes (of_int @@ Cstruct.length buffer));
+        B.write b buffer
+        >>= function
+        | Ok () ->
+          write_ops := Int64.succ !write_ops;
+          write_bytes := Int64.(add !write_bytes (of_int @@ Cstruct.length buffer));
+          loop ()
+        | Error e ->
+          finish := Some (Mirage_mtime.elapsed_ns ());
+          Lwt.return (Error (`B e))
+    in
+    {
+      read_bytes;
+      read_ops;
+      write_bytes;
+      write_ops;
+      finish;
+      start;
+      time = (fun () -> Mirage_mtime.elapsed_ns ());
+      t = loop ();
+    }
+
+  let wait t = t.t
+
+  let copy ~src:a ~dst:b =
+    let t = start a b in
+    wait t >|= function
+    | Ok ()   -> Ok (stats_lwt t)
+    | Error e -> Error e
+
+end
+
+module Proxy (A: Mirage_flow.S) (B: Mirage_flow.S) =
+struct
+
+  module A_to_B = Copy(A)(B)
+  module B_to_A = Copy(B)(A)
+
+  type error = [
+    | `A of A_to_B.error
+    | `B of B_to_A.error
+    | `A_and_B of A_to_B.error * B_to_A.error
+  ]
+
+  let pp_error ppf = function
+    | `A_and_B (e1, e2) ->
+      Fmt.pf ppf "flow proxy a: %a; flow proxy b: %a"
+        A_to_B.pp_error e1 B_to_A.pp_error e2
+    | `A e -> Fmt.pf ppf "flow proxy a: %a" A_to_B.pp_error e
+    | `B e -> Fmt.pf ppf "flow proxy b: %a" B_to_A.pp_error e
+
+  let proxy a b =
+    let a2b =
+      let t = A_to_B.start a b in
+      A_to_B.wait t >>= fun result ->
+      A.shutdown a `read >>= fun () ->
+      B.shutdown b `write >|= fun () ->
+      let stats = stats_lwt t in
+      match result with
+      | Ok ()   -> Ok stats
+      | Error e -> Error e
+    in
+    let b2a =
+      let t = B_to_A.start b a in
+      B_to_A.wait t >>= fun result ->
+      B.shutdown b `read >>= fun () ->
+      A.shutdown a `write >|= fun () ->
+      let stats = stats_lwt t in
+      match result with
+      | Ok ()   -> Ok stats
+      | Error e -> Error e
+    in
+    a2b >>= fun a_stats ->
+    b2a >|= fun b_stats ->
+    match a_stats, b_stats with
+    | Ok a_stats, Ok b_stats -> Ok (a_stats, b_stats)
+    | Error e1  , Error e2   -> Error (`A_and_B (e1, e2))
+    | Error e1  ,  _         -> Error (`A e1)
+    | _         , Error e2   -> Error (`B e2)
+
+end
+
+module F = struct
+
+  let (>>=) = Lwt.bind
+
+  type refill = Cstruct.t -> int -> int -> int Lwt.t
+
+  type error
+  let pp_error ppf (_:error) =
+    Fmt.string ppf "Mirage_flow_combinators.F.error"
+  type write_error = Mirage_flow.write_error
+  let pp_write_error = Mirage_flow.pp_write_error
+
+  let seq f1 f2 buf off len =
+    f1 buf off len >>= function
+    | 0 -> f2 buf off len
+    | n -> Lwt.return n
+
+  let zero _buf _off _len = Lwt.return 0
+
+  let rec iter fn = function
+    | []   -> zero
+    | h::t -> seq (fn h) (iter fn t)
+
+  type flow = {
+    close: unit -> unit Lwt.t;
+    input: refill;
+    output: refill;
+    mutable buf: Cstruct.t;
+    mutable ic_closed: bool;
+    mutable oc_closed: bool;
+  }
+
+  let default_buffer_size = 4096
+
+  let make ?(close=fun () -> Lwt.return_unit) ?input ?output () =
+    let buf = Cstruct.create default_buffer_size in
+    let ic_closed = input = None in
+    let oc_closed = output = None in
+    let input = match input with None -> zero | Some x -> x in
+    let output = match output with None -> zero | Some x -> x in
+    { close; input; output; buf; ic_closed; oc_closed; }
+
+  let input_fn len blit str =
+    let str_off = ref 0 in
+    let str_len = len str in
+    fun buf off len ->
+      if !str_off >= str_len then Lwt.return 0
+      else (
+        let len = min (str_len - !str_off) len in
+        blit str !str_off buf off len;
+        str_off := !str_off + len;
+        Lwt.return len
+      )
+
+  let output_fn len blit str =
+    let str_off = ref 0 in
+    let str_len = len str in
+    fun buf off len ->
+      if !str_off >= str_len then Lwt.return 0
+      else (
+        let len = min (str_len - !str_off) len in
+        blit buf off str !str_off len;
+        str_off := !str_off + len;
+        Lwt.return len
+      )
+
+  let mk fn_i fn_o ?input ?output () =
+    let input = match input with None -> None | Some x -> Some (fn_i x) in
+    let output = match output with None -> None | Some x -> Some (fn_o x) in
+    make ?input ?output ()
+
+  let input_string = input_fn String.length Cstruct.blit_from_string
+  let output_bytes = output_fn Bytes.length Cstruct.blit_to_bytes
+  let string = mk input_string output_bytes
+
+  let input_cstruct = input_fn Cstruct.length Cstruct.blit
+  let output_cstruct = output_fn Cstruct.length Cstruct.blit
+  let cstruct = mk input_cstruct output_cstruct
+
+  let input_strings = iter input_string
+  let output_bytess = iter output_bytes
+  let strings = mk input_strings output_bytess
+
+  let input_cstructs = iter input_cstruct
+  let output_cstructs = iter output_cstruct
+  let cstructs = mk input_cstructs output_cstructs
+
+  let refill ch =
+    if Cstruct.length ch.buf = 0 then (
+      let buf = Cstruct.create default_buffer_size in
+      ch.buf <- buf
+    )
+
+  let read ch =
+    if ch.ic_closed then Lwt.return @@ Ok `Eof
+    else (
+      refill ch;
+      ch.input ch.buf 0 default_buffer_size >>= fun n ->
+      if n = 0 then (
+        ch.ic_closed <- true;
+        Lwt.return (Ok `Eof);
+      ) else (
+        let ret = Cstruct.sub ch.buf 0 n in
+        let buf = Cstruct.shift ch.buf n in
+        ch.buf <- buf;
+        Lwt.return (Ok (`Data ret))
+      )
+    )
+
+  let write ch buf =
+    if ch.oc_closed then Lwt.return @@ Error `Closed
+    else (
+      let len = Cstruct.length buf in
+      let rec aux off =
+        if off = len then Lwt.return (Ok ())
+        else (
+          ch.output buf off (len - off) >>= fun n ->
+          if n = 0 then (
+            ch.oc_closed <- true;
+            Lwt.return @@ Error `Closed
+          ) else aux (off+n)
+        )
+      in
+      aux 0
+    )
+
+  let writev ch bufs =
+    if ch.oc_closed then Lwt.return @@ Error `Closed
+    else
+      let rec aux = function
+        | []   -> Lwt.return (Ok ())
+        | h::t ->
+          write ch h >>= function
+          | Error e -> Lwt.return (Error e)
+          | Ok ()   -> aux t
+      in
+      aux bufs
+
+  let shutdown ch mode =
+    (match mode with
+     | `read -> ch.ic_closed <- true
+     | `write -> ch.oc_closed <- true
+     | `read_write ->
+       ch.ic_closed <- true;
+       ch.oc_closed <- true);
+    Lwt.return_unit
+
+  let close ch =
+    ch.ic_closed <- true;
+    ch.oc_closed <- true;
+    ch.close ()
+
+end
+
+type error = [`Msg of string]
+type write_error = [ Mirage_flow.write_error | error ]
+let pp_error ppf (`Msg s) = Fmt.string ppf s
+
+let pp_write_error ppf = function
+  | #Mirage_flow.write_error as e -> Mirage_flow.pp_write_error ppf e
+  | #error as e                   -> pp_error ppf e
+
+type flow =
+  | Flow: string * (module CONCRETE with type flow = 'a) * 'a -> flow
+
+type t = flow
+
+let create (type a) (module M: Mirage_flow.S with type flow = a) t name =
+  let m = (module Concrete(M): CONCRETE with type flow = a) in
+  Flow (name, m , t)
+
+let read (Flow (_, (module F), flow)) = F.read flow
+let write (Flow (_, (module F), flow)) b = F.write flow b
+let writev (Flow (_, (module F), flow)) b = F.writev flow b
+let close (Flow (_, (module F), flow)) = F.close flow
+let shutdown (Flow (_, (module F), flow)) m = F.shutdown flow m
+let pp ppf (Flow (name, _, _)) = Fmt.string ppf name
+
+let forward ?(verbose=false) ~src ~dst () =
+  let rec loop () =
+    read src >>= function
+    | Ok `Eof ->
+      Log.err (fun l -> l "forward[%a => %a] EOF" pp src pp dst);
+      Lwt.return_unit
+    | Error e ->
+      Log.err (fun l -> l "forward[%a => %a] %a" pp src pp dst pp_error e);
+      Lwt.return_unit
+    | Ok (`Data buf) ->
+      Log.debug (fun l ->
+          let payload =
+            if verbose then Fmt.str "[%S]" @@ Cstruct.to_string buf
+            else Fmt.str "%d bytes" (Cstruct.length buf)
+          in
+          l "forward[%a => %a] %s" pp src pp dst payload);
+      write dst buf >>= function
+      | Ok ()   -> loop ()
+      | Error e ->
+        Log.err (fun l -> l "forward[%a => %a] %a"
+                    pp src pp dst pp_write_error e);
+        Lwt.return_unit
+  in
+  loop ()
+
+let proxy ?verbose f1 f2 =
+  Lwt.join [
+    forward ?verbose ~src:f1 ~dst:f2 ();
+    forward ?verbose ~src:f2 ~dst:f1 ();
+  ]