Add FD passing

Author: talex5

lib_eio/unix/eio_unix.mli

@@ -32,6 +32,7 @@ module Resource : sig
   module type FLOW = sig
     include Eio.Net.Pi.STREAM_SOCKET
     include Eio.File.Pi.WRITE with type t := t
+    include Net.Pi.STREAM_SOCKET with type t := t
 
     val fd : t -> Fd.t
   end

lib_eio/unix/fd.ml

@@ -46,6 +46,11 @@ let of_unix ~sw ?blocking ?seekable ~close_unix fd =
   t.release_hook <- Switch.on_release_cancellable sw (fun () -> close t);
   t
 
+let of_unix_list ~sw fds =
+  match Switch.get_error sw with
+  | Some e -> List.iter Unix.close fds; raise e
+  | None -> List.map (of_unix ~sw ~close_unix:true) fds
+
 external eio_is_blocking : Unix.file_descr -> bool = "eio_unix_is_blocking"
 
 let is_blocking t =

lib_eio/unix/fd.mli

@@ -14,6 +14,10 @@ val of_unix : sw:Switch.t -> ?blocking:bool -> ?seekable:bool -> close_unix:bool
     @param seekable The value to be returned by {!is_seekable}. Defaults to probing if needed.
     @param close_unix Whether {!close} also closes [fd] (this should normally be [true]). *)
 
+val of_unix_list : sw:Switch.t -> Unix.file_descr list -> t list
+(** [of_unix_list ~sw fds] is like [List.map (of_unix ~sw ~close_unix:true) fds],
+    except that if [sw] is off then it closes all the FDs. *)
+
 (** {2 Using FDs} *)
 
 val use : t -> (Unix.file_descr -> 'a) -> if_closed:(unit -> 'a) -> 'a

lib_eio/unix/net.ml

