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

use super::*;
use s2n_codec::EncoderValue;
use s2n_quic_core::frame::{self, FitError};

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

impl FrameWriter for Stream {
    type Context = VarInt;

    #[inline]
    fn write_chunk<W: WriteContext>(
        &self,
        offset: VarInt,
        data: &mut View,
        stream_id: Self::Context,
        context: &mut W,
    ) -> Result<(), FitError> {
        let remaining_capacity = context.remaining_capacity();

        debug_assert!(
            data.len() <= remaining_capacity,
            "the data sender should not pass a payload that exceeds the current capacity"
        );

        let mut frame = frame::Stream {
            stream_id,
            offset,
            // this will be updated by `try_fit`
            is_last_frame: false,
            // this will be updated after we've made sure the frame fits
            is_fin: false,
            data,
        };

        let len = frame.try_fit(remaining_capacity)?;
        if len == 0 {
            return Err(FitError);
        }

        frame.data.trim_off(frame.data.encoding_size() - len)?;
        frame.is_fin = frame.data.is_fin();

        context.write_fitted_frame(&frame);

        Ok(())
    }

    #[inline]
    fn write_fin<W: WriteContext>(
        &self,
        offset: VarInt,
        stream_id: Self::Context,
        context: &mut W,
    ) -> Result<(), FitError> {
        let mut frame = frame::Stream {
            stream_id,
            offset,
            is_last_frame: false,
            is_fin: true,
            data: &[0u8; 0][..],
        };

        // the length is always 0 so we don't need to trim the data
        frame.try_fit(context.remaining_capacity())?;
        context.write_fitted_frame(&frame);

        Ok(())
    }
}

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

impl FrameWriter for Crypto {
    type Context = ();

    //= https://www.rfc-editor.org/rfc/rfc9000#section-19.6
    //# The stream does not have an explicit end, so CRYPTO frames do not
    //# have a FIN bit.
    const WRITES_FIN: bool = false;

    //= https://www.rfc-editor.org/rfc/rfc9002#section-6.2.4
    //# An endpoint SHOULD include new data in packets that are sent on PTO
    //# expiration.  Previously sent data MAY be sent if no new data can be
    //# sent.
    // Allow already transmitted, unacked crypto frames to be included in
    // probe packets in anticipation the crypto frames were lost. This will
    // help the handshake recover from packet loss.
    const RETRANSMIT_IN_PROBE: bool = true;

    fn write_chunk<W: WriteContext>(
        &self,
        offset: VarInt,
        data: &mut View,
        _writer_context: Self::Context,
        context: &mut W,
    ) -> Result<(), FitError> {
        let remaining_capacity = context.remaining_capacity();
        debug_assert!(
            data.len() <= remaining_capacity,
            "the data sender should not pass a payload that exceeds the current capacity"
        );

        // Some QUIC implementations refuse to process empty CRYPTO frames so
        // make sure we never send them
        debug_assert_ne!(data.len(), 0u64);

        let frame = frame::Crypto { offset, data };

        let len = frame.try_fit(remaining_capacity)?;
        if len == 0 {
            return Err(FitError);
        }

        frame.data.trim_off(frame.data.encoding_size() - len)?;

        context.write_fitted_frame(&frame);

        Ok(())
    }

    fn write_fin<W: WriteContext>(
        &self,
        _offset: VarInt,
        _writer_context: Self::Context,
        _context: &mut W,
    ) -> Result<(), FitError> {
        //= https://www.rfc-editor.org/rfc/rfc9000#section-19.6
        //# The stream does not have an explicit end, so CRYPTO frames do not
        //# have a FIN bit.
        // do nothing
        Ok(())
    }
}


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

use crate::{
    ack::AckManager,
    connection,
    contexts::WriteContext,
    endpoint, path,
    path::mtu,
    recovery,
    space::{datagram, HandshakeStatus},
    stream::Manager as _,
    sync::{flag, flag::Ping},
    transmission::{self, Mode},
};
use core::ops::RangeInclusive;
use s2n_quic_core::packet::number::PacketNumberSpace;

pub enum Payload<'a, Config: endpoint::Config> {
    Normal(Normal<'a, Config>),
    MtuProbe(MtuProbe<'a>),
    /// For use on non-active paths where only path validation frames are sent.
    PathValidationOnly(PathValidationOnly<'a, Config>),
}

impl<'a, Config: endpoint::Config> Payload<'a, Config> {
    /// Constructs a transmission::application::Payload appropriate for the given
    /// `transmission::Mode` in the given `ConnectionTransmissionContext`
    #[allow(clippy::too_many_arguments)]
    pub fn new(
        path_id: path::Id,
        path_manager: &'a mut path::Manager<Config>,
        local_id_registry: &'a mut connection::LocalIdRegistry,
        transmission_mode: transmission::Mode,
        ack_manager: &'a mut AckManager,
        handshake_status: &'a mut HandshakeStatus,
        ping: &'a mut flag::Ping,
        stream_manager: &'a mut Config::StreamManager,
        recovery_manager: &'a mut recovery::Manager<Config>,
        datagram_manager: &'a mut datagram::Manager<Config>,
    ) -> Self {
        if transmission_mode != Mode::PathValidationOnly {
            debug_assert_eq!(path_id, path_manager.active_path_id());
        }

        match transmission_mode {
            Mode::LossRecoveryProbing | Mode::Normal => {
                transmission::application::Payload::Normal(Normal {
                    ack_manager,
                    handshake_status,
                    ping,
                    stream_manager,
                    local_id_registry,
                    path_manager,
                    recovery_manager,
                    datagram_manager,
                    prioritize_datagrams: false,
                })
            }
            Mode::MtuProbing => transmission::application::Payload::MtuProbe(MtuProbe {
                mtu_controller: &mut path_manager[path_id].mtu_controller,
            }),
            Mode::PathValidationOnly => {
                transmission::application::Payload::PathValidationOnly(PathValidationOnly {
                    path: &mut path_manager[path_id],
                })
            }
        }
    }
}

impl<'a, Config: endpoint::Config> super::Payload for Payload<'a, Config> {
    fn size_hint(&self, range: RangeInclusive<usize>) -> usize {
        // We need at least 1 byte to write a HANDSHAKE_DONE or PING frame
        (*range.start()).max(1)
    }

    fn on_transmit<W: WriteContext>(&mut self, context: &mut W) {
        match self {
            Payload::Normal(inner) => inner.on_transmit(context),
            Payload::MtuProbe(inner) => inner.on_transmit(context),
            Payload::PathValidationOnly(inner) => inner.on_transmit(context),
        }
    }

    fn packet_number_space(&self) -> PacketNumberSpace {
        PacketNumberSpace::ApplicationData
    }
}

impl<'a, Config: endpoint::Config> transmission::interest::Provider for Payload<'a, Config> {
    fn transmission_interest<Q: transmission::interest::Query>(
        &self,
        query: &mut Q,
    ) -> transmission::interest::Result {
        match self {
            Payload::Normal(inner) => inner.transmission_interest(query),
            Payload::MtuProbe(inner) => inner.transmission_interest(query),
            Payload::PathValidationOnly(inner) => inner.transmission_interest(query),
        }
    }
}

pub struct Normal<'a, Config: endpoint::Config> {
    ack_manager: &'a mut AckManager,
    handshake_status: &'a mut HandshakeStatus,
    ping: &'a mut Ping,
    stream_manager: &'a mut Config::StreamManager,
    local_id_registry: &'a mut connection::LocalIdRegistry,
    path_manager: &'a mut path::Manager<Config>,
    recovery_manager: &'a mut recovery::Manager<Config>,
    datagram_manager: &'a mut datagram::Manager<Config>,
    prioritize_datagrams: bool,
}

impl<'a, Config: endpoint::Config> Normal<'a, Config> {
    fn on_transmit<W: WriteContext>(&mut self, context: &mut W) {
        let can_transmit = context.transmission_constraint().can_transmit()
            || context.transmission_constraint().can_retransmit();

        //= https://www.rfc-editor.org/rfc/rfc9221#section-5
        //# DATAGRAM frames cannot be fragmented;
        //
        // We alternate between prioritizing filling the packet with datagrams
        // and filling the packet with other frames. This is because datagrams
        // cannot be fragmented across packets and we want to do the most to send
        // large datagrams.
        if self.prioritize_datagrams && can_transmit {
            self.datagram_manager.on_transmit(
                context,
                self.stream_manager,
                self.prioritize_datagrams,
            );
        }
        let did_send_ack = self.ack_manager.on_transmit(context);

        // Payloads can only transmit and retransmit
        if can_transmit {
            self.transmit_control_data(context);

            // If we did not prioritize datagrams in this packet, we send them just
            // before we send stream data.
            if !self.prioritize_datagrams {
                self.datagram_manager.on_transmit(
                    context,
                    self.stream_manager,
                    self.prioritize_datagrams,
                );
            }

            // The default sending behavior is to alternate between sending datagrams
            // and sending stream data. This can be configured by implementing a
            // custom datagram sender and choosing when to cede packet space for stream data.
            let _ = self.stream_manager.on_transmit(context);

            // send PINGs last, since they might not actually be needed if there's an ack-eliciting
            // frame already present in the payload
            self.recovery_manager.on_transmit(context);
            let _ = self.ping.on_transmit(context);
        }

        if did_send_ack {
            // inform the ack manager the packet is populated
            self.ack_manager.on_transmit_complete(context);
        }

        // Alternate between prioritizing datagrams or not each packet
        self.prioritize_datagrams = !self.prioritize_datagrams;
    }

    // Sends control data frames
    fn transmit_control_data<W: WriteContext>(&mut self, context: &mut W) {
        // send HANDSHAKE_DONE frames first, if needed, to ensure the handshake is confirmed as
        // soon as possible
        self.handshake_status.on_transmit(context);

        //= https://www.rfc-editor.org/rfc/rfc9000#section-8.2
        //# An endpoint MAY include other frames with the PATH_CHALLENGE and
        //# PATH_RESPONSE frames used for path validation.
        // prioritize PATH_CHALLENGE and PATH_RESPONSE frames higher than app data
        self.path_manager.active_path_mut().on_transmit(context);

        self.local_id_registry.on_transmit(context);

        self.path_manager.on_transmit(context);
    }
}

impl<'a, Config: endpoint::Config> transmission::interest::Provider for Normal<'a, Config> {
    fn transmission_interest<Q: transmission::interest::Query>(
        &self,
        query: &mut Q,
    ) -> transmission::interest::Result {
        self.ack_manager.transmission_interest(query)?;
        self.handshake_status.transmission_interest(query)?;
        self.stream_manager.transmission_interest(query)?;
        self.datagram_manager.transmission_interest(query)?;
        self.local_id_registry.transmission_interest(query)?;
        self.path_manager.transmission_interest(query)?;
        self.recovery_manager.transmission_interest(query)?;
        self.path_manager
            .active_path()
            .transmission_interest(query)?;
        self.ping.transmission_interest(query)?;
        Ok(())
    }
}

pub struct MtuProbe<'a> {
    mtu_controller: &'a mut mtu::Controller,
}

impl<'a> MtuProbe<'a> {
    fn on_transmit<W: WriteContext>(&mut self, context: &mut W) {
        if context.transmission_constraint().can_transmit() {
            self.mtu_controller.on_transmit(context)
        }
    }
}

impl<'a> transmission::interest::Provider for MtuProbe<'a> {
    fn transmission_interest<Q: transmission::interest::Query>(
        &self,
        query: &mut Q,
    ) -> transmission::interest::Result {
        self.mtu_controller.transmission_interest(query)
    }
}

pub struct PathValidationOnly<'a, Config: endpoint::Config> {
    path: &'a mut path::Path<Config>,
}

impl<'a, Config: endpoint::Config> PathValidationOnly<'a, Config> {
    fn on_transmit<W: WriteContext>(&mut self, context: &mut W) {
        if context.transmission_constraint().can_transmit() {
            self.path.on_transmit(context)
        }
    }
}

impl<'a, Config: endpoint::Config> transmission::interest::Provider
    for PathValidationOnly<'a, Config>
{
    fn transmission_interest<Q: transmission::interest::Query>(
        &self,
        query: &mut Q,
    ) -> transmission::interest::Result {
        self.path.transmission_interest(query)
    }
}


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

use crate::transmission::{self, WriteContext};
use core::ops::RangeInclusive;
use s2n_codec::EncoderValue;
use s2n_quic_core::{frame, packet::number::PacketNumberSpace};

pub struct Payload<'a> {
    pub connection_close: &'a frame::ConnectionClose<'a>,
    pub packet_number_space: PacketNumberSpace,
}

