Source file io.ml

open Import

exception Error of string

let () =
  Printexc.register_printer (function
    | Error msg -> Some ("Error: " ^ msg)
    | _ -> None)

let caseless_equal a b =
  if a == b then
    true
  else
    let len = String.length a in
    len = String.length b
    &&
    let stop = ref false in
    let idx = ref 0 in
    while (not !stop) && !idx < len do
      let c1 = String.unsafe_get a !idx in
      let c2 = String.unsafe_get b !idx in
      if Char.lowercase_ascii c1 <> Char.lowercase_ascii c2 then stop := true;
      incr idx
    done;
    not !stop

let content_type_lowercase =
  String.lowercase_ascii Header.Private.Key.content_type

let content_length_lowercase =
  String.lowercase_ascii Header.Private.Key.content_length

module Make (Io : sig
  type 'a t

  val return : 'a -> 'a t

  val raise : exn -> 'a t

  module O : sig
    val ( let+ ) : 'a t -> ('a -> 'b) -> 'b t

    val ( let* ) : 'a t -> ('a -> 'b t) -> 'b t
  end
end) (Chan : sig
  type input

  type output

  val read_line : input -> string option Io.t

  val read_exactly : input -> int -> string option Io.t

  val write : output -> string -> unit Io.t
end) =
struct
  open Io.O

  let read_header =
    let init_content_length = -1 in
    let rec loop chan content_length content_type =
      let* line = Chan.read_line chan in
      match line with
      | None -> Io.return None
      | Some ""
      | Some "\r" ->
        Io.return (Some (content_length, content_type))
      | Some line -> (
        match String.lsplit2 ~on:':' line with
        | None -> loop chan content_length content_type
        | Some (k, v) ->
          let k = String.trim k in
          if
            caseless_equal k content_length_lowercase
            && content_length = init_content_length
          then
            let content_length = int_of_string_opt (String.trim v) in
            match content_length with
            | None -> Io.raise (Error "Content-Length is invalid")
            | Some content_length -> loop chan content_length content_type
          else if caseless_equal k content_type_lowercase && content_type = None
          then
            let content_type = String.trim v in
            loop chan content_length (Some content_type)
          else
            loop chan content_length content_type)
    in
    fun chan ->
      let open Io.O in
      let* res = loop chan init_content_length None in
      match res with
      | None -> Io.return None
      | Some (content_length, content_type) ->
        let+ () =
          if content_length = init_content_length then
            Io.raise (Error "content length absent")
          else
            Io.return ()
        in
        Some (Header.create ?content_type ~content_length ())

  let read =
    let req json = Jsonrpc.Message (Jsonrpc.Message.either_of_yojson json) in
    let resp json = Jsonrpc.Response (Jsonrpc.Response.t_of_yojson json) in
    let packet =
      let open Json.O in
      req <|> resp
    in
    fun chan ->
      let* header = read_header chan in
      match header with
      | None -> Io.return None
      | Some header -> (
        let len = Header.content_length header in
        let* buf = Chan.read_exactly chan len in
        match buf with
        | None -> Io.raise (Error "unable to read json")
        | Some buf ->
          let json = Json.of_string buf in
          Io.return (Some (packet json)))

  let write chan packet =
    let json = Jsonrpc.yojson_of_packet packet in
    let data = Json.to_string json in
    let content_length = String.length data in
    let header = Header.create ~content_length () in
    let* () = Chan.write chan (Header.to_string header) in
    Chan.write chan data
end