Update client modules

lib/client_connection.ml

@@ -1,15 +1,25 @@
+module IOVec = Httpaf.IOVec
+
 type state =
   | Uninitialized
   | Handshake of Client_handshake.t
   | Websocket of Client_websocket.t
 
 type t = state ref
 
+type error =
+  [ Httpaf.Client_connection.error
+  | `Handshake_failure of Httpaf.Response.t * [`read] Httpaf.Body.t ]
+
+type input_handlers = Client_websocket.input_handlers =
+  { frame : opcode:Websocket.Opcode.t -> is_fin:bool -> Bigstring.t -> off:int -> len:int -> unit
+  ; eof   : unit                                                                          -> unit }
+
 let passes_scrutiny ~accept headers =
   let upgrade              = Httpaf.Headers.get headers "upgrade"    in
   let connection           = Httpaf.Headers.get headers "connection" in
   let sec_websocket_accept = Httpaf.Headers.get headers "sec-websocket-accept" in
-  sec_websocket_accept = accept
+  sec_websocket_accept = Some accept
   && (match upgrade with
      | None         -> false
      | Some upgrade -> String.lowercase_ascii upgrade = "websocket")
@@ -18,26 +28,36 @@ let passes_scrutiny ~accept headers =
      | Some connection -> String.lowercase_ascii connection = "upgrade")
 ;;
 
-let create 
-    ~nonce 
-    ~host 
-    ~port 
+let handshake_exn t =
+  match !t with
+  | Handshake handshake -> handshake
+  | Uninitialized
+  | Websocket _ -> assert false
+
+let create
+    ~nonce
+    ~host
+    ~port
     ~resource
     ~sha1
     ~error_handler
     ~websocket_handler
   =
   let t = ref Uninitialized in
+  let nonce = B64.encode nonce in
   let response_handler response response_body =
     let accept = sha1 (nonce ^ "258EAFA5-E914-47DA-95CA-C5AB0DC85B11") in
-    match response.code with
+    match response.Httpaf.Response.status with
     | `Switching_protocols when passes_scrutiny ~accept response.headers ->
-      Body.close response_body response_body;
-      t := Websocket (Client_websocket.create ~websocket_handler ~eof_handler)
+      Httpaf.Body.close_reader response_body;
+      let handshake = handshake_exn t in
+      t := Websocket (Client_websocket.create ~websocket_handler);
+      Client_handshake.close handshake
     | _                    ->
       error_handler (`Handshake_failure(response, response_body))
   in
-  let handshake = 
+  let handshake =
+    let error_handler = (error_handler :> Httpaf.Client_connection.error_handler) in
     Client_handshake.create
       ~nonce
       ~host
@@ -64,18 +84,11 @@ let read t bs ~off ~len =
   | Websocket websocket -> Client_websocket.read websocket bs ~off ~len
 ;;
 
-let yield_reader t f =
+let read_eof t bs ~off ~len =
   match !t with
   | Uninitialized       -> assert false
-  | Handshake handshake -> Client_handshake.yield_reader handshake f 
-  | Websocket websocket -> Client_websocket.yield_reader websocket f 
-;;
-
-let shutdown_reader t =
-  match !t with
-  | Uninitialized       -> assert false
-  | Handshake handshake -> Client_handshake.shutdown_reader handshake
-  | Websocket websocket -> Client_websocket.shutdown_reader websocket
+  | Handshake handshake -> Client_handshake.read handshake bs ~off ~len
+  | Websocket websocket -> Client_websocket.read_eof websocket bs ~off ~len
 ;;
 
 let next_write_operation t =
@@ -95,13 +108,13 @@ let report_write_result t result =
 let yield_writer t f =
   match !t with
   | Uninitialized       -> assert false
-  | Handshake handshake -> Client_handshake.yield_writer handshake f 
-  | Websocket websocket -> Client_websocket.yield_writer websocket f 
+  | Handshake handshake -> Client_handshake.yield_writer handshake f
+  | Websocket websocket -> Client_websocket.yield_writer websocket f
 ;;
 
 let close t =
   match !t with
   | Uninitialized       -> assert false