impl<'a> super::Payload for Payload<'a> {
    fn size_hint(&self, range: RangeInclusive<usize>) -> usize {
        (*range.start()).max(self.connection_close.encoding_size())
    }

    fn on_transmit<W: WriteContext>(&mut self, context: &mut W) {
        context.write_frame(self.connection_close);
    }

    fn packet_number_space(&self) -> PacketNumberSpace {
        self.packet_number_space
    }
}

impl<'a> transmission::interest::Provider for Payload<'a> {
    #[inline]
    fn transmission_interest<Q: transmission::interest::Query>(
        &self,
        query: &mut Q,
    ) -> transmission::interest::Result {
        //= https://www.rfc-editor.org/rfc/rfc9002#section-3
        //# Packets containing frames besides ACK or CONNECTION_CLOSE frames
        //# count toward congestion control limits and are considered to be in
        //# flight.

        // this packet only contains a CONNECTION_CLOSE so bypass the CC
        query.on_forced()
    }
}


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

use crate::{contexts::WriteContext, endpoint, path, transmission, transmission::Mode};
use core::marker::PhantomData;
use s2n_codec::{Encoder, EncoderBuffer, EncoderValue};
use s2n_quic_core::{
    event::{self, ConnectionPublisher as _, IntoEvent},
    frame::{ack::AckRanges as AckRangesTrait, ack_elicitation::AckElicitation, Ack, FrameTrait},
    packet::number::PacketNumber,
    time::Timestamp,
};

pub struct Context<'a, 'b, 'sub, Config: endpoint::Config> {
    pub outcome: &'a mut transmission::Outcome,
    pub buffer: &'a mut EncoderBuffer<'b>,
    pub packet_number: PacketNumber,
    pub transmission_constraint: transmission::Constraint,
    pub transmission_mode: transmission::Mode,
    pub timestamp: Timestamp,
    pub header_len: usize,
    pub tag_len: usize,
    pub config: PhantomData<Config>,
    pub path_id: path::Id,
    pub publisher: &'a mut event::ConnectionPublisherSubscriber<
        'sub,
        <Config as endpoint::Config>::EventSubscriber,
    >,
}

impl<'a, 'b, 'sub, Config: endpoint::Config> Context<'a, 'b, 'sub, Config> {
    #[inline]
    fn check_frame_constraint<Frame: FrameTrait>(&self, frame: &Frame) {
        // only apply checks with debug_assertions enabled
        if !cfg!(debug_assertions) {
            return;
        }

        //= https://www.rfc-editor.org/rfc/rfc9000#section-9
        //# Servers do not send non-
        //# probing packets (see Section 9.1) toward a client address until they
        //# see a non-probing packet from that address.
        //
        // The transmission_mode PathValidation is used by the non-active path
        // to only transmit probing frames. A packet containing only probing
        // frames is also a probing packet.
        if self.transmission_mode == Mode::PathValidationOnly {
            assert!(frame.path_validation().is_probing());
        }

        match self.transmission_constraint() {
            transmission::Constraint::AmplificationLimited => {
                unreachable!("frames should not be written when we're amplification limited")
            }
            transmission::Constraint::CongestionLimited => {
                assert!(!frame.is_congestion_controlled());
            }
            transmission::Constraint::RetransmissionOnly => {}
            transmission::Constraint::None => {}
        }
    }
}

impl<'a, 'b, 'sub, Config: endpoint::Config> WriteContext for Context<'a, 'b, 'sub, Config> {
    fn current_time(&self) -> Timestamp {
        self.timestamp
    }

    #[inline]
    fn transmission_constraint(&self) -> transmission::Constraint {
        self.transmission_constraint
    }

    #[inline]
    fn transmission_mode(&self) -> Mode {
        self.transmission_mode
    }

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

    #[inline]
    fn write_ack_frame<AckRanges: AckRangesTrait>(
        &mut self,
        ack_frame: &Ack<AckRanges>,
    ) -> Option<PacketNumber> {
        let res = self.write_frame(ack_frame);
        if res.is_some() {
            for range in ack_frame.ack_ranges.ack_ranges() {
                self.publisher
                    .on_ack_range_sent(event::builder::AckRangeSent {
                        packet_header: event::builder::PacketHeader::new(
                            self.packet_number,
                            self.publisher.quic_version(),
                        ),
                        path_id: self.path_id.into_event(),
                        ack_range: range.start().into_event()..=range.end().into_event(),
                    });
            }
        }
        res
    }

    #[inline]
    fn write_frame<Frame>(&mut self, frame: &Frame) -> Option<PacketNumber>
    where
        Frame: EncoderValue + FrameTrait,
        for<'frame> &'frame Frame: IntoEvent<event::builder::Frame>,
    {
        self.check_frame_constraint(frame);
        self.write_frame_forced(frame)
    }

    #[inline]
    fn write_fitted_frame<Frame>(&mut self, frame: &Frame) -> PacketNumber
    where
        Frame: EncoderValue + FrameTrait,
        for<'frame> &'frame Frame: IntoEvent<event::builder::Frame>,
    {
        self.check_frame_constraint(frame);
        debug_assert!(frame.encoding_size() <= self.buffer.remaining_capacity());

        self.buffer.encode(frame);
        self.outcome.ack_elicitation |= frame.ack_elicitation();
        self.outcome.is_congestion_controlled |= frame.is_congestion_controlled();

        self.publisher.on_frame_sent(event::builder::FrameSent {
            packet_header: event::builder::PacketHeader::new(
                self.packet_number,
                self.publisher.quic_version(),
            ),
            path_id: self.path_id.into_event(),
            frame: frame.into_event(),
        });
        self.packet_number
    }

    fn write_frame_forced<Frame>(&mut self, frame: &Frame) -> Option<PacketNumber>
    where
        Frame: EncoderValue + FrameTrait,
        for<'frame> &'frame Frame: IntoEvent<event::builder::Frame>,
    {
        if frame.encoding_size() > self.buffer.remaining_capacity() {
            return None;
        }

        self.buffer.encode(frame);
        self.outcome.ack_elicitation |= frame.ack_elicitation();
        self.outcome.is_congestion_controlled |= frame.is_congestion_controlled();

        self.publisher.on_frame_sent(event::builder::FrameSent {
            packet_header: event::builder::PacketHeader::new(
                self.packet_number,
                self.publisher.quic_version(),
            ),
            path_id: self.path_id.into_event(),
            frame: frame.into_event(),
        });
        Some(self.packet_number)
    }

    #[inline]
    fn ack_elicitation(&self) -> AckElicitation {
        self.outcome.ack_elicitation
    }

    #[inline]
    fn packet_number(&self) -> PacketNumber {
        self.packet_number
    }

    #[inline]
    fn local_endpoint_type(&self) -> endpoint::Type {
        Config::ENDPOINT_TYPE
    }

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

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

// Overrides a context's transmission constraint to allow only retransmissions to be written to
// packets
pub struct RetransmissionContext<'a, C: WriteContext> {
    context: &'a mut C,
}

impl<'a, C: WriteContext> RetransmissionContext<'a, C> {
    pub fn new(context: &'a mut C) -> Self {
        Self { context }
    }
}

impl<'a, C: WriteContext> WriteContext for RetransmissionContext<'a, C> {
    #[inline]
    fn current_time(&self) -> Timestamp {
        self.context.current_time()
    }

    #[inline]
    fn transmission_constraint(&self) -> transmission::Constraint {
        debug_assert!(
            self.context.transmission_constraint().can_retransmit(),
            "retransmission ability should be checked before using RetransmissionContext"
        );

        transmission::Constraint::RetransmissionOnly
    }

    #[inline]
    fn transmission_mode(&self) -> Mode {
        self.context.transmission_mode()
    }

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

    #[inline]
    fn write_ack_frame<AckRanges: AckRangesTrait>(
        &mut self,
        ack_frame: &Ack<AckRanges>,
    ) -> Option<PacketNumber> {
        self.context.write_ack_frame(ack_frame)
    }

    #[inline]
    fn write_frame<Frame>(&mut self, frame: &Frame) -> Option<PacketNumber>
    where
        Frame: EncoderValue + FrameTrait,
        for<'frame> &'frame Frame: IntoEvent<event::builder::Frame>,
    {
        self.context.write_frame(frame)
    }

    #[inline]
    fn write_fitted_frame<Frame>(&mut self, frame: &Frame) -> PacketNumber
    where
        Frame: EncoderValue + FrameTrait,
        for<'frame> &'frame Frame: IntoEvent<event::builder::Frame>,
    {
        self.context.write_fitted_frame(frame)
    }

    fn write_frame_forced<Frame>(&mut self, frame: &Frame) -> Option<PacketNumber>
    where
        Frame: EncoderValue + FrameTrait,
        for<'frame> &'frame Frame: IntoEvent<event::builder::Frame>,
    {
        self.context.write_frame_forced(frame)
    }

    #[inline]
    fn ack_elicitation(&self) -> AckElicitation {
        self.context.ack_elicitation()
    }

    #[inline]
    fn packet_number(&self) -> PacketNumber {
        self.context.packet_number()
    }

    #[inline]
    fn local_endpoint_type(&self) -> endpoint::Type {
        self.context.local_endpoint_type()
    }

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

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


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

use crate::{
    ack::AckManager, contexts::WriteContext, endpoint, recovery, space::CryptoStream, transmission,
};
use core::ops::RangeInclusive;
use s2n_quic_core::{frame::Ping, packet::number::PacketNumberSpace};

pub struct Payload<'a, Config: endpoint::Config> {
    pub ack_manager: &'a mut AckManager,
    pub crypto_stream: &'a mut CryptoStream,
    pub packet_number_space: PacketNumberSpace,
    pub recovery_manager: &'a mut recovery::Manager<Config>,
}

/// Rather than creating a packet with a very small CRYPTO frame (under 16 bytes), it would be
/// better to wait for another transmission and send something larger. This should be better for
/// performance, anyway, since you end up paying for encryption/decryption.
const MIN_SIZE: usize = 16;

impl<'a, Config: endpoint::Config> super::Payload for Payload<'a, Config> {
    fn size_hint(&self, range: RangeInclusive<usize>) -> usize {
        (*range.start()).max(MIN_SIZE)
    }

    fn on_transmit<W: WriteContext>(&mut self, context: &mut W) {
        debug_assert!(
            !context.transmission_mode().is_mtu_probing(),
            "Early transmissions should not be used for MTU probing"
        );

        // record the starting capacity
        let start_capacity = context.remaining_capacity();

        let did_send_ack = self.ack_manager.on_transmit(context);

        // Payloads can only transmit and retransmit
        if context.transmission_constraint().can_transmit()
            || context.transmission_constraint().can_retransmit()
        {
            let _ = self.crypto_stream.tx.on_transmit((), context);

            // send PINGs last, since they might not actually be needed if there's an ack-eliciting
            // frame already present in the payload
            self.recovery_manager.on_transmit(context);

            // In order to trigger the loss recovery mechanisms during the handshake make all packets
            // ack-eliciting. This is especially true for the client in order to give the server more
            // amplification credits.
            //
            // Only send a PING if:
            // * We're not congestion limited
            // * The packet isn't already ack-eliciting
            // * Another frame was written to the context
            if !context.ack_elicitation().is_ack_eliciting()
                && start_capacity != context.remaining_capacity()
            {
                let _ = context.write_frame(&Ping);
            }
        }

        if did_send_ack {
            // inform the ack manager the packet is populated
            self.ack_manager.on_transmit_complete(context);
        }
    }

    fn packet_number_space(&self) -> PacketNumberSpace {
        self.packet_number_space
    }
}

impl<'a, Config: endpoint::Config> transmission::interest::Provider for Payload<'a, Config> {
    fn transmission_interest<Q: transmission::interest::Query>(
        &self,
        query: &mut Q,
    ) -> transmission::interest::Result {
        self.ack_manager.transmission_interest(query)?;
        self.crypto_stream.transmission_interest(query)?;
        self.recovery_manager.transmission_interest(query)?;
        Ok(())
    }
}


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

use crate::transmission::Constraint;

#[derive(Debug, Copy, Clone, PartialEq, Eq, PartialOrd, Ord)]
pub enum Interest {
    None,
    NewData,
    LostData,
    Forced,
}

#[test]
fn ordering_test() {
    assert!(Interest::None < Interest::NewData);
    assert!(Interest::NewData < Interest::LostData);
    assert!(Interest::LostData < Interest::Forced);
}

impl Default for Interest {
    fn default() -> Self {
        Self::None
    }
}

