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

use crate::{features::gso, message::default as message, socket, syscall};
use s2n_quic_core::{
    endpoint::Endpoint,
    event::{self, EndpointPublisher as _},
    inet::{self, SocketAddress},
    io::event_loop::EventLoop,
    path::MaxMtu,
    time::Clock as ClockTrait,
};
use std::{convert::TryInto, io, io::ErrorKind};
use tokio::runtime::Handle;

mod builder;
mod clock;
pub(crate) mod task;
#[cfg(test)]
mod tests;

pub type PathHandle = message::Handle;
pub use builder::Builder;
pub(crate) use clock::Clock;

#[derive(Debug, Default)]
pub struct Io {
    builder: Builder,
}

impl Io {
    pub fn builder() -> Builder {
        Builder::default()
    }

    pub fn new<A: std::net::ToSocketAddrs>(addr: A) -> io::Result<Self> {
        let address = addr.to_socket_addrs()?.next().expect("missing address");
        let builder = Builder::default().with_receive_address(address)?;
        Ok(Self { builder })
    }

    pub fn start<E: Endpoint<PathHandle = PathHandle>>(
        self,
        mut endpoint: E,
    ) -> io::Result<(tokio::task::JoinHandle<()>, SocketAddress)> {
        let Builder {
            handle,
            rx_socket,
            tx_socket,
            recv_addr,
            send_addr,
            recv_buffer_size,
            send_buffer_size,
            mut max_mtu,
            max_segments,
            reuse_port,
        } = self.builder;

        let clock = Clock::default();

        let mut publisher = event::EndpointPublisherSubscriber::new(
            event::builder::EndpointMeta {
                endpoint_type: E::ENDPOINT_TYPE,
                timestamp: clock.get_time(),
            },
            None,
            endpoint.subscriber(),
        );

        publisher.on_platform_feature_configured(event::builder::PlatformFeatureConfigured {
            configuration: event::builder::PlatformFeatureConfiguration::Gso {
                max_segments: max_segments.into(),
            },
        });

        // try to use the tokio runtime handle if provided, otherwise try to use the implicit tokio
        // runtime in the current scope of the application.
        let handle = if let Some(handle) = handle {
            handle
        } else {
            Handle::try_current().map_err(|err| std::io::Error::new(io::ErrorKind::Other, err))?
        };

        let guard = handle.enter();

        let rx_socket = if let Some(rx_socket) = rx_socket {
            rx_socket
        } else if let Some(recv_addr) = recv_addr {
            syscall::bind_udp(recv_addr, reuse_port)?
        } else {
            return Err(io::Error::new(
                io::ErrorKind::InvalidInput,
                "missing bind address",
            ));
        };

        let rx_addr = convert_addr_to_std(rx_socket.local_addr()?)?;

        let tx_socket = if let Some(tx_socket) = tx_socket {
            tx_socket
        } else if let Some(send_addr) = send_addr {
            syscall::bind_udp(send_addr, reuse_port)?
        } else {
            // No tx_socket or send address was specified, so the tx socket
            // will be a handle to the rx socket.
            rx_socket.try_clone()?
        };

        if let Some(size) = send_buffer_size {
            tx_socket.set_send_buffer_size(size)?;
        }

        if let Some(size) = recv_buffer_size {
            rx_socket.set_recv_buffer_size(size)?;
        }

        // Configure MTU discovery
        if !syscall::configure_mtu_disc(&tx_socket) {
            // disable MTU probing if we can't prevent fragmentation
            max_mtu = MaxMtu::MIN;
        }

        publisher.on_platform_feature_configured(event::builder::PlatformFeatureConfigured {
            configuration: event::builder::PlatformFeatureConfiguration::MaxMtu {
                mtu: max_mtu.into(),
            },
        });

        // Configure the socket with GRO
        let gro_enabled = syscall::configure_gro(&rx_socket);

        publisher.on_platform_feature_configured(event::builder::PlatformFeatureConfigured {
            configuration: event::builder::PlatformFeatureConfiguration::Gro {
                enabled: gro_enabled,
            },
        });

        // Configure packet info CMSG
        syscall::configure_pktinfo(&rx_socket);

        // Configure TOS/ECN
        let tos_enabled = syscall::configure_tos(&rx_socket);

        publisher.on_platform_feature_configured(event::builder::PlatformFeatureConfigured {
            configuration: event::builder::PlatformFeatureConfiguration::Ecn {
                enabled: tos_enabled,
            },
        });

        let rx = {
            // if GRO is enabled, then we need to provide the syscall with the maximum size buffer
            let payload_len = if gro_enabled {
                u16::MAX
            } else {
                max_mtu.into()
            } as u32;

            // 8Mb - TODO make this configurable: https://github.com/aws/s2n-quic/issues/1811
            let rx_buffer_size = 8 * (1 << 20);
            let entries = rx_buffer_size / payload_len;
            let entries = if entries.is_power_of_two() {
                entries
            } else {
                // round up to the nearest power of two, since the ring buffers require it
                entries.next_power_of_two()
            };

            let mut consumers = vec![];

            let (producer, consumer) = socket::ring::pair(entries, payload_len);
            consumers.push(consumer);

            // spawn a task that actually reads from the socket into the ring buffer
            handle.spawn(task::rx(rx_socket, producer));

            // construct the RX side for the endpoint event loop
            let max_mtu = MaxMtu::try_from(payload_len as u16).unwrap();
            let addr: inet::SocketAddress = rx_addr.into();
            socket::io::rx::Rx::new(consumers, max_mtu, addr.into())
        };

        let tx = {
            let gso = crate::features::Gso::from(max_segments);

            // compute the payload size for each message from the number of GSO segments we can
            // fill
            let payload_len = {
                let max_mtu: u16 = max_mtu.into();
                (max_mtu as u32 * gso.max_segments() as u32).min(u16::MAX as u32)
            };

            // 8Mb - TODO make this configurable: https://github.com/aws/s2n-quic/issues/1811
            let tx_buffer_size = 8 * (1 << 20);
            let entries = tx_buffer_size / payload_len;
            let entries = if entries.is_power_of_two() {
                entries
            } else {
                // round up to the nearest power of two, since the ring buffers require it
                entries.next_power_of_two()
            };

            let mut producers = vec![];

            let (producer, consumer) = socket::ring::pair(entries, payload_len);
            producers.push(producer);

            // spawn a task that actually flushes the ring buffer to the socket
            handle.spawn(task::tx(tx_socket, consumer, gso.clone()));

            // construct the TX side for the endpoint event loop
            socket::io::tx::Tx::new(producers, gso, max_mtu)
        };

        // Notify the endpoint of the MTU that we chose
        endpoint.set_max_mtu(max_mtu);

        let task = handle.spawn(
            EventLoop {
                endpoint,
                clock,
                rx,
                tx,
            }
            .start(),
        );

        drop(guard);

        Ok((task, rx_addr.into()))
    }
}

fn convert_addr_to_std(addr: socket2::SockAddr) -> io::Result<std::net::SocketAddr> {
    addr.as_socket()
        .ok_or_else(|| io::Error::new(io::ErrorKind::InvalidInput, "invalid domain for socket"))
}


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

use crate::{
    io::tokio::Clock,
    message::{simple::Message, Message as _},
    socket::{
        io::{rx, tx},
        ring::{self, Consumer, Producer},
    },
};
use core::future::Future;
use futures::future::poll_fn;
use s2n_quic_core::{
    endpoint::Endpoint,
    inet::{self, SocketAddress},
    io::event_loop::{select::Select, EventLoop},
    path::{self, MaxMtu},
};
use std::{convert::TryInto, io, io::ErrorKind};
use tokio::runtime::Handle;
use turmoil::net::UdpSocket;

mod builder;
#[cfg(test)]
mod tests;

pub use builder::Builder;
pub type PathHandle = path::Tuple;

#[derive(Default)]
pub struct Io {
    builder: Builder,
}

impl Io {
    pub fn builder() -> Builder {
        Builder::default()
    }

    pub fn new<A: turmoil::ToSocketAddrs + Send + Sync + 'static>(addr: A) -> io::Result<Self> {
        let builder = Builder::default().with_address(addr)?;
        Ok(Self { builder })
    }

    async fn setup<E: Endpoint<PathHandle = PathHandle>>(
        self,
        mut endpoint: E,
    ) -> io::Result<(impl Future<Output = ()>, SocketAddress)> {
        let Builder {
            handle: _,
            socket,
            addr,
            max_mtu,
        } = self.builder;

        endpoint.set_max_mtu(max_mtu);

        let clock = Clock::default();

        let socket = if let Some(socket) = socket {
            socket
        } else if let Some(addr) = addr {
            UdpSocket::bind(&*addr).await?
        } else {
            return Err(io::Error::new(
                io::ErrorKind::InvalidInput,
                "missing bind address",
            ));
        };

        let local_addr = socket.local_addr()?;
        let local_addr: inet::SocketAddress = local_addr.into();
        let payload_len: usize = max_mtu.into();
        let payload_len = payload_len as u32;

        // This number is somewhat arbitrary but it's a decent number of messages without it consuming
        // large in memory. Eventually, it might be a good idea to expose this value in the
        // builder, but we'll wait until someone asks for it :).
        let entries = 1024;

        let (rx, rx_producer) = {
            let mut consumers = vec![];

            let (producer, consumer) = ring::pair(entries, payload_len);
            consumers.push(consumer);

            let rx = rx::Rx::new(consumers, max_mtu, local_addr.into());

            (rx, producer)
        };

        let (tx, tx_consumer) = {
            let mut producers = vec![];

            let (producer, consumer) = ring::pair(entries, payload_len);
            producers.push(producer);

            let gso = crate::features::Gso::default();

            // GSO is not supported by turmoil so disable it
            gso.disable();

            let tx = tx::Tx::new(producers, gso, max_mtu);

            (tx, consumer)
        };

        // Spawn a task that does the actual socket calls and coordinates with the event loop
        // through the ring buffers
        tokio::spawn(run_io(socket, rx_producer, tx_consumer));

        let event_loop = EventLoop {
            clock,
            rx,
            tx,
            endpoint,
        }
        .start();

        Ok((event_loop, local_addr))
    }

    pub fn start<E: Endpoint<PathHandle = PathHandle>>(
        mut self,
        endpoint: E,
    ) -> io::Result<(tokio::task::JoinHandle<()>, SocketAddress)> {
        let handle = if let Some(handle) = self.builder.handle.take() {
            handle
        } else {
            Handle::try_current().map_err(|err| std::io::Error::new(io::ErrorKind::Other, err))?
        };

        let guard = handle.enter();

        let task = handle.spawn(async move {
            let (instance, _local_addr) = self.setup(endpoint).await.unwrap();

            instance.await;
        });

        drop(guard);

        // TODO this is a potentially async operation - can we get this here?
        let local_addr = Default::default();

        Ok((task, local_addr))
    }
}

