// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
// SPDX-License-Identifier: Apache-2.0

use crate::message::{cmsg, cmsg::Encoder, Message as MessageTrait};
use core::{
    alloc::Layout,
    mem::{size_of, size_of_val},
};
use libc::{iovec, msghdr, sockaddr_in, sockaddr_in6, AF_INET, AF_INET6};
use s2n_quic_core::{
    inet::{
        datagram, ExplicitCongestionNotification, IpV4Address, IpV6Address, SocketAddress,
        SocketAddressV4, SocketAddressV6,
    },
    io::tx,
    path::{self, Handle as _},
};

mod ext;
mod handle;
#[cfg(test)]
mod tests;

use ext::Ext as _;

pub use handle::Handle;
pub use libc::msghdr as Message;

#[cfg(any(test, feature = "generator"))]
use bolero_generator::*;

impl MessageTrait for msghdr {
    type Handle = Handle;

    const SUPPORTS_GSO: bool = cfg!(s2n_quic_platform_gso);
    const SUPPORTS_ECN: bool = cfg!(s2n_quic_platform_tos);
    const SUPPORTS_FLOW_LABELS: bool = true;

    #[inline]
    fn alloc(entries: u32, payload_len: u32, offset: usize) -> super::Storage {
        unsafe { alloc(entries, payload_len, offset, |msg| msg) }
    }

    #[inline]
    fn payload_len(&self) -> usize {
        debug_assert!(!self.msg_iov.is_null());
        let len = unsafe { (*self.msg_iov).iov_len as _ };
        debug_assert!(len <= u16::MAX as usize);
        len
    }

    #[inline]
    unsafe fn set_payload_len(&mut self, payload_len: usize) {
        debug_assert!(payload_len <= u16::MAX as usize);
        debug_assert!(!self.msg_iov.is_null());
        (*self.msg_iov).iov_len = payload_len;
    }

    #[cfg(s2n_quic_platform_gso)]
    #[inline]
    fn set_segment_size(&mut self, size: usize) {
        self.encode_cmsg(libc::SOL_UDP, libc::UDP_SEGMENT, size as cmsg::UdpGso)
            .unwrap();
    }

    #[inline]
    unsafe fn reset(&mut self, mtu: usize) {
        // reset the payload
        self.set_payload_len(mtu);

        // reset the address
        self.set_remote_address(&SocketAddress::IpV6(Default::default()));

        #[inline]
        unsafe fn check_cmsg(msghdr: &msghdr) {
            if cfg!(debug_assertions) {
                let ptr = msghdr.msg_control as *mut u8;
                let cmsg = core::slice::from_raw_parts_mut(ptr, cmsg::MAX_LEN);
                // make sure nothing was written to the control message if it was set to 0
                #[cfg(not(kani))]
                {
                    assert!(cmsg.iter().all(|v| *v == 0), "msg_control was not cleared");
                }

                #[cfg(kani)]
                {
                    let index: usize = kani::any();
                    kani::assume(index < cmsg.len());
                    assert_eq!(cmsg[index], 0);
                }
            }
        }

        // make sure we didn't get any data written without setting the len
        if self.msg_controllen == 0 {
            check_cmsg(self);
        }

        // reset the control messages if it isn't set to the default value

        // some platforms encode lengths as `u32` so we cast everything to be safe
        #[allow(clippy::unnecessary_cast)]
        let msg_controllen = self.msg_controllen as usize;

        if msg_controllen != cmsg::MAX_LEN {
            core::slice::from_raw_parts_mut(self.msg_control as *mut u8, msg_controllen).fill(0);
        }

        check_cmsg(self);

        self.msg_controllen = cmsg::MAX_LEN as _;
    }

    #[inline]
    fn payload_ptr_mut(&mut self) -> *mut u8 {
        unsafe {
            let iovec = &mut *self.msg_iov;
            iovec.iov_base as *mut _
        }
    }

    #[inline]
    fn validate_replication(source: &Self, dest: &Self) {
        assert_eq!(source.msg_name, dest.msg_name);
        assert_eq!(source.msg_iov, dest.msg_iov);
        assert_eq!(source.msg_control, dest.msg_control);
    }

    #[inline]
    fn rx_read(
        &mut self,
        local_address: &path::LocalAddress,
    ) -> Option<super::RxMessage<Self::Handle>> {
        if cfg!(test) {
            assert_eq!(
                self.msg_flags & libc::MSG_CTRUNC,
                0,
                "control message buffers should always have enough capacity"
            );
        }

        let (mut header, cmsg) = self.header()?;

        // only copy the port if we are told the IP address
        if cfg!(s2n_quic_platform_pktinfo) {
            header.path.local_address.set_port(local_address.port());
        } else {
            header.path.local_address = *local_address;
        }

        let payload = self.payload_mut();

        let segment_size = if cmsg.segment_size == 0 {
            payload.len()
        } else {
            cmsg.segment_size as _
        };

        let message = crate::message::RxMessage {
            header,
            segment_size,
            payload,
        };

        Some(message)
    }

    #[inline]
    fn tx_write<M: tx::Message<Handle = Self::Handle>>(
        &mut self,
        mut message: M,
    ) -> Result<usize, tx::Error> {
        let payload = self.payload_mut();

        let max_len = payload.len();
        let len = message.write_payload(tx::PayloadBuffer::new(payload), 0)?;

        debug_assert_ne!(len, 0);
        debug_assert!(len <= max_len);
        let len = len.min(max_len);

        debug_assert_eq!(
            cmsg::MAX_LEN,
            self.msg_controllen as _,
            "message should be reset before writing"
        );
        self.msg_controllen = 0;

        unsafe {
            self.set_payload_len(len);
        }

        let handle = *message.path_handle();
        handle.update_msg_hdr(self);
        self.set_ecn(message.ecn(), &handle.remote_address.0);

        Ok(len)
    }
}

/// Allocates a region of memory holding `entries` number of `T` messages, each with `payload_len`
/// payloads.
///
/// # Safety
///
/// * `T` can be initialized with zero bytes and still be valid
#[inline]
pub(super) unsafe fn alloc<T: Copy + Sized, F: Fn(&mut T) -> &mut msghdr>(
    entries: u32,
    payload_len: u32,
    offset: usize,
    on_entry: F,
) -> super::Storage {
    // calculate the layout of the storage for the given configuration
    let (layout, entry_offset, header_offset, payload_offset) =
        layout::<T>(entries, payload_len, offset);

    // allocate a single contiguous block of memory
    let storage = super::Storage::new(layout);

    {
        let ptr = storage.as_ptr();

        // calculate each of the pointers we need to set up a message
        let mut entry_ptr = ptr.add(entry_offset) as *mut T;
        let mut header_ptr = ptr.add(header_offset) as *mut Header;
        let mut payload_ptr = ptr.add(payload_offset) as *mut u8;

        for _ in 0..entries {
            // for each message update all of the pointers to the correct locations

            let entry = on_entry(&mut *entry_ptr);
            (*header_ptr).update(entry, payload_ptr, payload_len);

            // increment the pointers for the next iteration
            entry_ptr = entry_ptr.add(1);
            header_ptr = header_ptr.add(1);
            payload_ptr = payload_ptr.add(payload_len as _);

            // make sure the pointers are within the bounds of the allocation
            storage.check_bounds(entry_ptr);
            storage.check_bounds(header_ptr);
            storage.check_bounds(payload_ptr);
        }

        // replicate the primary messages into the secondary region
        let primary = ptr.add(entry_offset) as *mut T;
        let secondary = primary.add(entries as _);
        storage.check_bounds(secondary.add(entries as _));
        core::ptr::copy_nonoverlapping(primary, secondary, entries as _);
    }

    storage
}

/// Computes the following layout
///
/// ```ignore
/// struct Storage {
///    cursor: Cursor,
///    headers: [Header; entries],
///    payloads: [[u8; payload_len]; entries],
///    entries: [T; entries * 2],
/// }
/// ```
fn layout<T: Copy + Sized>(
    entries: u32,
    payload_len: u32,
    offset: usize,
) -> (Layout, usize, usize, usize) {
    let cursor = Layout::array::<u8>(offset).unwrap();
    let headers = Layout::array::<Header>(entries as _).unwrap();
    let payloads = Layout::array::<u8>(entries as usize * payload_len as usize).unwrap();
    // double the number of entries we allocate to support the primary/secondary regions
    let entries = Layout::array::<T>((entries * 2) as usize).unwrap();
    let (layout, entry_offset) = cursor.extend(entries).unwrap();
    let (layout, header_offset) = layout.extend(headers).unwrap();
    let (layout, payload_offset) = layout.extend(payloads).unwrap();
    (layout, entry_offset, header_offset, payload_offset)
}

/// A structure for holding data pointed to in the [`libc::msghdr`] struct.
struct Header {
    pub iovec: iovec,
    pub msg_name: sockaddr_in6,
    pub cmsg: cmsg::Storage,
}

impl Header {
    /// sets all of the pointers of the provided `entry` to the correct locations
    unsafe fn update(&mut self, entry: &mut msghdr, payload: *mut u8, payload_len: u32) {
        let iovec = &mut self.iovec;

        iovec.iov_base = payload as *mut _;
        iovec.iov_len = payload_len as _;

        let entry = &mut *entry;

        entry.msg_name = &mut self.msg_name as *mut _ as *mut _;
        entry.msg_namelen = size_of_val(&self.msg_name) as _;
        entry.msg_iov = &mut self.iovec as *mut _;
        entry.msg_iovlen = 1;
        entry.msg_controllen = self.cmsg.len() as _;
        entry.msg_control = self.cmsg.as_mut_ptr() as *mut _;

        // make sure that the control pointer is well-aligned
        debug_assert_eq!(
            entry
                .msg_control
                .align_offset(core::mem::align_of::<cmsg::Storage>()),
            0
        );
    }
}


// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
// SPDX-License-Identifier: Apache-2.0

use crate::message::Message as MessageTrait;
use core::alloc::Layout;
use s2n_quic_core::{
    inet::{datagram, SocketAddress},
    io::tx,
    path,
};

/// A simple message type that holds an address and payload
///
/// All other fields are not supported by the platform.
#[derive(Clone, Copy, Debug)]
pub struct Message {
    address: SocketAddress,
    payload_ptr: *mut u8,
    payload_len: usize,
}

impl Message {
    #[inline]
    pub(crate) fn remote_address(&self) -> &SocketAddress {
        &self.address
    }

