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

use crate::{
    event,
    inet::datagram,
    io::{rx, tx},
    path::Tuple,
};
use core::task::{Context, Poll};
use std::sync::{Arc, Mutex};

pub type Handle = Tuple;

#[derive(Clone, Debug, Default)]
pub struct Channel {
    messages: Arc<Mutex<Vec<Message>>>,
}

impl Channel {
    pub fn push(&self, message: Message) {
        self.messages.lock().unwrap().push(message);
    }

    pub fn pop(&self) -> Option<Message> {
        self.messages.lock().unwrap().pop()
    }

    #[inline]
    fn queue<F: FnOnce(&mut Queue<'static>)>(&mut self, f: F) {
        if let Ok(mut messages) = self.messages.lock() {
            let messages: &mut Vec<_> = &mut *messages;
            let messages: &'static mut _ = 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(messages)
            };

            let mut queue = Queue { messages };
            f(&mut queue);
        }
    }
}

#[derive(Clone, Debug)]
pub struct Message {
    pub header: datagram::Header<Tuple>,
    pub payload: Vec<u8>,
}

impl Default for Message {
    fn default() -> Self {
        Self {
            header: datagram::Header {
                ecn: Default::default(),
                path: Tuple {
                    local_address: Default::default(),
                    remote_address: Default::default(),
                },
            },
            payload: Default::default(),
        }
    }
}

impl tx::Tx for Channel {
    type PathHandle = Tuple;
    type Queue = Queue<'static>;
    type Error = ();

    #[inline]
    fn poll_ready(&mut self, _cx: &mut Context) -> Poll<Result<(), Self::Error>> {
        // the TX channel is always ready
        Poll::Pending
    }

    #[inline]
    fn queue<F: FnOnce(&mut Self::Queue)>(&mut self, f: F) {
        Self::queue(self, f)
    }

    #[inline]
    fn handle_error<E: event::EndpointPublisher>(self, _error: Self::Error, _event: &mut E) {
        // nothing to do
    }
}

impl rx::Rx for Channel {
    type PathHandle = Tuple;
    type Queue = Queue<'static>;
    type Error = ();

    #[inline]
    fn poll_ready(&mut self, _cx: &mut Context) -> Poll<Result<(), Self::Error>> {
        let messages = self.messages.lock().map_err(|_| ())?;
        if messages.is_empty() {
            Poll::Pending
        } else {
            Poll::Ready(Ok(()))
        }
    }

    #[inline]
    fn queue<F: FnOnce(&mut Self::Queue)>(&mut self, f: F) {
        Self::queue(self, f)
    }

    #[inline]
    fn handle_error<E: event::EndpointPublisher>(self, _error: Self::Error, _event: &mut E) {
        // nothing to do
    }
}

pub struct Queue<'a> {
    messages: &'a mut Vec<Message>,
}

impl<'a> tx::Queue for Queue<'a> {
    type Handle = Tuple;

    #[inline]
    fn push<M: tx::Message<Handle = Self::Handle>>(
        &mut self,
        mut message: M,
    ) -> Result<tx::Outcome, tx::Error> {
        let mut out = Message::default();
        out.header.ecn = message.ecn();
        out.header.path = *message.path_handle();

        out.payload.resize(1500, 0);
        let buffer = tx::PayloadBuffer::new(&mut out.payload);
        let len = message.write_payload(buffer, 0)?;

        self.messages.push(out);

        let outcome = tx::Outcome { index: 0, len };

        Ok(outcome)
    }

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

impl<'a> rx::Queue for Queue<'a> {
    type Handle = Tuple;

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

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


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

use crate::{event, inet::ExplicitCongestionNotification, path};
use core::{
    task::{Context, Poll},
    time::Duration,
};

pub mod handle_map;
pub mod router;

pub trait Tx: Sized {
    type PathHandle;
    // TODO make this generic over lifetime
    // See https://github.com/aws/s2n-quic/issues/1742
    type Queue: Queue<Handle = Self::PathHandle>;
    type Error;

    /// Returns a future that yields after a packet is ready to be transmitted
    #[inline]
    fn ready(&mut self) -> TxReady<Self> {
        TxReady(self)
    }

    /// Polls the IO provider for capacity to send a packet
    fn poll_ready(&mut self, cx: &mut Context) -> Poll<Result<(), Self::Error>>;

    /// Calls the provided callback with the IO provider queue
    fn queue<F: FnOnce(&mut Self::Queue)>(&mut self, f: F);

    /// Handles the queue error and potentially publishes an event
    fn handle_error<E: event::EndpointPublisher>(self, error: Self::Error, event: &mut E);
}

impl_ready_future!(Tx, TxReady, Result<(), T::Error>);

/// Extension traits for Tx channels
pub trait TxExt: Tx {
    /// Routes messages into one channel or another
    #[inline]
    fn with_router<Router, Other>(
        self,
        router: Router,
        other: Other,
    ) -> router::Channel<Router, Self, Other>
    where
        Router: router::Router,
        Other: Tx,
    {
        router::Channel {
            router,
            a: self,
            b: other,
        }
    }

    /// Maps one type of handle to another with a mapping function
    #[inline]
    fn with_handle_map<Map, Handle>(self, map: Map) -> handle_map::Channel<Map, Self, Handle>
    where
        Map: Fn(&Handle) -> Self::PathHandle,
    {
        handle_map::Channel {
            map,
            tx: self,
            handle: Default::default(),
        }
    }
}

/// Implement the extension traits for all Tx queues
impl<T: Tx> TxExt for T {}

/// A structure capable of queueing and transmitting messages
pub trait Queue {
    type Handle: path::Handle;

    /// Set to true if the queue supports setting ECN markings
    const SUPPORTS_ECN: bool = false;

    /// Set to true if the queue supports pacing of sending messages
    const SUPPORTS_PACING: bool = false;

    /// Set to true if the queue supports setting IPv6 flow labels
    const SUPPORTS_FLOW_LABELS: bool = false;

    /// Pushes a message into the transmission queue
    ///
    /// The index of the message is returned to enable further operations to be
    /// performed, e.g. encryption.
    fn push<M: Message<Handle = Self::Handle>>(&mut self, message: M) -> Result<Outcome, Error>;

    /// Flushes any pending messages from the TX queue.
    ///
    /// This should be called between multiple connections to ensure GSO segments aren't shared.
    #[inline]
    fn flush(&mut self) {
        // default as no-op
    }

    /// Returns the number of remaining datagrams that can be transmitted
    fn capacity(&self) -> usize;

    /// Returns `true` if the queue will accept additional transmissions
    #[inline]
    fn has_capacity(&self) -> bool {
        self.capacity() != 0
    }
}

pub struct Outcome {
    pub len: usize,
    pub index: usize,
}

#[derive(Copy, Clone, Debug, PartialEq, PartialOrd, Eq)]
pub enum Error {
    /// The provided message did not write a payload
    EmptyPayload,

    /// The provided buffer was too small for the desired payload
    UndersizedBuffer,

    /// The transmission queue is at capacity
    AtCapacity,
}

/// Abstraction over a message to be sent on a socket
///
/// Instead of a concrete struct with eagerly evaluated fields,
/// using trait callbacks ensure messages only need to compute what
/// the actual transmission queue requires. For example, if the transmission
/// queue cannot set ECN markings, it will not call the [`Message::ecn`] function.
pub trait Message {
    type Handle: path::Handle;

    /// Returns the path handle on which this message should be sent
    fn path_handle(&self) -> &Self::Handle;

    /// Returns the ECN markings for the message
    fn ecn(&mut self) -> ExplicitCongestionNotification;

    /// Returns the Duration for which the message will be delayed.
    ///
    /// This is used in scenarios where packets need to be paced.
    fn delay(&mut self) -> Duration;

    /// Returns the IPv6 flow label for the message
    fn ipv6_flow_label(&mut self) -> u32;

    /// Returns true if the packet can be used in a GSO packet
    fn can_gso(&self, segment_len: usize, segment_count: usize) -> bool;

    /// Writes the payload of the message to an output buffer
    fn write_payload(&mut self, buffer: PayloadBuffer, gso_offset: usize) -> Result<usize, Error>;
}

#[derive(Debug)]
pub struct PayloadBuffer<'a>(&'a mut [u8]);

impl<'a> PayloadBuffer<'a> {
    #[inline]
    pub fn new(bytes: &'a mut [u8]) -> Self {
        Self(bytes)
    }

    /// # Safety
    ///
    /// This function should only be used in the case that the writer has its own safety checks in place
    #[inline]
    pub unsafe fn into_mut_slice(self) -> &'a mut [u8] {
        self.0
    }

    #[track_caller]
    #[inline]
    pub fn write(&mut self, bytes: &[u8]) -> Result<usize, Error> {
        if bytes.is_empty() {
            return Err(Error::EmptyPayload);
        }

        if let Some(buffer) = self.0.get_mut(0..bytes.len()) {
            buffer.copy_from_slice(bytes);
            Ok(bytes.len())
        } else {
            debug_assert!(
                false,
                "tried to write more bytes than was available in the buffer"
            );
            Err(Error::UndersizedBuffer)
        }
    }
}

impl<Handle: path::Handle, Payload: AsRef<[u8]>> Message for (Handle, Payload) {
    type Handle = Handle;

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

    fn ecn(&mut self) -> ExplicitCongestionNotification {
        Default::default()
    }

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

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

    fn can_gso(&self, segment_len: usize, _segment_count: usize) -> bool {
        segment_len >= self.1.as_ref().len()
    }

    fn write_payload(
        &mut self,
        mut buffer: PayloadBuffer,
        _gso_offset: usize,
    ) -> Result<usize, Error> {
        buffer.write(self.1.as_ref())
    }
}

impl<Handle: path::Handle, Payload: AsRef<[u8]>> Message
    for (Handle, ExplicitCongestionNotification, Payload)
{
    type Handle = Handle;

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

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

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

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

    fn can_gso(&self, segment_len: usize, _segment_count: usize) -> bool {
        segment_len >= self.2.as_ref().len()
    }

    fn write_payload(
        &mut self,
        mut buffer: PayloadBuffer,
        _gso_offset: usize,
    ) -> Result<usize, Error> {
        buffer.write(self.2.as_ref())
    }
}

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

    #[test]
    fn message_tuple_test() {
        let remote_address = SocketAddressV4::new([127, 0, 0, 1], 80).into();
        let local_address = SocketAddressV4::new([192, 168, 0, 1], 3000).into();
        let tuple = path::Tuple {
            remote_address,
            local_address,
        };
        let mut message = (tuple, [1u8, 2, 3]);

        let mut buffer = [0u8; 10];

        assert_eq!(*message.path_handle(), tuple);
        assert_eq!(message.ecn(), Default::default());
        assert_eq!(message.delay(), Default::default());
        assert_eq!(message.ipv6_flow_label(), 0);
        assert_eq!(
            message.write_payload(PayloadBuffer::new(&mut buffer[..]), 0),
            Ok(3)
        );
    }

    #[test]
    #[should_panic]
    fn message_tuple_undersized_test() {
        let remote_address = SocketAddressV4::new([127, 0, 0, 1], 80).into();
        let local_address = SocketAddressV4::new([192, 168, 0, 1], 3000).into();
        let tuple = path::Tuple {
            remote_address,
            local_address,
        };
        let mut message = (tuple, [1u8, 2, 3]);

        // assert an undersized buffer panics in debug
        let _ = message.write_payload(PayloadBuffer::new(&mut [][..]), 0);
    }
}


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