/// Turmoil doesn't allow to split sockets for Tx and Rx so we need to spawn a single task to
/// handle both jobs
async fn run_io(
    socket: UdpSocket,
    mut producer: Producer<Message>,
    mut consumer: Consumer<Message>,
) -> io::Result<()> {
    let mut poll_producer = false;

    loop {
        let socket_ready = socket.readable();
        let consumer_ready = poll_fn(|cx| consumer.poll_acquire(u32::MAX, cx));
        let producer_ready = async {
            // Only poll the producer if we need more capacity - otherwise we would constantly wake
            // up
            if poll_producer {
                poll_fn(|cx| producer.poll_acquire(u32::MAX, cx)).await
            } else {
                core::future::pending().await
            }
        };
        // The socket task doesn't have any application wakeups to handle so just make it pending
        let application_wakeup = core::future::pending();

        // We replace the timer future with the `socket_ready` instead, since we don't have a
        // timer here. Other than the application wakeup, Select doesn't really treat any of
        // the futures special.
        let is_readable = Select::new(
            consumer_ready,
            producer_ready,
            application_wakeup,
            socket_ready,
        )
        .await
        .unwrap()
        .timeout_expired;

        if is_readable {
            let mut count = 0;
            for entry in producer.data() {
                // Since UDP sockets are stateless, the only errors we should back is a WouldBlock.
                // If we get any errors, we'll try again later.
                if let Ok((len, addr)) = socket.try_recv_from(entry.payload_mut()) {
                    count += 1;
                    // update the packet information
                    entry.set_remote_address(&(addr.into()));
                    unsafe {
                        entry.set_payload_len(len);
                    }
                } else {
                    break;
                }
            }

            // release the received messages to the consumer
            producer.release(count);

            // only poll the producer if we need entries
            poll_producer = producer.data().is_empty();
        }

        {
            let mut count = 0;
            for entry in consumer.data() {
                let addr = *entry.remote_address();
                let addr: std::net::SocketAddr = addr.into();
                let payload = entry.payload_mut();
                // Since UDP sockets are stateless, the only errors we should back is a WouldBlock.
                // If we get any errors, we'll try again later.
                if socket.try_send_to(payload, addr).is_ok() {
                    count += 1;
                } else {
                    break;
                }
            }

            // release capacity back to the producer
            consumer.release(count);
        }

        // check to see if the rings are open, otherwise we need to shut down the task
        if !(producer.is_open() && consumer.is_open()) {
            return Ok(());
        }
    }
}


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

use crate::io::tokio::Clock;
use s2n_quic_core::{
    endpoint::Endpoint, inet::SocketAddress, io::event_loop::EventLoop, path::MaxMtu,
};

pub use s2n_quic_core::{
    io::rx,
    sync::{spsc, worker},
    xdp::path::Tuple as PathHandle,
};
pub use s2n_quic_xdp::*;

// export the encoder configuration for writing packets
pub mod encoder {
    pub use s2n_quic_core::xdp::encoder::State as Config;
}

// export socket types and helpers
pub mod socket {
    pub use s2n_quic_xdp::socket::*;

    /// Binds a UDP socket to a particular interface and socket address
    pub fn bind_udp(
        interface: &::std::ffi::CStr,
        addr: ::std::net::SocketAddr,
    ) -> ::std::io::Result<::std::net::UdpSocket> {
        let socket = crate::syscall::udp_socket(addr)?;

        // associate the socket with a single interface
        crate::syscall::bind_to_interface(&socket, interface)?;

        socket.bind(&addr.into())?;

        Ok(socket.into())
    }
}

pub mod tx {
    pub use s2n_quic_core::io::tx::*;

    pub fn channel(
        socket: ::std::net::UdpSocket,
    ) -> (
        impl Tx<PathHandle = crate::message::default::Handle>,
        impl core::future::Future<Output = ::std::io::Result<()>>,
    ) {
        // Initial packets don't need to be bigger than the minimum
        let max_mtu = s2n_quic_core::path::MaxMtu::MIN;

        // It's unlikely the initial packets will utilize GSO, so just disable it
        let gso = crate::features::Gso::default();
        gso.disable();

        // compute the payload size for each message from the MaxMtu
        let payload_len = {
            let max_mtu: u16 = max_mtu.into();
            max_mtu as u32
        };

        // 512Kb
        let tx_buffer_size: u32 = 1 << 19;
        let entries = tx_buffer_size / payload_len;
        let entries = if entries.is_power_of_two() {
            entries
        } else {
            // round up to the nearest power of two, since the ring buffers require it
            entries.next_power_of_two()
        };

        let mut producers = vec![];

        let (producer, consumer) = crate::socket::ring::pair(entries, payload_len);
        producers.push(producer);

        // spawn a task that actually flushes the ring buffer to the socket
        let task = crate::io::tokio::task::tx(socket, consumer, gso.clone());

        // construct the TX side for the endpoint event loop
        let io = crate::socket::io::tx::Tx::new(producers, gso, max_mtu);

        (io, task)
    }
}

impl From<PathHandle> for crate::message::msg::Handle {
    #[inline]
    fn from(handle: PathHandle) -> Self {
        let remote_address = handle.remote_address.into();
        let local_address = handle.local_address.into();
        crate::message::msg::Handle {
            remote_address,
            local_address,
        }
    }
}

impl From<&PathHandle> for crate::message::msg::Handle {
    #[inline]
    fn from(handle: &PathHandle) -> Self {
        let remote_address = handle.remote_address.into();
        let local_address = handle.local_address.into();
        crate::message::msg::Handle {
            remote_address,
            local_address,
        }
    }
}

mod builder;
pub use builder::Builder;

pub struct Provider<Rx, Tx> {
    rx: Rx,
    tx: Tx,
    max_mtu: MaxMtu,
    handle: Option<tokio::runtime::Handle>,
}

impl Provider<(), ()> {
    /// Creates a builder to construct an XDP provider
    pub fn builder() -> Builder {
        Builder::default()
    }
}

impl<Rx, Tx> Provider<Rx, Tx>
where
    Rx: 'static + rx::Rx + Send,
    Tx: 'static + tx::Tx<PathHandle = Rx::PathHandle> + Send,
{
    pub fn start<E: Endpoint<PathHandle = Rx::PathHandle>>(
        self,
        mut endpoint: E,
    ) -> std::io::Result<(tokio::task::JoinHandle<()>, SocketAddress)> {
        let Self {
            tx,
            rx,
            max_mtu,
            handle,
        } = self;

        // tell the endpoint what our MTU is
        endpoint.set_max_mtu(max_mtu);

        // create a tokio clock
        let clock = Clock::new();

        // create an event loop
        let event_loop = EventLoop {
            endpoint,
            clock,
            rx,
            tx,
        };

        // spawn the event loop on to the tokio handle
        let task = if let Some(handle) = handle {
            handle.spawn(event_loop.start())
        } else {
            tokio::spawn(event_loop.start())
        };

        Ok((task, SocketAddress::default()))
    }
}


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

use super::network::{Buffers, Network, Packet};
use core::time::Duration;
use s2n_quic_core::{havoc, path::MaxMtu};
use std::{
    borrow::Cow,
    sync::{
        atomic::{AtomicU16, AtomicU64, Ordering},
        Arc, Mutex,
    },
};

#[derive(Clone, Default)]
pub struct TxRecorder {
    packets: Arc<Mutex<Vec<Packet>>>,
}

impl TxRecorder {
    pub fn get_packets(&self) -> Arc<Mutex<Vec<Packet>>> {
        self.packets.clone()
    }
}

impl Network for TxRecorder {
    fn execute(&mut self, buffers: &Buffers) -> usize {
        let mut packets = self.packets.lock().unwrap();
        buffers.pending_transmission(|packet| {
            packets.push(packet.clone());
        });
        0
    }
}

#[derive(Clone, Default)]
pub struct Model(Arc<State>);

impl Model {
    pub fn jitter(&self) -> Duration {
        Duration::from_micros(self.0.jitter.load(Ordering::SeqCst))
    }

    /// The amount of time between sending packets
    ///
    /// Setting this value to `0` will transmit all allowed packets at the exact same time.
    pub fn set_jitter(&self, value: Duration) -> &Self {
        self.0
            .jitter
            .store(value.as_micros() as _, Ordering::SeqCst);
        self
    }

    pub fn network_jitter(&self) -> Duration {
        Duration::from_micros(self.0.network_jitter.load(Ordering::SeqCst))
    }

    /// The amount of jitter in the network itself
    ///
    /// Setting this value to `>0` will cause packets to be reordered.
    pub fn set_network_jitter(&self, value: Duration) -> &Self {
        self.0
            .network_jitter
            .store(value.as_micros() as _, Ordering::SeqCst);
        self
    }

    pub fn delay(&self) -> Duration {
        Duration::from_micros(self.0.delay.load(Ordering::SeqCst))
    }

    /// The amount of time a packet is delayed before the receiver is able to read it
    pub fn set_delay(&self, value: Duration) -> &Self {
        self.0.delay.store(value.as_micros() as _, Ordering::SeqCst);
        self
    }

    pub fn transmit_rate(&self) -> u64 {
        self.0.transmit_rate.load(Ordering::SeqCst)
    }

    /// The number of packets that can be transmitted in a single round.
    ///
    /// By default, all packet buffers will be cleared on every round.
    pub fn set_transmit_rate(&self, value: u64) -> &Self {
        self.0.transmit_rate.store(value, Ordering::SeqCst);
        self
    }

    fn retransmit_rate(&self) -> u64 {
        self.0.retransmit_rate.load(Ordering::SeqCst)
    }

    /// The odds a packet will be retransmitted.
    ///
    /// Each packet will make an independent decision with odds of `0.0..1.0`, with `0.0` having no
    /// chance and `1.0` occurring with each packet.
    pub fn set_retransmit_rate(&self, value: f64) -> &Self {
        let value = rate_to_u64(value);
        self.0.retransmit_rate.store(value, Ordering::SeqCst);
        self
    }

    fn corrupt_rate(&self) -> u64 {
        self.0.corrupt_rate.load(Ordering::SeqCst)
    }

    /// The odds a packet will be corrupted.
    ///
    /// Each packet will make an independent decision with odds of `0.0..1.0`, with `0.0` having no
    /// chance and `1.0` occurring with each packet.
    pub fn set_corrupt_rate(&self, value: f64) -> &Self {
        let value = rate_to_u64(value);
        self.0.corrupt_rate.store(value, Ordering::SeqCst);
        self
    }

    fn drop_rate(&self) -> u64 {
        self.0.drop_rate.load(Ordering::SeqCst)
    }

    /// The odds a packet will be dropped.
    ///
    /// Each packet will make an independent decision with odds of `0.0..1.0`, with `0.0` having no
    /// chance and `1.0` occurring with each packet.
    pub fn set_drop_rate(&self, value: f64) -> &Self {
        let value = rate_to_u64(value);
        self.0.drop_rate.store(value, Ordering::SeqCst);
        self
    }

    pub fn max_udp_payload(&self) -> u16 {
        self.0.max_udp_payload.load(Ordering::SeqCst)
    }

    /// The maximum payload size for the network
    pub fn set_max_udp_payload(&self, value: u16) -> &Self {
        self.0.max_udp_payload.store(value, Ordering::SeqCst);
        self
    }

    /// The number of inflight packets
    fn inflight(&self) -> u64 {
        self.0.current_inflight.load(Ordering::SeqCst)
    }

    pub fn max_inflight(&self) -> u64 {
        self.0.max_inflight.load(Ordering::SeqCst)
    }

    /// Sets the maximum number of packets that can be inflight for the network
    ///
    /// Any packets that exceed this amount will be dropped
    pub fn set_max_inflight(&self, value: u64) -> &Self {
        self.0.max_inflight.store(value, Ordering::SeqCst);
        self
    }

    pub fn inflight_delay(&self) -> Duration {
        Duration::from_micros(self.0.inflight_delay.load(Ordering::SeqCst))
    }

    /// Sets the delay for each packet above the inflight_delay
    pub fn set_inflight_delay(&self, value: Duration) -> &Self {
        self.0
            .inflight_delay
            .store(value.as_micros() as _, Ordering::SeqCst);
        self
    }

    pub fn inflight_delay_threshold(&self) -> u64 {
        self.0.inflight_delay_threshold.load(Ordering::SeqCst)
    }