    #[inline]
    pub(crate) fn set_remote_address(&mut self, remote_address: &SocketAddress) {
        let remote_address = *remote_address;

        self.address = remote_address;
    }
}

pub type Handle = path::Tuple;

impl MessageTrait for Message {
    type Handle = Handle;

    const SUPPORTS_GSO: bool = false;
    const SUPPORTS_ECN: bool = false;
    const SUPPORTS_FLOW_LABELS: bool = false;

    #[inline]
    fn alloc(entries: u32, payload_len: u32, offset: usize) -> super::Storage {
        unsafe { alloc(entries, payload_len, offset) }
    }

    #[inline]
    fn payload_len(&self) -> usize {
        self.payload_len
    }

    #[inline]
    unsafe fn set_payload_len(&mut self, len: usize) {
        self.payload_len = len;
    }

    #[inline]
    unsafe fn reset(&mut self, mtu: usize) {
        self.set_payload_len(mtu)
    }

    #[inline]
    fn payload_ptr_mut(&mut self) -> *mut u8 {
        self.payload_ptr
    }

    #[inline]
    fn validate_replication(source: &Self, dest: &Self) {
        assert_eq!(source.payload_ptr, dest.payload_ptr);
    }

    #[inline]
    fn rx_read(
        &mut self,
        local_address: &path::LocalAddress,
    ) -> Option<super::RxMessage<Self::Handle>> {
        let path = path::Tuple {
            remote_address: self.address.into(),
            local_address: *local_address,
        };
        let header = datagram::Header {
            path,
            ecn: Default::default(),
        };
        let payload = self.payload_mut();

        let message = super::RxMessage {
            header,
            segment_size: payload.len(),
            payload,
        };

        Some(message)
    }

    #[inline]
    fn tx_write<M: tx::Message<Handle = Self::Handle>>(
        &mut self,
        mut message: M,
    ) -> Result<usize, tx::Error> {
        let payload = self.payload_mut();

        let len = message.write_payload(tx::PayloadBuffer::new(payload), 0)?;

        unsafe {
            debug_assert!(len <= payload.len());
            let len = len.min(payload.len());
            self.set_payload_len(len);
        }

        let remote_address = message.path_handle().remote_address;
        self.address = remote_address.0;

        Ok(len)
    }
}

#[inline]
unsafe fn alloc(entries: u32, payload_len: u32, offset: usize) -> super::Storage {
    let (layout, entry_offset, payload_offset) = layout(entries, payload_len, offset);

    let storage = super::Storage::new(layout);

    {
        let ptr = storage.as_ptr();

        let mut entry_ptr = ptr.add(entry_offset) as *mut Message;
        let mut payload_ptr = ptr.add(payload_offset) as *mut u8;
        for _ in 0..entries {
            let entry = &mut *entry_ptr;
            entry.payload_ptr = payload_ptr;
            entry.payload_len = payload_len as _;

            entry_ptr = entry_ptr.add(1);
            payload_ptr = payload_ptr.add(payload_len as _);
            storage.check_bounds(entry_ptr);
            storage.check_bounds(payload_ptr);
        }

        let primary = ptr.add(entry_offset) as *mut Message;
        let secondary = primary.add(entries as _);
        storage.check_bounds(secondary.add(entries as _));
        core::ptr::copy_nonoverlapping(primary, secondary, entries as _);
    }

    storage
}

fn layout(entries: u32, payload_len: u32, offset: usize) -> (Layout, usize, usize) {
    let cursor = Layout::array::<u8>(offset).unwrap();
    let payloads = Layout::array::<u8>(entries as usize * payload_len as usize).unwrap();
    let entries = Layout::array::<Message>((entries * 2) as usize).unwrap();
    let (layout, entry_offset) = cursor.extend(entries).unwrap();
    let (layout, payload_offset) = layout.extend(payloads).unwrap();
    (layout, entry_offset, payload_offset)
}


// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
// SPDX-License-Identifier: Apache-2.0

use super::*;

pub trait Ext: cmsg::Encoder {
    fn header(&self) -> Option<(datagram::Header<Handle>, datagram::AncillaryData)>;
    fn set_ecn(&mut self, ecn: ExplicitCongestionNotification, remote_address: &SocketAddress);
    fn remote_address(&self) -> Option<SocketAddress>;
    fn set_remote_address(&mut self, remote_address: &SocketAddress);
}

impl Ext for msghdr {
    #[inline]
    fn header(&self) -> Option<(datagram::Header<Handle>, datagram::AncillaryData)> {
        let addr = self.remote_address()?;
        let mut path = Handle::from_remote_address(addr.into());

        let ancillary_data = cmsg::decode(self);
        let ecn = ancillary_data.ecn;

        path.with_ancillary_data(ancillary_data);

        let header = datagram::Header { path, ecn };

        Some((header, ancillary_data))
    }

    #[inline]
    fn set_ecn(&mut self, ecn: ExplicitCongestionNotification, remote_address: &SocketAddress) {
        if ecn == ExplicitCongestionNotification::NotEct {
            return;
        }

        let ecn = ecn as libc::c_int;

        // the remote address needs to be unmapped in order to set the appropriate cmsg
        match remote_address.unmap() {
            SocketAddress::IpV4(_) => {
                // FreeBSD uses an unsigned_char for IP_TOS
                // see https://svnweb.freebsd.org/base/stable/8/sys/netinet/ip_input.c?view=markup&pathrev=247944#l1716
                #[cfg(target_os = "freebsd")]
                let ecn = ecn as libc::c_uchar;

                self.encode_cmsg(libc::IPPROTO_IP, libc::IP_TOS, ecn)
                    .unwrap()
            }
            SocketAddress::IpV6(_) => self
                .encode_cmsg(libc::IPPROTO_IPV6, libc::IPV6_TCLASS, ecn)
                .unwrap(),
        };
    }

    #[inline]
    fn remote_address(&self) -> Option<SocketAddress> {
        debug_assert!(!self.msg_name.is_null());
        match self.msg_namelen as usize {
            size if size == size_of::<sockaddr_in>() => {
                let sockaddr: &sockaddr_in = unsafe { &*(self.msg_name as *const _) };
                let port = sockaddr.sin_port.to_be();
                let addr: IpV4Address = sockaddr.sin_addr.s_addr.to_ne_bytes().into();
                Some(SocketAddressV4::new(addr, port).into())
            }
            size if size == size_of::<sockaddr_in6>() => {
                let sockaddr: &sockaddr_in6 = unsafe { &*(self.msg_name as *const _) };
                let port = sockaddr.sin6_port.to_be();
                let addr: IpV6Address = sockaddr.sin6_addr.s6_addr.into();
                Some(SocketAddressV6::new(addr, port).into())
            }
            _ => None,
        }
    }

    #[inline]
    fn set_remote_address(&mut self, remote_address: &SocketAddress) {
        debug_assert!(!self.msg_name.is_null());

        match remote_address {
            SocketAddress::IpV4(addr) => {
                let sockaddr: &mut sockaddr_in = unsafe { &mut *(self.msg_name as *mut _) };
                sockaddr.sin_family = AF_INET as _;
                sockaddr.sin_port = addr.port().to_be();
                sockaddr.sin_addr.s_addr = u32::from_ne_bytes((*addr.ip()).into());
                self.msg_namelen = size_of::<sockaddr_in>() as _;
            }
            SocketAddress::IpV6(addr) => {
                let sockaddr: &mut sockaddr_in6 = unsafe { &mut *(self.msg_name as *mut _) };
                sockaddr.sin6_family = AF_INET6 as _;
                sockaddr.sin6_port = addr.port().to_be();
                sockaddr.sin6_addr.s6_addr = (*addr.ip()).into();
                self.msg_namelen = size_of::<sockaddr_in6>() as _;
            }
        }
    }
}


// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
// SPDX-License-Identifier: Apache-2.0

use super::ext::Ext as _;
use crate::message::cmsg::Encoder;
use libc::msghdr;
use s2n_quic_core::{
    inet::{AncillaryData, SocketAddress, SocketAddressV4},
    path::{self, LocalAddress, RemoteAddress},
};

#[cfg(any(test, feature = "generator"))]
use bolero_generator::*;

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
#[cfg_attr(any(test, feature = "generator"), derive(TypeGenerator))]
pub struct Handle {
    pub remote_address: RemoteAddress,
    pub local_address: LocalAddress,
}

impl Handle {
    #[inline]
    pub(super) fn with_ancillary_data(&mut self, ancillary_data: AncillaryData) {
        self.local_address = ancillary_data.local_address;
    }

    #[inline]
    pub(super) fn update_msg_hdr(&self, msghdr: &mut msghdr) {
        // when sending a packet, we start out with no cmsg items
        msghdr.msg_controllen = 0;

        msghdr.set_remote_address(&self.remote_address.0);

        #[cfg(s2n_quic_platform_pktinfo)]
        match self.local_address.0 {
            SocketAddress::IpV4(addr) => {
                use s2n_quic_core::inet::Unspecified;

                let ip = addr.ip();

                if ip.is_unspecified() {
                    return;
                }

                let mut pkt_info = unsafe { core::mem::zeroed::<libc::in_pktinfo>() };
                pkt_info.ipi_spec_dst.s_addr = u32::from_ne_bytes((*ip).into());

                msghdr
                    .encode_cmsg(libc::IPPROTO_IP, libc::IP_PKTINFO, pkt_info)
                    .unwrap();
            }
            SocketAddress::IpV6(addr) => {
                use s2n_quic_core::inet::Unspecified;

                let ip = addr.ip();

                if ip.is_unspecified() {
                    return;
                }

                let mut pkt_info = unsafe { core::mem::zeroed::<libc::in6_pktinfo>() };

                pkt_info.ipi6_addr.s6_addr = (*ip).into();

                msghdr
                    .encode_cmsg(libc::IPPROTO_IPV6, libc::IPV6_PKTINFO, pkt_info)
                    .unwrap();
            }
        }
    }
}

impl path::Handle for Handle {
    #[inline]
    fn from_remote_address(remote_address: RemoteAddress) -> Self {
        Self {
            remote_address,
            local_address: SocketAddressV4::UNSPECIFIED.into(),
        }
    }

    #[inline]
    fn remote_address(&self) -> RemoteAddress {
        self.remote_address
    }

    #[inline]
    fn local_address(&self) -> LocalAddress {
        self.local_address
    }

    #[inline]
    fn eq(&self, other: &Self) -> bool {
        let mut eq = true;

        // only compare local addresses if the OS returns them
        if cfg!(s2n_quic_platform_pktinfo) {
            eq &= self.local_address.eq(&other.local_address);
        }

        eq && path::Handle::eq(&self.remote_address, &other.remote_address)
    }