use crate::endpoint::CloseError;
use core::{
    future::Future,
    pin::Pin,
    task::{Context, Poll},
};
use pin_project_lite::pin_project;

pin_project!(
    /// The main event loop future for selecting readiness of sub-tasks
    ///
    /// This future ensures all sub-tasks are polled fairly by yielding once
    /// after completing any of the sub-tasks. This is especially important when the TX queue is
    /// flushed quickly and we never get notified of the RX socket having packets to read.
    pub struct Select<Rx, Tx, Wakeup, Sleep>
    where
        Rx: Future,
        Tx: Future,
        Wakeup: Future,
        Sleep: Future,
    {
        #[pin]
        rx: Rx,
        #[pin]
        tx: Tx,
        #[pin]
        wakeup: Wakeup,
        #[pin]
        sleep: Sleep,
    }
);

impl<Rx, Tx, Wakeup, Sleep> Select<Rx, Tx, Wakeup, Sleep>
where
    Rx: Future,
    Tx: Future,
    Wakeup: Future,
    Sleep: Future,
{
    #[inline(always)]
    pub fn new(rx: Rx, tx: Tx, wakeup: Wakeup, sleep: Sleep) -> Self {
        Self {
            rx,
            tx,
            wakeup,
            sleep,
        }
    }
}

#[derive(Debug)]
pub struct Outcome<Rx, Tx> {
    pub rx_result: Option<Rx>,
    pub tx_result: Option<Tx>,
    pub timeout_expired: bool,
    pub application_wakeup: bool,
}

pub type Result<Rx, Tx> = core::result::Result<Outcome<Rx, Tx>, CloseError>;

impl<Rx, Tx, Wakeup, Sleep> Future for Select<Rx, Tx, Wakeup, Sleep>
where
    Rx: Future,
    Tx: Future,
    Wakeup: Future<Output = core::result::Result<usize, CloseError>>,
    Sleep: Future,
{
    type Output = Result<Rx::Output, Tx::Output>;

    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
        let this = self.project();

        let mut should_wake = false;
        let mut application_wakeup = false;

        if let Poll::Ready(wakeup) = this.wakeup.poll(cx) {
            should_wake = true;
            application_wakeup = true;
            if let Err(err) = wakeup {
                return Poll::Ready(Err(err));
            }
        }

        let mut rx_result = None;
        if let Poll::Ready(v) = this.rx.poll(cx) {
            should_wake = true;
            rx_result = Some(v);
        }

        let mut tx_result = None;
        if let Poll::Ready(v) = this.tx.poll(cx) {
            should_wake = true;
            tx_result = Some(v);
        }

        let mut timeout_expired = false;

        if this.sleep.poll(cx).is_ready() {
            timeout_expired = true;
            should_wake = true;
        }

        // if none of the subtasks are ready, return
        if !should_wake {
            return Poll::Pending;
        }

        Poll::Ready(Ok(Outcome {
            rx_result,
            tx_result,
            timeout_expired,
            application_wakeup,
        }))
    }
}


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

use super::Rx;
use crate::{event, inet::datagram};
use core::task::{Context, Poll};

/// A pair of Rx channels that feed into the same endpoint
pub struct Channel<A, B>
where
    A: Rx,
    B: Rx<PathHandle = A::PathHandle>,
{
    pub(super) a: A,
    pub(super) b: B,
}

impl<A, B> Rx for Channel<A, B>
where
    A: Rx,
    B: Rx<PathHandle = A::PathHandle>,
    A::Queue: 'static,
    B::Queue: 'static,
{
    type PathHandle = A::PathHandle;
    type Queue = Queue<'static, A::Queue, B::Queue>;
    type Error = Error<A::Error, B::Error>;

    #[inline]
    fn poll_ready(&mut self, cx: &mut Context) -> Poll<Result<(), Self::Error>> {
        // assume we aren't ready until one channel returns ready
        let mut is_ready = false;

        macro_rules! ready {
            ($value:expr, $var:ident) => {
                match $value {
                    Poll::Ready(Ok(())) => is_ready = true,
                    Poll::Ready(Err(err)) => {
                        // one of the channels returned an error so shut down both
                        return Err(Error::$var(err)).into();
                    }
                    Poll::Pending => {}
                }
            };
        }

        ready!(self.a.poll_ready(cx), A);
        ready!(self.b.poll_ready(cx), B);

        if is_ready {
            Poll::Ready(Ok(()))
        } else {
            Poll::Pending
        }
    }

    #[inline]
    fn queue<F: FnOnce(&mut Self::Queue)>(&mut self, f: F) {
        let a = &mut self.a;
        let b = &mut self.b;
        a.queue(|a| {
            b.queue(|b| {
                let (a, b): (&'static mut _, &'static mut _) = 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(a), core::mem::transmute(b))
                };

                let mut queue = Queue { a, b };
                f(&mut queue);
            });
        });
    }

    #[inline]
    fn handle_error<E: event::EndpointPublisher>(self, error: Self::Error, event: &mut E) {
        // dispatch the error to the appropriate channel
        match error {
            Error::A(error) => self.a.handle_error(error, event),
            Error::B(error) => self.b.handle_error(error, event),
        }
    }
}

/// Tagged error for a pair of channels
pub enum Error<A, B> {
    A(A),
    B(B),
}

pub struct Queue<'a, A, B>
where
    A: super::Queue,
    B: super::Queue,
{
    a: &'a mut A,
    b: &'a mut B,
}

impl<'a, A, B> super::Queue for Queue<'a, A, B>
where
    A: super::Queue,
    B: super::Queue<Handle = A::Handle>,
{
    type Handle = A::Handle;

    #[inline]
    fn for_each<F: FnMut(datagram::Header<Self::Handle>, &mut [u8])>(&mut self, mut on_packet: F) {
        // drain both of the channels
        self.a.for_each(&mut on_packet);
        self.b.for_each(&mut on_packet);
    }

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

#[cfg(test)]
mod tests {
    use super::*;
    use crate::io::{
        rx::{Queue as _, RxExt as _},
        testing,
    };
    use futures_test::task::noop_waker;

    #[test]
    fn pair_test() {
        let channel_a = testing::Channel::default();
        let channel_b = testing::Channel::default();
        let mut merged = channel_a.clone().with_pair(channel_b.clone());

        let waker = noop_waker();
        let mut cx = Context::from_waker(&waker);
        let cx = &mut cx;

        for push_a in [false, true] {
            for push_b in [false, true] {
                assert!(merged.poll_ready(cx).is_pending());

                let mut expected = 0;

                if push_a {
                    expected += 1;
                    channel_a.push(Default::default());
                }

                if push_b {
                    expected += 1;
                    channel_b.push(Default::default());
                }

                assert_eq!(merged.poll_ready(cx).is_ready(), push_a || push_b);

                let mut actual = 0;
                merged.queue(|queue| {
                    assert_eq!(queue.is_empty(), !(push_a || push_b));

                    queue.for_each(|_header, _payload| {
                        actual += 1;
                    });
                });

                assert_eq!(expected, actual);
            }
        }
    }
}


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

use crate::{event, inet::ExplicitCongestionNotification, io::tx, path};
use core::{
    marker::PhantomData,
    task::{Context, Poll},
    time::Duration,
};

pub struct Channel<Map, Tx, U> {
    pub(super) map: Map,
    pub(super) tx: Tx,
    pub(super) handle: PhantomData<U>,
}

impl<Map, Tx, U> tx::Tx for Channel<Map, Tx, U>
where
    Map: 'static + Fn(&U) -> Tx::PathHandle,
    Tx: tx::Tx,
    Tx::Queue: 'static,
    U: path::Handle,
{
    type PathHandle = U;
    type Queue = Queue<'static, Map, Tx::Queue, U>;
    type Error = Tx::Error;

    #[inline]
    fn poll_ready(&mut self, cx: &mut Context) -> Poll<Result<(), Self::Error>> {
        self.tx.poll_ready(cx)
    }

    #[inline]
    fn queue<F: FnOnce(&mut Self::Queue)>(&mut self, f: F) {
        let map = &mut self.map;
        let tx = &mut self.tx;
        tx.queue(|tx| {
            let (map, tx): (&'static mut _, &'static mut _) = 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(map), core::mem::transmute(tx))
            };

            let mut queue = Queue {
                map,
                tx,
                handle: PhantomData,
            };
            f(&mut queue);
        });
    }

    #[inline]
    fn handle_error<E: event::EndpointPublisher>(self, error: Self::Error, events: &mut E) {
        self.tx.handle_error(error, events)
    }
}

pub struct Queue<'a, Map, Tx, U>
where
    Map: Fn(&U) -> Tx::Handle,
    Tx: tx::Queue,
{
    map: &'a Map,
    tx: &'a mut Tx,
    handle: PhantomData<U>,
}