    /// Sets the delay for each packet above the inflight_delay_threshold
    pub fn set_inflight_delay_threshold(&self, value: u64) -> &Self {
        self.0
            .inflight_delay_threshold
            .store(value, Ordering::SeqCst);
        self
    }
}

fn rate_to_u64(rate: f64) -> u64 {
    let value = rate.max(0.0).min(1.0);
    let value = value * u64::MAX as f64;
    value.round() as u64
}

struct State {
    delay: AtomicU64,
    jitter: AtomicU64,
    network_jitter: AtomicU64,
    transmit_rate: AtomicU64,
    retransmit_rate: AtomicU64,
    corrupt_rate: AtomicU64,
    drop_rate: AtomicU64,
    max_udp_payload: AtomicU16,
    max_inflight: AtomicU64,
    inflight_delay: AtomicU64,
    inflight_delay_threshold: AtomicU64,
    current_inflight: AtomicU64,
}

impl Default for State {
    fn default() -> Self {
        Self {
            delay: AtomicU64::new(Duration::from_millis(50).as_micros() as _),
            jitter: AtomicU64::new(0),
            network_jitter: AtomicU64::new(0),
            transmit_rate: AtomicU64::new(u64::MAX),
            retransmit_rate: AtomicU64::new(0),
            corrupt_rate: AtomicU64::new(0),
            drop_rate: AtomicU64::new(0),
            max_udp_payload: AtomicU16::new(MaxMtu::default().into()),
            max_inflight: AtomicU64::new(u64::MAX),
            inflight_delay: AtomicU64::new(0),
            inflight_delay_threshold: AtomicU64::new(u64::MAX),
            current_inflight: AtomicU64::new(0),
        }
    }
}

impl Network for Model {
    fn execute(&mut self, buffers: &Buffers) -> usize {
        let jitter = self.jitter();
        let network_jitter = self.network_jitter();
        let transmit_rate = self.transmit_rate();
        let retransmit_rate = self.retransmit_rate();
        let corrupt_rate = self.corrupt_rate();
        let drop_rate = self.drop_rate();
        let max_udp_payload = self.max_udp_payload() as usize;
        let inflight_delay = self.inflight_delay();
        let inflight_delay_threshold = self.inflight_delay_threshold();

        let now = super::time::now();
        let mut transmit_time = now + self.delay();
        let transmit_time = &mut transmit_time;

        #[inline]
        fn gen_rate(rate: u64) -> bool {
            // ensure the rate isn't 0 before actually generating a random number
            rate > 0 && super::rand::gen::<u64>() < rate
        }

        let mut transmit = |packet: Cow<Packet>| {
            // drop the packet if it's over the current MTU
            if packet.payload.len() > max_udp_payload {
                return 0;
            }

            // drop packets that exceed the maximum number of inflight packets for the network
            if self.inflight() >= self.max_inflight() {
                return 0;
            }

            // drop the packet if enabled
            if gen_rate(drop_rate) {
                return 0;
            }

            let mut packet = packet.into_owned();

            if !packet.payload.is_empty() && gen_rate(corrupt_rate) {
                use havoc::Strategy as _;

                let new_len = havoc::Truncate
                    .randomly()
                    .and_then(havoc::Swap.repeat(0..packet.payload.len()).randomly())
                    .and_then(havoc::Mutate.repeat(0..packet.payload.len()).randomly())
                    .havoc_slice(&mut super::rand::Havoc, &mut packet.payload);

                // if the len was changed, then update it
                if new_len != packet.payload.len() {
                    packet.payload.truncate(new_len);
                }
            }

            if !jitter.is_zero() {
                // add a delay for the next packet to be transmitted
                *transmit_time += gen_jitter(jitter);
            }

            // copy the transmit time for this packet
            let mut transmit_time = *transmit_time;

            if !network_jitter.is_zero() {
                transmit_time += gen_jitter(network_jitter);
            }

            let model = self.clone();
            let current_inflight = model.0.current_inflight.fetch_add(1, Ordering::SeqCst);

            // scale the inflight delay by the number above the delay threshold
            if let Some(mul) = current_inflight.checked_sub(inflight_delay_threshold) {
                transmit_time += inflight_delay * mul as u32;
            }

            // reverse the addresses so the dst/src are correct for the receiver
            packet.switch();

            let buffers = buffers.clone();

            // spawn a task that will push the packet onto the receiver queue at the transit time
            super::spawn(async move {
                // if the packet isn't scheduled to transmit immediately, wait until the computed
                // time
                if now != transmit_time {
                    super::time::delay_until(transmit_time).await;
                }

                buffers.rx(*packet.path.local_address, |queue| {
                    model.0.current_inflight.fetch_sub(1, Ordering::SeqCst);
                    queue.receive(packet);
                });
            });

            1
        };

        let mut transmission_count = 0;
        buffers.drain_pending_transmissions(|packet| {
            // retransmit the packet until the rate fails or we retransmit 5
            //
            // We limit retransmissions to 5 just so we don't endlessly iterate when the
            // `retransmit_rate` is high. This _should_ be high enough where we're getting
            // retransmission coverage without needlessly saturating the network.
            let mut count = 0;
            while count < 5 && gen_rate(retransmit_rate) {
                transmission_count += transmit(Cow::Borrowed(&packet));
                count += 1;
            }

            transmission_count += transmit(Cow::Owned(packet));

            // continue transmitting as long as we are under the rate
            if transmission_count < transmit_rate {
                Ok(())
            } else {
                Err(())
            }
        });

        transmission_count as usize
    }
}

fn gen_jitter(max_jitter: Duration) -> Duration {
    let micros = super::rand::gen_range(0..max_jitter.as_micros() as u64);
    let micros = micros as f64;
    // even though we're generated micros, we round to the nearest millisecond
    // so packets can be grouped together
    let millis = micros / 1000.0;
    let millis = f64::round(millis) as u64;
    Duration::from_millis(millis)
}


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

use core::task::{Context, Poll, Waker};
use s2n_quic_core::{
    event,
    inet::{datagram, ExplicitCongestionNotification, SocketAddress},
    io::{
        self, rx,
        tx::{self, Queue as _},
    },
    path::{LocalAddress, Tuple},
};
use std::{
    collections::{HashMap, VecDeque},
    sync::{
        atomic::{AtomicU16, AtomicU32, Ordering},
        Arc, Mutex,
    },
};

// This constant is used to size the buffer for packet payloads
// we use 10_000 since there are unit tests for jumbo frames, which
// have MTU's up to approximately 9_001
const MAX_TESTED_MTU: u16 = 10_000;

pub type PathHandle = Tuple;

pub trait Network {
    fn execute(&mut self, buffers: &Buffers) -> usize;
}

impl<A: Network, B: Network> Network for (A, B) {
    fn execute(&mut self, buffers: &Buffers) -> usize {
        let mut result = 0;
        result += self.0.execute(buffers);
        result += self.1.execute(buffers);
        result
    }
}

#[derive(Clone, Debug)]
pub struct Buffers {
    inner: Arc<Mutex<State>>,
    next_ip: Arc<AtomicU32>,
    next_port: Arc<AtomicU16>,
}

impl Default for Buffers {
    fn default() -> Self {
        Self {
            inner: Default::default(),
            next_ip: Arc::new(AtomicU32::new(u32::from_be_bytes([1, 0, 0, 0]))),
            //= https://www.rfc-editor.org/rfc/rfc6335#section-6
            //# o  the Dynamic Ports, also known as the Private or Ephemeral Ports,
            //#    from 49152-65535 (never assigned)
            next_port: Arc::new(AtomicU16::new(49152)),
        }
    }
}

impl Buffers {
    pub fn close(&self) {
        let mut lock = self.inner.lock().unwrap();
        lock.is_open = false;

        let state = &mut *lock;

        for entry in state.tx.values_mut().chain(state.rx.values_mut()) {
            if let Some(waker) = entry.waker.take() {
                waker.wake();
            }
        }
    }

    pub fn tx<F: FnOnce(&mut Queue)>(&self, handle: SocketAddress, f: F) {
        let mut lock = self.inner.lock().unwrap();
        if let Some(queue) = lock.tx.get_mut(&handle) {
            f(queue)
        }
    }

    pub fn rx<F: FnOnce(&mut Queue)>(&self, handle: SocketAddress, f: F) {
        let mut lock = self.inner.lock().unwrap();
        if let Some(queue) = lock.rx.get_mut(&handle) {
            f(queue)
        }
    }

    pub fn pending_transmission<F: FnMut(&Packet)>(&self, mut f: F) {
        let lock = self.inner.lock().unwrap();
        for queue in lock.tx.values() {
            for packet in &queue.packets {
                f(packet);
            }
        }
    }

    pub fn drain_pending_transmissions<F: FnMut(Packet) -> Result<(), ()>>(&self, mut f: F) {
        let mut lock = self.inner.lock().unwrap();

        let mut queues = vec![];

        // find all of the queues with at least one packet to transmit
        for queue in lock.tx.values_mut() {
            if queue.packets.is_empty() {
                continue;
            }

            queues.push(queue);
        }

        // shuffle the queue so each endpoint has a fair chance of transmitting
        super::rand::shuffle(&mut queues);

        loop {
            let mut has_result = false;
            for queue in &mut queues {
                // transmit a single packet at a time per queue so they are fairly
                // transmitted
                if let Some(packet) = queue.packets.pop_front() {
                    let result = f(packet);
                    has_result = true;

                    // notify the endpoint that it can send now
                    if let Some(waker) = queue.waker.take() {
                        waker.wake();
                    }

                    if result.is_err() {
                        return;
                    }
                }
            }

            // if all of the queues are empty then just return
            if !has_result {
                return;
            }
        }
    }

    pub fn execute<N: Network>(&self, n: &mut N) {
        n.execute(self);
    }

    /// Generate a unique address
    pub fn generate_addr(&self) -> SocketAddress {
        let ip = self
            .next_ip
            .fetch_add(1, Ordering::SeqCst)
            .to_be_bytes()
            .into();
        let port = self.next_port.fetch_add(1, Ordering::SeqCst);
        let addr = (ip, port);
        SocketAddress::IpV4(addr.into())
    }

    /// Register an address on the network
    pub fn register(&self, handle: SocketAddress) -> (TxIo, RxIo) {
        let mut lock = self.inner.lock().unwrap();

        let queue = Queue::new(handle);

        lock.tx.insert(handle, queue.clone());
        lock.rx.insert(handle, queue);

        let tx = TxIo {
            buffers: self.clone(),
            handle,
        };
        let rx = RxIo {
            buffers: self.clone(),
            handle,
        };

        (tx, rx)
    }
}

pub struct TxIo {
    buffers: Buffers,
    handle: SocketAddress,
}

impl tx::Tx for TxIo {
    type PathHandle = Tuple;
    type Queue = Queue;
    type Error = ();

    fn poll_ready(&mut self, cx: &mut Context) -> Poll<Result<(), Self::Error>> {
        let mut lock = self.buffers.inner.lock().unwrap();

        if !lock.is_open {
            return Err(()).into();
        }

        let tx = lock.tx.get_mut(&self.handle).unwrap();

        // If we weren't previously full, then return pending so we don't spin
        if !tx.is_blocked {
            return Poll::Pending;
        }

        // if we were blocked and now have capacity wake up the endpoint
        if tx.has_capacity() {
            tx.is_blocked = false;
            Poll::Ready(Ok(()))
        } else {
            tx.waker = Some(cx.waker().clone());
            Poll::Pending
        }
    }

    fn queue<F: FnOnce(&mut Self::Queue)>(&mut self, f: F) {
        self.buffers.tx(self.handle, f);
    }

