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

use s2n_quic_core::{
    io::tx::{self, router},
    path::Handle as _,
    xdp::path,
};

/// Routes TX messages based on if the local address port is non-zero
///
/// This can be used for client IO providers that wish to send the initial packet over a standard
/// UDP socket in order to offload address resolution to the operating system.
#[derive(Default)]
pub struct Router(());

impl router::Router for Router {
    type Handle = path::Tuple;

    /// If the local port is 0 then forward to `AddressUnknown`. Otherwise forward to
    /// `AddressKnown`.
    #[inline]
    fn route<M, AddressKnown, AddressUnknown>(
        &mut self,
        message: M,
        address_known: &mut AddressKnown,
        address_unknown: &mut AddressUnknown,
    ) -> Result<tx::Outcome, tx::Error>
    where
        M: tx::Message<Handle = Self::Handle>,
        AddressKnown: tx::Queue<Handle = Self::Handle>,
        AddressUnknown: tx::Queue<Handle = Self::Handle>,
    {
        if message.path_handle().local_address().port() == 0 {
            address_unknown.push(message)
        } else {
            address_known.push(message)
        }
    }
}


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

use crate::{if_xdp::RxTxDescriptor, ring, umem::Umem};
use core::task::{Context, Poll};
use s2n_codec::DecoderBufferMut;
use s2n_quic_core::{
    event,
    inet::datagram,
    io::rx,
    slice::zip,
    sync::atomic_waker,
    xdp::{decoder, path},
};

#[cfg(feature = "tokio")]
mod tokio_impl;

/// An interface to handle any errors that happen on the RX IO provider
pub trait ErrorLogger: Send {
    /// Called any time the packet could not be decoded.
    ///
    /// Assuming the correct BPF program is loaded, this should never happen. In case it does, the
    /// application should emit an alarm to diagnose the reason why it's happening.
    fn log_invalid_packet(&mut self, bytes: &[u8]);
}

/// Drives the Rx and Fill rings forward
pub trait Driver: 'static {
    fn poll(&mut self, rx: &mut ring::Rx, fill: &mut ring::Fill, cx: &mut Context) -> Option<u32>;
}

impl Driver for atomic_waker::Handle {
    #[inline]
    fn poll(&mut self, rx: &mut ring::Rx, fill: &mut ring::Fill, cx: &mut Context) -> Option<u32> {
        let mut count = 0;

        // iterate twice to avoid race conditions on the waker registration
        for i in 0..2 {
            count = rx.acquire(u32::MAX);
            count = fill.acquire(count).min(count);

            // we have items to receive and fill so return
            if count > 0 {
                break;
            }

            // check to see if the channel is open, if not return
            if !self.is_open() {
                return None;
            }

            // only register the waker the first iteration
            if i > 0 {
                continue;
            }

            trace!("registering waker");
            self.register(cx.waker());
            trace!("waking peer waker");
            self.wake();
        }

        Some(count)
    }
}

pub struct Channel<D: Driver> {
    pub rx: ring::Rx,
    pub fill: ring::Fill,
    pub driver: D,
}

impl<D: Driver> Channel<D> {
    /// Tries to acquire entries in the channel
    ///
    /// Returns None if the channel is closed
    #[inline]
    fn acquire(&mut self, cx: &mut Context) -> Option<u32> {
        self.driver.poll(&mut self.rx, &mut self.fill, cx)
    }

    /// Iterates over all of the acquired entries in the ring and calls `on_packet`
    #[inline]
    fn for_each<F: FnMut(RxTxDescriptor)>(&mut self, mut on_packet: F) {
        // one last effort to acquire any packets
        let len = self.rx.acquire(1);
        let len = self.fill.acquire(len);
        if len == 0 {
            return;
        }

        let rx = self.rx.data();
        let rx = [rx.0, rx.1];

        let fill = self.fill.data();
        let mut fill = [fill.0, fill.1];

        let count = zip(&rx, &mut fill, |rx, fill| {
            on_packet(*rx);
            // send the descriptor to the fill queue
            *fill = (*rx).into();
        });

        trace!("releasing {count} descriptors");

        self.rx.release(count as _);
        self.fill.release(count as _);
    }
}

pub struct Rx<D: Driver> {
    channels: Vec<Channel<D>>,
    umem: Umem,
    error_logger: Option<Box<dyn ErrorLogger>>,
}

impl<D: Driver> Rx<D> {
    /// Creates a RX IO interface for an s2n-quic endpoint
    pub fn new(channels: Vec<Channel<D>>, umem: Umem) -> Self {
        Self {
            channels,
            umem,
            error_logger: None,
        }
    }

    /// Sets the error logger on the RX IO provider
    pub fn with_error_logger(mut self, error_logger: Box<dyn ErrorLogger>) -> Self {
        self.error_logger = Some(error_logger);
        self
    }
}

impl<D: Driver> rx::Rx for Rx<D> {
    type PathHandle = path::Tuple;
    type Queue = Queue<'static, D>;
    type Error = ();

    #[inline]
    fn poll_ready(&mut self, cx: &mut Context) -> Poll<Result<(), Self::Error>> {
        // poll both channels to make sure we can make progress in both

        let mut is_any_ready = false;
        // assume all of the channels are closed until we get the first open one
        let mut is_all_closed = true;

        for (idx, channel) in self.channels.iter_mut().enumerate() {
            if let Some(count) = channel.acquire(cx) {
                trace!("acquired {count} items from channel {idx}");
                is_all_closed = false;
                is_any_ready |= count > 0;
            } else {
                trace!("channel {idx} closed");
            }
        }

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

        // wake the endpoint if any of the channels are ready to be processed
        if is_any_ready {
            trace!("rx ready");
            Poll::Ready(Ok(()))
        } else {
            Poll::Pending
        }
    }

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

        let channels = &mut this.channels;
        let umem = &mut this.umem;
        let error_logger = &mut this.error_logger;

        let mut queue = Queue {
            channels,
            umem,
            error_logger,
        };
        f(&mut queue);
    }

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

pub struct Queue<'a, D: Driver> {
    channels: &'a mut Vec<Channel<D>>,
    umem: &'a mut Umem,
    error_logger: &'a mut Option<Box<dyn ErrorLogger>>,
}

impl<'a, D: Driver> rx::Queue for Queue<'a, D> {
    type Handle = path::Tuple;

    #[inline]
    fn for_each<F: FnMut(datagram::Header<Self::Handle>, &mut [u8])>(&mut self, mut on_packet: F) {
        for (idx, channel) in self.channels.iter_mut().enumerate() {
            trace!("draining channel {idx}");

            channel.for_each(|descriptor| {
                trace!("received descriptor {descriptor:?}");

                let buffer = unsafe {
                    // Safety: this descriptor should be unique, assuming the tasks are functioning
                    // properly
                    self.umem.get_mut(descriptor)
                };

                // create a decoder from the descriptor's buffer
                let decoder = DecoderBufferMut::new(buffer);

                // try to decode the packet and emit the result
                match decoder::decode_packet(decoder) {
                    Ok(Some((header, payload))) => {
                        on_packet(header, payload.into_less_safe_slice());
                    }
                    Ok(None) | Err(_) => {
                        // This shouldn't happen. If it does, the BPF program isn't properly validating
                        // packets before they get to userspace.
                        if let Some(error_logger) = self.error_logger.as_mut() {
                            error_logger.log_invalid_packet(buffer);
                        }
                    }
                }
            });
        }
    }

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


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

use super::{
    rx::{self, Rx},
    tx::{self, Tx},
};
use crate::{if_xdp::RingFlags, ring, umem::Umem};
use core::{mem::size_of, time::Duration};
use rand::prelude::*;
use s2n_quic_core::{
    inet::ExplicitCongestionNotification,
    io::{
        rx::{Queue as _, Rx as _},
        tx::{Error, Message, PayloadBuffer, Queue as _, Tx as _},
    },
    sync::atomic_waker,
    xdp::path,
};

/// Tests the s2n-quic-core IO trait implementations by sending packets over spsc channels
#[tokio::test]
async fn tx_rx_test() {
    let frame_count = 16;
    let mut umem = Umem::builder();
    umem.frame_count = frame_count;
    umem.frame_size = 128;
    let umem = umem.build().unwrap();

    let iterations = 10;

    // send a various amount of packets for each test
    for packets in [1, 100, 1000, 10_000] {
        for input_counts in [1, 2] {
            for _ in 0..iterations {
                eprintln!("======================");
                eprintln!("packets: {packets}, input_counts: {input_counts}");

                let mut rx_inputs = vec![];
                let mut tx_outputs = vec![];

                let mut frames = umem.frames();

                for _ in 0..input_counts {
                    let (completion, mut fill) = ring::testing::completion_fill(32);
                    let (mut rx, tx) = ring::testing::rx_tx(16);
                    let (tx_waker, rx_waker) = atomic_waker::pair();

                    // we always need to wakeup
                    rx.set_flags(RingFlags::NEED_WAKEUP);

                    {
                        fill.acquire(u32::MAX);
                        let count = frame_count / input_counts;
                        let frames = (&mut frames).take(count as usize);
                        let (head, tail) = fill.data();
                        for (frame, fill) in frames.zip(head.iter_mut().chain(tail)) {
                            *fill = frame;
                        }
                        fill.release(count);
                    }

                    tx_outputs.push(tx::Channel {
                        tx,
                        driver: tx_waker,
                        completion,
                    });
                    rx_inputs.push(rx::Channel {
                        rx,
                        driver: rx_waker,
                        fill,
                    });
                }

                let send_task = tokio::spawn(send(packets, tx_outputs, umem.clone()));
                recv(packets, rx_inputs, umem.clone()).await;
                send_task.await.unwrap();
            }
        }
    }
}

/// Packets sent over the IO implementations
#[derive(Debug)]
struct Packet {
    pub path: path::Tuple,
    pub ecn: ExplicitCongestionNotification,
    pub counter: u32,
}

/// Make it easy to write the packet to the TX queue
impl Message for Packet {
    type Handle = path::Tuple;

    fn path_handle(&self) -> &Self::Handle {
        &self.path
    }

    fn ecn(&mut self) -> ExplicitCongestionNotification {
        self.ecn
    }

    fn delay(&mut self) -> Duration {
        Default::default()
    }

    fn ipv6_flow_label(&mut self) -> u32 {
        self.counter
    }

    fn can_gso(&self, _: usize, _: usize) -> bool {
        false
    }

