refactor: unify input handlers eof + websocket error handler (#72)

* refactor: unify input handlers eof + websocket error handler * wip
anmonteiro · Aug 27, 2024 · 9aacc12 · 9aacc12
1 parent f5580b4
commit 9aacc12
Show file tree

Hide file tree

Showing 25 changed files with 127 additions and 185 deletions.
diff --git a/async/httpun_ws_async.ml b/async/httpun_ws_async.ml
@@ -10,13 +10,11 @@ module Server = struct
   let create_connection_handler
     ?(config = Httpun.Config.default)
     ?error_handler
-    ?websocket_error_handler
     websocket_handler = fun client_addr socket ->
     let connection =
       Httpun_ws.Server_connection.create
         ~sha1
         ?error_handler:(Option.map ~f:(fun f -> f client_addr) error_handler)
-        ?websocket_error_handler:(Option.map ~f:(fun f -> f client_addr) websocket_error_handler)
         (websocket_handler client_addr)
     in
     Gluten_async.Server.create_connection_handler

diff --git a/async/httpun_ws_async.mli b/async/httpun_ws_async.mli
@@ -4,7 +4,6 @@ module Server : sig
   val create_connection_handler
     :  ?config : Httpun.Config.t
     -> ?error_handler: ('a -> Httpun.Server_connection.error_handler)
-    -> ?websocket_error_handler: ('a -> Httpun_ws.Server_connection.error_handler)
     -> ('a
          -> Httpun_ws.Wsd.t
          -> Httpun_ws.Websocket_connection.input_handlers)

diff --git a/eio/httpun_ws_eio.ml b/eio/httpun_ws_eio.ml
@@ -4,20 +4,13 @@ let sha1 s =
   |> Digestif.SHA1.to_raw_string
 
 module Server = struct
-  (* TODO: should this error handler be a websocket error handler or an HTTP
-   * error handler?*)
   let create_connection_handler
-    ?(config = Httpun.Config.default)
-    ?error_handler
-    ?websocket_error_handler
-    ~sw
-    websocket_handler =
+    ?(config = Httpun.Config.default) ?error_handler ~sw websocket_handler =
     fun client_addr socket ->
       let connection =
         Httpun_ws.Server_connection.create
           ~sha1
           ?error_handler:(Option.map (fun f -> f client_addr) error_handler)
-          ?websocket_error_handler:(Option.map (fun f -> f client_addr) websocket_error_handler)
           (websocket_handler client_addr)
       in
       Gluten_eio.Server.create_connection_handler

diff --git a/eio/httpun_ws_eio.mli b/eio/httpun_ws_eio.mli
@@ -1,11 +1,9 @@
 open Httpun_ws
 
 module Server : sig
-
   val create_connection_handler
-    :  ?config : Httpun.Config.t
-    -> ?error_handler : (Eio.Net.Sockaddr.stream -> Httpun.Server_connection.error_handler)
-    -> ?websocket_error_handler : (Eio.Net.Sockaddr.stream -> Server_connection.error_handler)
+    :  ?config: Httpun.Config.t
+    -> ?error_handler: (Eio.Net.Sockaddr.stream -> Httpun.Server_connection.error_handler)
     -> sw:Eio.Switch.t
     -> (Eio.Net.Sockaddr.stream -> Wsd.t -> Websocket_connection.input_handlers)
     -> (Eio.Net.Sockaddr.stream -> _ Eio.Net.stream_socket -> unit)
@@ -16,14 +14,14 @@ module Client : sig
 
   (* Perform HTTP/1.1 handshake and upgrade to WS. *)
   val connect
-    :  ?config : Httpun.Config.t
-    -> sw             : Eio.Switch.t
-    -> nonce          : string
-    -> host           : string
-    -> port           : int
-    -> resource       : string
-    -> error_handler : (Client_connection.error -> unit)
-    -> websocket_handler : (Wsd.t -> Websocket_connection.input_handlers)
+    :  ?config:Httpun.Config.t
+    -> sw:Eio.Switch.t
+    -> nonce:string
+    -> host:string
+    -> port:int
+    -> resource:string
+    -> error_handler:(Client_connection.error -> unit)
+    -> websocket_handler:(Wsd.t -> Websocket_connection.input_handlers)
     -> Eio_unix.Net.stream_socket_ty Eio.Net.stream_socket
     -> t
 

diff --git a/examples/async/echo_server.ml b/examples/async/echo_server.ml
@@ -26,26 +26,20 @@ let connection_handler : ([< Socket.Address.t] as 'a) -> ([`Active], 'a) Socket.
       | `Other _ ->
         ()
     in
-    let eof () =
-      Log.Global.error "EOF\n%!";
-      Httpun_ws.Wsd.close wsd
+    let eof ?error () =
+      match error with
+      | Some _ -> assert false
+      | None ->
+        Log.Global.error "EOF\n%!";
+        Httpun_ws.Wsd.close wsd
     in
     { Httpun_ws.Websocket_connection.frame
     ; eof
     }
   in
 
-  let error_handler _client_address wsd (`Exn exn) =
-    let message = Exn.to_string exn in
-    let payload = Bytes.of_string message in
-    Httpun_ws.Wsd.send_bytes wsd ~kind:`Text payload ~off:0
-      ~len:(Bytes.length payload);
-    Httpun_ws.Wsd.close wsd
-  in
-
   Httpun_ws_async.Server.create_connection_handler
     ?config:None
-    ~websocket_error_handler:error_handler
     websocket_handler
 
 let main port max_accepts_per_batch () =

diff --git a/examples/async/echo_server_upgrade.ml b/examples/async/echo_server_upgrade.ml
@@ -29,22 +29,18 @@ let connection_handler =
       | `Other _ ->
         ()
     in
-    let eof () =
-      Format.eprintf "EOF\n%!";
-      Httpun_ws.Wsd.close wsd
+    let eof ?error () =
+      match error with
+      | Some _ -> assert false
+      | None ->
+        Format.eprintf "EOF\n%!";
+        Httpun_ws.Wsd.close wsd
     in
     { Httpun_ws.Websocket_connection.frame
     ; eof
     }
   in
 
-  let error_handler wsd (`Exn exn) =
-    let message = Exn.to_string exn in
-    let payload = Bytes.of_string message in
-    Httpun_ws.Wsd.send_bytes wsd ~kind:`Text payload ~off:0
-      ~len:(Bytes.length payload);
-    Httpun_ws.Wsd.close wsd
-  in
   let http_error_handler _client_address ?request:_ error handle =
     let message =
       match error with
@@ -57,9 +53,7 @@ let connection_handler =
   in
   let upgrade_handler addr upgrade () =
     let ws_conn =
-      Httpun_ws.Server_connection.create_websocket
-        ~error_handler
-        (websocket_handler addr)
+      Httpun_ws.Server_connection.create_websocket (websocket_handler addr)
     in
     upgrade
       (Gluten.make (module Httpun_ws.Server_connection) ws_conn)

diff --git a/examples/async/wscat.ml b/examples/async/wscat.ml
@@ -23,8 +23,10 @@ let websocket_handler wsd =
     Log.Global.printf "%s\n%!" payload)
   in
 
-  let eof () =
-    Log.Global.error "[EOF]\n%!"
+  let eof ?error () =
+    match error with
+    | Some _ -> assert false
+    | None -> Log.Global.error "[EOF]\n%!"
   in
   { Httpun_ws.Websocket_connection.frame
   ; eof

diff --git a/examples/eio/echo_server.ml b/examples/eio/echo_server.ml
@@ -28,27 +28,26 @@ let connection_handler ~sw : Eio.Net.Sockaddr.stream -> _ Eio.Net.stream_socket
       | `Other _ ->
         ()
     in
-    let eof () =
-      Format.eprintf "EOF\n%!";
-      Httpun_ws.Wsd.close wsd
+    let eof ?error () =
+      match error with
+      Some (`Exn exn) ->
+        let message = Printexc.to_string exn in
+        let payload = Bytes.of_string message in
+        Httpun_ws.Wsd.send_bytes wsd ~kind:`Text payload ~off:0
+          ~len:(Bytes.length payload);
+        Httpun_ws.Wsd.close wsd
+      | None ->
+        Format.eprintf "EOF\n%!";
+        Httpun_ws.Wsd.close wsd
     in
     { Httpun_ws.Websocket_connection.frame
     ; eof
     }
   in
 
-  let error_handler _client_address wsd (`Exn exn) =
-    let message = Printexc.to_string exn in
-    let payload = Bytes.of_string message in
-    Httpun_ws.Wsd.send_bytes wsd ~kind:`Text payload ~off:0
-      ~len:(Bytes.length payload);
-    Httpun_ws.Wsd.close wsd
-  in
-
   Httpun_ws_eio.Server.create_connection_handler
     ?config:None
     ~sw
-    ~websocket_error_handler:error_handler
     websocket_handler
 
 
@@ -75,7 +74,7 @@ let () =
         Eio.Domain_manager.run domain_mgr (fun () ->
           Eio.Switch.run (fun sw ->
             while true do
-              Eio.Net.accept_fork socket ~sw ~on_error:(fun _ -> assert false) (fun client_sock client_addr ->
+              Eio.Net.accept_fork socket ~sw ~on_error:raise (fun client_sock client_addr ->
                   (* let p, u = Eio.Promise.create () in *)
                   connection_handler ~sw client_addr client_sock)
             done;

diff --git a/examples/eio/echo_server_upgrade.ml b/examples/eio/echo_server_upgrade.ml
@@ -26,22 +26,18 @@ let connection_handler =
       | `Other _ ->
         ()
     in
-    let eof () =
-      Format.eprintf "EOF\n%!";
-      Httpun_ws.Wsd.close wsd
+    let eof ?error () =
+      match error with
+      | Some _ -> assert false
+      | None ->
+        Format.eprintf "EOF\n%!";
+        Httpun_ws.Wsd.close wsd
     in
     { Httpun_ws.Websocket_connection.frame
     ; eof
     }
   in
 
-  let error_handler wsd (`Exn exn) =
-    let message = Printexc.to_string exn in
-    let payload = Bytes.of_string message in
-    Httpun_ws.Wsd.send_bytes wsd ~kind:`Text payload ~off:0
-      ~len:(Bytes.length payload);
-    Httpun_ws.Wsd.close wsd
-  in
   let http_error_handler _client_address ?request:_ error handle =
     let message =
       match error with
@@ -54,9 +50,7 @@ let connection_handler =
   in
   let upgrade_handler addr upgrade () =
     let ws_conn =
-      Httpun_ws.Server_connection.create_websocket
-        ~error_handler
-        (websocket_handler addr)
+      Httpun_ws.Server_connection.create_websocket (websocket_handler addr)
     in
     upgrade
       (Gluten.make (module Httpun_ws.Server_connection) ws_conn)

diff --git a/examples/eio/wscat.ml b/examples/eio/wscat.ml
@@ -22,8 +22,7 @@ let websocket_handler env ~sw u wsd =
     Bigstringaf.blit_to_bytes bs ~src_off:off payload ~dst_off:0 ~len;
     Format.printf "%s@." (Bytes.unsafe_to_string payload);)
   in
-
-  let eof () =
+  let eof ?error:_ () =
     Printf.eprintf "[EOF]\n%!";
     Promise.resolve u ()
   in

diff --git a/examples/eio/wscat_upgrade.ml b/examples/eio/wscat_upgrade.ml
@@ -23,9 +23,12 @@ let websocket_handler env ~sw u wsd =
     Format.printf "%s@." (Bytes.unsafe_to_string payload);)
   in
 
