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

use crate::{
    connection,
    crypto::{
        retry,
        retry::{IntegrityTag, RetryKey},
        CryptoError,
    },
    inet::SocketAddress,
    packet::{
        decoding::HeaderDecoder,
        initial::ProtectedInitial,
        long::{DestinationConnectionIdLen, SourceConnectionIdLen, Version},
        Tag,
    },
    random, token,
};
use core::{convert::TryInto, mem::size_of, ops::Range};
use retry::INTEGRITY_TAG_LEN;
use s2n_codec::{
    decoder_invariant, DecoderBufferMut, DecoderBufferMutResult, Encoder, EncoderBuffer,
    EncoderValue,
};

//= https://www.rfc-editor.org/rfc/rfc9000#section-17.2.5
//# Retry Packet {
//#   Header Form (1) = 1,
//#   Fixed Bit (1) = 1,
//#   Long Packet Type (2) = 3,
//#   Unused (4),
//#   Version (32),
//#   Destination Connection ID Length (8),
//#   Destination Connection ID (0..160),
//#   Source Connection ID Length (8),
//#   Source Connection ID (0..160),
//#   Retry Token (..),
//#   Retry Integrity Tag (128),
//# }

//= https://www.rfc-editor.org/rfc/rfc9000#section-17.2.5
//# a Retry packet uses a long packet header with a type value of 0x03.
macro_rules! retry_tag {
    () => {
        0b1111u8
    };
}

//= https://www.rfc-editor.org/rfc/rfc9000#section-17.2.5
//#   Retry Token:  An opaque token that the server can use to validate the
//#      client's address.

//= https://www.rfc-editor.org/rfc/rfc9000#section-17.2.5
//#   Retry Integrity Tag:  Defined in Section 5.8 ("Retry Packet
//#      Integrity") of [QUIC-TLS].

#[derive(Debug)]
pub struct Retry<'a> {
    pub tag: Tag,
    pub version: Version,
    pub destination_connection_id: &'a [u8],
    pub source_connection_id: &'a [u8],
    pub retry_token: &'a [u8],
    pub retry_integrity_tag: &'a IntegrityTag,
}

//= https://www.rfc-editor.org/rfc/rfc9001#section-5.8
//# Retry Pseudo-Packet {
//#   ODCID Length (8),
//#   Original Destination Connection ID (0..160),
//#   Header Form (1) = 1,
//#   Fixed Bit (1) = 1,
//#   Long Packet Type (2) = 3,
//#   Unused (4),
//#   Version (32),
//#   DCID Len (8),
//#   Destination Connection ID (0..160),
//#   SCID Len (8),
//#   Source Connection ID (0..160),
//#   Retry Token (..),
//# }

#[derive(Debug)]
pub struct PseudoRetry<'a> {
    pub original_destination_connection_id: &'a [u8],
    pub tag: Tag,
    pub version: Version,
    pub destination_connection_id: &'a [u8],
    pub source_connection_id: &'a [u8],
    pub retry_token: &'a [u8],
}

impl<'a> PseudoRetry<'a> {
    pub fn new(
        odcid: &'a [u8],
        tag: Tag,
        version: Version,
        destination_connection_id: &'a [u8],
        source_connection_id: &'a [u8],
        retry_token: &'a [u8],
    ) -> Self {
        Self {
            original_destination_connection_id: odcid,
            tag,
            version,
            destination_connection_id,
            source_connection_id,
            retry_token,
        }
    }
}

pub type ProtectedRetry<'a> = Retry<'a>;
pub type EncryptedRetry<'a> = Retry<'a>;
pub type CleartextRetry<'a> = Retry<'a>;

impl<'a> Retry<'a> {
    pub fn encode_packet<T: token::Format, C: RetryKey>(
        remote_address: &SocketAddress,
        packet: &ProtectedInitial,
        local_connection_id: &connection::LocalId,
        random: &mut dyn random::Generator,
        token_format: &mut T,
        packet_buf: &mut [u8],
    ) -> Option<Range<usize>> {
        //= https://www.rfc-editor.org/rfc/rfc9000#section-17.2.5.1
        //# This value MUST NOT be equal to the Destination
        //# Connection ID field of the packet sent by the client.
        debug_assert_ne!(
            local_connection_id.as_ref(),
            packet.destination_connection_id()
        );
        if local_connection_id.as_ref() == packet.destination_connection_id() {
            return None;
        }

        let retry_packet = Retry::from_initial(packet, local_connection_id.as_ref());
        let pseudo_packet = retry_packet.pseudo_packet(packet.destination_connection_id());

        let mut buffer = EncoderBuffer::new(packet_buf);
        pseudo_packet.encode(&mut buffer);

        let destination_connection_id =
            &connection::PeerId::try_from_bytes(retry_packet.destination_connection_id).unwrap();
        let mut context = token::Context::new(remote_address, destination_connection_id, random);

        let mut outcome = None;

        buffer.write_sized(T::TOKEN_LEN, |token_buf| {
            outcome = token_format.generate_retry_token(
                &mut context,
                &connection::InitialId::try_from_bytes(packet.destination_connection_id()).unwrap(),
                token_buf,
            );
        });

        outcome?;

        let tag = C::generate_tag(buffer.as_mut_slice());
        buffer.write_slice(&tag);
        let end = buffer.len();
        let start =
            packet.destination_connection_id().len() + size_of::<DestinationConnectionIdLen>();

        Some(start..end)
    }

    pub fn validate<Crypto, CreateBuf, Buf>(
        &self,
        odcid: &connection::InitialId,
        create_buf: CreateBuf,
    ) -> Result<(), CryptoError>
    where
        Crypto: RetryKey,
        CreateBuf: FnOnce(usize) -> Buf,
        Buf: AsMut<[u8]>,
    {
        let pseudo_packet = self.pseudo_packet(odcid.as_ref());
        let len = pseudo_packet.encoding_size();
        let mut buf = create_buf(len);
        let buf = buf.as_mut();

        let mut buffer = EncoderBuffer::new(buf);
        pseudo_packet.encode(&mut buffer);

        //= https://www.rfc-editor.org/rfc/rfc9001#section-5.8
        //# Retry packets (see Section 17.2.5 of [QUIC-TRANSPORT]) carry a Retry
        //# Integrity Tag that provides two properties: it allows the discarding
        //# of packets that have accidentally been corrupted by the network, and
        //# only an entity that observes an Initial packet can send a valid Retry
        //# packet.
        Crypto::validate(buf, *self.retry_integrity_tag)?;

        Ok(())
    }

    pub fn from_initial(
        initial_packet: &'a ProtectedInitial,
        local_connection_id: &'a [u8],
    ) -> Self {
        // The destination and source connection IDs are flipped because this packet is being sent
        // back to the client.

        Self {
            //= https://www.rfc-editor.org/rfc/rfc9000#section-17.2.5
            //# The value in the Unused field is set to an arbitrary value
            //# by the server; a client MUST ignore these bits.
            // The last 4 bits are unused. They are set to 0x0f here to allow easy testing with
            // example packets provided in the RFC.
            // https://www.rfc-editor.org/rfc/rfc9001#section-A.2
            tag: (retry_tag!() << 4) | 0x0f,
            version: initial_packet.version,
            destination_connection_id: initial_packet.source_connection_id(),
            source_connection_id: local_connection_id,
            retry_token: &[][..],
            retry_integrity_tag: {
                static EMPTY_TAG: IntegrityTag = [0u8; INTEGRITY_TAG_LEN];
                &EMPTY_TAG
            },
        }
    }

    #[inline]
    pub(crate) fn decode(
        tag: Tag,
        version: Version,
        buffer: DecoderBufferMut,
    ) -> DecoderBufferMutResult<Retry> {
        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)?;

        // split header and payload
        let header_len = decoder.decoded_len();
        let (header, buffer) = buffer.decode_slice(header_len)?;
        let header: &[u8] = header.into_less_safe_slice();

        // read borrowed slices
        let destination_connection_id = destination_connection_id.get(header);
        let source_connection_id = source_connection_id.get(header);

        let buffer_len = buffer.len().saturating_sub(retry::INTEGRITY_TAG_LEN);

        //= https://www.rfc-editor.org/rfc/rfc9000#section-17.2.5.2
        //# A client MUST discard a Retry packet with a zero-length
        //# Retry Token field.
        decoder_invariant!(buffer_len > 0, "Token cannot be empty");

        let (retry_token, buffer) = buffer.decode_slice(buffer_len)?;
        let retry_token: &[u8] = retry_token.into_less_safe_slice();

        let (retry_integrity_tag, buffer) = buffer.decode_slice(retry::INTEGRITY_TAG_LEN)?;
        let retry_integrity_tag: &[u8] = retry_integrity_tag.into_less_safe_slice();
        let retry_integrity_tag: &IntegrityTag = retry_integrity_tag
            .try_into()
            .expect("tag length already checked");

        let packet = Retry {
            tag,
            version,
            destination_connection_id,
            source_connection_id,
            retry_token,
            retry_integrity_tag,
        };

        Ok((packet, buffer))
    }

    #[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 retry_token(&self) -> &[u8] {
        self.retry_token
    }

    #[inline]
    fn pseudo_packet(&self, odcid: &'a [u8]) -> PseudoRetry {
        PseudoRetry {
            original_destination_connection_id: odcid,
            tag: self.tag,
            version: self.version,
            destination_connection_id: self.destination_connection_id,
            source_connection_id: self.source_connection_id,
            retry_token: self.retry_token,
        }
    }
}

impl<'a> EncoderValue for Retry<'a> {
    fn encode<E: Encoder>(&self, encoder: &mut E) {
        let tag: u8 = self.tag;
        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.retry_token.encode(encoder);
        self.retry_integrity_tag.as_ref().encode(encoder);
    }
}