    fn write_payload(&mut self, mut payload: PayloadBuffer, _gso: usize) -> Result<usize, Error> {
        payload.write(&self.counter.to_be_bytes())
    }
}

/// Sends `count` packets over the TX queue
async fn send(count: u32, outputs: Vec<tx::Channel<atomic_waker::Handle>>, umem: Umem) {
    let state = Default::default();
    let mut tx = Tx::new(outputs, umem, state);

    let mut counter = 0;
    let mut needs_poll = false;
    while counter < count {
        if core::mem::take(&mut needs_poll) {
            if tx.ready().await.is_err() {
                break;
            }
        }

        tx.queue(|queue| {
            let max = queue.capacity().min((count - counter) as usize);
            let count = thread_rng().gen_range(0..=max);
            trace!("max: {max}, count: {count}");

            for _ in 0..count {
                let path = path::Tuple::UNSPECIFIED;
                let ecn = ExplicitCongestionNotification::default();
                let packet = Packet { counter, ecn, path };
                counter += 1;
                queue.push(packet).unwrap();
            }

            needs_poll |= !queue.has_capacity();
        });

        // randomly yield to other tasks
        maybe_yield().await;
    }

    trace!("shutting down send");
}

/// Receives raw packets and converts them into [`Packet`]s, putting them on the `output` channel.
async fn recv(packets: u32, inputs: Vec<rx::Channel<atomic_waker::Handle>>, umem: Umem) {
    let mut rx = Rx::new(inputs, umem);
    let mut actual = s2n_quic_core::interval_set::IntervalSet::default();

    while rx.ready().await.is_ok() {
        rx.queue(|queue| {
            queue.for_each(|_datagram, payload| {
                assert_eq!(payload.len(), size_of::<u32>());
                let counter = u32::from_be_bytes(payload.try_into().unwrap());
                trace!("received packet {counter}");
                actual.insert_value(counter).unwrap();
            });
        });

        // randomly yield to other tasks
        maybe_yield().await;
    }

    assert_eq!(
        packets as usize,
        actual.count(),
        "total output packets does not match input"
    );

    trace!("shutting down recv");
}

/// Randomly yields to other tasks
async fn maybe_yield() {
    if thread_rng().gen() {
        trace!("yielding");
        tokio::task::yield_now().await;
    }
}


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

use crate::{ring, syscall, umem::Umem};
use core::task::{Context, Poll};
use s2n_codec::{Encoder as _, EncoderBuffer};
use s2n_quic_core::{
    event,
    io::tx,
    sync::atomic_waker,
    xdp::{encoder, path},
};

/// Drives the Tx and Completion rings forward
pub trait Driver: 'static {
    #[inline]
    fn poll(
        &mut self,
        tx: &mut ring::Tx,
        completion: &mut ring::Completion,
        cx: &mut Context,
    ) -> Option<bool> {
        // Default to doing nothing
        //
        // In order to keep the trait signature from having `_` prefixes in the name, discard the
        // variables in the body.
        let _ = tx;
        let _ = completion;
        let _ = cx;
        Some(false)
    }

    #[inline]
    fn wake(&mut self, tx: &mut ring::Tx, completion: &mut ring::Completion) {
        // Default to doing nothing
        //
        // In order to keep the trait signature from having `_` prefixes in the name, discard the
        // variables in the body.
        let _ = tx;
        let _ = completion;
    }
}

impl Driver for () {}

impl Driver for atomic_waker::Handle {
    #[inline]
    fn poll(
        &mut self,
        tx: &mut ring::Tx,
        completion: &mut ring::Completion,
        cx: &mut Context,
    ) -> Option<bool> {
        // record if either of the rings were empty to start
        let was_empty = tx.is_empty() || completion.is_empty();

        // iterate twice to avoid waker registration races
        for i in 0..2 {
            let count = completion.acquire(u32::MAX);
            let count = tx.acquire(count).min(count);

            trace!("acquired {count} entries");

            // return if we have entries in both rings
            if count > 0 {
                return Some(was_empty);
            }

            // if the peer's handle is closed, then shut down the task
            if !self.is_open() {
                return None;
            }

            // only register wakers if it's the first iteration and it started out empty
            if i > 0 || !was_empty {
                continue;
            }

            trace!("registering waker");
            self.register(cx.waker());
            trace!("waking waker");
            self.wake(tx, completion);
        }

        // we need to keep polling until we have at least one item here
        if tx.needs_wakeup() || completion.is_empty() || tx.is_empty() {
            atomic_waker::Handle::wake(self);
        }

        Some(false)
    }

    #[inline]
    fn wake(&mut self, tx: &mut ring::Tx, _completion: &mut ring::Completion) {
        if tx.needs_wakeup() {
            atomic_waker::Handle::wake(self);
        }
    }
}

pub struct BusyPoll;

impl Driver for BusyPoll {
    #[inline]
    fn poll(
        &mut self,
        tx: &mut ring::Tx,
        completion: &mut ring::Completion,
        cx: &mut Context,
    ) -> Option<bool> {
        // record if either of the rings were empty to start
        let was_empty = tx.is_empty() || completion.is_empty();

        // iterate twice to avoid waker registration races
        for i in 0..2 {
            let count = completion.acquire(u32::MAX);
            let count = tx.acquire(count).min(count);

            trace!("acquired {count} entries");

            // return if we have entries in both rings
            if count > 0 {
                return Some(was_empty);
            }

            // only wake the socket's driver if it's the first iteration
            if i == 0 {
                self.wake(tx, completion);
            }
        }

        // we need to keep polling until we have at least one item here
        if completion.is_empty() || tx.is_empty() {
            cx.waker().wake_by_ref();
        }

        Some(false)
    }

    #[inline]
    fn wake(&mut self, tx: &mut ring::Tx, _completion: &mut ring::Completion) {
        // wake up the driver if it's indicated we need to do so
        if tx.needs_wakeup() {
            let _ = syscall::wake_tx(tx.socket());
        }
    }
}

pub struct Channel<D: Driver> {
    pub tx: ring::Tx,
    pub completion: ring::Completion,
    pub driver: D,
}

impl<D: Driver> Channel<D> {
    #[inline]
    fn acquire(&mut self, cx: &mut Context) -> Option<bool> {
        // don't try to drive anything if the queues are both full
        if self.tx.is_full() && self.completion.is_full() {
            return Some(false);
        }

        trace!("acquiring channel capacity");

        self.driver.poll(&mut self.tx, &mut self.completion, cx)
    }

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

    #[inline]
    fn wake(&mut self) {
        self.driver.wake(&mut self.tx, &mut self.completion);
    }
}

pub struct Tx<D: Driver> {
    channels: Vec<Channel<D>>,
    umem: Umem,
    encoder: encoder::State,
}

impl<D: Driver> Tx<D> {
    /// Creates a TX IO interface for an s2n-quic endpoint
    pub fn new(channels: Vec<Channel<D>>, umem: Umem, encoder: encoder::State) -> Self {
        Self {
            channels,
            umem,
            encoder,
        }
    }

    /// Consumes the TX endpoint into the inner channels
    ///
    /// This is used for internal tests only.
    #[cfg(test)]
    pub fn consume(self) -> Vec<Channel<D>> {
        self.channels
    }
}

impl<D: Driver> tx::Tx for Tx<D> {
    type PathHandle = path::Tuple;
    type Queue = Queue<'static, D>;
    type Error = ();

    #[inline]
    fn poll_ready(&mut self, cx: &mut Context) -> Poll<Result<(), Self::Error>> {
        // poll both channels to make sure we can make progress in both
        let mut is_any_ready = false;
        // assume all of the channels are closed until we get one that's not
        let mut is_all_closed = true;

        for (idx, channel) in self.channels.iter_mut().enumerate() {
            if let Some(did_become_ready) = channel.acquire(cx) {
                if did_become_ready {
                    trace!("channel {idx} became ready");
                }

                is_all_closed = false;
                is_any_ready |= did_become_ready;
            } else {
                trace!("channel {idx} closed");
            }
        }

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

        if is_any_ready {
            trace!("tx ready");
            Poll::Ready(Ok(()))
        } else {
            Poll::Pending
        }
    }

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

        let mut capacity = 0;

        let mut first_channel_with_entries = None;

        for (idx, channel) in this.channels.iter_mut().enumerate() {
            // make one more effort to acquire entries in the rings
            let len = channel.tx.acquire(1);
            let len = channel.completion.acquire(len).min(len);
            trace!("acquired {len} entries for channel {idx}");
            capacity += len as usize;

            if len > 0 && first_channel_with_entries.is_none() {
                first_channel_with_entries = Some(idx);
            }
        }

        let channels = &mut this.channels;
        let umem = &mut this.umem;
        let encoder = &mut this.encoder;

        // use the first channel that had entries, otherwise return the length, which will indicate
        // the queue has no free items
        let channel_index = first_channel_with_entries.unwrap_or_else(|| channels.len());

        let mut queue = Queue {
            channels,
            channel_index,
            channel_needs_wake: false,
            capacity,
            umem,
            encoder,
        };
        f(&mut queue);
    }

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

pub struct Queue<'a, D: Driver> {
    channels: &'a mut Vec<Channel<D>>,
    /// The current index into the channels list
    channel_index: usize,
    /// Indicates if the current channel needs to be woken up
    channel_needs_wake: bool,
    /// The remaining capacity of the queue
    capacity: usize,
    umem: &'a mut Umem,
    encoder: &'a mut encoder::State,
}

impl<'a, D: Driver> tx::Queue for Queue<'a, D> {
    type Handle = path::Tuple;

    const SUPPORTS_ECN: bool = true;
    const SUPPORTS_FLOW_LABELS: bool = true;

    #[inline]
    fn push<M>(&mut self, mut message: M) -> Result<tx::Outcome, tx::Error>
    where
        M: tx::Message<Handle = Self::Handle>,
    {
        // if we're at capacity, then return an error
        if self.capacity == 0 {
            trace!("at capacity");
            return Err(tx::Error::AtCapacity);
        }

        let channel = loop {
            let channel = if let Some(channel) = self.channels.get_mut(self.channel_index) {
                channel
            } else {
                // we got to the end of the list without any more capacity
                return Err(tx::Error::AtCapacity);
            };

            // if this channel has entries, return it
            if !channel.is_empty() {
                trace!("selecting channel {}", self.channel_index);
                break channel;
            }

            // before moving on to the next channel, wake the current one if needed
            if core::mem::take(&mut self.channel_needs_wake) {
                trace!("waking channel {}", self.channel_index);
                channel.wake();
            }
            self.channel_index += 1;
        };

        // get the first descriptor in the ring
        let (entries, _) = channel.completion.data();
        let descriptor = entries[0];

        trace!("using descriptor {descriptor:?}");

        let buffer = unsafe {
            // Safety: this descriptor should be unique, assuming the tasks are functioning
            // properly
            self.umem.get_mut(descriptor)
        };

        // create an encoder for the descriptor region
        let mut buffer = EncoderBuffer::new(buffer);

        // write the message to the encoder using the configuration
        let payload_len = encoder::encode_packet(&mut buffer, &mut message, self.encoder)?;

        // take the length that we wrote and create a RxTxDescriptor with it
        let len = buffer.len();
        let descriptor = descriptor.with_len(len as _);

        trace!("packet written to {descriptor:?}");

        // push the descriptor on so it can be transmitted
        channel.tx.data().0[0] = descriptor;

        // make sure we give capacity back to the free queue
        channel.tx.release(1);
        channel.completion.release(1);

        // remember that we pushed something to the channel so it needs to be woken up
        self.channel_needs_wake = true;

        // decrement the total capacity after pushing another packet
        self.capacity -= 1;

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

        Ok(outcome)
    }

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

impl<'a, D: Driver> Drop for Queue<'a, D> {
    #[inline]
    fn drop(&mut self) {
        // if the current channel pushed some items it needs to be woken up
        if self.channel_needs_wake {
            if let Some(channel) = self.channels.get_mut(self.channel_index) {
                trace!("waking channel {}", self.channel_index);
                channel.wake();
            }
        }
    }
}


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

use super::Driver;
use crate::{ring, socket, syscall};
use core::task::{Context, Poll};

type Fd = tokio::io::unix::AsyncFd<socket::Fd>;

/// Polls read readiness for a tokio socket
///
/// In the case of a successful poll, the function will return `Some(count)`. Note that `count` may
/// still be `0`. If `None` is returned, an error was encountered and the task should be shut down.
#[inline]
fn poll(fd: &Fd, rx: &mut ring::Rx, fill: &mut ring::Fill, cx: &mut Context) -> Option<u32> {
    // iterate twice to avoid race conditions on waker registration
    for _ in 0..2 {
        let mut count = rx.acquire(u32::MAX);
        count = fill.acquire(count).min(count);

        // if we have entries in the rings, then return
        if count > 0 {
            return Some(count);
        }

        match fd.poll_read_ready(cx) {
            Poll::Ready(Ok(mut guard)) => {
                // since we don't have any entries, clear the readiness and try again
                guard.clear_ready();
                continue;
            }
            Poll::Ready(Err(_)) => {
                // the fd is no longer registered so shut down the task
                return None;
            }
            Poll::Pending => {
                // put the task to sleep until tokio wakes it up with Rx progress
                return Some(0);
            }
        }
    }

    // If we got here tokio said the socket was ready to read, even though the ring is pending. In
    // this case, we'll manually call the socket's busy poll method and wake up again to try to
    // acquire more items.
    //
    // It's very unlikely this happens, but it's good to have just in case so we don't occupy all
    // of the worker's cycles.
    let _ = syscall::busy_poll(fd.get_ref());
    cx.waker().wake_by_ref();

    Some(0)
}

/// Polling implementation for an asynchronous socket
impl Driver for Fd {
    #[inline]
    fn poll(&mut self, rx: &mut ring::Rx, fill: &mut ring::Fill, cx: &mut Context) -> Option<u32> {
        poll(self, rx, fill, cx)
    }
}

/// Polling implementation for a shared asynchronous socket
impl Driver for std::sync::Arc<Fd> {
    #[inline]
    fn poll(&mut self, rx: &mut ring::Rx, fill: &mut ring::Fill, cx: &mut Context) -> Option<u32> {
        poll(self, rx, fill, cx)
    }
}


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

use anyhow::Context;
use aya::{
    programs::{Xdp, XdpFlags},
    Bpf,
};
use aya_log::BpfLogger;
use clap::Parser;
use log::{info, warn};
use tokio::signal;

#[derive(Debug, Parser)]
struct Opt {
    /// The interface to run the program on
    #[clap(short, long, default_value = "lo")]
    iface: String,