-  | Handshake handshake -> Client_handshake.close handshake f 
-  | Websocket websocket -> Client_websocket.close websocket f 
+  | Handshake handshake -> Client_handshake.close handshake
+  | Websocket websocket -> Client_websocket.close websocket
 ;;

lib/client_connection.mli

@@ -1,32 +1,32 @@
 module IOVec = Httpaf.IOVec
 
-type t 
+type t
 
 type error =
   [ Httpaf.Client_connection.error
-  | `Handshake_failure of Httpaf.Request.t * [`read] Httpaf.Body.t ]
+  | `Handshake_failure of Httpaf.Response.t * [`read] Httpaf.Body.t ]
 
-type input_handlers = Client_websocket.input_handlers = 
-  { frame : Websocket.Opcode.t -> is_fin:bool -> Bigstring.t -> off:int -> len:int -> unit
-  ; eof   : unit                                                                   -> unit }
+type input_handlers = Client_websocket.input_handlers =
+  { frame : opcode:Websocket.Opcode.t -> is_fin:bool -> Bigstring.t -> off:int -> len:int -> unit
+  ; eof   : unit                                                                          -> unit }
 
-val create 
+val create
   :  nonce             : string
   -> host              : string
   -> port              : int
   -> resource          : string
   -> sha1              : (string -> string)
-  -> error_handler     : (Httpaf.Response.t -> [`read] Httpaf.Body.t -> unit)
-  -> websocket_handler : (Wsd.t                                      -> input_handlers)
+  -> error_handler     : (error -> unit)
+  -> websocket_handler : (Wsd.t -> input_handlers)
   -> t
 
-val next_read_operation  : t -> [ `Read | `Yield | `Close ]
+val next_read_operation  : t -> [ `Read | `Close ]
 val next_write_operation : t -> [ `Write of Bigstring.t IOVec.t list | `Yield | `Close of int ]
 
 val read : t -> Bigstring.t -> off:int -> len:int -> int
+val read_eof : t -> Bigstring.t -> off:int -> len:int -> int
 val report_write_result : t -> [`Ok of int | `Closed ] -> unit
 
-val yield_reader : t -> (unit -> unit) -> unit
 val yield_writer : t -> (unit -> unit) -> unit
 
 val close : t -> unit

lib/client_handshake.ml

@@ -1,29 +1,50 @@
 module IOVec = Httpaf.IOVec
-include Httpaf.Client_connection
 
-let create 
-    ~nonce 
-    ~host 
-    ~port 
+type t =
+  { connection : Httpaf.Client_connection.t
+  ; body       : [`write] Httpaf.Body.t }
+
+let create
+    ~nonce
+    ~host
+    ~port
     ~resource
     ~error_handler
-    ~response_handler 
+    ~response_handler
   =
-  let nonce = B64.encode nonce in
   let headers =
     [ "upgrade"              , "websocket"
     ; "connection"           , "upgrade"
     ; "host"                 , String.concat ":" [ host; string_of_int port ]
-    ; "sec-websocket-version", "13" 
+    ; "sec-websocket-version", "13"
     ; "sec-websocket-key"    , nonce
     ] |> Httpaf.Headers.of_list
   in