impl<'a, Map, Tx, U> tx::Queue for Queue<'a, Map, Tx, U>
where
    Map: Fn(&U) -> Tx::Handle,
    Tx: tx::Queue,
    U: path::Handle,
{
    type Handle = U;

    const SUPPORTS_ECN: bool = Tx::SUPPORTS_ECN;
    const SUPPORTS_PACING: bool = Tx::SUPPORTS_PACING;
    const SUPPORTS_FLOW_LABELS: bool = Tx::SUPPORTS_FLOW_LABELS;

    #[inline]
    fn push<M: tx::Message<Handle = Self::Handle>>(
        &mut self,
        inner: M,
    ) -> Result<tx::Outcome, tx::Error> {
        let handle = (self.map)(inner.path_handle());
        let message = Message { inner, handle };
        self.tx.push(message)
    }

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

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

pub struct Message<M, Handle> {
    inner: M,
    handle: Handle,
}

impl<M, Handle> tx::Message for Message<M, Handle>
where
    M: tx::Message,
    Handle: path::Handle,
{
    type Handle = Handle;

    #[inline]
    fn path_handle(&self) -> &Self::Handle {
        // use the mapped handle instead of the inner type
        &self.handle
    }

    #[inline]
    fn ecn(&mut self) -> ExplicitCongestionNotification {
        self.inner.ecn()
    }

    #[inline]
    fn delay(&mut self) -> Duration {
        self.inner.delay()
    }

    #[inline]
    fn ipv6_flow_label(&mut self) -> u32 {
        self.inner.ipv6_flow_label()
    }

    #[inline]
    fn can_gso(&self, segment_len: usize, segment_count: usize) -> bool {
        self.inner.can_gso(segment_len, segment_count)
    }

    #[inline]
    fn write_payload(
        &mut self,
        buffer: tx::PayloadBuffer,
        gso_offset: usize,
    ) -> Result<usize, tx::Error> {
        self.inner.write_payload(buffer, gso_offset)
    }
}

#[cfg(test)]
mod tests {
    use crate::{
        io::{
            testing,
            tx::{Queue as _, Tx as _, TxExt as _},
        },
        path::{Handle as _, RemoteAddress},
    };

    #[test]
    fn handle_map_test() {
        let channel = testing::Channel::default();
        let mut mapped = channel.clone().with_handle_map(|handle: &RemoteAddress| {
            let mut handle = testing::Handle::from_remote_address(*handle);
            handle.local_address.set_port(321);
            handle
        });

        mapped.queue(|queue| {
            let mut handle = RemoteAddress::default();
            handle.set_port(123);
            let msg = (handle, &[1, 2, 3][..]);
            queue.push(msg).unwrap();
        });

        let msg = channel.pop().unwrap();

        assert_eq!(msg.header.path.remote_address.port(), 123);
        assert_eq!(msg.header.path.local_address.port(), 321);
    }
}


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

use crate::{event, io::tx, path};
use core::task::{Context, Poll};

/// Defines how to route a message between two different channels
pub trait Router {
    type Handle: path::Handle;

    fn route<M, A, B>(
        &mut self,
        message: M,
        a: &mut A,
        b: &mut B,
    ) -> Result<tx::Outcome, tx::Error>
    where
        M: tx::Message<Handle = Self::Handle>,
        A: tx::Queue<Handle = Self::Handle>,
        B: tx::Queue<Handle = Self::Handle>;
}

pub struct Channel<R, A, B> {
    pub(super) router: R,
    pub(super) a: A,
    pub(super) b: B,
}

impl<R, A, B> tx::Tx for Channel<R, A, B>
where
    R: 'static + Router,
    A: tx::Tx<PathHandle = R::Handle>,
    B: tx::Tx<PathHandle = R::Handle>,
    A::Queue: 'static,
    B::Queue: 'static,
{
    type PathHandle = A::PathHandle;
    type Queue = Queue<'static, R, A::Queue, B::Queue>;
    type Error = Error<A::Error, B::Error>;

    #[inline]
    fn poll_ready(&mut self, cx: &mut Context) -> Poll<Result<(), Self::Error>> {
        // assume we aren't ready until one of the channels says it is
        let mut is_ready = false;

        macro_rules! ready {
            ($value:expr, $var:ident) => {
                match $value {
                    Poll::Ready(Ok(())) => is_ready = true,
                    Poll::Ready(Err(err)) => {
                        // one of the channels returned an error so shut down the task
                        return Err(Error::$var(err)).into();
                    }
                    Poll::Pending => {}
                }
            };
        }

        ready!(self.a.poll_ready(cx), A);
        ready!(self.b.poll_ready(cx), B);

        if is_ready {
            Poll::Ready(Ok(()))
        } else {
            Poll::Pending
        }
    }

    #[inline]
    fn queue<F: FnOnce(&mut Self::Queue)>(&mut self, f: F) {
        let router = &mut self.router;
        let a = &mut self.a;
        let b = &mut self.b;
        a.queue(|a| {
            b.queue(|b| {
                let (router, a, b): (&'static mut _, &'static mut _, &'static mut _) = 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(router),
                        core::mem::transmute(a),
                        core::mem::transmute(b),
                    )
                };

                let mut queue = Queue { router, a, b };
                f(&mut queue);
            });
        });
    }

    #[inline]
    fn handle_error<E: event::EndpointPublisher>(self, error: Self::Error, events: &mut E) {
        // dispatch the error into the sub-channels
        match error {
            Error::A(error) => self.a.handle_error(error, events),
            Error::B(error) => self.b.handle_error(error, events),
        }
    }
}

/// Tagged error for a pair of channels
pub enum Error<A, B> {
    A(A),
    B(B),
}

pub struct Queue<'a, R, A, B>
where
    R: Router,
    A: super::Queue<Handle = R::Handle>,
    B: super::Queue<Handle = R::Handle>,
{
    router: &'a mut R,
    a: &'a mut A,
    b: &'a mut B,
}

impl<'a, R, A, B> super::Queue for Queue<'a, R, A, B>
where
    R: Router,
    A: super::Queue<Handle = R::Handle>,
    B: super::Queue<Handle = R::Handle>,
{
    type Handle = R::Handle;

    const SUPPORTS_ECN: bool = A::SUPPORTS_ECN || B::SUPPORTS_ECN;
    const SUPPORTS_PACING: bool = A::SUPPORTS_PACING && B::SUPPORTS_PACING;
    const SUPPORTS_FLOW_LABELS: bool = A::SUPPORTS_FLOW_LABELS || B::SUPPORTS_FLOW_LABELS;

    #[inline]
    fn push<M: tx::Message<Handle = Self::Handle>>(
        &mut self,
        message: M,
    ) -> Result<tx::Outcome, tx::Error> {
        // route messages to the appropriate queue
        self.router.route(message, self.a, self.b)
    }

    #[inline]
    fn capacity(&self) -> usize {
        // take the minimum of the channel capacity, since we don't know where the next message
        // will go
        self.a.capacity().min(self.b.capacity())
    }

    #[inline]
    fn has_capacity(&self) -> bool {
        // we only have capacity if both channels do
        self.a.has_capacity() && self.b.has_capacity()
    }
}

#[cfg(test)]
mod tests {
    use crate::io::{
        testing,
        tx::{self, Queue as _, Tx as _, TxExt as _},
    };

    struct Router;

    impl super::Router for Router {
        type Handle = testing::Handle;

        fn route<M, A, B>(
            &mut self,
            message: M,
            a: &mut A,
            b: &mut B,
        ) -> Result<tx::Outcome, tx::Error>
        where
            M: tx::Message<Handle = Self::Handle>,
            A: tx::Queue<Handle = Self::Handle>,
            B: tx::Queue<Handle = Self::Handle>,
        {
            let handle = message.path_handle();
            if handle.remote_address.port() == 0 {
                dbg!("a", handle.remote_address);
                a.push(message)
            } else {
                dbg!("b", handle.remote_address);
                b.push(message)
            }
        }
    }

    #[test]
    fn router_test() {
        let channel_a = testing::Channel::default();
        let channel_b = testing::Channel::default();
        let mut merged = channel_a.clone().with_router(Router, channel_b.clone());

        for push_a in [false, true] {
            for push_b in [false, true] {
                dbg!((push_a, push_b));

                merged.queue(|queue| {
                    if push_a {
                        let handle = testing::Handle {
                            remote_address: Default::default(),
                            local_address: Default::default(),
                        };
                        let msg = (handle, &[1, 2, 3][..]);
                        queue.push(msg).unwrap();
                    }

                    if push_b {
                        let mut handle = testing::Handle {
                            remote_address: Default::default(),
                            local_address: Default::default(),
                        };
                        handle.remote_address.set_port(1);
                        let msg = (handle, &[1, 2, 3][..]);
                        queue.push(msg).unwrap();
                    }
                });

                assert_eq!(channel_a.pop().is_some(), push_a);
                assert_eq!(channel_b.pop().is_some(), push_b);
            }
        }
    }
}


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

use crate::{
    connection,
    connection::id::ConnectionInfo,
    crypto::ProtectedPayload,
    packet::{
        long::{
            validate_destination_connection_id_range, validate_source_connection_id_range,
            DestinationConnectionIdLen, SourceConnectionIdLen, Version,
        },
        number::ProtectedPacketNumber,
        Tag,
    },
    varint::VarInt,
};
use core::{convert::TryInto, mem::size_of};
use s2n_codec::{CheckedRange, DecoderBuffer, DecoderBufferMut, DecoderError, DecoderValue};

pub struct HeaderDecoder<'a> {
    initial_buffer_len: usize,
    peek: DecoderBuffer<'a>,
}

impl<'a> HeaderDecoder<'a> {
    pub fn new_long<'b>(buffer: &'a DecoderBufferMut<'b>) -> Self {
        let initial_buffer_len = buffer.len();
        let peek = buffer.peek();
        let peek = peek
            .skip(size_of::<Tag>() + size_of::<Version>())
            .expect("tag and version already verified");
        Self {
            initial_buffer_len,
            peek,
        }
    }

    pub fn new_short<'b>(buffer: &'a DecoderBufferMut<'b>) -> Self {
        let initial_buffer_len = buffer.len();
        let peek = buffer.peek();
        let peek = peek.skip(size_of::<Tag>()).expect("tag already verified");
        Self {
            initial_buffer_len,
            peek,
        }
    }

    pub fn decode_destination_connection_id(
        &mut self,
        buffer: &DecoderBufferMut<'_>,
    ) -> Result<CheckedRange, DecoderError> {
        let destination_connection_id =
            self.decode_checked_range::<DestinationConnectionIdLen>(buffer)?;
        validate_destination_connection_id_range(&destination_connection_id)?;
        Ok(destination_connection_id)
    }

    pub fn decode_short_destination_connection_id<Validator: connection::id::Validator>(
        &mut self,
        buffer: &DecoderBufferMut<'_>,
        connection_info: &ConnectionInfo,
        connection_id_validator: &Validator,
    ) -> Result<CheckedRange, DecoderError> {
        let destination_connection_id_len = if let Some(len) = connection_id_validator
            .validate(connection_info, self.peek.peek().into_less_safe_slice())
        {
            len
        } else {
            return Err(DecoderError::InvariantViolation("invalid connection id"));
        };

        let (destination_connection_id, peek) = self
            .peek
            .skip_into_range(destination_connection_id_len, buffer)?;
        self.peek = peek;
        validate_destination_connection_id_range(&destination_connection_id)?;
        Ok(destination_connection_id)
    }

    pub fn decode_source_connection_id(
        &mut self,
        buffer: &DecoderBufferMut<'_>,
    ) -> Result<CheckedRange, DecoderError> {
        let source_connection_id = self.decode_checked_range::<SourceConnectionIdLen>(buffer)?;
        validate_source_connection_id_range(&source_connection_id)?;
        Ok(source_connection_id)
    }

    pub fn decode_checked_range<Len: DecoderValue<'a> + TryInto<usize>>(
        &mut self,
        buffer: &DecoderBufferMut<'_>,
    ) -> Result<CheckedRange, DecoderError> {
        let (value, peek) = self.peek.skip_into_range_with_len_prefix::<Len>(buffer)?;
        self.peek = peek;
        Ok(value)
    }

    pub fn finish_long(mut self) -> Result<HeaderDecoderResult, DecoderError> {
        let (payload_len, peek) = self.peek.decode::<VarInt>()?;
        self.peek = peek;
        let header_len = self.decoded_len();

        self.peek = peek.skip(*payload_len as usize)?;
        let packet_len = self.decoded_len();

        Ok(HeaderDecoderResult {
            packet_len,
            header_len,
        })
    }

    pub fn finish_short(self) -> Result<HeaderDecoderResult, DecoderError> {
        let header_len = self.decoded_len();
        let packet_len = self.initial_buffer_len;

        Ok(HeaderDecoderResult {
            packet_len,
            header_len,
        })
    }

    pub fn decoded_len(&self) -> usize {
        self.initial_buffer_len - self.peek.len()
    }
}