    /// Traces BPF events
    #[clap(long)]
    trace: bool,

    /// Exits after attaching the BPF program
    ///
    /// This can be used to validate the correctness of the BPF program.
    #[clap(long)]
    exit_on_load: bool,
}

#[tokio::main]
async fn main() -> Result<(), anyhow::Error> {
    let opt = Opt::parse();

    env_logger::init();

    let bpf = if opt.trace {
        s2n_quic_xdp::bpf::DEFAULT_PROGRAM_TRACE
    } else {
        s2n_quic_xdp::bpf::DEFAULT_PROGRAM
    };

    let mut bpf = Bpf::load(bpf)?;

    if opt.trace {
        if let Err(e) = BpfLogger::init(&mut bpf) {
            warn!("failed to initialize eBPF logger: {}", e);
        }
    }

    let program: &mut Xdp = bpf
        .program_mut(s2n_quic_xdp::bpf::PROGRAM_NAME)
        .unwrap()
        .try_into()?;
    program.load()?;

    if opt.exit_on_load {
        return Ok(());
    }

    program.attach(&opt.iface, XdpFlags::default())
        .context("failed to attach the XDP program with default flags - try changing XdpFlags::default() to XdpFlags::SKB_MODE")?;

    info!("Waiting for Ctrl-C...");
    signal::ctrl_c().await?;
    info!("Exiting...");

    Ok(())
}


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

use std::path::Path;

pub fn run() -> Result<(), anyhow::Error> {
    let root = Path::new(env!("CARGO_MANIFEST_DIR"));

    let bindings = bindgen::Builder::default()
        .header(root.join("src/bindings/input.h").display().to_string())
        .allowlist_var("ETHTOOL_GCHANNELS")
        .allowlist_type("ethtool_channels")
        .rust_target(bindgen::RustTarget::Stable_1_47)
        .layout_tests(false)
        .raw_line(
            r#"
// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
// SPDX-License-Identifier: Apache-2.0

#![allow(non_camel_case_types)]
            "#
            .trim(),
        )
        .generate()?;

    let out = root.join("../s2n-quic-xdp/src/bindings.rs");
    bindings.write_to_file(out)?;

    Ok(())
}


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

use std::{path::PathBuf, process::Command};

pub fn run() -> Result<(), anyhow::Error> {
    for arch in ["bpfel", "bpfeb"] {
        let target = format!("{arch}-unknown-none");
        let target = ⌖

        for features in [None, Some("trace")] {
            eprintln!("    Building {target}");

            let mut command = Command::new("cargo");

            let dir = PathBuf::from("ebpf");
            command.current_dir(&dir).args([
                "build",
                "--target",
                target,
                "-Z",
                "build-std=core",
                "--release",
            ]);

            // Command::new creates a child process which inherits all env variables. This means env
            // vars set by the cargo xtask command are also inherited. RUSTUP_TOOLCHAIN is removed
            // so the rust-toolchain.toml file in the -ebpf folder is honored.
            command.env_remove("RUSTUP_TOOLCHAIN");

            let mut artifact_name = format!("s2n-quic-xdp-{arch}");

            let mut cargo_out = PathBuf::from("target")
                .join(target)
                .join("release")
                .join("s2n-quic-xdp");

            if let Some(features) = features {
                command.arg("--features");
                command.arg(features);
                command.arg("--target-dir");
                command.arg(PathBuf::from("../target").join(features));

                artifact_name += &format!("-{features}");

                cargo_out = PathBuf::from("target")
                    .join(features)
                    .join(target)
                    .join("release")
                    .join("s2n-quic-xdp");
            }

            let status = command.status().expect("failed to build bpf program");
            assert!(status.success());

            let out = PathBuf::from("s2n-quic-xdp/src/bpf")
                .join(artifact_name)
                .with_extension("ebpf");

            std::fs::create_dir_all(out.parent().unwrap())?;
            std::fs::copy(cargo_out, &out)?;
        }
    }
    Ok(())
}


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

use crate::{build_ebpf, disasm, run};
use anyhow::anyhow;
use std::{io, path::Path, process::Command};

pub fn run() -> Result<(), anyhow::Error> {
    let before = dump()?;
    build_ebpf::run()?;
    let after = dump()?;

    if before != after {
        let diff = similar::TextDiff::from_lines(&before, &after);
        let mut diff = diff.unified_diff();
        let out = diff.context_radius(10).header("committed", "current");

        println!("{out}");

        return Err(anyhow!("Difference detected"));
    }

    let mut opts = run::Options::default();

    opts.run_args.push("--exit-on-load".into());
    run::run(opts.clone())?;

    opts.run_args.push("--trace".into());
    run::run(opts)?;

    // run the normal tests first
    test()?;

    // run the CAP_NET_RAW tests after
    cap_net_raw_tests()?;

    Ok(())
}

fn dump() -> Result<String, anyhow::Error> {
    let mut output = vec![];
    for input in ["bpfel", "bpfeb", "bpfel-trace", "bpfeb-trace"] {
        disasm::piped(input, io::Cursor::new(&mut output))?;
    }
    let output = String::from_utf8(output)?;
    Ok(output)
}

fn test() -> Result<(), anyhow::Error> {
    let status = Command::new("cargo")
        .arg("test")
        .status()
        .expect("failed to run tests");
    assert!(status.success());
    Ok(())
}

fn cap_net_raw_tests() -> Result<(), anyhow::Error> {
    let mut at_least_one = false;

    for file in std::fs::read_dir("target/debug/deps")?.flatten() {
        let path = file.path();

        if path.extension().is_some()
            || !path
                .file_name()
                .and_then(|name| name.to_str())
                .map_or(false, |name| name.starts_with("s2n_quic_xdp"))
        {
            continue;
        }

        let status = Command::new("sudo")
            .arg(Path::new("../").join(path))
            .arg("--nocapture")
            .current_dir("s2n-quic-xdp")
            .env("CAP_NET_RAW_ENABLED", "1")
            .status()
            .expect("failed to run test case");
        assert!(status.success());

        at_least_one = true;
    }

    assert!(at_least_one);

    Ok(())
}


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

use crate::build_ebpf;
use anyhow::Context;
use std::io;

type File<'a> = elf::ElfBytes<'a, elf::endian::AnyEndian>;

pub fn run() -> Result<(), anyhow::Error> {
    build_ebpf::run().context("Error while building eBPF program")?;

    piped("bpfel", std::io::stdout())?;

    Ok(())
}

pub fn piped<O: io::Write>(input: &str, mut out: O) -> Result<(), anyhow::Error> {
    let prog = std::fs::read(format!("s2n-quic-xdp/src/bpf/s2n-quic-xdp-{input}.ebpf"))?;
    let file = File::minimal_parse(&prog)?;

    writeln!(out, "===== {input} =====")?;
    dump(&file, "xdp/s2n_quic_xdp", &mut out)?;
    dump(&file, ".text", &mut out)?;

    Ok(())
}

fn dump<O: io::Write>(file: &File, section: &str, mut out: O) -> Result<(), anyhow::Error> {
    if let Some(header) = file.section_header_by_name(section)? {
        if let Ok((data, _compression)) = file.section_data(&header) {
            writeln!(out, "----- {section} -----")?;
            for insn in rbpf::disassembler::to_insn_vec(data) {
                writeln!(out, "  {}", insn.desc)?;
            }
        }
    }

    Ok(())
}


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

mod bindings;
mod build_ebpf;
mod ci;
mod disasm;
mod run;

use std::process::exit;

use clap::Parser;

#[derive(Debug, Parser)]
pub struct Options {
    #[clap(subcommand)]
    command: Command,
}

#[derive(Debug, Parser)]
enum Command {
    BuildEbpf,
    Bindings,
    Ci,
    Disasm,
    Run(run::Options),
}

fn main() {
    let opts = Options::parse();

    use Command::*;
    let ret = match opts.command {
        BuildEbpf => build_ebpf::run(),
        Bindings => bindings::run(),
        Ci => ci::run(),
        Disasm => disasm::run(),
        Run(opts) => run::run(opts),
    };

    if let Err(e) = ret {
        eprintln!("{e:#}");
        exit(1);
    }
}


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

use crate::build_ebpf;
use anyhow::Context as _;
use clap::Parser;
use std::process::Command;

#[derive(Clone, Debug, Parser)]
pub struct Options {
    /// Build and run the release target
    #[clap(long)]
    pub release: bool,
    /// The command used to wrap your application
    #[clap(short, long, default_value = "sudo -E")]
    pub runner: String,
    /// Skips building the bpf programs
    #[clap(long)]
    pub no_build: bool,
    /// Arguments to pass to your application
    #[clap(name = "args", last = true)]
    pub run_args: Vec<String>,
}

impl Default for Options {
    fn default() -> Self {
        Options::parse_from::<_, &str>([])
    }
}

/// Build the project
fn build(opts: &Options) -> Result<(), anyhow::Error> {
    let mut args = vec!["build"];
    if opts.release {
        args.push("--release")
    }
    let status = Command::new("cargo")
        .args(&args)
        .status()
        .expect("failed to build userspace");
    assert!(status.success());
    Ok(())
}

/// Build and run the project
pub fn run(opts: Options) -> Result<(), anyhow::Error> {
    // build our ebpf program followed by our application
    if !opts.no_build {
        build_ebpf::run().context("Error while building eBPF program")?;
    }
    build(&opts).context("Error while building userspace application")?;

    // profile we are building (release or debug)
    let profile = if opts.release { "release" } else { "debug" };
    let bin_path = format!("target/{profile}/tester");

    // arguments to pass to the application
    let mut run_args: Vec<_> = opts.run_args.iter().map(String::as_str).collect();

    // configure args
    let mut args: Vec<_> = opts.runner.trim().split_terminator(' ').collect();
    args.push(bin_path.as_str());
    args.append(&mut run_args);

    // run the command
    let status = Command::new(args.first().expect("No first argument"))
        .args(args.iter().skip(1))
        .status()
        .expect("failed to run the command");

    if !status.success() {
        anyhow::bail!("Failed to run `{}`", args.join(" "));
    }
    Ok(())
}


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

use std::any::type_name;

use bench::{
    CipherSuite::*,
    CryptoConfig,
    ECGroup::{self, *},
    OpenSslHarness, RustlsHarness, S2NHarness, TlsBenchHarness,
};
use criterion::{
    criterion_group, criterion_main, measurement::WallTime, BatchSize, BenchmarkGroup, Criterion,
};

pub fn bench_handshake_key_exchange(c: &mut Criterion) {
    fn bench_handshake_for_library<T: TlsBenchHarness>(
        bench_group: &mut BenchmarkGroup<WallTime>,
        ec_group: &ECGroup,
    ) {
        bench_group.bench_function(type_name::<T>(), |b| {
            b.iter_batched_ref(
                || {
                    T::new(&CryptoConfig {
                        cipher_suite: AES_128_GCM_SHA256,
                        ec_group: *ec_group,
                    })
                    .unwrap()
                },
                |harness| {
                    harness.handshake().unwrap();
                },
                BatchSize::SmallInput,
            )
        });
    }

    for ec_group in [SECP256R1, X25519] {
        let mut bench_group = c.benchmark_group(format!("handshake-{:?}", ec_group));
        bench_handshake_for_library::<S2NHarness>(&mut bench_group, &ec_group);
        bench_handshake_for_library::<RustlsHarness>(&mut bench_group, &ec_group);
        bench_handshake_for_library::<OpenSslHarness>(&mut bench_group, &ec_group);
    }
}

criterion_group!(benches, bench_handshake_key_exchange);
criterion_main!(benches);


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

use std::{
    cell::RefCell,
    collections::VecDeque,
    error::Error,
    fs::read_to_string,
    io::{ErrorKind, Read, Write},
    rc::Rc,
};

pub fn read_to_bytes(path: &str) -> Vec<u8> {
    read_to_string(path).unwrap().into_bytes()
}

pub enum Mode {
    Client,
    Server,
}

// these parameters were the only ones readily usable for all three libaries:
// s2n-tls, rustls, and openssl
#[allow(non_camel_case_types)]
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum CipherSuite {
    AES_128_GCM_SHA256,
    AES_256_GCM_SHA384,
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum ECGroup {
    SECP256R1,
    X25519,
}

#[derive(Debug, PartialEq, Eq)]
pub struct CryptoConfig {
    pub cipher_suite: CipherSuite,
    pub ec_group: ECGroup,
}

pub trait TlsBenchHarness: Sized {
    /// Default harness
    fn default() -> Result<Self, Box<dyn Error>> {
        Self::new(&CryptoConfig {
            cipher_suite: CipherSuite::AES_128_GCM_SHA256,
            ec_group: ECGroup::X25519,
        })
    }

    /// Initialize buffers, configs, and connections (pre-handshake)
    fn new(crypto_config: &CryptoConfig) -> Result<Self, Box<dyn Error>>;

    /// Run handshake on initialized connection
    /// Returns error if handshake has already completed
    fn handshake(&mut self) -> Result<(), Box<dyn Error>>;

    /// Checks if handshake is finished for both client and server
    fn handshake_completed(&self) -> bool;

    /// Get negotiated cipher suite
    fn get_negotiated_cipher_suite(&self) -> CipherSuite;

    /// Get whether or negotiated version is TLS1.3
    fn negotiated_tls13(&self) -> bool;
}

/// Wrapper of two shared buffers to pass as stream
/// This wrapper `read()`s into one buffer and `write()`s to another
#[derive(Clone)]
pub struct ConnectedBuffer {
    recv: Rc<RefCell<VecDeque<u8>>>,
    send: Rc<RefCell<VecDeque<u8>>>,
}

impl ConnectedBuffer {
    /// Make a new struct with new internal buffers
    pub fn new() -> Self {
        ConnectedBuffer {
            recv: Rc::new(RefCell::new(VecDeque::new())),
            send: Rc::new(RefCell::new(VecDeque::new())),
        }
    }
    /// Make a new struct that shares internal buffers but swapped, ex.
    /// `write()` writes to the buffer that the inverse `read()`s from
    pub fn clone_inverse(&self) -> Self {
        ConnectedBuffer {
            recv: Rc::clone(&self.send),
            send: Rc::clone(&self.recv),
        }
    }
}

impl Read for ConnectedBuffer {
    fn read(&mut self, dest: &mut [u8]) -> Result<usize, std::io::Error> {
        match self.recv.borrow_mut().read(dest) {
            // rustls expects WouldBlock on read of length 0
            Ok(0) => Err(std::io::Error::new(ErrorKind::WouldBlock, "blocking")),
            Ok(len) => Ok(len),
            Err(err) => Err(err),
        }
    }
}

impl Write for ConnectedBuffer {
    fn write(&mut self, src: &[u8]) -> Result<usize, std::io::Error> {
        self.send.borrow_mut().write(src)
    }
    fn flush(&mut self) -> Result<(), std::io::Error> {
        Ok(()) // data already available to destination
    }
}

#[cfg(test)]
macro_rules! test_tls_bench_harnesses {
    ($($lib_name:ident: $harness_type:ty,)*) => {
    $(
        mod $lib_name {
            use super::*;

            #[test]
            fn test_handshake() {
                let mut harness = <$harness_type>::default().unwrap();
                assert!(!harness.handshake_completed());
                harness.handshake().unwrap();
                assert!(harness.handshake_completed());
                assert!(harness.negotiated_tls13());
            }

            #[test]
            fn test_different_crypto_config() {
                use CipherSuite::*;
                use ECGroup::*;

                let (mut harness, mut crypto_config);
                for cipher_suite in [AES_128_GCM_SHA256, AES_256_GCM_SHA384].iter() {
                    for ec_group in [SECP256R1, X25519].iter() {
                        crypto_config = CryptoConfig { cipher_suite: cipher_suite.clone(), ec_group: ec_group.clone() };
                        harness = <$harness_type>::new(&crypto_config).unwrap();
                        harness.handshake().unwrap();
                        assert_eq!(cipher_suite, &harness.get_negotiated_cipher_suite());
                    }
                }
            }
        }
    )*
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::{OpenSslHarness, RustlsHarness, S2NHarness, TlsBenchHarness};

    test_tls_bench_harnesses! {
        s2n_tls: S2NHarness,
        rustls: RustlsHarness,
        openssl: OpenSslHarness,
    }
}


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

pub mod harness;
pub mod openssl;
pub mod rustls;
pub mod s2n_tls;
pub use crate::{
    harness::{CipherSuite, CryptoConfig, ECGroup, TlsBenchHarness},
    openssl::OpenSslHarness,
    rustls::RustlsHarness,
    s2n_tls::S2NHarness,
};

const SERVER_KEY_PATH: &str = "certs/server-key.pem";
const SERVER_CERT_CHAIN_PATH: &str = "certs/fullchain.pem";
const CA_CERT_PATH: &str = "certs/ca-cert.pem";

#[cfg(test)]
mod tests {
    use std::path::Path;

    #[test]
    fn cert_paths_valid() {
        assert!(Path::new(crate::SERVER_KEY_PATH).exists());
        assert!(Path::new(crate::SERVER_CERT_CHAIN_PATH).exists());
        assert!(Path::new(crate::CA_CERT_PATH).exists());
    }
}


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

use crate::{
    harness::{CipherSuite, ConnectedBuffer, CryptoConfig, ECGroup, Mode, TlsBenchHarness},
    CA_CERT_PATH, SERVER_CERT_CHAIN_PATH, SERVER_KEY_PATH,
};
use openssl::ssl::{
    ErrorCode, Ssl, SslContext, SslContextBuilder, SslFiletype, SslMethod, SslStream, SslVersion,
};
use std::error::Error;

pub struct OpenSslHarness {
    client_conn: SslStream<ConnectedBuffer>,
    server_conn: SslStream<ConnectedBuffer>,
}

impl OpenSslHarness {
    fn common_config(
        builder: &mut SslContextBuilder,
        cipher_suite: &str,
        ec_key: &str,
    ) -> Result<(), Box<dyn Error>> {
        builder.set_min_proto_version(Some(SslVersion::TLS1_3))?;
        builder.set_ciphersuites(cipher_suite)?;
        builder.set_groups_list(ec_key)?;
        Ok(())
    }
    /// Process handshake for one connection, treating blocking errors as `Ok`
    fn handshake_conn(&mut self, mode: Mode) -> Result<(), Box<dyn Error>> {
        let result = match mode {
            Mode::Client => self.client_conn.connect(),
            Mode::Server => self.server_conn.accept(),
        };
        match result {
            Ok(_) => Ok(()),
            Err(err) => {
                if err.code() != ErrorCode::WANT_READ {
                    Err(err.into())
                } else {
                    Ok(())
                }
            }
        }
    }
}

impl TlsBenchHarness for OpenSslHarness {
    fn new(crypto_config: &CryptoConfig) -> Result<Self, Box<dyn Error>> {
        let client_buf = ConnectedBuffer::new();
        let server_buf = client_buf.clone_inverse();

        let cipher_suite = match crypto_config.cipher_suite {
            CipherSuite::AES_128_GCM_SHA256 => "TLS_AES_128_GCM_SHA256",
            CipherSuite::AES_256_GCM_SHA384 => "TLS_AES_256_GCM_SHA384",
        };

        let ec_key = match crypto_config.ec_group {
            ECGroup::SECP256R1 => "P-256",
            ECGroup::X25519 => "X25519",
        };

        let mut client_builder = SslContext::builder(SslMethod::tls_client())?;
        client_builder.set_ca_file(CA_CERT_PATH)?;
        Self::common_config(&mut client_builder, cipher_suite, ec_key)?;

        let mut server_builder = SslContext::builder(SslMethod::tls_server())?;
        server_builder.set_certificate_chain_file(SERVER_CERT_CHAIN_PATH)?;
        server_builder.set_private_key_file(SERVER_KEY_PATH, SslFiletype::PEM)?;
        Self::common_config(&mut server_builder, cipher_suite, ec_key)?;

        let client_config = client_builder.build();
        let server_config = server_builder.build();

        let client_conn = SslStream::new(Ssl::new(&client_config)?, client_buf)?;
        let server_conn = SslStream::new(Ssl::new(&server_config)?, server_buf)?;

        Ok(Self {
            client_conn,
            server_conn,
        })
    }

    fn handshake(&mut self) -> Result<(), Box<dyn Error>> {
        for _ in 0..2 {
            self.handshake_conn(Mode::Client)?;
            self.handshake_conn(Mode::Server)?;
        }
        Ok(())
    }

    fn handshake_completed(&self) -> bool {
        self.client_conn.ssl().is_init_finished() && self.server_conn.ssl().is_init_finished()
    }

    fn get_negotiated_cipher_suite(&self) -> CipherSuite {
        let cipher_suite = self
            .client_conn
            .ssl()
            .current_cipher()
            .expect("Handshake not completed")
            .name();
        match cipher_suite {
            "TLS_AES_128_GCM_SHA256" => CipherSuite::AES_128_GCM_SHA256,
            "TLS_AES_256_GCM_SHA384" => CipherSuite::AES_256_GCM_SHA384,
            _ => panic!("Unknown cipher suite"),
        }
    }

    fn negotiated_tls13(&self) -> bool {
        self.client_conn
            .ssl()
            .version2() // version() -> &str is deprecated, version2() returns an enum instead
            .expect("Handshake not completed")
            == SslVersion::TLS1_3
    }
}


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

use crate::{
    harness::{
        read_to_bytes, CipherSuite, ConnectedBuffer, CryptoConfig, ECGroup, Mode, TlsBenchHarness,
    },
    CA_CERT_PATH, SERVER_CERT_CHAIN_PATH, SERVER_KEY_PATH,
};
use rustls::{
    cipher_suite::{TLS13_AES_128_GCM_SHA256, TLS13_AES_256_GCM_SHA384},
    kx_group::{SECP256R1, X25519},
    version::TLS13,
    Certificate, ClientConfig, ClientConnection, PrivateKey,
    ProtocolVersion::TLSv1_3,
    RootCertStore, ServerConfig, ServerConnection, ServerName,
};
use rustls_pemfile::{certs, pkcs8_private_keys};
use std::{error::Error, io::BufReader, sync::Arc};

pub struct RustlsHarness {
    client_buf: ConnectedBuffer,
    server_buf: ConnectedBuffer,
    client_conn: ClientConnection,
    server_conn: ServerConnection,
}

impl RustlsHarness {
    fn get_root_cert_store() -> Result<RootCertStore, Box<dyn Error>> {
        let root_cert =
            Certificate(certs(&mut BufReader::new(&*read_to_bytes(CA_CERT_PATH)))?.remove(0));
        let mut root_certs = RootCertStore::empty();
        root_certs.add(&root_cert)?;
        Ok(root_certs)
    }

    fn get_cert_chain() -> Result<Vec<Certificate>, Box<dyn Error>> {
        let chain = certs(&mut BufReader::new(&*read_to_bytes(SERVER_CERT_CHAIN_PATH)))?;
        Ok(chain
            .iter()
            .map(|bytes| Certificate(bytes.to_vec()))
            .collect())
    }

    fn get_server_key() -> Result<PrivateKey, Box<dyn Error>> {
        Ok(PrivateKey(
            pkcs8_private_keys(&mut BufReader::new(&*read_to_bytes(SERVER_KEY_PATH)))?.remove(0),
        ))
    }

    /// Read all incoming data, process it, and write it out
    /// Works for handshaking and data transfer
    fn process_conn(&mut self, mode: Mode) -> Result<(), std::io::Error> {
        match match mode {
            Mode::Client => self.client_conn.complete_io(&mut self.client_buf),
            Mode::Server => self.server_conn.complete_io(&mut self.server_buf),
        } {
            Ok(_) => Ok(()),
            Err(err) => {
                if err.kind() == std::io::ErrorKind::WouldBlock {
                    Ok(())
                } else {
                    Err(err)
                }
            }
        }
    }
}

impl TlsBenchHarness for RustlsHarness {
    fn new(crypto_config: &CryptoConfig) -> Result<Self, Box<dyn Error>> {
        let client_buf = ConnectedBuffer::new();
        let server_buf = client_buf.clone_inverse();

        let cipher_suite = match crypto_config.cipher_suite {
            CipherSuite::AES_128_GCM_SHA256 => TLS13_AES_128_GCM_SHA256,
            CipherSuite::AES_256_GCM_SHA384 => TLS13_AES_256_GCM_SHA384,
        };

        let kx_group = match crypto_config.ec_group {
            ECGroup::SECP256R1 => &SECP256R1,
            ECGroup::X25519 => &X25519,
        };

        let client_config = Arc::new(
            ClientConfig::builder()
                .with_cipher_suites(&[cipher_suite])
                .with_kx_groups(&[kx_group])
                .with_protocol_versions(&[&TLS13])?
                .with_root_certificates(Self::get_root_cert_store()?)
                .with_no_client_auth(),
        );

        let server_config = Arc::new(
            ServerConfig::builder()
                .with_cipher_suites(&[cipher_suite])
                .with_kx_groups(&[kx_group])
                .with_protocol_versions(&[&TLS13])?
                .with_no_client_auth()
                .with_single_cert(Self::get_cert_chain()?, Self::get_server_key()?)?,
        );

        let client_conn = ClientConnection::new(client_config, ServerName::try_from("localhost")?)?;
        let server_conn = ServerConnection::new(server_config)?;

        Ok(Self {
            client_buf,
            server_buf,
            client_conn,
            server_conn,
        })
    }

    fn handshake(&mut self) -> Result<(), Box<dyn Error>> {
        for _ in 0..2 {
            self.process_conn(Mode::Client)?;
            self.process_conn(Mode::Server)?;
        }
        Ok(())
    }

    fn handshake_completed(&self) -> bool {
        !self.client_conn.is_handshaking() && !self.server_conn.is_handshaking()
    }

    fn get_negotiated_cipher_suite(&self) -> CipherSuite {
        match self.client_conn.negotiated_cipher_suite().unwrap().suite() {
            rustls::CipherSuite::TLS13_AES_128_GCM_SHA256 => CipherSuite::AES_128_GCM_SHA256,
            rustls::CipherSuite::TLS13_AES_256_GCM_SHA384 => CipherSuite::AES_256_GCM_SHA384,
            _ => panic!("Unknown cipher suite"),
        }
    }

    fn negotiated_tls13(&self) -> bool {
        self.client_conn
            .protocol_version()
            .expect("Handshake not completed")
            == TLSv1_3
    }
}


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

use crate::{
    harness::{read_to_bytes, CipherSuite, CryptoConfig, ECGroup, Mode, TlsBenchHarness},
    CA_CERT_PATH, SERVER_CERT_CHAIN_PATH, SERVER_KEY_PATH,
};
use s2n_tls::{
    callbacks::VerifyHostNameCallback,
    config::{Builder, Config},
    connection::Connection,
    enums::{Blinding, Version},
    security::Policy,
};
use std::{
    cell::UnsafeCell,
    collections::VecDeque,
    error::Error,
    ffi::c_void,
    io::{Read, Write},
    os::raw::c_int,
    pin::Pin,
    task::Poll::Ready,
};

pub struct S2NHarness {
    // UnsafeCell is needed b/c client and server share *mut to IO buffers
    // Pin<Box<T>> is to ensure long-term *mut to IO buffers remain valid
    client_to_server_buf: Pin<Box<UnsafeCell<VecDeque<u8>>>>,
    server_to_client_buf: Pin<Box<UnsafeCell<VecDeque<u8>>>>,
    client_config: Config,
    server_config: Config,
    client_conn: Connection,
    server_conn: Connection,
    client_handshake_completed: bool,
    server_handshake_completed: bool,
}

/// Custom callback for verifying hostnames. Rustls requires checking hostnames,
/// so this is to make a fair comparison
struct HostNameHandler<'a> {
    expected_server_name: &'a str,
}
impl VerifyHostNameCallback for HostNameHandler<'_> {
    fn verify_host_name(&self, hostname: &str) -> bool {
        self.expected_server_name == hostname
    }
}

impl S2NHarness {
    /// Unsafe callback for custom IO C API
    ///
    /// s2n-tls IO is usually used with file descriptors to a TCP socket, but we
    /// reduce overhead and outside noise with a local buffer for benchmarking
    unsafe extern "C" fn send_cb(context: *mut c_void, data: *const u8, len: u32) -> c_int {
        let context = &mut *(context as *mut VecDeque<u8>);
        let data = core::slice::from_raw_parts(data, len as _);
        context.write(data).unwrap() as _
    }

