expose client address to websocket

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

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
                           "@[<v>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
 

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) ->

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

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 *)

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");
       ()

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 :