internet 0.0.5

//! Packet headers encoding following [Section 17].
//!
//! > description
//!
//! This module provides the following primitives:
//!
//! - [`FixedBit`]
//! - [`LongHeaderPacketType`]
//! - [`ConnectionId`]
//! - [`PacketNumberLength`]
//! - [`PacketNumber`]
//! - [`Length`]
//! - [`RetryToken`]
//! - [`RetryIntegrityTag`]
//! - [`SpinBit`]
//! - [`KeyPhase`]
//!
//! And the following packet headers:
//!
//! - [`InitialPacket`]
//! - [`ZeroRttPacket`]
//! - [`HandshakePacket`]
//! - [`RetryPacket`]
//! - [`OneRttPacket`]
//!
//! [Section 17]: https://datatracker.ietf.org/doc/html/rfc9000#section-17

use crate::{
    Buf,
    BufError::{self},
    BufMut, BufResult, Codec, Cursor,
    quic::{HeaderForm, Version},
    quicv1::VariableLengthInteger,
};

/// A Fixed Bit of the Long Header Packet following [Section 17.2].
///
/// The bit is set to 1, unless the packet is a [Version Negotiation packet].
///
/// [Section 17.2]: https://datatracker.ietf.org/doc/html/rfc9000#section-17.2
/// [Version Negotiation packet]: VersionNegotiationPacket
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
#[repr(u8)]
pub enum FixedBit {
    /// Indicating any packet.
    One = 1,
    /// Indicating [Version Negotiation Packet](`VersionNegotiationPacket`).
    VersionNegotiationPacket = 0,
}

impl Codec for FixedBit {
    fn encode<W: BufMut>(&self, writer: &mut Cursor<W>, _: ()) -> BufResult<()> {
        let current_byte = writer.peek_u8().unwrap_or(0x00);
        let bit_mask = (*self as u8) << 6;
        let updated_byte = (current_byte & !0x40) | bit_mask;
        writer.poke_u8(updated_byte)
    }

    fn decode<R: Buf>(reader: &mut Cursor<R>, _: ()) -> BufResult<Self> {
        let byte = reader.peek_u8()?;
        match (byte & 0x40) >> 6 {
            0 => Ok(Self::VersionNegotiationPacket),
            _ => Ok(Self::One),
        }
    }
}

/// A Long Header Type following [Section 17.2, Table 5].
///
/// Indicating type of the Long Header Packet.
///
/// [Section 17.2, Table 5]: https://datatracker.ietf.org/doc/html/rfc9000#table-5
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
#[repr(u8)]
pub enum LongHeaderPacketType {
    /// [Initial Packet](InitialPacket) type.
    Initial = 0x00,

    /// [0-RTT Packet](ZeroRttPacket) type.
    ZeroRtt = 0x01,

    /// [Handshake Packet](HandshakePacket) type.
    Handshake = 0x02,

    /// [Retry Packet](RetryPacket) type.
    Retry = 0x03,
}

impl Codec for LongHeaderPacketType {
    fn encode<W: BufMut>(&self, writer: &mut Cursor<W>, _: ()) -> BufResult<()> {
        let current_byte = writer.peek_u8().unwrap_or(0x00);
        let bit_mask = (*self as u8) << 4;
        let updated_byte = (current_byte & !0x30) | bit_mask;
        writer.poke_u8(updated_byte)
    }

    fn decode<R: Buf>(reader: &mut Cursor<R>, _: ()) -> BufResult<Self> {
        let value = (reader.peek_u8()? & 0x30) >> 4;
        match value {
            0 => Ok(Self::Initial),
            1 => Ok(Self::ZeroRtt),
            2 => Ok(Self::Handshake),
            _ => Ok(Self::Retry),
        }
    }
}

/// A connection ID, specified for QUIC version 1.
///
/// Helper structure to contain the `Connection ID Length` and `Connection ID` fields.
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
pub struct ConnectionId {
    /// The length of the connection ID, must not exceed 20 in QUIC version 1.
    length: u8,
    /// The connection ID bytes.
    bytes: [u8; Self::MAXIMAL_LENGTH],
}

impl ConnectionId {
    ///
    pub const MAXIMAL_LENGTH: usize = 20;

    ///
    pub fn new(connection_id: &[u8]) -> Option<Self> {
        let length = connection_id.length();
        if length > Self::MAXIMAL_LENGTH as usize {
            return None;
        }
        let mut cid = [0u8; Self::MAXIMAL_LENGTH];
        cid[..length].copy_from_slice(connection_id);
        let length = length as u8;
        Some(Self { length, bytes: cid })
    }

    ///
    pub fn into_inner(self) -> (u8, [u8; Self::MAXIMAL_LENGTH]) {
        (self.length, self.bytes)
    }

    ///
    pub fn length(self) -> u8 {
        self.length
    }

    ///
    pub fn bytes(self) -> [u8; Self::MAXIMAL_LENGTH] {
        self.bytes
    }
}

impl Codec for ConnectionId {
    fn encode<W: BufMut>(&self, writer: &mut Cursor<W>, _: ()) -> BufResult<()> {
        let (cil, ci) = self.into_inner();
        cil.encode(writer, ())?;
        writer.write_slice(&ci[..cil as usize])
    }