    #[inline]
    fn strict_eq(&self, other: &Self) -> bool {
        PartialEq::eq(self, other)
    }

    #[inline]
    fn maybe_update(&mut self, other: &Self) {
        // once we discover our path, update the address local address
        if self.local_address.port() == 0 {
            self.local_address = other.local_address;
        }
    }
}


// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
// SPDX-License-Identifier: Apache-2.0

use super::*;
use bolero::check;
use core::mem::zeroed;
use s2n_quic_core::inet::{SocketAddress, Unspecified};

fn test_msghdr<F: FnOnce(&mut msghdr)>(f: F) {
    const PAYLOAD_LEN: usize = 16;

    let mut msghdr = unsafe { zeroed::<msghdr>() };

    let mut msgname = unsafe { zeroed::<sockaddr_in6>() };
    msghdr.msg_name = &mut msgname as *mut _ as *mut _;
    msghdr.msg_namelen = size_of::<sockaddr_in6>() as _;

    let mut iovec = unsafe { zeroed::<iovec>() };

    let mut payload = [0u8; PAYLOAD_LEN];
    iovec.iov_base = &mut payload as *mut _ as *mut _;
    iovec.iov_len = 1;

    msghdr.msg_iov = &mut iovec;

    let mut msg_control = cmsg::Storage::default();
    msghdr.msg_controllen = msg_control.len() as _;
    msghdr.msg_control = msg_control.as_mut_ptr() as *mut _;

    unsafe {
        msghdr.reset(PAYLOAD_LEN);
    }

    f(&mut msghdr);
}

#[test]
#[cfg_attr(kani, kani::proof, kani::solver(cadical), kani::unwind(17))]
fn address_inverse_pair_test() {
    check!()
        .with_type::<SocketAddress>()
        .cloned()
        .for_each(|addr| {
            test_msghdr(|message| {
                message.set_remote_address(&addr);

                assert_eq!(message.remote_address(), Some(addr));
            });
        });
}

#[test]
#[cfg_attr(kani, kani::proof, kani::solver(cadical), kani::unwind(65))]
#[cfg(any(not(kani), kani_slow))] // this test isn't able to terminate
fn handle_get_set_test() {
    check!()
        .with_generator((
            gen::<Handle>(),
            1..=crate::features::gso::MaxSegments::MAX.into(),
        ))
        .cloned()
        .for_each(|(handle, segment_size)| {
            test_msghdr(|message| {
                handle.update_msg_hdr(message);

                if segment_size > 1 {
                    message.set_segment_size(segment_size);
                }

                let (header, _cmsg) = message.header().unwrap();

                assert_eq!(header.path.remote_address, handle.remote_address);

                if cfg!(s2n_quic_platform_pktinfo) && !handle.local_address.ip().is_unspecified() {
                    assert_eq!(header.path.local_address.ip(), handle.local_address.ip());
                }

                // reset the message and ensure everything is zeroed
                unsafe {
                    message.reset(0);
                }

                // no need to check this on kani since it has to unwind the loop again
                #[cfg(not(kani))]
                {
                    let (header, _cmsg) = message.header().unwrap();
                    assert!(header.path.remote_address.is_unspecified());
                }
            });
        });
}


// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
// SPDX-License-Identifier: Apache-2.0

pub mod rx;
pub mod tx;


// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
// SPDX-License-Identifier: Apache-2.0

//! Structure for concurrently queueing network messages
//!
//! Two halves are created: `Producer` and `Consumer`. The producer's role is:
//!
//! * Acquire capacity to send messages
//! * Fill the messages with some data
//! * Release the filled messages to the consumer
//!
//! The consumer then:
//!
//! * Acquires filled messages
//! * Reads the messages
//! * Releases the read messages back to the producer to be filled with more messages
//!
//! Normally, ring buffers wrap around the data region and return 2 slices of data (see
//! [`std::collections::VecDeque::as_mut_slices`]). This causes issues for syscalls like
//! [`libc::sendmmsg`] where it expects a single contiguous region of memory to be passed
//! for the messages. This would result in either having to make 2 syscalls for
//! when both slices have items (one of the more expensive operations we do) or copying all of the
//! messages out of the ring buffer into a [`Vec`] and passing that to the syscall. Neither of
//! these are ideal. Instead, the ring buffer capacity is doubled and split into a primary and
//! secondary region. The messages are replicated by the producer between the regions to ensure a
//! single slice can be taken at any arbitrary index by the consumer.
//!
//! Looking at an example, let's assume we have created a ring with capacity of 4. The ring will
//! actually allocate a memory region of 8 messages, where the first 4 point to the same payload
//! buffer as the last 4:
//!
//! ```ignore
//! [ primary   ]
//!              [ secondary ]
//! [ 0, 1, 2, 3, 0, 1, 2, 3 ]
//! ```
//!
//! We call the first half of the messages the "primary" region and the second half "secondary".
//! Now, let's assume that the current index of the producer is `2`. The region of memory returned
//! in the [`Producer::data`] call would be:
//!
//! ```ignore
//! [ primary   ]
//!              [ secondary ]
//! [ 0, 1, 2, 3, 0, 1, 2, 3 ]
//!        [ data     ]
//! ```
//!
//! If the producer fills the `data` slice with messages it will have written into both the primary
//! and secondary regions and those values need to be copied to the areas that weren't written to
//! in order to maintain a consistent view of the data:
//!
//! ```ignore
//! [ primary   ]
//!              [ secondary ]
//! [ 0, 1, 2, 3, 0, 1, 2, 3 ]
//!        [ data     ]
//!        [src ]  ->  [ dst ]
//! [ dst ]  <-  [src ]
//! ```
//!
//! When the consumer goes to read the queue it can do so at its own index without having to split
//! the data, even if it wraps around the end.

use crate::message::{self, Message};
use alloc::sync::Arc;
use core::{
    mem::size_of,
    ptr::NonNull,
    sync::atomic::AtomicU32,
    task::{Context, Poll},
};
use s2n_quic_core::{
    assume,
    sync::{
        atomic_waker,
        cursor::{self, Cursor},
        CachePadded,
    },
};

const CURSOR_SIZE: usize = size_of::<CachePadded<AtomicU32>>();
const PRODUCER_OFFSET: usize = 0;
const CONSUMER_OFFSET: usize = CURSOR_SIZE;
const DATA_OFFSET: usize = CURSOR_SIZE * 2;

/// Creates a pair of rings for a given message type
pub fn pair<T: Message>(entries: u32, payload_len: u32) -> (Producer<T>, Consumer<T>) {
    let storage = T::alloc(entries, payload_len, DATA_OFFSET);

    let storage = Arc::new(storage);
    let ptr = storage.as_ptr();

    let wakers = atomic_waker::pair();

    let consumer = Consumer {
        cursor: unsafe { builder(ptr, entries).build_consumer() },
        wakers: wakers.0,
        storage: storage.clone(),
    };

    let producer = Producer {
        cursor: unsafe { builder(ptr, entries).build_producer() },
        wakers: wakers.1,
        storage,
    };

    (producer, consumer)
}

/// A consumer ring for messages
pub struct Consumer<T: Message> {
    cursor: Cursor<T>,
    wakers: atomic_waker::Handle,
    #[allow(dead_code)]
    storage: Arc<message::Storage>,
}

/// Safety: Storage is synchronized with the Cursor
unsafe impl<T: Message> Send for Consumer<T> {}
/// Safety: Storage is synchronized with the Cursor
unsafe impl<T: Message> Sync for Consumer<T> {}

impl<T: Message> Consumer<T> {
    /// Acquires ready-to-consume messages from the producer
    #[inline]
    pub fn acquire(&mut self, watermark: u32) -> u32 {
        self.cursor.acquire_consumer(watermark)
    }

    /// Polls ready-to-consume messages from the producer
    #[inline]
    pub fn poll_acquire(&mut self, watermark: u32, cx: &mut Context) -> Poll<u32> {
        macro_rules! try_acquire {
            () => {{
                let count = self.acquire(watermark);

                if count > 0 {
                    return Poll::Ready(count);
                }
            }};
        }

        // first try to acquire some messages
        try_acquire!();

        // if we couldn't acquire anything register our waker
        self.wakers.register(cx.waker());

        // try to acquire some messages in case we got some concurrently to waker registration
        try_acquire!();

        Poll::Pending
    }

    /// Releases consumed messages to the producer
    #[inline]
    pub fn release(&mut self, release_len: u32) {
        self.cursor.release_consumer(release_len);

        self.wakers.wake();
    }

    /// Returns the currently acquired messages
    #[inline]
    pub fn data(&mut self) -> &mut [T] {
        let idx = self.cursor.cached_consumer();
        let len = self.cursor.cached_consumer_len();
        let ptr = self.cursor.data_ptr();
        unsafe {
            let ptr = ptr.as_ptr().add(idx as _);
            core::slice::from_raw_parts_mut(ptr, len as _)
        }
    }

    /// Returns true if the producer is not closed
    #[inline]
    pub fn is_open(&self) -> bool {
        self.wakers.is_open()
    }
}

/// A producer ring for messages
pub struct Producer<T: Message> {
    cursor: Cursor<T>,
    wakers: atomic_waker::Handle,
    #[allow(dead_code)]
    storage: Arc<message::Storage>,
}

/// Safety: Storage is synchronized with the Cursor
unsafe impl<T: Message> Send for Producer<T> {}
/// Safety: Storage is synchronized with the Cursor
unsafe impl<T: Message> Sync for Producer<T> {}

impl<T: Message> Producer<T> {
    /// Acquires capacity for sending messages to the consumer
    #[inline]
    pub fn acquire(&mut self, watermark: u32) -> u32 {
        self.cursor.acquire_producer(watermark)
    }

    /// Polls capacity for sending messages to the consumer
    #[inline]
    pub fn poll_acquire(&mut self, watermark: u32, cx: &mut Context) -> Poll<u32> {
        macro_rules! try_acquire {
            () => {{
                let count = self.acquire(watermark);

                if count > 0 {
                    return Poll::Ready(count);
                }
            }};
        }

        // first try to acquire some messages
        try_acquire!();

        // if we couldn't acquire anything register our waker
        self.wakers.register(cx.waker());

        // try to acquire some messages in case we got some concurrently to waker registration
        try_acquire!();

        Poll::Pending
    }