impl<'a> EncoderValue for PseudoRetry<'a> {
    fn encode<E: Encoder>(&self, encoder: &mut E) {
        self.original_destination_connection_id
            .encode_with_len_prefix::<DestinationConnectionIdLen, E>(encoder);

        self.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.retry_token.encode(encoder);
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::{crypto::retry, inet, packet};
    use s2n_codec::EncoderBuffer;

    #[test]
    fn test_encode() {
        let packet = Retry {
            tag: (retry_tag!() << 4) | 0x0f,
            destination_connection_id: &retry::example::DCID,
            source_connection_id: &retry::example::SCID,
            retry_token: &retry::example::TOKEN,
            retry_integrity_tag: &retry::example::EXPECTED_TAG,
            version: retry::example::VERSION,
        };
        let mut buf = [0; retry::example::PACKET_LEN];
        let mut encoder = EncoderBuffer::new(&mut buf);
        packet.encode(&mut encoder);
        assert_eq!(retry::example::PACKET[..], buf[..]);
    }

    #[test]
    fn test_decode() {
        let mut buf = retry::example::PACKET;
        let decoder = DecoderBufferMut::new(&mut buf);
        let remote_address = inet::ip::SocketAddress::default();
        let connection_info = connection::id::ConnectionInfo::new(&remote_address);
        let (packet, _) = packet::ProtectedPacket::decode(decoder, &connection_info, &20).unwrap();
        let packet = match packet {
            packet::ProtectedPacket::Retry(retry) => retry,
            _ => panic!("expected retry packet type"),
        };

        //= https://www.rfc-editor.org/rfc/rfc9000#section-8.1.4
        //= type=test
        //# For this design to work,
        //# the token MUST be covered by integrity protection against
        //# modification or falsification by clients.
        assert_eq!(packet.retry_integrity_tag, &retry::example::EXPECTED_TAG);
        assert_eq!(packet.retry_token, retry::example::TOKEN);
        assert_eq!(packet.source_connection_id, retry::example::SCID);
        assert_eq!(packet.destination_connection_id, retry::example::DCID);
        assert_eq!(packet.version, retry::example::VERSION);
    }

    #[test]
    fn test_decode_no_token() {
        //= https://www.rfc-editor.org/rfc/rfc9000#section-17.2.5.2
        //= type=test
        //# A client MUST discard a Retry packet with a zero-length
        //# Retry Token field.
        let mut buf = retry::example::INVALID_PACKET_NO_TOKEN;
        let decoder = DecoderBufferMut::new(&mut buf);
        let remote_address = inet::ip::SocketAddress::default();
        let connection_info = connection::id::ConnectionInfo::new(&remote_address);
        assert!(packet::ProtectedPacket::decode(decoder, &connection_info, &20).is_err());
    }

    #[test]
    fn test_pseudo_decode() {
        let mut buf = retry::example::PACKET;
        let decoder = DecoderBufferMut::new(&mut buf);
        let remote_address = inet::ip::SocketAddress::default();
        let connection_info = connection::id::ConnectionInfo::new(&remote_address);
        let (packet, _) = packet::ProtectedPacket::decode(decoder, &connection_info, &20).unwrap();
        let packet = match packet {
            packet::ProtectedPacket::Retry(retry) => retry,
            _ => panic!("expected retry packet type"),
        };
        let pseudo_packet = packet.pseudo_packet(&retry::example::ODCID);

        assert_eq!(pseudo_packet.retry_token, retry::example::TOKEN);
        assert_eq!(pseudo_packet.source_connection_id, retry::example::SCID);
        assert_eq!(
            pseudo_packet.destination_connection_id,
            retry::example::DCID
        );
        assert_eq!(pseudo_packet.version, retry::example::VERSION);
        assert_eq!(
            pseudo_packet.original_destination_connection_id,
            retry::example::ODCID
        );

        let length = pseudo_packet.encoding_size();
        let mut pseudo_scratch: Vec<u8> = vec![0; length];
        let mut encoder = EncoderBuffer::new(&mut pseudo_scratch);
        pseudo_packet.encode(&mut encoder);

        assert_eq!(pseudo_scratch, retry::example::PSEUDO_PACKET);
    }
}


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

use crate::{
    connection,
    connection::id::ConnectionInfo,
    crypto::{CryptoError, EncryptedPayload, OneRttHeaderKey, OneRttKey, ProtectedPayload},
    packet::{
        decoding::HeaderDecoder,
        encoding::{PacketEncoder, PacketPayloadEncoder},
        number::{
            PacketNumber, PacketNumberLen, PacketNumberSpace, ProtectedPacketNumber,
            TruncatedPacketNumber,
        },
        KeyPhase, ProtectedKeyPhase, Tag,
    },
    transport,
};
use s2n_codec::{CheckedRange, DecoderBufferMut, DecoderBufferMutResult, Encoder, EncoderValue};

//= https://www.rfc-editor.org/rfc/rfc9000#section-17.3.1
//# 1-RTT Packet {
//#   Header Form (1) = 0,
//#   Fixed Bit (1) = 1,
//#   Spin Bit (1),
//#   Reserved Bits (2),
//#   Key Phase (1),
//#   Packet Number Length (2),
//#   Destination Connection ID (0..160),
//#   Packet Number (8..32),
//#   Packet Payload (..),
//# }

//= https://www.rfc-editor.org/rfc/rfc9000#section-17.3.1
//# Header Form:  The most significant bit (0x80) of byte 0 is set to 0
//#    for the short header.
//#
//# Fixed Bit:  The next bit (0x40) of byte 0 is set to 1.

macro_rules! short_tag {
    () => {
        0b0100u8..=0b0111u8
    };
}

const ENCODING_TAG: u8 = 0b0100_0000;

//= https://www.rfc-editor.org/rfc/rfc9000#section-17.3.1
//# Spin Bit:  The third most significant bit (0x20) of byte 0 is the
//#    latency spin bit, set as described in Section 17.4.

const SPIN_BIT_MASK: u8 = 0x20;

//= https://www.rfc-editor.org/rfc/rfc9000#section-17.3.1
//#  Reserved Bits:  The next two bits (those with a mask of 0x18) of byte
//#      0 are reserved.

const RESERVED_BITS_MASK: u8 = 0x18;

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

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

impl SpinBit {
    fn from_tag(tag: Tag) -> Self {
        if tag & SPIN_BIT_MASK == SPIN_BIT_MASK {
            Self::One
        } else {
            Self::Zero
        }
    }

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

//= https://www.rfc-editor.org/rfc/rfc9000#section-17.3.1
//# Reserved Bits:  The next two bits (those with a mask of 0x18) of byte
//#    0 are reserved.  These bits are protected using header protection;
//#    see Section 5.4 of [QUIC-TLS].

//= https://www.rfc-editor.org/rfc/rfc9000#section-17.3.1
//# Packet Number Length:  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.3.1
//# Destination Connection ID:  The Destination Connection ID is a
//#    connection ID that is chosen by the intended recipient of the
//#    packet.

//= https://www.rfc-editor.org/rfc/rfc9000#section-17.3.1
//# Packet Number:  The Packet Number 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 Packet Number Length field.  See Section 17.1 for
//#    details.

//= https://www.rfc-editor.org/rfc/rfc9000#section-17.3.1
//# Packet Payload:  1-RTT packets always include a 1-RTT protected
//#    payload.

#[derive(Debug)]
pub struct Short<DCID, KeyPhase, PacketNumber, Payload> {
    pub spin_bit: SpinBit,
    pub key_phase: KeyPhase,
    pub destination_connection_id: DCID,
    pub packet_number: PacketNumber,
    pub payload: Payload,
}

pub type ProtectedShort<'a> =
    Short<CheckedRange, ProtectedKeyPhase, ProtectedPacketNumber, ProtectedPayload<'a>>;
pub type EncryptedShort<'a> = Short<CheckedRange, KeyPhase, PacketNumber, EncryptedPayload<'a>>;
pub type CleartextShort<'a> = Short<&'a [u8], KeyPhase, PacketNumber, DecoderBufferMut<'a>>;

impl<'a> ProtectedShort<'a> {
    #[inline]
    pub(crate) fn decode<Validator: connection::id::Validator>(
        tag: Tag,
        buffer: DecoderBufferMut<'a>,
        connection_info: &ConnectionInfo,
        destination_connection_id_decoder: &Validator,
    ) -> DecoderBufferMutResult<'a, ProtectedShort<'a>> {
        let mut decoder = HeaderDecoder::new_short(&buffer);

        let spin_bit = SpinBit::from_tag(tag);
        let key_phase = ProtectedKeyPhase;

        let destination_connection_id = decoder.decode_short_destination_connection_id(
            &buffer,
            connection_info,
            destination_connection_id_decoder,
        )?;

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

        let packet = Short {
            spin_bit,
            key_phase,
            destination_connection_id,
            packet_number,
            payload,
        };

        Ok((packet, remaining))
    }

    pub fn unprotect<H: OneRttHeaderKey>(
        self,
        header_key: &H,
        largest_acknowledged_packet_number: PacketNumber,
    ) -> Result<EncryptedShort<'a>, CryptoError> {
        let Short {
            spin_bit,
            destination_connection_id,
            payload,
            ..
        } = self;

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

        let key_phase = KeyPhase::from_tag(payload.get_tag());

        let packet_number = truncated_packet_number.expand(largest_acknowledged_packet_number);

        Ok(Short {
            spin_bit,
            key_phase,
            destination_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()
    }
}

impl<'a> EncryptedShort<'a> {
    pub fn decrypt<C: OneRttKey>(self, crypto: &C) -> Result<CleartextShort<'a>, transport::Error> {
        let Short {
            spin_bit,
            key_phase,
            destination_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.3.1
        //# 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] & 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);

        Ok(Short {
            spin_bit,
            key_phase,
            destination_connection_id,
            packet_number,
            payload,
        })
    }

