From a405fb046d50de126be2fb7a9d667634315d1bc9 Mon Sep 17 00:00:00 2001 From: Simon Cruanes Date: Mon, 5 Feb 2024 00:36:32 -0500 Subject: [PATCH] expose client address to websocket --- examples/echo_ws.ml | 2 +- src/Tiny_httpd_server.ml | 25 ++++++++++++------------- src/Tiny_httpd_server.mli | 9 +++++++-- src/Tiny_httpd_util.ml | 5 +++++ src/Tiny_httpd_util.mli | 4 ++++ src/ws/tiny_httpd_ws.ml | 6 +++--- src/ws/tiny_httpd_ws.mli | 4 +--- 7 files changed, 33 insertions(+), 22 deletions(-) diff --git a/examples/echo_ws.ml b/examples/echo_ws.ml index fdb44411..68ec802e 100644 --- a/examples/echo_ws.ml +++ b/examples/echo_ws.ml @@ -30,7 +30,7 @@ let () = let server = S.create ~port:!port_ ~max_connections:!j () in Tiny_httpd_ws.add_route_handler server S.Route.(exact "echo" @/ return) - (fun ic oc -> + (fun addr ic oc -> Log.info (fun k -> k "new client connection"); let buf = Bytes.create 32 in let continue = ref true in diff --git a/src/Tiny_httpd_server.ml b/src/Tiny_httpd_server.ml index 56e9c930..7dcd466e 100644 --- a/src/Tiny_httpd_server.ml +++ b/src/Tiny_httpd_server.ml @@ -646,7 +646,8 @@ module type UPGRADE_HANDLER = sig (** Perform the handshake and upgrade the connection. The returned code is [101] alongside these headers. *) - val handle_connection : handshake_state -> IO.Input.t -> IO.Output.t -> unit + val handle_connection : + Unix.sockaddr -> handshake_state -> IO.Input.t -> IO.Output.t -> unit (** Take control of the connection and take it from there *) end @@ -893,11 +894,6 @@ let create_from ?(buf_size = 16 * 1_024) ?(middlewares = []) ~backend () : t = let is_ipv6_str addr : bool = String.contains addr ':' -let str_of_sockaddr = function - | Unix.ADDR_UNIX f -> f - | Unix.ADDR_INET (inet, port) -> - Printf.sprintf "%s:%d" (Unix.string_of_inet_addr inet) port - module Unix_tcp_server_ = struct type t = { addr: string; @@ -964,7 +960,8 @@ module Unix_tcp_server_ = struct let handle_client_unix_ (client_sock : Unix.file_descr) (client_addr : Unix.sockaddr) : unit = Log.info (fun k -> - k "serving new client on %s" (str_of_sockaddr client_addr)); + k "serving new client on %s" + (Tiny_httpd_util.show_sockaddr client_addr)); Unix.(setsockopt_float client_sock SO_RCVTIMEO self.timeout); Unix.(setsockopt_float client_sock SO_SNDTIMEO self.timeout); let oc = @@ -974,14 +971,14 @@ module Unix_tcp_server_ = struct handle.handle ~client_addr ic oc; Log.info (fun k -> k "done with client on %s, exiting" - @@ str_of_sockaddr client_addr); + @@ Tiny_httpd_util.show_sockaddr client_addr); (try Unix.shutdown client_sock Unix.SHUTDOWN_ALL; Unix.close client_sock with e -> Log.error (fun k -> k "error when closing sock for client %s: %s" - (str_of_sockaddr client_addr) + (Tiny_httpd_util.show_sockaddr client_addr) (Printexc.to_string e))); () in @@ -1010,7 +1007,7 @@ module Unix_tcp_server_ = struct k "@[Handler: uncaught exception for client %s:@ \ %s@ %s@]" - (str_of_sockaddr client_addr) + (Tiny_httpd_util.show_sockaddr client_addr) (Printexc.to_string e) (Printexc.raw_backtrace_to_string bt))); ignore Unix.(sigprocmask SIG_UNBLOCK Sys.[ sigint; sighup ]) @@ -1096,7 +1093,7 @@ let client_handle_for (self : t) ~client_addr ic oc : unit = let elapsed = B.get_time_s () -. req.start_time in k ("response to=%s code=%d time=%.3fs path=%S" : _ format4) - (str_of_sockaddr client_addr) + (Tiny_httpd_util.show_sockaddr client_addr) resp.code elapsed req.path in if Response_code.is_success resp.code then @@ -1113,7 +1110,9 @@ let client_handle_for (self : t) ~client_addr ic oc : unit = in if not Log.dummy then Log.error (fun k -> - k "response to %s code=%d" (str_of_sockaddr client_addr) resp.code); + k "response to %s code=%d" + (Tiny_httpd_util.show_sockaddr client_addr) + resp.code); Response.output_ ~buf:buf_res oc resp in @@ -1166,7 +1165,7 @@ let client_handle_for (self : t) ~client_addr ic oc : unit = ic in - UP.handle_connection handshake_st ic oc + UP.handle_connection client_addr handshake_st ic oc with e -> handle_bad_req req e in diff --git a/src/Tiny_httpd_server.mli b/src/Tiny_httpd_server.mli index 3842060f..591fff94 100644 --- a/src/Tiny_httpd_server.mli +++ b/src/Tiny_httpd_server.mli @@ -666,11 +666,16 @@ module type UPGRADE_HANDLER = sig The connection is closed without further ado. *) val handle_connection : - handshake_state -> Tiny_httpd_io.Input.t -> Tiny_httpd_io.Output.t -> unit - (** Take control of the connection and take it from there *) + Unix.sockaddr -> + handshake_state -> + Tiny_httpd_io.Input.t -> + Tiny_httpd_io.Output.t -> + unit + (** Take control of the connection and take it from ther.e *) end type upgrade_handler = (module UPGRADE_HANDLER) +(** @since NEXT_RELEASE *) val add_upgrade_handler : ?accept:(unit Request.t -> (unit, Response_code.t * string) result) -> diff --git a/src/Tiny_httpd_util.ml b/src/Tiny_httpd_util.ml index c87adbe4..9ec935ae 100644 --- a/src/Tiny_httpd_util.ml +++ b/src/Tiny_httpd_util.ml @@ -107,3 +107,8 @@ let parse_query s : (_ list, string) result = | Invalid_argument _ | Not_found | Failure _ -> Error (Printf.sprintf "error in parse_query for %S: i=%d,j=%d" s !i !j) | Invalid_query -> Error ("invalid query string: " ^ s) + +let show_sockaddr = function + | Unix.ADDR_UNIX f -> f + | Unix.ADDR_INET (inet, port) -> + Printf.sprintf "%s:%d" (Unix.string_of_inet_addr inet) port diff --git a/src/Tiny_httpd_util.mli b/src/Tiny_httpd_util.mli index f29209ce..ac996855 100644 --- a/src/Tiny_httpd_util.mli +++ b/src/Tiny_httpd_util.mli @@ -34,3 +34,7 @@ val parse_query : string -> ((string * string) list, string) result The order might not be preserved. @since 0.3 *) + +val show_sockaddr : Unix.sockaddr -> string +(** Simple printer for socket addresses. + @since NEXT_RELEASE *) diff --git a/src/ws/tiny_httpd_ws.ml b/src/ws/tiny_httpd_ws.ml index 869d5378..3637612b 100644 --- a/src/ws/tiny_httpd_ws.ml +++ b/src/ws/tiny_httpd_ws.ml @@ -6,7 +6,7 @@ module IO = Tiny_httpd_io let spf = Printf.sprintf let ( let@ ) = ( @@ ) -type handler = IO.Input.t -> IO.Output.t -> unit +type handler = Unix.sockaddr -> IO.Input.t -> IO.Output.t -> unit module Frame_type = struct type t = int @@ -427,7 +427,7 @@ end) : UPGRADE_HANDLER = struct let handshake req : _ result = try Ok (handshake_ req) with Bad_req s -> Error s - let handle_connection () ic oc = + let handle_connection addr () ic oc = let writer = Writer.create ~oc () in let reader = Reader.create ~ic ~writer () in let ws_ic : IO.Input.t = @@ -447,7 +447,7 @@ end) : UPGRADE_HANDLER = struct close = (fun () -> Writer.close writer); } in - try X.handler ws_ic ws_oc + try X.handler addr 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 f1265412..8ad8cd0f 100644 --- a/src/ws/tiny_httpd_ws.mli +++ b/src/ws/tiny_httpd_ws.mli @@ -2,9 +2,7 @@ open Common_ open Tiny_httpd_server module IO = Tiny_httpd_io -(* FIXME: also pass client address to the handler *) - -type handler = IO.Input.t -> IO.Output.t -> unit +type handler = Unix.sockaddr -> IO.Input.t -> IO.Output.t -> unit (** Websocket handler *) val add_route_handler :