@@ -30,6 +30,33 @@ let sockaddr_of_unix_datagram = function
     let host = Ipaddr.of_unix host in
     `Udp (host, port)
 
+module Pi = struct
+  module type STREAM_SOCKET = sig
+    type t
+
+    val send_msg : t -> fds:Fd.t list -> Cstruct.t list -> int
+
+    val recv_msg_with_fds : t -> sw:Switch.t -> max_fds:int -> Cstruct.t list -> int * Fd.t list
+  end
+
+  type (_, _, _) Eio.Resource.pi +=
+    | Stream_socket : ('t, (module STREAM_SOCKET with type t = 't), [> `Platform of [> `Unix] | `Socket | `Stream]) Eio.Resource.pi
+end
+
+let send_msg (Eio.Resource.T (t, ops)) ?(fds=[]) bufs =
+  let module X = (val (Eio.Resource.get ops Pi.Stream_socket)) in
+  let rec aux ~fds bufs =
+    let sent = X.send_msg t ~fds bufs in
+    match Cstruct.shiftv bufs sent with
+    | [] -> ()
+    | bufs -> aux bufs ~fds:[]
+  in
+  aux ~fds bufs
+
+let recv_msg_with_fds (Eio.Resource.T (t, ops)) ~sw ~max_fds bufs =
+  let module X = (val (Eio.Resource.get ops Pi.Stream_socket)) in
+  X.recv_msg_with_fds t ~sw ~max_fds bufs
+
 let getnameinfo (sockaddr : Eio.Net.Sockaddr.t) =
   let options =
     match sockaddr with

lib_eio/unix/net.mli

@@ -13,6 +13,24 @@ type 'a listening_socket = ([> listening_socket_ty] as 'a) r
 
 type t = [`Generic | `Unix] Eio.Net.ty r
 
+(** {2 Passing file descriptors} *)
+
+val send_msg :
+  [> `Platform of [>`Unix] | `Socket | `Stream] r ->
+  ?fds:Fd.t list ->
+  Cstruct.t list -> unit
+(** Like {!Eio.Flow.write}, but allows passing file descriptors (for Unix-domain sockets). *)
+
+val recv_msg_with_fds :
+  [> `Platform of [>`Unix] | `Socket | `Stream] r ->
+  sw:Switch.t ->
+  max_fds:int ->
+  Cstruct.t list ->
+  int * Fd.t list
+(** Like {!Eio.Flow.single_read}, but also allows receiving file descriptors (for Unix-domain sockets).
+
+    @param max_fds The maximum number of file descriptors to accept (additional ones will be closed). *)
+
 (** {2 Unix address conversions}
 
     Note: OCaml's {!Unix.sockaddr} type considers e.g. TCP port 80 and UDP port
@@ -71,6 +89,19 @@ val socketpair_datagram :
     This creates OS-level resources using [socketpair(2)].
     Note that, like all FDs created by Eio, they are both marked as close-on-exec by default. *)
 
+module Pi : sig
+  module type STREAM_SOCKET = sig
+    type t
+
+    val send_msg : t -> fds:Fd.t list -> Cstruct.t list -> int
+
+    val recv_msg_with_fds : t -> sw:Switch.t -> max_fds:int -> Cstruct.t list -> int * Fd.t list
+  end
+
+  type (_, _, _) Eio.Resource.pi +=
+    | Stream_socket : ('t, (module STREAM_SOCKET with type t = 't), [> `Platform of [> `Unix] | `Socket | `Stream]) Eio.Resource.pi
+end
+
 (** {2 Private API for backends} *)
 
 val getnameinfo : Eio.Net.Sockaddr.t -> (string * string)

lib_eio/unix/resource.ml

@@ -11,6 +11,7 @@ let fd_opt (Eio.Resource.T (t, ops)) =
 module type FLOW = sig
   include Eio.Net.Pi.STREAM_SOCKET
   include Eio.File.Pi.WRITE with type t := t
+  include Net.Pi.STREAM_SOCKET with type t := t
 
   val fd : t -> Fd.t
 end
@@ -20,6 +21,7 @@ let flow_handler (type t tag) (module X : FLOW with type t = t and type tag = ta
   Eio.Resource.bindings (Eio.Net.Pi.stream_socket (module X)) @
   Eio.Resource.bindings (Eio.File.Pi.rw (module X)) @ [
     H (T, X.fd);
+    H (Net.Pi.Stream_socket, (module X));
   ]
 
 module type DATAGRAM_SOCKET = sig

lib_eio_linux/eio_linux.ml

@@ -185,6 +185,13 @@ module Flow = struct
     | `Receive -> Unix.SHUTDOWN_RECEIVE
     | `Send -> Unix.SHUTDOWN_SEND
     | `All -> Unix.SHUTDOWN_ALL
+
+  let send_msg t ~fds data =
+    Low_level.send_msg t ~fds data
+
+  let recv_msg_with_fds t ~sw ~max_fds data =
+    let _addr, n, fds = Low_level.recv_msg_with_fds t ~sw ~max_fds data in
+    n, fds
 end
 
 let flow_handler = Eio_unix.Resource.flow_handler (module Flow)

lib_eio_linux/low_level.ml

@@ -263,10 +263,7 @@ let recv_msg_with_fds ~sw ~max_fds fd buf =
   if res < 0 then (
     raise @@ Err.wrap (Uring.error_of_errno res) "recv_msg" ""
   );
-  let fds =
-    Uring.Msghdr.get_fds msghdr
-    |> List.map (fun fd -> Fd.of_unix ~sw ~close_unix:true fd)
-  in
+  let fds = Uring.Msghdr.get_fds msghdr |> Fd.of_unix_list ~sw in
   addr, res, fds
 
 let with_chunk ~fallback fn =

lib_eio_linux/tests/fd_passing.md

@@ -212,23 +212,51 @@ CAMLprim value caml_eio_posix_spawn(value v_errors, value v_actions) {
   CAMLreturn(Val_long(child_pid));
 }
 
-CAMLprim value caml_eio_posix_send_msg(value v_fd, value v_dst_opt, value v_bufs) {
-  CAMLparam2(v_dst_opt, v_bufs);
+/* Copy [n_fds] from [v_fds] to [msg]. */
+static void fill_fds(struct msghdr *msg, int n_fds, value v_fds) {
+  if (n_fds > 0) {
+    int i;
+    struct cmsghdr *cm;
+    cm = CMSG_FIRSTHDR(msg);
+    cm->cmsg_level = SOL_SOCKET;
+    cm->cmsg_type = SCM_RIGHTS;
+    cm->cmsg_len = CMSG_LEN(n_fds * sizeof(int));
+    for (i = 0; i < n_fds; i++) {
+      int fd = -1;
+      if (Is_block(v_fds)) {
+	fd = Int_val(Field(v_fds, 0));
+	v_fds = Field(v_fds, 1);
+      }
+      ((int *)CMSG_DATA(cm))[i] = fd;
+    }
+  }
+}
+
+CAMLprim value caml_eio_posix_send_msg(value v_fd, value v_n_fds, value v_fds, value v_dst_opt, value v_bufs) {
+  CAMLparam3(v_fds, v_dst_opt, v_bufs);
   int n_bufs = Wosize_val(v_bufs);
+  int n_fds = Int_val(v_n_fds);
   struct iovec iov[n_bufs];
   union sock_addr_union dst_addr;
+  int controllen = n_fds > 0 ? CMSG_SPACE(sizeof(int) * n_fds) : 0;
+  char cmsg[controllen];
   struct msghdr msg = {
     .msg_iov = iov,
     .msg_iovlen = n_bufs,
+    .msg_control = n_fds > 0 ? cmsg : NULL,
+    .msg_controllen = controllen,
   };
   ssize_t r;
 
+  memset(cmsg, 0, controllen);
+
   if (Is_some(v_dst_opt)) {
     caml_unix_get_sockaddr(Some_val(v_dst_opt), &dst_addr, &msg.msg_namelen);
     msg.msg_name = &dst_addr;
   }
 
   fill_iov(iov, v_bufs);
+  fill_fds(&msg, n_fds, v_fds);
 
   caml_enter_blocking_section();
   r = sendmsg(Int_val(v_fd), &msg, 0);
@@ -238,20 +266,49 @@ CAMLprim value caml_eio_posix_send_msg(value v_fd, value v_dst_opt, value v_bufs
   CAMLreturn(Val_long(r));
 }
 
-CAMLprim value caml_eio_posix_recv_msg(value v_fd, value v_bufs) {
+static value get_msghdr_fds(struct msghdr *msg) {
+  CAMLparam0();
+  CAMLlocal2(v_list, v_cons);
+  struct cmsghdr *cm;
+  v_list = Val_int(0);
+  for (cm = CMSG_FIRSTHDR(msg); cm; cm = CMSG_NXTHDR(msg, cm)) {
+    if (cm->cmsg_level == SOL_SOCKET && cm->cmsg_type == SCM_RIGHTS) {
+      int *fds = (int *) CMSG_DATA(cm);
+      int n_fds = (cm->cmsg_len - CMSG_LEN(0)) / sizeof(int);
+      int i;
+      for (i = n_fds - 1; i >= 0; i--) {
+	value fd = Val_int(fds[i]);
+	v_cons = caml_alloc_tuple(2);
+	Store_field(v_cons, 0, fd);
+	Store_field(v_cons, 1, v_list);
+	v_list = v_cons;
+      }
+    }
+  }
+  CAMLreturn(v_list);
+}
+
+CAMLprim value caml_eio_posix_recv_msg(value v_fd, value v_max_fds, value v_bufs) {
   CAMLparam1(v_bufs);
   CAMLlocal2(v_result, v_addr);
+  int max_fds = Int_val(v_max_fds);
   int n_bufs = Wosize_val(v_bufs);
   struct iovec iov[n_bufs];
   union sock_addr_union source_addr;
+  int controllen = max_fds > 0 ? CMSG_SPACE(sizeof(int) * max_fds) : 0;
+  char cmsg[controllen];
   struct msghdr msg = {
     .msg_name = &source_addr,
     .msg_namelen = sizeof(source_addr),
     .msg_iov = iov,
     .msg_iovlen = n_bufs,
+    .msg_control = max_fds > 0 ? cmsg : NULL,
+    .msg_controllen = controllen,
   };
   ssize_t r;
 
+  memset(cmsg, 0, controllen);
+
   fill_iov(iov, v_bufs);
 
   caml_enter_blocking_section();
@@ -261,9 +318,10 @@ CAMLprim value caml_eio_posix_recv_msg(value v_fd, value v_bufs) {
 
   v_addr = caml_unix_alloc_sockaddr(&source_addr, msg.msg_namelen, -1);
 
-  v_result = caml_alloc_tuple(2);
+  v_result = caml_alloc_tuple(3);
   Store_field(v_result, 0, v_addr);
   Store_field(v_result, 1, Val_long(r));
+  Store_field(v_result, 2, get_msghdr_fds(&msg));
 
   CAMLreturn(v_result);
 }