    /// Unsafe callback for custom IO C API
    unsafe extern "C" fn recv_cb(context: *mut c_void, data: *mut u8, len: u32) -> c_int {
        let context = &mut *(context as *mut VecDeque<u8>);
        let data = core::slice::from_raw_parts_mut(data, len as _);
        context.flush().unwrap();
        let len = context.read(data).unwrap();
        if len == 0 {
            errno::set_errno(errno::Errno(libc::EWOULDBLOCK));
            -1
        } else {
            len as _
        }
    }

    fn create_config(mode: Mode, crypto_config: &CryptoConfig) -> Result<Config, Box<dyn Error>> {
        let security_policy = match (&crypto_config.cipher_suite, &crypto_config.ec_group) {
            (CipherSuite::AES_128_GCM_SHA256, ECGroup::SECP256R1) => "20230317",
            (CipherSuite::AES_256_GCM_SHA384, ECGroup::SECP256R1) => "20190802",
            (CipherSuite::AES_128_GCM_SHA256, ECGroup::X25519) => "default_tls13",
            (CipherSuite::AES_256_GCM_SHA384, ECGroup::X25519) => "20190801",
        };

        let mut builder = Builder::new();
        builder.set_security_policy(&Policy::from_version(security_policy)?)?;

        match mode {
            Mode::Server => builder.load_pem(
                read_to_bytes(SERVER_CERT_CHAIN_PATH).as_slice(),
                read_to_bytes(SERVER_KEY_PATH).as_slice(),
            )?,
            Mode::Client => builder
                .trust_pem(read_to_bytes(CA_CERT_PATH).as_slice())?
                .set_verify_host_callback(HostNameHandler {
                    expected_server_name: "localhost",
                })?,
        };

        Ok(builder.build()?)
    }