    #[inline]
    pub fn key_phase(&self) -> KeyPhase {
        self.key_phase
    }

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

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

impl<DCID: EncoderValue, Payload: EncoderValue> EncoderValue
    for Short<DCID, KeyPhase, TruncatedPacketNumber, Payload>
{
    #[inline]
    fn encode<E: Encoder>(&self, encoder: &mut E) {
        self.encode_header(self.packet_number.len(), encoder);
        self.packet_number.encode(encoder);
        self.payload.encode(encoder);
    }
}

impl<DCID: EncoderValue, PacketNumber, Payload> Short<DCID, KeyPhase, PacketNumber, Payload> {
    #[inline]
    fn encode_header<E: Encoder>(&self, packet_number_len: PacketNumberLen, encoder: &mut E) {
        (ENCODING_TAG
            | self.spin_bit.into_packet_tag_mask()
            | self.key_phase.into_packet_tag_mask()
            | packet_number_len.into_packet_tag_mask())
        .encode(encoder);

        self.destination_connection_id.encode(encoder);
    }
}

impl<DCID: EncoderValue, Payload: PacketPayloadEncoder, K: OneRttKey, H: OneRttHeaderKey>
    PacketEncoder<K, H, Payload> for Short<DCID, KeyPhase, PacketNumber, Payload>
{
    type PayloadLenCursor = ();

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

    #[inline]
    fn encode_header<E: Encoder>(&self, packet_number_len: PacketNumberLen, encoder: &mut E) {
        Short::encode_header(self, packet_number_len, encoder);
    }

    #[inline]
    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::{
    connection,
    packet::{number::PacketNumberLen, Tag},
    random, stateless_reset,
};

//= https://www.rfc-editor.org/rfc/rfc9000#section-10.3
//# Stateless Reset {
//#   Fixed Bits (2) = 1,
//#   Unpredictable Bits (38..),
//#   Stateless Reset Token (128),
//# }

//= https://www.rfc-editor.org/rfc/rfc9000#section-10.3
//# Endpoints MUST send Stateless Resets formatted as a packet
//# with a short header.
const TAG: u8 = 0b0100_0000;
const TAG_OFFSET: u8 = 2;

// This value represents the minimum packet size of a packet that will be indistinguishable from
// valid QUIC version 1 packets, including one byte for the minimal payload and excluding the
// authentication tag (which may be variable and should be added to this constant). Since the
// connection ID length is either determined by a provider or by the peer, connection::id::MAX_LEN
// is used to ensure this value is valid no matter what length connection ID is used.
const MIN_INDISTINGUISHABLE_PACKET_LEN_WITHOUT_TAG: usize =
    core::mem::size_of::<Tag>() + PacketNumberLen::MAX_LEN + connection::id::MAX_LEN + 1;

/// Calculates the minimum packet length required such that a packet is indistinguishable from
/// other valid QUIC version 1 packets.
pub fn min_indistinguishable_packet_len(max_tag_len: usize) -> usize {
    MIN_INDISTINGUISHABLE_PACKET_LEN_WITHOUT_TAG + max_tag_len
}

/// Encodes a stateless reset packet into the given packet buffer.
pub fn encode_packet(
    token: stateless_reset::Token,
    max_tag_len: usize,
    triggering_packet_len: usize,
    random_generator: &mut dyn random::Generator,
    packet_buf: &mut [u8],
) -> Option<usize> {
    //= https://www.rfc-editor.org/rfc/rfc9000#section-10.3
    //# These values assume that the stateless reset token is the same length
    //# as the minimum expansion of the packet protection AEAD.  Additional
    //# unpredictable bytes are necessary if the endpoint could have
    //# negotiated a packet protection scheme with a larger minimum
    //# expansion.
    // The tag length for all cipher suites defined in TLS 1.3 is 16 bytes, but
    // we will calculate based on a given max tag length to allow for future cipher
    // suites with larger tags.
    let min_len = min_indistinguishable_packet_len(max_tag_len);

    //= https://www.rfc-editor.org/rfc/rfc9000#section-10.3
    //# An endpoint MUST NOT send a Stateless Reset that is three times or
    //# more larger than the packet it receives to avoid being used for
    //# amplification.

    //= https://www.rfc-editor.org/rfc/rfc9000#section-10.3.3
    //# An endpoint MUST ensure that every Stateless Reset that it sends is
    //# smaller than the packet that triggered it, unless it maintains state
    //# sufficient to prevent looping.
    let max_len = triggering_packet_len
        .saturating_sub(1)
        .min(packet_buf.len());

    // The packet that triggered this stateless reset was too small to send a stateless reset
    // that would be indistinguishable from a valid short header packet, so we'll just drop the
    // packet instead of sending a stateless reset.
    if max_len < min_len {
        return None;
    }

    // Generate unpredictable bits, leaving room for the stateless reset token
    let unpredictable_bits_min_len = min_len - stateless_reset::token::LEN;
    let unpredictable_bits_max_len = max_len - stateless_reset::token::LEN;

    let unpredictable_bits_len = generate_unpredictable_bits(
        random_generator,
        unpredictable_bits_min_len,
        &mut packet_buf[..unpredictable_bits_max_len],
    );
    // Write the short header tag over the first two bits
    packet_buf[0] = packet_buf[0] >> TAG_OFFSET | TAG;

    let packet_len = unpredictable_bits_len + stateless_reset::token::LEN;

    packet_buf[unpredictable_bits_len..packet_len].copy_from_slice(token.as_ref());

    if cfg!(debug_assertions) {
        assert!(packet_len >= min_len);
        assert!(packet_len <= max_len);
        assert!(packet_len < triggering_packet_len);
    }

    Some(packet_len)
}

/// Fills the given buffer with a random amount of random data at least of the
/// given `min_len`. Returns the length of the unpredictable bits that were generated.
fn generate_unpredictable_bits(
    random_generator: &mut dyn random::Generator,
    min_len: usize,
    buffer: &mut [u8],
) -> usize {
    // Generate a random amount of unpredictable bits within the valid range
    // to further decrease the likelihood a stateless reset could be distinguished
    // from a valid packet. This will have slight bias towards the lower end of the range,
    // but this bias does not result in any reduction in security for this usage and is actually
    // welcome as it results in reaching the minimal stateless reset size and thus
    // existing stateless reset loops sooner.
    let len = random::gen_range_biased(random_generator, min_len..=buffer.len());

    //= https://www.rfc-editor.org/rfc/rfc9000#section-10.3
    //# The remainder of the first byte
    //# and an arbitrary number of bytes following it are set to values that
    //# SHOULD be indistinguishable from random.
    random_generator.public_random_fill(&mut buffer[..len]);

    len
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::{path::MINIMUM_MTU, stateless_reset::token::testing::TEST_TOKEN_1};

    #[test]
    #[cfg_attr(miri, ignore)] // This test is too expensive for miri to complete in a reasonable amount of time
    fn generate_unpredictable_bits_test() {
        bolero::check!()
            .with_type::<(u8, u16, u16)>()
            .cloned()
            .for_each(|(seed, mut min, mut max)| {
                if min > max {
                    core::mem::swap(&mut min, &mut max);
                }
                let mut generator = random::testing::Generator(seed);
                let mut buffer = vec![0; max.into()];
                let len = generate_unpredictable_bits(&mut generator, min.into(), &mut buffer);
                assert!(len >= min.into());
                assert!(len <= max.into());
            });
    }

    //= https://www.rfc-editor.org/rfc/rfc9000#section-10.3
    //= type=test
    //# The remainder of the first byte
    //# and an arbitrary number of bytes following it are set to values that
    //# SHOULD be indistinguishable from random.
    #[test]
    fn unpredictable_bits_are_indistinguishable_from_random() {
        const MIN_LEN: usize = 100;
        const MAX_LEN: usize = 1000;

        let mut generator = random::testing::Generator(123);
        let mut buffer = [0; MAX_LEN];
        let mut buffer_2 = [0; MAX_LEN];
        generate_unpredictable_bits(&mut generator, MIN_LEN, &mut buffer);
        generate_unpredictable_bits(&mut generator, MIN_LEN, &mut buffer_2);

        assert_ne!(buffer[0..32], buffer_2[0..32]);
    }

    #[test]
    fn encode_packet_test() {
        let max_tag_len = 16;
        let triggering_packet_len = 600;
        let mut generator = random::testing::Generator(123);

        let mut buffer = [0; MINIMUM_MTU as usize];

        let packet_len = encode_packet(
            TEST_TOKEN_1,
            max_tag_len,
            triggering_packet_len,
            &mut generator,
            &mut buffer,
        )
        .unwrap();

        //= https://www.rfc-editor.org/rfc/rfc9000#section-10.3
        //= type=test
        //# An endpoint MUST NOT send a Stateless Reset that is three times or
        //# more larger than the packet it receives to avoid being used for
        //# amplification.

        //= https://www.rfc-editor.org/rfc/rfc9000#section-10.3.3
        //= type=test
        //# An endpoint MUST ensure that every Stateless Reset that it sends is
        //# smaller than the packet that triggered it, unless it maintains state
        //# sufficient to prevent looping.
        assert!(packet_len < triggering_packet_len);

        //= https://www.rfc-editor.org/rfc/rfc9000#section-10.3
        //= type=test
        //# Endpoints MUST send Stateless Resets formatted as a packet
        //# with a short header.
        assert!(matches!(&buffer[0] >> 4, short_tag!()));

        assert_eq!(
            TEST_TOKEN_1.into_inner(),
            buffer[packet_len - stateless_reset::token::LEN..packet_len]
        );
    }

    #[test]
    fn min_packet_test() {
        let max_tag_len = 16;
        let mut triggering_packet_len = min_indistinguishable_packet_len(max_tag_len) + 1;
        let mut generator = random::testing::Generator(123);
        let mut buffer = [0; MINIMUM_MTU as usize];

        let packet_len = encode_packet(
            TEST_TOKEN_1,
            max_tag_len,
            triggering_packet_len,
            &mut generator,
            &mut buffer,
        );

        assert_eq!(packet_len, Some(triggering_packet_len - 1));

        triggering_packet_len -= 1;

        let packet_len = encode_packet(
            TEST_TOKEN_1,
            max_tag_len,
            triggering_packet_len,
            &mut generator,
            &mut buffer,
        );

        assert!(packet_len.is_none());

        triggering_packet_len = 0;

        let packet_len = encode_packet(
            TEST_TOKEN_1,
            max_tag_len,
            triggering_packet_len,
            &mut generator,
            &mut buffer,
        );

        assert!(packet_len.is_none());
    }

    #[test]
    fn max_packet_test() {
        let max_tag_len = 16;
        let triggering_packet_len = (MINIMUM_MTU * 2) as usize;
        let mut generator = random::testing::Generator(123);
        let mut buffer = [0; MINIMUM_MTU as usize];

        let packet_len = encode_packet(
            TEST_TOKEN_1,
            max_tag_len,
            triggering_packet_len,
            &mut generator,
            &mut buffer,
        );

        assert!(packet_len.is_some());

        assert!(packet_len.unwrap() <= MINIMUM_MTU as usize);
    }

    #[test]
    #[cfg_attr(miri, ignore)] // This test breaks in CI but can't be reproduced locally - https://github.com/aws/s2n-quic/issues/867
    fn packet_encoding_test() {
        let mut buffer = [0; MINIMUM_MTU as usize];

        bolero::check!()
            .with_type::<(u8, usize, u16)>()
            .cloned()
            .for_each(|(seed, triggering_packet_len, max_tag_len)| {
                let mut generator = random::testing::Generator(seed);
                let packet_len = encode_packet(
                    TEST_TOKEN_1,
                    max_tag_len.into(),
                    triggering_packet_len,
                    &mut generator,
                    &mut buffer,
                );

                let min_len = MIN_INDISTINGUISHABLE_PACKET_LEN_WITHOUT_TAG + max_tag_len as usize;
                let max_len = triggering_packet_len.saturating_sub(1).min(buffer.len());

                if min_len <= max_len {
                    assert!(packet_len.is_some());
                    let packet_len = packet_len.unwrap();
                    assert!(packet_len <= max_len);
                    assert!(packet_len >= min_len);

                    //= https://www.rfc-editor.org/rfc/rfc9000#section-10.3
                    //= type=test
                    //# Endpoints MUST send Stateless Resets formatted as a packet
                    //# with a short header.
                    assert!(matches!(&buffer[0] >> 4, short_tag!()));

                    assert_eq!(
                        TEST_TOKEN_1.into_inner(),
                        buffer[packet_len - stateless_reset::token::LEN..packet_len]
                    );
                } else {
                    assert!(packet_len.is_none());
                }
            })
    }
}


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

use crate::{
    connection::id::ConnectionInfo,
    crypto::key::testing,
    inet::SocketAddress,
    packet::{
        encoding::PacketEncoder, number::PacketNumberSpace, CleartextPacket, ProtectedPacket,
    },
    transport,
};
use bolero::check;
use s2n_codec::{DecoderBufferMut, Encoder, EncoderBuffer};

#[test]
fn round_trip() {
    let mut encoder_data = vec![];
    check!().for_each(move |data| {
        let mut data = data.to_vec();
        // add a few bytes to the end for padding
        encoder_data.resize(data.len() * 2, 0);

        let mut decoder_buffer = DecoderBufferMut::new(&mut data);
        let mut encoder_buffer = EncoderBuffer::new(&mut encoder_data);
        let remote_address = SocketAddress::default();
        let connection_info = ConnectionInfo::new(&remote_address);

        while let Ok((packet, remaining)) =
            ProtectedPacket::decode(decoder_buffer, &connection_info, &20)
        {
            if let Ok(cleartext_packet) = decrypt_packet(packet) {
                encoder_buffer = encode_packet(cleartext_packet, encoder_buffer);
            }
            decoder_buffer = remaining;
        }
    });
}

fn decrypt_packet(packet: ProtectedPacket) -> Result<CleartextPacket, transport::Error> {
    use ProtectedPacket::*;
    match packet {
        Handshake(packet) => {
            let dcid = packet.destination_connection_id().to_vec();
            let scid = packet.source_connection_id().to_vec();

            let packet = packet.unprotect(
                &testing::HeaderKey::new(),
                PacketNumberSpace::Handshake.new_packet_number(Default::default()),
            )?;
            assert_eq!(dcid, packet.destination_connection_id());
            assert_eq!(scid, packet.source_connection_id());

            let packet = packet.decrypt(&testing::Key::new())?;
            assert_eq!(dcid, packet.destination_connection_id());
            assert_eq!(scid, packet.source_connection_id());

            Ok(CleartextPacket::Handshake(packet))
        }
        Initial(packet) => {
            let dcid = packet.destination_connection_id().to_vec();
            let scid = packet.source_connection_id().to_vec();
            let token = packet.token().to_vec();

            let packet = packet.unprotect(
                &testing::HeaderKey::new(),
                PacketNumberSpace::Initial.new_packet_number(Default::default()),
            )?;
            assert_eq!(dcid, packet.destination_connection_id());
            assert_eq!(scid, packet.source_connection_id());
            assert_eq!(token, packet.token());

            let packet = packet.decrypt(&testing::Key::new())?;
            assert_eq!(dcid, packet.destination_connection_id());
            assert_eq!(scid, packet.source_connection_id());
            assert_eq!(token, packet.token());

            Ok(CleartextPacket::Initial(packet))
        }
        Retry(packet) => {
            let _ = packet.destination_connection_id();
            let _ = packet.source_connection_id();

            Ok(CleartextPacket::Retry(packet))
        }
        Short(packet) => {
            let dcid = packet.destination_connection_id().to_vec();

            let packet = packet.unprotect(
                &testing::HeaderKey::new(),
                PacketNumberSpace::ApplicationData.new_packet_number(Default::default()),
            )?;
            assert_eq!(dcid, packet.destination_connection_id());

            let packet = packet.decrypt(&testing::Key::new())?;
            assert_eq!(dcid, packet.destination_connection_id());

            Ok(CleartextPacket::Short(packet))
        }
        ZeroRtt(packet) => {
            let dcid = packet.destination_connection_id().to_vec();
            let scid = packet.source_connection_id().to_vec();

            let packet = packet.unprotect(
                &testing::HeaderKey::new(),
                PacketNumberSpace::ApplicationData.new_packet_number(Default::default()),
            )?;
            assert_eq!(dcid, packet.destination_connection_id());
            assert_eq!(scid, packet.source_connection_id());

            let packet = packet.decrypt(&testing::Key::new())?;
            assert_eq!(dcid, packet.destination_connection_id());
            assert_eq!(scid, packet.source_connection_id());

            Ok(CleartextPacket::ZeroRtt(packet))
        }
        VersionNegotiation(packet) => {
            let _: Vec<_> = packet.iter().collect();

            Ok(CleartextPacket::VersionNegotiation(packet))
        }
    }
}

fn encode_packet<'a>(packet: CleartextPacket, mut encoder: EncoderBuffer<'a>) -> EncoderBuffer<'a> {
    use CleartextPacket::*;
    let result = match packet {
        Handshake(packet) => packet.encode_packet(
            &testing::Key::new(),
            &testing::HeaderKey::new(),
            PacketNumberSpace::Handshake.new_packet_number(Default::default()),
            None,
            encoder,
        ),
        Initial(packet) => packet.encode_packet(
            &testing::Key::new(),
            &testing::HeaderKey::new(),
            PacketNumberSpace::Initial.new_packet_number(Default::default()),
            None,
            encoder,
        ),
        Retry(packet) => {
            encoder.encode(&packet);
            return encoder;
        }
        Short(packet) => packet.encode_packet(
            &testing::Key::new(),
            &testing::HeaderKey::new(),
            PacketNumberSpace::ApplicationData.new_packet_number(Default::default()),
            None,
            encoder,
        ),
        ZeroRtt(packet) => packet.encode_packet(
            &testing::Key::new(),
            &testing::HeaderKey::new(),
            PacketNumberSpace::ApplicationData.new_packet_number(Default::default()),
            None,
            encoder,
        ),
        VersionNegotiation(packet) => {
            encoder.encode(&packet);
            return encoder;
        }
    };