    fn handle_error<E: event::EndpointPublisher>(self, _error: Self::Error, _events: &mut E) {}
}

pub struct RxIo {
    buffers: Buffers,
    handle: SocketAddress,
}

impl rx::Rx for RxIo {
    type PathHandle = Tuple;
    type Queue = Queue;
    type Error = ();

    fn poll_ready(&mut self, cx: &mut Context) -> Poll<Result<(), Self::Error>> {
        let mut lock = self.buffers.inner.lock().unwrap();

        if !lock.is_open {
            return Err(()).into();
        }

        let rx = lock.rx.get_mut(&self.handle).unwrap();

        // wake up the endpoint if we have an rx message
        if !io::rx::Queue::is_empty(rx) {
            return Poll::Ready(Ok(()));
        }

        // store the waker for later notifications
        rx.waker = Some(cx.waker().clone());
        Poll::Pending
    }

    fn queue<F: FnOnce(&mut Self::Queue)>(&mut self, f: F) {
        self.buffers.rx(self.handle, f);
    }

    fn handle_error<E: event::EndpointPublisher>(self, _error: Self::Error, _events: &mut E) {}
}

#[derive(Debug)]
struct State {
    is_open: bool,
    tx: HashMap<SocketAddress, Queue>,
    rx: HashMap<SocketAddress, Queue>,
}

impl Default for State {
    fn default() -> Self {
        Self {
            is_open: true,
            tx: Default::default(),
            rx: Default::default(),
        }
    }
}

#[derive(Clone, Debug)]
pub struct Queue {
    capacity: usize,
    mtu: u16,
    packets: VecDeque<Packet>,
    pending: Packet,
    local_address: LocalAddress,
    is_blocked: bool,
    waker: Option<Waker>,
}

impl Queue {
    fn new(addr: SocketAddress) -> Self {
        let mtu = MAX_TESTED_MTU;
        let local_address = addr.into();
        Self {
            capacity: 1024,
            mtu,
            packets: VecDeque::new(),
            pending: Packet::new(mtu, local_address),
            local_address,
            is_blocked: false,
            waker: None,
        }
    }

    pub fn receive(&mut self, packet: Packet) {
        if self.packets.len() == self.capacity {
            // drop old packets if we're at capacity
            let _ = self.packets.pop_front();
        }

        self.packets.push_back(packet);

        if let Some(w) = self.waker.take() {
            w.wake();
        }
    }

    pub fn take(&mut self, count: usize) -> impl Iterator<Item = Packet> + '_ {
        let count = self.packets.len().min(count);
        self.packets.drain(..count)
    }

    pub fn drain(&mut self) -> impl Iterator<Item = Packet> + '_ {
        self.packets.drain(..)
    }
}

impl io::tx::Queue for Queue {
    type Handle = Tuple;

    const SUPPORTS_ECN: bool = true;

    fn push<M: io::tx::Message<Handle = Self::Handle>>(
        &mut self,
        message: M,
    ) -> Result<io::tx::Outcome, io::tx::Error> {
        if !self.has_capacity() {
            self.is_blocked = true;
            return Err(io::tx::Error::AtCapacity);
        }

        let len = self.pending.write(message)?;

        // create a packet for the next transmission
        let next = Packet::new(self.mtu, self.local_address);
        let packet = core::mem::replace(&mut self.pending, next);

        self.packets.push_back(packet);

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

    fn capacity(&self) -> usize {
        self.capacity - self.packets.len()
    }
}

impl io::rx::Queue for Queue {
    type Handle = Tuple;

    #[inline]
    fn for_each<F: FnMut(datagram::Header<Self::Handle>, &mut [u8])>(&mut self, mut on_packet: F) {
        for mut packet in self.packets.drain(..) {
            let header = datagram::Header {
                path: packet.path,
                ecn: packet.ecn,
            };
            let payload = &mut packet.payload;
            on_packet(header, payload);
        }
    }

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

#[derive(Clone, Debug)]
pub struct Packet {
    pub path: Tuple,
    pub ecn: ExplicitCongestionNotification,
    pub payload: Vec<u8>,
}

impl Packet {
    fn new(mtu: u16, local_address: LocalAddress) -> Self {
        Self {
            path: Tuple {
                local_address,
                remote_address: Default::default(),
            },
            ecn: Default::default(),
            payload: vec![0u8; mtu as usize],
        }
    }

    pub fn switch(&mut self) {
        let path = self.path;
        let remote_address = path.local_address.0.into();
        let local_address = path.remote_address.0.into();
        self.path = Tuple {
            remote_address,
            local_address,
        };
    }

    fn write<M: tx::Message<Handle = Tuple>>(
        &mut self,
        mut message: M,
    ) -> Result<usize, tx::Error> {
        let buffer = tx::PayloadBuffer::new(&mut self.payload);

        let len = message.write_payload(buffer, 0)?;

        self.payload.truncate(len);
        self.path.remote_address = message.path_handle().remote_address;
        self.ecn = message.ecn();

        Ok(len)
    }
}

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

    #[test]
    fn address_generator() {
        let buffers = Buffers::default();

        let mut addrs = vec![];
        for _ in 0..10 {
            addrs.push(buffers.generate_addr());
        }

        if !cfg!(miri) {
            assert_debug_snapshot!(addrs);
        }
    }
}


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

use bach::time::{self, scheduler};
use core::{
    future::Future,
    pin::Pin,
    task::{Context, Poll},
    time::Duration,
};
use futures::ready;
use s2n_quic_core::time::{clock, Timestamp};

pub fn now() -> Timestamp {
    unsafe { Timestamp::from_duration(time::now()) }
}

pub fn delay(duration: Duration) -> Timer {
    Timer::new(now() + duration, duration)
}

pub fn delay_until(deadline: Timestamp) -> Timer {
    let delay = deadline.saturating_duration_since(now());
    Timer::new(deadline, delay)
}

#[derive(Debug, Default)]
pub(crate) struct Clock(());

impl clock::Clock for Clock {
    fn get_time(&self) -> Timestamp {
        now()
    }
}

impl clock::ClockWithTimer for Clock {
    type Timer = Timer;

    fn timer(&self) -> Timer {
        Default::default()
    }
}

#[must_use = "futures do nothing unless you `.await` or poll them"]
pub struct Timer {
    timer: scheduler::Timer,
    deadline: Option<Timestamp>,
}

impl Default for Timer {
    fn default() -> Self {
        let timer = time::delay(Duration::ZERO);

        Self {
            timer,
            deadline: None,
        }
    }
}

impl Timer {
    fn new(deadline: Timestamp, delay: Duration) -> Self {
        Self {
            timer: time::delay(delay),
            deadline: Some(deadline),
        }
    }

    pub fn cancel(&mut self) {
        self.deadline = None;
        self.timer.cancel()
    }
}

impl clock::Timer for Timer {
    fn poll_ready(&mut self, cx: &mut Context) -> Poll<()> {
        if self.deadline.is_none() {
            return Poll::Pending;
        }

        ready!(Pin::new(&mut self.timer).poll(cx));

        self.deadline = None;
        Poll::Ready(())
    }

    fn update(&mut self, deadline: Timestamp) {
        if self.deadline != Some(deadline) {
            self.cancel();
            *self = delay_until(deadline);
        }
    }
}

impl Future for Timer {
    type Output = ();

    fn poll(mut self: Pin<&mut Self>, cx: &mut Context) -> Poll<()> {
        clock::Timer::poll_ready(&mut *self, cx)
    }
}


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

use super::*;

#[derive(Debug, Default)]
pub struct Builder {
    pub(super) handle: Option<Handle>,
    pub(super) rx_socket: Option<socket2::Socket>,
    pub(super) tx_socket: Option<socket2::Socket>,
    pub(super) recv_addr: Option<std::net::SocketAddr>,
    pub(super) send_addr: Option<std::net::SocketAddr>,
    pub(super) recv_buffer_size: Option<usize>,
    pub(super) send_buffer_size: Option<usize>,
    pub(super) max_mtu: MaxMtu,
    pub(super) max_segments: gso::MaxSegments,
    pub(super) reuse_port: bool,
}

impl Builder {
    #[must_use]
    pub fn with_handle(mut self, handle: Handle) -> Self {
        self.handle = Some(handle);
        self
    }

    /// Sets the local address for the runtime to listen on. If no send address
    /// or tx socket is specified, this address will also be used for transmitting from.
    ///
    /// NOTE: this method is mutually exclusive with `with_rx_socket`
    pub fn with_receive_address(mut self, addr: std::net::SocketAddr) -> io::Result<Self> {
        debug_assert!(self.rx_socket.is_none(), "rx socket has already been set");
        self.recv_addr = Some(addr);
        Ok(self)
    }

    /// Sets the local address for the runtime to transmit from. If no send address
    /// or tx socket is specified, the receive_address will be used for transmitting.
    ///
    /// NOTE: this method is mutually exclusive with `with_tx_socket`
    pub fn with_send_address(mut self, addr: std::net::SocketAddr) -> io::Result<Self> {
        debug_assert!(self.tx_socket.is_none(), "tx socket has already been set");
        self.send_addr = Some(addr);
        Ok(self)
    }

    /// Sets the socket used for receiving for the runtime. If no tx_socket or send address is
    /// specified, this socket will be used for transmitting.
    ///
    /// NOTE: this method is mutually exclusive with `with_receive_address`
    pub fn with_rx_socket(mut self, socket: std::net::UdpSocket) -> io::Result<Self> {
        debug_assert!(
            self.recv_addr.is_none(),
            "recv address has already been set"
        );
        self.rx_socket = Some(socket.into());
        Ok(self)
    }

    /// Sets the socket used for transmitting on for the runtime. If no tx_socket or send address is
    /// specified, the rx_socket will be used for transmitting.
    ///
    /// NOTE: this method is mutually exclusive with `with_send_address`
    pub fn with_tx_socket(mut self, socket: std::net::UdpSocket) -> io::Result<Self> {
        debug_assert!(
            self.send_addr.is_none(),
            "send address has already been set"
        );
        self.tx_socket = Some(socket.into());
        Ok(self)
    }

    /// Sets the size of the operating systemâ€™s send buffer associated with the tx socket
    pub fn with_send_buffer_size(mut self, send_buffer_size: usize) -> io::Result<Self> {
        self.send_buffer_size = Some(send_buffer_size);
        Ok(self)
    }

    /// Sets the size of the operating systemâ€™s receive buffer associated with the rx socket
    pub fn with_recv_buffer_size(mut self, recv_buffer_size: usize) -> io::Result<Self> {
        self.recv_buffer_size = Some(recv_buffer_size);
        Ok(self)
    }

    /// Sets the largest maximum transmission unit (MTU) that can be sent on a path
    pub fn with_max_mtu(mut self, max_mtu: u16) -> io::Result<Self> {
        self.max_mtu = max_mtu
            .try_into()
            .map_err(|err| io::Error::new(ErrorKind::InvalidInput, format!("{err}")))?;
        Ok(self)
    }

    /// Disables Generic Segmentation Offload (GSO)
    ///
    /// By default, GSO will be used unless the platform does not support it or an attempt to use
    /// GSO fails. If it is known that GSO is not available, set this option to explicitly disable it.
    pub fn with_gso_disabled(mut self) -> io::Result<Self> {
        self.max_segments = 1.try_into().expect("1 is always a valid MaxSegments value");
        Ok(self)
    }

    /// Enables the port reuse (SO_REUSEPORT) socket option
    pub fn with_reuse_port(mut self) -> io::Result<Self> {
        if !cfg!(unix) {
            return Err(io::Error::new(
                io::ErrorKind::InvalidInput,
                "reuse_port is not supported on the current platform",
            ));
        }
        self.reuse_port = true;
        Ok(self)
    }