    /// Set up connections with config and custom IO
    fn init_conn(&mut self, mode: Mode) -> Result<(), Box<dyn Error>> {
        let client_to_server_ptr = self.client_to_server_buf.get() as *mut c_void;
        let server_to_client_ptr = self.server_to_client_buf.get() as *mut c_void;
        let (read_ptr, write_ptr, config, conn) = match mode {
            Mode::Client => (
                server_to_client_ptr,
                client_to_server_ptr,
                &self.client_config,
                &mut self.client_conn,
            ),
            Mode::Server => (
                client_to_server_ptr,
                server_to_client_ptr,
                &self.server_config,
                &mut self.server_conn,
            ),
        };

        conn.set_blinding(Blinding::SelfService)?
            .set_config(config.clone())?
            .set_send_callback(Some(Self::send_cb))?
            .set_receive_callback(Some(Self::recv_cb))?;
        unsafe {
            conn.set_send_context(write_ptr)?
                .set_receive_context(read_ptr)?;
        }

        Ok(())
    }

    /// Handshake step for one connection
    fn handshake_conn(&mut self, mode: Mode) -> Result<(), Box<dyn Error>> {
        let (conn, handshake_completed) = match mode {
            Mode::Client => (&mut self.client_conn, &mut self.client_handshake_completed),
            Mode::Server => (&mut self.server_conn, &mut self.server_handshake_completed),
        };

        if let Ready(res) = conn.poll_negotiate() {
            res?;
            *handshake_completed = true;
        } else {
            *handshake_completed = false;
        }
        Ok(())
    }
}

impl TlsBenchHarness for S2NHarness {
    fn new(crypto_config: &CryptoConfig) -> Result<Self, Box<dyn Error>> {
        let client_to_server_buf = Box::pin(UnsafeCell::new(VecDeque::new()));
        let server_to_client_buf = Box::pin(UnsafeCell::new(VecDeque::new()));

        let client_config = Self::create_config(Mode::Client, crypto_config)?;
        let server_config = Self::create_config(Mode::Server, crypto_config)?;

        let mut harness = Self {
            client_to_server_buf,
            server_to_client_buf,
            client_config,
            server_config,
            client_conn: Connection::new_client(),
            server_conn: Connection::new_server(),
            client_handshake_completed: false,
            server_handshake_completed: false,
        };

        harness.init_conn(Mode::Client)?;
        harness.init_conn(Mode::Server)?;

        Ok(harness)
    }

    fn handshake(&mut self) -> Result<(), Box<dyn Error>> {
        for _ in 0..2 {
            self.handshake_conn(Mode::Client)?;
            self.handshake_conn(Mode::Server)?;
        }
        Ok(())
    }

    fn handshake_completed(&self) -> bool {
        self.client_handshake_completed && self.server_handshake_completed
    }