impl Interest {
    #[inline]
    pub fn can_transmit(self, limit: Constraint) -> bool {
        match (self, limit) {
            // nothing can be transmitted when we're at amplification limits
            (_, Constraint::AmplificationLimited) => false,

            // a component wants to try to recover so ignore limits
            (Interest::Forced, _) => true,

            // transmit lost data when we're either not limited, probing, or we want to do a fast
            // retransmission to try to recover
            (Interest::LostData, _) => limit.can_retransmit(),

            // new data may only be transmitted when we're not limited or probing
            (Interest::NewData, _) => limit.can_transmit(),

            // nothing is interested in transmitting anything
            (Interest::None, _) => false,
        }
    }

    #[inline]
    pub fn is_none(self) -> bool {
        matches!(self, Interest::None)
    }

    #[inline]
    pub fn merge_with<F: FnOnce(&mut Self) -> Result>(&mut self, f: F) -> Result {
        f(self)
    }
}

pub trait Provider {
    fn transmission_interest<Q: Query>(&self, query: &mut Q) -> Result;

    #[inline]
    fn get_transmission_interest(&self) -> Interest {
        let mut interest = Interest::None;
        let _ = self.transmission_interest(&mut interest);
        interest
    }

    #[inline]
    fn has_transmission_interest(&self) -> bool {
        let mut query = HasTransmissionInterestQuery;
        self.transmission_interest(&mut query).is_err()
    }

    #[inline]
    fn can_transmit(&self, mut constraint: Constraint) -> bool {
        self.transmission_interest(&mut constraint).is_err()
    }
}

pub trait Query {
    fn on_interest(&mut self, interest: Interest) -> Result;

    #[inline]
    fn on_new_data(&mut self) -> Result {
        self.on_interest(Interest::NewData)
    }

    #[inline]
    fn on_lost_data(&mut self) -> Result {
        self.on_interest(Interest::LostData)
    }

    #[inline]
    fn on_forced(&mut self) -> Result {
        self.on_interest(Interest::Forced)
    }
}

impl Query for Interest {
    #[inline]
    fn on_interest(&mut self, interest: Interest) -> Result {
        match (*self, interest) {
            // we don't need to keep querying if we're already at the max interest
            (Interest::Forced, _) | (_, Interest::Forced) => {
                *self = Interest::Forced;
                return Err(QueryBreak);
            }
            (Interest::LostData, _) | (_, Interest::LostData) => *self = Interest::LostData,
            (Interest::NewData, _) | (_, Interest::NewData) => *self = Interest::NewData,
            (Interest::None, _) => {}
        }

        Ok(())
    }
}

#[test]
fn interest_query_test() {
    use Interest::*;

    let levels = [None, NewData, LostData, Forced];
    for a in levels.iter().copied() {
        for b in levels.iter().copied() {
            let mut query = a;
            let result = query.on_interest(b);

            assert_eq!(query, a.max(b));
            assert_eq!(matches!(a, Forced) || matches!(b, Forced), result.is_err());
        }
    }
}

impl Query for Constraint {
    #[inline]
    fn on_interest(&mut self, interest: Interest) -> Result {
        // If we can transmit with the given constraint bail since we now have an answer
        if interest.can_transmit(*self) {
            return Err(QueryBreak);
        }

        Ok(())
    }
}

pub struct HasTransmissionInterestQuery;

impl Query for HasTransmissionInterestQuery {
    #[inline]
    fn on_interest(&mut self, interest: Interest) -> Result {
        if interest.is_none() {
            Ok(())
        } else {
            // If we've got anything other than `None` then bail since we now have an answer
            Err(QueryBreak)
        }
    }

    // any calls to interest should bail the query
    #[inline]
    fn on_new_data(&mut self) -> Result {
        Err(QueryBreak)
    }

    #[inline]
    fn on_lost_data(&mut self) -> Result {
        Err(QueryBreak)
    }

    #[inline]
    fn on_forced(&mut self) -> Result {
        Err(QueryBreak)
    }
}

pub struct Debugger;

impl Query for Debugger {
    #[inline]
    #[track_caller]
    fn on_interest(&mut self, interest: Interest) -> Result {
        eprintln!("  {} - {:?}", core::panic::Location::caller(), interest);
        Ok(())
    }

    #[inline]
    #[track_caller]
    fn on_new_data(&mut self) -> Result {
        eprintln!(
            "  {} - {:?}",
            core::panic::Location::caller(),
            Interest::NewData
        );
        Ok(())
    }

    #[inline]
    #[track_caller]
    fn on_lost_data(&mut self) -> Result {
        eprintln!(
            "  {} - {:?}",
            core::panic::Location::caller(),
            Interest::LostData
        );
        Ok(())
    }

    #[inline]
    #[track_caller]
    fn on_forced(&mut self) -> Result {
        eprintln!(
            "  {} - {:?}",
            core::panic::Location::caller(),
            Interest::Forced
        );
        Ok(())
    }
}

pub struct QueryBreak;

pub type Result<T = (), E = QueryBreak> = core::result::Result<T, E>;

#[cfg(test)]
mod test {
    use crate::transmission::{
        Constraint,
        Constraint::*,
        Interest::{None, *},
    };

    #[test]
    fn can_transmit() {
        // Amplification Limited
        assert!(!None.can_transmit(AmplificationLimited));
        assert!(!NewData.can_transmit(AmplificationLimited));
        assert!(!LostData.can_transmit(AmplificationLimited));
        assert!(!Forced.can_transmit(AmplificationLimited));

        // Congestion Limited
        assert!(!None.can_transmit(CongestionLimited));
        assert!(!NewData.can_transmit(CongestionLimited));
        assert!(!LostData.can_transmit(CongestionLimited));
        assert!(Forced.can_transmit(CongestionLimited));

        // Retransmission Only
        assert!(!None.can_transmit(RetransmissionOnly));
        assert!(!NewData.can_transmit(RetransmissionOnly));
        assert!(LostData.can_transmit(RetransmissionOnly));
        assert!(Forced.can_transmit(RetransmissionOnly));

        // No Constraint
        assert!(!None.can_transmit(Constraint::None));
        assert!(NewData.can_transmit(Constraint::None));
        assert!(LostData.can_transmit(Constraint::None));
        assert!(Forced.can_transmit(Constraint::None));
    }
}


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

pub mod context;
use context::Context;

pub mod application;
pub mod connection_close;
pub mod early;
pub mod interest;

pub use crate::contexts::WriteContext;
pub use interest::Interest;

/// re-export core
pub use s2n_quic_core::transmission::*;

use crate::{
    endpoint, path,
    space::TxPacketNumbers,
    transmission::{self, interest::Provider as _},
};
use core::{marker::PhantomData, ops::RangeInclusive};
use s2n_codec::{Encoder, EncoderBuffer};
use s2n_quic_core::{
    event,
    frame::Padding,
    packet::{
        encoding::PacketPayloadEncoder,
        number::{PacketNumber, PacketNumberSpace},
        stateless_reset,
    },
    time::Timestamp,
};

pub trait Payload: interest::Provider {
    fn size_hint(&self, payload_range: RangeInclusive<usize>) -> usize;
    fn on_transmit<W: WriteContext>(&mut self, context: &mut W);
    fn packet_number_space(&self) -> PacketNumberSpace;
}

pub struct Transmission<'a, 'sub, Config: endpoint::Config, P: Payload> {
    pub config: PhantomData<Config>,
    pub outcome: &'a mut transmission::Outcome,
    pub payload: P,
    pub packet_number: PacketNumber,
    pub timestamp: Timestamp,
    pub transmission_constraint: transmission::Constraint,
    pub transmission_mode: transmission::Mode,
    pub tx_packet_numbers: &'a mut TxPacketNumbers,
    pub path_id: path::Id,
    pub publisher: &'a mut event::ConnectionPublisherSubscriber<
        'sub,
        <Config as endpoint::Config>::EventSubscriber,
    >,
    pub packet_interceptor: &'a mut <Config as endpoint::Config>::PacketInterceptor,
}

impl<'a, 'sub, Config: endpoint::Config, P: Payload> PacketPayloadEncoder
    for Transmission<'a, 'sub, Config, P>
{
    fn encoding_size_hint<E: Encoder>(&mut self, encoder: &E, minimum_len: usize) -> usize {
        if self.has_transmission_interest() {
            self.payload.size_hint(minimum_len..=encoder.capacity())
        } else {
            0
        }
    }

    fn encode(
        &mut self,
        buffer: &mut EncoderBuffer,
        minimum_len: usize,
        header_len: usize,
        tag_len: usize,
    ) {
        debug_assert!(
            buffer.is_empty(),
            "the implementation assumes an empty buffer"
        );

        let mut context: Context<Config> = Context {
            outcome: self.outcome,
            buffer,
            packet_number: self.packet_number,
            transmission_constraint: self.transmission_constraint,
            transmission_mode: self.transmission_mode,
            timestamp: self.timestamp,
            header_len,
            tag_len,
            config: Default::default(),
            path_id: self.path_id,
            publisher: self.publisher,
        };

        self.payload.on_transmit(&mut context);

        if !context.buffer.is_empty() {
            // Add padding up to minimum_len
            let mut length = minimum_len.saturating_sub(context.buffer.len());

            // if we've only got a few bytes left in the buffer may as well pad it to full
            // capacity
            let remaining_capacity = context.buffer.remaining_capacity();
            if remaining_capacity < stateless_reset::min_indistinguishable_packet_len(tag_len) {
                length = remaining_capacity;
            }

            if length > 0 {
                context.write_frame(&Padding { length });
            }

            {
                use s2n_quic_core::{
                    event::ConnectionPublisher,
                    packet::interceptor::{Interceptor, Packet},
                };

                // intercept the payload before it is encrypted
                self.packet_interceptor.intercept_tx_payload(
                    &self.publisher.subject(),
                    &Packet {
                        number: self.packet_number,
                        timestamp: self.timestamp,
                    },
                    buffer,
                );
            }

            self.tx_packet_numbers.on_transmit(self.packet_number);
            self.outcome.bytes_sent = header_len + tag_len + buffer.len();
        }
    }
}

impl<'a, 'sub, Config: endpoint::Config, P: Payload> transmission::interest::Provider
    for Transmission<'a, 'sub, Config, P>
{
    fn transmission_interest<Q: transmission::interest::Query>(
        &self,
        query: &mut Q,
    ) -> transmission::interest::Result {
        self.payload.transmission_interest(query)
    }
}


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

use s2n_quic::{Client, Server};
use s2n_quic_core::stream::testing::Data;

#[derive(Clone, Copy, Debug, PartialEq)]
struct Snapshot {
    total: u64,
    rss: u64,
    max: u64,
}

impl Snapshot {
    pub fn new() -> Self {
        let stats = dhat::HeapStats::get();
        Self {
            total: stats.total_bytes,
            rss: stats.curr_bytes as _,
            max: stats.max_bytes as _,
        }
    }

    pub fn print_diff(&self, event: &str, streams: usize) {
        let (alloc, rss) = self.diff(Self::new());

        // make some assertions about the amount of memory use
        if streams > 0 {
            match event {
                "post-handshake" => {
                    assert!(rss < 12_000, "{rss}");
                }
                "post-transfer" => {
                    assert!(rss < 30_000, "{rss}");
                }
                "post-close" => {
                    assert!(rss < 128, "{rss}");
                }
                e => unimplemented!("{}", e),
            }
        }

        println!("{event}\t{alloc}\t{rss}\t{streams}");
    }

    pub fn diff(&self, other: Self) -> (u64, i64) {
        let alloc = self.alloc_diff(other);
        let rss = self.rss_diff(other);
        (alloc, rss)
    }

    pub fn rss_diff(&self, other: Self) -> i64 {
        let before = self.rss as i64;
        let after = other.rss as i64;
        after - before
    }

    pub fn alloc_diff(&self, other: Self) -> u64 {
        other.total - self.total
    }
}

#[global_allocator]
static ALLOC: dhat::Alloc = dhat::Alloc;

type Error = Box<dyn std::error::Error>;
type Result<T = (), E = Error> = core::result::Result<T, E>;

fn main() -> Result {
    let mut args = std::env::args();
    let _ = args.next();
    let arg = args.next();
    let _profiler = dhat::Profiler::new_heap();
    run(arg.as_deref())
}

#[tokio::main]
async fn run(arg: Option<&str>) -> Result {
    match arg {
        Some("server") => server().await,
        Some("client") => client().await,
        _ => Err("memory-report server|client".into()),
    }
}