    fn decode<R: Buf>(reader: &mut Cursor<R>, _: ()) -> BufResult<Self> {
        let ci = &mut [0u8; Self::MAXIMAL_LENGTH];
        let cil = u8::decode(reader, ())?;
        reader.read_into(&mut ci[..cil as usize])?;
        Ok(Self::new(&ci[..cil as usize]).ok_or(BufError::UnexpectedValue)?)
    }
}

impl Codec<u8> for ConnectionId {
    fn encode<W: BufMut>(&self, writer: &mut Cursor<W>, _: u8) -> BufResult<()> {
        let (cil, ci) = self.into_inner();
        writer.write_slice(&ci[..cil as usize])
    }

    fn decode<R: Buf>(reader: &mut Cursor<R>, cil: u8) -> BufResult<Self> {
        let ci = &mut [0u8; Self::MAXIMAL_LENGTH];
        reader.read_into(&mut ci[..cil as usize])?;
        Ok(Self::new(&ci[..cil as usize]).ok_or(BufError::UnexpectedValue)?)
    }
}

/// A Packet Number Length.
///
/// > description
///
/// Defined in [IETF RFC 9000, Section 17.2](https://datatracker.ietf.org/doc/html/rfc9000#section-17.2).
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
pub struct PacketNumberLength(pub u8);

impl PacketNumberLength {
    /// Выбирает минимальную длину, достаточную для кодирования данного номера пакета.
    pub fn from_packet_number(packet_number: PacketNumber) -> BufResult<Self> {
        // Предполагается, что VariableLengthInteger имеет поле .0 типа u64
        let pn = packet_number.0.0;

        if pn <= u8::MAX as u64 {
            Ok(Self(0)) // 1 byte
        } else if pn <= u16::MAX as u64 {
            Ok(Self(1)) // 2 bytes
        } else if pn <= 0xFFFFFF {
            Ok(Self(2)) // 3 bytes
        } else {
            Ok(Self(3)) // 4 bytes
        }
    }

    /// Возвращает фактическое количество байт, занимаемое номером пакета (1..=4).
    pub fn byte_length(&self) -> u8 {
        self.0 + 1
    }
}

impl Codec for PacketNumberLength {
    fn encode<W: BufMut>(&self, writer: &mut Cursor<W>, _: ()) -> BufResult<()> {
        let current_byte = writer.peek_u8().unwrap_or(0x00);
        let bit_mask = self.0 & 0x03;
        let updated_byte = (current_byte & !0x03) | bit_mask;

        writer.write_u8(updated_byte)
    }

    fn decode<R: Buf>(reader: &mut Cursor<R>, _: ()) -> BufResult<Self> {
        Ok(Self(reader.read_u8()? & 0x03))
    }
}

/// A Packet Number.
///
/// > description
///
/// Defined in [IETF RFC 9000, Section 17.2](https://datatracker.ietf.org/doc/html/rfc9000#section-17.2).
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
pub struct PacketNumber(pub VariableLengthInteger);
impl PacketNumber {
    ///
    pub fn new(last_sent: PacketNumber) -> Self {
        let mut last_sent = last_sent;
        last_sent.0.0 += 1;
        Self(last_sent.0)
    }

    ///
    pub fn length(&self) -> BufResult<PacketNumberLength> {
        PacketNumberLength::from_packet_number(*self)
    }
}
impl Codec<PacketNumberLength> for PacketNumber {
    fn encode<W: BufMut>(
        &self,
        writer: &mut Cursor<W>,
        length: PacketNumberLength,
    ) -> BufResult<()> {
        let pn = self.0.0;

        match length.0 {
            0 => writer.write_u8(pn as u8),
            1 => writer.write_u16_be(pn as u16), // Big-endian согласно RFC
            2 => {
                // 3 bytes: записываем по одному байту в big-endian порядке
                writer.write_u8((pn >> 16) as u8)?;
                writer.write_u8((pn >> 8) as u8)?;
                writer.write_u8(pn as u8)
            }
            3 => writer.write_u32_be(pn as u32),
            _ => Err(BufError::UnexpectedValue),
        }
    }

    fn decode<R: Buf>(reader: &mut Cursor<R>, length: PacketNumberLength) -> BufResult<Self> {
        let pn = match length.0 {
            0 => reader.read_u8()? as u64,
            1 => reader.read_u16_be()? as u64,
            2 => {
                // Читаем 3 байта и собираем в u64 в big-endian порядке
                let b1 = reader.read_u8()? as u64;
                let b2 = reader.read_u8()? as u64;
                let b3 = reader.read_u8()? as u64;
                (b1 << 16) | (b2 << 8) | b3
            }
            3 => reader.read_u32_be()? as u64,
            _ => return Err(BufError::UnexpectedValue),
        };

        Ok(Self(VariableLengthInteger(pn)))
    }
}

/// A length of the [PacketNumber] and Payload fields in bytes following [Section 17.2].
///
/// > description
///
/// Defined in [Section 17.2](https://datatracker.ietf.org/doc/html/rfc9000#section-17.2).
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
pub struct Length(pub VariableLengthInteger);
impl Length {
    ///
    pub fn calculate(
        packet_payload_length: usize,
        packet_number_length: PacketNumberLength,
    ) -> BufResult<Self> {
        Ok(Self(VariableLengthInteger::new(
            packet_payload_length as u64 + packet_number_length.0 as u64,
        )?))
    }
}

