unikernel/duniverse/ocaml-tls/eio/tests/mock_socket.ml

@@ -0,0 +1,94 @@
+open Eio.Std
+
+module W = Eio.Buf_write
+
+let src = Logs.Src.create "mock-socket" ~doc:"Test socket"
+module Log = (val Logs.src_log src : Logs.LOG)
+
+type transmit_amount = [`Bytes of int | `Drain]
+
+type ty = [`Mock_tls | Eio.Flow.two_way_ty | Eio.Resource.close_ty]
+type t = ty r
+
+let rec takev len = function
+  | [] -> []
+  | x :: xs ->
+    if len = 0 then []
+    else if Cstruct.length x >= len then [Cstruct.sub x 0 len]
+    else x :: takev (len - Cstruct.length x) xs
+
+module Impl = struct
+  type t = {
+    to_peer : W.t;
+    from_peer : W.t;
+    label : string;
+    output_sizes : transmit_amount Eio.Stream.t;
+  }
+
+  let create ~to_peer ~from_peer label = {
+    to_peer;
+    from_peer;
+    label;
+    output_sizes = Eio.Stream.create max_int;
+  }
+
+  let transmit t x =
+    Eio.Stream.add t.output_sizes x
+
+  let single_write t bufs =
+    let size =
+      match Eio.Stream.take t.output_sizes with
+      | `Drain -> Eio.Stream.add t.output_sizes `Drain; Cstruct.lenv bufs
+      | `Bytes size -> size
+    in
+    let bufs = takev size bufs in
+    List.iter (W.cstruct t.to_peer) bufs;
+    let len = Cstruct.lenv bufs in
+    Log.info (fun f -> f "%s: wrote %d bytes to network" t.label len);
+    len
+
+  let copy t ~src = Eio.Flow.Pi.simple_copy ~single_write t ~src
+
+  let single_read t buf =
+    let batch = W.await_batch t.from_peer in
+    let got, _ = Cstruct.fillv ~src:batch ~dst:buf in
+    Log.info (fun f -> f "%s: read %d bytes from network" t.label got);
+    W.shift t.from_peer got;
+    got
+
+  let shutdown t = function
+    | `Send -> 
+      Log.info (fun f -> f "%s: close writer" t.label);
+      W.close t.to_peer
+    | _ -> failwith "Not implemented"
+
+  let close t =
+    Log.info (fun f -> f "%s: close connection" t.label)
+
+  let read_methods = []
+
+  type (_, _, _) Eio.Resource.pi += Raw : ('t, 't -> t, ty) Eio.Resource.pi
+  let raw (Eio.Resource.T (t, ops)) = Eio.Resource.get ops Raw t
+end
+
+let handler =
+  Eio.Resource.handler (
+    H (Impl.Raw, Fun.id) ::
+    H (Eio.Resource.Close, Impl.close) ::
+    Eio.Resource.bindings (Eio.Flow.Pi.two_way (module Impl))
+  )
+
+let transmit t x =
+  let t = Impl.raw t in
+  Impl.transmit t x
+
+let create ~from_peer ~to_peer label =
+  let t = Impl.create ~from_peer ~to_peer label in
+  Eio.Resource.T (t, handler)
+
+let create_pair () =
+  let to_a = W.create 100 in
+  let to_b = W.create 100 in
+  let a = create ~from_peer:to_a ~to_peer:to_b "client" in
+  let b = create ~from_peer:to_b ~to_peer:to_a "server" in
+  a, b