async fn client() -> Result {
    let io = ("0.0.0.0", 0);

    let tls = s2n_quic_core::crypto::tls::testing::certificates::CERT_PEM;

    let client = Client::builder()
        .with_io(io)?
        .with_tls(tls)?
        .start()
        .unwrap();

    println!("event\talloc_diff\trss_diff\tstreams");

    for stream_count in 0..10 {
        // wait for a bit to have the allocations settle
        tokio::time::sleep(core::time::Duration::from_millis(1000)).await;

        let snapshot = Snapshot::new();
        let connect = s2n_quic::client::Connect::new(("127.0.0.1".parse()?, 4433))
            .with_server_name("localhost");

        let mut connection = client.connect(connect).await?;

        // wait for a bit to have the allocations settle
        tokio::time::sleep(core::time::Duration::from_millis(1000)).await;

        snapshot.print_diff("post-handshake", stream_count);

        for _ in 0..stream_count {
            let mut stream = connection.open_bidirectional_stream().await?;

            let mut data = Data::new(5 * 1_000_000);

            while let Some(chunk) = data.send_one(usize::MAX) {
                stream.send(chunk).await?;
            }

            stream.close().await?;
        }

        tokio::time::sleep(core::time::Duration::from_millis(5000)).await;

        snapshot.print_diff("post-transfer", stream_count);

        connection.close(123u8.into());
        drop(connection);

        tokio::time::sleep(core::time::Duration::from_millis(5000)).await;

        snapshot.print_diff("post-close", stream_count);
    }

    Ok(())
}

async fn server() -> Result {
    let io = ("127.0.0.1", 4433);

    let tls = (
        s2n_quic_core::crypto::tls::testing::certificates::CERT_PEM,
        s2n_quic_core::crypto::tls::testing::certificates::KEY_PEM,
    );

    let mut server = Server::builder()
        .with_io(io)?
        .with_tls(tls)?
        .start()
        .unwrap();

    eprintln!("Server listening on port {}", io.1);

    while let Some(mut connection) = server.accept().await {
        tokio::spawn(async move {
            while let Ok(Some(mut stream)) = connection.accept_bidirectional_stream().await {
                tokio::spawn(async move {
                    while let Ok(Some(data)) = stream.receive().await {
                        let _ = data;
                    }
                });
            }
        });
    }

    Ok(())
}


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

#![no_std]
#![no_main]

use aya_bpf::{
    bindings::xdp_action,
    macros::{map, xdp},
    maps::{HashMap, XskMap},
    programs::XdpContext,
};
use s2n_quic_core::{
    inet::udp,
    xdp::{
        bpf::DecoderBufferMut,
        decoder::{self, EventHandler},
    },
};

#[map(name = "S2N_QUIC_XDP_SOCKETS")]
static SOCKETS: XskMap = XskMap::with_max_entries(1024, 0);

#[map(name = "S2N_QUIC_XDP_PORTS")]
static PORTS: HashMap<u16, u8> = HashMap::with_max_entries(1024, 0);

#[xdp(name = "s2n_quic_xdp")]
pub fn s2n_quic_xdp(ctx: XdpContext) -> u32 {
    let action = handle_packet(&ctx);

    #[cfg(feature = "trace")]
    {
        use aya_log_ebpf as log;
        match action {
            xdp_action::XDP_DROP => log::trace!(&ctx, "ACTION: DROP"),
            xdp_action::XDP_PASS => log::trace!(&ctx, "ACTION: PASS"),
            xdp_action::XDP_REDIRECT => log::trace!(&ctx, "ACTION: REDIRECT"),
            xdp_action::XDP_ABORTED => log::trace!(&ctx, "ACTION: ABORTED"),
            _ => (),
        }
    }

    action
}

#[inline(always)]
fn handle_packet(ctx: &XdpContext) -> u32 {
    let start = ctx.data() as *mut u8;
    let end = ctx.data_end() as *mut u8;
    let buffer = unsafe {
        // Safety: start and end come from the caller and have been validated
        DecoderBufferMut::new(start, end)
    };
    match Validator.decode_packet(buffer) {
        Ok(Some(payload)) => {
            // if the payload is empty there isn't much we can do with it
            if payload.is_empty() {
                return xdp_action::XDP_DROP;
            }

            // if the packet is valid forward it on to the associated AF_XDP socket
            let queue_id = unsafe { (*ctx.ctx).rx_queue_index };
            let not_found_action = xdp_action::XDP_PASS as _;
            SOCKETS.redirect(queue_id, not_found_action)
        }
        Ok(None) => xdp_action::XDP_PASS,
        Err(_) => xdp_action::XDP_ABORTED,
    }
}

struct Validator;

impl EventHandler for Validator {
    #[inline(always)]
    fn on_udp_header(&mut self, header: &udp::Header) -> decoder::Result {
        // Make sure the port is in the port map. Otherwise, forward the packet to the OS.
        if PORTS.get_ptr(&header.destination().get()).is_some() {
            Ok(Some(()))
        } else {
            Ok(None)
        }
    }
}

/// Define a no-op panic handler
///
/// The implementation shouldn't panic at all. But we do need to define one in
/// `#[no_std]` builds.
#[panic_handler]
fn panic(_info: &core::panic::PanicInfo) -> ! {
    unsafe { core::hint::unreachable_unchecked() }
}


/* automatically generated by rust-bindgen 0.65.1 */

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

#![allow(non_camel_case_types)]

pub const ETHTOOL_GCHANNELS: u32 = 60;
pub type __u32 = ::std::os::raw::c_uint;
#[repr(C)]
#[derive(Debug, Copy, Clone)]
pub struct ethtool_channels {
    pub cmd: __u32,
    pub max_rx: __u32,
    pub max_tx: __u32,
    pub max_other: __u32,
    pub max_combined: __u32,
    pub rx_count: __u32,
    pub tx_count: __u32,
    pub other_count: __u32,
    pub combined_count: __u32,
}


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

use aya::include_bytes_aligned;

/// The default BPF program to direct QUIC traffic
pub static DEFAULT_PROGRAM: &[u8] = {
    #[cfg(target_endian = "little")]
    let prog = include_bytes_aligned!(concat!(
        env!("CARGO_MANIFEST_DIR"),
        "/src/bpf/s2n-quic-xdp-bpfel.ebpf"
    ));

    #[cfg(target_endian = "big")]
    let prog = include_bytes_aligned!(concat!(
        env!("CARGO_MANIFEST_DIR"),
        "/src/bpf/s2n-quic-xdp-bpfeb.ebpf"
    ));

    prog
};

/// The default BPF program to direct QUIC traffic with tracing enabled
pub static DEFAULT_PROGRAM_TRACE: &[u8] = {
    #[cfg(target_endian = "little")]
    let prog = include_bytes_aligned!(concat!(
        env!("CARGO_MANIFEST_DIR"),
        "/src/bpf/s2n-quic-xdp-bpfel-trace.ebpf"
    ));

    #[cfg(target_endian = "big")]
    let prog = include_bytes_aligned!(concat!(
        env!("CARGO_MANIFEST_DIR"),
        "/src/bpf/s2n-quic-xdp-bpfeb-trace.ebpf"
    ));

    prog
};

/// The name of the default XDP program
pub static PROGRAM_NAME: &str = "s2n_quic_xdp";

/// The name of the AF_XDP socket map
pub static XSK_MAP_NAME: &str = "S2N_QUIC_XDP_SOCKETS";

/// The name of the port map
pub static PORT_MAP_NAME: &str = "S2N_QUIC_XDP_PORTS";


// Copyright Intel Corporation
// SPDX-License-Identifier: GPL-2.0 WITH Linux-syscall-note
// Modifications copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
// SPDX-License-Identifier: Apache-2.0

use crate::Result;
use bitflags::bitflags;
use core::mem::size_of;
use std::{ffi::CStr, os::unix::io::AsRawFd};

bitflags!(
    /// Options for the `flags` field in [`Address`]
    ///
    /// See [if_xdp.h](https://github.com/torvalds/linux/blob/2bac7dc169af3cd4a0cb5200aa1f7b89affa042a/include/uapi/linux/if_xdp.h#L16-L27)
    ///
    /// Note that the size of the value is specified in the [sxdp_flags field](
    /// https://github.com/torvalds/linux/blob/0d3eb744aed40ffce820cded61d7eac515199165/include/uapi/linux/if_xdp.h#L34).
    #[derive(Copy, Clone, Debug, Default)]
    #[repr(transparent)]
    pub struct XdpFlags: u16 {
        const SHARED_UMEM = 1 << 0;
        /// Force copy mode
        const COPY = 1 << 1;
        /// Force zero-copy mode
        const ZEROCOPY = 1 << 2;
        /// If this option is set, the driver might go to sleep and in that case the
        /// XDP_RING_NEED_WAKEUP flag in the fill and/or Tx rings will be set.
        ///
        /// If it is set, the application needs to explicitly wake up the driver with a `poll()` for
        /// Rx or `sendto()` for Tx. If you are running the driver and the application on the same
        /// core, you should use this option so that the kernel will yield to the user space
        /// application.
        const USE_NEED_WAKEUP = 1 << 3;
    }
);

bitflags!(
    /// Flags for the umem config
    ///
    /// See [if_xdp.h](https://github.com/torvalds/linux/blob/2bac7dc169af3cd4a0cb5200aa1f7b89affa042a/include/uapi/linux/if_xdp.h#L30>)
    ///
    /// Note that in `if_xdp.h`, the size is left unspecified. However, when it's used in libxdp,
    /// the size is set to `u32`: [xsk.h](https://github.com/xdp-project/xdp-tools/blob/a76e7a2b156b8cfe38992206abe9df1df0a29e38/headers/xdp/xsk.h#L203).
    #[derive(Copy, Clone, Debug, Default)]
    #[repr(transparent)]
    pub struct UmemFlags: u32 {
        const UNALIGNED_CHUNK_FLAG = 1 << 0;
    }
);

/// A structure for representing the address of an AF_XDP socket
///
/// See [if_xdp.h](https://github.com/torvalds/linux/blob/2bac7dc169af3cd4a0cb5200aa1f7b89affa042a/include/uapi/linux/if_xdp.h#L32-L38)
#[derive(Clone, Copy, Debug)]
#[repr(C)]
pub struct Address {
    pub family: u16,
    pub flags: XdpFlags,
    pub ifindex: u32,
    pub queue_id: u32,
    pub shared_umem_fd: u32,
}

impl Default for Address {
    fn default() -> Self {
        Self {
            family: libc::PF_XDP as _,
            flags: Default::default(),
            ifindex: 0,
            queue_id: 0,
            shared_umem_fd: 0,
        }
    }
}

impl Address {
    /// Resolves the `ifindex` for the provided `if_name` and associates it with the address.
    ///
    /// If the device does not exist, then an error is returned.
    #[inline]
    pub fn set_if_name(&mut self, name: &CStr) -> Result<&mut Self> {
        self.ifindex = crate::syscall::if_nametoindex(name)?;
        Ok(self)
    }

    /// Configures the address with a shared UMEM region from a previous socket file descriptor
    pub fn set_shared_umem<F: AsRawFd>(&mut self, fd: &F) -> &mut Self {
        // clear out all of the previous flags except the shared - the others will be inherited
        self.flags = XdpFlags::SHARED_UMEM;
        self.shared_umem_fd = fd.as_raw_fd() as _;
        self
    }
}

bitflags!(
    /// Flags set on a particular ring state `flags` field
    ///
    /// See [if_xdp.h](https://github.com/torvalds/linux/blob/2bac7dc169af3cd4a0cb5200aa1f7b89affa042a/include/uapi/linux/if_xdp.h#L41)
    ///
    /// Note that in `if_xdp.h`, the size is left unspecified. However, when it's used in libxdp,
    /// the size is set to `u32`: [xsk.h](
    /// https://github.com/xdp-project/xdp-tools/blob/a76e7a2b156b8cfe38992206abe9df1df0a29e38/headers/xdp/xsk.h#L42)
    #[derive(Default)]
    #[repr(transparent)]
    pub struct RingFlags: u32 {
        const NEED_WAKEUP = 1 << 0;
    }
);

/// A structure to communicate the offsets for an individual ring for an AF_XDP socket
///
/// Note that the `flags` field was added in kernel 5.4.
///
/// See [if_xdp.h](https://github.com/torvalds/linux/blob/9116e5e2b1fff71dce501d971e86a3695acc3dba/include/uapi/linux/if_xdp.h#L28-L32)
#[derive(Clone, Copy, Debug)]
#[repr(C)]
pub struct RingOffsetV1 {
    pub producer: u64,
    pub consumer: u64,
    pub desc: u64,
}