    match result {
        Ok((_, encoder)) => encoder,
        Err(err) => err.take_buffer(),
    }
}


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

use crate::packet::{
    initial::ProtectedInitial,
    long::{
        validate_destination_connection_id_len, validate_source_connection_id_len,
        DestinationConnectionIdLen, SourceConnectionIdLen, Version,
    },
    Tag,
};
use core::mem::size_of;
use s2n_codec::{
    decoder_invariant, DecoderBuffer, DecoderBufferMut, DecoderBufferMutResult, Encoder,
    EncoderValue,
};

//= https://www.rfc-editor.org/rfc/rfc9000#section-17.2.1
//# Version Negotiation Packet {
//#   Header Form (1) = 1,
//#   Unused (7),
//#   Version (32) = 0,
//#   Destination Connection ID Length (8),
//#   Destination Connection ID (0..2040),
//#   Source Connection ID Length (8),
//#   Source Connection ID (0..2040),
//#   Supported Version (32) ...,
//# }

//= https://www.rfc-editor.org/rfc/rfc9000#section-17.2.1
//# The value in the Unused field is set to an arbitrary value by the
//# server.  Clients MUST ignore the value of this field.  Where QUIC
//# might be multiplexed with other protocols (see [RFC7983]), servers
//# SHOULD set the most significant bit of this field (0x40) to 1 so that
//# Version Negotiation packets appear to have the Fixed Bit field.

macro_rules! version_negotiation_no_fixed_bit_tag {
    () => {
        0b1000u8..=0b1011u8
    };
}

const ENCODING_TAG: u8 = 0b1100_0000;

//= https://www.rfc-editor.org/rfc/rfc9000#section-17.2.1
//# The Version field of a Version Negotiation packet MUST be set to
//# 0x00000000.

pub(crate) const VERSION: u32 = 0x0000_0000;

#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)]
pub struct VersionNegotiation<'a, SupportedVersions> {
    pub tag: Tag,
    pub destination_connection_id: &'a [u8],
    pub source_connection_id: &'a [u8],
    pub supported_versions: SupportedVersions,
}

pub type ProtectedVersionNegotiation<'a> = VersionNegotiation<'a, &'a [u8]>;
pub type EncryptedVersionNegotiation<'a> = VersionNegotiation<'a, &'a [u8]>;
pub type CleartextVersionNegotiation<'a> = VersionNegotiation<'a, &'a [u8]>;

impl<'a> ProtectedVersionNegotiation<'a> {
    #[inline]
    pub fn decode(
        tag: Tag,
        _version: Version,
        buffer: DecoderBufferMut,
    ) -> DecoderBufferMutResult<VersionNegotiation<&[u8]>> {
        let buffer = buffer
            .skip(size_of::<Tag>() + size_of::<Version>())
            .expect("tag and version already verified");

        let (destination_connection_id, buffer) =
            buffer.decode_slice_with_len_prefix::<DestinationConnectionIdLen>()?;
        let destination_connection_id = destination_connection_id.into_less_safe_slice();
        validate_destination_connection_id_len(destination_connection_id.len())?;

        let (source_connection_id, buffer) =
            buffer.decode_slice_with_len_prefix::<SourceConnectionIdLen>()?;
        let source_connection_id = source_connection_id.into_less_safe_slice();
        validate_source_connection_id_len(source_connection_id.len())?;

        let (supported_versions, buffer) = buffer.decode::<DecoderBufferMut>()?;
        let supported_versions: &[u8] = supported_versions.into_less_safe_slice();

        // ensure at least one version is present
        decoder_invariant!(
            supported_versions.len() >= size_of::<u32>(),
            "missing at least one version"
        );

        // ensure payload length can successfully decode a list of u32s
        decoder_invariant!(
            supported_versions.len() % size_of::<u32>() == 0,
            "invalid payload length"
        );

        let packet = VersionNegotiation {
            tag,
            destination_connection_id,
            source_connection_id,
            supported_versions,
        };

        Ok((packet, buffer))
    }

