fix(server): transfer handoff descriptors in batches (#3411)

Live handoff refused any session with more than 64 panes. The pane count
was checked twice against MAX_FDS_PER_HANDOFF, and the transfer itself put
every pane's pty master into one SCM_RIGHTS control message, so the guard
was the only thing keeping the send inside the kernel's per-message limit.
A session past the limit could only be updated by closing panes or by a
normal restart, which ends every pane process.

Send the descriptors in batches of 64 instead and drop both guards. The
receiving side accumulates across recvmsg calls until the expected count
arrives, bounds every SCM_RIGHTS payload it reads by the control bytes the
kernel returned, rejects a batch that carries more descriptors than it
asked for, and closes the descriptors it already holds on any failure. A
session of 64 panes or fewer still produces one batch, so the bytes on the
wire are unchanged and HANDOFF_VERSION stays at 1.

refs #3393

Co-authored-by: kataokatsuki <kataokatsuki@users.noreply.github.com>
Co-authored-by: JJ Liebig <jonathan.liebig@gmail.com>

2 files changed, 83 insertions, 40 deletions

diff --git a/src/server/handoff.rs b/src/server/handoff.rs
index 2cd17144..aaa87f38 100644
--- a/src/server/handoff.rs
+++ b/src/server/handoff.rs
@@ -22,8 +22,11 @@ const HANDOFF_VERSION: u32 = 1;
 const READY_TIMEOUT: Duration = Duration::from_secs(30);
 #[cfg(unix)]
 const OWNED_ACK_TIMEOUT: Duration = Duration::from_millis(500);
+// Descriptors are transferred in batches of this size. A single SCM_RIGHTS
+// control message caps out at 253 descriptors on Linux and 254 on macOS, so the
+// batch stays well below both limits and the number of panes stays unbounded.
 #[cfg(unix)]
-pub(crate) const MAX_FDS_PER_HANDOFF: usize = 64;
+const FDS_PER_MESSAGE: usize = 64;
 #[cfg(unix)]
 pub(crate) const MAX_REPLAY_BYTES_PER_PANE: usize = 8 * 1024;
 #[cfg(unix)]
@@ -174,12 +177,6 @@ pub(crate) fn accept_and_validate_on(
 
 #[cfg(unix)]
 pub(crate) fn send_fds_and_wait_restored(stream: &mut UnixStream, fds: &[RawFd]) -> io::Result<()> {
-    if fds.len() > MAX_FDS_PER_HANDOFF {
-        return Err(io::Error::new(
-            io::ErrorKind::InvalidInput,
-            format!("handoff supports at most {MAX_FDS_PER_HANDOFF} pane file descriptors at once"),
-        ));
-    }
     send_fds(stream, fds)?;
 
     stream.set_read_timeout(Some(READY_TIMEOUT))?;
@@ -382,6 +379,14 @@ fn read_line_unbuffered(stream: &mut UnixStream) -> io::Result<String> {
 
 #[cfg(unix)]
 fn send_fds(stream: &UnixStream, fds: &[RawFd]) -> io::Result<()> {
+    for batch in fds.chunks(FDS_PER_MESSAGE) {
+        send_fd_batch(stream, batch)?;
+    }
+    Ok(())
+}
+
+#[cfg(unix)]
+fn send_fd_batch(stream: &UnixStream, fds: &[RawFd]) -> io::Result<()> {
     if fds.is_empty() {
         return Ok(());
     }
@@ -415,16 +420,47 @@ fn send_fds(stream: &UnixStream, fds: &[RawFd]) -> io::Result<()> {
 }
 
 #[cfg(unix)]
+fn close_raw_fds(fds: &[RawFd]) {
+    for fd in fds {
+        let _ = unsafe { libc::close(*fd) };
+    }
+}
+
+#[cfg(unix)]
 fn recv_fds(stream: &UnixStream, expected: usize) -> io::Result<Vec<RawFd>> {
-    if expected == 0 {
-        return Ok(Vec::new());
+    let mut out: Vec<RawFd> = Vec::with_capacity(expected);
+    while out.len() < expected {
+        let wanted = (expected - out.len()).min(FDS_PER_MESSAGE);
+        let batch = match recv_fd_batch(stream, wanted) {
+            Ok(batch) => batch,
+            Err(err) => {
+                close_raw_fds(&out);
+                return Err(err);
+            }
+        };
+        if batch.is_empty() {
+            let received = out.len();
+            close_raw_fds(&out);
+            return Err(io::Error::new(
+                io::ErrorKind::UnexpectedEof,
+                format!(
+                    "handoff stream closed after {received} of {expected} pane file descriptors"
+                ),
+            ));
+        }
+        out.extend(batch);
     }
+    Ok(out)
+}
+
+#[cfg(unix)]
+fn recv_fd_batch(stream: &UnixStream, wanted: usize) -> io::Result<Vec<RawFd>> {
     let mut byte = [0u8; 1];
     let mut iov = [libc::iovec {
         iov_base: byte.as_mut_ptr() as *mut libc::c_void,
         iov_len: byte.len(),
     }];
-    let fd_bytes = expected * std::mem::size_of::<RawFd>();
+    let fd_bytes = wanted * std::mem::size_of::<RawFd>();
     let mut control = vec![0u8; unsafe { libc::CMSG_SPACE(fd_bytes as u32) as usize }];
     let mut msg: libc::msghdr = unsafe { std::mem::zeroed() };
     msg.msg_iov = iov.as_mut_ptr();
@@ -436,34 +472,51 @@ fn recv_fds(stream: &UnixStream, expected: usize) -> io::Result<Vec<RawFd>> {
     if read < 0 {
         return Err(io::Error::last_os_error());
     }
-    if msg.msg_flags & libc::MSG_CTRUNC != 0 {
-        return Err(io::Error::other("handoff fd control message was truncated"));
-    }
 
     let mut out = Vec::new();
     unsafe {
-        let cmsg = libc::CMSG_FIRSTHDR(&msg);
-        if cmsg.is_null()
-            || (*cmsg).cmsg_level != libc::SOL_SOCKET
-            || (*cmsg).cmsg_type != libc::SCM_RIGHTS
-        {
-            return Err(io::Error::other("handoff fd message missing SCM_RIGHTS"));
-        }
-        let data_len = ((*cmsg).cmsg_len as usize).saturating_sub(libc::CMSG_LEN(0) as usize);
-        let count = data_len / std::mem::size_of::<RawFd>();
-        let data = libc::CMSG_DATA(cmsg) as *const RawFd;
-        for idx in 0..count {
-            out.push(*data.add(idx));
+        let control_end = control.as_ptr() as usize + msg.msg_controllen as usize;
+        let mut cmsg = libc::CMSG_FIRSTHDR(&msg);
+        while !cmsg.is_null() {
+            if (*cmsg).cmsg_level == libc::SOL_SOCKET && (*cmsg).cmsg_type == libc::SCM_RIGHTS {
+                let data = libc::CMSG_DATA(cmsg);
+                // Bound the payload by both the header's own length and the
+                // bytes the kernel wrote into `control`, so the read below can
+                // never run past the buffer.
+                let available = control_end.saturating_sub(data as usize);
+                let data_len = ((*cmsg).cmsg_len as usize)
+                    .saturating_sub(libc::CMSG_LEN(0) as usize)
+                    .min(available);
+                let count = data_len / std::mem::size_of::<RawFd>();
+                let data = data as *const RawFd;
+                for idx in 0..count {
+                    out.push(*data.add(idx));
+                }
+            }
+            cmsg = libc::CMSG_NXTHDR(&msg, cmsg);
         }
     }
-    if out.len() != expected {
-        for fd in out {
-            let _ = unsafe { libc::close(fd) };
-        }
+
+    // Truncation means the kernel closed the descriptors that did not fit, so
+    // the batch is unrecoverable rather than merely short.
+    if msg.msg_flags & libc::MSG_CTRUNC != 0 {
+        close_raw_fds(&out);
+        return Err(io::Error::other("handoff fd control message was truncated"));
+    }
+    if read == 0 {
+        close_raw_fds(&out);
+        return Ok(Vec::new());
+    }
+    if out.len() > wanted {
+        let received = out.len();
+        close_raw_fds(&out);
         return Err(io::Error::other(format!(
-            "expected {expected} handoff fds, received fewer"
+            "handoff fd message carried {received} descriptors, expected at most {wanted}"
         )));
     }
+    if out.is_empty() {
+        return Err(io::Error::other("handoff fd message missing SCM_RIGHTS"));
+    }
     Ok(out)
 }
 
diff --git a/src/server/headless/lifecycle.rs b/src/server/headless/lifecycle.rs
index ab1448c3..f7ad87dd 100644
--- a/src/server/headless/lifecycle.rs
+++ b/src/server/headless/lifecycle.rs
@@ -53,16 +53,6 @@ impl HeadlessServer {
                 }
             }
         }
-        if pane_by_terminal.len() > crate::server::handoff::MAX_FDS_PER_HANDOFF {
-            let _ = std::fs::remove_file(&socket_path);
-            return Err(io::Error::new(
-                io::ErrorKind::InvalidInput,
-                format!(
-                    "live handoff supports at most {} panes in one update; close panes or restart herdr normally",
-                    crate::server::handoff::MAX_FDS_PER_HANDOFF
-                ),
-            ));
-        }
 
         self.handoff_in_progress = true;
         self.disconnect_all_clients_for_handoff();