compositor: add common parent for both compostior and forker

src/utils/clone3.rs

@@ -1,6 +1,10 @@
 use {
-    crate::forker::ForkerError,
-    uapi::{OwnedFd, c},
+    crate::{
+        forker::ForkerError,
+        utils::{errorfmt::ErrorFmt, on_drop::OnDrop, process_name::set_process_name},
+    },
+    std::{env, mem::MaybeUninit, process, slice, str::FromStr},
+    uapi::{Msghdr, MsghdrMut, OwnedFd, c},
 };
 
 #[derive(Default, Copy, Clone)]
@@ -21,9 +25,11 @@ struct clone_args {
 
 pub enum Forked {
     Parent { pid: c::pid_t, pidfd: OwnedFd },
-    Child { _pidfd: Option<OwnedFd> },
+    Child { pidfd: Option<OwnedFd> },
 }
 
+const REAPER_VAR: &str = "JAY_REAPER_PID";
+
 pub fn fork_with_pidfd(pidfd_for_child: bool) -> Result<Forked, ForkerError> {
     let mut pidfd: c::c_int = 0;
     let mut args = clone_args {
@@ -46,9 +52,8 @@ pub fn fork_with_pidfd(pidfd_for_child: bool) -> Result<Forked, ForkerError> {
             return Err(ForkerError::Fork(e.into()));
         }
         let res = if pid == 0 {
-            Forked::Child {
-                _pidfd: child_pidfd,
-            }
+            env::remove_var(REAPER_VAR);
+            Forked::Child { pidfd: child_pidfd }
         } else {
             Forked::Parent {
                 pid: pid as _,
@@ -58,3 +63,109 @@ pub fn fork_with_pidfd(pidfd_for_child: bool) -> Result<Forked, ForkerError> {
         Ok(res)
     }
 }
+
+pub fn double_fork() -> Result<Option<OwnedFd>, ForkerError> {
+    let (p, c) = uapi::socketpair(c::AF_UNIX, c::SOCK_DGRAM | c::SOCK_CLOEXEC, 0)
+        .map_err(|e| ForkerError::Socketpair(e.into()))?;
+    match fork_with_pidfd(false)? {
+        Forked::Parent { pid, .. } => {
+            drop(c);
+            let mut buf = [MaybeUninit::<u8>::uninit(); 128];
+            let iov: &mut [&mut [u8]] = &mut [];
+            let mut msghdr = MsghdrMut {
+                iov,
+                control: Some(&mut buf),
+                name: uapi::sockaddr_none_mut(),
+                flags: 0,
+            };
+            let _wait = OnDrop(|| {
+                let _ = uapi::waitpid(pid, 0);
+            });
+            let (_, _, mut ctrl) = uapi::recvmsg(p.raw(), &mut msghdr, c::MSG_CMSG_CLOEXEC)
+                .map_err(|e| ForkerError::RecvPidfd(e.into()))?;
+            let (_, hdr, data) =
+                uapi::cmsg_read(&mut ctrl).map_err(|e| ForkerError::CmsgRead(e.into()))?;
+            if hdr.cmsg_level != c::SOL_SOCKET || hdr.cmsg_type != c::SCM_RIGHTS {
+                return Err(ForkerError::InvalidCmsg);
+            }
+            let Ok(fd) = uapi::pod_read(data) else {
+                return Err(ForkerError::InvalidCmsg);
+            };
+            Ok(Some(OwnedFd::new(fd)))
+        }
+        Forked::Child { .. } => {
+            drop(p);
+            if let Ok(f) = fork_with_pidfd(true) {
+                match f {
+                    Forked::Parent { pidfd, .. } => {
+                        let pidfd = pidfd.raw();
+                        let mut buf = [MaybeUninit::uninit(); 128];
+                        let hdr = c::cmsghdr {
+                            cmsg_len: 0,
+                            cmsg_level: c::SOL_SOCKET,
+                            cmsg_type: c::SCM_RIGHTS,
+                        };
+                        let _ = uapi::cmsg_write(&mut &mut buf[..], hdr, &pidfd);
+                        let iov: &[&[u8]] = &[];
+                        let msghdr = Msghdr {
+                            iov,
+                            control: Some(&buf[..uapi::cmsg_space(size_of_val(&pidfd))]),
+                            name: uapi::sockaddr_none_ref(),
+                        };
+                        let _ = uapi::sendmsg(c.raw(), &msghdr, 0);
+                    }
+                    Forked::Child { pidfd } => {
+                        let pidfd = pidfd.unwrap();
+                        let mut pollfd = c::pollfd {
+                            fd: pidfd.raw(),
+                            events: c::POLLIN as _,
+                            revents: 0,
+                        };
+                        let _ = uapi::poll(slice::from_mut(&mut pollfd), -1);
+                        return Ok(None);
+                    }
+                }
+            };
+            unsafe {
+                c::_exit(0);
+            }
+        }
+    }
+}
+
+pub fn ensure_reaper() -> c::pid_t {
+    if let Ok(id) = env::var(REAPER_VAR) {
+        if let Ok(id) = c::pid_t::from_str(&id) {
+            if uapi::getppid() == id {
+                return id;
+            }
+        }
+    }
+    let reaper_pid = uapi::getpid();
+    unsafe {
+        c::prctl(c::PR_SET_CHILD_SUBREAPER, 1);
+    }
+    let res = match fork_with_pidfd(false) {
+        Ok(r) => r,
+        Err(e) => {
+            fatal!("Could not fork reaper: {}", ErrorFmt(e));
+        }
+    };
+    let Forked::Parent {
+        pid: main_process_id,
+        ..
+    } = res
+    else {
+        unsafe {
+            env::set_var(REAPER_VAR, reaper_pid.to_string());
+        }
+        return reaper_pid;
+    };
+    set_process_name("jay reaper");
+    while let Ok((pid, status)) = uapi::wait() {
+        if pid == main_process_id {
+            process::exit(uapi::WEXITSTATUS(status));
+        }
+    }
+    process::exit(1);
+}