#[derive(Debug)]
pub struct HeaderDecoderResult {
    pub packet_len: usize,
    pub header_len: usize,
}

impl HeaderDecoderResult {
    pub fn split_off_packet<'a>(
        &self,
        buffer: DecoderBufferMut<'a>,
    ) -> Result<
        (
            ProtectedPayload<'a>,
            ProtectedPacketNumber,
            DecoderBufferMut<'a>,
        ),
        DecoderError,
    > {
        let (payload, remaining) = buffer.decode_slice(self.packet_len)?;
        let packet_number = ProtectedPacketNumber::default();
        let payload = ProtectedPayload::new(self.header_len, payload.into_less_safe_slice());

        Ok((payload, packet_number, remaining))
    }
}


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

use crate::{
    crypto::{HeaderKey, Key as CryptoKey, ProtectedPayload},
    packet::{
        number::{PacketNumber, PacketNumberLen},
        stateless_reset,
    },
};
use s2n_codec::{Encoder, EncoderBuffer, EncoderLenEstimator, EncoderValue};

pub trait PacketPayloadLenCursor: EncoderValue {
    fn new() -> Self;
    fn update(&self, buffer: &mut EncoderBuffer, actual_len: usize);
}

/// used for short packets that don't use a payload len
impl PacketPayloadLenCursor for () {
    fn new() {}

    fn update(&self, _buffer: &mut EncoderBuffer, _actual_len: usize) {
        // noop
    }
}

pub trait PacketPayloadEncoder {
    /// Returns an estimate of the encoding size of the payload. This
    /// may be inaccurate from what actually is encoded. Estimates should
    /// be less than or equal to what is actually written.
    /// Implementations can return 0 to skip encoding.
    fn encoding_size_hint<E: Encoder>(&mut self, encoder: &E, minimum_len: usize) -> usize;

    /// Encodes the payload into the buffer. Implementations should ensure
    /// the encoding len is at least the minimum_len, otherwise the packet
    /// writing will panic.
    fn encode(
        &mut self,
        buffer: &mut EncoderBuffer,
        minimum_len: usize,
        header_len: usize,
        tag_len: usize,
    );
}

impl<T: EncoderValue> PacketPayloadEncoder for T {
    fn encoding_size_hint<E: Encoder>(&mut self, encoder: &E, minimum_len: usize) -> usize {
        let len = self.encoding_size_for_encoder(encoder);
        if len < minimum_len {
            0
        } else {
            len
        }
    }

    fn encode(
        &mut self,
        buffer: &mut EncoderBuffer,
        _minimum_len: usize,
        _header_len: usize,
        _tag_len: usize,
    ) {
        // the minimum len check is not needed, as it was already performed
        // in encoding_size_hint
        self.encode_mut(buffer);
    }
}

#[derive(Debug)]
pub enum PacketEncodingError<'a> {
    /// The packet number could not be truncated with the
    /// current largest_acknowledged_packet_number
    PacketNumberTruncationError(EncoderBuffer<'a>),

    /// The buffer does not have enough space to hold
    /// the packet encoding.
    InsufficientSpace(EncoderBuffer<'a>),

    /// The payload did not write anything
    EmptyPayload(EncoderBuffer<'a>),

    /// The key used to encrypt the buffer has exceeded the confidentiality limit
    AeadLimitReached(EncoderBuffer<'a>),
}

impl<'a> PacketEncodingError<'a> {
    /// Returns the buffer that experienced an encoding error
    pub fn take_buffer(self) -> EncoderBuffer<'a> {
        match self {
            Self::PacketNumberTruncationError(buffer) => buffer,
            Self::InsufficientSpace(buffer) => buffer,
            Self::EmptyPayload(buffer) => buffer,
            Self::AeadLimitReached(buffer) => buffer,
        }
    }
}

pub trait PacketEncoder<K: CryptoKey, H: HeaderKey, Payload: PacketPayloadEncoder>: Sized {
    type PayloadLenCursor: PacketPayloadLenCursor;

    /// Encodes the current packet's header into the provided encoder
    fn encode_header<E: Encoder>(&self, packet_number_len: PacketNumberLen, encoder: &mut E);

    /// Returns the payload for the current packet
    fn payload(&mut self) -> &mut Payload;

    /// Returns the packet number for the current packet
    fn packet_number(&self) -> PacketNumber;

    // Encodes, encrypts, and header-protects a packet into a buffer
    fn encode_packet<'a>(
        mut self,
        key: &K,
        header_key: &H,
        largest_acknowledged_packet_number: PacketNumber,
        min_packet_len: Option<usize>,
        mut buffer: EncoderBuffer<'a>,
    ) -> Result<(ProtectedPayload<'a>, EncoderBuffer<'a>), PacketEncodingError<'a>> {
        let packet_number = self.packet_number();

        // Truncate the packet number from the largest_acknowledged_packet_number.
        let truncated_packet_number =
            if let Some(tpn) = packet_number.truncate(largest_acknowledged_packet_number) {
                tpn
            } else {
                return Err(PacketEncodingError::PacketNumberTruncationError(buffer));
            };

        let packet_number_len = truncated_packet_number.len();

        // We need to build an estimate of how large this packet is before writing it
        let mut estimator = EncoderLenEstimator::new(buffer.remaining_capacity());

        // Start by measuring the header len
        self.encode_header(packet_number_len, &mut estimator);

        // Create a payload_len cursor so we can update it with the actual value
        // after writing the payload
        let mut payload_len_cursor = Self::PayloadLenCursor::new();

        // `encode_mut` is called here to initialize the len based on the
        // remaining buffer capacity
        payload_len_cursor.encode_mut(&mut estimator);

        // Save the header_len for later use
        let header_len = estimator.len();

        // Record the truncated_packet_number encoding size
        truncated_packet_number.encode(&mut estimator);

        // Make sure the crypto tag can be written.
        // We want to write this before the payload in the
        // estimator so the payload size hint has an accurate
        // view of remaining capacity.
        estimator.write_repeated(key.tag_len(), 0);

        //= https://www.rfc-editor.org/rfc/rfc9000#section-10.3
        //# To achieve that end,
        //# the endpoint SHOULD ensure that all packets it sends are at least 22
        //# bytes longer than the minimum connection ID length that it requests
        //# the peer to include in its packets, adding PADDING frames as
        //# necessary.
        // One additional byte is added so that a stateless reset sent in response to this packet
        // (which is required to be smaller than this packet) is large enough to be
        // indistinguishable from a valid packet.
        let minimum_packet_len = min_packet_len
            .unwrap_or(0)
            .max(stateless_reset::min_indistinguishable_packet_len(key.tag_len()) + 1);

        // Compute how much the payload will need to write to satisfy the
        // minimum_packet_len
        let minimum_payload_len = minimum_packet_len.saturating_sub(estimator.len());

        //= https://www.rfc-editor.org/rfc/rfc9001#section-5.4.2
        //# in sampling packet ciphertext for header protection,
        //# the Packet Number field is assumed to be 4 bytes long

        // Header protection sampling assumes a packet number length of 4 bytes,
        // but the actual packet number may be smaller than that, so we need to ensure
        // there is still enough payload to sample from given the actual packet number length.
        let minimum_payload_len = minimum_payload_len.max(
            PacketNumberLen::MAX_LEN - truncated_packet_number.len().bytesize()
                + header_key.sealing_sample_len(),
        );

        // Try to estimate the payload size - it may be inaccurate
        // but this provides some checks to save writing the packet
        // header
        let estimated_payload_len = self
            .payload()
            .encoding_size_hint(&estimator, minimum_payload_len);

        // The payload is not interested in writing to this packet
        if estimated_payload_len == 0 {
            return Err(PacketEncodingError::EmptyPayload(buffer));
        }

        // Use the estimated_payload_len to check if we're
        // going to have enough room for it.
        estimator.write_repeated(estimated_payload_len, 0);

        // We don't have enough room to write this packet
        if estimator.overflowed() {
            return Err(PacketEncodingError::InsufficientSpace(buffer));
        }

        // After computing the minimum length we actually start writing to the buffer

        // Now we actually encode the header
        self.encode_header(packet_number_len, &mut buffer);

        // Write the estimated payload len. This will be updated
        // with the accurate value after everything is written.
        payload_len_cursor.encode(&mut buffer);

        // Write the packet number
        truncated_packet_number.encode(&mut buffer);

        let payload_len = {
            // Create a temporary buffer for writing the payload
            let (header_buffer, payload_buffer) = buffer.split_mut();

            // Payloads should not be able to write into the crypto tag space
            let payload_len = payload_buffer.len() - key.tag_len();
            let mut payload_buffer = EncoderBuffer::new(&mut payload_buffer[..payload_len]);

            // Try to encode the payload into the buffer
            self.payload().encode(
                &mut payload_buffer,
                minimum_payload_len,
                header_buffer.len(),
                key.tag_len(),
            );

            // read how much was written
            payload_buffer.len()
        };

        // The payload didn't have anything to write so rewind the cursor
        if payload_len == 0 {
            buffer.set_position(0);
            return Err(PacketEncodingError::EmptyPayload(buffer));
        }

        debug_assert!(
            payload_len >= minimum_payload_len,
            "payloads should write at least the minimum_len"
        );

        // Advance the buffer cursor by what the payload wrote
        buffer.advance_position(payload_len);

        // Update the payload_len cursor with the actual payload len
        let actual_payload_len = buffer.len() + key.tag_len() - header_len;
        payload_len_cursor.update(&mut buffer, actual_payload_len);

        // Encrypt the written payload. Note that the tag is appended to the
        // buffer in the `encrypt` function.
        let (encrypted_payload, remaining) =
            crate::crypto::encrypt(key, packet_number, packet_number_len, header_len, buffer)
                .expect("encryption should always work");

        // Protect the packet
        let protected_payload = crate::crypto::protect(header_key, encrypted_payload)
            .expect("header protection should always work");

        // SUCCESS!!!

        Ok((protected_payload, remaining))
    }
}


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

