v0.17~preview.128.40+46

Author: public-release

async_rpc/src/rpc.ml

@@ -143,6 +143,16 @@ module Connection = struct
     Result.ok_exn res
   ;;
 
+  let connection_description ~server_addr ~client_addr =
+    let server_addr = (server_addr :> Socket.Address.t) in
+    let client_addr = (client_addr :> Socket.Address.t) in
+    Info.create_s
+      [%message
+        "TCP server" (server_addr : Socket.Address.t) (client_addr : Socket.Address.t)]
+  ;;
+
+  let default_on_handshake_error = `Ignore
+
   let make_serve_func
     serve_with_transport_handler
     ~implementations
@@ -157,7 +167,7 @@ module Connection = struct
     ?handshake_timeout
     ?heartbeat_config
     ?auth
-    ?(on_handshake_error = `Ignore)
+    ?(on_handshake_error = default_on_handshake_error)
     ?on_handler_error
     ()
     =
@@ -172,18 +182,11 @@ module Connection = struct
       ?auth
       ?on_handler_error
       (fun ~client_addr ~server_addr transport ->
-      let description =
-        let server_addr = (server_addr :> Socket.Address.t) in
-        let client_addr = (client_addr :> Socket.Address.t) in
-        Info.create_s
-          [%message
-            "TCP server" (server_addr : Socket.Address.t) (client_addr : Socket.Address.t)]
-      in
       serve_with_transport
         ~handshake_timeout
         ~heartbeat_config
         ~implementations
-        ~description
+        ~description:(connection_description ~server_addr ~client_addr)
         ~connection_state:(fun conn -> initial_connection_state client_addr conn)
         ~on_handshake_error
         ~client_addr
@@ -198,6 +201,46 @@ module Connection = struct
     make_serve_func Rpc_transport.Tcp.serve_inet ~implementations
   ;;
 
+  let serve_unix
+    ~implementations
+    ~initial_connection_state
+    ~where_to_listen
+    ?max_connections
+    ?backlog
+    ?drop_incoming_connections
+    ?time_source
+    ?max_message_size
+    ?make_transport
+    ?handshake_timeout
+    ?heartbeat_config
+    ?auth
+    ?(on_handshake_error = default_on_handshake_error)
+    ?on_handler_error
+    ()
+    =
+    Rpc_transport.Tcp.serve_unix
+      ~where_to_listen
+      ?max_connections
+      ?backlog
+      ?drop_incoming_connections
+      ?time_source
+      ?max_message_size
+      ?make_transport
+      ?auth
+      ?on_handler_error
+      (fun ~client_addr ~server_addr peer_creds transport ->
+      serve_with_transport
+        ~handshake_timeout
+        ~heartbeat_config
+        ~implementations
+        ~description:(connection_description ~server_addr ~client_addr)
+        ~connection_state:(fun conn ->
+          initial_connection_state client_addr peer_creds conn)
+        ~on_handshake_error
+        ~client_addr
+        transport)
+  ;;
+
   let default_handshake_timeout_float =
     Time_ns.Span.to_span_float_round_nearest
       Async_rpc_kernel.Async_rpc_kernel_private.default_handshake_timeout

async_rpc/src/rpc.mli

@@ -140,6 +140,30 @@ module Connection : sig
     -> unit
     -> (Socket.Address.Inet.t, int) Tcp.Server.t
 
+  (** As [serve], but only accepts Unix sockets; provides peer credentials of the socket
+      to [initial_connection_state]. *)
+  val serve_unix
+    :  implementations:'s Implementations.t
+    -> initial_connection_state:
+         (Socket.Address.Unix.t -> Linux_ext.Peer_credentials.t -> t -> 's)
+    -> where_to_listen:Tcp.Where_to_listen.unix
+    -> ?max_connections:int
+    -> ?backlog:int
+    -> ?drop_incoming_connections:bool
+    -> ?time_source:[> read ] Time_source.T1.t
+    -> ?max_message_size:int
+    -> ?make_transport:transport_maker
+    -> ?handshake_timeout:Time_float.Span.t
+    -> ?heartbeat_config:Heartbeat_config.t
+    -> ?auth:(Socket.Address.Unix.t -> bool) (** default is [`Ignore] *)
+    -> ?on_handshake_error:
+         [ `Raise | `Ignore | `Call of Socket.Address.Unix.t -> exn -> unit ]
+         (** default is [`Ignore] *)
+    -> ?on_handler_error:
+         [ `Raise | `Ignore | `Call of Socket.Address.Unix.t -> exn -> unit ]
+    -> unit
+    -> Tcp.Server.unix Deferred.t
+
   (** [client where_to_connect ()] connects to the server at [where_to_connect] and
       returns the connection or an Error if a connection could not be made. It is the
       responsibility of the caller to eventually call [close].

async_rpc/src/rpc_transport.ml

@@ -5,12 +5,12 @@ module Header = Kernel_transport.Header
 module Handler_result = Kernel_transport.Handler_result
 module Send_result = Kernel_transport.Send_result
 
-let environment_variable = "ASYNC_RPC_MAX_MESSAGE_SIZE"
+let max_message_size_env_var = "ASYNC_RPC_MAX_MESSAGE_SIZE"
 
 let max_message_size_from_environment =
   lazy
     (Option.try_with_join (fun () ->
-       Sys.getenv environment_variable |> Option.map ~f:Int.of_string))
+       Sys.getenv max_message_size_env_var |> Option.map ~f:Int.of_string))
 ;;
 
 let aux_effective_max_message_size ~max_message_size_from_environment ~proposed_max =
@@ -99,7 +99,9 @@ end = struct
     then
       failwiths
         ~here:[%here]
-        "Rpc_transport: message too small or too big"
+        [%string
+          "Rpc_transport: message is too large or has negative size. Try increasing the \
+           size limit by setting the %{max_message_size_env_var} env var"]
         (`Message_size payload_len, `Max_message_size t.max_message_size)
         [%sexp_of: [ `Message_size of int ] * [ `Max_message_size of int ]]
   ;;
@@ -297,7 +299,7 @@ let of_fd ?buffer_age_limit ?reader_buffer_size ?writer_buffer_size ~max_message
 module Tcp = struct
   let default_transport_maker fd ~max_message_size = of_fd fd ~max_message_size
 
-  let make_serve_func
+  let make_serve_func_with_fd
     tcp_creator
     ~where_to_listen
     ?max_connections
@@ -324,10 +326,12 @@ module Tcp = struct
       | false -> return ()
       | true ->
         let max_message_size = effective_max_message_size ~proposed_max in
-        let transport = make_transport ~max_message_size (Socket.fd socket) in
+        let fd = Socket.fd socket in
+        let transport = make_transport ~max_message_size fd in
         let%bind result =
           Monitor.try_with ~run:`Schedule ~rest:`Raise (fun () ->
             handle_transport
+              fd
               ~client_addr
               ~server_addr:(Socket.getsockname socket)
               transport)
@@ -338,6 +342,34 @@ module Tcp = struct
          | Error exn -> raise exn))
   ;;
 
+  let make_serve_func
+    tcp_creator
+    ~where_to_listen
+    ?max_connections
+    ?backlog
+    ?drop_incoming_connections
+    ?time_source
+    ?max_message_size
+    ?make_transport
+    ?auth
+    ?on_handler_error
+    handle_transport
+    =
+    make_serve_func_with_fd
+      tcp_creator
+      ~where_to_listen
+      ?max_connections
+      ?backlog
+      ?drop_incoming_connections
+      ?time_source
+      ?max_message_size
+      ?make_transport
+      ?auth
+      ?on_handler_error
+      (fun (_ : Fd.t) ~client_addr ~server_addr transport ->
+      handle_transport ~client_addr ~server_addr transport)
+  ;;
+
   (* eta-expand [where_to_listen] to avoid value restriction. *)
   let serve ~where_to_listen = make_serve_func Tcp.Server.create_sock ~where_to_listen
 
@@ -346,6 +378,39 @@ module Tcp = struct
     make_serve_func Tcp.Server.create_sock_inet ~where_to_listen
   ;;
 
+  let serve_unix
+    ~(where_to_listen : Tcp.Where_to_listen.unix)
+    ?max_connections
+    ?backlog
+    ?drop_incoming_connections
+    ?time_source
+    ?max_message_size
+    ?make_transport
+    ?auth
+    ?on_handler_error
+    handle_transport
+    =
+    make_serve_func_with_fd
+      Tcp.Server.create_sock
+      ~where_to_listen
+      ?max_connections
+      ?backlog
+      ?drop_incoming_connections
+      ?time_source
+      ?max_message_size
+      ?make_transport
+      ?auth
+      ?on_handler_error
+      (fun fd ~client_addr ~server_addr transport ->
+      let peer_credentials =
+        Or_error.try_with (fun () ->
+          (ok_exn Linux_ext.peer_credentials) (Fd.file_descr_exn fd))
+        |> Or_error.tag ~tag:"Error getting peer credentials of unix socket"
+        |> ok_exn
+      in
+      handle_transport ~client_addr ~server_addr peer_credentials transport)
+  ;;
+
   let connect
     ?max_message_size:proposed_max
     ?(make_transport = default_transport_maker)

async_rpc/src/rpc_transport.mli

@@ -85,6 +85,26 @@ module Tcp : sig
         -> unit Deferred.t)
     -> (Socket.Address.Inet.t, int) Tcp.Server.t
 
