Fix client. Add server. Add Lwt support.

dune-project

@@ -0,0 +1 @@
+(lang dune 1.1)

examples/lwt/dune

@@ -0,0 +1,3 @@
+(executables
+ (names wscat echo_server)
+ (libraries websocketaf websocketaf-lwt lwt lwt.unix))

examples/lwt/echo_server.ml

@@ -0,0 +1,71 @@
+let connection_handler : Unix.sockaddr -> Lwt_unix.file_descr -> unit Lwt.t =
+  let module Body = Httpaf.Body in
+  let module Headers = Httpaf.Headers in
+  let module Reqd = Httpaf.Reqd in
+  let module Response = Httpaf.Response in
+  let module Status = Httpaf.Status in
+
+  let websocket_handler _ wsd =
+    let frame ~opcode ~is_fin:_ bs ~off ~len =
+      match opcode with
+      | `Continuation
+      | `Text
+      | `Binary ->
+        Websocketaf.Wsd.schedule wsd bs ~kind:`Text ~off ~len
+      | `Connection_close ->
+        Websocketaf.Wsd.close wsd
+      | `Ping ->
+        Websocketaf.Wsd.send_ping wsd
+      | `Pong
+      | `Other _ ->
+        ()
+    in
+    let eof () = ()
+    in
+    { Websocketaf.Server_connection.frame
+    ; eof
+    }
+  in
+
+  let error_handler _client_address ?request:_ error start_response =
+    let response_body = start_response Headers.empty in
+
+    begin match error with
+    | `Exn exn ->
+      Body.write_string response_body (Printexc.to_string exn);
+      Body.write_string response_body "\n";
+
+    | #Status.standard as error ->
+      Body.write_string response_body (Status.default_reason_phrase error)
+    end;
+  in
+
+  Websocketaf_lwt.Server.create_connection_handler
+    ?config:None
+    ~websocket_handler
+    ~error_handler
+
+
+
+let () =
+  let open Lwt.Infix in
+
+  let port = ref 8080 in
+  Arg.parse
+    ["-p", Arg.Set_int port, " Listening port number (8080 by default)"]
+    ignore
+    "Echoes websocket messages. Runs forever.";
+
+  let listen_address = Unix.(ADDR_INET (inet_addr_loopback, !port)) in
+
+  Lwt.async begin fun () ->
+    Lwt_io.establish_server_with_client_socket
+      listen_address connection_handler
+    >>= fun _server ->
+      Printf.printf "Listening on port %i and echoing websocket messages.\n" !port;
+      flush stdout;
+      Lwt.return_unit
+  end;
+
+  let forever, _ = Lwt.wait () in
+  Lwt_main.run forever

examples/lwt/wscat.ml

@@ -0,0 +1,66 @@
+open Lwt.Infix
+
+let websocket_handler wsd =
+  let rec input_loop wsd () =
+    Lwt_io.(read_line stdin) >>= fun line ->
+    let payload = Bytes.of_string line in
+    Websocketaf.Wsd.send_bytes wsd ~kind:`Text payload ~off:0 ~len:(Bytes.length payload);
+    input_loop wsd ()
+  in
+  Lwt.async (input_loop wsd);
+  let frame ~opcode:_ ~is_fin:_ bs ~off ~len =
+    let payload = Bytes.create len in
+    Lwt_bytes.blit_to_bytes
+      bs off
+      payload 0
+      len;
+    Printf.printf "%s\n" (Bytes.unsafe_to_string payload);
+    flush stdout
+  in
+  let eof () =
+    Printf.printf "[EOF]\n"
+  in
+  { Websocketaf.Client_connection.frame
+  ; eof
+  }
+
+let error_handler = function
+  | `Handshake_failure (rsp, _body) ->
+    Format.printf "Handshake failure: %a" Httpaf.Response.pp_hum rsp
+  | _ -> assert false
+
+let () =
+  let host = ref None in
+  let port = ref 80 in
+
+  Arg.parse
+    ["-p", Set_int port, " Port number (80 by default)"]
+    (fun host_argument -> host := Some host_argument)
+    "wscat.exe [-p N] HOST";
+
+  let host =
+    match !host with
+    | None -> failwith "No hostname provided"
+    | Some host -> host
+  in
+
+  Lwt_main.run begin
+    Lwt_unix.getaddrinfo host (string_of_int !port) [Unix.(AI_FAMILY PF_INET)]
+    >>= fun addresses ->
+
+    let socket = Lwt_unix.socket Unix.PF_INET Unix.SOCK_STREAM 0 in
+    Lwt_unix.connect socket (List.hd addresses).Unix.ai_addr
+    >>= fun () ->
+
+    let nonce = "0123456789ABCDEF" in
+    let resource = "/" in
+    let port = !port in
+    Websocketaf_lwt.Client.connect
+      socket
+      ~nonce
+      ~host
+      ~port
+      ~resource
+      ~error_handler
+      ~websocket_handler
+  end

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 -> Bigstringaf.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
 ;;