    /// Releases ready-to-consume messages to the consumer
    #[inline]
    pub fn release(&mut self, release_len: u32) {
        if release_len == 0 {
            return;
        }

        debug_assert!(
            release_len <= self.cursor.cached_producer_len(),
            "cannot release more messages than acquired"
        );

        let idx = self.cursor.cached_producer();
        let ring_size = self.cursor.capacity();

        // replicate any written items to the secondary region
        unsafe {
            assume!(ring_size > idx, "idx should never exceed the ring size");

            // calculate the maximum number of replications we need to perform for the primary ->
            // secondary
            let max_possible_replications = ring_size - idx;
            // the replication count should exceed the number that we're releasing
            let replication_count = max_possible_replications.min(release_len);

            assume!(
                replication_count != 0,
                "we should always be releasing at least 1 item"
            );

            // calculate the data pointer based on the current message index
            let primary = self.cursor.data_ptr().as_ptr().add(idx as _);
            // add the size of the ring to the primary pointer to get into the secondary message
            let secondary = primary.add(ring_size as _);

            // copy the primary into the secondary
            self.replicate(primary, secondary, replication_count as _);

            // if messages were also written to the secondary region, we need to copy them back to the
            // primary region
            assume!(
                idx.checked_add(release_len).is_some(),
                "overflow amount should not exceed u32::MAX"
            );
            assume!(
                idx + release_len < ring_size * 2,
                "overflow amount should not extend beyond the secondary replica"
            );

            let overflow_amount = (idx + release_len).checked_sub(ring_size).filter(|v| {
                // we didn't overflow if the count is 0
                *v > 0
            });

            if let Some(replication_count) = overflow_amount {
                // secondary -> primary replication always happens at the beginning of the data
                let primary = self.cursor.data_ptr().as_ptr();
                // add the size of the ring to the primary pointer to get into the secondary
                // message
                let secondary = primary.add(ring_size as _);

                // copy the secondary into the primary
                self.replicate(secondary, primary, replication_count as _);
            }
        }

        // finally release the len to the consumer
        self.cursor.release_producer(release_len);

        // wake up the consumer to notify it of progress
        self.wakers.wake();
    }

    /// Returns the empty messages for the producer
    #[inline]
    pub fn data(&mut self) -> &mut [T] {
        let idx = self.cursor.cached_producer();
        let len = self.cursor.cached_producer_len();
        let ptr = self.cursor.data_ptr();
        unsafe {
            let ptr = ptr.as_ptr().add(idx as _);
            core::slice::from_raw_parts_mut(ptr, len as _)
        }
    }

    /// Returns true if the consumer is not closed
    #[inline]
    pub fn is_open(&self) -> bool {
        self.wakers.is_open()
    }

    /// Replicates messages from the primary to secondary memory regions
    #[inline]
    unsafe fn replicate(&self, primary: *mut T, secondary: *mut T, len: usize) {
        debug_assert_ne!(len, 0);

        #[cfg(debug_assertions)]
        {
            let primary = core::slice::from_raw_parts(primary, len as _);
            let secondary = core::slice::from_raw_parts(secondary, len as _);
            for (primary, secondary) in primary.iter().zip(secondary) {
                T::validate_replication(primary, secondary);
            }
        }

        core::ptr::copy_nonoverlapping(primary, secondary, len as _);
    }
}

#[inline]
unsafe fn builder<T: Message>(ptr: *mut u8, size: u32) -> cursor::Builder<T> {
    let producer = ptr.add(PRODUCER_OFFSET) as *mut _;
    let producer = NonNull::new(producer).unwrap();
    let consumer = ptr.add(CONSUMER_OFFSET) as *mut _;
    let consumer = NonNull::new(consumer).unwrap();
    let data = ptr.add(DATA_OFFSET) as *mut _;
    let data = NonNull::new(data).unwrap();

    cursor::Builder {
        producer,
        consumer,
        data,
        size,
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use bolero::check;
    use s2n_quic_core::{
        inet::{ExplicitCongestionNotification, SocketAddress},
        path::{Handle as _, LocalAddress, RemoteAddress},
    };

    #[cfg(not(kani))]
    type Counts = Vec<u32>;
    #[cfg(kani)]
    type Counts = s2n_quic_core::testing::InlineVec<u32, 2>;

    macro_rules! replication_test {
        ($name:ident, $msg:ty) => {
            #[test]
            #[cfg_attr(kani, kani::proof, kani::solver(cadical), kani::unwind(3))]
            #[cfg(any(not(kani), kani_slow))] // this test takes too much memory for our CI
                                              // environment
            fn $name() {
                check!().with_type::<Counts>().for_each(|counts| {
                    let entries = if cfg!(kani) { 2 } else { 16 };
                    let payload_len = if cfg!(kani) { 2 } else { 128 };

                    let (mut producer, mut consumer) = pair::<$msg>(entries, payload_len);

                    let mut counter = 0;

                    for count in counts.iter().copied() {
                        let count = producer.acquire(count);

                        for entry in &mut producer.data()[..count as usize] {
                            unsafe {
                                entry.set_payload_len(counter);
                            }
                            counter += 1;
                        }

                        producer.release(count);

                        #[cfg(kani)]
                        let ids_to_check = {
                            let idx: u32 = kani::any();
                            kani::assume(idx < entries);
                            idx..idx + 1
                        };

                        #[cfg(not(kani))]
                        let ids_to_check = 0..entries;

                        for idx in ids_to_check {
                            let ptr = producer.cursor.data_ptr().as_ptr();
                            unsafe {
                                let primary = &*ptr.add(idx as _);
                                let secondary = &*ptr.add((idx + entries) as _);

                                assert_eq!(primary.payload_len(), secondary.payload_len());
                            }
                        }

                        let count = consumer.acquire(count);
                        consumer.release(count);
                    }
                });
            }
        };
    }

    replication_test!(simple_replication, crate::message::simple::Message);
    #[cfg(s2n_quic_platform_socket_msg)]
    replication_test!(msg_replication, crate::message::msg::Message);
    #[cfg(s2n_quic_platform_socket_mmsg)]
    replication_test!(mmsg_replication, crate::message::mmsg::Message);

    macro_rules! send_recv_test {
        ($name:ident, $msg:ty) => {
            #[test]
            fn $name() {
                check!().with_type::<Counts>().for_each(|counts| {
                    let entries = if cfg!(miri) { 2 } else { 16 };
                    let payload_len = if cfg!(miri) { 4 } else { 128 };

                    let (mut producer, mut consumer) = pair::<$msg>(entries, payload_len);

                    let mut tx_counter = 0u32;
                    let mut rx_counter = 0u32;

                    let local_address = LocalAddress::from(SocketAddress::default());

                    for count in counts.iter().copied() {
                        let count = producer.acquire(count);

                        for entry in &mut producer.data()[..count as usize] {
                            unsafe {
                                entry.reset(payload_len as _);
                            }

                            let mut remote_address = SocketAddress::default();
                            remote_address.set_port(tx_counter as _);
                            let remote_address = RemoteAddress::from(remote_address);
                            let handle =
                                <$msg as Message>::Handle::from_remote_address(remote_address);
                            let ecn = ExplicitCongestionNotification::new(tx_counter as _);
                            let payload = tx_counter.to_le_bytes();
                            let msg = (handle, ecn, &payload[..]);
                            entry.tx_write(msg).unwrap();
                            tx_counter += 1;
                        }

                        producer.release(count);

                        let count = consumer.acquire(count);
                        for entry in consumer.data() {
                            let message = entry.rx_read(&local_address).unwrap();
                            message.for_each(|header, payload| {
                                if <$msg>::SUPPORTS_ECN {
                                    let ecn = ExplicitCongestionNotification::new(rx_counter as _);
                                    assert_eq!(header.ecn, ecn);
                                }

                                let counter: &[u8; 4] = (&*payload).try_into().unwrap();
                                let counter = u32::from_le_bytes(*counter);
                                assert_eq!(counter, rx_counter);

                                rx_counter += 1;
                            });
                        }
                        consumer.release(count);
                    }
                });
            }
        };
    }

    send_recv_test!(simple_send_recv, crate::message::simple::Message);
    #[cfg(s2n_quic_platform_socket_msg)]
    send_recv_test!(msg_send_recv, crate::message::msg::Message);
    #[cfg(s2n_quic_platform_socket_mmsg)]
    send_recv_test!(mmsg_send_recv, crate::message::mmsg::Message);
}


// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
// SPDX-License-Identifier: Apache-2.0

use crate::{
    buffer::Buffer,
    message::{
        queue,
        simple::{self, Message, Ring},
        Message as _,
    },
};
use errno::errno;
use s2n_quic_core::{event, inet::SocketAddress, path::LocalAddress};

pub use simple::Handle;

pub trait Socket {
    type Error: Error;

    /// Receives a payload and returns the length and source address
    fn recv_from(&self, buf: &mut [u8]) -> Result<(usize, Option<SocketAddress>), Self::Error>;

    /// Sends a payload to the given address and returns the length of the sent payload
    fn send_to(&self, buf: &[u8], addr: &SocketAddress) -> Result<usize, Self::Error>;
}

#[cfg(feature = "std")]
impl Socket for std::net::UdpSocket {
    type Error = std::io::Error;

    fn recv_from(&self, buf: &mut [u8]) -> Result<(usize, Option<SocketAddress>), Self::Error> {
        debug_assert!(!buf.is_empty());
        let (len, addr) = self.recv_from(buf)?;
        Ok((len, Some(addr.into())))
    }

    fn send_to(&self, buf: &[u8], addr: &SocketAddress) -> Result<usize, Self::Error> {
        debug_assert!(!buf.is_empty());
        let addr: std::net::SocketAddr = (*addr).into();
        self.send_to(buf, addr)
    }
}

pub trait Error {
    fn would_block(&self) -> bool;
    fn was_interrupted(&self) -> bool;
    fn permission_denied(&self) -> bool;
    fn connection_reset(&self) -> bool;
}

#[cfg(feature = "std")]
impl Error for std::io::Error {
    fn would_block(&self) -> bool {
        self.kind() == std::io::ErrorKind::WouldBlock
    }

    fn was_interrupted(&self) -> bool {
        self.kind() == std::io::ErrorKind::Interrupted
    }

    fn permission_denied(&self) -> bool {
        self.kind() == std::io::ErrorKind::PermissionDenied
    }

    fn connection_reset(&self) -> bool {
        self.kind() == std::io::ErrorKind::ConnectionReset
    }
}

#[derive(Debug, Default)]
pub struct Queue<B: Buffer>(queue::Queue<Ring<B>>);

impl<B: Buffer> Queue<B> {
    pub fn new(buffer: B) -> Self {
        let queue = queue::Queue::new(Ring::new(buffer, 1));

        Self(queue)
    }

    pub fn free_len(&self) -> usize {
        self.0.free_len()
    }

    pub fn occupied_len(&self) -> usize {
        self.0.occupied_len()
    }

    pub fn set_local_address(&mut self, local_address: LocalAddress) {
        self.0.set_local_address(local_address)
    }

    pub fn tx<S: Socket, Publisher: event::EndpointPublisher>(
        &mut self,
        socket: &S,
        publisher: &mut Publisher,
    ) -> Result<usize, S::Error> {
        let mut count = 0;
        let mut entries = self.0.occupied_mut();

        for entry in entries.as_mut() {
            let remote_address = *entry.remote_address();
            match socket.send_to(entry.payload_mut(), &remote_address) {
                Ok(_) => {
                    count += 1;

                    publisher.on_platform_tx(event::builder::PlatformTx { count: 1 });
                }
                Err(err) if count > 0 && err.would_block() => {
                    break;
                }
                Err(err) if err.was_interrupted() || err.permission_denied() => {
                    break;
                }
                Err(err) => {
                    entries.finish(count);

                    publisher
                        .on_platform_tx_error(event::builder::PlatformTxError { errno: errno().0 });

                    return Err(err);
                }
            }
        }

        entries.finish(count);

        Ok(count)
    }

    pub fn rx<S: Socket, Publisher: event::EndpointPublisher>(
        &mut self,
        socket: &S,
        publisher: &mut Publisher,
    ) -> Result<usize, S::Error> {
        let mut count = 0;
        let mut entries = self.0.free_mut();

        while let Some(entry) = entries.get_mut(count) {
            match socket.recv_from(entry.payload_mut()) {
                Ok((payload_len, Some(remote_address))) => {
                    entry.set_remote_address(&remote_address);
                    unsafe {
                        // Safety: The payload_len should not be bigger than the number of
                        // allocated bytes.

                        debug_assert!(payload_len < entry.payload_len());
                        let payload_len = payload_len.min(entry.payload_len());

                        entry.set_payload_len(payload_len);
                    }

                    count += 1;

                    publisher.on_platform_rx(event::builder::PlatformRx { count: 1 });
                }
                Ok((_payload_len, None)) => {}
                Err(err) if count > 0 && err.would_block() => {
                    break;
                }
                Err(err) if err.was_interrupted() => {
                    break;
                }
                Err(err) if err.connection_reset() => {
                    count += 1;
                }
                Err(err) => {
                    entries.finish(count);

                    publisher
                        .on_platform_rx_error(event::builder::PlatformRxError { errno: errno().0 });

                    return Err(err);
                }
            }
        }

        entries.finish(count);

        Ok(count)
    }

    pub fn rx_queue(&mut self) -> queue::Occupied<Message> {
        self.0.occupied_mut()
    }

    pub fn tx_queue(&mut self) -> queue::Free<Message> {
        self.0.free_mut()
    }
}


// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
// SPDX-License-Identifier: Apache-2.0

mod events;
pub mod rx;
pub mod tx;

pub use rx::Receiver;
pub use tx::Sender;


// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
// SPDX-License-Identifier: Apache-2.0

use crate::{message::Message, socket::ring::Consumer};
use core::task::{Context, Poll};
use s2n_quic_core::{
    event,
    inet::datagram,
    io::rx,
    path::{LocalAddress, MaxMtu},
};

/// Structure for receiving messages from consumer channels
pub struct Rx<T: Message> {
    channels: Vec<Consumer<T>>,
    max_mtu: MaxMtu,
    local_address: LocalAddress,
}

impl<T: Message> Rx<T> {
    #[inline]
    pub fn new(channels: Vec<Consumer<T>>, max_mtu: MaxMtu, local_address: LocalAddress) -> Self {
        Self {
            channels,
            max_mtu,
            local_address,
        }
    }
}

impl<T: Message> rx::Rx for Rx<T> {
    type PathHandle = T::Handle;
    type Queue = RxQueue<'static, T>;
    type Error = ();

    #[inline]
    fn poll_ready(&mut self, cx: &mut Context) -> Poll<Result<(), Self::Error>> {
        let mut is_any_ready = false;
        let mut is_all_closed = true;

        // try to acquire any messages we can from the set of channels
        for channel in &mut self.channels {
            match channel.poll_acquire(u32::MAX, cx) {
                Poll::Ready(_) => {
                    is_all_closed = false;
                    is_any_ready = true;
                }
                Poll::Pending => {
                    is_all_closed &= !channel.is_open();
                }
            }
        }

        // if all of the channels are closed then shut down the task
        if is_all_closed {
            return Err(()).into();
        }

        // if any have items to be consumed the wake the endpoint up
        if is_any_ready {
            Poll::Ready(Ok(()))
        } else {
            Poll::Pending
        }
    }

    #[inline]
    fn queue<F: FnOnce(&mut Self::Queue)>(&mut self, f: F) {
        let this: &'static mut Self = unsafe {
            // Safety: As noted in the [transmute examples](https://doc.rust-lang.org/std/mem/fn.transmute.html#examples)
            // it can be used to temporarily extend the lifetime of a reference. In this case, we
            // don't want to use GATs until the MSRV is >=1.65.0, which means `Self::Queue` is not
            // allowed to take generic lifetimes.
            //
            // We are left with using a `'static` lifetime here and encapsulating it in a private
            // field. The `Self::Queue` struct is then borrowed for the lifetime of the `F`
            // function. This will prevent the value from escaping beyond the lifetime of `&mut
            // self`.
            //
            // See https://play.rust-lang.org/?version=stable&mode=debug&edition=2021&gist=9a32abe85c666f36fb2ec86496cc41b4
            //
            // Once https://github.com/aws/s2n-quic/issues/1742 is resolved this code can go away
            core::mem::transmute(self)
        };

        let mut queue = RxQueue {
            channels: &mut this.channels,
            max_mtu: this.max_mtu,
            local_address: &this.local_address,
        };

        f(&mut queue);
    }

    #[inline]
    fn handle_error<E: event::EndpointPublisher>(self, _error: Self::Error, _events: &mut E) {
        // The only reason we would be returning an error is if a channel closed. This could either
        // be because the endpoint is shutting down or one of the tasks panicked. Either way, we
        // don't know what the cause is here so we don't have any events to emit.
    }
}

