From d8ff243e8d1079a3fbdd2d6324b0e4a51cd6be35 Mon Sep 17 00:00:00 2001 From: Simon Cruanes Date: Tue, 2 Apr 2024 14:35:57 -0400 Subject: [PATCH] feat ws: pass the whole request to the handler --- examples/echo_ws.ml | 4 ++-- src/core/server.ml | 12 +++++++----- src/core/server.mli | 18 +++++++++++------- src/ws/tiny_httpd_ws.ml | 15 ++++++++------- src/ws/tiny_httpd_ws.mli | 2 +- 5 files changed, 29 insertions(+), 22 deletions(-) diff --git a/examples/echo_ws.ml b/examples/echo_ws.ml index f24cf283..ccbe4695 100644 --- a/examples/echo_ws.ml +++ b/examples/echo_ws.ml @@ -10,9 +10,9 @@ let setup_logging ~debug () = else Logs.Info) -let handle_ws _client_addr ic oc = +let handle_ws (req : unit Request.t) ic oc = Log.info (fun k -> - k "new client connection from %s" (Util.show_sockaddr _client_addr)); + k "new client connection from %s" (Util.show_sockaddr req.client_addr)); let (_ : Thread.t) = Thread.create diff --git a/src/core/server.ml b/src/core/server.ml index 1eb1715b..0bef5d26 100644 --- a/src/core/server.ml +++ b/src/core/server.ml @@ -35,12 +35,14 @@ module type UPGRADE_HANDLER = sig val name : string (** Name in the "upgrade" header *) - val handshake : unit Request.t -> (Headers.t * handshake_state, string) result + val handshake : + Unix.sockaddr -> + unit Request.t -> + (Headers.t * handshake_state, string) result (** Perform the handshake and upgrade the connection. The returned code is [101] alongside these headers. *) - val handle_connection : - Unix.sockaddr -> handshake_state -> IO.Input.t -> IO.Output.t -> unit + val handle_connection : handshake_state -> IO.Input.t -> IO.Output.t -> unit (** Take control of the connection and take it from there *) end @@ -362,7 +364,7 @@ let client_handle_for (self : t) ~client_addr ic oc : unit = | None -> bad_reqf 426 "expected 'connection: upgrade' header"); (* ok, this is the upgrade we expected *) - match UP.handshake req with + match UP.handshake client_addr req with | Error msg -> (* fail the upgrade *) Log.error (fun k -> k "upgrade failed: %s" msg); @@ -378,7 +380,7 @@ let client_handle_for (self : t) ~client_addr ic oc : unit = log_response req resp; Response.Private_.output_ ~bytes:bytes_res oc resp; - UP.handle_connection client_addr handshake_st ic oc + UP.handle_connection handshake_st ic oc with e -> let bt = Printexc.get_raw_backtrace () in handle_bad_req req e bt diff --git a/src/core/server.mli b/src/core/server.mli index e856c7e4..b79284f3 100644 --- a/src/core/server.mli +++ b/src/core/server.mli @@ -250,14 +250,18 @@ module type UPGRADE_HANDLER = sig val name : string (** Name in the "upgrade" header *) - val handshake : unit Request.t -> (Headers.t * handshake_state, string) result - (** Perform the handshake and upgrade the connection. The returned - code is [101] alongside these headers. - In case the handshake fails, this only returns [Error log_msg]. - The connection is closed without further ado. *) + val handshake : + Unix.sockaddr -> + unit Request.t -> + (Headers.t * handshake_state, string) result + (** Perform the handshake and upgrade the connection. This returns either + [Ok (resp_headers, state)] in case of success, in which case the + server sends a [101] response with [resp_headers]; + or it returns [Error log_msg] if the the handshake fails, in which case + the connection is closed without further ado and [log_msg] is logged + locally (but not returned to the client). *) - val handle_connection : - Unix.sockaddr -> handshake_state -> IO.Input.t -> IO.Output.t -> unit + val handle_connection : handshake_state -> IO.Input.t -> IO.Output.t -> unit (** Take control of the connection and take it from ther.e *) end diff --git a/src/ws/tiny_httpd_ws.ml b/src/ws/tiny_httpd_ws.ml index 54dc1481..276c853c 100644 --- a/src/ws/tiny_httpd_ws.ml +++ b/src/ws/tiny_httpd_ws.ml @@ -1,6 +1,6 @@ open Common_ws_ -type handler = Unix.sockaddr -> IO.Input.t -> IO.Output.t -> unit +type handler = unit Request.t -> IO.Input.t -> IO.Output.t -> unit module Frame_type = struct type t = int @@ -407,8 +407,9 @@ let upgrade ic oc : _ * _ = module Make_upgrade_handler (X : sig val accept_ws_protocol : string -> bool val handler : handler -end) : Server.UPGRADE_HANDLER = struct - type handshake_state = unit +end) : Server.UPGRADE_HANDLER with type handshake_state = unit Request.t = +struct + type handshake_state = unit Request.t let name = "websocket" @@ -443,14 +444,14 @@ end) : Server.UPGRADE_HANDLER = struct let headers = [ "sec-websocket-accept", accept ] in Log.debug (fun k -> k "websocket: upgrade successful, accept key is %S" accept); - headers, () + headers, req - let handshake req : _ result = + let handshake _addr req : _ result = try Ok (handshake_ req) with Bad_req s -> Error s - let handle_connection addr () ic oc = + let handle_connection req ic oc = let ws_ic, ws_oc = upgrade ic oc in - try X.handler addr ws_ic ws_oc + try X.handler req ws_ic ws_oc with Close_connection -> Log.debug (fun k -> k "websocket: requested to close the connection"); () diff --git a/src/ws/tiny_httpd_ws.mli b/src/ws/tiny_httpd_ws.mli index 2bd30f70..c9d0e9b2 100644 --- a/src/ws/tiny_httpd_ws.mli +++ b/src/ws/tiny_httpd_ws.mli @@ -4,7 +4,7 @@ for a websocket server. It has no additional dependencies. *) -type handler = Unix.sockaddr -> IO.Input.t -> IO.Output.t -> unit +type handler = unit Request.t -> IO.Input.t -> IO.Output.t -> unit (** Websocket handler *) val upgrade : IO.Input.t -> IO.Output.t -> IO.Input.t * IO.Output.t