use smallvec::SmallVec;
use std::os::unix::io::RawFd;
#[derive(Debug, Default)]
pub struct SocketAncillary {
recv_buf: Vec<u8>,
send_fds: SmallVec<[RawFd; 8]>,
recv_fds: SmallVec<[RawFd; 8]>,
truncated: bool,
}
impl SocketAncillary {
#[must_use]
pub fn new(recv_capacity: usize) -> Self {
Self {
recv_buf: Vec::with_capacity(recv_capacity),
send_fds: SmallVec::new(),
recv_fds: SmallVec::new(),
truncated: false,
}
}
#[must_use]
pub fn capacity(&self) -> usize {
self.recv_buf.capacity()
}
#[must_use]
pub fn is_empty(&self) -> bool {
self.send_fds.is_empty() && self.recv_fds.is_empty()
}
#[must_use]
pub fn is_truncated(&self) -> bool {
self.truncated
}
pub fn clear(&mut self) {
self.send_fds.clear();
self.recv_fds.clear();
self.truncated = false;
self.recv_buf.clear();
}
pub fn add_fds(&mut self, fds: &[RawFd]) -> bool {
self.send_fds.extend_from_slice(fds);
true
}
#[must_use]
pub fn messages(&self) -> AncillaryMessages<'_> {
AncillaryMessages {
yielded: false,
recv_fds: &self.recv_fds,
}
}
pub(crate) fn send_fds(&self) -> &[RawFd] {
&self.send_fds
}
pub(crate) fn clear_send_fds(&mut self) {
self.send_fds.clear();
}
pub(crate) fn prepare_for_recv(&mut self) -> &mut [u8] {
self.recv_fds.clear();
self.truncated = false;
self.recv_buf.resize(self.recv_buf.capacity(), 0);
self.recv_buf.as_mut_slice()
}
pub(crate) fn push_received_fds(&mut self, fds: &[RawFd]) {
self.recv_fds.extend_from_slice(fds);
}
pub(crate) fn mark_truncated(&mut self) {
self.truncated = true;
}
}
#[derive(Debug)]
pub struct AncillaryMessages<'a> {
yielded: bool,
recv_fds: &'a [RawFd],
}
impl<'a> Iterator for AncillaryMessages<'a> {
type Item = AncillaryMessage<'a>;
fn next(&mut self) -> Option<Self::Item> {
if self.yielded {
return None;
}
self.yielded = true;
if self.recv_fds.is_empty() {
None
} else {
Some(AncillaryMessage::ScmRights(ScmRights {
fds: self.recv_fds,
}))
}
}
}
#[derive(Debug)]
pub enum AncillaryMessage<'a> {
ScmRights(ScmRights<'a>),
}
#[derive(Debug)]
pub struct ScmRights<'a> {
fds: &'a [RawFd],
}
impl Iterator for ScmRights<'_> {
type Item = RawFd;
fn next(&mut self) -> Option<Self::Item> {
self.fds.split_first().map(|(fd, rest)| {
self.fds = rest;
*fd
})
}
fn size_hint(&self) -> (usize, Option<usize>) {
(self.fds.len(), Some(self.fds.len()))
}
}
impl ExactSizeIterator for ScmRights<'_> {}
#[must_use]
pub fn ancillary_space_for_fds(fd_count: usize) -> usize {
if fd_count == 0 {
0
} else {
let per = nix::sys::socket::cmsg_space::<RawFd>();
fd_count.saturating_mul(per)
}
}
#[cfg(test)]
mod tests {
use super::*;
fn init_test(name: &str) {
crate::test_utils::init_test_logging();
crate::test_phase!(name);
}
#[test]
fn test_ancillary_new() {
init_test("test_ancillary_new");
let ancillary = SocketAncillary::new(128);
crate::assert_with_log!(
ancillary.capacity() == 128,
"capacity",
128,
ancillary.capacity()
);
crate::assert_with_log!(ancillary.is_empty(), "is_empty", true, ancillary.is_empty());
crate::assert_with_log!(
!ancillary.is_truncated(),
"truncated",
false,
ancillary.is_truncated()
);
crate::test_complete!("test_ancillary_new");
}
#[test]
fn test_add_fds() {
init_test("test_add_fds");
let mut ancillary = SocketAncillary::new(128);
let fds = [3, 4, 5];
let added = ancillary.add_fds(&fds);
crate::assert_with_log!(added, "added", true, added);
crate::assert_with_log!(
!ancillary.is_empty(),
"not empty",
false,
ancillary.is_empty()
);
crate::test_complete!("test_add_fds");
}
#[test]
fn test_add_fds_too_small_buffer() {
init_test("test_add_fds_too_small");
let mut ancillary = SocketAncillary::new(4);
let fds = [3, 4, 5];
let added = ancillary.add_fds(&fds);
crate::assert_with_log!(added, "added", true, added);
crate::assert_with_log!(
!ancillary.is_empty(),
"not empty",
false,
ancillary.is_empty()
);
crate::test_complete!("test_add_fds_too_small");
}
#[test]
fn test_clear() {
init_test("test_ancillary_clear");
let mut ancillary = SocketAncillary::new(128);
ancillary.add_fds(&[3, 4]);
crate::assert_with_log!(
!ancillary.is_empty(),
"not empty",
false,
ancillary.is_empty()
);
ancillary.clear();
crate::assert_with_log!(
ancillary.is_empty(),
"empty after clear",
true,
ancillary.is_empty()
);
crate::test_complete!("test_ancillary_clear");
}
#[test]
fn test_ancillary_space_for_fds() {
init_test("test_ancillary_space_for_fds");
let space0 = ancillary_space_for_fds(0);
crate::assert_with_log!(space0 == 0, "space for 0", 0, space0);
let space1 = ancillary_space_for_fds(1);
crate::assert_with_log!(space1 > 0, "space for 1 > 0", true, space1 > 0);
let space3 = ancillary_space_for_fds(3);
crate::assert_with_log!(space3 > space1, "space for 3 > 1", true, space3 > space1);
crate::test_complete!("test_ancillary_space_for_fds");
}
#[test]
fn test_prepare_for_recv_exposes_full_buffer_len() {
init_test("test_prepare_for_recv_exposes_full_buffer_len");
let mut ancillary = SocketAncillary::new(128);
let recv_buf = ancillary.prepare_for_recv();
crate::assert_with_log!(recv_buf.len() == 128, "recv buf len", 128, recv_buf.len());
crate::test_complete!("test_prepare_for_recv_exposes_full_buffer_len");
}
#[test]
fn socket_ancillary_debug_default() {
let anc = SocketAncillary::default();
let dbg = format!("{anc:?}");
assert!(dbg.contains("SocketAncillary"), "{dbg}");
assert_eq!(anc.capacity(), 0);
assert!(anc.is_empty());
}
#[test]
fn ancillary_messages_debug() {
let anc = SocketAncillary::new(0);
let msgs = anc.messages();
let dbg = format!("{msgs:?}");
assert!(dbg.contains("AncillaryMessages"), "{dbg}");
}
}