-  let request_body, t =
+  let body, connection =
     Httpaf.Client_connection.request
       (Httpaf.Request.create ~headers `GET resource)
       ~error_handler
       ~response_handler
   in
-  Httpaf.Body.close request_body;
-  t
+  { connection
+  ; body
+  }
 ;;
+
+let next_read_operation t =
+  Httpaf.Client_connection.next_read_operation t.connection
+
+let next_write_operation t =
+  Httpaf.Client_connection.next_write_operation t.connection
+
+let read t =
+  Httpaf.Client_connection.read t.connection
+
+let report_write_result t =
+  Httpaf.Client_connection.report_write_result t.connection
+
+let yield_writer t =
+  Httpaf.Client_connection.yield_writer t.connection
+
+let close t =
+  Httpaf.Body.close_writer t.body

lib/client_handshake.mli

@@ -1,8 +1,8 @@
 module IOVec = Httpaf.IOVec
 
-type t 
+type t
 
-val create 
+val create
   :  nonce            : string
   -> host             : string
   -> port             : int
@@ -15,7 +15,8 @@ val next_read_operation  : t -> [ `Read | `Close ]
 val next_write_operation : t -> [ `Write of Bigstring.t IOVec.t list | `Yield | `Close of int ]
 
 val read : t -> Bigstring.t -> off:int -> len:int -> int
-val shutdown_reader : t -> unit
 val report_write_result : t -> [`Ok of int | `Closed ] -> unit
 
 val yield_writer : t -> (unit -> unit) -> unit
+
+val close : t -> unit

lib/client_websocket.ml

@@ -1,14 +1,46 @@
+module IOVec = Httpaf.IOVec
+
 type t =
-  { reader : Reader.t
+  { reader : [`Parse of string list * string] Reader.t
   ; wsd    : Wsd.t }
 
 type input_handlers =
-  { frame : Websocket.Opcode.t -> is_fin:bool -> Bigstring.t -> off:int -> len:int -> unit
-  ; eof   : unit                                                                   -> unit }
+  { frame : opcode:Websocket.Opcode.t -> is_fin:bool -> Bigstring.t -> off:int -> len:int -> unit
+  ; eof   : unit                                                                          -> unit }
+
+let random_int32 () =
+  Random.int32 Int32.max_int
 
 let create ~websocket_handler =
-  let wsd            = Wsd.create () in
-  let { frame; eof } = websocket_handler wsd in
+  let mode         = `Client random_int32 in
+  let wsd          = Wsd.create mode in
+  let { frame; _ } = websocket_handler wsd in
   { reader = Reader.create frame
-  ; wds 
+  ; wsd
   }
+
+let next_read_operation t =
+  Reader.next t.reader
+
+let next_write_operation t =
+  Wsd.next t.wsd
+
+let read t bs ~off ~len =
+  Reader.read_with_more t.reader bs ~off ~len Incomplete
+
+let read_eof t bs ~off ~len =
+  Reader.read_with_more t.reader bs ~off ~len Complete
+
+let report_write_result t result =
+  Wsd.report_result t.wsd result
+
+let yield_writer t k =
+  if Wsd.is_closed t.wsd
+  then begin
+    Wsd.close t.wsd;
+    k ()
+  end else
+    Wsd.when_ready_to_write t.wsd k
+
+let close { wsd; _ } =
+  Wsd.close wsd

lib/client_websocket.mli

@@ -3,20 +3,20 @@ module IOVec = Httpaf.IOVec
 type t
 
 type input_handlers =
-  { frame : Websocket.Opcode.t -> is_fin:bool -> Bigstring.t -> off:int -> len:int -> unit
-  ; eof   : unit                                                                   -> unit }
+  { frame : opcode:Websocket.Opcode.t -> is_fin:bool -> Bigstring.t -> off:int -> len:int -> unit
+  ; eof   : unit                                                                          -> unit }
 
-val create 
+val create
   :  websocket_handler : (Wsd.t -> input_handlers)
   -> t
 
-val next_read_operation  : t -> [ `Read | `Yield | `Close ]
+val next_read_operation  : t -> [ `Read | `Close ]
 val next_write_operation : t -> [ `Write of Bigstring.t IOVec.t list | `Yield | `Close of int ]
 
 val read : t -> Bigstring.t -> off:int -> len:int -> int
+val read_eof : t -> Bigstring.t -> off:int -> len:int -> int
 val report_write_result : t -> [`Ok of int | `Closed ] -> unit
 
-val yield_reader : t -> (unit -> unit) -> unit
 val yield_writer : t -> (unit -> unit) -> unit
 
 val close : t -> unit