    pub fn build(self) -> io::Result<Io> {
        Ok(Io { builder: self })
    }
}


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

use core::{
    future::Future,
    pin::Pin,
    task::{Context, Poll},
    time::Duration,
};
use s2n_quic_core::time::{self, Timestamp};
use tokio::time::{sleep_until, Instant, Sleep};

#[derive(Clone, Debug)]
pub struct Clock(Instant);

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

impl Clock {
    pub fn new() -> Self {
        Self(Instant::now())
    }
}

impl time::Clock for Clock {
    #[inline]
    fn get_time(&self) -> time::Timestamp {
        let duration = self.0.elapsed();
        unsafe {
            // Safety: time duration is only derived from a single `Instant`
            time::Timestamp::from_duration(duration)
        }
    }
}

impl time::ClockWithTimer for Clock {
    type Timer = Timer;

    #[inline]
    fn timer(&self) -> Timer {
        Timer::new(self.clone())
    }
}

#[derive(Debug)]
pub struct Timer {
    /// A reference to the current clock
    clock: Clock,
    /// The `Instant` at which the timer should expire
    target: Option<Instant>,
    /// The handle to the timer entry in the tokio runtime
    sleep: Pin<Box<Sleep>>,
}

impl Timer {
    fn new(clock: Clock) -> Self {
        /// We can't create a timer without first arming it to something, so just set it to 1s in
        /// the future.
        const INITIAL_TIMEOUT: Duration = Duration::from_secs(1);

        let target = clock.0 + INITIAL_TIMEOUT;
        let sleep = Box::pin(sleep_until(target));
        Self {
            clock,
            target: Some(target),
            sleep,
        }
    }
}

impl time::clock::Timer for Timer {
    #[inline]
    fn poll_ready(&mut self, cx: &mut Context) -> Poll<()> {
        // Only poll the inner timer if we have a target set
        if self.target.is_none() {
            return Poll::Pending;
        }

        let res = self.sleep.as_mut().poll(cx);

        if res.is_ready() {
            // clear the target after it fires, otherwise we'll endlessly wake up the task
            self.target = None;
        }

        res
    }

    #[inline]
    fn update(&mut self, timestamp: Timestamp) {
        let delay = unsafe {
            // Safety: the same clock epoch is being used
            timestamp.as_duration()
        };

        // floor the delay to milliseconds to reduce timer churn
        let delay = Duration::from_millis(delay.as_millis() as u64);

        // add the delay to the clock's epoch
        let next_time = self.clock.0 + delay;

        // If the target hasn't changed then don't do anything
        if Some(next_time) == self.target {
            return;
        }

        // if the clock has changed let the sleep future know
        self.sleep.as_mut().reset(next_time);
        self.target = Some(next_time);
    }
}


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

// depending on the platform, some of these implementations aren't used
#![allow(dead_code)]

mod simple;
#[cfg(unix)]
mod unix;

cfg_if::cfg_if! {
    if #[cfg(s2n_quic_platform_socket_mmsg)] {
        pub use mmsg::{rx, tx};
    } else if #[cfg(s2n_quic_platform_socket_msg)] {
        pub use msg::{rx, tx};
    } else {
        pub use simple::{rx, tx};
    }
}

macro_rules! libc_msg {
    ($message:ident, $cfg:ident) => {
        #[cfg($cfg)]
        mod $message {
            use super::unix;
            use crate::{features::Gso, message::$message::Message, socket::ring};

            pub async fn rx<S: Into<std::net::UdpSocket>>(
                socket: S,
                producer: ring::Producer<Message>,
            ) -> std::io::Result<()> {
                unix::rx(socket, producer).await
            }

            pub async fn tx<S: Into<std::net::UdpSocket>>(
                socket: S,
                consumer: ring::Consumer<Message>,
                gso: Gso,
            ) -> std::io::Result<()> {
                unix::tx(socket, consumer, gso).await
            }
        }
    };
}

libc_msg!(msg, s2n_quic_platform_socket_msg);
libc_msg!(mmsg, s2n_quic_platform_socket_mmsg);


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

use super::*;
use core::{
    convert::TryInto,
    task::{Context, Poll},
};
use s2n_quic_core::{
    endpoint::{self, CloseError},
    event,
    inet::ExplicitCongestionNotification,
    io::{rx, tx},
    path::Handle as _,
    time::{Clock, Duration, Timestamp},
};
use std::{collections::BTreeMap, net::ToSocketAddrs};

struct TestEndpoint<const IS_SERVER: bool> {
    handle: PathHandle,
    messages: BTreeMap<u32, Option<Timestamp>>,
    now: Option<Timestamp>,
    subscriber: NoopSubscriber,
}

impl<const IS_SERVER: bool> TestEndpoint<IS_SERVER> {
    fn new(handle: PathHandle) -> Self {
        let messages = if IS_SERVER { 0 } else { 1000 };
        let messages = (0..messages).map(|id| (id, None)).collect();
        Self {
            handle,
            messages,
            now: None,
            subscriber: Default::default(),
        }
    }
}

#[derive(Debug, Default)]
struct NoopSubscriber;

impl event::Subscriber for NoopSubscriber {
    type ConnectionContext = ();

    fn create_connection_context(
        &mut self,
        _meta: &event::api::ConnectionMeta,
        _info: &event::api::ConnectionInfo,
    ) -> Self::ConnectionContext {
    }
}

impl<const IS_SERVER: bool> Endpoint for TestEndpoint<IS_SERVER> {
    type PathHandle = PathHandle;
    type Subscriber = NoopSubscriber;

    const ENDPOINT_TYPE: endpoint::Type = if IS_SERVER {
        endpoint::Type::Server
    } else {
        endpoint::Type::Client
    };

    fn transmit<Tx: tx::Queue<Handle = PathHandle>, C: Clock>(
        &mut self,
        queue: &mut Tx,
        clock: &C,
    ) {
        let now = clock.get_time();
        self.now = Some(now);

        for (id, tx_time) in &mut self.messages {
            match tx_time {
                Some(time) if now.saturating_duration_since(*time) < Duration::from_millis(50) => {
                    continue
                }
                _ => {
                    let payload = id.to_be_bytes();
                    let addr = self.handle;
                    let ecn = ExplicitCongestionNotification::Ect0;
                    let msg = (addr, ecn, payload);
                    if queue.push(msg).is_ok() {
                        *tx_time = Some(now);
                    } else {
                        // no more capacity
                        return;
                    }
                }
            }
        }
    }

    fn receive<Rx: rx::Queue<Handle = PathHandle>, C: Clock>(&mut self, queue: &mut Rx, clock: &C) {
        let now = clock.get_time();
        self.now = Some(now);

        queue.for_each(|_header, payload| {
            // we should only be receiving u32 values
            if payload.len() != 4 {
                return;
            }

            let id = (&*payload).try_into().unwrap();
            let id = u32::from_be_bytes(id);

            if IS_SERVER {
                self.messages.insert(id, None);
            } else {
                self.messages.remove(&id);
            }
        });
    }

    fn poll_wakeups<C: Clock>(
        &mut self,
        _cx: &mut Context<'_>,
        clock: &C,
    ) -> Poll<Result<usize, CloseError>> {
        let now = clock.get_time();
        self.now = Some(now);

        if !IS_SERVER && self.messages.is_empty() {
            return Err(CloseError).into();
        }

        Poll::Pending
    }

    fn timeout(&self) -> Option<Timestamp> {
        self.now.map(|now| now + Duration::from_millis(50))
    }

    fn set_max_mtu(&mut self, _max_mtu: MaxMtu) {
        // noop
    }

    fn subscriber(&mut self) -> &mut Self::Subscriber {
        &mut self.subscriber
    }
}

async fn runtime<A: ToSocketAddrs>(
    receive_addr: A,
    send_addr: Option<A>,
) -> io::Result<(super::Io, SocketAddress)> {
    let rx_socket = syscall::bind_udp(receive_addr, false)?;
    rx_socket.set_nonblocking(true)?;
    let rx_socket: std::net::UdpSocket = rx_socket.into();
    let rx_addr = rx_socket.local_addr()?;

    let mut io_builder = Io::builder().with_rx_socket(rx_socket)?;

    if let Some(tx_addr) = send_addr {
        let tx_socket = syscall::bind_udp(tx_addr, false)?;
        tx_socket.set_nonblocking(true)?;
        let tx_socket: std::net::UdpSocket = tx_socket.into();
        io_builder = io_builder.with_tx_socket(tx_socket)?
    }

    let io = io_builder.build()?;

    let rx_addr = if rx_addr.is_ipv6() {
        ("::1", rx_addr.port())
    } else {
        ("127.0.0.1", rx_addr.port())
    }
    .to_socket_addrs()?
    .next()
    .unwrap();

    Ok((io, rx_addr.into()))
}

/// The tokio IO provider allows the application to configure different sockets for rx
/// and tx. This function will accept optional TX addresses to test this functionality.
async fn test<A: ToSocketAddrs>(
    server_rx_addr: A,
    server_tx_addr: Option<A>,
    client_rx_addr: A,
    client_tx_addr: Option<A>,
) -> io::Result<()> {
    let (server_io, server_addr) = runtime(server_rx_addr, server_tx_addr).await?;
    let (client_io, client_addr) = runtime(client_rx_addr, client_tx_addr).await?;

    let server_endpoint = {
        let mut handle = PathHandle::from_remote_address(client_addr.into());
        handle.local_address = server_addr.into();
        TestEndpoint::<true>::new(handle)
    };

    let client_endpoint = {
        let mut handle = PathHandle::from_remote_address(server_addr.into());
        handle.local_address = client_addr.into();
        TestEndpoint::<false>::new(handle)
    };

    let (server_task, actual_server_addr) = server_io.start(server_endpoint)?;
    assert_eq!(actual_server_addr, server_addr);

    let (client_task, actual_client_addr) = client_io.start(client_endpoint)?;
    assert_eq!(actual_client_addr, client_addr);

    tokio::time::timeout(core::time::Duration::from_secs(10), client_task).await??;

    server_task.abort();

    Ok(())
}

static IPV4_LOCALHOST: &str = "127.0.0.1:0";
static IPV6_LOCALHOST: &str = "[::1]:0";

#[tokio::test]
#[cfg_attr(miri, ignore)]
async fn ipv4_test() -> io::Result<()> {
    test(IPV4_LOCALHOST, None, IPV4_LOCALHOST, None).await
}

#[tokio::test]
#[cfg_attr(miri, ignore)]
async fn ipv4_two_socket_test() -> io::Result<()> {
    test(
        IPV4_LOCALHOST,
        Some(IPV4_LOCALHOST),
        IPV4_LOCALHOST,
        Some(IPV4_LOCALHOST),
    )
    .await
}

#[tokio::test]
#[cfg_attr(miri, ignore)]
async fn ipv6_test() -> io::Result<()> {
    let result = test(IPV6_LOCALHOST, None, IPV6_LOCALHOST, None).await;

    match result {
        Err(err) if err.kind() == io::ErrorKind::AddrNotAvailable => {
            eprintln!("The current environment does not support IPv6; skipping");
            Ok(())
        }
        other => other,
    }
}

#[tokio::test]
#[cfg_attr(miri, ignore)]
async fn ipv6_two_socket_test() -> io::Result<()> {
    let result = test(
        IPV6_LOCALHOST,
        Some(IPV6_LOCALHOST),
        IPV6_LOCALHOST,
        Some(IPV6_LOCALHOST),
    )
    .await;

    match result {
        Err(err) if err.kind() == io::ErrorKind::AddrNotAvailable => {
            eprintln!("The current environment does not support IPv6; skipping");
            Ok(())
        }
        other => other,
    }
}


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

