h2_mirage.ml

open Lwt.Infix

module type S =
  sig
    module Mirage : Mirage_flow.S
    type data = (Cstruct.t Mirage_flow.or_eof, Mirage.error) result
    type t
    val create : Mirage.flow -> t
    val mirage_flow : t -> Mirage.flow
    val read :
      t -> int -> data Lwt.t
  end

module Gluten_flow (Flow : Mirage_flow.S) :
  S with module Mirage = Flow = struct

  module Mirage = Flow

  type data = (Cstruct.t Mirage_flow.or_eof, Mirage.error) result

  type t = {
    flow: Flow.flow;
    mutable buf: Cstruct.t;
  }

  let create flow = {
    flow = flow;
    buf = Cstruct.empty;
  }

  let mirage_flow flow = flow.flow

  let read flow len =
    let trunc buf =
      match Cstruct.length buf > len with
      | false -> buf
      | true ->
        let head, rest = Cstruct.split buf len in
        flow.buf <- rest;
        head
    in
    let buffered_data =
      match Cstruct.is_empty flow.buf with
      | true -> None
      | false -> 
        let buf = flow.buf in
        flow.buf <- Cstruct.empty;
        Some (Ok (`Data (trunc buf)))
    in
    match buffered_data with
    | Some data ->
      Lwt.return data
    | None -> Flow.read flow.flow >|= fun data ->
      assert (Cstruct.is_empty flow.buf);
      match data with
      | Ok (`Data buf) -> Ok (`Data (trunc buf))
      | x -> x

end

module Make_IO (Flow : S) :
  Gluten_lwt.IO with type socket = Flow.t and type addr = unit = struct

  type socket = Flow.t

  type addr = unit

  let shutdown flow = Flow.(Mirage.close (mirage_flow flow))

  let shutdown_receive flow = Lwt.async (fun () -> shutdown flow)

  let shutdown_send flow = Lwt.async (fun () -> shutdown flow)

  let close = shutdown

  let read flow bigstring ~off ~len =
    Lwt.catch
      (fun () ->
        Flow.read flow len >|= function
        | Ok (`Data buf) ->
          Bigstringaf.blit
            buf.buffer
            ~src_off:buf.off
            bigstring
            ~dst_off:off
            ~len:buf.len;
          `Ok buf.len
        | Ok `Eof ->
          `Eof
        | Error error ->
          failwith (Format.asprintf "%a" Flow.Mirage.pp_error error))
      (fun exn -> shutdown flow >>= fun () -> Lwt.fail exn)

  let writev flow iovecs =
    let cstruct_iovecs =
      List.map
        (fun { Faraday.buffer; off; len } -> Cstruct.of_bigarray ~off ~len buffer)
        iovecs
    in
    let len = Cstruct.lenv cstruct_iovecs in
    let data = Cstruct.create_unsafe len in
    let n, _ = Cstruct.fillv cstruct_iovecs data in
    assert (n = len);
    Lwt.catch
      (fun () ->
        Flow.Mirage.write (Flow.mirage_flow flow) data >|= fun x ->
        match x with
        | Ok () ->
          `Ok (Cstruct.lenv cstruct_iovecs)
        | Error `Closed ->
          `Closed
        | Error other_error ->
          raise (Failure (Format.asprintf "%a" Flow.Mirage.pp_write_error other_error)))
      (fun exn -> shutdown flow >>= fun () -> Lwt.fail exn)
end

module Client (Flow : S) = H2_lwt.Client (Gluten_lwt.Client (Make_IO (Flow)))