+  (** [serve_inet] is like [serve] but only for unix sockets (not inet sockets), and
+      returns the identity of the peer on the socket. *)
+  val serve_unix
+    :  where_to_listen:Tcp.Where_to_listen.unix
+    -> ?max_connections:int
+    -> ?backlog:int
+    -> ?drop_incoming_connections:bool
+    -> ?time_source:[> read ] Time_source.T1.t
+    -> ?max_message_size:int
+    -> ?make_transport:transport_maker
+    -> ?auth:(Socket.Address.Unix.t -> bool)
+    -> ?on_handler_error:
+         [ `Raise | `Ignore | `Call of Socket.Address.Unix.t -> exn -> unit ]
+    -> (client_addr:Socket.Address.Unix.t
+        -> server_addr:Socket.Address.Unix.t
+        -> Linux_ext.Peer_credentials.t
+        -> Rpc_kernel.Transport.t
+        -> unit Deferred.t)
+    -> (Socket.Address.Unix.t, string) Tcp.Server.t Deferred.t
+
   (** [connect ?make_transport where_to_connect ()] connects to the server at
       [where_to_connect]. On success, it returns the transport created using
       [make_transport] and the [Socket.Address.t] that it connected to, otherwise it

async_rpc/test/dune

@@ -1,2 +1,3 @@
-(library (name async_rpc_test) (libraries async re2 sexp_select)
+(library (name async_rpc_test)
+ (libraries async expect_test_helpers_async re2 sexp_select)
  (preprocess (pps ppx_jane)))

async_rpc/test/test_rpc.ml

@@ -203,6 +203,35 @@ let%test_unit "Open dispatches see connection closed error" =
     Tcp.Server.close server)
 ;;
 
+let%expect_test "serve_unix returns the identity of the caller" =
+  Expect_test_helpers_async.with_temp_dir
+    ~in_dir:
+      (let force_tmp_to_avoid_unix_socket_path_length_problems = "/tmp" in
+       force_tmp_to_avoid_unix_socket_path_length_problems)
+    (fun tempdir ->
+      let socket_path = tempdir ^/ "socket" in
+      let%bind server =
+        Rpc.Connection.serve_unix
+          ~initial_connection_state:
+            (fun
+              (_ : Socket.Address.Unix.t)
+              (peer_creds : Linux_ext.Peer_credentials.t)
+              (_ : Rpc.Connection.t)
+              ->
+            [%test_result: int] peer_creds.uid ~expect:(Unix.getuid ());
+            [%test_result: int] peer_creds.gid ~expect:(Unix.getgid ());
+            [%test_result: Pid.t] peer_creds.pid ~expect:(Unix.getpid ()))
+          ~implementations:(Rpc.Implementations.null ())
+          ~where_to_listen:(Tcp.Where_to_listen.of_file socket_path)
+          ()
+      in
+      let%bind connection =
+        Rpc.Connection.client (Tcp.Where_to_connect.of_file socket_path) >>| Result.ok_exn
+      in
+      let%bind () = Rpc.Connection.close connection in
+      Tcp.Server.close server)
+;;
+
 let%test_module "Exception handling" =
   (module struct
     let on_exception ~callback_triggered ~expect_close_connection =