/// A structure to communicate the offsets for an individual ring for an AF_XDP socket
///
/// See [if_xdp.h](https://github.com/torvalds/linux/blob/2bac7dc169af3cd4a0cb5200aa1f7b89affa042a/include/uapi/linux/if_xdp.h#L43)
#[derive(Clone, Copy, Debug, Default)]
#[repr(C)]
pub struct RingOffsetV2 {
    pub producer: u64,
    pub consumer: u64,
    pub desc: u64,
    pub flags: u64,
}

impl RingOffsetV2 {
    /// Converts an offset description from an older version
    #[inline]
    fn set_v1(&mut self, v1: RingOffsetV1) {
        // Logic from https://github.com/xdp-project/xdp-tools/blob/a76e7a2b156b8cfe38992206abe9df1df0a29e38/lib/libxdp/xsk.c#L197
        self.producer = v1.producer;
        self.consumer = v1.consumer;
        self.desc = v1.desc;
        self.flags = v1.consumer + (size_of::<u32>() as u64);
    }
}

/// A structure to communicate the offsets of each of the 4 rings for an AF_XDP socket
///
/// See [if_xdp.h](https://github.com/torvalds/linux/blob/2bac7dc169af3cd4a0cb5200aa1f7b89affa042a/include/uapi/linux/if_xdp.h#L50)
#[derive(Clone, Copy, Debug, Default)]
#[repr(C)]
pub struct MmapOffsets<O = RingOffsetV2> {
    pub rx: O,
    pub tx: O,
    pub fill: O,
    pub completion: O,
}

impl MmapOffsets {
    /// See [if_xdp.h](https://github.com/torvalds/linux/blob/2bac7dc169af3cd4a0cb5200aa1f7b89affa042a/include/uapi/linux/if_xdp.h#L92)
    pub const RX_RING: usize = 0;
    /// See [if_xdp.h](https://github.com/torvalds/linux/blob/2bac7dc169af3cd4a0cb5200aa1f7b89affa042a/include/uapi/linux/if_xdp.h#L93)
    pub const TX_RING: usize = 0x8000_0000;
    /// See [if_xdp.h](https://github.com/torvalds/linux/blob/2bac7dc169af3cd4a0cb5200aa1f7b89affa042a/include/uapi/linux/if_xdp.h#L94)
    pub const FILL_RING: usize = 0x1_0000_0000;
    /// See [if_xdp.h](https://github.com/torvalds/linux/blob/2bac7dc169af3cd4a0cb5200aa1f7b89affa042a/include/uapi/linux/if_xdp.h#L95)
    pub const COMPLETION_RING: usize = 0x1_8000_0000;

    #[inline]
    pub(crate) fn as_v1(mut self) -> Self {
        // getsockopt on a kernel <= 5.3 has no flags fields.
        // Copy over the offsets to the correct places in the >=5.4 format
        // and put the flags where they would have been on that kernel.

        // Logic from https://github.com/xdp-project/xdp-tools/blob/a76e7a2b156b8cfe38992206abe9df1df0a29e38/lib/libxdp/xsk.c#L197
        let v1 = unsafe { *(&self as *const Self as *const MmapOffsets<RingOffsetV1>) };
        self.rx.set_v1(v1.rx);
        self.tx.set_v1(v1.tx);
        self.fill.set_v1(v1.fill);
        self.completion.set_v1(v1.completion);

        self
    }
}

/// Socket option ids passed to an AF_XDP socket with the sockopt syscalls.
///
/// See [if_xdp.h](https://github.com/torvalds/linux/blob/2bac7dc169af3cd4a0cb5200aa1f7b89affa042a/include/uapi/linux/if_xdp.h#L58-L65)
///
/// Note that in `if_xdp.h`, the size is left unspecified. However, it is used in the `name`
/// argument for [`libc::setsockopt`] and [`libc::getsockopt`], which is a [`libc::c_int`] (i32).
#[derive(Clone, Copy, Debug)]
#[repr(i32)]
pub enum SocketOptions {
    MmapOffsets = 1,
    RxRing = 2,
    TxRing = 3,
    UmemReg = 4,
    UmemFillRing = 5,
    UmemCompletionRing = 6,
    Statistics = 7,
    Options = 8,
}

/// Umem configuration on an AF_XDP socket
///
/// This struct is used with the [`libc::setsockopt`] call with the `name` parameter set to
/// [`SocketOptions::UmemReg`].
///
/// See [if_xdp.h](https://github.com/torvalds/linux/blob/2bac7dc169af3cd4a0cb5200aa1f7b89affa042a/include/uapi/linux/if_xdp.h#L67-L73).
#[derive(Clone, Copy, Debug)]
#[repr(C)]
pub struct UmemReg {
    /// Start of packet data area
    pub addr: u64,
    /// Length of packet data area
    pub len: u64,
    pub chunk_size: u32,
    /// The length reserved before packets
    pub headroom: u32,
    /// The flag value for the Umem
    pub flags: UmemFlags,
}

/// Statistics returned from an AF_XDP socket
///
/// This struct is used with the [`libc::getsockopt`] call with the `name` parameter set to
/// [`SocketOptions::Statistics`].
///
/// See [if_xdp.h](https://github.com/torvalds/linux/blob/2bac7dc169af3cd4a0cb5200aa1f7b89affa042a/include/uapi/linux/if_xdp.h#L75-L82)
#[derive(Clone, Copy, Debug, Default)]
#[repr(C)]
pub struct Statistics {
    /// Dropped for other reasons
    pub rx_dropped: u64,
    /// Dropped due to invalid descriptor
    pub rx_invalid_descriptors: u64,
    /// Dropped due to invalid descriptor
    pub tx_invalid_descriptors: u64,
    /// Dropped due to rx ring being full
    pub rx_ring_full: u64,
    /// Failed to retrieve item from fill ring
    pub rx_fill_ring_empty_descriptors: u64,
    /// Failed to retrieve item from tx ring
    pub tx_ring_empty_descriptors: u64,
}

bitflags!(
    /// Options returned from an AF_XDP socket
    ///
    /// This struct is used with the [`libc::getsockopt`] call with the `name` parameter set to
    /// [`SocketOptions::Options`].
    ///
    /// See [if_xdp.h](https://github.com/torvalds/linux/blob/2bac7dc169af3cd4a0cb5200aa1f7b89affa042a/include/uapi/linux/if_xdp.h#L89)
    #[derive(Clone, Copy, Debug, Default)]
    #[repr(transparent)]
    pub struct XdpOptions: u32 {
        const ZEROCOPY = 1 << 0;
    }
);

/// Masks for unaligned chunks mode
///
/// See [if_xdp](https://github.com/torvalds/linux/blob/0d3eb744aed40ffce820cded61d7eac515199165/include/uapi/linux/if_xdp.h#L97-L100).
pub const XSK_UNALIGNED_BUF_ADDR_MASK: u64 = (1 << XSK_UNALIGNED_BUF_OFFSET_SHIFT) - 1;
pub const XSK_UNALIGNED_BUF_OFFSET_SHIFT: u64 = 48;

/// Rx/Tx descriptor
///
/// See [if_xdp.h](https://github.com/torvalds/linux/blob/2bac7dc169af3cd4a0cb5200aa1f7b89affa042a/include/uapi/linux/if_xdp.h#L103-L107)
#[derive(Clone, Copy, Debug, Default)]
#[repr(C)]
pub struct RxTxDescriptor {
    /// Offset into the umem where the packet starts
    pub address: u64,
    /// Length of the packet
    pub len: u32,
    /// Options set on the descriptor
    pub options: u32,
}

/// Umem Descriptor
///
/// See [if_xdp.h](https://github.com/torvalds/linux/blob/0d3eb744aed40ffce820cded61d7eac515199165/include/uapi/linux/if_xdp.h#L109).
#[derive(Clone, Copy, Debug, Default)]
#[repr(C)]
pub struct UmemDescriptor {
    /// Offset into the umem where the packet starts
    pub address: u64,
}

impl UmemDescriptor {
    /// Sets the length for the descriptor and converts it into a [`RxTxDescriptor`]
    #[inline]
    pub fn with_len(self, len: u32) -> RxTxDescriptor {
        RxTxDescriptor {
            address: self.address,
            len,
            options: 0,
        }
    }
}

impl From<RxTxDescriptor> for UmemDescriptor {
    #[inline]
    fn from(desc: RxTxDescriptor) -> Self {
        Self {
            address: desc.address,
        }
    }
}


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

pub mod router;
pub mod rx;
pub mod tx;

#[cfg(test)]
mod tests;


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

type Result<T = (), E = std::io::Error> = core::result::Result<T, E>;

/// Emits a log line if the `s2n_quic_xdp_trace` cfg option is enabled. Otherwise, the trace is a
/// no-op.
macro_rules! trace {
    ($($fmt:tt)*) => {{
        if cfg!(s2n_quic_xdp_trace) {
            let args = format!($($fmt)*);
            println!("{}:{}: {}", module_path!(), line!(), args);
        }
    }}
}

/// Low-level bindings to various linux userspace APIs
mod bindings;

/// Default BPF programs to direct QUIC traffic
pub mod bpf;
/// Primitive types for AF-XDP kernel APIs
pub mod if_xdp;
/// Implementations of the IO traits from [`s2n_quic_core::io`]
pub mod io;
/// Helpers for creating mmap'd regions
pub mod mmap;
/// Structures for tracking ring cursors and synchronizing with the kernel
pub mod ring;
/// Structure for opening and reference counting an AF-XDP socket
pub mod socket;
/// Helpers for making API calls to AF-XDP sockets
pub mod syscall;
/// A shared region of memory for holding frame (packet) data
pub mod umem;


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

use crate::{
    syscall::{mmap, munmap},
    Result,
};
use core::{
    ffi::c_void,
    ops::{Deref, DerefMut},
    ptr::NonNull,
};
use std::os::unix::io::RawFd;

/// A mmap'd region in memory
#[derive(Debug)]
pub struct Mmap {
    addr: NonNull<c_void>,
    len: usize,
}

/// Safety: Mmap pointer can be sent between threads
unsafe impl Send for Mmap {}

/// Safety: Mmap pointer can be shared between threads
unsafe impl Sync for Mmap {}

impl Mmap {
    /// Creates a new mmap'd region, with an optional file descriptor.
    #[inline]
    pub fn new(len: usize, offset: usize, fd: Option<RawFd>) -> Result<Self> {
        let addr = mmap(len, offset, fd)?;
        Ok(Self { addr, len })
    }

    /// Returns the raw address for the mmap region
    #[inline]
    pub fn addr(&self) -> NonNull<c_void> {
        self.addr
    }
}

impl Deref for Mmap {
    type Target = [u8];

    #[inline]
    fn deref(&self) -> &Self::Target {
        unsafe { core::slice::from_raw_parts(self.addr.as_ptr() as _, self.len) }
    }
}

impl DerefMut for Mmap {
    #[inline]
    fn deref_mut(&mut self) -> &mut Self::Target {
        unsafe { core::slice::from_raw_parts_mut(self.addr.as_ptr() as _, self.len) }
    }
}

impl Drop for Mmap {
    #[inline]
    fn drop(&mut self) {
        let _ = unsafe {
            // Safety: the len is the same value as on creation
            munmap(self.addr, self.len)
        };
    }
}

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

    #[test]
    fn mmap_test() {
        let mut area = Mmap::new(32, 0, None).unwrap();
        assert_eq!(area.len(), 32);
        let _ = &area[..];
        let _ = &mut area[..];
    }
}


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

use crate::{
    if_xdp::{MmapOffsets, RingFlags, RingOffsetV2, RxTxDescriptor, UmemDescriptor},
    mmap::Mmap,
    socket, syscall,
};
use core::{fmt, mem::size_of, ptr::NonNull};
use s2n_quic_core::sync::cursor::{self, Cursor};
use std::{io, os::unix::io::AsRawFd};

#[derive(Debug)]
#[allow(dead_code)] // we hold on to `area` and `socket` to ensure they live long enough
struct Ring<T: Copy + fmt::Debug> {
    cursor: Cursor<T>,
    flags: NonNull<RingFlags>,
    // make the area clonable in test mode
    #[cfg(test)]
    area: std::sync::Arc<Mmap>,
    #[cfg(not(test))]
    area: Mmap,
    socket: socket::Fd,
}