use crate::{
    crypto::{CryptoError, EncryptedPayload, HandshakeHeaderKey, HandshakeKey, ProtectedPayload},
    packet::{
        decoding::HeaderDecoder,
        encoding::{PacketEncoder, PacketPayloadEncoder},
        long::{
            DestinationConnectionIdLen, LongPayloadEncoder, LongPayloadLenCursor,
            SourceConnectionIdLen, Version,
        },
        number::{
            PacketNumber, PacketNumberLen, PacketNumberSpace, ProtectedPacketNumber,
            TruncatedPacketNumber,
        },
        KeyPhase, Tag,
    },
    transport,
    varint::VarInt,
};
use s2n_codec::{CheckedRange, DecoderBufferMut, DecoderBufferMutResult, Encoder, EncoderValue};

//= https://www.rfc-editor.org/rfc/rfc9000#section-17.2.4
//# Handshake Packet {
//#   Header Form (1) = 1,
//#   Fixed Bit (1) = 1,
//#   Long Packet Type (2) = 2,
//#   Reserved Bits (2),
//#   Packet Number Length (2),
//#   Version (32),
//#   Destination Connection ID Length (8),
//#   Destination Connection ID (0..160),
//#   Source Connection ID Length (8),
//#   Source Connection ID (0..160),
//#   Length (i),
//#   Packet Number (8..32),
//#   Packet Payload (..),
//# }

//= https://www.rfc-editor.org/rfc/rfc9000#section-17.2.4
//# A Handshake packet uses long headers with a type value of 0x2
macro_rules! handshake_tag {
    () => {
        0b1110u8
    };
}

#[derive(Debug)]
pub struct Handshake<DCID, SCID, PacketNumber, Payload> {
    pub version: Version,
    pub destination_connection_id: DCID,
    pub source_connection_id: SCID,
    pub packet_number: PacketNumber,
    pub payload: Payload,
}

pub type ProtectedHandshake<'a> =
    Handshake<CheckedRange, CheckedRange, ProtectedPacketNumber, ProtectedPayload<'a>>;
pub type EncryptedHandshake<'a> =
    Handshake<CheckedRange, CheckedRange, PacketNumber, EncryptedPayload<'a>>;
pub type CleartextHandshake<'a> = Handshake<&'a [u8], &'a [u8], PacketNumber, DecoderBufferMut<'a>>;

impl<'a> ProtectedHandshake<'a> {
    #[inline]
    pub(crate) fn decode(
        _tag: Tag,
        version: Version,
        buffer: DecoderBufferMut,
    ) -> DecoderBufferMutResult<ProtectedHandshake> {
        let mut decoder = HeaderDecoder::new_long(&buffer);

        //= https://www.rfc-editor.org/rfc/rfc9000#section-17.2
        //# Endpoints that receive a version 1 long header
        //# with a value larger than 20 MUST drop the packet.
        let destination_connection_id = decoder.decode_destination_connection_id(&buffer)?;
        let source_connection_id = decoder.decode_source_connection_id(&buffer)?;

        let (payload, packet_number, remaining) =
            decoder.finish_long()?.split_off_packet(buffer)?;

        let packet = Handshake {
            version,
            destination_connection_id,
            source_connection_id,
            packet_number,
            payload,
        };

        Ok((packet, remaining))
    }

    pub fn unprotect<K: HandshakeHeaderKey>(
        self,
        key: &K,
        largest_acknowledged_packet_number: PacketNumber,
    ) -> Result<EncryptedHandshake<'a>, CryptoError> {
        let Handshake {
            version,
            destination_connection_id,
            source_connection_id,
            payload,
            ..
        } = self;

        let (truncated_packet_number, payload) =
            crate::crypto::unprotect(key, PacketNumberSpace::Handshake, payload)?;

        let packet_number = truncated_packet_number.expand(largest_acknowledged_packet_number);

        Ok(Handshake {
            version,
            destination_connection_id,
            source_connection_id,
            packet_number,
            payload,
        })
    }

    #[inline]
    pub fn destination_connection_id(&self) -> &[u8] {
        self.payload
            .get_checked_range(&self.destination_connection_id)
            .into_less_safe_slice()
    }

    #[inline]
    pub fn source_connection_id(&self) -> &[u8] {
        self.payload
            .get_checked_range(&self.source_connection_id)
            .into_less_safe_slice()
    }
}

impl<'a> EncryptedHandshake<'a> {
    pub fn decrypt<C: HandshakeKey>(
        self,
        crypto: &C,
    ) -> Result<CleartextHandshake<'a>, transport::Error> {
        let Handshake {
            version,
            destination_connection_id,
            source_connection_id,
            packet_number,
            payload,
        } = self;

        let (header, payload) = crate::crypto::decrypt(crypto, packet_number, payload)?;

        let header = header.into_less_safe_slice();

        //= https://www.rfc-editor.org/rfc/rfc9000#section-17.2
        //# The value
        //# included prior to protection MUST be set to 0.  An endpoint MUST
        //# treat receipt of a packet that has a non-zero value for these bits
        //# after removing both packet and header protection as a connection
        //# error of type PROTOCOL_VIOLATION.
        if header[0] & super::long::RESERVED_BITS_MASK != 0 {
            return Err(
                transport::Error::PROTOCOL_VIOLATION.with_reason("reserved bits are non-zero")
            );
        }

        let destination_connection_id = destination_connection_id.get(header);
        let source_connection_id = source_connection_id.get(header);

        Ok(Handshake {
            version,
            destination_connection_id,
            source_connection_id,
            packet_number,
            payload,
        })
    }

    #[inline]
    pub fn destination_connection_id(&self) -> &[u8] {
        self.payload
            .get_checked_range(&self.destination_connection_id)
            .into_less_safe_slice()
    }

    #[inline]
    pub fn source_connection_id(&self) -> &[u8] {
        self.payload
            .get_checked_range(&self.source_connection_id)
            .into_less_safe_slice()
    }

    // HandshakePackets do not have a KeyPhase
    #[inline]
    pub fn key_phase(&self) -> KeyPhase {
        KeyPhase::Zero
    }
}

impl<'a> CleartextHandshake<'a> {
    #[inline]
    pub fn destination_connection_id(&self) -> &[u8] {
        self.destination_connection_id
    }

    #[inline]
    pub fn source_connection_id(&self) -> &[u8] {
        self.source_connection_id
    }
}

impl<DCID: EncoderValue, SCID: EncoderValue, Payload: EncoderValue> EncoderValue
    for Handshake<DCID, SCID, TruncatedPacketNumber, Payload>
{
    fn encode<E: Encoder>(&self, encoder: &mut E) {
        self.encode_header(self.packet_number.len(), encoder);
        LongPayloadEncoder {
            packet_number: self.packet_number,
            payload: &self.payload,
        }
        .encode_with_len_prefix::<VarInt, E>(encoder)
    }
}

impl<DCID: EncoderValue, SCID: EncoderValue, PacketNumber, Payload>
    Handshake<DCID, SCID, PacketNumber, Payload>
{
    fn encode_header<E: Encoder>(&self, packet_number_len: PacketNumberLen, encoder: &mut E) {
        let mut tag: u8 = handshake_tag!() << 4;
        tag |= packet_number_len.into_packet_tag_mask();
        tag.encode(encoder);

        self.version.encode(encoder);
        self.destination_connection_id
            .encode_with_len_prefix::<DestinationConnectionIdLen, E>(encoder);
        self.source_connection_id
            .encode_with_len_prefix::<SourceConnectionIdLen, E>(encoder);
    }
}

impl<
        DCID: EncoderValue,
        SCID: EncoderValue,
        Payload: PacketPayloadEncoder,
        K: HandshakeKey,
        H: HandshakeHeaderKey,
    > PacketEncoder<K, H, Payload> for Handshake<DCID, SCID, PacketNumber, Payload>
{
    type PayloadLenCursor = LongPayloadLenCursor;

    fn packet_number(&self) -> PacketNumber {
        self.packet_number
    }

    fn encode_header<E: Encoder>(&self, packet_number_len: PacketNumberLen, encoder: &mut E) {
        Handshake::encode_header(self, packet_number_len, encoder);
    }

    fn payload(&mut self) -> &mut Payload {
        &mut self.payload
    }
}


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

use crate::{
    crypto::{CryptoError, EncryptedPayload, InitialHeaderKey, InitialKey, ProtectedPayload},
    packet::{
        decoding::HeaderDecoder,
        encoding::{PacketEncoder, PacketPayloadEncoder},
        long::{
            DestinationConnectionIdLen, LongPayloadEncoder, LongPayloadLenCursor,
            SourceConnectionIdLen, Version,
        },
        number::{
            PacketNumber, PacketNumberLen, PacketNumberSpace, ProtectedPacketNumber,
            TruncatedPacketNumber,
        },
        KeyPhase, Tag,
    },
    varint::VarInt,
};
use s2n_codec::{CheckedRange, DecoderBufferMut, DecoderBufferMutResult, Encoder, EncoderValue};

//= https://www.rfc-editor.org/rfc/rfc9000#section-17.2.2
//# Initial Packet {
//#   Header Form (1) = 1,
//#   Fixed Bit (1) = 1,
//#   Long Packet Type (2) = 0,
//#   Reserved Bits (2),
//#   Packet Number Length (2),
//#   Version (32),
//#   Destination Connection ID Length (8),
//#   Destination Connection ID (0..160),
//#   Source Connection ID Length (8),
//#   Source Connection ID (0..160),
//#   Token Length (i),
//#   Token (..),
//#   Length (i),
//#   Packet Number (8..32),
//#   Packet Payload (..),
//# }

//= https://www.rfc-editor.org/rfc/rfc9000#section-17.2.2
//# An Initial packet uses long headers with a type value of 0x0.
macro_rules! initial_tag {
    () => {
        0b1100u8
    };
}

//= https://www.rfc-editor.org/rfc/rfc9000#section-17.2.2
//# Token Length:  A variable-length integer specifying the length of the
//# Token field, in bytes.  This value is 0 if no token is present.

//= https://www.rfc-editor.org/rfc/rfc9000#section-17.2.2
//# Token:  The value of the token that was previously provided in a
//#    Retry packet or NEW_TOKEN frame; see Section 8.1.

#[derive(Debug)]
pub struct Initial<DCID, SCID, Token, PacketNumber, Payload> {
    pub version: Version,
    pub destination_connection_id: DCID,
    pub source_connection_id: SCID,
    pub token: Token,
    pub packet_number: PacketNumber,
    pub payload: Payload,
}

pub type ProtectedInitial<'a> =
    Initial<CheckedRange, CheckedRange, CheckedRange, ProtectedPacketNumber, ProtectedPayload<'a>>;
pub type EncryptedInitial<'a> =
    Initial<CheckedRange, CheckedRange, CheckedRange, PacketNumber, EncryptedPayload<'a>>;
pub type CleartextInitial<'a> =
    Initial<&'a [u8], &'a [u8], &'a [u8], PacketNumber, DecoderBufferMut<'a>>;

