Implement exact bootstrap recovery and artifact GC

crates/lazaret-linux/src/ipc.rs

@@ -1,12 +1,16 @@
-use std::io;
+use std::io::{self, IoSlice, IoSliceMut};
+use std::mem::MaybeUninit;
 use std::os::fd::{AsFd, OwnedFd};
 
 use lazaret_core::codec::{decode, encode};
 use rustix::net::{
-    AddressFamily, RecvFlags, SendFlags, SocketFlags, SocketType, recv, send, socketpair,
+    AddressFamily, RecvAncillaryBuffer, RecvAncillaryMessage, RecvFlags, ReturnFlags,
+    SendAncillaryBuffer, SendAncillaryMessage, SendFlags, SocketFlags, SocketType, recv, recvmsg,
+    send, sendmsg, socketpair,
 };
 
 const MAX_PRIVATE_PACKET: usize = 64 * 1024;
+const MAX_TRANSFERRED_DESCRIPTORS: usize = 8;
 
 /// # Errors
 /// Returns an error when the private sequenced-packet pair cannot be created.
@@ -61,13 +65,85 @@ where
     decode(&encoded[..received]).map_err(|error| io::Error::other(error.to_string()))
 }
 
+/// Sends one nonempty bounded descriptor set with a one-byte packet.
+///
+/// # Errors
+/// Refuses empty, oversized, incomplete, or failed descriptor transfer.
+pub fn send_descriptors(
+    socket: impl AsFd,
+    descriptors: &[std::os::fd::BorrowedFd<'_>],
+) -> io::Result<()> {
+    if descriptors.is_empty() || descriptors.len() > MAX_TRANSFERRED_DESCRIPTORS {
+        return Err(io::Error::other("private descriptor set has invalid size"));
+    }
+    let mut space =
+        [MaybeUninit::uninit(); rustix::cmsg_space!(ScmRights(MAX_TRANSFERRED_DESCRIPTORS))];
+    let mut control = SendAncillaryBuffer::new(&mut space);
+    if !control.push(SendAncillaryMessage::ScmRights(descriptors)) {
+        return Err(io::Error::other(
+            "private descriptor control buffer is too small",
+        ));
+    }
+    let bytes = [1_u8];
+    let sent = sendmsg(
+        socket,
+        &[IoSlice::new(&bytes)],
+        &mut control,
+        SendFlags::empty(),
+    )?;
+    if sent == bytes.len() {
+        Ok(())
+    } else {
+        Err(io::Error::other("private descriptor packet was incomplete"))
+    }
+}
+
+/// Receives exactly `expected` descriptors with close-on-exec set atomically.
+///
+/// # Errors
+/// Refuses missing, extra, truncated, malformed, or failed descriptor transfer.
+pub fn receive_descriptors(socket: impl AsFd, expected: usize) -> io::Result<Vec<OwnedFd>> {
+    if expected == 0 || expected > MAX_TRANSFERRED_DESCRIPTORS {
+        return Err(io::Error::other("expected descriptor count is invalid"));
+    }
+    let mut byte = [0_u8];
+    let mut slices = [IoSliceMut::new(&mut byte)];
+    let mut space =
+        [MaybeUninit::uninit(); rustix::cmsg_space!(ScmRights(MAX_TRANSFERRED_DESCRIPTORS))];
+    let mut control = RecvAncillaryBuffer::new(&mut space);
+    let received = recvmsg(socket, &mut slices, &mut control, RecvFlags::CMSG_CLOEXEC)?;
+    if received.bytes != 1
+        || byte != [1]
+        || received
+            .flags
+            .intersects(ReturnFlags::TRUNC | ReturnFlags::CTRUNC)
+    {
+        return Err(io::Error::other("private descriptor packet is malformed"));
+    }
+    let mut descriptors = Vec::new();
+    for message in control.drain() {
+        match message {
+            RecvAncillaryMessage::ScmRights(rights) => descriptors.extend(rights),
+            _ => return Err(io::Error::other("private ancillary message is unexpected")),
+        }
+    }
+    if descriptors.len() == expected {
+        Ok(descriptors)
+    } else {
+        Err(io::Error::other(
+            "private descriptor count mismatches request",
+        ))
+    }
+}
+
 #[cfg(test)]
 mod tests {
-    use std::os::fd::AsFd;
+    use std::fs::File;
+    use std::os::fd::{AsFd, AsRawFd};
 
     use rustix::net::{RecvFlags, SendFlags, recv, send};
 
-    use super::private_seqpacket_pair;
+    use super::{private_seqpacket_pair, receive_descriptors, send_descriptors};
 
     #[test]
     fn preserves_message_boundaries() {
@@ -80,4 +156,22 @@ mod tests {
         let (_, second) = recv(right.as_fd(), &mut buffer, RecvFlags::empty()).unwrap();
         assert_eq!(&buffer[..second], b"second");
     }
+
+    #[test]
+    fn transfers_exact_close_on_exec_descriptors() {
+        let (left, right) = private_seqpacket_pair().unwrap();
+        let first = File::open("/dev/null").unwrap();
+        let second = File::open("/dev/zero").unwrap();
+        send_descriptors(left, &[first.as_fd(), second.as_fd()]).unwrap();
+        let received = receive_descriptors(right, 2).unwrap();
+        assert_ne!(received[0].as_raw_fd(), first.as_raw_fd());
+        assert_ne!(received[1].as_raw_fd(), second.as_raw_fd());
+        for descriptor in received {
+            assert!(
+                rustix::io::fcntl_getfd(&descriptor)
+                    .unwrap()
+                    .contains(rustix::io::FdFlags::CLOEXEC)
+            );
+        }
+    }
 }