impl Codec for Length {
    fn encode<W: BufMut>(&self, writer: &mut Cursor<W>, _: ()) -> BufResult<()> {
        self.0.encode(writer, ())
    }

    fn decode<R: Buf>(reader: &mut Cursor<R>, _: ()) -> BufResult<Self> {
        Ok(Self(VariableLengthInteger::decode(reader, ())?))
    }
}

/// A Token.
///
/// > description
#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)]
pub struct RetryToken(pub Vec<u8>);

impl Codec for RetryToken {
    ///
    fn encode<W: BufMut>(&self, writer: &mut Cursor<W>, _: ()) -> BufResult<()> {
        self.0.encode(writer, ())
    }

    ///
    fn decode<R: Buf>(reader: &mut Cursor<R>, _: ()) -> BufResult<Self> {
        let remaining = reader.remaining();
        if remaining < 16 {
            return Err(BufError::OutOfBounds);
        }
        let length = remaining - 16;
        Ok(Self(Vec::decode(reader, length)?))
    }
}

impl Codec<VariableLengthInteger> for RetryToken {
    ///
    fn encode<W: BufMut>(&self, writer: &mut Cursor<W>, _: VariableLengthInteger) -> BufResult<()> {
        self.0.encode(writer, ())
    }

    ///
    fn decode<R: Buf>(reader: &mut Cursor<R>, length: VariableLengthInteger) -> BufResult<Self> {
        Ok(Self(Vec::decode(reader, length.0 as usize)?))
    }
}

/// A Retry Integrity Tag.
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
pub struct RetryIntegrityTag(pub [u8; 16]);

impl Codec for RetryIntegrityTag {
    fn encode<W: BufMut>(&self, writer: &mut Cursor<W>, _: ()) -> BufResult<()> {
        writer.write_array(&self.0)
    }

    fn decode<R: Buf>(reader: &mut Cursor<R>, _: ()) -> BufResult<Self> {
        Ok(Self(reader.read_array::<16>()?))
    }
}

/// A Spin Bit of the [1-RTT packet].
///
/// > description
///
/// [1-RTT packet]: OneRttPacket
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
#[repr(u8)]
pub enum SpinBit {
    ///
    Zero = 0,
    ///
    One = 1,
}

impl SpinBit {
    ///
    pub fn spin(&mut self) {
        *self = match self {
            SpinBit::Zero => SpinBit::One,
            SpinBit::One => SpinBit::Zero,
        }
    }
}

impl Codec for SpinBit {
    fn encode<W: BufMut>(&self, writer: &mut Cursor<W>, _: ()) -> BufResult<()> {
        let current_byte = writer.peek_u8().unwrap_or(0x00);
        let bit_mask = (*self as u8) << 5;
        let updated_byte = (current_byte & !0b100000) | bit_mask;
        writer.poke_u8(updated_byte)
    }

    fn decode<R: Buf>(reader: &mut Cursor<R>, _: ()) -> BufResult<Self> {
        let value = reader.peek_u8()? & 0b100000;
        match value {
            0 => Ok(Self::Zero),
            _ => Ok(Self::One),
        }
    }
}

/// A Key Phase of the [1-RTT packet].
///
/// > description
///
/// [1-RTT packet]: OneRttPacket
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
#[repr(u8)]
pub enum KeyPhase {
    /// .
    Zero = 0,
    /// .
    One = 1,
}

impl Codec for KeyPhase {
    fn encode<W: BufMut>(&self, writer: &mut Cursor<W>, _: ()) -> BufResult<()> {
        let current_byte = writer.peek_u8().unwrap_or(0x00);
        let bit_mask = (*self as u8) << 2;
        let updated_byte = (current_byte & !0b100) | bit_mask;
        writer.poke_u8(updated_byte)
    }

    fn decode<R: Buf>(reader: &mut Cursor<R>, _: ()) -> BufResult<Self> {
        let value = reader.peek_u8()? & 0b100;
        match value {
            0 => Ok(Self::Zero),
            _ => Ok(Self::One),
        }
    }
}

/// An Initial Packet following [Section 17.2.2].
///
/// Carries the first [CRYPTO frames] from client and [ACK frames] in either direction.
///
/// [CRYPTO frames]: crate::quicv1::CryptoFrame
/// [ACK frames]: crate::quicv1::AckFrame
/// [Section 17.2.2]: https://datatracker.ietf.org/doc/html/rfc9000#section-17.2.2
#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)]
pub struct InitialPacket {
    /// The version.
    pub version: Version,
    /// The destination connection ID.
    pub destination_connection_id: ConnectionId,
    /// The source connection ID.
    pub source_connection_id: ConnectionId,
    /// The token received from the the server [Retry packet](RetryPacket).
    pub token: RetryToken,
    /// The Packet Number.
    pub packet_number: PacketNumber,
    /// The Packet Payload.
    pub packet_payload: Vec<u8>,
}