use crate::{
    features::Gso,
    message::{simple::Message, Message as _},
    socket::{
        ring, task,
        task::{rx, tx},
    },
    syscall::SocketEvents,
};
use core::task::{Context, Poll};
use tokio::{io, net::UdpSocket};

pub async fn rx<S: Into<std::net::UdpSocket>>(
    socket: S,
    producer: ring::Producer<Message>,
) -> io::Result<()> {
    let socket = socket.into();
    socket.set_nonblocking(true).unwrap();

    let socket = UdpSocket::from_std(socket).unwrap();
    let result = task::Receiver::new(producer, socket).await;
    if let Some(err) = result {
        Err(err)
    } else {
        Ok(())
    }
}

pub async fn tx<S: Into<std::net::UdpSocket>>(
    socket: S,
    consumer: ring::Consumer<Message>,
    gso: Gso,
) -> io::Result<()> {
    let socket = socket.into();
    socket.set_nonblocking(true).unwrap();

    let socket = UdpSocket::from_std(socket).unwrap();
    let result = task::Sender::new(consumer, socket, gso).await;
    if let Some(err) = result {
        Err(err)
    } else {
        Ok(())
    }
}

impl tx::Socket<Message> for UdpSocket {
    type Error = io::Error;

    #[inline]
    fn send(
        &mut self,
        cx: &mut Context,
        entries: &mut [Message],
        events: &mut tx::Events,
    ) -> io::Result<()> {
        for entry in entries {
            let target = (*entry.remote_address()).into();
            let payload = entry.payload_mut();
            match self.poll_send_to(cx, payload, target) {
                Poll::Ready(Ok(_)) => {
                    if events.on_complete(1).is_break() {
                        return Ok(());
                    }
                }
                Poll::Ready(Err(err)) => {
                    if events.on_error(err).is_break() {
                        return Ok(());
                    }
                }
                Poll::Pending => {
                    events.blocked();
                    break;
                }
            }
        }

        Ok(())
    }
}

impl rx::Socket<Message> for UdpSocket {
    type Error = io::Error;

    #[inline]
    fn recv(
        &mut self,
        cx: &mut Context,
        entries: &mut [Message],
        events: &mut rx::Events,
    ) -> io::Result<()> {
        for entry in entries {
            let payload = entry.payload_mut();
            let mut buf = io::ReadBuf::new(payload);
            match self.poll_recv_from(cx, &mut buf) {
                Poll::Ready(Ok(addr)) => {
                    unsafe {
                        let len = buf.filled().len();
                        entry.set_payload_len(len);
                    }
                    entry.set_remote_address(&(addr.into()));

                    if events.on_complete(1).is_break() {
                        return Ok(());
                    }
                }
                Poll::Ready(Err(err)) => {
                    if events.on_error(err).is_break() {
                        return Ok(());
                    }
                }
                Poll::Pending => {
                    events.blocked();
                    break;
                }
            }
        }

        Ok(())
    }
}


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

use crate::{
    features::Gso,
    socket::{
        ring,
        task::{rx, tx},
    },
    syscall::{SocketType, UnixMessage},
};
use core::task::{Context, Poll};
use std::{io, os::unix::io::AsRawFd};
use tokio::io::unix::AsyncFd;

pub async fn rx<S: Into<std::net::UdpSocket>, M: UnixMessage + Unpin>(
    socket: S,
    producer: ring::Producer<M>,
) -> io::Result<()> {
    let socket = socket.into();
    socket.set_nonblocking(true).unwrap();

    let socket = AsyncFd::new(socket).unwrap();
    let result = rx::Receiver::new(producer, socket).await;
    if let Some(err) = result {
        Err(err)
    } else {
        Ok(())
    }
}

pub async fn tx<S: Into<std::net::UdpSocket>, M: UnixMessage + Unpin>(
    socket: S,
    consumer: ring::Consumer<M>,
    gso: Gso,
) -> io::Result<()> {
    let socket = socket.into();
    socket.set_nonblocking(true).unwrap();

    let socket = AsyncFd::new(socket).unwrap();
    let result = tx::Sender::new(consumer, socket, gso).await;
    if let Some(err) = result {
        Err(err)
    } else {
        Ok(())
    }
}

impl<S: AsRawFd, M: UnixMessage> tx::Socket<M> for AsyncFd<S> {
    type Error = io::Error;

    #[inline]
    fn send(
        &mut self,
        cx: &mut Context,
        entries: &mut [M],
        events: &mut tx::Events,
    ) -> io::Result<()> {
        // Call the syscall for the socket
        //
        // NOTE: we usually wrap this in a `AsyncFdReadyGuard::try_io`. However, here we just
        //       assume the socket is ready in the general case and then fall back to querying
        //       socket readiness if it's not. This can avoid some things like having to construct
        //       a `std::io::Error` with `WouldBlock` and dereferencing the registration.
        M::send(self.get_ref().as_raw_fd(), entries, events);

        // yield back if we weren't blocked
        if !events.is_blocked() {
            return Ok(());
        }

        // * First iteration we need to clear socket readiness since the `send` call returned a
        // `WouldBlock`.
        // * Second iteration we need to register the waker, assuming the socket readiness was
        // cleared.
        //   * If we got a `Ready` anyway, then clear the blocked status and have the caller try
        //   again.
        for i in 0..2 {
            match self.poll_write_ready(cx) {
                Poll::Ready(guard) => {
                    let mut guard = guard?;
                    if i == 0 {
                        guard.clear_ready();
                    } else {
                        events.take_blocked();
                    }
                }
                Poll::Pending => {
                    return Ok(());
                }
            }
        }

        Ok(())
    }
}

impl<S: AsRawFd, M: UnixMessage> rx::Socket<M> for AsyncFd<S> {
    type Error = io::Error;

    #[inline]
    fn recv(
        &mut self,
        cx: &mut Context,
        entries: &mut [M],
        events: &mut rx::Events,
    ) -> io::Result<()> {
        // Call the syscall for the socket
        //
        // NOTE: we usually wrap this in a `AsyncFdReadyGuard::try_io`. However, here we just
        //       assume the socket is ready in the general case and then fall back to querying
        //       socket readiness if it's not. This can avoid some things like having to construct
        //       a `std::io::Error` with `WouldBlock` and dereferencing the registration.
        M::recv(
            self.get_ref().as_raw_fd(),
            SocketType::NonBlocking,
            entries,
            events,
        );

        // yield back if we weren't blocked
        if !events.is_blocked() {
            return Ok(());
        }

        // * First iteration we need to clear socket readiness since the `recv` call returned a
        // `WouldBlock`.
        // * Second iteration we need to register the waker, assuming the socket readiness was
        // cleared.
        //   * If we got a `Ready` anyway, then clear the blocked status and have the caller try
        //   again.
        for i in 0..2 {
            match self.poll_read_ready(cx) {
                Poll::Ready(guard) => {
                    let mut guard = guard?;
                    if i == 0 {
                        guard.clear_ready();
                    } else {
                        events.take_blocked();
                    }
                }
                Poll::Pending => {
                    return Ok(());
                }
            }
        }

        Ok(())
    }
}


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

use super::*;

#[derive(Default)]
pub struct Builder {
    pub(super) handle: Option<Handle>,
    pub(super) socket: Option<UdpSocket>,
    pub(super) addr: Option<Box<dyn turmoil::ToSocketAddrs + Send + Sync + 'static>>,
    pub(super) max_mtu: MaxMtu,
}

impl Builder {
    #[must_use]
    pub fn with_handle(mut self, handle: Handle) -> Self {
        self.handle = Some(handle);
        self
    }

    /// Sets the local address for the runtime to listen on.
    ///
    /// NOTE: this method is mutually exclusive with `with_socket`
    pub fn with_address<A: turmoil::ToSocketAddrs + Send + Sync + 'static>(
        mut self,
        addr: A,
    ) -> io::Result<Self> {
        debug_assert!(self.socket.is_none(), "socket has already been set");
        self.addr = Some(Box::new(addr));
        Ok(self)
    }

    /// Sets the socket used for sending and receiving for the runtime.
    ///
    /// NOTE: this method is mutually exclusive with `with_address`
    pub fn with_socket(mut self, socket: UdpSocket) -> io::Result<Self> {
        debug_assert!(self.addr.is_none(), "address has already been set");
        self.socket = Some(socket);
        Ok(self)
    }

    /// Sets the largest maximum transmission unit (MTU) that can be sent on a path
    pub fn with_max_mtu(mut self, max_mtu: u16) -> io::Result<Self> {
        self.max_mtu = max_mtu
            .try_into()
            .map_err(|err| io::Error::new(ErrorKind::InvalidInput, format!("{err}")))?;
        Ok(self)
    }

    pub fn build(self) -> io::Result<Io> {
        Ok(Io { builder: self })
    }
}


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

use super::*;
use core::{
    convert::TryInto,
    task::{Context, Poll},
};
use s2n_quic_core::{
    endpoint::{self, CloseError},
    event,
    inet::SocketAddress,
    io::{rx, tx},
    path::Handle as _,
    time::{timer::Provider as _, Clock, Duration, Timer, Timestamp},
};
use std::collections::BTreeMap;

struct TestEndpoint {
    addr: SocketAddress,
    tx_message_id: u32,
    rx_messages: BTreeMap<u32, Timestamp>,
    total_messages: u32,
    subscriber: NoopSubscriber,
    close_timer: Timer,
}

impl TestEndpoint {
    fn new(addr: SocketAddress) -> Self {
        Self {
            addr,
            tx_message_id: 0,
            rx_messages: BTreeMap::new(),
            total_messages: 1000,
            subscriber: Default::default(),
            close_timer: Default::default(),
        }
    }
}

#[derive(Debug, Default)]
struct NoopSubscriber;

impl event::Subscriber for NoopSubscriber {
    type ConnectionContext = ();

    fn create_connection_context(
        &mut self,
        _meta: &event::api::ConnectionMeta,
        _info: &event::api::ConnectionInfo,
    ) -> Self::ConnectionContext {
    }
}

impl Endpoint for TestEndpoint {
    type PathHandle = PathHandle;
    type Subscriber = NoopSubscriber;

    const ENDPOINT_TYPE: endpoint::Type = endpoint::Type::Server;

    fn transmit<Tx: tx::Queue<Handle = PathHandle>, C: Clock>(
        &mut self,
        queue: &mut Tx,
        _clock: &C,
    ) {
        while self.tx_message_id < self.total_messages {
            let payload = self.tx_message_id.to_be_bytes();
            let addr = PathHandle::from_remote_address(self.addr.into());
            let msg = (addr, payload);
            if queue.push(msg).is_ok() {
                self.tx_message_id += 1;
            } else {
                // no more capacity
                return;
            }
        }
    }

    fn receive<Rx: rx::Queue<Handle = PathHandle>, C: Clock>(&mut self, queue: &mut Rx, clock: &C) {
        let now = clock.get_time();
        queue.for_each(|_header, payload| {
            assert_eq!(payload.len(), 4, "invalid payload {:?}", payload);

            let id = (&*payload).try_into().unwrap();
            let id = u32::from_be_bytes(id);
            self.rx_messages.insert(id, now);
        });
    }

    fn poll_wakeups<C: Clock>(
        &mut self,
        _cx: &mut Context<'_>,
        clock: &C,
    ) -> Poll<Result<usize, CloseError>> {
        let now = clock.get_time();

        if self.close_timer.poll_expiration(now).is_ready() {
            assert!(self.rx_messages.len() as u32 * 4 > self.total_messages);
            return Err(CloseError).into();
        }

        if !self.close_timer.is_armed()
            && self.total_messages <= self.tx_message_id
            && !self.rx_messages.is_empty()
        {
            self.close_timer.set(now + Duration::from_millis(100));
        }

        Poll::Pending
    }

