Close the transfer socket and set CLOEXEC on received listener fds

pingora-core/src/server/transfer_fd/mod.rs

@@ -25,7 +25,7 @@ use std::io::Write;
 #[cfg(target_os = "linux")]
 use std::io::{IoSlice, IoSliceMut};
 #[cfg(target_os = "linux")]
-use std::os::fd::{AsRawFd, BorrowedFd};
+use std::os::fd::{AsRawFd, BorrowedFd, FromRawFd, OwnedFd};
 use std::os::unix::io::RawFd;
 #[cfg(target_os = "linux")]
 use std::{thread, time};
@@ -188,13 +188,24 @@ where
         }
     };
 
+    // The accepted connection is only needed for this transfer. Take ownership of it so
+    // that it is closed on every path out of this function, including the early return
+    // from `cmsgs()?` below; otherwise every graceful upgrade leaks one unix socket.
+    //
+    // SAFETY: `fd` was just returned by accept(2) and is not owned or closed anywhere else.
+    let conn = unsafe { OwnedFd::from_raw_fd(fd) };
+
     let mut io_vec = [IoSliceMut::new(payload); 1];
     let mut cmsg_buf = nix::cmsg_space!([RawFd; MAX_FDS]);
+    // MSG_CMSG_CLOEXEC sets FD_CLOEXEC on the received descriptors atomically. These
+    // listening sockets are kept for the rest of the process's lifetime, so without it any
+    // subprocess an application built on pingora execs inherits them, and can keep the
+    // port bound after the server itself is gone.
     let msg: RecvMsg<UnixAddr> = socket::recvmsg(
-        fd,
+        conn.as_raw_fd(),
         &mut io_vec,
         Some(&mut cmsg_buf),
-        socket::MsgFlags::empty(),
+        socket::MsgFlags::MSG_CMSG_CLOEXEC,
     )
     .unwrap();
 
@@ -387,6 +398,7 @@ mod tests {
 
     use super::*;
     use log::{debug, error};
+    use nix::fcntl;
 
     fn init_log() {
         let _ = env_logger::builder().is_test(true).try_init();
@@ -468,6 +480,90 @@ mod tests {
         child.join().unwrap();
     }
 
+    /// How many fds in this process refer to a unix socket bound to `path`.
+    ///
+    /// Reads the socket inodes bound to `path` from /proc/net/unix, then scans
+    /// /proc/self/fd for descriptors pointing at them. Filtering by path keeps this
+    /// unaffected by unrelated descriptors opened by tests running in parallel.
+    fn unix_socket_fds_bound_to(path: &str) -> usize {
+        let unix = std::fs::read_to_string("/proc/net/unix").unwrap();
+        let inodes: HashSet<&str> = unix
+            .lines()
+            .filter_map(|line| {
+                let mut cols = line.split_whitespace();
+                let inode = cols.nth(6)?;
+                (cols.next() == Some(path)).then_some(inode)
+            })
+            .collect();
+        if inodes.is_empty() {
+            return 0;
+        }
+        std::fs::read_dir("/proc/self/fd")
+            .unwrap()
+            .filter_map(|entry| std::fs::read_link(entry.ok()?.path()).ok())
+            .filter(|target| {
+                target
+                    .to_str()
+                    .and_then(|t| t.strip_prefix("socket:["))
+                    .and_then(|t| t.strip_suffix(']'))
+                    .is_some_and(|inode| inodes.contains(inode))
+            })
+            .count()
+    }
+
+    #[test]
+    fn test_receive_does_not_leak_fds() {
+        init_log();
+        const SOCK: &str = "/tmp/pingora_fds_receive3.sock";
+
+        let dumb_fd = socket::socket(
+            AddressFamily::Unix,
+            SockType::Stream,
+            SockFlag::empty(),
+            None,
+        )
+        .unwrap();
+
+        // receiver need to start in another thread since it is blocking
+        let child = thread::spawn(move || {
+            let mut buf: [u8; 32] = [0; 32];
+            let (fds, _) = get_fds_from(SOCK, &mut buf, None).unwrap();
+            assert_eq!(1, fds.len());
+
+            // The received listener is kept for the lifetime of the process, so it must
+            // not be inherited by unrelated children across exec().
+            let flags = fcntl::fcntl(
+                // SAFETY: the fd was just received and stays open for this call.
+                unsafe { BorrowedFd::borrow_raw(fds[0]) },
+                fcntl::FcntlArg::F_GETFD,
+            )
+            .unwrap();
+            assert!(
+                fcntl::FdFlag::from_bits_truncate(flags).contains(fcntl::FdFlag::FD_CLOEXEC),
+                "fd received over SCM_RIGHTS is missing FD_CLOEXEC"
+            );
+
+            // The accepted connection is only needed during the transfer itself.
+            assert_eq!(
+                0,
+                unix_socket_fds_bound_to(SOCK),
+                "the accepted transfer socket was left open"
+            );
+
+            // Don't leak the descriptors this test just received.
+            for fd in fds {
+                // SAFETY: received over SCM_RIGHTS just above and not owned anywhere else.
+                drop(unsafe { OwnedFd::from_raw_fd(fd) });
+            }
+        });
+
+        let fds = vec![dumb_fd.as_raw_fd()];
+        let buf: [u8; 32] = [1; 32];
+        send_fds_to(fds, &buf, SOCK, None).unwrap();
+
+        child.join().unwrap();
+    }
+
     #[test]
     fn test_serde_via_socket() {
         init_log();