impl<T: Copy + fmt::Debug> Ring<T> {
    /// Returns `true` if the ring needs to be notified when entries are updated
    ///
    /// See [xsk.h](https://github.com/xdp-project/xdp-tools/blob/a76e7a2b156b8cfe38992206abe9df1df0a29e38/headers/xdp/xsk.h#L87).
    #[inline]
    pub fn needs_wakeup(&self) -> bool {
        self.flags().contains(RingFlags::NEED_WAKEUP)
    }

    /// Returns a reference to the flags on the ring
    #[inline]
    pub fn flags(&self) -> &RingFlags {
        unsafe { &*self.flags.as_ptr() }
    }

    /// Returns a mutable reference to the flags on the ring
    #[inline]
    #[cfg(test)]
    pub fn flags_mut(&mut self) -> &mut RingFlags {
        unsafe { &mut *self.flags.as_ptr() }
    }
}

/// Safety: the Mmap area is held for as long as the Cursor
unsafe impl<T: Copy + fmt::Debug> Send for Ring<T> {}

/// Safety: the Mmap area is held for as long as the Cursor
unsafe impl<T: Copy + fmt::Debug> Sync for Ring<T> {}

#[inline]
unsafe fn builder<T: Copy>(
    area: &Mmap,
    offsets: &RingOffsetV2,
    size: u32,
) -> (cursor::Builder<T>, NonNull<RingFlags>) {
    let producer = area.addr().as_ptr().add(offsets.producer as _);
    let producer = NonNull::new_unchecked(producer as _);
    let consumer = area.addr().as_ptr().add(offsets.consumer as _);
    let consumer = NonNull::new_unchecked(consumer as _);

    let flags = area.addr().as_ptr().add(offsets.flags as _);
    let flags = NonNull::new_unchecked(flags as *mut RingFlags);

    let descriptors = area.addr().as_ptr().add(offsets.desc as _);
    let descriptors = NonNull::new_unchecked(descriptors as *mut T);

    let builder = cursor::Builder {
        producer,
        consumer,
        data: descriptors,
        size,
    };

    (builder, flags)
}

macro_rules! impl_producer {
    ($T:ty, $syscall:ident, $field:ident, $offset:ident) => {
        /// Creates a new ring with the given configuration
        pub fn new(socket: socket::Fd, offsets: &MmapOffsets, size: u32) -> io::Result<Self> {
            syscall::$syscall(&socket, size)?;
            let offsets = &offsets.$field;

            // start with the descriptor offset as the total length
            let mut len = offsets.desc as usize;
            // extend the length by the `size` multiplied the entry size
            len += size as usize * size_of::<$T>();

            // Use the hard-coded offset of the ring type
            let offset = MmapOffsets::$offset;

            let area = Mmap::new(len, offset, Some(socket.as_raw_fd()))?;

            let (cursor, flags) = unsafe {
                // Safety: `area` lives as long as `cursor`
                let (builder, flags) = builder(&area, offsets, size);
                (builder.build_producer(), flags)
            };

            #[cfg(test)]
            let area = std::sync::Arc::new(area);

            Ok(Self(Ring {
                cursor,
                flags,
                area,
                socket,
            }))
        }

        /// Acquire a number of entries for the ring
        ///
        /// If the cached length is lower than the provided `watermark` no synchronization will be
        /// performed.
        #[inline]
        pub fn acquire(&mut self, watermark: u32) -> u32 {
            self.0.cursor.acquire_producer(watermark)
        }

        /// Releases `len` number of entries to the consumer side of the ring.
        ///
        /// # Panics
        ///
        /// If the `len` exceeds the number of acquired entries, this function will panic.
        #[inline]
        pub fn release(&mut self, len: u32) {
            self.0.cursor.release_producer(len)
        }

        /// Returns `true` if the ring needs to be woken up in order to notify the kernel
        #[inline]
        pub fn needs_wakeup(&self) -> bool {
            self.0.needs_wakeup()
        }

        /// Returns the unfilled entries for the producer
        ///
        /// After filling `n` entries, the `release` function should be called with `n`.
        #[inline]
        pub fn data(&mut self) -> (&mut [$T], &mut [$T]) {
            unsafe {
                // Safety: this is the producer for the ring
                self.0.cursor.producer_data()
            }
        }

        /// Returns the overall size of the ring
        #[inline]
        pub fn capacity(&self) -> usize {
            self.0.cursor.capacity() as _
        }

        #[inline]
        pub fn len(&self) -> u32 {
            self.0.cursor.cached_producer_len()
        }

        #[inline]
        pub fn is_empty(&self) -> bool {
            self.0.cursor.cached_producer_len() == 0
        }

        #[inline]
        pub fn is_full(&self) -> bool {
            self.0.cursor.cached_producer_len() == self.0.cursor.capacity()
        }

        /// Returns the socket associated with the ring
        #[inline]
        pub fn socket(&self) -> &socket::Fd {
            &self.0.socket
        }
    };
}

macro_rules! impl_consumer {
    ($T:ty, $syscall:ident, $field:ident, $offset:ident) => {
        /// Creates a new ring with the given configuration
        pub fn new(socket: socket::Fd, offsets: &MmapOffsets, size: u32) -> io::Result<Self> {
            syscall::$syscall(&socket, size)?;
            let offsets = &offsets.$field;

            // start with the descriptor offset as the total length
            let mut len = offsets.desc as usize;
            // extend the length by the `size` multiplied the entry size
            len += size as usize * size_of::<$T>();

            // Use the hard-coded offset of the ring type
            let offset = MmapOffsets::$offset;

            let area = Mmap::new(len, offset, Some(socket.as_raw_fd()))?;

            let (cursor, flags) = unsafe {
                // Safety: `area` lives as long as `cursor`
                let (builder, flags) = builder(&area, offsets, size);
                (builder.build_consumer(), flags)
            };

            #[cfg(test)]
            let area = std::sync::Arc::new(area);

            Ok(Self(Ring {
                cursor,
                flags,
                area,
                socket,
            }))
        }

        /// Acquire a number of entries for the ring
        ///
        /// If the cached length is lower than the provided `watermark` no synchronization will be
        /// performed.
        #[inline]
        pub fn acquire(&mut self, watermark: u32) -> u32 {
            self.0.cursor.acquire_consumer(watermark)
        }

        /// Releases `len` number of entries to the producer side of the ring.
        ///
        /// # Panics
        ///
        /// If the `len` exceeds the number of acquired entries, this function will panic.
        #[inline]
        pub fn release(&mut self, len: u32) {
            self.0.cursor.release_consumer(len)
        }

        /// Returns the filled entries for the consumer
        ///
        /// After filling `n` entries, the `release` function should be called with `n`.
        #[inline]
        pub fn data(&mut self) -> (&mut [$T], &mut [$T]) {
            unsafe {
                // Safety: this is the consumer for the ring
                self.0.cursor.consumer_data()
            }
        }

        /// Returns the overall size of the ring
        #[inline]
        pub fn capacity(&self) -> usize {
            self.0.cursor.capacity() as _
        }

        #[inline]
        pub fn len(&self) -> u32 {
            self.0.cursor.cached_consumer_len()
        }

        #[inline]
        pub fn is_empty(&self) -> bool {
            self.0.cursor.cached_consumer_len() == 0
        }

        #[inline]
        pub fn is_full(&self) -> bool {
            self.0.cursor.cached_consumer_len() == self.0.cursor.capacity()
        }

        /// Returns the socket associated with the ring
        #[inline]
        pub fn socket(&self) -> &socket::Fd {
            &self.0.socket
        }

        #[cfg(test)]
        pub fn set_flags(&mut self, flags: crate::if_xdp::RingFlags) {
            *self.0.flags_mut() = flags;
        }
    };
}

/// A transmission ring for entries to be transmitted
#[derive(Debug)]
pub struct Tx(Ring<RxTxDescriptor>);

impl Tx {
    impl_producer!(RxTxDescriptor, set_tx_ring_size, tx, TX_RING);
}

/// A receive ring for entries to be processed
#[derive(Debug)]
pub struct Rx(Ring<RxTxDescriptor>);

impl Rx {
    impl_consumer!(RxTxDescriptor, set_rx_ring_size, rx, RX_RING);
}

/// The fill ring for entries to be populated
#[derive(Debug)]
pub struct Fill(Ring<UmemDescriptor>);

impl Fill {
    impl_producer!(UmemDescriptor, set_fill_ring_size, fill, FILL_RING);

    /// Initializes the ring with the given Umem descriptors
    ///
    /// # Panics
    ///
    /// This should only be called at initialization and will panic if called on a non-full ring.
    #[inline]
    pub fn init<I: Iterator<Item = UmemDescriptor>>(&mut self, descriptors: I) {
        assert!(self.is_full());

        let (head, tail) = self.data();
        let items = head.iter_mut().chain(tail);
        let mut count = 0;
        for (item, desc) in items.zip(descriptors) {
            *item = desc;
            count += 1;
        }
        self.release(count);
    }
}

/// The completion ring for entries to be reused for transmission
#[derive(Debug)]
pub struct Completion(Ring<UmemDescriptor>);

impl Completion {
    impl_consumer!(
        UmemDescriptor,
        set_completion_ring_size,
        completion,
        COMPLETION_RING
    );

    /// Initializes the ring with the given Umem descriptors
    ///
    /// # Panics
    ///
    /// This should only be called at initialization and will panic if called on a non-empty ring.
    #[inline]
    pub fn init<I: Iterator<Item = UmemDescriptor>>(&mut self, descriptors: I) {
        assert!(self.is_empty());

        {
            // pretend we're the producer so we can push items to ourselves
            let size = self.capacity() as u32;
            self.0
                .cursor
                .producer()
                .fetch_add(size, core::sync::atomic::Ordering::SeqCst);

            self.acquire(size);
        }

        let (head, tail) = self.data();
        let items = head.iter_mut().chain(tail);
        for (item, desc) in items.zip(descriptors) {
            *item = desc;
        }
    }
}

#[cfg(test)]
pub mod testing {
    use super::*;
    use crate::{if_xdp, socket::Fd};

    fn offsets() -> if_xdp::RingOffsetV2 {
        if_xdp::RingOffsetV2 {
            producer: 0,
            consumer: core::mem::size_of::<usize>() as _,
            flags: (core::mem::size_of::<usize>() * 2) as _,
            desc: (core::mem::size_of::<usize>() * 3) as _,
        }
    }

    macro_rules! impl_pair {
        ($name:ident, $consumer:ident, $producer:ident, $T:ident) => {
            /// Creates a pair of rings used for testing
            pub fn $name(size: u32) -> ($consumer, $producer) {
                assert!(size.is_power_of_two());

                let offsets = offsets();

                // start with the descriptor offset as the total length
                let mut len = offsets.desc as usize;
                // extend the length by the `size` multiplied the entry size
                len += size as usize * size_of::<$T>();

                let area = Mmap::new(len, 0, None).unwrap();

                let (consumer_cursor, flags) = unsafe {
                    // Safety: `area` lives as long as `cursor`
                    let (builder, flags) = builder(&area, &offsets, size);
                    (builder.build_consumer(), flags)
                };

                let producer_cursor = unsafe {
                    // Safety: `area` lives as long as `cursor`
                    let (builder, _flags) = builder(&area, &offsets, size);
                    builder.build_producer()
                };

                let area = std::sync::Arc::new(area);

                let cons = $consumer(Ring {
                    cursor: consumer_cursor,
                    flags,
                    area: area.clone(),
                    socket: Fd::from_raw(-1),
                });

                let prod = $producer(Ring {
                    cursor: producer_cursor,
                    flags,
                    area,
                    socket: Fd::from_raw(-1),
                });

                (cons, prod)
            }
        };
    }

    impl_pair!(rx_tx, Rx, Tx, RxTxDescriptor);
    impl_pair!(completion_fill, Completion, Fill, UmemDescriptor);

    #[test]
    fn rx_tx_test() {
        let _ = rx_tx(16);
    }

    #[test]
    fn comp_fill_test() {
        let _ = completion_fill(16);
    }
}


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

use crate::{syscall, Result};
use core::fmt;
use std::{
    os::unix::io::{AsRawFd, RawFd},
    sync::Arc,
};

/// A structure for reference counting an AF-XDP socket
#[derive(Clone, PartialEq, Eq)]
pub struct Fd(Arc<Inner>);

impl fmt::Debug for Fd {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        f.debug_tuple("Fd").field(&(self.0).0).finish()
    }
}

impl Fd {
    /// Opens an AF_XDP socket
    ///
    /// This call requires `CAP_NET_RAW` capabilities to succeed.
    #[inline]
    pub fn open() -> Result<Self> {
        let fd = syscall::open()?;
        let fd = Arc::new(Inner(fd));
        Ok(Self(fd))
    }