    fn timeout(&self) -> Option<Timestamp> {
        self.close_timer.next_expiration()
    }

    fn set_max_mtu(&mut self, _max_mtu: MaxMtu) {
        // noop
    }

    fn subscriber(&mut self) -> &mut Self::Subscriber {
        &mut self.subscriber
    }
}

fn bind(port: u16) -> std::net::SocketAddr {
    use std::net::Ipv4Addr;
    (Ipv4Addr::UNSPECIFIED, port).into()
}

#[test]
fn sim_test() -> io::Result<()> {
    use turmoil::lookup;

    let mut sim = turmoil::Builder::new().build();

    sim.client("client", async move {
        let io = Io::builder().with_address(bind(123))?.build()?;

        let endpoint = TestEndpoint::new((lookup("server"), 456).into());

        let (task, _) = io.start(endpoint)?;

        task.await?;

        Ok(())
    });

    sim.client("server", async move {
        let io = Io::builder().with_address(bind(456))?.build()?;

        let endpoint = TestEndpoint::new((lookup("client"), 123).into());

        let (task, _) = io.start(endpoint)?;

        task.await?;

        Ok(())
    });

    sim.run().unwrap();

    Ok(())
}


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

use core::mem::size_of;
use s2n_quic_core::{
    inet::{ethernet, ipv4, udp},
    path::{MaxMtu, MaxMtuError},
};
use s2n_quic_xdp::umem::DEFAULT_FRAME_SIZE;
use tokio::runtime::Handle;

/// Calculate how much a packet will need for fixed-size headers
const MIN_FRAME_OVERHEAD: u16 =
    (size_of::<ethernet::Header>() + size_of::<ipv4::Header>() + size_of::<udp::Header>()) as _;

#[derive(Debug)]
#[must_use = "Builders do nothing without calling `build`"]
pub struct Builder<Rx = (), Tx = ()> {
    rx: Rx,
    tx: Tx,
    max_mtu: MaxMtu,
    handle: Option<Handle>,
}

impl Default for Builder<(), ()> {
    fn default() -> Self {
        Self {
            rx: (),
            tx: (),
            max_mtu: MaxMtu::try_from(DEFAULT_FRAME_SIZE as u16 - MIN_FRAME_OVERHEAD).unwrap(),
            handle: None,
        }
    }
}

impl<Rx, Tx> Builder<Rx, Tx> {
    /// Sets the tokio runtime handle for the provider
    pub fn with_handle(mut self, handle: Handle) -> Self {
        self.handle = Some(handle);
        self
    }

    /// Sets the UMEM frame size for the provider
    pub fn with_frame_size(mut self, frame_size: u16) -> Result<Self, MaxMtuError> {
        self.max_mtu = frame_size.saturating_sub(MIN_FRAME_OVERHEAD).try_into()?;
        Ok(self)
    }

    /// Sets the RX implementation for the provider
    pub fn with_rx<NewRx>(self, rx: NewRx) -> Builder<NewRx, Tx>
    where
        NewRx: super::rx::Rx,
    {
        let Self {
            tx,
            handle,
            max_mtu,
            ..
        } = self;
        Builder {
            rx,
            tx,
            handle,
            max_mtu,
        }
    }

    /// Sets the TX implementation for the provider
    pub fn with_tx<NewTx>(self, tx: NewTx) -> Builder<Rx, NewTx>
    where
        NewTx: super::tx::Tx,
    {
        let Self {
            rx,
            handle,
            max_mtu,
            ..
        } = self;
        Builder {
            rx,
            tx,
            handle,
            max_mtu,
        }
    }
}

impl<Rx, Tx> Builder<Rx, Tx>
where
    Rx: 'static + super::rx::Rx + Send,
    Tx: 'static + super::tx::Tx<PathHandle = Rx::PathHandle> + Send,
{
    pub fn build(self) -> super::Provider<Rx, Tx> {
        let Self {
            rx,
            tx,
            handle,
            max_mtu,
        } = self;
        super::Provider {
            rx,
            tx,
            handle,
            max_mtu,
        }
    }
}


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

#![allow(clippy::unnecessary_cast)] // some platforms encode lengths as `u32` so we cast everything to be safe

use core::mem::{align_of, size_of};
use libc::cmsghdr;
use s2n_quic_core::inet::{AncillaryData, ExplicitCongestionNotification};

const fn size_of_cmsg<T: Copy + Sized>() -> usize {
    unsafe { libc::CMSG_SPACE(size_of::<T>() as _) as _ }
}

const fn const_max(a: usize, b: usize) -> usize {
    if a > b {
        a
    } else {
        b
    }
}

/// The maximum number of bytes allocated for cmsg data
///
/// This should be enough for UDP_SEGMENT + IP_TOS + IP_PKTINFO. It may need to be increased
/// to allow for future control messages.
pub const MAX_LEN: usize = {
    let tos_size = size_of_cmsg::<IpTos>();

    #[cfg(s2n_quic_platform_gso)]
    let gso_size = size_of_cmsg::<UdpGso>();
    #[cfg(not(s2n_quic_platform_gso))]
    let gso_size = 0;

    #[cfg(s2n_quic_platform_gro)]
    let gro_size = size_of_cmsg::<UdpGro>();
    #[cfg(not(s2n_quic_platform_gro))]
    let gro_size = 0;

    let segment_offload_size = const_max(gso_size, gro_size);

    // rather than taking the max, we add these in case the OS gives us both
    #[cfg(s2n_quic_platform_pktinfo)]
    let pktinfo_size = size_of_cmsg::<libc::in_pktinfo>() + size_of_cmsg::<libc::in6_pktinfo>();
    #[cfg(not(s2n_quic_platform_pktinfo))]
    let pktinfo_size = 0;

    // This is currently needed due to how we detect if CMSG data has been written or not.
    //
    // TODO remove this once we split the `reset` traits into TX and RX types
    let padding = size_of::<cmsghdr>();

    tos_size + segment_offload_size + pktinfo_size + padding
};

#[cfg(s2n_quic_platform_gso)]
pub type UdpGso = u16;
#[cfg(s2n_quic_platform_gro)]
pub type UdpGro = libc::c_int;
pub type IpTos = libc::c_int;

#[repr(align(8))] // the storage needs to be aligned to the same as `cmsghdr`
#[derive(Clone, Debug)]
pub struct Storage([u8; MAX_LEN]);

impl Default for Storage {
    fn default() -> Self {
        Self([0; MAX_LEN])
    }
}

impl Storage {
    #[inline]
    pub fn len(&self) -> usize {
        self.0.len()
    }

    #[inline]
    #[allow(dead_code)] // clippy wants this to exist but we don't use it
    pub fn is_empty(&self) -> bool {
        self.0.is_empty()
    }

    #[inline]
    pub fn as_mut_ptr(&mut self) -> *mut u8 {
        self.0.as_mut_ptr()
    }
}

#[derive(Clone, Copy, Debug)]
pub struct OutOfSpace;

pub struct SliceEncoder<'a> {
    storage: &'a mut [u8],
    cursor: usize,
}

impl<'a> Encoder for SliceEncoder<'a> {
    #[inline]
    fn encode_cmsg<T: Copy + ?Sized>(
        &mut self,
        level: libc::c_int,
        ty: libc::c_int,
        value: T,
    ) -> Result<usize, OutOfSpace> {
        unsafe {
            debug_assert!(
                align_of::<T>() <= align_of::<cmsghdr>(),
                "alignment of T should be less than or equal to cmsghdr"
            );

            // CMSG_SPACE() returns the number of bytes an ancillary element
            // with payload of the passed data length occupies.
            let element_len = size_of_cmsg::<T>();
            debug_assert_ne!(element_len, 0);
            debug_assert_eq!(libc::CMSG_SPACE(size_of::<T>() as _) as usize, element_len);

            let new_cursor = self.cursor.checked_add(element_len).ok_or(OutOfSpace)?;

            self.storage
                .len()
                .checked_sub(new_cursor)
                .ok_or(OutOfSpace)?;

            let cmsg_ptr = {
                // Safety: the msg_control buffer should always be allocated to MAX_LEN
                let msg_controllen = self.cursor;
                let msg_control = self.storage.as_mut_ptr().add(msg_controllen as _);
                msg_control as *mut cmsghdr
            };

            {
                let cmsg = &mut *cmsg_ptr;

                // interpret the start of cmsg as a cmsghdr
                // Safety: the cmsg slice should already be zero-initialized and aligned

                // Indicate the type of cmsg
                cmsg.cmsg_level = level;
                cmsg.cmsg_type = ty;

                // CMSG_LEN() returns the value to store in the cmsg_len member
                // of the cmsghdr structure, taking into account any necessary
                // alignment.  It takes the data length as an argument.
                cmsg.cmsg_len = libc::CMSG_LEN(size_of::<T>() as _) as _;
            }

            {
                // Write the actual value in the data space of the cmsg
                // Safety: we asserted we had enough space in the cmsg buffer above
                // CMSG_DATA() returns a pointer to the data portion of a
                // cmsghdr. The pointer returned cannot be assumed to be
                // suitably aligned for accessing arbitrary payload data types.
                // Applications should not cast it to a pointer type matching the
                // payload, but should instead use memcpy(3) to copy data to or
                // from a suitably declared object.
                let data_ptr = cmsg_ptr.add(1);

                debug_assert_eq!(data_ptr as *mut u8, libc::CMSG_DATA(cmsg_ptr) as *mut u8);

                core::ptr::copy_nonoverlapping(
                    &value as *const T as *const u8,
                    data_ptr as *mut u8,
                    size_of::<T>(),
                );
            }

            // add the values as a usize to make sure we work cross-platform
            self.cursor = new_cursor;
            debug_assert!(
                self.cursor <= self.storage.len(),
                "msg should not exceed max allocated"
            );

            Ok(self.cursor)
        }
    }
}

pub trait Encoder {
    /// Encodes the given value as a control message in the cmsg buffer.
    ///
    /// The msghdr.msg_control should be zero-initialized and aligned and contain enough
    /// room for the value to be written.
    fn encode_cmsg<T: Copy + ?Sized>(
        &mut self,
        level: libc::c_int,
        ty: libc::c_int,
        value: T,
    ) -> Result<usize, OutOfSpace>;
}

impl Encoder for libc::msghdr {
    #[inline]
    fn encode_cmsg<T: Copy + ?Sized>(
        &mut self,
        level: libc::c_int,
        ty: libc::c_int,
        value: T,
    ) -> Result<usize, OutOfSpace> {
        let storage = unsafe { &mut *(self.msg_control as *mut Storage) };

        let mut encoder = SliceEncoder {
            storage: &mut storage.0,
            cursor: self.msg_controllen as _,
        };

        let cursor = encoder.encode_cmsg(level, ty, value)?;

        self.msg_controllen = cursor as _;

        Ok(cursor)
    }
}

