use libc::c_uint;use serde::Deserialize;use serde::Serialize;use socket2::Domain;use socket2::MaybeUninitSlice;use socket2::MsgHdr;use socket2::MsgHdrMut;use socket2::Socket;use socket2::Type;use std::io::IoSlice;use std::mem::MaybeUninit;use std::os::fd::AsRawFd;use std::os::fd::FromRawFd;use std::os::fd::OwnedFd;use std::os::fd::RawFd;use tokio::io::Interest;use tokio::io::unix::AsyncFd;const MAX_FDS_PER_MESSAGE: usize = 16;const LENGTH_PREFIX_SIZE: usize = size_of::<u32>();const MAX_DATAGRAM_SIZE: usize = 8192;/// Converts a slice of MaybeUninit<T> to a slice of T.////// The caller guarantees that every element of `buf` is initialized.fn assume_init<T>(buf: &[MaybeUninit<T>]) -> &[T] {    unsafe { std::slice::from_raw_parts(buf.as_ptr().cast(), buf.len()) }}fn assume_init_slice<T, const N: usize>(buf: &[MaybeUninit<T>; N]) -> &[T; N] {    unsafe { &*(buf as *const [MaybeUninit<T>; N] as *const [T; N]) }}fn assume_init_vec<T>(mut buf: Vec<MaybeUninit<T>>) -> Vec<T> {    unsafe {        let ptr = buf.as_mut_ptr() as *mut T;        let len = buf.len();        let cap = buf.capacity();        std::mem::forget(buf);        Vec::from_raw_parts(ptr, len, cap)    }}fn control_space_for_fds(count: usize) -> usize {    unsafe { libc::CMSG_SPACE((count * size_of::<RawFd>()) as _) as usize }}/// Extracts the FDs from a SCM_RIGHTS control message.fn extract_fds(control: &[u8]) -> Vec<OwnedFd> {    let mut fds = Vec::new();    let mut hdr: libc::msghdr = unsafe { std::mem::zeroed() };    hdr.msg_control = control.as_ptr() as *mut libc::c_void;    hdr.msg_controllen = control.len() as _;    let hdr = hdr; // drop mut    let mut cmsg = unsafe { libc::CMSG_FIRSTHDR(&hdr) as *const libc::cmsghdr };    while !cmsg.is_null() {        let level = unsafe { (*cmsg).cmsg_level };        let ty = unsafe { (*cmsg).cmsg_type };        if level == libc::SOL_SOCKET && ty == libc::SCM_RIGHTS {            let data_ptr = unsafe { libc::CMSG_DATA(cmsg).cast::<RawFd>() };            let fd_count: usize = {                // `cmsghdr::cmsg_len` is not typed consistently across targets, so normalize it                // before doing the size arithmetic.                #[allow(clippy::useless_conversion, clippy::expect_used)]                let cmsg_data_len = usize::try_from(unsafe { (*cmsg).cmsg_len })                    .expect("cmsghdr length fits")                    - unsafe { libc::CMSG_LEN(0) as usize };                cmsg_data_len / size_of::<RawFd>()            };            for i in 0..fd_count {                let fd = unsafe { data_ptr.add(i).read() };                fds.push(unsafe { OwnedFd::from_raw_fd(fd) });            }        }        cmsg = unsafe { libc::CMSG_NXTHDR(&hdr, cmsg) };    }    fds}/// Read a frame from a SOCK_STREAM socket.////// A frame is a message length prefix followed by a payload. FDs may be included in the control/// message when receiving the frame header.async fn read_frame(async_socket: &AsyncFd<Socket>) -> std::io::Result<(Vec<u8>, Vec<OwnedFd>)> {    let (message_len, fds) = read_frame_header(async_socket).await?;    let payload = read_frame_payload(async_socket, message_len).await?;    Ok((payload, fds))}/// Read the frame header (i.e. length) and any FDs from a SOCK_STREAM socket.async fn read_frame_header(    async_socket: &AsyncFd<Socket>,) -> std::io::Result<(usize, Vec<OwnedFd>)> {    let mut header = [MaybeUninit::<u8>::uninit(); LENGTH_PREFIX_SIZE];    let mut filled = 0;    let mut control = vec![MaybeUninit::<u8>::uninit(); control_space_for_fds(MAX_FDS_PER_MESSAGE)];    let mut captured_control = false;    while filled < LENGTH_PREFIX_SIZE {        let mut guard = async_socket.readable().await?;        // The first read should come with a control message containing any FDs.        let read = if !captured_control {            match guard.try_io(|inner| {                let mut bufs = [MaybeUninitSlice::new(&mut header[filled..])];                let (read, control_len) = {                    let mut msg = MsgHdrMut::new()                        .with_buffers(&mut bufs)                        .with_control(&mut control);                    let read = inner.get_ref().recvmsg(&mut msg, 0)?;                    (read, msg.control_len())                };                control.truncate(control_len);                captured_control = true;                Ok(read)            }) {                Ok(Ok(read)) => read,                Ok(Err(err)) => return Err(err),                Err(_would_block) => continue,            }        } else {            match guard.try_io(|inner| inner.get_ref().recv(&mut header[filled..])) {                Ok(Ok(read)) => read,                Ok(Err(err)) => return Err(err),                Err(_would_block) => continue,            }        };        if read == 0 {            return Err(std::io::Error::new(                std::io::ErrorKind::UnexpectedEof,                "socket closed while receiving frame header",            ));        }        filled += read;        assert!(filled <= LENGTH_PREFIX_SIZE);        if filled == LENGTH_PREFIX_SIZE {            let len_bytes = assume_init_slice(&header);            let payload_len = u32::from_le_bytes(*len_bytes) as usize;            let fds = extract_fds(assume_init(&control));            return Ok((payload_len, fds));        }    }    unreachable!("header loop always returns")}/// Read `message_len` bytes from a SOCK_STREAM socket.async fn read_frame_payload(    async_socket: &AsyncFd<Socket>,    message_len: usize,) -> std::io::Result<Vec<u8>> {    if message_len == 0 {        return Ok(Vec::new());    }    let mut payload = vec![MaybeUninit::<u8>::uninit(); message_len];    let mut filled = 0;    while filled < message_len {        let mut guard = async_socket.readable().await?;        let read = match guard.try_io(|inner| inner.get_ref().recv(&mut payload[filled..])) {            Ok(Ok(read)) => read,            Ok(Err(err)) => return Err(err),            Err(_would_block) => continue,        };        if read == 0 {            return Err(std::io::Error::new(                std::io::ErrorKind::UnexpectedEof,                "socket closed while receiving frame payload",            ));        }        filled += read;        assert!(filled <= message_len);        if filled == message_len {            return Ok(assume_init_vec(payload));        }    }    unreachable!("loop exits only after returning payload")}fn send_datagram_bytes(socket: &Socket, data: &[u8], fds: &[OwnedFd]) -> std::io::Result<()> {    let control = make_control_message(fds)?;    let payload = [IoSlice::new(data)];    let msg = if control.is_empty() {        MsgHdr::new().with_buffers(&payload)    } else {        MsgHdr::new().with_buffers(&payload).with_control(&control)    };    let written = socket.sendmsg(&msg, 0)?;    if written != data.len() {        return Err(std::io::Error::new(            std::io::ErrorKind::WriteZero,            format!(                "short datagram write: wrote {written} bytes out of {}",                data.len()            ),        ));    }    Ok(())}fn encode_length(len: usize) -> std::io::Result<[u8; LENGTH_PREFIX_SIZE]> {    let len_u32 = u32::try_from(len).map_err(|_| {        std::io::Error::new(            std::io::ErrorKind::InvalidInput,            format!("message too large: {len}"),        )    })?;    Ok(len_u32.to_le_bytes())}fn make_control_message(fds: &[OwnedFd]) -> std::io::Result<Vec<u8>> {    if fds.len() > MAX_FDS_PER_MESSAGE {        Err(std::io::Error::new(            std::io::ErrorKind::InvalidInput,            format!("too many fds: {}", fds.len()),        ))    } else if fds.is_empty() {        Ok(Vec::new())    } else {        let mut control = vec![0u8; control_space_for_fds(fds.len())];        unsafe {            let cmsg = control.as_mut_ptr().cast::<libc::cmsghdr>();            (*cmsg).cmsg_len =                libc::CMSG_LEN(size_of::<RawFd>() as c_uint * fds.len() as c_uint) as _;            (*cmsg).cmsg_level = libc::SOL_SOCKET;            (*cmsg).cmsg_type = libc::SCM_RIGHTS;            let data_ptr = libc::CMSG_DATA(cmsg).cast::<RawFd>();            for (i, fd) in fds.iter().enumerate() {                data_ptr.add(i).write(fd.as_raw_fd());            }        }        Ok(control)    }}fn receive_datagram_bytes(socket: &Socket) -> std::io::Result<(Vec<u8>, Vec<OwnedFd>)> {    let mut buffer = vec![MaybeUninit::<u8>::uninit(); MAX_DATAGRAM_SIZE];    let mut control = vec![MaybeUninit::<u8>::uninit(); control_space_for_fds(MAX_FDS_PER_MESSAGE)];    let (read, control_len) = {        let mut bufs = [MaybeUninitSlice::new(&mut buffer)];        let mut msg = MsgHdrMut::new()            .with_buffers(&mut bufs)            .with_control(&mut control);        let read = socket.recvmsg(&mut msg, 0)?;        (read, msg.control_len())    };    let data = assume_init(&buffer[..read]).to_vec();    let fds = extract_fds(assume_init(&control[..control_len]));    Ok((data, fds))}pub(crate) struct AsyncSocket {    inner: AsyncFd<Socket>,}impl AsyncSocket {    fn new(socket: Socket) -> std::io::Result<AsyncSocket> {        socket.set_nonblocking(true)?;        let async_socket = AsyncFd::new(socket)?;        Ok(AsyncSocket {            inner: async_socket,        })    }    pub fn from_fd(fd: OwnedFd) -> std::io::Result<AsyncSocket> {        AsyncSocket::new(Socket::from(fd))    }    pub fn pair() -> std::io::Result<(AsyncSocket, AsyncSocket)> {        // `socket2::Socket::pair()` also applies "common flags" (including        // `SO_NOSIGPIPE` on Apple platforms), which can fail for AF_UNIX sockets.        // Use `pair_raw()` to avoid those side effects, then restore `CLOEXEC`        // explicitly on both endpoints.        let (server, client) = Socket::pair_raw(Domain::UNIX, Type::STREAM, None)?;        server.set_cloexec(true)?;        client.set_cloexec(true)?;        Ok((AsyncSocket::new(server)?, AsyncSocket::new(client)?))    }    pub async fn send_with_fds<T: Serialize>(        &self,        msg: T,        fds: &[OwnedFd],    ) -> std::io::Result<()> {        let payload = serde_json::to_vec(&msg)?;        let mut frame = Vec::with_capacity(LENGTH_PREFIX_SIZE + payload.len());        frame.extend_from_slice(&encode_length(payload.len())?);        frame.extend_from_slice(&payload);        send_stream_frame(&self.inner, &frame, fds).await    }    pub async fn receive_with_fds<T: for<'de> Deserialize<'de>>(        &self,    ) -> std::io::Result<(T, Vec<OwnedFd>)> {        let (payload, fds) = read_frame(&self.inner).await?;        let message: T = serde_json::from_slice(&payload)?;        Ok((message, fds))    }    pub async fn send<T>(&self, msg: T) -> std::io::Result<()>    where        T: Serialize,    {        self.send_with_fds(&msg, &[]).await    }    pub async fn receive<T: for<'de> Deserialize<'de>>(&self) -> std::io::Result<T> {        let (msg, fds) = self.receive_with_fds().await?;        if !fds.is_empty() {            tracing::warn!("unexpected fds in receive: {}", fds.len());        }        Ok(msg)    }    pub fn into_inner(self) -> Socket {        self.inner.into_inner()    }}async fn send_stream_frame(    socket: &AsyncFd<Socket>,    frame: &[u8],    fds: &[OwnedFd],) -> std::io::Result<()> {    let mut written = 0;    let mut include_fds = !fds.is_empty();    while written < frame.len() {        let mut guard = socket.writable().await?;        let bytes_written = match guard            .try_io(|inner| send_stream_chunk(inner.get_ref(), &frame[written..], fds, include_fds))        {            Ok(result) => result?,            Err(_would_block) => continue,        };        if bytes_written == 0 {            return Err(std::io::Error::new(                std::io::ErrorKind::WriteZero,                "socket closed while sending frame payload",            ));        }        written += bytes_written;        include_fds = false;    }    Ok(())}fn send_stream_chunk(    socket: &Socket,    frame: &[u8],    fds: &[OwnedFd],    include_fds: bool,) -> std::io::Result<usize> {    let control = if include_fds {        make_control_message(fds)?    } else {        Vec::new()    };    let payload = [IoSlice::new(frame)];    let msg = if control.is_empty() {        MsgHdr::new().with_buffers(&payload)    } else {        MsgHdr::new().with_buffers(&payload).with_control(&control)    };    socket.sendmsg(&msg, 0)}pub(crate) struct AsyncDatagramSocket {    inner: AsyncFd<Socket>,}impl AsyncDatagramSocket {    fn new(socket: Socket) -> std::io::Result<Self> {        socket.set_nonblocking(true)?;        Ok(Self {            inner: AsyncFd::new(socket)?,        })    }    pub unsafe fn from_raw_fd(fd: RawFd) -> std::io::Result<Self> {        Self::new(unsafe { Socket::from_raw_fd(fd) })    }    pub fn pair() -> std::io::Result<(Self, Self)> {        // `socket2::Socket::pair()` also applies "common flags" (including        // `SO_NOSIGPIPE` on Apple platforms), which can fail for AF_UNIX sockets.        // Use `pair_raw()` to avoid those side effects, then restore `CLOEXEC`        // explicitly on both endpoints.        let (server, client) = Socket::pair_raw(Domain::UNIX, Type::DGRAM, None)?;        server.set_cloexec(true)?;        client.set_cloexec(true)?;        Ok((Self::new(server)?, Self::new(client)?))    }    pub async fn send_with_fds(&self, data: &[u8], fds: &[OwnedFd]) -> std::io::Result<()> {        self.inner            .async_io(Interest::WRITABLE, |socket| {                send_datagram_bytes(socket, data, fds)            })            .await    }    pub async fn receive_with_fds(&self) -> std::io::Result<(Vec<u8>, Vec<OwnedFd>)> {        self.inner            .async_io(Interest::READABLE, receive_datagram_bytes)            .await    }    pub fn into_inner(self) -> Socket {        self.inner.into_inner()    }}#[cfg(test)]mod tests {    use super::*;    use pretty_assertions::assert_eq;    use serde::Deserialize;    use serde::Serialize;    use std::os::fd::AsFd;    use std::os::fd::AsRawFd;    use tempfile::NamedTempFile;    #[derive(Debug, Serialize, Deserialize, PartialEq, Eq, Clone)]    struct TestPayload {        id: i32,        label: String,    }    fn fd_list(count: usize) -> std::io::Result<Vec<OwnedFd>> {        let file = NamedTempFile::new()?;        let mut fds = Vec::new();        for _ in 0..count {            fds.push(file.as_fd().try_clone_to_owned()?);        }        Ok(fds)    }    #[tokio::test]    async fn async_socket_round_trips_payload_and_fds() -> std::io::Result<()> {        let (server, client) = AsyncSocket::pair()?;        let payload = TestPayload {            id: 7,            label: "round-trip".to_string(),        };        let send_fds = fd_list(/*count*/ 1)?;        let receive_task =            tokio::spawn(async move { server.receive_with_fds::<TestPayload>().await });        client.send_with_fds(payload.clone(), &send_fds).await?;        drop(send_fds);        let (received_payload, received_fds) = receive_task.await.unwrap()?;        assert_eq!(payload, received_payload);        assert_eq!(1, received_fds.len());        let fd_status = unsafe { libc::fcntl(received_fds[0].as_raw_fd(), libc::F_GETFD) };        assert!(            fd_status >= 0,            "expected received file descriptor to be valid, but got {fd_status}",        );        Ok(())    }    #[tokio::test]    async fn async_socket_handles_large_payload() -> std::io::Result<()> {        let (server, client) = AsyncSocket::pair()?;        let payload = vec![b'A'; 10_000];        let receive_task = tokio::spawn(async move { server.receive::<Vec<u8>>().await });        client.send(payload.clone()).await?;        let received_payload = receive_task.await.unwrap()?;        assert_eq!(payload, received_payload);        Ok(())    }    #[tokio::test]    async fn async_datagram_sockets_round_trip_messages() -> std::io::Result<()> {        let (server, client) = AsyncDatagramSocket::pair()?;        let data = b"datagram payload".to_vec();        let send_fds = fd_list(/*count*/ 1)?;        let receive_task = tokio::spawn(async move { server.receive_with_fds().await });        client.send_with_fds(&data, &send_fds).await?;        drop(send_fds);        let (received_bytes, received_fds) = receive_task.await.unwrap()?;        assert_eq!(data, received_bytes);        assert_eq!(1, received_fds.len());        Ok(())    }    #[test]    fn send_datagram_bytes_rejects_excessive_fd_counts() -> std::io::Result<()> {        let (socket, _peer) = Socket::pair_raw(Domain::UNIX, Type::DGRAM, None)?;        let fds = fd_list(MAX_FDS_PER_MESSAGE + 1)?;        let err = send_datagram_bytes(&socket, b"hi", &fds).unwrap_err();        assert_eq!(std::io::ErrorKind::InvalidInput, err.kind());        Ok(())    }    #[test]    fn send_stream_chunk_rejects_excessive_fd_counts() -> std::io::Result<()> {        let (socket, _peer) = Socket::pair_raw(Domain::UNIX, Type::STREAM, None)?;        let fds = fd_list(MAX_FDS_PER_MESSAGE + 1)?;        let err = send_stream_chunk(&socket, b"hello", &fds, /*include_fds*/ true).unwrap_err();        assert_eq!(std::io::ErrorKind::InvalidInput, err.kind());        Ok(())    }    #[test]    fn encode_length_errors_for_oversized_messages() {        let err = encode_length(usize::MAX).unwrap_err();        assert_eq!(std::io::ErrorKind::InvalidInput, err.kind());    }    #[tokio::test]    async fn receive_fails_when_peer_closes_before_header() {        let (server, client) = AsyncSocket::pair().expect("failed to create socket pair");        drop(client);        let err = server            .receive::<serde_json::Value>()            .await            .expect_err("expected read failure");        assert_eq!(std::io::ErrorKind::UnexpectedEof, err.kind());    }}