    pub fn attach_umem(&self, umem: &crate::umem::Umem) -> Result<()> {
        umem.attach(self)?;
        // TODO store the umem
        Ok(())
    }

    /// Creates a socket from a raw file descriptor
    ///
    /// This can be useful for automatically cleaning up a socket on drop
    pub(crate) fn from_raw(value: RawFd) -> Self {
        Self(Arc::new(Inner(value)))
    }
}

impl AsRawFd for Fd {
    #[inline]
    fn as_raw_fd(&self) -> RawFd {
        (self.0).0
    }
}

/// Wrap the RawFd in a structure that automatically closes the socket on drop
#[derive(PartialEq, Eq)]
struct Inner(RawFd);

impl Drop for Inner {
    fn drop(&mut self) {
        unsafe {
            let _ = libc::close(self.0);
        }
    }
}


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

use crate::{
    if_xdp::{Address, MmapOffsets, RingOffsetV1, SocketOptions, Statistics, UmemReg},
    Result,
};
use core::{mem::size_of, ptr::NonNull};
use libc::{c_void, AF_XDP, SOCK_RAW, SOL_XDP};
use std::{
    ffi::CStr,
    io,
    os::unix::io::{AsRawFd, RawFd},
    path::Path,
};

/// Calls the given libc function and wraps the result in an `io::Result`.
macro_rules! libc {
    ($fn: ident ( $($arg: expr),* $(,)* ) ) => {{
        let res = libc::$fn($($arg, )*);
        if res < 0 {
            Err(std::io::Error::last_os_error())
        } else {
            Ok(res)
        }
    }};
}

/// Opens an AF_XDP socket
///
/// This call requires `CAP_NET_RAW` capabilities to succeed.
///
/// See [xsk.c](https://github.com/xdp-project/xdp-tools/blob/a76e7a2b156b8cfe38992206abe9df1df0a29e38/lib/libxdp/xsk.c#L327).
pub fn open() -> io::Result<RawFd> {
    unsafe { libc!(socket(AF_XDP, SOCK_RAW, 0)) }
}

/// Returns all of the [`MmapOffsets`] configured for the AF_XDP socket
///
/// See [xsk.c](https://github.com/xdp-project/xdp-tools/blob/a76e7a2b156b8cfe38992206abe9df1df0a29e38/lib/libxdp/xsk.c#L228).
#[inline]
pub fn offsets<Fd: AsRawFd>(fd: &Fd) -> Result<MmapOffsets> {
    let mut offsets = MmapOffsets::default();

    let optlen = xdp_option(fd, SocketOptions::MmapOffsets, &mut offsets)?;

    // if the written size matched the V2 size, then return
    if optlen == size_of::<MmapOffsets>() {
        return Ok(offsets);
    }

    // adapt older versions of the kernel from v1 to v2
    if optlen == size_of::<MmapOffsets<RingOffsetV1>>() {
        return Ok(offsets.as_v1());
    }

    // an invalid size was returned
    Err(io::Error::new(
        io::ErrorKind::Other,
        format!("invalid mmap offset size: {optlen}"),
    ))
}

/// Returns the collected statistics for the provided AF_XDP socket
#[inline]
pub fn statistics<Fd: AsRawFd>(fd: &Fd) -> Result<Statistics> {
    let mut stats = Statistics::default();
    xdp_option(fd, SocketOptions::Statistics, &mut stats)?;
    Ok(stats)
}

/// Returns the netns cookie associated with the provided AF_XDP socket
///
/// See [xsk.c](https://github.com/xdp-project/xdp-tools/blob/a76e7a2b156b8cfe38992206abe9df1df0a29e38/lib/libxdp/xsk.c#L1055).
pub fn netns_cookie<Fd: AsRawFd>(fd: &Fd) -> Result<u64> {
    // Rust's `libc` doesn't have this defined so we need to define it here
    // https://github.com/torvalds/linux/blob/e62252bc55b6d4eddc6c2bdbf95a448180d6a08d/include/uapi/asm-generic/socket.h#L125
    const SO_NETNS_COOKIE: libc::c_int = 71;

    let mut cookie = 0u64;
    let mut optlen = size_of::<u64>() as libc::socklen_t;

    let ret = unsafe {
        libc::getsockopt(
            fd.as_raw_fd(),
            libc::SOL_SOCKET,
            SO_NETNS_COOKIE,
            &mut cookie as *mut _ as _,
            &mut optlen,
        )
    };

    if ret == 0 {
        return Ok(cookie);
    }

    // the cookie syscall is not supported here so return the default value of 1
    // https://github.com/xdp-project/xdp-tools/blob/a76e7a2b156b8cfe38992206abe9df1df0a29e38/lib/libxdp/xsk.c#L55
    if errno::errno().0 == libc::ENOPROTOOPT {
        return Ok(1);
    }

    Err(io::Error::last_os_error())
}

/// Configures the fill ring size of the provided AF_XDP socket
///
/// See [xsk.c](https://github.com/xdp-project/xdp-tools/blob/a76e7a2b156b8cfe38992206abe9df1df0a29e38/lib/libxdp/xsk.c#L257).
#[inline]
pub fn set_fill_ring_size<Fd: AsRawFd>(fd: &Fd, len: u32) -> Result<()> {
    set_xdp_option(fd, SocketOptions::UmemFillRing, &len)
}

/// Configures the completion ring size of the provided AF_XDP socket
///
/// See [xsk.c](https://github.com/xdp-project/xdp-tools/blob/a76e7a2b156b8cfe38992206abe9df1df0a29e38/lib/libxdp/xsk.c#L263).
#[inline]
pub fn set_completion_ring_size<Fd: AsRawFd>(fd: &Fd, len: u32) -> Result<()> {
    set_xdp_option(fd, SocketOptions::UmemCompletionRing, &len)
}

/// Configures the RX ring size of the provided AF_XDP socket
///
/// See [xsk.c](https://github.com/xdp-project/xdp-tools/blob/a76e7a2b156b8cfe38992206abe9df1df0a29e38/lib/libxdp/xsk.c#L1081).
#[inline]
pub fn set_rx_ring_size<Fd: AsRawFd>(fd: &Fd, len: u32) -> Result<()> {
    set_xdp_option(fd, SocketOptions::RxRing, &len)
}

/// Configures the TX ring size of the provided AF_XDP socket
///
/// See [xsk.c](https://github.com/xdp-project/xdp-tools/blob/a76e7a2b156b8cfe38992206abe9df1df0a29e38/lib/libxdp/xsk.c#L1093).
#[inline]
pub fn set_tx_ring_size<Fd: AsRawFd>(fd: &Fd, len: u32) -> Result<()> {
    set_xdp_option(fd, SocketOptions::TxRing, &len)
}

/// Configures the UMEM options for the provided AF_XDP socket
///
/// See [xsk.c](https://github.com/xdp-project/xdp-tools/blob/a76e7a2b156b8cfe38992206abe9df1df0a29e38/lib/libxdp/xsk.c#L344).
#[inline]
pub fn set_umem<Fd: AsRawFd>(fd: &Fd, reg: &UmemReg) -> Result<()> {
    set_xdp_option(fd, SocketOptions::UmemReg, reg)
}

/// Binds the provided AF_XDP socket to an address
#[inline]
pub fn bind<Fd: AsRawFd>(fd: &Fd, addr: &mut Address) -> Result<()> {
    unsafe {
        libc!(bind(
            fd.as_raw_fd(),
            addr as *mut _ as _,
            size_of::<Address>() as _
        ))?;
    }
    Ok(())
}

/// Notifies the kernel to send any packets on the TX ring
///
/// This should be called after checking if the TX ring needs a wake up.
#[inline]
pub fn wake_tx<Fd: AsRawFd>(fd: &Fd) -> Result<()> {
    unsafe {
        // after some testing, `sendto` is better than `sengmsg` here since it doesn't have to copy
        // the msghdr from userspace, which would be zeroed and meaningless
        libc!(sendto(
            fd.as_raw_fd(),
            core::ptr::null_mut(),
            0,
            libc::MSG_DONTWAIT,
            core::ptr::null_mut(),
            0,
        ))?;
    };

    Ok(())
}

/// Tries to receive packets on the RX ring
///
/// Note that this call is non-blocking and is usually used with busy polling. Use [`libc::poll`]
/// instead to block a task on socket readiness.
#[inline]
pub fn busy_poll<Fd: AsRawFd>(fd: &Fd) -> Result<u32> {
    let mut msg = unsafe {
        // Safety: msghdr is zeroable
        core::mem::zeroed()
    };
    let count = unsafe { libc!(recvmsg(fd.as_raw_fd(), &mut msg, libc::MSG_DONTWAIT,))? };
    Ok(count as u32)
}

/// Opens a mmap region, with an optional associated file descriptor.
///
/// See [xsk.c](https://github.com/xdp-project/xdp-tools/blob/a76e7a2b156b8cfe38992206abe9df1df0a29e38/lib/libxdp/xsk.c#L273).
#[inline]
pub fn mmap(len: usize, offset: usize, fd: Option<RawFd>) -> Result<NonNull<c_void>> {
    let flags = if fd.is_some() {
        libc::MAP_SHARED | libc::MAP_POPULATE
    } else {
        libc::MAP_PRIVATE | libc::MAP_ANONYMOUS
    };

    // See:
    // * Fill https://github.com/xdp-project/xdp-tools/blob/a76e7a2b156b8cfe38992206abe9df1df0a29e38/lib/libxdp/xsk.c#L273
    // * Completion https://github.com/xdp-project/xdp-tools/blob/a76e7a2b156b8cfe38992206abe9df1df0a29e38/lib/libxdp/xsk.c#L287
    // * RX https://github.com/xdp-project/xdp-tools/blob/a76e7a2b156b8cfe38992206abe9df1df0a29e38/lib/libxdp/xsk.c#L1111
    // * TX https://github.com/xdp-project/xdp-tools/blob/a76e7a2b156b8cfe38992206abe9df1df0a29e38/lib/libxdp/xsk.c#L1132
    let addr = unsafe {
        libc::mmap(
            core::ptr::null_mut(),
            len as _,
            libc::PROT_READ | libc::PROT_WRITE,
            flags,
            fd.unwrap_or(-1),
            offset as _,
        )
    };

    if addr == libc::MAP_FAILED {
        return Err(io::Error::last_os_error());
    }

    let addr = NonNull::new(addr)
        .ok_or_else(|| io::Error::new(io::ErrorKind::Other, "mmap returned null pointer"))?;

    Ok(addr)
}

/// Unmaps a mmap region
///
/// # Safety
///
/// The following must be true:
///
/// * The provided `addr` MUST be a mmap'd region
/// * The original mmap length MUST match the provided `len` parameter.
pub unsafe fn munmap(addr: NonNull<c_void>, len: usize) -> Result<()> {
    libc!(munmap(addr.as_ptr(), len as _))?;
    Ok(())
}

/// Converts an interface name to its index
///
/// Returns an error if the name could not be resolved.
pub fn if_nametoindex(name: &CStr) -> Result<u32> {
    unsafe {
        let ifindex = libc::if_nametoindex(name.as_ptr());

        // https://man7.org/linux/man-pages/man3/if_nametoindex.3.html
        // > On success, if_nametoindex() returns the index number of the
        // > network interface; on error, 0 is returned and errno is set to
        // > indicate the error.
        if ifindex == 0 {
            return Err(std::io::Error::last_os_error());
        }

        Ok(ifindex)
    }
}

/// Returns the maximum number of queues that can be opened on a particular network device
pub fn max_queues(ifname: &str) -> u32 {
    // See https://github.com/xdp-project/xdp-tools/blob/1c8662f4bb44445454c8f66df56ccd11274d4e30/lib/libxdp/xsk.c#L436

    // First try the ethtool API
    if let Some(queues) = ethtool_queues(ifname).filter(|v| *v > 0) {
        return queues;
    }

    // Then try to query the number from sysfs
    if let Some(queues) = sysfs_queues(ifname).filter(|v| *v > 0) {
        return queues;
    }

    // If none of the previous methods worked, then just default to a single queue
    1
}