-  let eof () =
-    Printf.eprintf "[EOF]\n%!";
-    Promise.resolve u ()
+  let eof ?error () =
+    match error with
+    | Some _ -> assert false
+    | None ->
+      Printf.eprintf "[EOF]\n%!";
+      Promise.resolve u ()
   in
   { Httpun_ws.Websocket_connection.frame
   ; eof

diff --git a/examples/lwt/echo_server.ml b/examples/lwt/echo_server.ml
@@ -21,26 +21,20 @@ let connection_handler : Unix.sockaddr -> Lwt_unix.file_descr -> unit Lwt.t =
       | `Other _ ->
         ()
     in
-    let eof () =
-      Format.eprintf "EOF\n%!";
-      Httpun_ws.Wsd.close wsd
+    let eof ?error () =
+      match error with
+      | Some _ -> assert false
+      | None ->
+        Format.eprintf "EOF\n%!";
+        Httpun_ws.Wsd.close wsd
     in
     { Httpun_ws.Websocket_connection.frame
     ; eof
     }
   in
 
-  let error_handler _client_address wsd (`Exn exn) =
-    let message = Printexc.to_string exn in
-    let payload = Bytes.of_string message in
-    Httpun_ws.Wsd.send_bytes wsd ~kind:`Text payload ~off:0
-      ~len:(Bytes.length payload);
-    Httpun_ws.Wsd.close wsd
-  in
-
   Httpun_ws_lwt_unix.Server.create_connection_handler
     ?config:None
-    ~websocket_error_handler:error_handler
     websocket_handler
 
 

diff --git a/examples/lwt/echo_server_upgrade.ml b/examples/lwt/echo_server_upgrade.ml
@@ -26,22 +26,18 @@ let connection_handler =
       | `Other _ ->
         ()
     in