impl<'a> ProtectedInitial<'a> {
    #[inline]
    pub(crate) fn decode(
        _tag: Tag,
        version: Version,
        buffer: DecoderBufferMut,
    ) -> DecoderBufferMutResult<ProtectedInitial> {
        let mut decoder = HeaderDecoder::new_long(&buffer);

        //= https://www.rfc-editor.org/rfc/rfc9000#section-17.2
        //# In order to
        //# properly form a Version Negotiation packet, servers SHOULD be able
        //# to read longer connection IDs from other QUIC versions.
        // Connection ID validation for Initial packets occurs after version
        // negotiation has determined the specified version is supported.
        let destination_connection_id =
            decoder.decode_checked_range::<DestinationConnectionIdLen>(&buffer)?;
        let source_connection_id =
            decoder.decode_checked_range::<SourceConnectionIdLen>(&buffer)?;
        let token = decoder.decode_checked_range::<VarInt>(&buffer)?;

        let (payload, packet_number, remaining) =
            decoder.finish_long()?.split_off_packet(buffer)?;

        let packet = Initial {
            version,
            destination_connection_id,
            source_connection_id,
            token,
            packet_number,
            payload,
        };

        Ok((packet, remaining))
    }

    pub fn unprotect<H: InitialHeaderKey>(
        self,
        header_key: &H,
        largest_acknowledged_packet_number: PacketNumber,
    ) -> Result<EncryptedInitial<'a>, CryptoError> {
        let Initial {
            version,
            destination_connection_id,
            source_connection_id,
            token,
            payload,
            ..
        } = self;

        let (truncated_packet_number, payload) =
            crate::crypto::unprotect(header_key, PacketNumberSpace::Initial, payload)?;

        let packet_number = truncated_packet_number.expand(largest_acknowledged_packet_number);

        Ok(Initial {
            version,
            destination_connection_id,
            source_connection_id,
            token,
            packet_number,
            payload,
        })
    }

    #[inline]
    pub fn destination_connection_id(&self) -> &[u8] {
        self.payload
            .get_checked_range(&self.destination_connection_id)
            .into_less_safe_slice()
    }

    #[inline]
    pub fn source_connection_id(&self) -> &[u8] {
        self.payload
            .get_checked_range(&self.source_connection_id)
            .into_less_safe_slice()
    }

    #[inline]
    pub fn token(&self) -> &[u8] {
        self.payload
            .get_checked_range(&self.token)
            .into_less_safe_slice()
    }
}

impl<'a> EncryptedInitial<'a> {
    pub fn decrypt<C: InitialKey>(self, crypto: &C) -> Result<CleartextInitial<'a>, CryptoError> {
        let Initial {
            version,
            destination_connection_id,
            source_connection_id,
            token,
            packet_number,
            payload,
        } = self;

        let (header, payload) = crate::crypto::decrypt(crypto, packet_number, payload)?;

        let header = header.into_less_safe_slice();

        let destination_connection_id = destination_connection_id.get(header);
        let source_connection_id = source_connection_id.get(header);
        let token = token.get(header);

        Ok(Initial {
            version,
            destination_connection_id,
            source_connection_id,
            token,
            packet_number,
            payload,
        })
    }

    #[inline]
    pub fn destination_connection_id(&self) -> &[u8] {
        self.payload
            .get_checked_range(&self.destination_connection_id)
            .into_less_safe_slice()
    }

    #[inline]
    pub fn source_connection_id(&self) -> &[u8] {
        self.payload
            .get_checked_range(&self.source_connection_id)
            .into_less_safe_slice()
    }

    #[inline]
    pub fn token(&self) -> &[u8] {
        self.payload
            .get_checked_range(&self.token)
            .into_less_safe_slice()
    }

    // InitialPackets do not have a KeyPhase
    #[inline]
    pub fn key_phase(&self) -> KeyPhase {
        KeyPhase::Zero
    }
}

impl<'a> CleartextInitial<'a> {
    #[inline]
    pub fn destination_connection_id(&self) -> &[u8] {
        self.destination_connection_id
    }

    #[inline]
    pub fn source_connection_id(&self) -> &[u8] {
        self.source_connection_id
    }

    #[inline]
    pub fn token(&self) -> &[u8] {
        self.token
    }
}

impl<DCID: EncoderValue, SCID: EncoderValue, Token: EncoderValue, Payload: EncoderValue>
    EncoderValue for Initial<DCID, SCID, Token, TruncatedPacketNumber, Payload>
{
    fn encode<E: Encoder>(&self, encoder: &mut E) {
        self.encode_header(self.packet_number.len(), encoder);
        LongPayloadEncoder {
            packet_number: self.packet_number,
            payload: &self.payload,
        }
        .encode_with_len_prefix::<VarInt, E>(encoder)
    }
}

impl<DCID: EncoderValue, SCID: EncoderValue, Token: EncoderValue, PacketNumber, Payload>
    Initial<DCID, SCID, Token, PacketNumber, Payload>
{
    fn encode_header<E: Encoder>(&self, packet_number_len: PacketNumberLen, encoder: &mut E) {
        let mut tag: u8 = initial_tag!() << 4;
        tag |= packet_number_len.into_packet_tag_mask();
        tag.encode(encoder);

        self.version.encode(encoder);

        self.destination_connection_id
            .encode_with_len_prefix::<DestinationConnectionIdLen, E>(encoder);
        self.source_connection_id
            .encode_with_len_prefix::<SourceConnectionIdLen, E>(encoder);
        self.token.encode_with_len_prefix::<VarInt, E>(encoder);
    }
}

impl<
        DCID: EncoderValue,
        SCID: EncoderValue,
        Token: EncoderValue,
        Payload: PacketPayloadEncoder,
        K: InitialKey,
        H: InitialHeaderKey,
    > PacketEncoder<K, H, Payload> for Initial<DCID, SCID, Token, PacketNumber, Payload>
{
    type PayloadLenCursor = LongPayloadLenCursor;

    fn packet_number(&self) -> PacketNumber {
        self.packet_number
    }

    fn encode_header<E: Encoder>(&self, packet_number_len: PacketNumberLen, encoder: &mut E) {
        Initial::encode_header(self, packet_number_len, encoder);
    }

    fn payload(&mut self) -> &mut Payload {
        &mut self.payload
    }
}


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

use crate::{
    event::api::{SocketAddress, Subject},
    havoc,
    packet::number::PacketNumber,
    time::Timestamp,
};
use s2n_codec::{DecoderBufferMut, EncoderBuffer};

pub mod loss;
pub use loss::Loss;

/// TODO add `non_exhaustive` once/if this feature is stable
#[derive(Debug)]
pub struct Packet {
    pub number: PacketNumber,
    pub timestamp: Timestamp,
}

/// TODO add `non_exhaustive` once/if this feature is stable
#[derive(Debug)]
pub struct Datagram<'a> {
    pub remote_address: SocketAddress<'a>,
    pub local_address: SocketAddress<'a>,
    pub timestamp: Timestamp,
}

/// Trait which enables an application to intercept packets that are transmitted and received
pub trait Interceptor: 'static + Send {
    #[inline(always)]
    fn intercept_rx_datagram<'a>(
        &mut self,
        subject: &Subject,
        datagram: &Datagram,
        payload: DecoderBufferMut<'a>,
    ) -> DecoderBufferMut<'a> {
        let _ = subject;
        let _ = datagram;
        payload
    }

    #[inline(always)]
    fn intercept_rx_payload<'a>(
        &mut self,
        subject: &Subject,
        packet: &Packet,
        payload: DecoderBufferMut<'a>,
    ) -> DecoderBufferMut<'a> {
        let _ = subject;
        let _ = packet;
        payload
    }

    #[inline(always)]
    fn intercept_tx_datagram(
        &mut self,
        subject: &Subject,
        datagram: &Datagram,
        payload: &mut EncoderBuffer,
    ) {
        let _ = subject;
        let _ = datagram;
        let _ = payload;
    }

    #[inline(always)]
    fn intercept_tx_payload(
        &mut self,
        subject: &Subject,
        packet: &Packet,
        payload: &mut EncoderBuffer,
    ) {
        let _ = subject;
        let _ = packet;
        let _ = payload;
    }
}

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

impl Interceptor for Disabled {}

impl<A, B> Interceptor for (A, B)
where
    A: Interceptor,
    B: Interceptor,
{
    #[inline(always)]
    fn intercept_rx_datagram<'a>(
        &mut self,
        subject: &Subject,
        datagram: &Datagram,
        payload: DecoderBufferMut<'a>,
    ) -> DecoderBufferMut<'a> {
        let payload = self.0.intercept_rx_datagram(subject, datagram, payload);
        self.1.intercept_rx_datagram(subject, datagram, payload)
    }

    #[inline(always)]
    fn intercept_rx_payload<'a>(
        &mut self,
        subject: &Subject,
        packet: &Packet,
        payload: DecoderBufferMut<'a>,
    ) -> DecoderBufferMut<'a> {
        let payload = self.0.intercept_rx_payload(subject, packet, payload);
        self.1.intercept_rx_payload(subject, packet, payload)
    }

    #[inline(always)]
    fn intercept_tx_datagram(
        &mut self,
        subject: &Subject,
        datagram: &Datagram,
        payload: &mut EncoderBuffer,
    ) {
        self.0.intercept_tx_datagram(subject, datagram, payload);
        self.1.intercept_tx_datagram(subject, datagram, payload);
    }

    #[inline(always)]
    fn intercept_tx_payload(
        &mut self,
        subject: &Subject,
        packet: &Packet,
        payload: &mut EncoderBuffer,
    ) {
        self.0.intercept_tx_payload(subject, packet, payload);
        self.1.intercept_tx_payload(subject, packet, payload);
    }
}

#[derive(Debug, Default)]
pub struct Havoc<Rx, Tx, R>
where
    Rx: 'static + Send + havoc::Strategy,
    Tx: 'static + Send + havoc::Strategy,
    R: 'static + Send + havoc::Random,
{
    pub rx: Rx,
    pub tx: Tx,
    pub random: R,
}

impl<Rx, Tx, R> Interceptor for Havoc<Rx, Tx, R>
where
    Rx: 'static + Send + havoc::Strategy,
    Tx: 'static + Send + havoc::Strategy,
    R: 'static + Send + havoc::Random,
{
    #[inline]
    fn intercept_rx_payload<'a>(
        &mut self,
        _subject: &Subject,
        _packet: &Packet,
        payload: DecoderBufferMut<'a>,
    ) -> DecoderBufferMut<'a> {
        let payload = payload.into_less_safe_slice();
        let len = payload.len();

        let len = {
            use s2n_codec::Encoder;
            let mut payload = EncoderBuffer::new(payload);
            payload.set_position(len);
            self.rx.havoc(&mut self.random, &mut payload);
            payload.len()
        };

        let payload = &mut payload[..len];

        DecoderBufferMut::new(payload)
    }

    #[inline]
    fn intercept_tx_payload(
        &mut self,
        _subject: &Subject,
        _packet: &Packet,
        payload: &mut EncoderBuffer,
    ) {
        self.tx.havoc(&mut self.random, payload);
    }
}


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