impl InitialPacket {
    /// The header form.
    pub const HEADER_FORM: HeaderForm = HeaderForm::LongHeader;
    /// The fixed bit.
    pub const FIXED_BIT: FixedBit = FixedBit::One;
    /// The long packet type.
    pub const LONG_PACKET_TYPE: LongHeaderPacketType = LongHeaderPacketType::Initial;
}

impl Codec for InitialPacket {
    fn encode<W: BufMut>(&self, writer: &mut Cursor<W>, _: ()) -> BufResult<()> {
        Self::HEADER_FORM.encode(writer, ())?;
        Self::FIXED_BIT.encode(writer, ())?;
        Self::LONG_PACKET_TYPE.encode(writer, ())?;
        let packet_number_length = self.packet_number.length()?;
        packet_number_length.encode(writer, ())?;
        self.version.encode(writer, ())?;
        self.destination_connection_id.encode(writer, ())?;
        self.source_connection_id.encode(writer, ())?;
        VariableLengthInteger::new(self.token.0.len() as u64)?.encode(writer, ())?;
        self.token.0.encode(writer, ())?;
        Length::calculate(self.packet_payload.len(), packet_number_length)?.encode(writer, ())?;
        self.packet_number.encode(writer, packet_number_length)?;
        self.packet_payload.encode(writer, ())
    }

    fn decode<R: Buf>(reader: &mut Cursor<R>, _: ()) -> BufResult<Self> {
        if HeaderForm::decode(reader, ())? != Self::HEADER_FORM {
            return Err(BufError::UnexpectedValue);
        }
        if FixedBit::decode(reader, ())? != Self::FIXED_BIT {
            return Err(BufError::UnexpectedValue);
        }
        if LongHeaderPacketType::decode(reader, ())? != Self::LONG_PACKET_TYPE {
            return Err(BufError::UnexpectedValue);
        }
        let packet_number_length = PacketNumberLength::decode(reader, ())?;
        let version = Version::decode(reader, ())?;
        let destination_connection_id = ConnectionId::decode(reader, ())?;
        let source_connection_id = ConnectionId::decode(reader, ())?;
        let token_length = VariableLengthInteger::decode(reader, ())?;
        let token_bytes = &mut [0u8; 16];
        reader.read_into(&mut token_bytes[..token_length.0 as usize])?;
        let token = RetryToken((&token_bytes[..token_length.0 as usize]).to_vec());
        let _length = VariableLengthInteger::decode(reader, ())?;
        let packet_number = PacketNumber::decode(reader, packet_number_length)?;
        let packet_payload = Vec::decode(reader, ())?;

        Ok(Self {
            version,
            destination_connection_id,
            source_connection_id,
            token,
            packet_number,
            packet_payload,
        })
    }
}

/// A 0-RTT Packet following [Section 17.2.3].
///
/// Carries early data from the client to the server prior to handshake completion.
///
/// [Section 17.2.3]: https://datatracker.ietf.org/doc/html/rfc9000#section-17.2.3
#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)]
pub struct ZeroRttPacket {
    /// The version.
    pub version: Version,
    /// The destination connection ID.
    pub destination_connection_id: ConnectionId,
    /// The source connection ID.
    pub source_connection_id: ConnectionId,
    /// The packet number.
    pub packet_number: PacketNumber,
    /// The packet payload.
    pub packet_payload: Vec<u8>,
}

impl ZeroRttPacket {
    /// The header form.
    pub const HEADER_FORM: HeaderForm = HeaderForm::LongHeader;
    /// The fixed bit.
    pub const FIXED_BIT: FixedBit = FixedBit::One;
    /// The long packet type.
    pub const LONG_PACKET_TYPE: LongHeaderPacketType = LongHeaderPacketType::ZeroRtt;
}

impl Codec for ZeroRttPacket {
    fn encode<W: BufMut>(&self, writer: &mut Cursor<W>, _: ()) -> BufResult<()> {
        Self::HEADER_FORM.encode(writer, ())?;
        Self::FIXED_BIT.encode(writer, ())?;
        Self::LONG_PACKET_TYPE.encode(writer, ())?;
        let packet_number_length = self.packet_number.length()?;
        packet_number_length.encode(writer, ())?;
        self.version.encode(writer, ())?;
        self.destination_connection_id.encode(writer, ())?;
        self.source_connection_id.encode(writer, ())?;
        Length::calculate(self.packet_payload.len(), packet_number_length)?.encode(writer, ())?;
        self.packet_number.encode(writer, packet_number_length)?;
        self.packet_payload.encode(writer, ())
    }

    fn decode<R: Buf>(reader: &mut Cursor<R>, _: ()) -> BufResult<Self> {
        if HeaderForm::decode(reader, ())? != Self::HEADER_FORM {
            return Err(BufError::UnexpectedValue);
        }
        if FixedBit::decode(reader, ())? != Self::FIXED_BIT {
            return Err(BufError::UnexpectedValue);
        }
        if LongHeaderPacketType::decode(reader, ())? != Self::LONG_PACKET_TYPE {
            return Err(BufError::UnexpectedValue);
        }
        let packet_number_length = PacketNumberLength::decode(reader, ())?;
        let version = Version::decode(reader, ())?;
        let destination_connection_id = ConnectionId::decode(reader, ())?;
        let source_connection_id = ConnectionId::decode(reader, ())?;
        let _length = VariableLengthInteger::decode(reader, ())?;
        let packet_number = PacketNumber::decode(reader, packet_number_length)?;
        let packet_payload = Vec::decode(reader, ())?;

        Ok(Self {
            version,
            destination_connection_id,
            source_connection_id,
            packet_number,
            packet_payload,
        })
    }
}