    fn get_negotiated_cipher_suite(&self) -> CipherSuite {
        match self.client_conn.cipher_suite().unwrap() {
            "TLS_AES_128_GCM_SHA256" => CipherSuite::AES_128_GCM_SHA256,
            "TLS_AES_256_GCM_SHA384" => CipherSuite::AES_256_GCM_SHA384,
            _ => panic!("Unknown cipher suite"),
        }
    }

    fn negotiated_tls13(&self) -> bool {
        self.client_conn.actual_protocol_version().unwrap() == Version::TLS13
    }
}


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

use std::{
    collections::BTreeSet,
    fs::{self, read_to_string},
    io,
    path::Path,
    sync::{Arc, Mutex},
};

use bindgen::callbacks::ItemKind;

/// This is a placeholder that is replaced with the appropriate "feature token".
/// The placeholder is found in the *.template files of s2n-tls-sys/templates.
/// In Cargo.template this is replaced with the autogenerated list of features,
/// and in features.template this is replaced with the autogenerated list of
/// modules.
const FEATURE_TOKEN_PLACEHOLDER: &str = "<TOKEN_REPLACED_WITH_UNSTABLE_FEATURES>";

/// This binary is only expected to run in the context of the generate.sh script
/// which handles certain behaviors such as copying header files to
/// s2n-tls-sys/lib and other sundry actions.
fn main() {
    let out_dir = std::env::args().nth(1).expect("missing sys dir");
    let out_dir = Path::new(&out_dir);

    let functions = FunctionCallbacks::default();

    gen_bindings(
        "#include <s2n.h>",
        &out_dir.join("lib"),
        functions.with_feature(None),
    )
    .allowlist_type("s2n_.*")
    .allowlist_function("s2n_.*")
    .allowlist_var("s2n_.*")
    .generate()
    .unwrap()
    .write_to_file(out_dir.join("src/api.rs"))
    .unwrap();

    write_feature_bindings(
        out_dir.join("lib/tls/s2n_internal.h"),
        "internal",
        out_dir,
        out_dir.join("src/features/internal.rs"),
        functions.clone(),
    );
    write_feature_bindings(
        out_dir.join("lib/tls/s2n_quic_support.h"),
        "quic",
        out_dir,
        out_dir.join("src/features/quic.rs"),
        functions.clone(),
    );

    // get all of the files in the unstable folder
    let unstable_api = out_dir.join("lib/api/unstable");
    let unstable_headers: Vec<(String, fs::DirEntry)> = fs::read_dir(unstable_api)
        .expect("unable to iterate through files in unstable api folder")
        .into_iter()
        .map(|dir_entry| dir_entry.expect("failed to read header"))
        .map(|dir_entry| {
            (
                dir_entry
                    .path()
                    .file_stem()
                    .unwrap()
                    .to_str()
                    .unwrap()
                    .to_owned(),
                dir_entry,
            )
        })
        .collect();

    // write unstable bindings for them
    for (header_name, header) in unstable_headers.iter() {
        let feature_name = format!("unstable-{}", header_name);
        write_feature_bindings(
            header.path(),
            &feature_name,
            out_dir,
            out_dir.join(format!("src/features/{}.rs", header_name)),
            functions.clone(),
        );
    }

    // generate a cargo.toml that defines the correct features
    let features_definition_token = unstable_headers
        .iter()
        .map(|(header_name, _header)| format!("unstable-{header_name} = []"))
        .collect::<Vec<String>>()
        .join("\n");
    let cargo_template = out_dir.join("templates/Cargo.template");
    let cargo_template = read_to_string(cargo_template).expect("unable to read cargo template");
    let cargo_toml = cargo_template.replace(FEATURE_TOKEN_PLACEHOLDER, &features_definition_token);
    fs::write(out_dir.join("Cargo.toml"), cargo_toml).unwrap();

    // generate a features.rs that includes the correct modules
    let features_module_token = unstable_headers
        .iter()
        .map(|(header_name, _header)| {
            format!("conditional_module!({header_name}, \"unstable-{header_name}\");")
        })
        .collect::<Vec<String>>()
        .join("\n");
    let features_template = out_dir.join("templates/features.template");
    let features_template = read_to_string(features_template).expect("unable to features template");
    let features_rs = features_template.replace(FEATURE_TOKEN_PLACEHOLDER, &features_module_token);
    std::fs::write(out_dir.join("src/features.rs"), features_rs).unwrap();

    functions.tests(&out_dir.join("src/tests.rs")).unwrap();

    gen_files(&out_dir.join("lib"), &out_dir.join("files.rs")).unwrap();
}

const COPYRIGHT: &str = r#"
// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
// SPDX-License-Identifier: Apache-2.0
"#;

const PRELUDE: &str = r#"
#![allow(unused_imports, non_camel_case_types)]

use libc::{iovec, FILE};
"#;

fn base_builder() -> bindgen::Builder {
    bindgen::Builder::default()
        .use_core()
        .layout_tests(true)
        .detect_include_paths(true)
        .size_t_is_usize(true)
        .enable_function_attribute_detection()
        .default_enum_style(bindgen::EnumVariation::ModuleConsts)
        .rust_target(bindgen::RustTarget::Stable_1_47)
        // rust can't access thread-local variables
        // https://github.com/rust-lang/rust/issues/29594
        .blocklist_item("s2n_errno")
        .raw_line(COPYRIGHT)
        .raw_line(PRELUDE)
        .ctypes_prefix("::libc")
}

fn gen_bindings(entry: &str, s2n_dir: &Path, functions: FunctionCallbacks) -> bindgen::Builder {
    base_builder()
        .header_contents("s2n-sys.h", entry)
        // only export s2n-related stuff
        .blocklist_type("iovec")
        .blocklist_type("FILE")
        .blocklist_type("_IO_.*")
        .blocklist_type("__.*")
        .blocklist_type("fpos_t")
        .parse_callbacks(Box::new(functions))
        .clang_arg(format!("-I{}/api", s2n_dir.display()))
        .clang_arg(format!("-I{}", s2n_dir.display()))
}

fn write_feature_bindings(
    header_path: impl AsRef<Path>,
    feature_flag: &str,
    s2n_tls_sys_dir: &Path,
    output_path: impl AsRef<Path>,
    functions: FunctionCallbacks,
) {
    let header_path_str = format!("{}", header_path.as_ref().display());
    let lib_path = s2n_tls_sys_dir.join("lib");
    base_builder()
        .header(&header_path_str)
        .parse_callbacks(Box::new(
            functions.with_feature(Some(feature_flag.to_owned())),
        ))
        // manually include header contents
        .clang_arg(format!("-I{}/api", lib_path.display()))
        .clang_arg(format!("-I{}", lib_path.display()))
        .allowlist_recursively(false)
        .allowlist_file(&header_path_str)
        // s2n_internal.h defines opaque handles to these structs, but we want
        // them to be imported from the main api module
        .blocklist_type("s2n_connection")
        .blocklist_type("s2n_config")
        .raw_line("use crate::api::*;\n")
        .generate()
        .unwrap()
        .write_to_file(output_path)
        .unwrap();
}

fn gen_files(input: &Path, out: &Path) -> io::Result<()> {
    use io::Write;

    let mut files = std::fs::File::create(out)?;
    let mut o = io::BufWriter::new(&mut files);

    let pattern = format!("{}/**/*.c", input.display());

    writeln!(o, "{}", COPYRIGHT)?;
    writeln!(o, "[")?;
    for file in glob::glob(&pattern).unwrap() {
        let file = file.unwrap();
        let file = file.strip_prefix(input).unwrap();
        // don't include tests
        if file.starts_with("tests") {
            continue;
        }
        writeln!(o, "    {:?},", Path::new("lib").join(file).display())?;
    }
    writeln!(o, "]")?;
    Ok(())
}

type SharedBTreeSet<T> = Arc<Mutex<BTreeSet<T>>>;

#[derive(Clone, Debug, Default)]
struct FunctionCallbacks {
    /// the current feature that is having bindings generated
    feature: Arc<Mutex<Option<String>>>,
    /// a list of all functions that have had bindings generated for them
    functions: SharedBTreeSet<(Option<String>, String)>,
}

impl FunctionCallbacks {
    fn with_feature(&self, feature: Option<String>) -> Self {
        *self.feature.lock().unwrap() = feature;
        self.clone()
    }

    fn tests(&self, out: &Path) -> io::Result<()> {
        use io::Write;
        let functions = self.functions.lock().unwrap();
        let mut tests = std::fs::File::create(out)?;
        let mut o = io::BufWriter::new(&mut tests);

        writeln!(o, "{}", COPYRIGHT)?;
        let iter = functions.iter();
        for (feature, function) in iter {
            // don't generate a test if it's enabled without a feature
            if feature.is_some() && functions.contains(&(None, function.to_string())) {
                continue;
            }

            writeln!(o, "#[test]")?;

            // if the function is behind a feature, gate it with `cfg`
            if let Some(feature) = feature {
                writeln!(o, "#[cfg(feature = {:?})]", feature)?;
            };

            writeln!(o, "fn {} () {{", function)?;
            writeln!(o, "    let ptr = crate::{} as *const ();", function)?;
            writeln!(o, "    assert!(!ptr.is_null());")?;
            writeln!(o, "}}")?;
            writeln!(o)?;
        }

        Ok(())
    }
}

impl bindgen::callbacks::ParseCallbacks for FunctionCallbacks {
    fn enum_variant_name(
        &self,
        _name: Option<&str>,
        variant_name: &str,
        _variant_value: bindgen::callbacks::EnumVariantValue,
    ) -> Option<String> {
        if !variant_name.starts_with("S2N_") {
            return None;
        }

        let variant_name = variant_name
            .trim_start_matches("S2N_ERR_T_")
            .trim_start_matches("S2N_EXTENSION_")
            // keep the LEN_ so it's a valid identifier
            .trim_start_matches("S2N_TLS_MAX_FRAG_")
            .trim_start_matches("S2N_ALERT_")
            .trim_start_matches("S2N_CT_SUPPORT_")
            .trim_start_matches("S2N_STATUS_REQUEST_")
            .trim_start_matches("S2N_CERT_AUTH_")
            .trim_start_matches("S2N_CLIENT_HELLO_CB_")
            .trim_start_matches("S2N_TLS_SIGNATURE_")
            .trim_start_matches("S2N_TLS_HASH_")
            .trim_start_matches("S2N_PSK_HMAC_")
            .trim_start_matches("S2N_PSK_MODE_")
            .trim_start_matches("S2N_ASYNC_PKEY_VALIDATION_")
            .trim_start_matches("S2N_ASYNC_")
            .trim_start_matches("S2N_EARLY_DATA_STATUS_")
            // match everything else
            .trim_start_matches("S2N_");

        Some(variant_name.to_owned())
    }

