unikernel/duniverse/httpun/lib/body.ml

@@ -0,0 +1,249 @@
+(*----------------------------------------------------------------------------
+    Copyright (c) 2018 Inhabited Type LLC.
+
+    All rights reserved.
+
+    Redistribution and use in source and binary forms, with or without
+    modification, are permitted provided that the following conditions
+    are met:
+
+    1. Redistributions of source code must retain the above copyright
+       notice, this list of conditions and the following disclaimer.
+
+    2. Redistributions in binary form must reproduce the above copyright
+       notice, this list of conditions and the following disclaimer in the
+       documentation and/or other materials provided with the distribution.
+
+    3. Neither the name of the author nor the names of his contributors
+       may be used to endorse or promote products derived from this software
+       without specific prior written permission.
+
+    THIS SOFTWARE IS PROVIDED BY THE CONTRIBUTORS ``AS IS'' AND ANY EXPRESS
+    OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
+    WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+    DISCLAIMED.  IN NO EVENT SHALL THE AUTHORS OR CONTRIBUTORS BE LIABLE FOR
+    ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+    DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS
+    OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION)
+    HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT,
+    STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN
+    ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
+    POSSIBILITY OF SUCH DAMAGE.
+  ----------------------------------------------------------------------------*)
+
+module Reader = struct
+  type t =
+    { faraday                        : Faraday.t
+    ; mutable read_scheduled         : bool
+    ; mutable on_eof                 : unit -> unit
+    ; mutable eof_has_been_called    : bool
+    ; mutable on_read                : Bigstringaf.t -> off:int -> len:int -> unit
+    ; when_ready_to_read             : Optional_thunk.t
+    }
+
+  let default_on_eof         = Sys.opaque_identity (fun () -> ())
+  let default_on_read        = Sys.opaque_identity (fun _ ~off:_ ~len:_ -> ())
+
+  let create buffer ~when_ready_to_read =
+    { faraday                = Faraday.of_bigstring buffer
+    ; read_scheduled         = false
+    ; eof_has_been_called    = false
+    ; on_eof                 = default_on_eof
+    ; on_read                = default_on_read
+    ; when_ready_to_read
+    }
+
+  let create_empty () =
+    let t = create Bigstringaf.empty ~when_ready_to_read:Optional_thunk.none in
+    Faraday.close t.faraday;
+    t
+
+  let is_closed t =
+    Faraday.is_closed t.faraday
+
+  let unsafe_faraday t =
+    t.faraday
+
+  let ready_to_read t = Optional_thunk.call_if_some t.when_ready_to_read
+
+  let rec do_execute_read t on_eof on_read =
+    match Faraday.operation t.faraday with
+    | `Yield           -> ()
+    | `Close           ->
+      t.read_scheduled <- false;
+      t.on_eof         <- default_on_eof;
+      t.on_read        <- default_on_read;
+      if not t.eof_has_been_called then begin
+        t.eof_has_been_called <- true;
+        on_eof ()
+      end
+    (* [Faraday.operation] never returns an empty list of iovecs *)
+    | `Writev []       -> assert false
+    | `Writev (iovec::_) ->
+      t.read_scheduled <- false;
+      t.on_eof         <- default_on_eof;
+      t.on_read        <- default_on_read;
+      let { IOVec.buffer; off; len } = iovec in
+      Faraday.shift t.faraday len;
+      on_read buffer ~off ~len;
+      execute_read t
+  and execute_read t =
+    if t.read_scheduled then do_execute_read t t.on_eof t.on_read
+
+let schedule_read t ~on_eof ~on_read =
+  if t.read_scheduled
+  then failwith "Body.schedule_read: reader already scheduled";
+  if not (is_closed t)
+  then begin
+    t.read_scheduled <- true;
+    t.on_eof         <- on_eof;
+    t.on_read        <- on_read;
+  end;
+  do_execute_read t on_eof on_read;
+  ready_to_read t
+
+  let close t =
+    Faraday.close t.faraday;
+    execute_read t;
+    ready_to_read t
+  ;;
+
+  let has_pending_output t = Faraday.has_pending_output t.faraday
+
+  let is_read_scheduled t = t.read_scheduled
+end
+
+module Writer = struct
+  type encoding =
+    | Identity
+    | Chunked of { mutable written_final_chunk : bool }
+
+  type t =
+    { faraday             : Faraday.t
+    ; encoding            : encoding
+    ; writer              : Serialize.Writer.t
+    ; mutable buffered_bytes : int
+    }
+
+  let of_faraday faraday ~encoding ~writer =
+    let encoding =
+      match encoding with
+      | `Fixed _ | `Close_delimited -> Identity
+      | `Chunked -> Chunked { written_final_chunk = false }
+    in
+    { faraday
+    ; encoding
+    ; writer
+    ; buffered_bytes = 0
+    }
+
+  let create buffer ~encoding =
+    of_faraday (Faraday.of_bigstring buffer) ~encoding
+
+  let create_empty ~writer =
+    let t =
+      create
+        Bigstringaf.empty
+        ~encoding:(`Fixed 0)
+        ~writer
+    in
+    Faraday.close t.faraday;
+    t
+
+  let write_char t c =
+    if not (Faraday.is_closed t.faraday) then
+      Faraday.write_char t.faraday c
+
+  let write_string t ?off ?len s =
+    if not (Faraday.is_closed t.faraday) then
+      Faraday.write_string ?off ?len t.faraday s
+
+  let write_bigstring t ?off ?len b =
+    if not (Faraday.is_closed t.faraday) then
+      Faraday.write_bigstring ?off ?len t.faraday b
+
+  let schedule_bigstring t ?off ?len (b:Bigstringaf.t) =
+    if not (Faraday.is_closed t.faraday) then
+      Faraday.schedule_bigstring ?off ?len t.faraday b
+
+  let ready_to_write t = Serialize.Writer.wakeup t.writer
+
+  let flush t kontinue =
+    if Serialize.Writer.is_closed t.writer then
+      kontinue `Closed
+    else begin
+      Faraday.flush_with_reason t.faraday (function
+        | Drain -> kontinue `Closed
+        | Nothing_pending | Shift -> Serialize.Writer.flush t.writer kontinue);
+      ready_to_write t
+    end
+
+  let is_closed t =
+    Faraday.is_closed t.faraday
+
+  let close_and_drain t =
+    Faraday.close t.faraday;
+    (* Resolve all pending flushes *)
+    ignore (Faraday.drain t.faraday : int)
+
+  let close t =
+    Serialize.Writer.unyield t.writer;
+    Faraday.close t.faraday;
+    ready_to_write t;
+  ;;
+
+  let force_close t =
+    begin match t.encoding with
+    | Chunked t -> t.written_final_chunk <- true
+    | Identity -> ()
+    end;
+    close t
+
+  let has_pending_output t =
+    (* Force another write poll to make sure that the final chunk is emitted for
+       chunk-encoded bodies. *)
+    let faraday_has_output = Faraday.has_pending_output t.faraday in
+    let additional_encoding_output =
+      match t.encoding with
+      | Identity -> false
+      | Chunked { written_final_chunk } ->
+        Faraday.is_closed t.faraday && not written_final_chunk
+    in
+    faraday_has_output || additional_encoding_output
+
+  let requires_output t =
+    not (is_closed t) || has_pending_output t
+
+  let transfer_to_writer t =
+    let faraday = t.faraday in
+    if Serialize.Writer.is_closed t.writer then
+      close_and_drain t
+    else
+    begin match Faraday.operation faraday with
+    | `Yield -> ()
+    | `Close ->
+      (match t.encoding with
+       | Identity -> ()
+       | Chunked ({ written_final_chunk } as chunked) ->
+         if not written_final_chunk then begin
+           chunked.written_final_chunk <- true;
+           Serialize.Writer.schedule_chunk t.writer [];
+         end);
+    | `Writev iovecs ->
+      begin match IOVec.shiftv iovecs t.buffered_bytes with
+      | []     -> ()
+      | iovecs ->
+        let lengthv  = IOVec.lengthv iovecs in
+        t.buffered_bytes <- t.buffered_bytes + lengthv;
+        begin match t.encoding with
+        | Identity  -> Serialize.Writer.schedule_fixed t.writer iovecs
+        | Chunked _ -> Serialize.Writer.schedule_chunk t.writer iovecs
+        end;
+        Serialize.Writer.flush t.writer (function
+          | `Closed -> close_and_drain t
+          | `Written ->
+            Faraday.shift faraday lengthv;
+            t.buffered_bytes <- t.buffered_bytes - lengthv)
+      end
+    end
+end