/// A Handshake Packet following [Section 17.2.4].
///
/// Carries cryptographic handshake messages and acknowledgments from the server and client.
///
/// [Section 17.2.4]: https://datatracker.ietf.org/doc/html/rfc9000#section-17.2.4
#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)]
pub struct HandshakePacket {
    /// The version.
    pub version: Version,
    /// The destination connection ID.
    pub destination_connection_id: ConnectionId,
    /// The source connection ID.
    pub source_connection_id: ConnectionId,
    /// The packet number.
    pub packet_number: PacketNumber,
    /// The packet payload.
    pub packet_payload: Vec<u8>,
}

impl HandshakePacket {
    /// The header form.
    pub const HEADER_FORM: HeaderForm = HeaderForm::LongHeader;
    /// The fixed bit.
    pub const FIXED_BIT: FixedBit = FixedBit::One;
    /// The long packet type.
    pub const LONG_PACKET_TYPE: LongHeaderPacketType = LongHeaderPacketType::Handshake;
}

impl Codec for HandshakePacket {
    fn encode<W: BufMut>(&self, writer: &mut Cursor<W>, _: ()) -> BufResult<()> {
        Self::HEADER_FORM.encode(writer, ())?;
        Self::FIXED_BIT.encode(writer, ())?;
        Self::LONG_PACKET_TYPE.encode(writer, ())?;
        let packet_number_length = self.packet_number.length()?;
        packet_number_length.encode(writer, ())?;
        self.version.encode(writer, ())?;
        self.destination_connection_id.encode(writer, ())?;
        self.source_connection_id.encode(writer, ())?;
        Length::calculate(self.packet_payload.len(), packet_number_length)?.encode(writer, ())?;
        self.packet_number.encode(writer, packet_number_length)?;
        self.packet_payload.encode(writer, ())
    }

    fn decode<R: Buf>(reader: &mut Cursor<R>, _: ()) -> BufResult<Self> {
        if HeaderForm::decode(reader, ())? != Self::HEADER_FORM {
            return Err(BufError::UnexpectedValue);
        }
        if FixedBit::decode(reader, ())? != Self::FIXED_BIT {
            return Err(BufError::UnexpectedValue);
        }
        if LongHeaderPacketType::decode(reader, ())? != Self::LONG_PACKET_TYPE {
            return Err(BufError::UnexpectedValue);
        }
        let packet_number_length = PacketNumberLength::decode(reader, ())?;
        let version = Version::decode(reader, ())?;
        let destination_connection_id = ConnectionId::decode(reader, ())?;
        let source_connection_id = ConnectionId::decode(reader, ())?;
        let _length = VariableLengthInteger::decode(reader, ())?;
        let packet_number = PacketNumber::decode(reader, packet_number_length)?;
        let packet_payload = Vec::decode(reader, ())?;

        Ok(Self {
            version,
            destination_connection_id,
            source_connection_id,
            packet_number,
            packet_payload,
        })
    }
}

/// A Retry Packet following [Section 17.2.5].
///
/// Carries an [address validation token] created by the server to the client.
///
/// [Section 17.2.5]: https://datatracker.ietf.org/doc/html/rfc9000#section-17.2.5
/// [address validation token]: RetryToken
#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)]
pub struct RetryPacket {
    /// The version.
    pub version: Version,
    /// The destination connection ID.
    pub destination_connection_id: ConnectionId,
    /// The source connection ID.
    pub source_connection_id: ConnectionId,
    /// The retry token.
    pub retry_token: RetryToken,
    /// The retry integrity tag.
    pub retry_integrity_tag: RetryIntegrityTag,
}

impl RetryPacket {
    /// The header form.
    pub const HEADER_FORM: HeaderForm = HeaderForm::LongHeader;
    /// The fixed bit.
    pub const FIXED_BIT: FixedBit = FixedBit::One;
    /// The long packet type.
    pub const LONG_PACKET_TYPE: LongHeaderPacketType = LongHeaderPacketType::Retry;
}

impl Codec for RetryPacket {
    fn encode<W: BufMut>(&self, writer: &mut Cursor<W>, _: ()) -> BufResult<()> {
        Self::HEADER_FORM.encode(writer, ())?;
        Self::FIXED_BIT.encode(writer, ())?;
        Self::LONG_PACKET_TYPE.encode(writer, ())?;
        writer.advance(1)?;
        self.version.encode(writer, ())?;
        self.destination_connection_id.encode(writer, ())?;
        self.source_connection_id.encode(writer, ())?;
        self.retry_token.encode(writer, ())?;
        self.retry_integrity_tag.encode(writer, ())
    }