-    let eof () =
-      Format.eprintf "EOF\n%!";
-      Httpun_ws.Wsd.close wsd
+    let eof ?error () =
+      match error with
+      | Some _ -> assert false
+      | None ->
+        Format.eprintf "EOF\n%!";
+        Httpun_ws.Wsd.close wsd
     in
     { Httpun_ws.Websocket_connection.frame
     ; eof
     }
   in
 
-  let error_handler wsd (`Exn exn) =
-    let message = Printexc.to_string exn in
-    let payload = Bytes.of_string message in
-    Httpun_ws.Wsd.send_bytes wsd ~kind:`Text payload ~off:0
-      ~len:(Bytes.length payload);
-    Httpun_ws.Wsd.close wsd
-  in
   let http_error_handler _client_address ?request:_ error handle =
     let message =
       match error with
@@ -54,9 +50,7 @@ let connection_handler =
   in
   let upgrade_handler addr upgrade () =
     let ws_conn =
-      Httpun_ws.Server_connection.create_websocket
-        ~error_handler
-        (websocket_handler addr)
+      Httpun_ws.Server_connection.create_websocket (websocket_handler addr)
     in
     upgrade
       (Gluten.make (module Httpun_ws.Server_connection) ws_conn)

diff --git a/examples/lwt/wscat.ml b/examples/lwt/wscat.ml
@@ -24,9 +24,12 @@ let websocket_handler u wsd =
     Format.printf "%s@." (Bytes.unsafe_to_string payload);)
   in
 