/// Decodes all recognized control messages in the given `msghdr` into `AncillaryData`
#[inline]
pub fn decode(msghdr: &libc::msghdr) -> AncillaryData {
    let mut result = AncillaryData::default();

    let iter = unsafe { Iter::from_msghdr(msghdr) };

    for (cmsg, value) in iter {
        unsafe {
            match (cmsg.cmsg_level, cmsg.cmsg_type) {
                // Linux uses IP_TOS, FreeBSD uses IP_RECVTOS
                (libc::IPPROTO_IP, libc::IP_TOS)
                | (libc::IPPROTO_IP, libc::IP_RECVTOS)
                | (libc::IPPROTO_IPV6, libc::IPV6_TCLASS) => {
                    // IP_TOS cmsgs should be 1 byte, but occasionally are reported as 4 bytes
                    let value = match value.len() {
                        1 => decode_value::<u8>(value),
                        4 => decode_value::<u32>(value) as u8,
                        len => {
                            if cfg!(test) {
                                panic!(
                                    "invalid size for ECN marking. len: {len}, value: {value:?}"
                                );
                            }
                            continue;
                        }
                    };

                    result.ecn = ExplicitCongestionNotification::new(value);
                }
                #[cfg(s2n_quic_platform_pktinfo)]
                (libc::IPPROTO_IP, libc::IP_PKTINFO) => {
                    let pkt_info = decode_value::<libc::in_pktinfo>(value);

                    // read from both fields in case only one is set and not the other
                    //
                    // from https://man7.org/linux/man-pages/man7/ip.7.html:
                    //
                    // > ipi_spec_dst is the local address
                    // > of the packet and ipi_addr is the destination address in
                    // > the packet header.
                    let local_address =
                        match (pkt_info.ipi_addr.s_addr, pkt_info.ipi_spec_dst.s_addr) {
                            (0, v) => v.to_ne_bytes(),
                            (v, _) => v.to_ne_bytes(),
                        };

                    // The port should be specified by a different layer that has that information
                    let port = 0;
                    let local_address =
                        s2n_quic_core::inet::SocketAddressV4::new(local_address, port);
                    result.local_address = local_address.into();
                    result.local_interface = Some(pkt_info.ipi_ifindex as _);
                }
                #[cfg(s2n_quic_platform_pktinfo)]
                (libc::IPPROTO_IPV6, libc::IPV6_PKTINFO) => {
                    let pkt_info = decode_value::<libc::in6_pktinfo>(value);
                    let local_address = pkt_info.ipi6_addr.s6_addr;
                    // The port should be specified by a different layer that has that information
                    let port = 0;
                    let local_address =
                        s2n_quic_core::inet::SocketAddressV6::new(local_address, port);
                    result.local_address = local_address.into();
                    result.local_interface = Some(pkt_info.ipi6_ifindex as _);
                }
                #[cfg(s2n_quic_platform_gso)]
                (libc::SOL_UDP, libc::UDP_SEGMENT) => {
                    // ignore GSO settings when reading
                    continue;
                }
                #[cfg(s2n_quic_platform_gro)]
                (libc::SOL_UDP, libc::UDP_GRO) => {
                    let segment_size = decode_value::<UdpGro>(value);
                    result.segment_size = segment_size as _;
                }
                (level, ty) if cfg!(test) => {
                    // if we're getting an unexpected cmsg we should know about it in testing
                    panic!("unexpected cmsghdr {{ level: {level}, type: {ty}, value: {value:?} }}");
                }
                _ => {}
            }
        }
    }

    result
}

/// Decodes a value of type `T` from the given `cmsghdr`
/// # Safety
///
/// `cmsghdr` must refer to a cmsg containing a payload of type `T`
#[inline]
unsafe fn decode_value<T: Copy>(value: &[u8]) -> T {
    use core::mem;

    debug_assert!(mem::align_of::<T>() <= mem::align_of::<cmsghdr>());
    debug_assert!(value.len() >= size_of::<T>());

    let mut v = mem::zeroed::<T>();

    core::ptr::copy_nonoverlapping(value.as_ptr(), &mut v as *mut T as *mut u8, size_of::<T>());

    v
}

struct Iter<'a> {
    cursor: *const u8,
    len: usize,
    contents: core::marker::PhantomData<&'a [u8]>,
}

impl<'a> Iter<'a> {
    /// Creates a new cmsg::Iter used for iterating over control message headers in the given
    /// slice of bytes.
    ///
    /// # Safety
    ///
    /// * `contents` must be aligned to cmsghdr
    #[inline]
    unsafe fn new(contents: &'a [u8]) -> Iter<'a> {
        let cursor = contents.as_ptr();
        let len = contents.len();

        debug_assert_eq!(
            cursor.align_offset(align_of::<cmsghdr>()),
            0,
            "contents must be aligned to cmsghdr"
        );

        Self {
            cursor,
            len,
            contents: Default::default(),
        }
    }

    #[inline]
    unsafe fn from_msghdr(msghdr: &'a libc::msghdr) -> Self {
        let ptr = msghdr.msg_control as *const u8;
        let len = msghdr.msg_controllen as usize;
        let slice = core::slice::from_raw_parts(ptr, len);
        Self::new(slice)
    }
}

impl<'a> Iterator for Iter<'a> {
    type Item = (&'a cmsghdr, &'a [u8]);

    #[inline]
    fn next(&mut self) -> Option<Self::Item> {
        unsafe {
            let cursor = self.cursor;

            // make sure we can decode a cmsghdr
            self.len.checked_sub(size_of::<cmsghdr>())?;
            let cmsg = &*(cursor as *const cmsghdr);
            let data_ptr = cursor.add(size_of::<cmsghdr>());

            let cmsg_len = cmsg.cmsg_len as usize;

            // make sure we have capacity to decode the provided cmsg_len
            self.len.checked_sub(cmsg_len)?;

            // the cmsg_len includes the header itself so it needs to be subtracted off
            let data_len = cmsg_len.checked_sub(size_of::<cmsghdr>())?;
            // construct a slice with the provided data len
            let data = core::slice::from_raw_parts(data_ptr, data_len);

            // empty messages are invalid
            if data.is_empty() {
                return None;
            }

            // calculate the next message and update the cursor/len
            {
                let space = libc::CMSG_SPACE(data_len as _) as usize;
                debug_assert!(
                    space >= data_len,
                    "space ({space}) should be at least of size len ({data_len})"
                );
                self.len = self.len.saturating_sub(space);
                self.cursor = cursor.add(space);
            }

            Some((cmsg, data))
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use bolero::{check, TypeGenerator};
    use libc::c_int;

    /// Ensures the cmsg iterator doesn't crash or segfault
    #[test]
    #[cfg_attr(kani, kani::proof, kani::solver(cadical), kani::unwind(17))]
    fn iter_test() {
        check!().for_each(|cmsg| {
            // the bytes needs to be aligned to a cmsghdr
            let offset = cmsg.as_ptr().align_offset(align_of::<cmsghdr>());

            if let Some(cmsg) = cmsg.get(offset..) {
                for (cmsghdr, value) in unsafe { Iter::new(cmsg) } {
                    let _ = cmsghdr;
                    let _ = value;
                }
            }
        });
    }

    #[derive(Clone, Copy, Debug, TypeGenerator)]
    struct Op {
        level: c_int,
        ty: c_int,
        value: Value,
    }

    #[derive(Clone, Copy, Debug, TypeGenerator)]
    enum Value {
        U8(u8),
        U16(u16),
        U32(u32),
        // alignment can't exceed that of cmsghdr
        U64([u32; 2]),
        U128([u32; 4]),
    }

    impl Value {
        fn check_value(&self, bytes: &[u8]) {
            let expected_len = match self {
                Self::U8(_) => 1,
                Self::U16(_) => 2,
                Self::U32(_) => 4,
                Self::U64(_) => 8,
                Self::U128(_) => 16,
            };
            assert_eq!(expected_len, bytes.len());
        }
    }

    fn round_trip(ops: &[Op]) {
        let mut storage = Storage::default();
        let mut encoder = SliceEncoder {
            storage: &mut storage.0,
            cursor: 0,
        };

        let mut expected_encoded_count = 0;

        for op in ops {
            let res = match op.value {
                Value::U8(value) => encoder.encode_cmsg(op.level, op.ty, value),
                Value::U16(value) => encoder.encode_cmsg(op.level, op.ty, value),
                Value::U32(value) => encoder.encode_cmsg(op.level, op.ty, value),
                Value::U64(value) => encoder.encode_cmsg(op.level, op.ty, value),
                Value::U128(value) => encoder.encode_cmsg(op.level, op.ty, value),
            };

            match res {
                Ok(_) => expected_encoded_count += 1,
                Err(_) => break,
            }
        }

        let cursor = encoder.cursor;
        let mut actual_decoded_count = 0;
        let mut iter = unsafe { Iter::new(&storage.0[..cursor]) };

        for (op, (cmsghdr, value)) in ops.iter().zip(&mut iter) {
            assert_eq!(op.level, cmsghdr.cmsg_level);
            assert_eq!(op.ty, cmsghdr.cmsg_type);
            op.value.check_value(value);
            actual_decoded_count += 1;
        }

        assert_eq!(expected_encoded_count, actual_decoded_count);
        assert!(iter.next().is_none());
    }

    #[cfg(not(kani))]
    type Ops = Vec<Op>;
    #[cfg(kani)]
    type Ops = s2n_quic_core::testing::InlineVec<Op, 8>;

    #[test]
    #[cfg_attr(kani, kani::proof, kani::solver(cadical), kani::unwind(9))]
    fn round_trip_test() {
        check!().with_type::<Ops>().for_each(|ops| round_trip(ops));
    }
}


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

use crate::message::{msg, Message as MessageTrait};
use libc::mmsghdr;
use s2n_quic_core::{io::tx, path};

pub use libc::mmsghdr as Message;
pub type Handle = msg::Handle;

impl MessageTrait for mmsghdr {
    type Handle = Handle;

    const SUPPORTS_GSO: bool = libc::msghdr::SUPPORTS_GSO;
    const SUPPORTS_ECN: bool = libc::msghdr::SUPPORTS_ECN;
    const SUPPORTS_FLOW_LABELS: bool = libc::msghdr::SUPPORTS_FLOW_LABELS;

    #[inline]
    fn alloc(entries: u32, payload_len: u32, offset: usize) -> super::Storage {
        unsafe {
            msg::alloc(entries, payload_len, offset, |mmsghdr: &mut mmsghdr| {
                mmsghdr.msg_len = payload_len as _;
                &mut mmsghdr.msg_hdr
            })
        }
    }

    #[inline]
    fn payload_len(&self) -> usize {
        let payload_len = self.msg_len as usize;
        debug_assert!(payload_len <= core::u16::MAX as usize);
        payload_len
    }

    #[inline]
    unsafe fn set_payload_len(&mut self, len: usize) {
        debug_assert!(len <= core::u16::MAX as usize);
        self.msg_len = len as _;
        self.msg_hdr.set_payload_len(len);
    }

    #[inline]
    fn set_segment_size(&mut self, size: usize) {
        self.msg_hdr.set_segment_size(size)
    }

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

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

    #[inline]
    fn validate_replication(source: &Self, dest: &Self) {
        libc::msghdr::validate_replication(&source.msg_hdr, &dest.msg_hdr)
    }

    #[inline]
    fn rx_read(
        &mut self,
        local_address: &path::LocalAddress,
    ) -> Option<super::RxMessage<Self::Handle>> {
        unsafe {
            // We need to replicate the `msg_len` field to the inner type before delegating
            // Safety: The `msg_len` is associated with the same buffer as the `msg_hdr`
            self.msg_hdr.set_payload_len(self.msg_len as _);
        }
        self.msg_hdr.rx_read(local_address)
    }

    #[inline]
    fn tx_write<M: tx::Message<Handle = Self::Handle>>(
        &mut self,
        message: M,
    ) -> Result<usize, tx::Error> {
        let len = self.msg_hdr.tx_write(message)?;
        // We need to replicate the len with the `msg_len` field after delegating to `msg_hdr`
        debug_assert!(len <= core::u16::MAX as usize);
        self.msg_len = len as _;
        Ok(len)
    }
}