ruc 11.0.0

Rust Util Collections
Documentation
//!
//! # UAU
//!
//! Unix(Unix Domain Socket) Abstract Udp
//!
//! An abstract socket address is distinguished (from a pathname socket)
//! by the fact that sun_path[0] is a null byte ('\0').
//! The socket's address in this namespace is given
//! by the additional bytes in sun_path that are covered
//! by the specified length of the address structure.
//! Null bytes in the name have no special  significance.
//! The name has no connection with filesystem pathnames.
//! When the address of an abstract socket is returned,
//! the returned addrlen is greater than sizeof(sa_family_t) (i.e., greater than 2),
//! and the name of the socket is contained in the first (addrlen - sizeof(sa_family_t)) bytes of sun_path.
//!
//! `man unix(7)` for more information.
//!

use crate::*;
use nix::sys::{
    socket::{
        AddressFamily, MsgFlags, SockFlag, SockType, UnixAddr, bind, recvfrom,
        sendto, setsockopt, socket, sockopt,
    },
    time::{TimeVal, TimeValLike},
};
use std::os::{fd::OwnedFd, unix::io::AsRawFd};

/// Wrap raw data
pub struct UauSock {
    // `OwnedFd` closes the fd automatically on drop
    fd: OwnedFd,
    sa: UnixAddr,
}

impl UauSock {
    /// NOTE:
    ///
    /// The Unix Socket that needs to receive messages must be bound to an explicit address;
    /// If you send msg anonymously, you will not receive the reply msg from the peer.
    ///
    /// So the `addr` parameter should not be empty.
    pub fn new(addr: &[u8], recv_timeout: Option<i64>) -> Result<Self> {
        let fd = socket(
            AddressFamily::Unix,
            SockType::Datagram,
            SockFlag::empty(),
            None,
        )
        .c(d!())?;

        setsockopt(&fd, sockopt::ReuseAddr, &true).c(d!())?;
        if let Some(to) = recv_timeout {
            setsockopt(
                &fd,
                sockopt::ReceiveTimeout,
                &TimeVal::milliseconds(to),
            )
            .c(d!())?;
        }

        let sa = UnixAddr::new_abstract(addr).c(d!())?;
        bind(fd.as_raw_fd(), &sa).c(d!())?;

        Ok(UauSock { fd, sa })
    }

    /// Generate a random instance
    #[inline(always)]
    pub fn create(recv_timeout: Option<i64>) -> Result<Self> {
        let addr = (ts!() as u32 ^ rand::random::<u32>()).to_ne_bytes();
        Self::new(&addr, recv_timeout).c(d!())
    }

    /// Get the addr of UauSock
    #[inline(always)]
    pub fn addr(&self) -> &UnixAddr {
        &self.sa
    }

    /// Send msg to another peer
    #[inline(always)]
    pub fn send(&self, msg: &[u8], peeraddr: &UnixAddr) -> Result<()> {
        sendto(self.fd.as_raw_fd(), msg, peeraddr, MsgFlags::empty())
            .c(d!())
            .map(|_| ())
    }

    /// Receive msg with a const-generic sized buffer, returning data and peer address.
    ///
    /// Returns an error if the incoming datagram is larger than `N` bytes
    /// (oversized datagrams are never silently truncated) — choose a
    /// larger `N` in that case.
    #[inline(always)]
    pub fn recv_buf<const N: usize>(&self) -> Result<(Vec<u8>, UnixAddr)> {
        // one extra probe byte detects datagrams larger than `N`
        let mut buf = vec![0u8; N + 1];
        let (n, peer) = self.recv(&mut buf).c(d!())?;
        ensure!(
            n <= N,
            "datagram larger than the {N}-byte buffer, choose a larger `N`"
        );
        buf.truncate(n);
        Ok((buf, peer))
    }

    /// Receive msg with a const-generic sized buffer, discarding peer address.
    ///
    /// Returns an error if the incoming datagram is larger than `N` bytes,
    /// see [`recv_buf`](Self::recv_buf).
    #[inline(always)]
    pub fn recvonly_buf<const N: usize>(&self) -> Result<Vec<u8>> {
        self.recv_buf::<N>().map(|(b, _)| b)
    }

    /// Receive msg with a given buffer.
    ///
    /// NOTE: a datagram larger than the buffer is silently truncated
    /// to the buffer size; use [`recv_buf`](Self::recv_buf) for
    /// truncation detection.
    #[inline(always)]
    pub fn recv(&self, buf: &mut [u8]) -> Result<(usize, UnixAddr)> {
        match recvfrom::<UnixAddr>(self.fd.as_raw_fd(), buf) {
            Ok((n, Some(peer))) => Ok((n, peer)),
            Err(e) => Err(eg!(e)),
            _ => Err(eg!("peer address is unknown")),
        }
    }

    /// Try to convert a user-given addr to SockAddr(unix sock)
    pub fn addr_to_sock(addr: &[u8]) -> Result<UnixAddr> {
        UnixAddr::new_abstract(addr).c(d!())
    }
}

#[cfg(test)]
mod test {
    use super::*;

    #[test]
    fn t_send_recv() {
        let sender = pnk!(UauSock::create(None));
        let receiver = pnk!(UauSock::create(None));

        for _ in 0..4 {
            pnk!(
                sender.send(&987654321_u32.to_ne_bytes()[..], receiver.addr())
            );
            assert_eq!(
                &987654321_u32.to_ne_bytes()[..],
                &pnk!(receiver.recvonly_buf::<64>())
            );
        }
    }

    #[test]
    fn t_recv_buf_rejects_oversized_datagram() {
        let sender = pnk!(UauSock::create(None));
        let receiver = pnk!(UauSock::create(None));

        pnk!(sender.send(&[0u8; 8][..], receiver.addr()));
        assert!(receiver.recvonly_buf::<4>().is_err());

        // exactly-fitting datagrams still succeed
        pnk!(sender.send(&[7u8; 4][..], receiver.addr()));
        assert_eq!(&[7u8; 4][..], &pnk!(receiver.recvonly_buf::<4>()));
    }

    #[test]
    fn t_send_recv_const_generic() {
        let sender = pnk!(UauSock::create(None));
        let receiver = pnk!(UauSock::create(None));

        pnk!(sender.send(&987654321_u32.to_ne_bytes()[..], receiver.addr()));
        assert_eq!(
            &987654321_u32.to_ne_bytes()[..],
            &pnk!(receiver.recvonly_buf::<64>())
        );

        pnk!(sender.send(&987654321_u32.to_ne_bytes()[..], receiver.addr()));
        assert_eq!(
            &987654321_u32.to_ne_bytes()[..],
            &pnk!(receiver.recvonly_buf::<256>())
        );
    }
}