torrust-actix 4.2.15

A rich, fast and efficient Bittorrent Tracker.
use crate::udp::udp::MAX_PACKET_SIZE;
use std::net::{
    IpAddr,
    Ipv4Addr,
    Ipv6Addr,
    SocketAddr
};
use std::os::unix::io::RawFd;

pub const BATCH: usize = 64;

pub struct RecvBatch {
    bufs: Box<[[u8; MAX_PACKET_SIZE]; BATCH]>,
    addrs: Box<[libc::sockaddr_storage; BATCH]>,
    iovecs: Box<[libc::iovec; BATCH]>,
    msgs: Box<[libc::mmsghdr; BATCH]>,
}

unsafe impl Send for RecvBatch {}

impl Default for RecvBatch {
    fn default() -> Self {
        Self::new()
    }
}

impl RecvBatch {
    /// Allocates the reusable buffers for one `recvmmsg` batch.
    pub fn new() -> Self {
        RecvBatch {
            bufs: Box::new([[0u8; MAX_PACKET_SIZE]; BATCH]),
            addrs: Box::new(unsafe { std::mem::zeroed() }),
            iovecs: Box::new(unsafe { std::mem::zeroed() }),
            msgs: Box::new(unsafe { std::mem::zeroed() }),
        }
    }

    /// Receives up to a full batch of datagrams in one `recvmmsg` call.
    ///
    /// # Errors
    ///
    /// Returns the socket error reported by the kernel.
    pub fn recv(&mut self, fd: RawFd) -> std::io::Result<usize> {
        for i in 0..BATCH {
            let buf_ptr = self.bufs[i].as_mut_ptr() as *mut libc::c_void;
            let addr_ptr = (&mut self.addrs[i] as *mut libc::sockaddr_storage) as *mut libc::c_void;
            self.iovecs[i].iov_base = buf_ptr;
            self.iovecs[i].iov_len = MAX_PACKET_SIZE;
            let iov_ptr = &mut self.iovecs[i] as *mut libc::iovec;
            let hdr = &mut self.msgs[i].msg_hdr;
            hdr.msg_name = addr_ptr;
            hdr.msg_namelen = std::mem::size_of::<libc::sockaddr_storage>() as libc::socklen_t;
            hdr.msg_iov = iov_ptr;
            hdr.msg_iovlen = 1;
            hdr.msg_control = std::ptr::null_mut();
            hdr.msg_controllen = 0;
            hdr.msg_flags = 0;
            self.msgs[i].msg_len = 0;
        }
        let ret = unsafe {
            libc::recvmmsg(
                fd,
                self.msgs.as_mut_ptr(),
                BATCH as libc::c_uint,
                libc::MSG_DONTWAIT as _,
                std::ptr::null_mut(),
            )
        };
        if ret < 0 {
            return Err(std::io::Error::last_os_error());
        }
        Ok(ret as usize)
    }

    /// Returns the payload slice and source address of the `index`-th datagram of the last batch.
    pub fn datagram(&self, index: usize) -> Option<(&[u8], SocketAddr)> {
        let len = (self.msgs[index].msg_len as usize).min(MAX_PACKET_SIZE);
        let buf = &self.bufs[index][..len];
        let addr = sockaddr_to_socketaddr(&self.addrs[index])?;
        Some((buf, addr))
    }
}

/// Converts a raw `sockaddr_storage` filled by the kernel into a Rust [`SocketAddr`].
pub(crate) fn sockaddr_to_socketaddr(storage: &libc::sockaddr_storage) -> Option<SocketAddr> {
    match storage.ss_family as libc::c_int {
        libc::AF_INET => {
            let addr = unsafe { &*(storage as *const libc::sockaddr_storage as *const libc::sockaddr_in) };
            let ip = Ipv4Addr::from(u32::from_be(addr.sin_addr.s_addr));
            let port = u16::from_be(addr.sin_port);
            Some(SocketAddr::new(IpAddr::V4(ip), port))
        }
        libc::AF_INET6 => {
            let addr = unsafe { &*(storage as *const libc::sockaddr_storage as *const libc::sockaddr_in6) };
            let ip = Ipv6Addr::from(addr.sin6_addr.s6_addr);
            let port = u16::from_be(addr.sin6_port);
            Some(SocketAddr::new(IpAddr::V6(ip), port))
        }
        _ => None,
    }
}