internet 0.1.0

//! Packet headers encoding following [Section 17].
//!
//! > description
//!
//! 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,
    ietf::quic::{HeaderForm, Version},
    ietf::quicv1::{
        ConnectionId, FixedBit, KeyPhase, Length, PacketNumber, PacketNumberLength, RetryToken,
        SpinBit, VariableLengthInteger,
    },
};

/// 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 = 0b01,

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

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

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

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 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>()?))
    }
}

/// An Initial Packet following [Section 17.2.2].
///
/// Carries the first [CRYPTO frames] from client and [ACK frames] in either direction.
///
/// [CRYPTO frames]: crate::ietf::quicv1::CryptoFrame
/// [ACK frames]: crate::ietf::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,
        })
    }
}