    #[inline]
    pub fn iter(&'a self) -> VersionNegotiationIterator<'a> {
        self.into_iter()
    }

    #[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<'a, SupportedVersions: EncoderValue> VersionNegotiation<'a, SupportedVersions> {
    pub fn from_initial(
        initial_packet: &'a ProtectedInitial,
        supported_versions: SupportedVersions,
    ) -> Self {
        // The destination and source connection IDs are flipped because this packet is being sent
        // back to the client.
        Self {
            tag: 0,
            //= https://www.rfc-editor.org/rfc/rfc9000#section-17.2.1
            //# The server MUST include the value from the Source Connection ID field
            //# of the packet it receives in the Destination Connection ID field.
            destination_connection_id: initial_packet.source_connection_id(),
            //= https://www.rfc-editor.org/rfc/rfc9000#section-17.2.1
            //# The value for Source Connection ID MUST be copied from the
            //# Destination Connection ID of the received packet, which is initially
            //# randomly selected by a client.
            source_connection_id: initial_packet.destination_connection_id(),
            supported_versions,
        }
    }
}

impl<'a> IntoIterator for ProtectedVersionNegotiation<'a> {
    type IntoIter = VersionNegotiationIterator<'a>;
    type Item = u32;

    fn into_iter(self) -> Self::IntoIter {
        VersionNegotiationIterator(DecoderBuffer::new(self.supported_versions))
    }
}

#[derive(Clone, Copy, Debug)]
pub struct VersionNegotiationIterator<'a>(DecoderBuffer<'a>);

impl<'a> Iterator for VersionNegotiationIterator<'a> {
    type Item = u32;

    fn next(&mut self) -> Option<Self::Item> {
        if let Ok((value, buffer)) = self.0.decode() {
            self.0 = buffer;
            Some(value)
        } else {
            None
        }
    }
}

impl<'a, SupportedVersions: EncoderValue> EncoderValue
    for VersionNegotiation<'a, SupportedVersions>
{
    fn encode<E: Encoder>(&self, encoder: &mut E) {
        (self.tag | ENCODING_TAG).encode(encoder);
        VERSION.encode(encoder);
        self.destination_connection_id
            .encode_with_len_prefix::<DestinationConnectionIdLen, _>(encoder);
        self.source_connection_id
            .encode_with_len_prefix::<SourceConnectionIdLen, _>(encoder);
        self.supported_versions.encode(encoder);
    }
}


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

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

//= https://www.rfc-editor.org/rfc/rfc9000#section-17.2.3
//# 0-RTT Packet {
//#   Header Form (1) = 1,
//#   Fixed Bit (1) = 1,
//#   Long Packet Type (2) = 1,
//#   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.3
//# A 0-RTT packet uses long headers with a type value of 0x1,
macro_rules! zero_rtt_tag {
    () => {
        0b1101u8
    };
}

#[derive(Debug)]
pub struct ZeroRtt<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 ProtectedZeroRtt<'a> =
    ZeroRtt<CheckedRange, CheckedRange, ProtectedPacketNumber, ProtectedPayload<'a>>;
pub type EncryptedZeroRtt<'a> =
    ZeroRtt<CheckedRange, CheckedRange, PacketNumber, EncryptedPayload<'a>>;
pub type CleartextZeroRtt<'a> = ZeroRtt<&'a [u8], &'a [u8], PacketNumber, DecoderBufferMut<'a>>;

impl<'a> ProtectedZeroRtt<'a> {
    #[inline]
    pub(crate) fn decode(
        _tag: Tag,
        version: Version,
        buffer: DecoderBufferMut,
    ) -> DecoderBufferMutResult<ProtectedZeroRtt> {
        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 = ZeroRtt {
            version,
            destination_connection_id,
            source_connection_id,
            packet_number,
            payload,
        };

        Ok((packet, remaining))
    }

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

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

        let packet_number = truncated_packet_number.expand(largest_acknowledged_packet_number);