    /// This doesn't actually rename anything, and is just used to get a list of
    /// all the functions that we generate bindings for, which is used for test
    /// generation later in the process.
    fn generated_name_override(
        &self,
        item_info: bindgen::callbacks::ItemInfo<'_>,
    ) -> Option<String> {
        if !item_info.name.starts_with("s2n_") {
            return None;
        }

        if let ItemKind::Function = item_info.kind {
            let feature = self.feature.lock().unwrap().clone();
            self.functions
                .lock()
                .unwrap()
                .insert((feature, item_info.name.to_owned()));
        }
        None
    }
}


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

use std::path::PathBuf;

fn main() {
    // ensure the build script exports the include directory
    let include_dir = std::env::var("DEP_S2N_TLS_INCLUDE").expect("missing DEP_S2N_TLS_INCLUDE");
    let include_dir = PathBuf::from(include_dir);

    // make sure that `s2n.h` is available
    let api = std::fs::read_to_string(include_dir.join("s2n.h")).unwrap();
    assert!(api.contains("s2n_negotiate"));
}


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

use criterion::{criterion_group, criterion_main, Criterion};
use s2n_tls::{
    security,
    testing::{build_config, establish_connection},
};

pub fn handshake(c: &mut Criterion) {
    let mut group = c.benchmark_group("s2n-tls_client_server");

    for policy in security::ALL_POLICIES {
        let config = build_config(policy).unwrap();
        group.bench_function(format!("handshake_{:?}", policy), move |b| {
            // This does include connection initialization overhead.
            // TODO: create a separate benchmark that excludes this step.
            b.iter(|| establish_connection(config.clone()));
        });
    }

    group.finish();
}

criterion_group!(benches, handshake);
criterion_main!(benches);


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

use criterion::{criterion_group, criterion_main, Criterion};
use std::{env, process::Command, time::Duration};

pub fn s2nc(c: &mut Criterion) {
    let mut group = c.benchmark_group("s2nc");
    let s2nc_env: &str = &env::var("S2NC_ARGS").unwrap();
    let s2nc_test_name: &str = &env::var("S2NC_TEST_NAME").unwrap();
    let test_name = format!("s2nc_{}", s2nc_test_name);
    let s2nc_split = s2nc_env.split(' ').collect::<Vec<&str>>();
    group.bench_function(test_name, move |b| {
        b.iter(|| {
            let s2nc_argvec = s2nc_split.clone();
            let status = Command::new("/usr/local/bin/s2nc")
                .args(s2nc_argvec)
                .status()
                .expect("failed to execute process");
            assert!(status.success());
        });
    });

    group.finish();
}

criterion_group!(name = benches;
                 config = Criterion::default().sample_size(10).measurement_time(Duration::from_secs(1));
                 targets = s2nc);
criterion_main!(benches);


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

use criterion::{criterion_group, criterion_main, Criterion};
use std::{env, process::Command, time::Duration};

pub fn s2nd(c: &mut Criterion) {
    let mut group = c.benchmark_group("s2nd");
    let s2nd_env: &str = &env::var("S2ND_ARGS").unwrap();
    let s2nd_test_name: &str = &env::var("S2ND_TEST_NAME").unwrap();
    let test_name = format!("s2nd_{}", s2nd_test_name);
    let s2nd_split = s2nd_env.split(' ').collect::<Vec<&str>>();
    group.bench_function(test_name, move |b| {
        b.iter(|| {
            let s2nd_argvec = s2nd_split.clone();
            let status = Command::new("/usr/local/bin/s2nd")
                .args(s2nd_argvec)
                .status()
                .expect("failed to execute process");
            assert!(status.success());
        });
    });

    group.finish();
}

criterion_group!(name = benches;
                 config = Criterion::default().sample_size(10).measurement_time(Duration::from_secs(1));
                 targets = s2nd);
criterion_main!(benches);


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

use s2n_tls_sys::*;

fn main() {
    unsafe {
        s2n_init();
        let conn = s2n_connection_new(s2n_mode::SERVER);

        if !conn.is_null() {
            s2n_connection_free(conn);
        }
    }
}


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

//! Utilities to handle passing Rust code to s2n-tls's C callbacks.
//!
//! s2n-tls uses callbacks to temporarily return control to the application
//! and allow the application to execute custom code.
//!
//! To use a callback in your application, just implement the trait for the
//! target callback type and pass your implementation to the appropriate
//! connection or config method. For example, you can implement
//! [`ClientHelloCallback`] and pass that implementation to
//! [config::Builder::set_client_hello_callback()](`crate::config::Builder::set_client_hello_callback()`)
//! in order to execute custom logic after an s2n-tls server receives a client hello.
//!
//! s2n-tls callbacks come in two flavors:
//! * "sync" callbacks return an immediate result and will block the task
//!   performing the handshake until they return success or failure. See
//!   [`VerifyHostNameCallback`] as an example.
//! * "async" callbacks return a [Poll](`core::task::Poll`) and should not block the task performing the handshake.
//!   They will be polled until they return [Poll::Ready](`core::task::Poll::Ready`).
//!   [Connection::waker()](`crate::connection::Connection::waker()`)
//!   can be used to register the task for wakeup. See [`ClientHelloCallback`] as an example.

use crate::{config::Context, connection::Connection};
use core::{mem::ManuallyDrop, ptr::NonNull, time::Duration};
use s2n_tls_sys::s2n_connection;

mod async_cb;
pub use async_cb::*;

mod client_hello;
pub use client_hello::*;

mod pkey;
pub use pkey::*;

/// Convert the connection pointer provided to a callback into a Connection
/// and Context useable with the Rust bindings.
///
/// # Safety
///
/// This must ONLY be used for connection pointers provided to callbacks,
/// which can be assumed to point to valid Connections because the
/// callbacks were configured through the Rust bindings.
pub(crate) unsafe fn with_context<F, T>(conn_ptr: *mut s2n_connection, action: F) -> T
where
    F: FnOnce(&mut Connection, &mut Context) -> T,
{
    let raw = NonNull::new(conn_ptr).expect("connection should not be null");
    let mut conn = Connection::from_raw(raw);
    let mut config = conn.config().expect("config should not be null");
    let context = config.context_mut();
    let r = action(&mut conn, context);
    // Since this is a callback, it receives a pointer to the connection
    // but doesn't own that connection or control its lifecycle.
    // Do not drop / free the connection.
    let _ = ManuallyDrop::new(conn);
    r
}

/// A trait for the callback used to verify host name(s) during X509
/// verification.
///
/// The implementation should verify the certificate host name and return `true`
/// if the name is valid, `false` otherwise.
pub trait VerifyHostNameCallback {
    fn verify_host_name(&self, host_name: &str) -> bool;
}

/// A trait for the callback used to retrieve the system / wall clock time.
pub trait WallClock {
    fn get_time_since_epoch(&self) -> Duration;
}

/// A trait for the callback used to retrieve the monotonic time.
pub trait MonotonicClock {
    fn get_time(&self) -> Duration;
}

/// Invoke the user provided VerifyHostNameCallback on the host_name.
///
/// # Safety
///
/// The caller must ensure that the memory underlying host_name is a valid
/// slice.
pub(crate) unsafe fn verify_host(
    host_name: *const ::libc::c_char,
    host_name_len: usize,
    handler: &mut Box<dyn VerifyHostNameCallback>,
) -> u8 {
    let host_name = host_name as *const u8;
    let host_name = core::slice::from_raw_parts(host_name, host_name_len);

    match core::str::from_utf8(host_name) {
        Ok(host_name_str) => handler.verify_host_name(host_name_str) as u8,
        Err(_) => 0, // If the host name can't be parsed, fail closed.
    }
}


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

use crate::error::{Error, Fallible};
use s2n_tls_sys::*;
use std::fmt;

#[derive(Copy, Clone)]
pub enum FingerprintType {
    JA3,
}

// this is the size of the MD5 hash digest that is used for the JA3 fingerprint
const MD5_HASH_SIZE: u32 = 16;

impl From<FingerprintType> for s2n_tls_sys::s2n_fingerprint_type::Type {
    fn from(value: FingerprintType) -> Self {
        match value {
            FingerprintType::JA3 => s2n_tls_sys::s2n_fingerprint_type::FINGERPRINT_JA3,
        }
    }
}

// ClientHello is an opaque wrapper struct around `s2n_client_hello`. Note that
// the size of this type is not known, and as such it can only be used through
// references and pointers.
//
// This implementation is motivated by the different memory management required
// for different s2n_client_hello pointers. `s2n_client_hello_parse_message`
// returns a `*mut s2n_client_hello` which owns its own memory. This neatly fits
//  the "smart pointer" pattern and can be represented as a `Box<T>`.
//
// `s2n_connection_get_client_hello` returns a `*mut s2n_client_hello` which
// references memory owned by the connection, and therefore must not outlive
// the connection struct. This is best represented as a reference tied to the
// lifetime of the `Connection` struct.

/// ```no_run
/// use s2n_tls::client_hello::{ClientHello, FingerprintType};
/// use s2n_tls::connection::Connection;
/// use s2n_tls::enums::Mode;
///
/// let mut conn = Connection::new(Mode::Server);
/// // handshake happens
/// let mut client_hello: &ClientHello = conn.client_hello().unwrap();
/// let mut hash = Vec::new();
/// let string_size = client_hello.fingerprint_hash(FingerprintType::JA3, &mut hash).unwrap();
/// // hash has been resized so that it can store the fingerprint hash
///
/// let mut string = String::with_capacity(string_size as usize);
/// // string will not be resized, and the method will fail with
/// // ErrorType::UsageError if the string doesn't have enough capacity
/// client_hello.fingerprint_string(FingerprintType::JA3, &mut string).unwrap();
/// ```
pub struct ClientHello(s2n_client_hello);

impl ClientHello {
    pub fn parse_client_hello(hello: &[u8]) -> Result<Box<Self>, crate::error::Error> {
        crate::init::init();
        let handle = unsafe {
            s2n_client_hello_parse_message(hello.as_ptr(), hello.len() as u32).into_result()?
        };
        let client_hello = handle.as_ptr() as *mut ClientHello;
        // safety: s2n_client_hello_parse_message returns a pointer that "owns"
        // its memory. This memory must be cleaned up by the application. The
        // Box<Self> will call Self::Drop when it goes out of scope so memory
        // will be automatically managed.
        unsafe { Ok(Box::from_raw(client_hello)) }
    }

    // this accepts a mut ref instead of a pointer, so that lifetimes are nicely
    // calculated for us. As is always the case, the reference must not be null.
    // this is marked "pub(crate)" to expose it to the connection module but
    // prevent it from being used externally.
    pub(crate) fn from_ptr(hello: &s2n_client_hello) -> &Self {
        // SAFETY: casting *s2n_client_hello <-> *ClientHello: For repr(Rust),
        // repr(packed(N)), repr(align(N)), and repr(C) structs: if all fields of a
        // struct have size 0, then the struct has size 0.
        // https://rust-lang.github.io/unsafe-code-guidelines/layout/structs-and-tuples.html#zero-sized-structs
        unsafe { &*(hello as *const s2n_client_hello as *const ClientHello) }
    }

    // SAFETY: casting *const s2n_client_hello -> *mut s2n_client_hello: This is
    // safe as long as the data is not actually mutated. As authors of s2n-tls,
    // we know that the get_hash and get_fingerprint methods do not mutate the
    // data, and use mut pointers as a matter of convention because it makes
    // working with s2n_stuffers and s2n_blobs easier.
    fn deref_mut_ptr(&self) -> *mut s2n_client_hello {
        &self.0 as *const s2n_client_hello as *mut s2n_client_hello
    }

    /// `fingerprint_hash` calculates the hash, and also returns the size
    /// required for the full fingerprint string. The return value can be used
    /// to construct a string of appropriate capacity to call
    /// `fingerprint_string`. `output` will be extended if necessary to store
    /// the full hash.
    pub fn fingerprint_hash(
        &self,
        hash: FingerprintType,
        output: &mut Vec<u8>,
    ) -> Result<u32, Error> {
        let mut hash_size: u32 = 0;
        let mut str_size: u32 = 0;
        // make sure the vec has sufficient space for the hash
        if output.capacity() < MD5_HASH_SIZE as usize {
            output.reserve_exact(MD5_HASH_SIZE as usize - output.len());
        }
        unsafe {
            s2n_client_hello_get_fingerprint_hash(
                self.deref_mut_ptr(),
                hash.into(),
                MD5_HASH_SIZE,
                output.as_mut_ptr(),
                &mut hash_size,
                &mut str_size,
            )
            .into_result()?;
            // SAFETY: we wrote to the raw vec (using the mut pointer), and need
            // to update the state of the vec to reflect the changes we made.
            output.set_len(hash_size as usize);
        };
        Ok(str_size)
    }