/// Queries the number of queues for a given interface with ethtool APIs
fn ethtool_queues(ifname: &str) -> Option<u32> {
    // See https://github.com/xdp-project/xdp-tools/blob/1c8662f4bb44445454c8f66df56ccd11274d4e30/lib/libxdp/xsk.c#L436

    let fd = unsafe { libc!(socket(libc::AF_LOCAL, libc::SOCK_DGRAM, 0)).ok()? };
    // close the FD on drop
    let fd = crate::socket::Fd::from_raw(fd);

    let mut channels = unsafe { core::mem::zeroed::<crate::bindings::ethtool_channels>() };
    channels.cmd = crate::bindings::ETHTOOL_GCHANNELS;

    let mut ifreq = unsafe { core::mem::zeroed::<libc::ifreq>() };
    ifreq.ifr_ifru.ifru_data = (&mut channels) as *mut _ as *mut _;

    assert!(ifname.len() < ifreq.ifr_name.len());
    unsafe {
        core::ptr::copy_nonoverlapping(
            ifname.as_bytes().as_ptr(),
            &mut ifreq.ifr_name as *mut _ as *mut u8,
            ifname.len(),
        );
    }

    unsafe {
        libc!(ioctl(fd.as_raw_fd(), libc::SIOCETHTOOL, &mut ifreq)).ok()?;
    }

    // Take the max of the max values, each driver returns in a different way
    let queues = channels
        .max_rx
        .max(channels.max_tx)
        .max(channels.max_combined);

    Some(queues)
}

/// Queries the number of queues for a given interface with sysfs
fn sysfs_queues(ifname: &str) -> Option<u32> {
    // See https://github.com/xdp-project/xdp-tools/blob/1c8662f4bb44445454c8f66df56ccd11274d4e30/lib/libxdp/xsk.c#L408

    let mut rx = 0;
    let mut tx = 0;

    let path = Path::new("/sys/class/net").join(ifname).join("queues");

    for entry in path.read_dir().ok()?.flatten() {
        let path = entry.path();

        if let Some(path) = path.file_name().map(|p| p.to_string_lossy()) {
            if path.starts_with("rx") {
                rx += 1;
            } else if path.starts_with("tx") {
                tx += 1;
            }
        }
    }

    let queues = rx.max(tx).max(1);
    Some(queues)
}

#[inline]
fn xdp_option<Fd: AsRawFd, T: Sized>(fd: &Fd, opt: SocketOptions, value: &mut T) -> Result<usize> {
    let mut optlen = size_of::<T>() as libc::socklen_t;

    unsafe {
        libc!(getsockopt(
            fd.as_raw_fd(),
            SOL_XDP,
            opt as _,
            value as *mut _ as _,
            &mut optlen,
        ))?;
    }

    Ok(optlen as usize)
}

#[inline]
fn set_xdp_option<Fd: AsRawFd, T: Sized>(fd: &Fd, opt: SocketOptions, value: &T) -> Result<()> {
    let optlen = size_of::<T>() as libc::socklen_t;

    unsafe {
        libc!(setsockopt(
            fd.as_raw_fd(),
            SOL_XDP,
            opt as _,
            value as *const _ as _,
            optlen,
        ))?;
    }

    Ok(())
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::mmap::Mmap;
    use core::ffi::CStr;

    #[test]
    fn max_queues_test() {
        // we can't make any assumptions about the test environment but we can at least make sure
        // it returns >=1
        assert!(max_queues("lo") >= 1);
    }

    #[test]
    fn syscall_test() {
        // This call requires `CAP_NET_RAW`. If the test doesn't have this set, then log and skip
        // the next calls
        let fd = if let Ok(fd) = open() {
            fd
        } else {
            assert!(
                std::env::var("CAP_NET_RAW_ENABLED").is_err(),
                "expected the syscall test to be executed"
            );

            use std::io::Write;
            let _ = writeln!(
                std::io::stdout(),
                "WARNING: CAP_NET_RAW capabilities missing; skipping `syscall_test`",
            );
            return;
        };
        // close the FD on drop
        let _owned_fd = crate::socket::Fd::from_raw(fd);

        // the ring sizes need to be a power of 2
        let ring_size = 32u32;
        let frame_size = 4096;
        let frame_count = (ring_size * 4) as usize;

        // Set up the ring sizes
        {
            set_fill_ring_size(&fd, ring_size).unwrap();
            set_completion_ring_size(&fd, ring_size).unwrap();
            set_rx_ring_size(&fd, ring_size).unwrap();
            set_tx_ring_size(&fd, ring_size).unwrap();
        }

        // Set up the UMEM region
        let umem_len = frame_size * frame_count;
        let umem = Mmap::new(umem_len, 0, None).unwrap();
        {
            let umem_conf = UmemReg {
                addr: umem.addr().as_ptr() as *mut _ as _,
                chunk_size: frame_size as _,
                flags: Default::default(),
                headroom: 0,
                len: umem_len as _,
            };

            set_umem(&fd, &umem_conf).unwrap();
        }

        // Print out the getter calls
        {
            dbg!(offsets(&fd).unwrap());
            dbg!(statistics(&fd).unwrap());
            dbg!(netns_cookie(&fd).unwrap());
        }

        // try binding to the loopback interface and calling the IO syscalls
        {
            let mut addr = Address::default();
            let if_name = unsafe { CStr::from_ptr(b"lo\0" as *const _ as _) };

            if addr.set_if_name(if_name).is_ok() {
                eprintln!("using ifindex {}", addr.ifindex);
                if bind(&fd, &mut addr).is_ok() {
                    wake_tx(&fd).unwrap();
                    dbg!(busy_poll(&fd).unwrap());
                }
            }
        }
    }
}


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

use crate::{
    if_xdp::{RxTxDescriptor, UmemDescriptor, UmemFlags, UmemReg},
    mmap::Mmap,
    syscall, Result,
};
use core::ptr::NonNull;
use std::{os::unix::io::AsRawFd, sync::Arc};

/// The default value for frame sizes
pub const DEFAULT_FRAME_SIZE: u32 = 4096;

#[derive(Clone, Copy, Debug)]
pub struct Builder {
    /// The maximum number of bytes a frame can hold (MTU)
    pub frame_size: u32,
    /// The number of frames that should be allocated
    pub frame_count: u32,
    /// The headroom size for each frame
    pub frame_headroom: u32,
    /// The flags for the Umem
    pub flags: UmemFlags,
}

impl Default for Builder {
    fn default() -> Self {
        Self {
            frame_size: DEFAULT_FRAME_SIZE,
            frame_count: 1024,
            frame_headroom: 0,
            flags: Default::default(),
        }
    }
}

impl Builder {
    pub fn build(self) -> Result<Umem> {
        let len = self.frame_size as usize * self.frame_count as usize;
        let area = Mmap::new(len, 0, None)?;
        let area = Arc::new(area);
        let mem = area.addr().cast();

        Ok(Umem {
            area,
            mem,
            frame_size: self.frame_size,
            frame_count: self.frame_count,
            flags: self.flags,
            frame_headroom: self.frame_headroom,
        })
    }
}

/// A shared region of memory for holding frame (packet) data
///
/// Callers are responsible for correct descriptor allocation. This means only one descriptor
/// should be alive for each frame. If this invariant is not held, the borrowing rules will be
/// violated and potentially result in UB.
#[derive(Clone)]
pub struct Umem {
    area: Arc<Mmap>,
    mem: NonNull<u8>,
    frame_size: u32,
    frame_count: u32,
    frame_headroom: u32,
    flags: UmemFlags,
}

/// Safety: The umem mmap region can be sent to other threads
unsafe impl Send for Umem {}

/// Safety: The umem mmap region is synchronized by Rings and the data getters are marked as
/// `unsafe`.
unsafe impl Sync for Umem {}

impl Umem {
    /// Creates a Umem builder with defaults
    pub fn builder() -> Builder {
        Builder::default()
    }

    /// Returns the configured size for each frame
    #[inline]
    pub fn frame_size(&self) -> u32 {
        self.frame_size
    }

    /// Returns the total number of frames in the Umem
    #[inline]
    pub fn frame_count(&self) -> u32 {
        self.frame_count
    }

    /// Returns the configured headroom for each frame
    #[inline]
    pub fn frame_headroom(&self) -> u32 {
        self.frame_headroom
    }

    /// Returns the flags for the Umem
    #[inline]
    pub fn flags(&self) -> UmemFlags {
        self.flags
    }

    /// Returns an iterator over all of the frame descriptors
    ///
    /// This can be used to initialize a frame allocator
    pub fn frames(&self) -> impl Iterator<Item = UmemDescriptor> {
        let size = self.frame_size as u64;
        (0..self.frame_count as u64).map(move |idx| UmemDescriptor {
            address: idx * size,
        })
    }

    /// Returns the number of bytes in the Umem
    #[inline]
    pub fn len(&self) -> usize {
        self.area.len()
    }

    /// Returns `true` if the Umem is empty
    #[inline]
    pub fn is_empty(&self) -> bool {
        self.area.is_empty()
    }

    /// Returns the region of memory as specified by the index type
    ///
    /// # Safety
    ///
    /// The caller MUST ensure that this index is not already mutably borrowed and that the index
    /// is in bounds for this Umem.
    #[inline]
    pub unsafe fn get<T>(&self, idx: T) -> &[u8]
    where
        Self: UnsafeIndex<T>,
    {
        self.index(idx)
    }

    /// Returns the mutable region of memory as specified by the index type
    ///
    /// # Safety
    ///
    /// The caller MUST ensure that this index is not already mutably borrowed and that the index
    /// is in bounds for this Umem.
    #[inline]
    #[allow(clippy::mut_from_ref)] // interior mutability safety is enforced by the caller
    pub unsafe fn get_mut<T>(&self, idx: T) -> &mut [u8]
    where
        Self: UnsafeIndex<T>,
    {
        self.index_mut(idx)
    }

    /// Attaches the Umem to the specified socket
    pub(crate) fn attach<Fd: AsRawFd>(&self, socket: &Fd) -> Result<()> {
        let umem_conf = UmemReg {
            addr: self.area.addr().as_ptr() as _,
            chunk_size: self.frame_size,
            flags: self.flags,
            headroom: self.frame_headroom,
            len: self.area.len() as _,
        };

        syscall::set_umem(socket, &umem_conf)?;

        Ok(())
    }

    #[inline]
    fn validate_rx_tx_descriptor(&self, desc: RxTxDescriptor) -> *mut u8 {
        debug_assert!(desc.len <= self.frame_size, "frame too large");
        debug_assert!(
            desc.address + desc.len as u64 <= self.area.len() as u64,
            "pointer out of bounds"
        );
        unsafe { self.mem.as_ptr().add(desc.address as _) }
    }

    #[inline]
    fn validate_umem_descriptor(&self, desc: UmemDescriptor) -> *mut u8 {
        debug_assert!(
            desc.address + self.frame_size as u64 <= self.area.len() as u64,
            "pointer out of bounds"
        );
        unsafe { self.mem.as_ptr().add(desc.address as _) }
    }
}

/// Specifies an indexable value, which relies on the caller to guarantee borrowing rules are not
/// violated.
pub trait UnsafeIndex<T> {
    unsafe fn index(&self, idx: T) -> &[u8];
    #[allow(clippy::mut_from_ref)] // interior mutability safety is enforced by the caller
    unsafe fn index_mut(&self, idx: T) -> &mut [u8];
}

impl UnsafeIndex<RxTxDescriptor> for Umem {
    #[inline]
    unsafe fn index(&self, idx: RxTxDescriptor) -> &[u8] {
        let ptr = self.validate_rx_tx_descriptor(idx);
        core::slice::from_raw_parts(ptr, idx.len as _)
    }

    #[inline]
    unsafe fn index_mut(&self, idx: RxTxDescriptor) -> &mut [u8] {
        let ptr = self.validate_rx_tx_descriptor(idx);
        core::slice::from_raw_parts_mut(ptr, idx.len as _)
    }
}

impl UnsafeIndex<UmemDescriptor> for Umem {
    #[inline]
    unsafe fn index(&self, idx: UmemDescriptor) -> &[u8] {
        let ptr = self.validate_umem_descriptor(idx);
        core::slice::from_raw_parts(ptr, self.frame_size as _)
    }

    #[inline]
    unsafe fn index_mut(&self, idx: UmemDescriptor) -> &mut [u8] {
        let ptr = self.validate_umem_descriptor(idx);
        core::slice::from_raw_parts_mut(ptr, self.frame_size as _)
    }
}

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

    #[test]
    fn index_test() {
        let umem = Umem::builder().build().unwrap();

        for descriptor in umem.frames() {
            unsafe {
                let _ = umem.get(descriptor);
                let _ = umem.get_mut(descriptor);
            }

            let rx = RxTxDescriptor {
                address: descriptor.address,
                len: 1,
                options: Default::default(),
            };

            unsafe {
                let _ = umem.get(rx);
                let _ = umem.get_mut(rx);
            }
        }
    }
}