403 lines
15 KiB
OCaml
403 lines
15 KiB
OCaml
module type RUNTIME = sig
|
|
type t
|
|
|
|
val next_read_operation : t -> [ `Read | `Yield | `Close | `Upgrade ]
|
|
(** [next_read_connection t] returns a value describing the next operation
|
|
that the caller should conduit on behalf of the connection. *)
|
|
|
|
val read : t -> Bigstringaf.t -> off:int -> len:int -> int
|
|
(** [read t bigstring ~off ~len] reads bytes of input from the provided range
|
|
of [bigstring] an returns the number of bytes consumed by the connection.
|
|
{!read} should be called after {!next_read_operation} returns a [`Read]
|
|
value an additional input is available for the connection to consume. *)
|
|
|
|
val read_eof : t -> Bigstringaf.t -> off:int -> len:int -> int
|
|
(** [read_eof t bigstring ~off ~len] reads bytes of input from the provided
|
|
range of [bigstring] and returns the number of bytes consumed by the
|
|
connection. {!read_eof} should be called after {!next_read_operation}
|
|
returns a [`Read] and an EOF has been received from the communication
|
|
channel. The connection will attempt to consume any buffered input and
|
|
then shutdown the HTTP parser for the connection. *)
|
|
|
|
val yield_reader : t -> (unit -> unit) -> unit
|
|
(** [yield_reader t continue] registers with the connection to call [continue]
|
|
when reading should resume. {!yield_reader} should be called after
|
|
{!next_read_operation} returns a [`Yield] value. *)
|
|
|
|
val next_write_operation :
|
|
t ->
|
|
[ `Write of Bigstringaf.t Faraday.iovec list
|
|
| `Yield
|
|
| `Close of int
|
|
| `Upgrade ]
|
|
(** [next_write_operation t] returns a value describing the next operation
|
|
that the caller should conduct on behalf the connection. *)
|
|
|
|
val report_write_result : t -> [ `Ok of int | `Closed ] -> unit
|
|
(** [report_write_result t result] reports the result of the latest write
|
|
attempt to the connection. {!report_write_result} should be called after a
|
|
call to {!next_write_operation} that returns a [`Write buffer] value.
|
|
|
|
- [`Ok n] indicates that the caller successfully wrote [n] bytes of output
|
|
from the buffer that the caller was provided by {!next_write_operation}
|
|
that returns a [`Write buffer] value.
|
|
- [`Closed] indicates that the output destination will no longer accept
|
|
bytes from the write processor. *)
|
|
|
|
val yield_writer : t -> (unit -> unit) -> unit
|
|
(** [yield_writer t continue] registers with the connection to call [continue]
|
|
when writing should resume. {!yield_writer} should be called after
|
|
{!next_write_operation} returns a [`Yield] value. *)
|
|
|
|
val report_exn : t -> exn -> unit
|
|
(** [report_exn t exn] reports that an error [exn] has been caught and that it
|
|
has been attributed to [t]. Calling this function will switch [t] into an
|
|
error state. Depending on the tate [t] is transitioning from, it may call
|
|
its error handler before terminating the connection. *)
|
|
|
|
val is_closed : t -> bool
|
|
(** [is_closed t] is [true] if both the read and write processors have been
|
|
shutdown. When this is the case {!next_read_operation} will return
|
|
[`Close _] and {!next_write_operation} will return a [`Write _] until all
|
|
buffered output has been flushed, at which point it will return [`Close].
|
|
*)
|
|
|
|
val shutdown : t -> unit
|
|
(** [shutdown t] asks to shutdown the connection. *)
|
|
end
|
|
|
|
type 'conn runtime = (module RUNTIME with type t = 'conn)
|
|
|
|
exception Flow of string
|
|
exception Flow_write of string
|
|
|
|
let src = Logs.Src.create "paf-flow"
|
|
|
|
module Log_flow = (val Logs.src_log src : Logs.LOG)
|
|
|
|
module Make (Flow : Mirage_flow.S) = struct
|
|
type flow = {
|
|
flow : Flow.flow;
|
|
queue : (char, Bigarray.int8_unsigned_elt) Ke.Rke.t;
|
|
mutable rd_closed : bool;
|
|
mutable wr_closed : bool;
|
|
}
|
|
|
|
let create flow =
|
|
let queue = Ke.Rke.create ~capacity:0x1000 Bigarray.char in
|
|
Lwt.return { flow; queue; rd_closed = false; wr_closed = false }
|
|
|
|
let safely_close flow =
|
|
if flow.rd_closed && flow.wr_closed
|
|
then (
|
|
Log_flow.debug (fun m -> m "Close the connection.") ;
|
|
Flow.close flow.flow)
|
|
else Lwt.return ()
|
|
|
|
let blit src src_off dst dst_off len =
|
|
let dst = Cstruct.of_bigarray ~off:dst_off ~len dst in
|
|
Cstruct.blit src src_off dst 0 len
|
|
|
|
open Lwt.Infix
|
|
|
|
type eof = [ `Eof ]
|
|
|
|
let recv flow ~report_error ~report_closed ~read ~read_eof =
|
|
Ke.Rke.compress flow.queue ;
|
|
Flow.read flow.flow >>= function
|
|
| (Error _ | Ok #eof) as v ->
|
|
flow.rd_closed <- true ;
|
|
safely_close flow >>= fun () ->
|
|
let _shift =
|
|
match
|
|
Ke.Rke.compress flow.queue ;
|
|
Ke.Rke.N.peek flow.queue
|
|
with
|
|
| [] -> read_eof Bigstringaf.empty ~off:0 ~len:0
|
|
| [ slice ] -> read_eof slice ~off:0 ~len:(Bigstringaf.length slice)
|
|
| _ -> assert false
|
|
(* XXX(dinosaure): impossible due to [compress]. *) in
|
|
(match v with
|
|
| Ok `Eof -> report_closed ()
|
|
| Error err -> report_error err) ;
|
|
Lwt.return `Closed
|
|
| Ok (`Data v) ->
|
|
let len = Cstruct.length v in
|
|
Ke.Rke.N.push flow.queue ~blit ~length:Cstruct.length ~off:0 ~len v ;
|
|
let[@warning "-8"] (slice :: _) = Ke.Rke.N.peek flow.queue in
|
|
let shift = read slice ~off:0 ~len:(Bigstringaf.length slice) in
|
|
Ke.Rke.N.shift_exn flow.queue shift ;
|
|
Lwt.return `Continue
|
|
|
|
let writev ~report_error flow iovecs =
|
|
let iovecs =
|
|
List.map
|
|
(fun { Faraday.buffer; off; len } ->
|
|
Cstruct.to_string (Cstruct.of_bigarray buffer ~off ~len) ~off:0 ~len)
|
|
iovecs in
|
|
let iovecs = List.map Cstruct.of_string iovecs in
|
|
(* XXX(dinosaure): the copy is needed:
|
|
1) [Mirage_flow.S] explicitly says that [write] takes the ownership on
|
|
the given [Cstruct.t]
|
|
2) [ocaml-h2] wants to keep the ownership on given [Faraday.iovec]s
|
|
|
|
To protect one from the other, copying is necessary. *)
|
|
Log_flow.debug (fun m ->
|
|
m "Start to write %d byte(s)."
|
|
(List.fold_left (fun acc cs -> Cstruct.length cs + acc) 0 iovecs)) ;
|
|
Flow.writev flow.flow iovecs >>= function
|
|
| Ok () ->
|
|
Lwt.return
|
|
(`Ok (List.fold_left (fun acc cs -> acc + Cstruct.length cs) 0 iovecs))
|
|
| Error err ->
|
|
Log_flow.err (fun m ->
|
|
m "Got an error when we wrote something: %a." Flow.pp_write_error
|
|
err) ;
|
|
report_error err ;
|
|
flow.wr_closed <- true ;
|
|
safely_close flow >>= fun () -> Lwt.return `Closed
|
|
|
|
let send ~report_error flow iovecs =
|
|
if flow.wr_closed
|
|
then safely_close flow >>= fun () -> Lwt.return `Closed
|
|
else writev ~report_error flow iovecs
|
|
|
|
let close flow =
|
|
match (flow.rd_closed, flow.wr_closed) with
|
|
| true, true -> Lwt.return_unit
|
|
| _ ->
|
|
flow.rd_closed <- true ;
|
|
flow.wr_closed <- true ;
|
|
Flow.close flow.flow
|
|
end
|
|
|
|
let src = Logs.Src.create "paf-server"
|
|
|
|
module Log_server = (val Logs.src_log src : Logs.LOG)
|
|
|
|
module Server (Flow : Mirage_flow.S) (Runtime : RUNTIME) : sig
|
|
val server : Runtime.t -> Flow.flow -> unit Lwt.t
|
|
end = struct
|
|
module Easy_flow = Make (Flow)
|
|
open Lwt.Infix
|
|
|
|
let to_flow_exception err : exn = Flow (Fmt.str "%a" Flow.pp_error err)
|
|
|
|
let to_flow_write_exception err : exn =
|
|
Flow_write (Fmt.str "%a" Flow.pp_write_error err)
|
|
|
|
let server connection flow =
|
|
Easy_flow.create flow >>= fun flow ->
|
|
let rd_exit, notify_rd_exit = Lwt.wait () in
|
|
let wr_exit, notify_wr_exit = Lwt.wait () in
|
|
|
|
let rec rd_fiber () =
|
|
let report_error err =
|
|
Runtime.report_exn connection (to_flow_exception err) in
|
|
let rec go () =
|
|
Log_server.debug (fun m -> m "Compute next read operation.") ;
|
|
match Runtime.next_read_operation connection with
|
|
| `Upgrade -> failwith "Unimplemented"
|
|
| `Read ->
|
|
Log_server.debug (fun m -> m "next read operation: `read") ;
|
|
Easy_flow.recv flow ~report_error ~report_closed:ignore
|
|
~read:(Runtime.read connection)
|
|
~read_eof:(Runtime.read_eof connection)
|
|
>>= fun _ -> Lwt.pause () >>= go
|
|
| `Yield ->
|
|
Log_server.debug (fun m -> m "next read operation: `yield") ;
|
|
Runtime.yield_reader connection rd_fiber ;
|
|
Lwt.pause ()
|
|
| `Close ->
|
|
Log_server.debug (fun m -> m "next read operation: `close") ;
|
|
Lwt.wakeup_later notify_rd_exit () ;
|
|
Flow.shutdown flow.flow `read in
|
|
Lwt.async @@ fun () ->
|
|
Lwt.catch go (fun exn ->
|
|
Runtime.report_exn connection exn ;
|
|
Lwt.return_unit) in
|
|
let rec wr_fiber () =
|
|
let report_error err =
|
|
Runtime.report_exn connection (to_flow_write_exception err) in
|
|
let rec go () =
|
|
Log_server.debug (fun m -> m "Compute next write operation.") ;
|
|
match Runtime.next_write_operation connection with
|
|
| `Upgrade -> failwith "Unimplemented"
|
|
| `Write iovecs ->
|
|
Log_server.debug (fun m -> m "next write operation: `write") ;
|
|
Easy_flow.send ~report_error flow iovecs >>= fun res ->
|
|
Runtime.report_write_result connection res ;
|
|
Lwt.pause () >>= go
|
|
| `Yield ->
|
|
Log_server.debug (fun m -> m "next write operation: `yield") ;
|
|
Runtime.yield_writer connection wr_fiber ;
|
|
Lwt.pause ()
|
|
| `Close _ ->
|
|
Log_server.debug (fun m -> m "next write operation: `close") ;
|
|
Lwt.wakeup_later notify_wr_exit () ;
|
|
Flow.shutdown flow.flow `write in
|
|
Lwt.async @@ fun () ->
|
|
Lwt.catch go (fun exn ->
|
|
(* Runtime.report_write_result connection `Closed ; *)
|
|
Runtime.report_exn connection exn ;
|
|
Lwt.return_unit) in
|
|
rd_fiber () ;
|
|
wr_fiber () ;
|
|
Lwt.join [ rd_exit; wr_exit ] >>= fun () ->
|
|
Log_server.debug (fun m -> m "End of transmission.") ;
|
|
Easy_flow.close flow
|
|
end
|
|
|
|
let src = Logs.Src.create "paf-client"
|
|
|
|
module Log_client = (val Logs.src_log src : Logs.LOG)
|
|
|
|
module Client (Flow : Mirage_flow.S) (Runtime : RUNTIME) : sig
|
|
val run : Runtime.t -> Flow.flow -> unit Lwt.t
|
|
end = struct
|
|
open Lwt.Infix
|
|
module Easy_flow = Make (Flow)
|
|
|
|
let to_flow_exception err : exn = Flow (Fmt.str "%a" Flow.pp_error err)
|
|
|
|
let to_flow_write_exception err : exn =
|
|
Flow_write (Fmt.str "%a" Flow.pp_write_error err)
|
|
|
|
let run connection flow =
|
|
Easy_flow.create flow >>= fun flow ->
|
|
let rd_exit, notify_rd_exit = Lwt.wait () in
|
|
let wr_exit, notify_wr_exit = Lwt.wait () in
|
|
|
|
let rec rd_fiber () =
|
|
let report_error err =
|
|
Runtime.report_exn connection (to_flow_exception err) in
|
|
let rec go () =
|
|
match Runtime.next_read_operation connection with
|
|
| `Upgrade -> failwith "Unimplemented"
|
|
| `Read ->
|
|
Log_client.debug (fun m -> m "next read operation: `read") ;
|
|
Easy_flow.recv flow ~report_error ~report_closed:ignore
|
|
~read:(Runtime.read connection)
|
|
~read_eof:(Runtime.read_eof connection)
|
|
>>= fun _ -> Lwt.pause () >>= go
|
|
| `Yield ->
|
|
Log_client.debug (fun m -> m "next read operation: `yield") ;
|
|
Runtime.yield_reader connection rd_fiber ;
|
|
Lwt.pause ()
|
|
| `Close ->
|
|
Log_client.debug (fun m -> m "next read operation: `close.") ;
|
|
Lwt.wakeup_later notify_rd_exit () ;
|
|
flow.Easy_flow.rd_closed <- true ;
|
|
Easy_flow.safely_close flow in
|
|
Lwt.async @@ fun () ->
|
|
Lwt.catch go (fun exn ->
|
|
Runtime.report_exn connection exn ;
|
|
Lwt.return_unit) in
|
|
let rec wr_fiber () =
|
|
let report_error err =
|
|
Runtime.report_exn connection (to_flow_write_exception err) in
|
|
let rec go () =
|
|
match Runtime.next_write_operation connection with
|
|
| `Upgrade -> failwith "Unimplemented"
|
|
| `Write iovecs ->
|
|
Log_client.debug (fun m -> m "next write operation: `write.") ;
|
|
Easy_flow.send ~report_error flow iovecs >>= fun res ->
|
|
Runtime.report_write_result connection res ;
|
|
Lwt.pause () >>= go
|
|
| `Yield ->
|
|
Log_client.debug (fun m -> m "next write operation: `yield.") ;
|
|
Runtime.yield_writer connection wr_fiber ;
|
|
Lwt.pause ()
|
|
| `Close _ ->
|
|
Log_client.debug (fun m -> m "next write operation: `close.") ;
|
|
Lwt.wakeup_later notify_wr_exit () ;
|
|
flow.Easy_flow.wr_closed <- true ;
|
|
Easy_flow.safely_close flow in
|
|
Lwt.async @@ fun () ->
|
|
Lwt.catch go (fun exn ->
|
|
Runtime.report_exn connection exn ;
|
|
Lwt.return ()) in
|
|
wr_fiber () ;
|
|
rd_fiber () ;
|
|
Lwt.join [ rd_exit; wr_exit ] >>= fun () ->
|
|
Log_client.debug (fun m -> m "End of transmission.") ;
|
|
Easy_flow.close flow
|
|
end
|
|
|
|
type impl = Runtime : 'conn runtime * 'conn -> impl
|
|
|
|
type 't service =
|
|
| Service : {
|
|
accept : 't -> ('socket, ([> `Closed ] as 'error)) result Lwt.t;
|
|
handshake : 'socket -> ('flow, ([> `Closed ] as 'error)) result Lwt.t;
|
|
connection : 'flow -> (Mimic.flow * impl, 'error) result Lwt.t;
|
|
close : 't -> unit Lwt.t;
|
|
}
|
|
-> 't service
|
|
|
|
and ('t, 'socket, 'flow, 'error) posix = {
|
|
accept : 't -> ('socket, 'error) result Lwt.t;
|
|
handshake : 'socket -> ('flow, 'error) result Lwt.t;
|
|
close : 't -> unit Lwt.t;
|
|
}
|
|
constraint 'error = [> `Closed ]
|
|
|
|
let service connection handshake accept close =
|
|
Service { accept; connection; handshake; close }
|
|
|
|
open Lwt.Infix
|
|
|
|
let serve_when_ready : type t socket flow.
|
|
(t, socket, flow, _) posix ->
|
|
?stop:Lwt_switch.t ->
|
|
handler:(flow -> unit Lwt.t) ->
|
|
t ->
|
|
[ `Initialized of unit Lwt.t ] =
|
|
fun service ?stop ~handler t ->
|
|
let { accept; handshake; close } = service in
|
|
`Initialized
|
|
(let switched_off =
|
|
let t, u = Lwt.wait () in
|
|
Lwt_switch.add_hook stop (fun () ->
|
|
Lwt.wakeup_later u (Ok `Stopped) ;
|
|
Lwt.return_unit) ;
|
|
t in
|
|
let rec loop () =
|
|
accept t >>= function
|
|
| Ok socket ->
|
|
Lwt.async (fun () ->
|
|
handshake socket >>= function
|
|
| Ok flow -> handler flow
|
|
| Error `Closed ->
|
|
Logs.info (fun m -> m "Connection closed by peer") ;
|
|
Lwt.return ()
|
|
| Error _err ->
|
|
Logs.err (fun m ->
|
|
m "Got an error from a TCP/IP connection.") ;
|
|
Lwt.return ()) ;
|
|
loop ()
|
|
| Error `Closed -> Lwt.return_error `Closed
|
|
| Error _ -> Lwt.pause () >>= loop in
|
|
let stop_result =
|
|
Lwt.pick [ switched_off; loop () ] >>= function
|
|
| Ok `Stopped -> close t >>= fun () -> Lwt.return_ok ()
|
|
| Error _ as err -> close t >>= fun () -> Lwt.return err in
|
|
stop_result >>= function Ok () | Error `Closed -> Lwt.return_unit)
|
|
|
|
let server : type t. t runtime -> t -> Mimic.flow -> unit Lwt.t =
|
|
fun (module Runtime) conn flow ->
|
|
let module Server = Server (Mimic) (Runtime) in
|
|
Server.server conn flow
|
|
|
|
let serve ?stop service t =
|
|
let (Service { accept; handshake; connection; close }) = service in
|
|
let handler flow =
|
|
connection flow >>= function
|
|
| Ok (flow, Runtime (runtime, conn)) -> server runtime conn flow
|
|
| Error _ -> Lwt.return_unit in
|
|
serve_when_ready ?stop ~handler { accept; handshake; close } t
|
|
|
|
let run : type t. t runtime -> t -> Mimic.flow -> unit Lwt.t =
|
|
fun (module Runtime) conn flow ->
|
|
let module Client = Client (Mimic) (Runtime) in
|
|
Client.run conn flow
|