use super::Tag;

//= https://www.rfc-editor.org/rfc/rfc9000#section-17.3.1
//# Key Phase:  The next bit (0x04) of byte 0 indicates the key phase,
//# which allows a recipient of a packet to identify the packet
//# protection keys that are used to protect the packet.

const KEY_PHASE_MASK: u8 = 0x04;

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

#[derive(Clone, Copy, Debug, PartialEq)]
pub enum KeyPhase {
    Zero,
    One,
}

impl Default for KeyPhase {
    fn default() -> Self {
        Self::Zero
    }
}

const PHASES: [KeyPhase; 2] = [KeyPhase::Zero, KeyPhase::One];
impl From<u8> for KeyPhase {
    fn from(v: u8) -> Self {
        // Will only be 0 or 1. Invalid phases may result in a failed decryption, still in constant
        // time.
        PHASES[(v & 0x01) as usize]
    }
}

impl KeyPhase {
    pub fn from_tag(tag: Tag) -> Self {
        let phase = (tag & KEY_PHASE_MASK) >> 2;
        PHASES[phase as usize]
    }

    pub fn into_packet_tag_mask(self) -> u8 {
        match self {
            Self::One => KEY_PHASE_MASK,
            Self::Zero => 0,
        }
    }

    #[must_use]
    pub fn next_phase(self) -> Self {
        PHASES[(((self as u8) + 1) % 2) as usize]
    }
}

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

    #[test]
    fn test_keyphase_from_tag() {
        for i in 0..255 {
            let phase = KeyPhase::from_tag(i);
            assert_eq!(phase.into_packet_tag_mask(), i & KEY_PHASE_MASK);
        }
    }

    #[test]
    fn test_next_phase() {
        for i in 0..254 {
            let phase = KeyPhase::from(i);
            let next_phase = KeyPhase::from(i + 1);
            assert_eq!(phase.next_phase(), next_phase);
        }
    }
}


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

use crate::{
    packet::{encoding::PacketPayloadLenCursor, number::TruncatedPacketNumber},
    varint::VarInt,
};
use core::convert::TryFrom;
use s2n_codec::{
    decoder_invariant, CheckedRange, DecoderError, Encoder, EncoderBuffer, EncoderValue,
};

//= https://www.rfc-editor.org/rfc/rfc9000#section-17.2
//# Long Header Packet {
//#   Header Form (1) = 1,
//#   Fixed Bit (1) = 1,
//#   Long Packet Type (2),
//#   Type-Specific Bits (4),
//#   Version (32),
//#   Destination Connection ID Length (8),
//#   Destination Connection ID (0..160),
//#   Source Connection ID Length (8),
//#   Source Connection ID (0..160),
//# }

//= https://www.rfc-editor.org/rfc/rfc9000#section-17.2
//# Header Form:  The most significant bit (0x80) of byte 0 (the first
//#   byte) is set to 1 for long headers.

//= https://www.rfc-editor.org/rfc/rfc9000#section-17.2
//# Fixed Bit:  The next bit (0x40) of byte 0 is set to 1.

//= https://www.rfc-editor.org/rfc/rfc9000#section-17.2
//# Long Packet Type:  The next two bits (those with a mask of 0x30)
//#    of byte 0 contain a packet type.  Packet types are listed in
//#    Table 5.

pub(crate) const PACKET_TYPE_MASK: u8 = 0x30;
const PACKET_TYPE_OFFSET: u8 = 4;

//= https://www.rfc-editor.org/rfc/rfc9000#section-17.2
//# Type-Specific Bits:  The semantics of the lower four bits (those with
//# a mask of 0x0f) of byte 0 are determined by the packet type.

//= https://www.rfc-editor.org/rfc/rfc9000#section-17.2
//# Version: The QUIC Version is a 32-bit field that follows the first
//#    byte.  This field indicates the version of QUIC that is in use and
//#    determines how the rest of the protocol fields are interpreted.

pub(crate) type Version = u32;

//= https://www.rfc-editor.org/rfc/rfc9000#section-17.2
//# Destination Connection ID Length:  The byte following the version
//#    contains the length in bytes of the Destination Connection ID
//#    field that follows it.  This length is encoded as an 8-bit
//#    unsigned integer.  In QUIC version 1, this value MUST NOT exceed
//#    20.

pub(crate) type DestinationConnectionIdLen = u8;
pub(crate) const DESTINATION_CONNECTION_ID_MAX_LEN: usize = 20;

//= https://www.rfc-editor.org/rfc/rfc9000#section-17.2
//# Destination Connection ID:  The Destination Connection ID field
//#   follows the Destination Connection ID Length field, which
//#   indicates the length of this field.  Section 7.2 describes the use
//#   of this field in more detail.

pub(crate) fn validate_destination_connection_id_range(
    range: &CheckedRange,
) -> Result<(), DecoderError> {
    validate_destination_connection_id_len(range.len())
}

pub(crate) fn validate_destination_connection_id_len(len: usize) -> Result<(), DecoderError> {
    //= https://www.rfc-editor.org/rfc/rfc9000#section-17.2
    //# Endpoints that receive a version 1 long header with a value
    //# larger than 20 MUST drop the packet.
    decoder_invariant!(
        len <= DESTINATION_CONNECTION_ID_MAX_LEN,
        "destination connection exceeds max length"
    );
    Ok(())
}

//= https://www.rfc-editor.org/rfc/rfc9000#section-17.2
//# Source Connection ID Length:  The byte following the Destination
//#   Connection ID contains the length in bytes of the Source
//#   Connection ID field that follows it.  This length is encoded as a
//#   8-bit unsigned integer.  In QUIC version 1, this value MUST NOT
//#   exceed 20 bytes.

pub(crate) type SourceConnectionIdLen = u8;
pub(crate) const SOURCE_CONNECTION_ID_MAX_LEN: usize = 20;

//= https://www.rfc-editor.org/rfc/rfc9000#section-17.2
//# Source Connection ID:  The Source Connection ID field follows the
//#   Source Connection ID Length field, which indicates the length of
//#   this field.  Section 7.2 describes the use of this field in more
//#   detail.

pub(crate) fn validate_source_connection_id_range(
    range: &CheckedRange,
) -> Result<(), DecoderError> {
    //= https://www.rfc-editor.org/rfc/rfc9000#section-17.2
    //# Endpoints that receive a version 1 long header
    //# with a value larger than 20 MUST drop the packet.
    validate_source_connection_id_len(range.len())
}

pub(crate) fn validate_source_connection_id_len(len: usize) -> Result<(), DecoderError> {
    decoder_invariant!(
        len <= SOURCE_CONNECTION_ID_MAX_LEN,
        "source connection exceeds max length"
    );
    Ok(())
}

//= https://www.rfc-editor.org/rfc/rfc9000#section-17.2
//# In this version of QUIC, the following packet types with the long
//# header are defined:
//#
//#                 +======+===========+================+
//#                 | Type | Name      | Section        |
//#                 +======+===========+================+
//#                 | 0x00 | Initial   | Section 17.2.2 |
//#                 +------+-----------+----------------+
//#                 | 0x01 | 0-RTT     | Section 17.2.3 |
//#                 +------+-----------+----------------+
//#                 | 0x02 | Handshake | Section 17.2.4 |
//#                 +------+-----------+----------------+
//#                 | 0x03 | Retry     | Section 17.2.5 |
//#                 +------+-----------+----------------+
//#
//#                   Table 5: Long Header Packet Types

#[repr(u8)]
#[derive(Clone, Copy, Debug)]
pub enum PacketType {
    Initial = 0x0,
    ZeroRtt = 0x1,
    Handshake = 0x2,
    Retry = 0x3,
}

impl PacketType {
    pub const fn into_bits(self) -> u8 {
        (self as u8) << PACKET_TYPE_OFFSET & PACKET_TYPE_MASK
    }

    pub fn from_bits(bits: u8) -> Self {
        (bits & PACKET_TYPE_MASK >> PACKET_TYPE_OFFSET).into()
    }
}

impl From<u8> for PacketType {
    fn from(bits: u8) -> Self {
        match bits {
            0x0 => PacketType::Initial,
            0x1 => PacketType::ZeroRtt,
            0x2 => PacketType::Handshake,
            0x3 => PacketType::Retry,
            _ => PacketType::Initial,
        }
    }
}

impl From<PacketType> for u8 {
    fn from(v: PacketType) -> Self {
        v.into_bits()
    }
}

//= https://www.rfc-editor.org/rfc/rfc9000#section-17.2
//# Reserved Bits:  Two bits (those with a mask of 0x0c) of byte 0 are
//#    reserved across multiple packet types.  These bits are protected
//#    using header protection; see Section 5.4 of [QUIC-TLS].
pub const RESERVED_BITS_MASK: u8 = 0x0c;

//= https://www.rfc-editor.org/rfc/rfc9000#section-17.2
//# Packet Number Length:  In packet types that contain a Packet Number
//# field, the least significant two bits (those with a mask of 0x03)
//# of byte 0 contain the length of the Packet Number field, encoded
//# as an unsigned two-bit integer that is one less than the length of
//# the Packet Number field in bytes.

//= https://www.rfc-editor.org/rfc/rfc9000#section-17.2
//# Length:  This is the length of the remainder of the packet (that is,
//# the Packet Number and Payload fields) in bytes, encoded as a
//# variable-length integer (Section 16).

//= https://www.rfc-editor.org/rfc/rfc9000#section-17.2
//# Packet Number:  This field is 1 to 4 bytes long.  The packet number
//# is protected using header protection; see Section 5.4 of
//# [QUIC-TLS].  The length of the Packet Number field is encoded in
//# the Packet Number Length bits of byte 0; see above.

pub(crate) struct LongPayloadEncoder<Payload> {
    pub packet_number: TruncatedPacketNumber,
    pub payload: Payload,
}

impl<'a, Payload: EncoderValue> EncoderValue for LongPayloadEncoder<&'a Payload> {
    fn encode<E: Encoder>(&self, encoder: &mut E) {
        self.packet_number.encode(encoder);
        self.payload.encode(encoder);
    }
}

impl<'a, Payload: EncoderValue> EncoderValue for LongPayloadEncoder<&'a mut Payload> {
    fn encode<E: Encoder>(&self, encoder: &mut E) {
        self.packet_number.encode(encoder);
        self.payload.encode(encoder);
    }

    fn encode_mut<E: Encoder>(&mut self, encoder: &mut E) {
        self.packet_number.encode_mut(encoder);
        self.payload.encode_mut(encoder);
    }
}