    /// `fingerprint_string` will try to calculate the fingerprint and store the
    /// resulting string in `output`. If `output` does not have sufficient
    /// capacity an Error of `ErrorType::UsageError` will be returned.
    pub fn fingerprint_string(
        &self,
        hash: FingerprintType,
        output: &mut String,
    ) -> Result<(), Error> {
        let mut output_size = 0;
        unsafe {
            s2n_tls_sys::s2n_client_hello_get_fingerprint_string(
                self.deref_mut_ptr(),
                hash.into(),
                output.capacity() as u32,
                output.as_mut_ptr(),
                &mut output_size,
            )
            .into_result()?;
            // SAFETY: update internal state of string to match the data written
            // into it.
            output.as_mut_vec().set_len(output_size as usize);
        };
        Ok(())
    }

    fn session_id(&self) -> Result<Vec<u8>, Error> {
        let mut session_id_length = 0;
        unsafe {
            s2n_client_hello_get_session_id_length(self.deref_mut_ptr(), &mut session_id_length)
                .into_result()?;
        }

        let mut session_id = vec![0; session_id_length as usize];
        let mut out_length = 0;
        unsafe {
            s2n_client_hello_get_session_id(
                self.deref_mut_ptr(),
                session_id.as_mut_ptr(),
                &mut out_length,
                session_id_length,
            )
            .into_result()?;
        }
        Ok(session_id)
    }

    fn raw_message(&self) -> Result<Vec<u8>, Error> {
        let message_length =
            unsafe { s2n_client_hello_get_raw_message_length(self.deref_mut_ptr()).into_result()? };

        let mut raw_message = vec![0; message_length];
        unsafe {
            s2n_client_hello_get_raw_message(
                self.deref_mut_ptr(),
                raw_message.as_mut_ptr(),
                message_length as u32,
            )
            .into_result()?
        };
        Ok(raw_message)
    }
}

impl Drop for ClientHello {
    fn drop(&mut self) {
        let mut client_hello: *mut s2n_client_hello = &mut self.0;
        // ignore failures. There isn't anything to be done to handle them, but
        // allowing the program to continue is preferable to crashing.
        let _ = unsafe {
            s2n_tls_sys::s2n_client_hello_free(std::ptr::addr_of_mut!(client_hello)).into_result()
        };
    }
}

impl fmt::Debug for ClientHello {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        let session_id = self.session_id().map_err(|_| fmt::Error)?;
        let session_id = hex::encode(session_id);
        let message_head = self.raw_message().map_err(|_| fmt::Error)?;
        let mut hash = Vec::new();
        self.fingerprint_hash(FingerprintType::JA3, &mut hash)
            .map_err(|_| fmt::Error)?;
        f.debug_struct("ClientHello")
            .field("session_id", &session_id)
            .field("message_len", &(message_head.len()))
            .field("ja3_fingerprint", &hex::encode(hash))
            .finish_non_exhaustive()
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::{
        connection::Connection,
        error::{Error, ErrorType},
        security,
        testing::{poll_tls_pair, tls_pair},
    };

    /// This function is a test fixture used a generate a valid ClientHello so
    /// that we don't have to copy and paste the raw bytes for test fixtures
    fn get_client_hello_bytes() -> Vec<u8> {
        let config = crate::testing::config_builder(&security::DEFAULT_TLS13)
            .unwrap()
            .build()
            .unwrap();
        let pair = tls_pair(config);
        let pair = poll_tls_pair(pair);
        // this doesn't have the handshake header
        let client_hello_message = pair
            .server
            .0
            .connection()
            .client_hello()
            .unwrap()
            .raw_message()
            .unwrap();
        // handshake header is {tag: u8, client_hello_length: u24}
        let mut client_hello = vec![0; 4];
        // As long as the client hello is small, no bit fiddling is required
        assert!(client_hello_message.len() < u8::MAX as usize);
        // tag for handshake header
        client_hello[0] = 1;
        client_hello[3] = client_hello_message.len() as u8;
        client_hello.extend(client_hello_message.iter());
        client_hello
    }

    fn get_client_hello() -> Box<ClientHello> {
        // sets up connection and handshakes
        let raw_client_hello = get_client_hello_bytes();
        ClientHello::parse_client_hello(raw_client_hello.as_slice()).unwrap()
    }

    // test that a fingerprint can successfully be calculated from ClientHellos
    // returned from a connection
    #[checkers::test]
    fn io_fingerprint_test() {
        let config = crate::testing::config_builder(&security::DEFAULT_TLS13)
            .unwrap()
            .build()
            .unwrap();
        let pair = crate::testing::tls_pair(config);

        // client_hellos can not be accessed before the handshake
        assert!(pair.client.0.connection().client_hello().is_err());
        assert!(pair.server.0.connection().client_hello().is_err());

        let pair = poll_tls_pair(pair);
        let server_conn = pair.server.0.connection();
        let client_conn = pair.server.0.connection();

        let check_client_hello = |conn: &Connection| -> Result<(), Error> {
            let client_hello = conn.client_hello().unwrap();
            let mut hash = Vec::new();
            let fingerprint_size =
                client_hello.fingerprint_hash(FingerprintType::JA3, &mut hash)?;
            let mut string = String::with_capacity(fingerprint_size as usize);
            client_hello.fingerprint_string(FingerprintType::JA3, &mut string)?;
            Ok(())
        };

        assert!(check_client_hello(server_conn).is_ok());
        assert!(check_client_hello(client_conn).is_ok());
    }

    fn known_test_case(
        raw_client_hello: Vec<u8>,
        expected_string: &str,
        expected_hash_hex: &str,
    ) -> Result<(), Error> {
        let expected_hash: Vec<u8> = hex::decode(expected_hash_hex).unwrap();
        let client_hello = ClientHello::parse_client_hello(raw_client_hello.as_slice()).unwrap();

        let mut hash = Vec::new();
        let string_size = client_hello
            .fingerprint_hash(FingerprintType::JA3, &mut hash)
            .unwrap();
        assert_eq!(hash, expected_hash);

        let mut string = String::with_capacity(string_size as usize);
        client_hello
            .fingerprint_string(FingerprintType::JA3, &mut string)
            .unwrap();
        assert_eq!(string, expected_string);
        Ok(())
    }

    #[test]
    fn invalid_client_bytes() {
        let raw_client_hello_bytes =
            "random_value_that_is_unlikely_to_be_valid_client_hello".as_bytes();
        let result = ClientHello::parse_client_hello(raw_client_hello_bytes);
        assert!(result.is_err());
    }

    // known value test case copied from s2n_fingerprint_ja3_test.c
    #[checkers::test]
    fn valid_client_bytes() {
        let raw_client_hello = vec![
            0x01, 0x00, 0x00, 0xEC, 0x03, 0x03, 0x90, 0xe8, 0xcc, 0xee, 0xe5, 0x70, 0xa2, 0xa1,
            0x2f, 0x6b, 0x69, 0xd2, 0x66, 0x96, 0x0f, 0xcf, 0x20, 0xd5, 0x32, 0x6e, 0xc4, 0xb2,
            0x8c, 0xc7, 0xbd, 0x0a, 0x06, 0xc2, 0xa5, 0x14, 0xfc, 0x34, 0x20, 0xaf, 0x72, 0xbf,
            0x39, 0x99, 0xfb, 0x20, 0x70, 0xc3, 0x10, 0x83, 0x0c, 0xee, 0xfb, 0xfa, 0x72, 0xcc,
            0x5d, 0xa8, 0x99, 0xb4, 0xc5, 0x53, 0xd6, 0x3d, 0xa0, 0x53, 0x7a, 0x5c, 0xbc, 0xf5,
            0x0b, 0x00, 0x1e, 0xc0, 0x2b, 0xc0, 0x2f, 0xcc, 0xa9, 0xcc, 0xa8, 0xc0, 0x2c, 0xc0,
            0x30, 0xc0, 0x0a, 0xc0, 0x09, 0xc0, 0x13, 0xc0, 0x14, 0x00, 0x33, 0x00, 0x39, 0x00,
            0x2f, 0x00, 0x35, 0x00, 0x0a, 0x01, 0x00, 0x00, 0x85, 0x00, 0x00, 0x00, 0x23, 0x00,
            0x21, 0x00, 0x00, 0x1e, 0x69, 0x6e, 0x63, 0x6f, 0x6d, 0x69, 0x6e, 0x67, 0x2e, 0x74,
            0x65, 0x6c, 0x65, 0x6d, 0x65, 0x74, 0x72, 0x79, 0x2e, 0x6d, 0x6f, 0x7a, 0x69, 0x6c,
            0x6c, 0x61, 0x2e, 0x6f, 0x72, 0x67, 0x00, 0x17, 0x00, 0x00, 0xff, 0x01, 0x00, 0x01,
            0x00, 0x00, 0x0a, 0x00, 0x0a, 0x00, 0x08, 0x00, 0x1d, 0x00, 0x17, 0x00, 0x18, 0x00,
            0x19, 0x00, 0x0b, 0x00, 0x02, 0x01, 0x00, 0x00, 0x23, 0x00, 0x00, 0x00, 0x10, 0x00,
            0x0e, 0x00, 0x0c, 0x02, 0x68, 0x32, 0x08, 0x68, 0x74, 0x74, 0x70, 0x2f, 0x31, 0x2e,
            0x31, 0x00, 0x05, 0x00, 0x05, 0x01, 0x00, 0x00, 0x00, 0x00, 0x00, 0x0d, 0x00, 0x18,
            0x00, 0x16, 0x04, 0x03, 0x05, 0x03, 0x06, 0x03, 0x08, 0x04, 0x08, 0x05, 0x08, 0x06,
            0x04, 0x01, 0x05, 0x01, 0x06, 0x01, 0x02, 0x03, 0x02, 0x01, 0x00, 0x1c, 0x00, 0x02,
            0x40, 0x00,
        ];
        let expected_fingerprint = "771,49195-49199-52393-52392-49196-49200-\
                                    49162-49161-49171-49172-51-57-47-53-10,0-\
                                    23-65281-10-11-35-16-5-13-28,29-23-24-25,0";
        let expected_hash_hex = "839bbe3ed07fed922ded5aaf714d6842";
        known_test_case(raw_client_hello, expected_fingerprint, expected_hash_hex).unwrap();
    }

    #[test]
    fn hash_output_resizing() {
        let client_hello = get_client_hello();
        let hash_capacities = vec![0, MD5_HASH_SIZE, 1_000];
        for initial_size in hash_capacities {
            let mut hash = Vec::with_capacity(initial_size as usize);
            client_hello
                .fingerprint_hash(FingerprintType::JA3, &mut hash)
                .unwrap();
            assert_eq!(hash.len(), MD5_HASH_SIZE as usize);
        }
    }

    #[test]
    fn string_output_too_small() {
        let client_hello = get_client_hello();
        let mut fingerprint_string = String::with_capacity(0);
        let fingerprint_err = client_hello
            .fingerprint_string(FingerprintType::JA3, &mut fingerprint_string)
            .unwrap_err();
        assert_eq!(fingerprint_err.kind(), ErrorType::UsageError);
    }

    // make sure that debug doesn't panic and seems reasonable
    #[test]
    fn debug() {
        let client_hello = get_client_hello();
        let mut hash = Vec::new();
        client_hello
            .fingerprint_hash(FingerprintType::JA3, &mut hash)
            .unwrap();
        let client_hello_debug = format!("{:?}", client_hello);
        assert!(client_hello_debug.contains(&hex::encode(hash)));
    }
}