        Ok(ZeroRtt {
            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> EncryptedZeroRtt<'a> {
    pub fn decrypt<C: ZeroRttKey>(self, crypto: &C) -> Result<CleartextZeroRtt<'a>, CryptoError> {
        let ZeroRtt {
            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();

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

        Ok(ZeroRtt {
            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> CleartextZeroRtt<'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 ZeroRtt<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>
    ZeroRtt<DCID, SCID, PacketNumber, Payload>
{
    fn encode_header<E: Encoder>(&self, packet_number_len: PacketNumberLen, encoder: &mut E) {
        let mut tag: u8 = zero_rtt_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: ZeroRttKey,
        H: ZeroRttHeaderKey,
    > PacketEncoder<K, H, Payload> for ZeroRtt<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) {
        ZeroRtt::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 super::{havoc, Interceptor};
use core::ops::Range;

#[derive(Debug)]
struct Direction {
    loss: Range<u64>,
    pass: Range<u64>,
    mode: Mode,
}

impl Default for Direction {
    fn default() -> Self {
        Self {
            loss: 0..0,
            pass: u64::MAX..u64::MAX,
            mode: Mode::Loss { remaining: 0 },
        }
    }
}

impl Direction {
    #[inline]
    fn should_pass<R: havoc::Random>(&mut self, random: &mut R) -> bool {
        let (remaining, mut is_pass) = match &mut self.mode {
            Mode::Pass { remaining } => (remaining, true),
            Mode::Loss { remaining } => (remaining, false),
        };

        // if we can decrement the remaining, then stay on the current mode
        if let Some(new_value) = remaining.checked_sub(1) {
            *remaining = new_value;
            return is_pass;
        }

        // try 3 times to generate the next value. It's good to be an odd number so we're at least
        // alternating if the ranges keep returning 0.
        for _ in 0..3 {
            // go to the next mode
            is_pass = !is_pass;

            let remaining = if is_pass {
                Self::gen_range(&self.pass, random)
            } else {
                Self::gen_range(&self.loss, random)
            };

            // if this round picked 0 then try again
            if remaining == 0 {
                continue;
            }

            if is_pass {
                self.mode = Mode::Pass { remaining };
            } else {
                self.mode = Mode::Loss { remaining };
            }
        }

        is_pass
    }

    #[inline]
    fn gen_range<R: havoc::Random>(range: &Range<u64>, random: &mut R) -> u64 {
        if range.start == range.end {
            return range.start;
        }

        random.gen_range(range.clone())
    }
}

#[derive(Debug)]
enum Mode {
    Loss { remaining: u64 },
    Pass { remaining: u64 },
}

#[derive(Debug, Default)]
pub struct Builder<R> {
    tx: Direction,
    rx: Direction,
    random: R,
}

impl<R> Builder<R>
where
    R: 'static + Send + havoc::Random,
{
    pub fn new(random: R) -> Self {
        Self {
            tx: Direction::default(),
            rx: Direction::default(),
            random,
        }
    }

    pub fn with_tx_pass(mut self, range: Range<u64>) -> Self {
        self.tx.pass = range;
        self
    }

    pub fn with_tx_loss(mut self, range: Range<u64>) -> Self {
        self.tx.loss = range;
        self
    }

    pub fn with_rx_pass(mut self, range: Range<u64>) -> Self {
        self.rx.pass = range;
        self
    }

    pub fn with_rx_loss(mut self, range: Range<u64>) -> Self {
        self.rx.loss = range;
        self
    }

    pub fn build(self) -> Loss<R> {
        Loss {
            tx: self.tx,
            rx: self.rx,
            random: self.random,
        }
    }
}

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

impl<R> Loss<R>
where
    R: 'static + Send + havoc::Random,
{
    pub fn builder(random: R) -> Builder<R> {
        Builder::new(random)
    }
}

impl<R> Interceptor for Loss<R>
where
    R: 'static + Send + havoc::Random,
{
    #[inline]
    fn intercept_rx_datagram<'a>(
        &mut self,
        _subject: &crate::event::api::Subject,
        _datagram: &super::Datagram,
        payload: s2n_codec::DecoderBufferMut<'a>,
    ) -> s2n_codec::DecoderBufferMut<'a> {
        if !self.rx.should_pass(&mut self.random) {
            return s2n_codec::DecoderBufferMut::new(&mut payload.into_less_safe_slice()[..0]);
        }

        payload
    }

    #[inline]
    fn intercept_tx_datagram(
        &mut self,
        _subject: &crate::event::api::Subject,
        _datagram: &super::Datagram,
        payload: &mut s2n_codec::EncoderBuffer,
    ) {
        if !self.tx.should_pass(&mut self.random) {
            payload.set_position(0);
        }
    }
}

#[cfg(test)]
mod tests {
    use super::{havoc::testing::RandomSlice, *};

    #[test]
    fn alternate_test() {
        static SLICE: &[u8] = &{
            let mut slice = [0u8; 256];
            let mut i = 0;
            while i < slice.len() {
                slice[i] = i as _;
                i += 1;
            }
            slice
        };

        let mut rand = RandomSlice::new(SLICE);

        let mut rx = Direction {
            loss: 0..10,
            pass: 1..10,
            ..Default::default()
        };

        let mut passed = 0;
        let mut dropped = 0;

        for _ in 0..256 {
            if rx.should_pass(&mut rand) {
                passed += 1;
            } else {
                dropped += 1;
            }
        }

        // these values will always be the same since the Random generator is deterministic
        assert_eq!(passed, 143);
        assert_eq!(dropped, 113);
    }
}


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

use crate::packet::number::{PacketNumber, PacketNumberRange, PacketNumberSpace};
use alloc::{boxed::Box, vec::Vec};
use core::fmt;

/// A data structure for tracking packets that are pending acknowledgement
///
/// The following assumptions are made and exploited
///
/// * Packet numbers are monotonically generated and inserted
/// * Packet numbers will mostly be removed in ranges
/// * Packet numbers that are deemed lost will also be removed and retransmitted
///
/// This is implemented as a buffer ring with a moving range for the lower and upper bound of
/// contained packet numbers. The following example illustrates how each field tracks state:
///
/// ```ignore
/// packets = [ PN(2), None, PN(0), PN(1) ]
///                           ^ index = 2
/// start = PN(0)
/// end = PN(2)
/// ```
///
/// Upon inserting `PN(3)` the state is now:
///
/// ```ignore
/// packets = [ PN(2), PN(3), PN(0), PN(1) ]
///                           ^ index = 2
/// start = PN(0)
/// end = PN(3)
/// ```
///
/// Upon removing `PN(0)` the state is now:
///
/// ```ignore
/// packets = [ PN(2), PN(3), None, PN(1) ]
///                                 ^ index = 3
/// start = PN(1)
/// end = PN(3)
/// ```
#[derive(Clone)]
pub struct Map<V> {
    /// The stored values for each packet number
    values: Box<[Option<V>]>,
    /// The smallest contained inclusive packet number in the map
    start: PacketNumber,
    /// The largest contained inclusive packet number in the map
    end: PacketNumber,
    /// The starting index of the first occupied packet
    ///
    /// This field will be set to the `packets.len()` if the map is empty
    index: usize,
}

impl<V: fmt::Debug> fmt::Debug for Map<V> {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.debug_map().entries(self.iter()).finish()
    }
}

/// Start with 8 sent packets at a time
///
/// Capacity will grow exponentially as more packet number entries are added
const DEFAULT_CAPACITY: usize = 8;

impl<V> Default for Map<V> {
    fn default() -> Self {
        // we use the Initial packet number space as a filler until an actual
        // packet number is inserted
        let base = PacketNumberSpace::Initial.new_packet_number(0u8.into());

        let mut values = Vec::with_capacity(DEFAULT_CAPACITY);
        while values.len() != values.capacity() {
            values.push(None);
        }
        let values = values.into_boxed_slice();

        // Set the index to the len (OOB) to indicate that it's empty
        let index = values.len();

        Self {
            values,
            start: base,
            end: base,
            index,
        }
    }
}

impl<V> Map<V> {
    /// Inserts the given `value`
    pub fn insert(&mut self, packet_number: PacketNumber, value: V) {
        if self.is_empty() {
            self.start = packet_number;
            self.end = packet_number;
            self.values[0] = Some(value);
            self.index = 0;
            return;
        }

        // The implementation assumes monotonicity of insertion
        debug_assert!(
            packet_number > self.start && packet_number > self.end,
            "packet numbers should be monotonic: {:?} > {:?} && {:?}",
            packet_number,
            self.start,
            self.end
        );

        // check if we need to increase capacity
        let distance = (packet_number.as_u64() - self.start.as_u64()) as usize;

        let index = if distance >= self.values.len() {
            self.resize(distance);

            // use the distance as the index since we've already resized beyond it
            distance
        } else {
            (self.index + distance) % self.values.len()
        };

        self.values[index] = Some(value);
        self.end = packet_number;
    }

    /// Inserts the given `value` into the map or updates the existing entry
    pub fn insert_or_update<F: FnOnce(&mut V)>(
        &mut self,
        packet_number: PacketNumber,
        value: V,
        update: F,
    ) {
        if self.is_empty() {
            self.start = packet_number;
            self.end = packet_number;
            self.values[0] = Some(value);
            self.index = 0;
            return;
        }

        // The implementation assumes insertion is not lower than the start
        debug_assert!(
            packet_number >= self.start,
            "packet numbers should be monotonic: {:?} > {:?}",
            packet_number,
            self.start,
        );

        // check if we need to increase capacity
        let distance = (packet_number.as_u64() - self.start.as_u64()) as usize;

        let index = if distance >= self.values.len() {
            self.resize(distance);

            // use the distance as the index since we've already resized beyond it
            distance
        } else {
            (self.index + distance) % self.values.len()
        };

        let entry = &mut self.values[index];

        if let Some(prev) = entry.as_mut() {
            update(prev);
        } else {
            *entry = Some(value);
        }

        self.end = self.end.max(packet_number);
    }

    /// Returns a reference to the `V` associated with the given `packet_number`
    #[inline]
    pub fn get(&self, packet_number: PacketNumber) -> Option<&V> {
        let index = self.pn_index(packet_number)?;
        self.values[index].as_ref()
    }

    /// Removes the value associated with the given `packet_number`
    /// and returns the value if it was present
    pub fn remove(&mut self, packet_number: PacketNumber) -> Option<V> {
        let index = self.pn_index(packet_number)?;
        let info = self.values[index].take()?;

        // update the bounds
        match (self.start == packet_number, self.end == packet_number) {
            // the bounds are inclusive so the map is now empty, reset it
            //              [_, _, _, 3]
            // remove(3) => [_, _, _, _]
            (true, true) => {
                self.clear();
            }
            // the packet was removed from the front
            //              [0, 1, _, 3, 4]
            // remove(0) => [_, 1, _, 3, 4]
            // remove(1) => [_, _, _, 3, 4]
            // remove(3) => [_, _, _, _, 4]
            (true, false) => {
                self.set_start(packet_number.next().unwrap());
            }
            // the packet was removed from the back
            //              [0, 1, _, 3, 4]
            // remove(4) => [0, 1, _, 3, _]
            // remove(3) => [0, 1, _, _, _]
            // remove(1) => [0, _, _, _, _]
            (false, true) => {
                self.set_end(packet_number.prev().unwrap());
            }
            // the packet was removed from the middle
            //              [0, 1, 2]
            // remove(2) => [0, _, 2]
            (false, false) => {
                // nothing to do
            }
        }

        Some(info)
    }

    /// Removes a range of packets from the map and returns their value
    #[inline]
    pub fn remove_range(&mut self, range: PacketNumberRange) -> RemoveIter<V> {
        RemoveIter::new(self, range)
    }

    /// Get the inclusive PacketNumberRange
    #[inline]
    pub fn get_range(&self) -> PacketNumberRange {
        PacketNumberRange::new(self.start, self.end)
    }

    /// Gets an iterator over the sent packet entries, sorted by PacketNumber
    #[inline]
    pub fn iter(&self) -> Iter<V> {
        Iter::new(self)
    }

    /// Returns true if there are no entries
    #[inline]
    pub fn is_empty(&self) -> bool {
        self.index == self.values.len()
    }

    /// Clears all of the packet information in the sent
    #[inline]
    pub fn clear(&mut self) {
        self.index = self.values.len();
    }

    #[inline]
    fn pn_index(&self, packet_number: PacketNumber) -> Option<usize> {
        // the map is empty so there are no valid entries
        if self.is_empty() {
            return None;
        }

        // make sure it's within the inserted packet numbers
        if packet_number > self.end {
            return None;
        }

        let offset = packet_number.checked_distance(self.start)?;
        let index = self.index.checked_add(offset as usize)?;
        let index = index % self.values.len();
        Some(index)
    }

    #[inline]
    fn set_start(&mut self, packet_number: PacketNumber) {
        // this function assumes we have at least one element
        debug_assert!(!self.is_empty());
        debug_assert!(packet_number >= self.start);
        debug_assert!(packet_number <= self.end);

        // find the next occupied slot
        for packet_number in PacketNumberRange::new(packet_number, self.end) {
            if self.get(packet_number).is_some() {
                let index = self
                    .pn_index(packet_number)
                    .expect("packet should be in bounds");

                self.index = index;
                self.start = packet_number;
                debug_assert!(self.start <= self.end);
                debug_assert_eq!(self.pn_index(packet_number), Some(index));
                return;
            }
        }

        unreachable!("could not find an occupied entry; map should be empty");
    }

    #[inline]
    fn set_end(&mut self, packet_number: PacketNumber) {
        // this function assumes we have at least one element
        debug_assert!(!self.is_empty());
        debug_assert!(packet_number >= self.start);
        debug_assert!(packet_number <= self.end);

        // find the next occupied slot
        for packet_number in PacketNumberRange::new(self.start, packet_number).rev() {
            if self.get(packet_number).is_some() {
                self.end = packet_number;
                debug_assert!(self.start <= self.end);
                return;
            }
        }

        unreachable!("could not find an occupied entry; map should be empty");
    }

    fn resize(&mut self, len: usize) {
        let mut new_len = self.values.len();

        // grow capacity until we can fit the inserted PN
        loop {
            new_len *= 2;
            if len < new_len {
                break;
            }
        }

        // allocate a new packet buffer and copy the previous values
        let mut values = Vec::with_capacity(new_len);
        // The packets are stored in a ring so we copy from the index
        // to the end, then from the start to the index
        values.extend(self.values[self.index..].iter_mut().map(|v| v.take()));
        values.extend(self.values[..self.index].iter_mut().map(|v| v.take()));
        while values.len() != values.capacity() {
            values.push(None);
        }

        // reset the index to the beginning of the buffer
        self.index = 0;
        self.values = values.into_boxed_slice();
    }
}

/// An iterator over all of the contained packet numbers
///
/// This iterator is optimized to reduce the amount of bounds checks being performed
#[derive(Debug)]
pub struct Iter<'a, V> {
    packets: &'a Map<V>,
    packet_number: Option<PacketNumber>,
    index: usize,
    remaining: usize,
}

impl<'a, V> Iter<'a, V> {
    #[inline]
    fn new(packets: &'a Map<V>) -> Self {
        let start = packets.start;
        let end = packets.end;
        let index = packets.index;

        let mut iter = Self {
            packets,
            packet_number: Some(start),
            index,
            // start with an empty iterator
            remaining: 0,
        };

        // make sure we have at least one packet
        if iter.packets.is_empty() {
            return iter;
        }

        // set the number of remaining entries based on the bounded range
        iter.remaining = (end.as_u64() - start.as_u64()) as usize;
        // we always have at least 1 items since the range is inclusive
        iter.remaining += 1;

        debug_assert!(iter.remaining <= iter.packets.values.len());

        iter
    }
}

impl<'a, V> Iterator for Iter<'a, V> {
    type Item = (PacketNumber, &'a V);

    #[inline]
    fn next(&mut self) -> Option<Self::Item> {
        while self.remaining > 0 {
            self.remaining -= 1;

            let packet_number = self.packet_number?;
            self.packet_number = packet_number.next();

            let index = self.index;
            self.index = (index + 1) % self.packets.values.len();

            if let Some(info) = self.packets.values[index].as_ref() {
                return Some((packet_number, info));
            }
        }

        None
    }
}

/// An iterator which removes a set of packet numbers in a range
///
/// This iterator is optimized to reduce the amount of bounds checks being performed
#[derive(Debug)]
pub struct RemoveIter<'a, V> {
    packets: &'a mut Map<V>,
    packet_number: Option<PacketNumber>,
    index: usize,
    remaining: usize,
}

impl<'a, V> RemoveIter<'a, V> {
    #[inline]
    fn new(packets: &'a mut Map<V>, range: PacketNumberRange) -> Self {
        let mut start = packets.start;
        let mut end = packets.end;

        let index = packets.index;

        let mut iter = Self {
            packets,
            packet_number: None,
            index,
            // start with an empty iterator
            remaining: 0,
        };

        // make sure we have at least one packet
        if iter.packets.is_empty() {
            return iter;
        }

        // ensure the range overlaps with the contained items
        if range.end() < start || range.start() > end {
            return iter;
        }

        use core::cmp::Ordering::*;

        match (range.start().cmp(&start), range.end().cmp(&end)) {
            (Less, Equal) | (Less, Greater) | (Equal, Greater) | (Equal, Equal) => {
                // deleting all entries

                // clear the sent packets
                //
                // NOTE: this doesn't actually delete anything in the buffer
                iter.packets.clear();

                // no need to update index as it's already set to the lower bound
            }
            (Less, Less) | (Equal, Less) => {
                // deleting start
                end = range.end();

                iter.packets.set_start(end.next().unwrap());
            }
            (Greater, Greater) | (Greater, Equal) => {
                // deleting end
                start = range.start();

                iter.index = iter
                    .packets
                    .pn_index(start)
                    .expect("packet number bounds have already been checked");

                iter.packets.set_end(start.prev().unwrap());
            }
            (Greater, Less) => {
                // deleting middle part
                start = range.start();
                end = range.end();

                iter.index = iter
                    .packets
                    .pn_index(start)
                    .expect("packet number bounds have already been checked");
            }
        }

        // Update the starting packet number
        iter.packet_number = Some(start);
        // set the number of remaining entries based on the bounded range
        iter.remaining = (end.as_u64() - start.as_u64()) as usize;
        // we always have at least 1 items since the range is inclusive
        iter.remaining += 1;

        debug_assert!(iter.remaining <= iter.packets.values.len());

        iter
    }
}

impl<'a, V> Iterator for RemoveIter<'a, V> {
    type Item = (PacketNumber, V);

    #[inline]
    fn next(&mut self) -> Option<Self::Item> {
        while self.remaining > 0 {
            self.remaining -= 1;

            let packet_number = self.packet_number?;
            self.packet_number = packet_number.next();

            let index = self.index;
            self.index = (index + 1) % self.packets.values.len();

            if let Some(info) = self.packets.values[index].take() {
                return Some((packet_number, info));
            }
        }

        None
    }
}

impl<'a, V> Drop for RemoveIter<'a, V> {
    fn drop(&mut self) {
        // make sure the iterator is drained, otherwise the entries might dangle
        while self.next().is_some() {}
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::{
        packet::number::{PacketNumber, PacketNumberRange, PacketNumberSpace},
        varint::VarInt,
    };
    use alloc::collections::BTreeMap;
    use bolero::{check, generator::*};

    type TestMap = Map<u64>;

    #[test]
    fn insert_get_range() {
        let mut sent_packets = TestMap::default();

        let packet_number_1 = PacketNumberSpace::Initial.new_packet_number(VarInt::from_u8(1));
        let packet_number_2 = packet_number_1.next().unwrap();
        let packet_number_3 = packet_number_2.next().unwrap();

        sent_packets.insert(packet_number_1, 1);
        sent_packets.insert(packet_number_2, 2);

        assert!(sent_packets.get(packet_number_1).is_some());
        assert!(sent_packets.get(packet_number_2).is_some());
        assert!(sent_packets.get(packet_number_3).is_none());

        assert_eq!(sent_packets.get(packet_number_1).unwrap(), &1);
        assert_eq!(sent_packets.get(packet_number_2).unwrap(), &2);

        sent_packets.insert(packet_number_3, 3);

        assert!(sent_packets.get(packet_number_3).is_some());
        assert_eq!(sent_packets.get(packet_number_3).unwrap(), &3);

        for (packet_number, sent_packet_info) in sent_packets.iter() {
            assert_eq!(sent_packets.get(packet_number).unwrap(), sent_packet_info);
        }
    }

    #[test]
    fn remove() {
        let mut sent_packets = TestMap::default();
        let packet_number = PacketNumberSpace::Initial.new_packet_number(VarInt::from_u8(1));
        sent_packets.insert(packet_number, 1);

        assert!(sent_packets.get(packet_number).is_some());
        assert_eq!(sent_packets.get(packet_number).unwrap(), &1);

        assert_eq!(Some(1), sent_packets.remove(packet_number));

        assert!(sent_packets.get(packet_number).is_none());

        // Removing a packet that was already removed doesn't panic
        assert_eq!(None, sent_packets.remove(packet_number));
    }

    #[test]
    fn empty() {
        let mut sent_packets = TestMap::default();
        assert!(sent_packets.is_empty());

        let packet_number = PacketNumberSpace::Initial.new_packet_number(VarInt::from_u8(1));
        sent_packets.insert(packet_number, 1);
        assert!(!sent_packets.is_empty());
    }

    #[test]
    #[should_panic]
    fn wrong_packet_space_on_insert() {
        let mut sent_packets = new_sent_packets(PacketNumberSpace::Initial);

        let packet_number =
            PacketNumberSpace::ApplicationData.new_packet_number(VarInt::from_u8(1));
        sent_packets.insert(packet_number, 1);
    }

    #[test]
    #[should_panic]
    fn wrong_packet_space_on_get() {
        let sent_packets = new_sent_packets(PacketNumberSpace::Initial);

        let packet_number =
            PacketNumberSpace::ApplicationData.new_packet_number(VarInt::from_u8(1));
        sent_packets.get(packet_number);
    }

    #[test]
    #[should_panic]
    fn wrong_packet_space_on_remove_range() {
        let mut sent_packets = new_sent_packets(PacketNumberSpace::Initial);

        let packet_number_start =
            PacketNumberSpace::ApplicationData.new_packet_number(VarInt::from_u8(1));
        let packet_number_end =
            PacketNumberSpace::ApplicationData.new_packet_number(VarInt::from_u8(2));
        sent_packets
            .remove_range(PacketNumberRange::new(
                packet_number_start,
                packet_number_end,
            ))
            .for_each(|_| ());
    }

    #[test]
    #[should_panic]
    fn wrong_packet_space_on_remove() {
        let mut sent_packets = new_sent_packets(PacketNumberSpace::Initial);

        let packet_number =
            PacketNumberSpace::ApplicationData.new_packet_number(VarInt::from_u8(1));
        sent_packets.remove(packet_number);
    }

    fn new_sent_packets(space: PacketNumberSpace) -> TestMap {
        let mut sent_packets = TestMap::default();
        let packet_number = space.new_packet_number(VarInt::from_u8(0));
        sent_packets.insert(packet_number, 0);
        sent_packets
    }

    /// An operation to be performed against a model
    #[derive(Clone, Copy, Debug, TypeGenerator)]
    enum Operation {
        // Inserts the current packet number
        Insert,
        // Skips the packet number
        Skip,
        // Removes a packet number
        Remove(VarInt),
        // Removes a range of packet numbers
        RemoveRange(VarInt, VarInt),
    }

    fn model(ops: &[Operation]) {
        let mut current = PacketNumberSpace::ApplicationData.new_packet_number(VarInt::from_u8(0));

        /// Tracks the subject against an oracle to ensure differential equivalency
        #[derive(Debug, Default)]
        struct Model {
            subject: TestMap,
            oracle: BTreeMap<PacketNumber, u64>,
        }

        impl Model {
            pub fn insert(&mut self, packet_number: PacketNumber) {
                let value = packet_number.as_u64();

                self.subject.insert(packet_number, value);
                self.oracle.insert(packet_number, value);
                self.check_consistency();
            }

            pub fn remove(&mut self, packet_number: PacketNumber) {
                assert_eq!(
                    self.subject.remove(packet_number),
                    self.oracle.remove(&packet_number)
                );
                self.check_consistency();
            }

            pub fn remove_range(&mut self, range: PacketNumberRange) {
                // trim range so we're not slamming the BTreeMap
                let range = if self.subject.is_empty() {
                    PacketNumberRange::new(range.start(), range.start())
                } else {
                    let start = range.start().max(self.subject.start);
                    let end = range.end().min(self.subject.end);
                    if start > end {
                        PacketNumberRange::new(start, start)
                    } else {
                        PacketNumberRange::new(start, end)
                    }
                };

                let actual: Vec<_> = self.subject.remove_range(range).collect();
                let mut expected = vec![];

                for pn in range {
                    if let Some(value) = self.oracle.remove(&pn) {
                        expected.push((pn, value));
                    }
                }

                assert_eq!(expected, actual);

                self.check_consistency();
            }

            fn check_consistency(&self) {
                let mut subject = self.subject.iter();
                let mut oracle = self.oracle.iter();
                loop {
                    match (subject.next(), oracle.next()) {
                        (Some(actual), Some((expected_pn, expected_info))) => {
                            assert_eq!((*expected_pn, expected_info), actual);
                        }
                        (None, None) => break,
                        (actual, expected) => {
                            panic!("expected: {:?}, actual: {:?}", expected, actual);
                        }
                    }
                }
            }
        }

        let mut model = Model::default();

        for op in ops.iter().copied() {
            match op {
                Operation::Insert => {
                    model.insert(current);
                    current = current.next().unwrap();
                }
                Operation::Skip => {
                    current = current.next().unwrap();
                }
                Operation::Remove(pn) => {
                    let pn = PacketNumberSpace::ApplicationData.new_packet_number(pn);

                    model.remove(pn);
                }
                Operation::RemoveRange(start, end) => {
                    let (start, end) = if start > end {
                        (end, start)
                    } else {
                        (start, end)
                    };
                    let start = PacketNumberSpace::ApplicationData.new_packet_number(start);
                    let end = PacketNumberSpace::ApplicationData.new_packet_number(end);
                    let range = PacketNumberRange::new(start, end);

                    model.remove_range(range);
                }
            }
        }
    }

    #[test]
    fn differential_test() {
        check!()
            .with_type::<Vec<Operation>>()
            .for_each(|ops| model(ops))
    }

    #[test]
    #[cfg_attr(kani, kani::proof, kani::unwind(9), kani::solver(kissat))]
    fn insert_value() {
        // Confirm that a value is inserted
        check!().with_type().cloned().for_each(|pn| {
            let space = PacketNumberSpace::ApplicationData;
            let mut map = Map::default();
            assert!(map.is_empty());
            let pn = space.new_packet_number(pn);

            map.insert(pn, ());

            assert!(map.get(pn).is_some());
            assert!(!map.is_empty());
        });
    }
}


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

mod sliding_window;
pub use sliding_window::{SlidingWindow, SlidingWindowError};

mod protected_packet_number;
pub use protected_packet_number::ProtectedPacketNumber;

//= https://www.rfc-editor.org/rfc/rfc9000#section-12.3
//# The packet number is an integer in the range 0 to 2^62-1.  This
//# number is used in determining the cryptographic nonce for packet
//# protection.  Each endpoint maintains a separate packet number for
//# sending and receiving.

use crate::varint::VarInt;

mod packet_number;
pub use packet_number::PacketNumber;

//= https://www.rfc-editor.org/rfc/rfc9000#section-12.3
//# Packet numbers are limited to this range because they need to be
//# representable in whole in the Largest Acknowledged field of an ACK
//# frame (Section 19.3).  When present in a long or short header
//# however, packet numbers are reduced and encoded in 1 to 4 bytes; see
//# Section 17.1.

mod truncated_packet_number;
pub use truncated_packet_number::TruncatedPacketNumber;

//= https://www.rfc-editor.org/rfc/rfc9000#section-12.3
//# Initial space:  All Initial packets (Section 17.2.2) are in this
//# space.
//#
//# Handshake space:  All Handshake packets (Section 17.2.4) are in this
//# space.
//#
//# Application data space:  All 0-RTT (Section 17.2.3) and 1-RTT
//# (Section 17.3.1) packets are in this space.

mod packet_number_space;
pub use packet_number_space::PacketNumberSpace;

//= https://www.rfc-editor.org/rfc/rfc9000#section-17.1
//# Packet numbers are integers in the range 0 to 2^62-1 (Section 12.3).
//# When present in long or short packet headers, they are encoded in 1
//# to 4 bytes.  The number of bits required to represent the packet
//# number is reduced by including only the least significant bits of the
//# packet number.

/// The packet number len is the two least significant bits of the packet tag
pub(crate) const PACKET_NUMBER_LEN_MASK: u8 = 0b11;

mod packet_number_len;
pub use packet_number_len::PacketNumberLen;

mod packet_number_range;
pub use packet_number_range::PacketNumberRange;

#[cfg(feature = "alloc")]
pub mod map;
#[cfg(feature = "alloc")]
pub use map::Map;

#[cfg(test)]
mod tests;

//= https://www.rfc-editor.org/rfc/rfc9000#section-17.1
//# the sender MUST use a packet number size able to represent more than
//# twice as large a range as the difference between the largest
//# acknowledged packet number and the packet number being sent.  A peer
//# receiving the packet will then correctly decode the packet number,
//# unless the packet is delayed in transit such that it arrives after
//# many higher-numbered packets have been received.  An endpoint SHOULD
//# use a large enough packet number encoding to allow the packet number
//# to be recovered even if the packet arrives after packets that are
//# sent afterwards.

fn derive_truncation_range(
    largest_acknowledged_packet_number: PacketNumber,
    packet_number: PacketNumber,
) -> Option<PacketNumberLen> {
    let space = packet_number.space();
    space.assert_eq(largest_acknowledged_packet_number.space());
    packet_number
        .as_u64()
        .checked_sub(largest_acknowledged_packet_number.as_u64())
        .and_then(|value| value.checked_mul(2))
        .and_then(|value| VarInt::new(value).ok())
        .and_then(|value| PacketNumberLen::from_varint(value, space))
}

//= https://www.rfc-editor.org/rfc/rfc9000#section-17.1
//# As a result, the size of the packet number encoding is at least one
//# bit more than the base-2 logarithm of the number of contiguous
//# unacknowledged packet numbers, including the new packet.

//= https://www.rfc-editor.org/rfc/rfc9000#appendix-A.2
//# For example, if an endpoint has received an acknowledgment for packet
//# 0xabe8b3 and is sending a packet with a number of 0xac5c02, there are
//# 29,519 (0x734f) outstanding packet numbers.  In order to represent at
//# least twice this range (59,038 packets, or 0xe69e), 16 bits are
//# required.
//#
//# In the same state, sending a packet with a number of 0xace8fe uses
//# the 24-bit encoding, because at least 18 bits are required to
//# represent twice the range (131,222 packets, or 0x020096).

#[test]
fn packet_number_len_example_test() {
    let largest_acknowledged_packet_number =
        PacketNumberSpace::default().new_packet_number(VarInt::from_u32(0x00ab_e8bc));

    assert_eq!(
        PacketNumberSpace::default()
            .new_packet_number(VarInt::from_u32(0x00ac_5c02))
            .truncate(largest_acknowledged_packet_number)
            .unwrap()
            .bitsize(),
        16,
    );

    assert_eq!(
        PacketNumberSpace::default()
            .new_packet_number(VarInt::from_u32(0x00ac_e8fe))
            .truncate(largest_acknowledged_packet_number)
            .unwrap()
            .bitsize(),
        24,
    );
}

//= https://www.rfc-editor.org/rfc/rfc9000#section-17.1
//# At a receiver, protection of the packet number is removed prior to
//# recovering the full packet number.  The full packet number is then
//# reconstructed based on the number of significant bits present, the
//# value of those bits, and the largest packet number received in a
//# successfully authenticated packet.  Recovering the full packet number
//# is necessary to successfully complete the removal of packet
//# protection.
//
//# Once header protection is removed, the packet number is decoded by
//# finding the packet number value that is closest to the next expected
//# packet.  The next expected packet is the highest received packet
//# number plus one.  Pseudocode and an example for packet number
//# decoding can be found in Appendix A.3.

#[test]
fn packet_decoding_example_test() {
    let space = PacketNumberSpace::default();
    let largest_packet_number = space.new_packet_number(VarInt::from_u32(0xa82f_30ea));
    let truncated_packet_number = TruncatedPacketNumber::new(0x9b32u16, space);
    let expected = space.new_packet_number(VarInt::from_u32(0xa82f_9b32));
    let actual = decode_packet_number(largest_packet_number, truncated_packet_number);
    assert_eq!(actual, expected);
    assert_eq!(
        expected.truncate(largest_packet_number).unwrap(),
        truncated_packet_number
    );
}

//= https://www.rfc-editor.org/rfc/rfc9000#appendix-A.3
//# DecodePacketNumber(largest_pn, truncated_pn, pn_nbits):
//#   expected_pn  = largest_pn + 1
//#   pn_win       = 1 << pn_nbits
//#   pn_hwin      = pn_win / 2
//#   pn_mask      = pn_win - 1
//#   // The incoming packet number should be greater than
//#   // expected_pn - pn_hwin and less than or equal to
//#   // expected_pn + pn_hwin
//#   //
//#   // This means we cannot just strip the trailing bits from
//#   // expected_pn and add the truncated_pn because that might
//#   // yield a value outside the window.
//#   //
//#   // The following code calculates a candidate value and
//#   // makes sure it's within the packet number window.
//#   // Note the extra checks to prevent overflow and underflow.
//#   candidate_pn = (expected_pn & ~pn_mask) | truncated_pn
//#   if candidate_pn <= expected_pn - pn_hwin and
//#     candidate_pn < (1 << 62) - pn_win:
//#     return candidate_pn + pn_win
//#   if candidate_pn > expected_pn + pn_hwin and
//#     candidate_pn >= pn_win:
//#     return candidate_pn - pn_win
//#   return candidate_pn

#[inline(never)] // prevent the compiler from optimizing call size values and making this non-constant time
fn decode_packet_number(
    largest_pn: PacketNumber,
    truncated_pn: TruncatedPacketNumber,
) -> PacketNumber {
    let space = largest_pn.space();
    space.assert_eq(truncated_pn.space());

    let pn_nbits = truncated_pn.bitsize();
    // deref to u64 so we have enough room
    let expected_pn = largest_pn.as_u64() + 1;
    let pn_win = 1 << pn_nbits;
    let pn_hwin = pn_win / 2;
    let pn_mask = pn_win - 1;
    let mut candidate_pn = (expected_pn & !pn_mask) | truncated_pn.into_u64();

    let a = expected_pn
        .checked_sub(pn_hwin)
        .filter(|v| candidate_pn <= *v)
        .is_some();
    let b = (1u64 << 62)
        .checked_sub(pn_win)
        .filter(|v| candidate_pn < *v)
        .is_some();
    let c = expected_pn
        .checked_add(pn_hwin)
        .filter(|v| candidate_pn > *v)
        .is_some();
    let d = candidate_pn >= pn_win;

    let ab = a && b;
    let cd = !ab && c && d;

    // Make sure the compiler doesn't try and optimize the if statements
    // See https://godbolt.org/z/348WbbbzM
    core::sync::atomic::compiler_fence(core::sync::atomic::Ordering::SeqCst);

    if ab {
        candidate_pn += pn_win;
    }

    if cd {
        candidate_pn -= pn_win
    }

    let candidate_pn = VarInt::new(candidate_pn).unwrap_or(VarInt::MAX);

    PacketNumber::from_varint(candidate_pn, space)
}