    fn decode<R: Buf>(reader: &mut Cursor<R>, _: ()) -> BufResult<Self> {
        if HeaderForm::decode(reader, ())? != Self::HEADER_FORM {
            return Err(BufError::UnexpectedValue);
        }
        if FixedBit::decode(reader, ())? != Self::FIXED_BIT {
            return Err(BufError::UnexpectedValue);
        }
        if LongHeaderPacketType::decode(reader, ())? != Self::LONG_PACKET_TYPE {
            return Err(BufError::UnexpectedValue);
        }
        reader.advance(1)?;
        let version = Version::decode(reader, ())?;
        let destination_connection_id = ConnectionId::decode(reader, ())?;
        let source_connection_id = ConnectionId::decode(reader, ())?;
        let mut left = Vec::decode(reader, ())?;
        let length = left.len();
        let (retry_token, _retry_integrity_tag) = left.split_at_mut_checked(length - 16).unwrap();
        let retry_token = RetryToken(retry_token.to_vec());
        let retry_integrity_tag = match _retry_integrity_tag.try_into() {
            Ok(x) => x,
            Err(_) => return Err(BufError::UnexpectedValue),
        };
        let retry_integrity_tag = RetryIntegrityTag(retry_integrity_tag);
        Ok(Self {
            version,
            destination_connection_id,
            source_connection_id,
            retry_token,
            retry_integrity_tag,
        })
    }
}

/// A 1-RTT Packet by [Section 17.3.1].
///
/// Used after the version and 1-RTT keys are negotiated.
///
/// [Section 17.3.1]: https://datatracker.ietf.org/doc/html/rfc9000#section-17.3.1
#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)]
pub struct OneRttPacket {
    /// The Spin Bit.
    pub spin_bit: SpinBit,
    /// The Key Phase.
    pub key_phase: KeyPhase,
    /// The negotiated destination connection ID.
    pub destination_connection_id: ConnectionId,
    /// The number of this packet.
    pub packet_number: PacketNumber,
    /// The Packet Payload.
    pub packet_payload: Vec<u8>,
}

impl OneRttPacket {
    /// The header form.
    pub const HEADER_FORM: HeaderForm = HeaderForm::ShortHeader;
    /// The fixed bit.
    pub const FIXED_BIT: FixedBit = FixedBit::One;
}

impl Codec<u8> for OneRttPacket {
    fn encode<W: BufMut>(&self, writer: &mut Cursor<W>, dcid_length: u8) -> BufResult<()> {
        Self::HEADER_FORM.encode(writer, ())?;
        Self::FIXED_BIT.encode(writer, ())?;
        self.spin_bit.encode(writer, ())?;
        self.key_phase.encode(writer, ())?;
        let packet_number_length = self.packet_number.length()?;
        packet_number_length.encode(writer, ())?;
        self.destination_connection_id.encode(writer, dcid_length)?;
        self.packet_number.encode(writer, packet_number_length)?;
        self.packet_payload.encode(writer, ())
    }

    fn decode<R: Buf>(reader: &mut Cursor<R>, dcid_length: u8) -> BufResult<Self> {
        if HeaderForm::decode(reader, ())? != Self::HEADER_FORM {
            return Err(BufError::UnexpectedValue);
        }
        if FixedBit::decode(reader, ())? != Self::FIXED_BIT {
            return Err(BufError::UnexpectedValue);
        }
        let spin_bit = SpinBit::decode(reader, ())?;
        let key_phase = KeyPhase::decode(reader, ())?;
        let packet_number_length = PacketNumberLength::decode(reader, ())?;
        let destination_connection_id = ConnectionId::decode(reader, dcid_length)?;
        let packet_number = PacketNumber::decode(reader, packet_number_length)?;
        let packet_payload = Vec::decode(reader, ())?;

        Ok(Self {
            spin_bit,
            key_phase,
            destination_connection_id,
            packet_number,
            packet_payload,
        })
    }
}

#[cfg(test)]
mod tests {
    use core::fmt::Debug;

    use crate::{
        Codec, Cursor,
        quic::Version,
        quicv1::{
            ConnectionId, FixedBit, HandshakePacket, InitialPacket, KeyPhase, LongHeaderPacketType,
            OneRttPacket, PacketNumber, PacketNumberLength, RetryIntegrityTag, RetryPacket,
            RetryToken, SpinBit, VariableLengthInteger, ZeroRttPacket,
        },
    };

    fn codec_roundtrip<T: Codec<C> + Debug + Eq, C: Copy>(
        etalon_struct: T,
        etalon_bytes: &[u8],
        context: C,
    ) {
        let mut encoded_bytes = vec![];
        {
            let writer = &mut Cursor::new(&mut encoded_bytes);
            etalon_struct.encode(writer, context).unwrap();
        }
        assert_eq!(etalon_bytes, &encoded_bytes);

        let decoded_struct = {
            let reader = &mut Cursor::new(&mut encoded_bytes);
            T::decode(reader, context).unwrap()
        };
        assert_eq!(etalon_struct, decoded_struct);

        encoded_bytes.fill(0x00);
        {
            let writer = &mut Cursor::new(&mut encoded_bytes);
            decoded_struct.encode(writer, context).unwrap();
        }
        assert_eq!(etalon_bytes, &encoded_bytes);
    }