pub struct RxQueue<'a, T: Message> {
    channels: &'a mut [Consumer<T>],
    max_mtu: MaxMtu,
    local_address: &'a LocalAddress,
}

impl<'a, T: Message> rx::Queue for RxQueue<'a, T> {
    type Handle = T::Handle;

    #[inline]
    fn for_each<F: FnMut(datagram::Header<Self::Handle>, &mut [u8])>(&mut self, mut on_packet: F) {
        for channel in self.channels.iter_mut() {
            // one last effort to acquire items if some were received since we last polled
            let len = channel.acquire(u32::MAX);

            let data = channel.data();
            debug_assert_eq!(data.len(), len as usize);

            for message in data {
                // call the `on_packet` function for each message received
                //
                // NOTE: it's important that we process all of the messages in the queue as the
                //       channel is completely drained here.
                if let Some(message) = message.rx_read(self.local_address) {
                    message.for_each(&mut on_packet);
                }

                unsafe {
                    // Safety: the message was allocated with the configured MaxMtu
                    message.reset(self.max_mtu.into());
                }
            }

            // release the messages back to the producer
            channel.release(len);
        }
    }

    #[inline]
    fn is_empty(&self) -> bool {
        false
    }
}


// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
// SPDX-License-Identifier: Apache-2.0

use crate::{features::Gso, message::Message, socket::ring::Producer};
use core::task::{Context, Poll};
use s2n_quic_core::{
    event,
    inet::ExplicitCongestionNotification,
    io::tx,
    path::{Handle as _, MaxMtu},
};

/// Structure for sending messages to producer channels
pub struct Tx<T: Message> {
    channels: Vec<Producer<T>>,
    gso: Gso,
    max_mtu: usize,
    is_full: bool,
}

impl<T: Message> Tx<T> {
    #[inline]
    pub fn new(channels: Vec<Producer<T>>, gso: Gso, max_mtu: MaxMtu) -> Self {
        Self {
            channels,
            gso,
            max_mtu: max_mtu.into(),
            is_full: true,
        }
    }
}

impl<T: Message> tx::Tx for Tx<T> {
    type PathHandle = T::Handle;
    type Queue = TxQueue<'static, T>;
    type Error = ();

    #[inline]
    fn poll_ready(&mut self, cx: &mut Context) -> Poll<Result<(), Self::Error>> {
        // We only need to poll for capacity if we completely filled up all of the channels.
        // If we always polled, this would cause the endpoint to spin since most of the time it has
        // capacity for sending.
        if !self.is_full {
            return Poll::Pending;
        }

        let mut is_any_ready = false;
        let mut is_all_closed = true;

        for channel in &mut self.channels {
            match channel.poll_acquire(1, cx) {
                Poll::Ready(_) => {
                    is_all_closed = false;
                    is_any_ready = true;
                }
                Poll::Pending => {
                    is_all_closed &= !channel.is_open();
                }
            }
        }

        // if all of the channels were closed then shut the task down
        if is_all_closed {
            return Err(()).into();
        }

        // if any of the channels became ready then wake the endpoint up
        if is_any_ready {
            Poll::Ready(Ok(()))
        } else {
            Poll::Pending
        }
    }

    #[inline]
    fn queue<F: FnOnce(&mut Self::Queue)>(&mut self, f: F) {
        let this: &'static mut Self = unsafe {
            // Safety: As noted in the [transmute examples](https://doc.rust-lang.org/std/mem/fn.transmute.html#examples)
            // it can be used to temporarily extend the lifetime of a reference. In this case, we
            // don't want to use GATs until the MSRV is >=1.65.0, which means `Self::Queue` is not
            // allowed to take generic lifetimes.
            //
            // We are left with using a `'static` lifetime here and encapsulating it in a private
            // field. The `Self::Queue` struct is then borrowed for the lifetime of the `F`
            // function. This will prevent the value from escaping beyond the lifetime of `&mut
            // self`.
            //
            // See https://play.rust-lang.org/?version=stable&mode=debug&edition=2021&gist=9a32abe85c666f36fb2ec86496cc41b4
            //
            // Once https://github.com/aws/s2n-quic/issues/1742 is resolved this code can go away
            core::mem::transmute(self)
        };