// used internally for estimating long packet payload len values.
#[doc(hidden)]
#[derive(Clone, Copy, Debug, PartialEq)]
pub struct LongPayloadLenCursor {
    position: usize,
    max_value: VarInt,
}

impl PacketPayloadLenCursor for LongPayloadLenCursor {
    fn new() -> Self {
        let max_value = VarInt::MAX;
        Self {
            position: 0,
            max_value,
        }
    }

    fn update(&self, buffer: &mut EncoderBuffer, actual_len: usize) {
        debug_assert!(
            self.position != 0,
            "position cursor was not updated. encode_mut should be called instead of encode"
        );

        let actual_value =
            VarInt::try_from(actual_len).expect("packets should not be larger than VarInt::MAX");
        let max_value = self.max_value;

        let prev_pos = buffer.len();
        buffer.set_position(self.position);
        max_value.encode_updated(actual_value, buffer);
        buffer.set_position(prev_pos);
    }
}

impl EncoderValue for LongPayloadLenCursor {
    fn encode<E: Encoder>(&self, encoder: &mut E) {
        self.max_value.encode(encoder)
    }

    fn encode_mut<E: Encoder>(&mut self, encoder: &mut E) {
        self.position = encoder.len();
        self.max_value = VarInt::try_from(encoder.remaining_capacity()).unwrap_or(VarInt::MAX);
        self.max_value.encode(encoder)
    }
}


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

use crate::connection;
use s2n_codec::{decoder_invariant, DecoderBufferMut, DecoderBufferMutResult, DecoderError};

pub(crate) type Tag = u8;

#[macro_use]
pub mod short;
#[macro_use]
pub mod version_negotiation;
#[macro_use]
pub mod initial;
#[macro_use]
pub mod zero_rtt;
#[macro_use]
pub mod handshake;
#[macro_use]
pub mod retry;

pub mod decoding;
pub mod encoding;
pub mod interceptor;
pub mod key_phase;
pub mod long;

pub mod number;
pub mod stateless_reset;

#[cfg(test)]
mod tests;

pub use key_phase::{KeyPhase, ProtectedKeyPhase};

use connection::id::ConnectionInfo;
use handshake::ProtectedHandshake;
use initial::ProtectedInitial;
use retry::ProtectedRetry;
use short::ProtectedShort;
use version_negotiation::ProtectedVersionNegotiation;
use zero_rtt::ProtectedZeroRtt;

// === API ===

pub type RemainingBuffer<'a> = Option<DecoderBufferMut<'a>>;

#[derive(Debug)]
pub enum ProtectedPacket<'a> {
    Short(ProtectedShort<'a>),
    VersionNegotiation(ProtectedVersionNegotiation<'a>),
    Initial(ProtectedInitial<'a>),
    ZeroRtt(ProtectedZeroRtt<'a>),
    Handshake(ProtectedHandshake<'a>),
    Retry(ProtectedRetry<'a>),
}

impl<'a> ProtectedPacket<'a> {
    pub fn decode<Validator: connection::id::Validator>(
        buffer: DecoderBufferMut<'a>,
        connection_info: &ConnectionInfo,
        connection_id_validator: &Validator,
    ) -> DecoderBufferMutResult<'a, Self> {
        BasicPacketDecoder.decode_packet(buffer, connection_info, connection_id_validator)
    }

    /// Returns the packet's destination connection ID
    pub fn destination_connection_id(&self) -> &[u8] {
        match self {
            ProtectedPacket::Short(packet) => packet.destination_connection_id(),
            ProtectedPacket::VersionNegotiation(packet) => packet.destination_connection_id(),
            ProtectedPacket::Initial(packet) => packet.destination_connection_id(),
            ProtectedPacket::ZeroRtt(packet) => packet.destination_connection_id(),
            ProtectedPacket::Handshake(packet) => packet.destination_connection_id(),
            ProtectedPacket::Retry(packet) => packet.destination_connection_id(),
        }
    }

    /// Returns the packet's source connection ID
    pub fn source_connection_id(&self) -> Option<&[u8]> {
        match self {
            ProtectedPacket::Short(_packet) => None,
            ProtectedPacket::VersionNegotiation(packet) => Some(packet.source_connection_id()),
            ProtectedPacket::Initial(packet) => Some(packet.source_connection_id()),
            ProtectedPacket::ZeroRtt(packet) => Some(packet.source_connection_id()),
            ProtectedPacket::Handshake(packet) => Some(packet.source_connection_id()),
            ProtectedPacket::Retry(packet) => Some(packet.source_connection_id()),
        }
    }

    pub fn version(&self) -> Option<u32> {
        match self {
            ProtectedPacket::Short(_) => None,
            ProtectedPacket::VersionNegotiation(_) => None,
            ProtectedPacket::Initial(packet) => Some(packet.version),
            ProtectedPacket::ZeroRtt(packet) => Some(packet.version),
            ProtectedPacket::Handshake(packet) => Some(packet.version),
            ProtectedPacket::Retry(packet) => Some(packet.version),
        }
    }
}

#[derive(Debug)]
pub enum CleartextPacket<'a> {
    Short(short::CleartextShort<'a>),
    VersionNegotiation(version_negotiation::CleartextVersionNegotiation<'a>),
    Initial(initial::CleartextInitial<'a>),
    ZeroRtt(zero_rtt::CleartextZeroRtt<'a>),
    Handshake(handshake::CleartextHandshake<'a>),
    Retry(retry::CleartextRetry<'a>),
}

struct BasicPacketDecoder;

impl<'a> PacketDecoder<'a> for BasicPacketDecoder {
    type Error = DecoderError;
    type Output = ProtectedPacket<'a>;

    fn handle_short_packet(
        &mut self,
        packet: ProtectedShort<'a>,
    ) -> Result<Self::Output, DecoderError> {
        Ok(ProtectedPacket::Short(packet))
    }

    fn handle_version_negotiation_packet(
        &mut self,
        packet: ProtectedVersionNegotiation<'a>,
    ) -> Result<Self::Output, DecoderError> {
        Ok(ProtectedPacket::VersionNegotiation(packet))
    }

    fn handle_initial_packet(
        &mut self,
        packet: ProtectedInitial<'a>,
    ) -> Result<Self::Output, DecoderError> {
        Ok(ProtectedPacket::Initial(packet))
    }

    fn handle_zero_rtt_packet(
        &mut self,
        packet: ProtectedZeroRtt<'a>,
    ) -> Result<Self::Output, DecoderError> {
        Ok(ProtectedPacket::ZeroRtt(packet))
    }

    fn handle_handshake_packet(
        &mut self,
        packet: ProtectedHandshake<'a>,
    ) -> Result<Self::Output, DecoderError> {
        Ok(ProtectedPacket::Handshake(packet))
    }

    fn handle_retry_packet(
        &mut self,
        packet: ProtectedRetry<'a>,
    ) -> Result<Self::Output, DecoderError> {
        Ok(ProtectedPacket::Retry(packet))
    }
}

pub trait PacketDecoder<'a> {
    type Output;
    type Error: From<DecoderError>;

    fn handle_short_packet(
        &mut self,
        packet: ProtectedShort<'a>,
    ) -> Result<Self::Output, Self::Error>;

    fn handle_version_negotiation_packet(
        &mut self,
        packet: ProtectedVersionNegotiation<'a>,
    ) -> Result<Self::Output, Self::Error>;

    fn handle_initial_packet(
        &mut self,
        packet: ProtectedInitial<'a>,
    ) -> Result<Self::Output, Self::Error>;

    fn handle_zero_rtt_packet(
        &mut self,
        packet: ProtectedZeroRtt<'a>,
    ) -> Result<Self::Output, Self::Error>;

    fn handle_handshake_packet(
        &mut self,
        packet: ProtectedHandshake<'a>,
    ) -> Result<Self::Output, Self::Error>;

    fn handle_retry_packet(
        &mut self,
        packet: ProtectedRetry<'a>,
    ) -> Result<Self::Output, Self::Error>;

    fn decode_packet<Validator: connection::id::Validator>(
        &mut self,
        buffer: DecoderBufferMut<'a>,
        connection_info: &ConnectionInfo,
        connection_id_validator: &Validator,
    ) -> Result<(Self::Output, DecoderBufferMut<'a>), Self::Error> {
        let peek = buffer.peek();

        let (tag, peek) = peek.decode()?;

        macro_rules! version_negotiation {
            ($version:ident) => {{
                let (packet, buffer) = ProtectedVersionNegotiation::decode(tag, $version, buffer)?;
                let output = self.handle_version_negotiation_packet(packet)?;
                Ok((output, buffer))
            }};
        }

        macro_rules! long_packet {
            ($struct:ident, $handler:ident) => {{
                let (version, _peek) = peek.decode()?;
                if version == version_negotiation::VERSION {
                    version_negotiation!(version)
                } else {
                    let (packet, buffer) = $struct::decode(tag, version, buffer)?;
                    let output = self.$handler(packet)?;
                    Ok((output, buffer))
                }
            }};
        }

        match tag >> 4 {
            short_tag!() => {
                let (packet, buffer) = short::ProtectedShort::decode(
                    tag,
                    buffer,
                    connection_info,
                    connection_id_validator,
                )?;
                let output = self.handle_short_packet(packet)?;
                Ok((output, buffer))
            }
            version_negotiation_no_fixed_bit_tag!() => {
                let (version, _peek) = peek.decode()?;
                decoder_invariant!(
                    version_negotiation::VERSION == version,
                    "invalid version negotiation packet"
                );
                version_negotiation!(version)
            }
            initial_tag!() => long_packet!(ProtectedInitial, handle_initial_packet),
            zero_rtt_tag!() => long_packet!(ProtectedZeroRtt, handle_zero_rtt_packet),
            handshake_tag!() => long_packet!(ProtectedHandshake, handle_handshake_packet),
            retry_tag!() => long_packet!(ProtectedRetry, handle_retry_packet),
            _ => Err(DecoderError::InvariantViolation("invalid packet").into()),
        }
    }
}

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

    macro_rules! snapshot {
        ($name:ident) => {
            #[test]
            fn $name() {
                s2n_codec::assert_codec_round_trip_sample_file!(
                    crate::packet::ProtectedPacket,
                    concat!("src/packet/test_samples/", stringify!($name), ".bin"),
                    |buffer| {
                        let remote_address = crate::inet::ip::SocketAddress::default();
                        let connection_info =
                            crate::connection::id::ConnectionInfo::new(&remote_address);
                        crate::packet::ProtectedPacket::decode(
                            buffer,
                            &connection_info,
                            &long::DESTINATION_CONNECTION_ID_MAX_LEN,
                        )
                        .unwrap()
                    }
                );
            }
        };
    }

    snapshot!(short);
    snapshot!(initial);
    snapshot!(zero_rtt);
    snapshot!(handshake);
    snapshot!(retry);
    snapshot!(version_negotiation);
}