    #[test]
    fn fixed_bit() {
        let etalon_bytes = &[0b01000000];
        let etalon_struct = FixedBit::One;
        codec_roundtrip(etalon_struct, etalon_bytes, ());

        let etalon_bytes = &[0b00000000];
        let etalon_struct = FixedBit::VersionNegotiationPacket;
        codec_roundtrip(etalon_struct, etalon_bytes, ());
    }

    #[test]
    fn long_header_packet_type() {
        let etalon_bytes = &[0b00000000];
        let etalon_struct = LongHeaderPacketType::Initial;
        codec_roundtrip(etalon_struct, etalon_bytes, ());

        let etalon_bytes = &[0b00010000];
        let etalon_struct = LongHeaderPacketType::ZeroRtt;
        codec_roundtrip(etalon_struct, etalon_bytes, ());

        let etalon_bytes = &[0b00100000];
        let etalon_struct = LongHeaderPacketType::Handshake;
        codec_roundtrip(etalon_struct, etalon_bytes, ());

        let etalon_bytes = &[0b00110000];
        let etalon_struct = LongHeaderPacketType::Retry;
        codec_roundtrip(etalon_struct, etalon_bytes, ());
    }

    #[test]
    fn connection_id() {
        assert_eq!(ConnectionId::new(&[0x08; 21]), None);

        let etalon_bytes = &[0x08; 9];
        let etalon_struct = ConnectionId::new(&[0x08; 8]).unwrap();
        codec_roundtrip(etalon_struct, etalon_bytes, ());

        let etalon_bytes = &[0x010; 17];
        let etalon_struct = ConnectionId::new(&[0x010; 16]).unwrap();
        codec_roundtrip(etalon_struct, etalon_bytes, ());

        let etalon_bytes = &[0x014; 21];
        let etalon_struct = ConnectionId::new(&[0x014; 20]).unwrap();
        codec_roundtrip(etalon_struct, etalon_bytes, ());
    }

    #[test]
    fn packet_number_length() {
        let etalon_bytes = &[0x00];
        let etalon_struct = PacketNumberLength::from_packet_number(PacketNumber(
            VariableLengthInteger::new(0).unwrap(),
        ))
        .unwrap();
        codec_roundtrip(etalon_struct, etalon_bytes, ());
    }

    #[test]
    fn packet_number() {
        let etalon_bytes = &[0x00];
        let etalon_struct = PacketNumber(VariableLengthInteger::new(0).unwrap());
        codec_roundtrip(etalon_struct, etalon_bytes, etalon_struct.length().unwrap());
    }

    #[test]
    fn token() {
        let etalon_bytes = &[
            0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
            0x00, 0x00, // token
            0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
            0x00, 0x00, // retry integrity tag
        ];
        let etalon_struct = RetryToken(vec![0x00; 16]);
        {
            let context = ();
            let mut encoded_bytes = vec![];
            {
                let writer = &mut Cursor::new(&mut encoded_bytes);
                etalon_struct.encode(writer, context).unwrap();
            }
            encoded_bytes.extend_from_slice(&[0x00; 16]);
            assert_eq!(&etalon_bytes.to_vec(), &encoded_bytes);

            let decoded_struct = {
                let reader = &mut Cursor::new(&mut encoded_bytes);
                RetryToken::decode(reader, context).unwrap()
            };
            assert_eq!(etalon_struct, decoded_struct);

            encoded_bytes.fill(0x00);
            {
                let writer = &mut Cursor::new(&mut encoded_bytes);
                decoded_struct.encode(writer, context).unwrap();
            }
            assert_eq!(&etalon_bytes.to_vec(), &encoded_bytes);
        };

        let etalon_bytes = &[
            0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
            0x00, 0x00,
        ];
        let etalon_struct = RetryToken(vec![0x00; 16]);
        codec_roundtrip(
            etalon_struct,
            etalon_bytes,
            VariableLengthInteger::new_const(16),
        );
    }

    #[test]
    fn retry_integrity_tag() {
        let etalon_bytes = &[0; 16];
        let etalon_struct = RetryIntegrityTag([0; 16]);

        codec_roundtrip(etalon_struct, etalon_bytes, ());
    }

    #[test]
    fn spin_bit() {
        let etalon_bytes = &[0b00000000];
        let etalon_struct = SpinBit::Zero;
        codec_roundtrip(etalon_struct, etalon_bytes, ());

        let etalon_bytes = &[0b00100000];
        let etalon_struct = SpinBit::One;
        codec_roundtrip(etalon_struct, etalon_bytes, ());
    }

    #[test]
    fn key_phase() {
        let etalon_bytes = &[0b00000000];
        let etalon_struct = KeyPhase::Zero;
        codec_roundtrip(etalon_struct, etalon_bytes, ());

        let etalon_bytes = &[0b100];
        let etalon_struct = KeyPhase::One;
        codec_roundtrip(etalon_struct, etalon_bytes, ());
    }