        let mut capacity = 0;
        let mut first_with_free_slots = None;
        for (idx, channel) in this.channels.iter_mut().enumerate() {
            // try to make one more effort to acquire capacity for sending
            let count = channel.acquire(u32::MAX) as usize;

            if count > 0 && first_with_free_slots.is_none() {
                // find the first channel that had capacity
                first_with_free_slots = Some(idx);
            }

            capacity += count;
        }

        // mark that we're still full so we need to poll and wake up next iteration
        this.is_full = capacity == 0;

        // start with the first queue that has free slots, otherwise set the index to the length,
        // which will return an AtCapacity error immediately.
        let channel_index = first_with_free_slots.unwrap_or(this.channels.len());

        // query the maximum number of segments we can fill at this point in time
        //
        // NOTE: this value could be lowered in the case the TX task encounters an error with GSO
        //       so we do need to query it each iteration.
        let max_segments = this.gso.max_segments();

        let mut queue = TxQueue {
            channels: &mut this.channels,
            channel_index,
            message_index: 0,
            pending_release: 0,
            gso_segment: None,
            max_segments,
            max_mtu: this.max_mtu,
            capacity,
            is_full: &mut this.is_full,
        };

        f(&mut queue);
    }

    #[inline]
    fn handle_error<E: event::EndpointPublisher>(self, _error: Self::Error, _events: &mut E) {
        // The only reason we would be returning an error is if a channel closed. This could either
        // be because the endpoint is shutting down or one of the tasks panicked. Either way, we
        // don't know what the cause is here so we don't have any events to emit.
    }
}

/// Tracks the current state of a GSO message
#[derive(Debug, Default)]
pub struct GsoSegment<Handle> {
    /// The path handle of the current GSO segment being written
    ///
    /// This is used to determine if future messages should be included in this payload or need a
    /// separate packet.
    handle: Handle,
    /// The value of the ecn markings for the current GSO segment being written.
    ///
    /// This is used to determine if future messages should be included in this payload or need a
    /// separate packet.
    ecn: ExplicitCongestionNotification,
    /// The number of segments that have been written
    count: usize,
    /// The size of each segment.
    ///
    /// Note that the last segment can be smaller than the previous ones and will result in a flush
    size: usize,
}

pub struct TxQueue<'a, T: Message> {
    channels: &'a mut [Producer<T>],
    /// The channel index that we are currently operating on.
    ///
    /// This will be incremented after each channel is filled until it exceeds the len of `channels`.
    channel_index: usize,
    /// The message index into the current channel that we are operating on.
    ///
    /// This is incremented after each message is finished until it exceeds the acquired free
    /// slots, after which the `channel_index` is incremented (and message_index is reset to zero).
    message_index: usize,
    /// The number of messages in the current channel that need to be released to notify the
    /// consumer.
    ///
    /// This is to avoid calling `release` for each message and waking up the socket task too much.
    pending_release: u32,
    /// The current GSO segment we are filling, if any
    gso_segment: Option<GsoSegment<T::Handle>>,
    /// The maximum number of GSO segments that can be written
    max_segments: usize,
    /// The maximum MTU for any given packet
    max_mtu: usize,
    /// The maximum number of packets that can be sent in the current iteration
    capacity: usize,
    /// Used to track if we have filled up the producer queue and waiting on free slots to be
    /// released by the consumer.
    is_full: &'a mut bool,
}

impl<'a, T: Message> TxQueue<'a, T> {
    /// Tries to send a message as a GSO segment
    ///
    /// Returns the Err(Message) if it was not able to. Otherwise, the index of the GSO'd message is returned.
    #[inline]
    fn try_gso<M: tx::Message<Handle = T::Handle>>(
        &mut self,
        mut message: M,
    ) -> Result<Result<tx::Outcome, M>, tx::Error> {
        // the message doesn't support GSO to return it
        if !T::SUPPORTS_GSO {
            return Ok(Err(message));
        }

        let max_segments = self.max_segments;

        let (prev_message, gso) = if let Some(gso) = self.gso_message() {
            gso
        } else {
            return Ok(Err(message));
        };

        debug_assert!(
            max_segments > 1,
            "gso_segment should only be set when max_gso > 1"
        );

        // check to make sure the message can be GSO'd and can be included in the same
        // GSO payload as the previous message
        let can_gso = message.can_gso(gso.size, gso.count)
            && message.path_handle().strict_eq(&gso.handle)
            && message.ecn() == gso.ecn;

        // if we can't use GSO then flush the current message
        if !can_gso {
            self.flush_gso();
            return Ok(Err(message));
        }

        debug_assert!(
            gso.count < max_segments,
            "{} cannot exceed {}",
            gso.count,
            max_segments
        );

        let payload_len = prev_message.payload_len();

        let buffer = unsafe {
            // Create a slice the `message` can write into. This avoids having to update the
            // payload length and worrying about panic safety.

            let payload = prev_message.payload_ptr_mut();

            // Safety: all payloads should have enough capacity to extend max_segments *
            // gso.size
            let current_payload = payload.add(payload_len);
            core::slice::from_raw_parts_mut(current_payload, gso.size)
        };
        let buffer = tx::PayloadBuffer::new(buffer);

        let size = message.write_payload(buffer, gso.count)?;

        // we don't want to send empty packets
        if size == 0 {
            return Err(tx::Error::EmptyPayload);
        }

        unsafe {
            debug_assert!(
                gso.size >= size,
                "the payload tried to write more than available"
            );
            // Set the len to the actual amount written to the payload. In case there is a bug,
            // take the min anyway so we don't have errors elsewhere.
            prev_message.set_payload_len(payload_len + size.min(gso.size));
        }
        // increment the number of segments that we've written
        gso.count += 1;

        debug_assert!(
            gso.count <= max_segments,
            "{} cannot exceed {}",
            gso.count,
            max_segments
        );

        // the last segment can be smaller but we can't write any more if it is
        let size_mismatch = gso.size != size;

        // we're bounded by the max_segments amount
        let at_segment_limit = gso.count >= max_segments;

        // we also can't write more data than u16::MAX
        let at_payload_limit = gso.size * (gso.count + 1) > u16::MAX as usize;

        // if we've hit any limits, then flush the GSO information to the message
        if size_mismatch || at_segment_limit || at_payload_limit {
            self.flush_gso();
        }

        Ok(Ok(tx::Outcome {
            len: size,
            index: 0,
        }))
    }

    /// Flushes the current GSO message, if any
    ///
    /// In the `gso_segment` field, we track which message is currently being
    /// built. If there ended up being multiple payloads written to the single message
    /// we need to set the msg_control values to indicate the GSO size.
    #[inline]
    fn flush_gso(&mut self) {
        // no need to flush if the message type doesn't support GSO
        if !T::SUPPORTS_GSO {
            debug_assert!(
                self.gso_segment.is_none(),
                "gso_segment should not be set if GSO is unsupported"
            );
            return;
        }

        if let Some((message, gso)) = self.gso_message() {
            // only need to set the segment size if there was more than one payload written to the message
            if gso.count > 1 {
                message.set_segment_size(gso.size);
            }

            // clear out the current state and release the message
            self.gso_segment = None;
            self.release_message();
        }
    }

    /// Returns the current GSO message waiting for more segments
    #[inline]
    fn gso_message(&mut self) -> Option<(&mut T, &mut GsoSegment<T::Handle>)> {
        let gso = self.gso_segment.as_mut()?;

        let channel = unsafe {
            // Safety: the channel_index should always be in-bound if gso_segment is set
            s2n_quic_core::assume!(self.channels.len() > self.channel_index);
            &mut self.channels[self.channel_index]
        };

        let message = unsafe {
            // Safety: the message_index should always be in-bound if gso_segment is set
            let data = channel.data();
            s2n_quic_core::assume!(data.len() > self.message_index);
            &mut data[self.message_index]
        };

        Some((message, gso))
    }

    /// Releases the current message and marks it pending for release
    #[inline]
    fn release_message(&mut self) {
        self.capacity -= 1;
        *self.is_full = self.capacity == 0;
        self.message_index += 1;
        self.pending_release += 1;
    }

    /// Flushes the current channel and releases any pending messages
    #[inline]
    fn flush_channel(&mut self) {
        if let Some(channel) = self.channels.get_mut(self.channel_index) {
            channel.release(self.pending_release);
            self.message_index = 0;
            self.pending_release = 0;
        }
    }
}

impl<'a, T: Message> tx::Queue for TxQueue<'a, T> {
    type Handle = T::Handle;

    const SUPPORTS_ECN: bool = T::SUPPORTS_ECN;
    const SUPPORTS_FLOW_LABELS: bool = T::SUPPORTS_FLOW_LABELS;

    #[inline]
    fn push<M>(&mut self, message: M) -> Result<tx::Outcome, tx::Error>
    where
        M: tx::Message<Handle = Self::Handle>,
    {
        // first try to write a GSO payload, if supported
        let mut message = match self.try_gso(message)? {
            Ok(outcome) => return Ok(outcome),
            Err(message) => message,
        };

        // find the next free entry, if any
        let entry = loop {
            let channel = self
                .channels
                .get_mut(self.channel_index)
                .ok_or(tx::Error::AtCapacity)?;

            if let Some(entry) = channel.data().get_mut(self.message_index) {
                break entry;
            } else {
                // this channel is out of free messages so flush it and move to the next channel
                self.flush_channel();
                self.channel_index += 1;
            };
        };

        // prepare the entry for writing and reset all of the fields
        unsafe {
            // Safety: the entries should have been allocated with the MaxMtu
            entry.reset(self.max_mtu);
        }

        // query the values that we use for GSO before we write the message to the entry
        let handle = *message.path_handle();
        let ecn = message.ecn();
        let can_gso = message.can_gso(self.max_mtu, 0);

        // write the message to the entry
        let payload_len = entry.tx_write(message)?;

        // if GSO is supported and we are allowed to have additional segments, store the GSO state
        // for another potential message to be written later
        if T::SUPPORTS_GSO && self.max_segments > 1 && can_gso {
            self.gso_segment = Some(GsoSegment {
                handle,
                ecn,
                count: 1,
                size: payload_len,
            });
        } else {
            // otherwise, release the message to the consumer
            self.release_message();
        }

        // let the caller know how big the payload was
        let outcome = tx::Outcome {
            len: payload_len,
            index: 0,
        };

        Ok(outcome)
    }