-  let eof () =
-    Printf.eprintf "[EOF]\n%!";
-    Lwt.wakeup_later u ()
+  let eof ?error () =
+    match error with
+    | Some _ -> assert false
+    | None ->
+      Printf.eprintf "[EOF]\n%!";
+      Lwt.wakeup_later u ()
   in
   { Httpun_ws.Websocket_connection.frame
   ; eof

diff --git a/examples/lwt/wscat_upgrade.ml b/examples/lwt/wscat_upgrade.ml
@@ -23,9 +23,12 @@ let websocket_handler u wsd =
       len;
     Format.printf "%s@." (Bytes.unsafe_to_string payload);)
   in
-  let eof () =
-    Printf.eprintf "[EOF]\n%!";
-    Lwt.wakeup_later u ()
+  let eof ?error () =
+    match error with
+    | Some _ -> assert false
+    | None ->
+      Printf.eprintf "[EOF]\n%!";
+      Lwt.wakeup_later u ()
   in
   { Httpun_ws.Websocket_connection.frame
   ; eof

diff --git a/lib/client_connection.ml b/lib/client_connection.ml
@@ -104,12 +104,11 @@ let connect
   in
   Lazy.force t
 
-let create ?error_handler websocket_handler =
+let create websocket_handler =
   { state =
       Websocket
         (Websocket_connection.create
           ~mode:(`Client Websocket_connection.random_int32)
-          ?error_handler
           websocket_handler) }
 
 let next_read_operation t =