    #[test]
    fn initial_packet() {
        let etalon_struct = InitialPacket {
            version: Version(1),
            destination_connection_id: ConnectionId::new(&[0x82; 8]).unwrap(),
            source_connection_id: ConnectionId::new(&[0x41; 8]).unwrap(),
            token: RetryToken(vec![0]),
            packet_number: PacketNumber(VariableLengthInteger::new(0).unwrap()),
            packet_payload: vec![],
        };

        // 0xC0 = Long Header, Initial packet type, packet number length = 1 byte.
        // Length = 0x01, так как payload пустой и packet number занимает 1 байт.
        let etalon_bytes: &[u8] = &[
            0xC0, // Version = 1
            0x00, 0x00, 0x00, 0x01, // Destination Connection ID Length + Value
            0x08, 0x82, 0x82, 0x82, 0x82, 0x82, 0x82, 0x82, 0x82,
            // Source Connection ID Length + Value
            0x08, 0x41, 0x41, 0x41, 0x41, 0x41, 0x41, 0x41, 0x41, // Token Length = 1
            0x01, // Token
            0x00, // Length = Packet Number length + Payload length = 1 + 0
            0x00, // Packet Number
            0x00,
        ];

        codec_roundtrip(etalon_struct, etalon_bytes, ());
    }

    #[test]
    fn zero_rtt_packet() {
        let etalon_struct = ZeroRttPacket {
            version: Version(1),
            destination_connection_id: ConnectionId::new(&[0x82; 8]).unwrap(),
            source_connection_id: ConnectionId::new(&[0x41; 8]).unwrap(),
            packet_number: PacketNumber(VariableLengthInteger::new(0).unwrap()),
            packet_payload: vec![],
        };

        // 0xD0 = Long Header, 0-RTT packet type, packet number length = 1 byte.
        let etalon_bytes: &[u8] = &[
            // First byte
            0xD0, // Version = 1
            0x00, 0x00, 0x00, 0x01, // Destination Connection ID Length + Value
            0x08, 0x82, 0x82, 0x82, 0x82, 0x82, 0x82, 0x82, 0x82,
            // Source Connection ID Length + Value
            0x08, 0x41, 0x41, 0x41, 0x41, 0x41, 0x41, 0x41, 0x41,
            // Length = Packet Number length + Payload length = 1 + 0
            0x00, // Packet Number
            0x00,
        ];

        codec_roundtrip(etalon_struct, etalon_bytes, ());
    }

    #[test]
    fn handshake_packet() {
        let etalon_struct = HandshakePacket {
            version: Version(1),
            destination_connection_id: ConnectionId::new(&[0x82; 8]).unwrap(),
            source_connection_id: ConnectionId::new(&[0x41; 8]).unwrap(),
            packet_number: PacketNumber(VariableLengthInteger::new(0).unwrap()),
            packet_payload: vec![],
        };

        // 0xE0 = Long Header, Handshake packet type, packet number length = 1 byte.
        let etalon_bytes: &[u8] = &[
            // First byte
            0xE0, // Version = 1
            0x00, 0x00, 0x00, 0x01, // Destination Connection ID Length + Value
            0x08, 0x82, 0x82, 0x82, 0x82, 0x82, 0x82, 0x82, 0x82,
            // Source Connection ID Length + Value
            0x08, 0x41, 0x41, 0x41, 0x41, 0x41, 0x41, 0x41, 0x41,
            // Length = Packet Number length + Payload length = 1 + 0
            0x00, 0x00, // Packet Number
        ];

        codec_roundtrip(etalon_struct, etalon_bytes, ());
    }

    #[test]
    fn retry_packet() {
        let etalon_struct = RetryPacket {
            version: Version(1),
            destination_connection_id: ConnectionId::new(&[0x82; 8]).unwrap(),
            source_connection_id: ConnectionId::new(&[0x41; 8]).unwrap(),
            retry_token: RetryToken(vec![0]),
            retry_integrity_tag: RetryIntegrityTag([0x55; 16]),
        };

        let etalon_bytes: &[u8] = &[
            // First byte
            0xF0, // Version = 1
            0x00, 0x00, 0x00, 0x01, // Destination Connection ID Length + Value
            0x08, 0x82, 0x82, 0x82, 0x82, 0x82, 0x82, 0x82, 0x82,
            // Source Connection ID Length + Value
            0x08, 0x41, 0x41, 0x41, 0x41, 0x41, 0x41, 0x41, 0x41, // Retry Token
            0x00, // Retry Integrity Tag, 16 bytes
            0x55, 0x55, 0x55, 0x55, 0x55, 0x55, 0x55, 0x55, 0x55, 0x55, 0x55, 0x55, 0x55, 0x55,
            0x55, 0x55,
        ];

        codec_roundtrip(etalon_struct, etalon_bytes, ());
    }

    #[test]
    fn one_rtt_packet() {
        let etalon_struct = OneRttPacket {
            spin_bit: SpinBit::Zero,
            key_phase: KeyPhase::Zero,
            destination_connection_id: ConnectionId::new(&[0x82; 8]).unwrap(),
            packet_number: PacketNumber(VariableLengthInteger::new(0).unwrap()),
            packet_payload: vec![],
        };

        let etalon_bytes: &[u8] = &[
            0b01000000, // First byte
            0x82, 0x82, 0x82, 0x82, 0x82, 0x82, 0x82, 0x82, // Destination Connection ID
            0x00, // Packet Number
        ];

        codec_roundtrip(etalon_struct, etalon_bytes, 8);
    }
}