    #[inline]
    fn flush(&mut self) {
        // flush GSO segments between connections
        self.flush_gso();
    }

    #[inline]
    fn capacity(&self) -> usize {
        self.capacity
    }
}

impl<'a, T: Message> Drop for TxQueue<'a, T> {
    #[inline]
    fn drop(&mut self) {
        // flush the current GSO message, if possible
        self.flush_gso();
        // flush the pending messages for the channel
        self.flush_channel();
    }
}


// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
// SPDX-License-Identifier: Apache-2.0

// Some of the functions in these impls are not used on non-unix systems
#![cfg_attr(not(unix), allow(dead_code))]

use crate::features::Gso;
use core::ops::ControlFlow;

#[derive(Debug)]
pub struct TxEvents {
    count: usize,
    is_blocked: bool,
    #[cfg_attr(not(s2n_quic_platform_gso), allow(dead_code))]
    gso: Gso,
}

impl TxEvents {
    #[inline]
    pub fn new(gso: Gso) -> Self {
        Self {
            count: 0,
            is_blocked: false,
            gso,
        }
    }

    /// Returns if the task is blocked
    #[inline]
    pub fn is_blocked(&self) -> bool {
        self.is_blocked
    }

    /// Returns if the task was blocked and resets the value
    #[inline]
    pub fn take_blocked(&mut self) -> bool {
        core::mem::take(&mut self.is_blocked)
    }

    /// Sets the task to blocked
    #[inline]
    pub fn blocked(&mut self) {
        self.is_blocked = true;
    }

    /// Returns and resets the number of messages sent
    #[inline]
    pub fn take_count(&mut self) -> usize {
        core::mem::take(&mut self.count)
    }
}

impl crate::syscall::SocketEvents for TxEvents {
    #[inline]
    fn on_complete(&mut self, count: usize) -> ControlFlow<(), ()> {
        // increment the total sent packets and reset our blocked status
        self.count += count;
        self.is_blocked = false;
        ControlFlow::Continue(())
    }

    #[inline]
    fn on_error(&mut self, error: ::std::io::Error) -> ControlFlow<(), ()> {
        use std::io::ErrorKind::*;

        match error.kind() {
            WouldBlock => {
                // record that we're blocked
                self.is_blocked = true;
                ControlFlow::Break(())
            }
            Interrupted => {
                // if we got interrupted break and have the task try again
                ControlFlow::Break(())
            }
            #[cfg(s2n_quic_platform_gso)]
            _ if error.raw_os_error() == Some(libc::EIO) => {
                // on platforms that don't support GSO we need to disable it and mark the packet as
                // "sent" even though we weren't able to.
                self.count += 1;

                self.gso.disable();

                // We `continue` instead of break because it's very unlikely the message would be
                // accepted at a later time, so we just discard the packet.
                ControlFlow::Continue(())
            }
            _ => {
                // ignore all other errors and just consider the packet sent
                self.count += 1;

                // We `continue` instead of break because it's very unlikely the message would be
                // accepted at a later time, so we just discard the packet.
                ControlFlow::Continue(())
            }
        }
    }
}

#[derive(Debug, Default)]
pub struct RxEvents {
    count: usize,
    is_blocked: bool,
}

impl RxEvents {
    /// Returns if the task is blocked
    #[inline]
    pub fn is_blocked(&self) -> bool {
        self.is_blocked
    }

    /// Returns if the task was blocked and resets the value
    #[inline]
    pub fn take_blocked(&mut self) -> bool {
        core::mem::take(&mut self.is_blocked)
    }

    /// Sets the task to blocked
    #[inline]
    pub fn blocked(&mut self) {
        self.is_blocked = true;
    }

    /// Returns and resets the number of messages sent
    #[inline]
    pub fn take_count(&mut self) -> usize {
        core::mem::take(&mut self.count)
    }
}

impl crate::syscall::SocketEvents for RxEvents {
    #[inline]
    fn on_complete(&mut self, count: usize) -> ControlFlow<(), ()> {
        // increment the total sent packets and reset our blocked status
        self.count += count;
        self.is_blocked = false;
        ControlFlow::Continue(())
    }

    #[inline]
    fn on_error(&mut self, error: ::std::io::Error) -> ControlFlow<(), ()> {
        use std::io::ErrorKind::*;

        match error.kind() {
            WouldBlock => {
                // record that we're blocked
                self.is_blocked = true;
                ControlFlow::Break(())
            }
            Interrupted => {
                // if we got interrupted break and have the task try again
                ControlFlow::Break(())
            }
            _ => {
                // ignore all other errors and have the task try again
                ControlFlow::Break(())
            }
        }
    }
}


// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
// SPDX-License-Identifier: Apache-2.0

use crate::{
    message::Message,
    socket::{ring::Producer, task::events},
};
use core::{
    future::Future,
    pin::Pin,
    task::{Context, Poll},
};
use futures::ready;

pub use events::RxEvents as Events;

pub trait Socket<T: Message> {
    type Error;

    fn recv(
        &mut self,
        cx: &mut Context,
        entries: &mut [T],
        events: &mut Events,
    ) -> Result<(), Self::Error>;
}

pub struct Receiver<T: Message, S: Socket<T>> {
    ring: Producer<T>,
    /// Implementation of a socket that fills free slots in the ring buffer
    rx: S,
    /// The number of messages that have been filled but not yet released to the consumer.
    ///
    /// This value is to avoid calling `release` too much and excessively waking up the consumer.
    pending: u32,
}

impl<T, S> Receiver<T, S>
where
    T: Message + Unpin,
    S: Socket<T> + Unpin,
{
    #[inline]
    pub fn new(ring: Producer<T>, rx: S) -> Self {
        Self {
            ring,
            rx,
            pending: 0,
        }
    }

    #[inline]
    fn poll_ring(&mut self, watermark: u32, cx: &mut Context) -> Poll<Result<(), ()>> {
        loop {
            let count = match self.ring.poll_acquire(watermark, cx) {
                Poll::Ready(count) => count,
                Poll::Pending if self.pending == 0 => {
                    return if !self.ring.is_open() {
                        Err(()).into()
                    } else {
                        Poll::Pending
                    };
                }
                Poll::Pending => 0,
            };

            // if the number of free slots increased since last time then yield
            if count > self.pending {
                return Ok(()).into();
            }

            // If there is no additional capacity available (i.e. we have filled all slots),
            // then release those filled slots for the consumer to read from. Once
            // the consumer reads, we will have spare capacity to populate again.
            self.release();
        }
    }

    #[inline]
    fn release(&mut self) {
        let to_release = core::mem::take(&mut self.pending);
        self.ring.release(to_release);
    }
}

impl<T, S> Future for Receiver<T, S>
where
    T: Message + Unpin,
    S: Socket<T> + Unpin,
{
    type Output = Option<S::Error>;

    #[inline]
    fn poll(self: Pin<&mut Self>, cx: &mut Context) -> Poll<Self::Output> {
        let this = self.get_mut();

        let mut events = Events::default();

        while !events.take_blocked() {
            if ready!(this.poll_ring(u32::MAX, cx)).is_err() {
                return None.into();
            }

            // slice the ring data by the number of slots we've already filled
            let entries = &mut this.ring.data()[this.pending as usize..];

            // perform the recv syscall
            match this.rx.recv(cx, entries, &mut events) {
                Ok(()) => {
                    // increment the number of received messages
                    this.pending += events.take_count() as u32
                }
                Err(err) => return Some(err).into(),
            }
        }

        // release any of the messages we wrote back to the consumer
        this.release();

        Poll::Pending
    }
}


// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
// SPDX-License-Identifier: Apache-2.0

use crate::{
    features::Gso,
    message::Message,
    socket::{ring::Consumer, task::events},
};
use core::{
    future::Future,
    pin::Pin,
    task::{Context, Poll},
};
use futures::ready;

pub use events::TxEvents as Events;

pub trait Socket<T: Message> {
    type Error;

    fn send(
        &mut self,
        cx: &mut Context,
        entries: &mut [T],
        events: &mut Events,
    ) -> Result<(), Self::Error>;
}

pub struct Sender<T: Message, S: Socket<T>> {
    ring: Consumer<T>,
    /// Implementation of a socket that transmits filled slots in the ring buffer
    tx: S,
    /// The number of messages that have been transmitted but not yet released to the producer.
    ///
    /// This value is to avoid calling `release` too much and excessively waking up the producer.
    pending: u32,
    events: Events,
}

impl<T, S> Sender<T, S>
where
    T: Message + Unpin,
    S: Socket<T> + Unpin,
{
    #[inline]
    pub fn new(ring: Consumer<T>, tx: S, gso: Gso) -> Self {
        Self {
            ring,
            tx,
            pending: 0,
            events: Events::new(gso),
        }
    }

    #[inline]
    fn poll_ring(&mut self, watermark: u32, cx: &mut Context) -> Poll<Result<(), ()>> {
        loop {
            let count = match self.ring.poll_acquire(watermark, cx) {
                Poll::Ready(count) => count,
                Poll::Pending if self.pending == 0 => {
                    return if !self.ring.is_open() {
                        Err(()).into()
                    } else {
                        Poll::Pending
                    };
                }
                Poll::Pending => 0,
            };

            // if the number of free slots increased since last time then yield
            if count > self.pending {
                return Ok(()).into();
            }

            // If there is no additional capacity available (i.e. we have filled all slots),
            // then release those filled slots for the consumer to read from. Once
            // the consumer reads, we will have spare capacity to populate again.
            self.release();
        }
    }

    #[inline]
    fn release(&mut self) {
        let to_release = core::mem::take(&mut self.pending);
        self.ring.release(to_release);
    }
}

impl<T, S> Future for Sender<T, S>
where
    T: Message + Unpin,
    S: Socket<T> + Unpin,
{
    type Output = Option<S::Error>;

    #[inline]
    fn poll(self: Pin<&mut Self>, cx: &mut Context) -> Poll<Self::Output> {
        let this = self.get_mut();

        while !this.events.take_blocked() {
            if ready!(this.poll_ring(u32::MAX, cx)).is_err() {
                return None.into();
            }

            // slice the ring data by the number of items we've already received
            let entries = &mut this.ring.data()[this.pending as usize..];

            // perform the send syscall
            match this.tx.send(cx, entries, &mut this.events) {
                Ok(()) => {
                    // increment the number of received messages
                    this.pending += this.events.take_count() as u32
                }
                Err(err) => return Some(err).into(),
            }
        }

        // release any of the messages we wrote back to the consumer
        this.release();

        Poll::Pending
    }
}


// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
// SPDX-License-Identifier: Apache-2.0

use super::{SocketEvents, SocketType, UnixMessage};
use libc::mmsghdr;
use std::os::unix::io::{AsRawFd, RawFd};

impl UnixMessage for mmsghdr {
    #[inline]
    fn send<E: SocketEvents>(fd: RawFd, entries: &mut [Self], events: &mut E) {
        send(&fd, entries, events)
    }

    #[inline]
    fn recv<E: SocketEvents>(fd: RawFd, ty: SocketType, entries: &mut [Self], events: &mut E) {
        recv(&fd, ty, entries, events)
    }
}

#[inline]
pub fn send<Sock: AsRawFd, E: SocketEvents>(
    socket: &Sock,
    packets: &mut [mmsghdr],
    events: &mut E,
) {
    if packets.is_empty() {
        return;
    }

    // Safety: calling a libc function is inherently unsafe as rust cannot
    // make any invariant guarantees. This has to be reviewed by humans instead
    // so the [docs](https://linux.die.net/man/2/sendmmsg) are inlined here:

    // > The sockfd argument is the file descriptor of the socket on which data
    // > is to be transmitted.
    let sockfd = socket.as_raw_fd();

    // > The msgvec argument is a pointer to an array of mmsghdr structures.
    //
    // > The msg_hdr field is a msghdr structure, as described in sendmsg(2).
    // > The msg_len field is used to return the number of bytes sent from the
    // > message in msg_hdr.
    let msgvec = packets.as_mut_ptr();

    // > The size of this array is specified in vlen.
    //
    // > The value specified in vlen is capped to UIO_MAXIOV (1024).
    let vlen = packets.len().min(1024) as _;

    // > The flags argument contains flags ORed together.
    //
    // No flags are currently set
    let flags = Default::default();

    // > The sendmmsg() system call is an extension of sendmsg(2) that allows
    // > the caller to transmit multiple messages on a socket using a single
    // > system call. (This has performance benefits for some applications.)
    //
    // > A nonblocking call sends as many messages as possible (up to the limit
    // > specified by vlen) and returns immediately.
    //
    // > On return from sendmmsg(), the msg_len fields of successive elements
    // > of msgvec are updated to contain the number of bytes transmitted from
    // > the corresponding msg_hdr. The return value of the call indicates the
    // > number of elements of msgvec that have been updated.
    //
    // > On success, sendmmsg() returns the number of messages sent from msgvec;
    // > if this is less than vlen, the caller can retry with a further sendmmsg()
    // > call to send the remaining messages.
    //
    // > On error, -1 is returned, and errno is set to indicate the error.

    let res = libc!(sendmmsg(sockfd, msgvec, vlen, flags));

    let _ = match res {
        Ok(count) => events.on_complete(count as _),
        Err(error) => events.on_error(error),
    };
}

#[inline]
pub fn recv<Sock: AsRawFd, E: SocketEvents>(
    socket: &Sock,
    socket_type: SocketType,
    packets: &mut [mmsghdr],
    events: &mut E,
) {
    if packets.is_empty() {
        return;
    }

    // Safety: calling a libc function is inherently unsafe as rust cannot
    // make any invariant guarantees. This has to be reviewed by humans instead
    // so the [docs](https://linux.die.net/man/2/recvmmsg) are inlined here:

    // > The sockfd argument is the file descriptor of the socket to receive data from.
    let sockfd = socket.as_raw_fd();

    // > The msgvec argument is a pointer to an array of mmsghdr structures.
    //
    // > The msg_len field is the number of bytes returned for the message in the entry.
    let msgvec = packets.as_mut_ptr();

    // > The size of this array is specified in vlen.
    let vlen = packets.len() as _;

    // > The flags argument contains flags ORed together.
    //
    // If the socket is blocking, set the MSG_WAITFORONE flag so we don't hang until the entire
    // buffer is full.
    let flags = match socket_type {
        SocketType::Blocking => libc::MSG_WAITFORONE,
        SocketType::NonBlocking => libc::MSG_DONTWAIT,
    };

    // > The timeout argument points to a struct timespec defining a timeout
    // > (seconds plus nanoseconds) for the receive operation.
    //
    // Since we currently only use non-blocking sockets, this isn't needed.
    // If support is added for non-blocking sockets, this will need to be
    // updated.
    let timeout = core::ptr::null_mut();

    // > The recvmmsg() system call is an extension of recvmsg(2)
    // > that allows the caller to receive multiple messages from a
    // > socket using a single system call.
    //
    // > A nonblocking call reads as many messages as are available
    // > (up to the limit specified by vlen) and returns immediately.
    //
    // > On return from recvmmsg(), successive elements of msgvec are
    // > updated to contain information about each received message:
    // > msg_len contains the size of the received message;
    // > the subfields of msg_hdr are updated as described in recvmsg(2).
    // > The return value of the call indicates the number of elements of
    // > msgvec that have been updated.
    //
    // > On success, recvmmsg() returns the number of messages received in
    // > msgvec; on error, -1 is returned, and errno is set to indicate the error.
    let res = libc!(recvmmsg(sockfd, msgvec, vlen, flags, timeout));

    let _ = match res {
        Ok(count) => events.on_complete(count as _),
        Err(error) => events.on_error(error),
    };
}


// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
// SPDX-License-Identifier: Apache-2.0

use super::{SocketEvents, SocketType, UnixMessage};
use crate::message::Message as _;
use libc::msghdr;
use std::os::unix::io::{AsRawFd, RawFd};

impl UnixMessage for msghdr {
    #[inline]
    fn send<E: SocketEvents>(fd: RawFd, entries: &mut [Self], events: &mut E) {
        send(&fd, entries, events)
    }

    #[inline]
    fn recv<E: SocketEvents>(fd: RawFd, ty: SocketType, entries: &mut [Self], events: &mut E) {
        recv(&fd, ty, entries, events)
    }
}

#[inline]
pub fn send<'a, Sock: AsRawFd, P: IntoIterator<Item = &'a mut msghdr>, E: SocketEvents>(
    socket: &Sock,
    packets: P,
    events: &mut E,
) {
    for packet in packets {
        #[cfg(debug_assertions)]
        let prev_msg_control_ptr = packet.msg_control;

        // macOS doesn't like when msg_control have valid pointers but the len is 0
        //
        // If that's the case here, then set the `msg_control` to null and restore it after
        // calling sendmsg.
        #[cfg(any(target_os = "macos", target_os = "ios", test))]
        let msg_control = {
            let msg_control = packet.msg_control;

            if packet.msg_controllen == 0 {
                packet.msg_control = core::ptr::null_mut();
            }

            msg_control
        };

        // Safety: calling a libc function is inherently unsafe as rust cannot
        // make any invariant guarantees. This has to be reviewed by humans instead
        // so the [docs](https://linux.die.net/man/2/sendmsg) are inlined here:

        // > The argument sockfd is the file descriptor of the sending socket.
        let sockfd = socket.as_raw_fd();

        // > The address of the target is given by msg.msg_name, with msg.msg_namelen
        // > specifying its size.
        //
        // > The message is pointed to by the elements of the array msg.msg_iov.
        // > The sendmsg() call also allows sending ancillary data (also known as
        // > control information).
        let msg = packet;

        // > The flags argument is the bitwise OR of zero or more flags.
        //
        // No flags are currently set
        let flags = Default::default();

        // > On success, these calls return the number of characters sent.
        // > On error, -1 is returned, and errno is set appropriately.
        let result = libc!(sendmsg(sockfd, msg, flags));

        // restore the msg_control pointer if needed
        #[cfg(any(target_os = "macos", target_os = "ios", test))]
        {
            msg.msg_control = msg_control;
        }

        #[cfg(debug_assertions)]
        {
            assert_eq!(
                prev_msg_control_ptr, msg.msg_control,
                "msg_control pointer was modified by the OS"
            );
        }

        let cf = match result {
            Ok(_) => events.on_complete(1),
            Err(err) => events.on_error(err),
        };

        if cf.is_break() {
            return;
        }
    }
}

#[inline]
pub fn recv<'a, Sock: AsRawFd, P: IntoIterator<Item = &'a mut msghdr>, E: SocketEvents>(
    socket: &Sock,
    socket_type: SocketType,
    packets: P,
    events: &mut E,
) {
    let mut flags = match socket_type {
        SocketType::Blocking => Default::default(),
        SocketType::NonBlocking => libc::MSG_DONTWAIT,
    };

    for packet in packets {
        #[cfg(debug_assertions)]
        let prev_msg_control_ptr = packet.msg_control;

        // Safety: calling a libc function is inherently unsafe as rust cannot
        // make any invariant guarantees. This has to be reviewed by humans instead
        // so the [docs](https://linux.die.net/man/2/recmsg) are inlined here:

        // > The argument sockfd is the file descriptor of the receiving socket.
        let sockfd = socket.as_raw_fd();

        // > The recvmsg() call uses a msghdr structure to minimize the number of
        // > directly supplied arguments.
        //
        // > Here msg_name and msg_namelen specify the source address if the
        // > socket is unconnected.
        //
        // > The fields msg_iov and msg_iovlen describe scatter-gather locations
        //
        // > When recvmsg() is called, msg_controllen should contain the length
        // > of the available buffer in msg_control; upon return from a successful
        // > call it will contain the length of the control message sequence.
        let msg = packet;

        // > The flags argument to a recv() call is formed by ORing one or more flags
        //
        // We set MSG_DONTWAIT if it's nonblocking or there is more than one call

        // > recvmsg() calls are used to receive messages from a socket
        //
        // > All three routines return the length of the message on successful completion.
        // > If a message is too long to fit in the supplied buffer, excess bytes may be
        // > discarded depending on the type of socket the message is received from.
        //
        // > These calls return the number of bytes received, or -1 if an error occurred.
        let result = libc!(recvmsg(sockfd, msg, flags));

        #[cfg(debug_assertions)]
        {
            assert_eq!(
                prev_msg_control_ptr, msg.msg_control,
                "msg_control pointer was modified by the OS"
            );
        }

        let cf = match result {
            Ok(payload_len) => {
                // update the message based on the return size of the syscall
                unsafe {
                    msg.set_payload_len(payload_len.min(u16::MAX as _).max(0) as _);
                }
                events.on_complete(1)
            }
            Err(err) => events.on_error(err),
        };

        if cf.is_break() {
            return;
        }

        // don't block the follow-up calls
        flags = libc::MSG_DONTWAIT;
    }
}