rustls/

quic.rs

use alloc::boxed::Box;
use alloc::vec::Vec;
use core::ops::{Deref, DerefMut};
use core::{fmt, mem};

use pki_types::{DnsName, FipsStatus, ServerName};

use crate::TlsInputBuffer;
use crate::client::{ClientConfig, ClientSide};
pub use crate::common_state::Side;
use crate::common_state::{CommonState, ConnectionOutputs, Protocol};
use crate::conn::{ConnectionCommon, KeyingMaterialExporter, MessageIter, SideData, StateMachine};
use crate::crypto::cipher::{AeadKey, Iv, Payload};
use crate::crypto::tls13::{Hkdf, HkdfExpander, OkmBlock};
use crate::enums::ApplicationProtocol;
use crate::error::{ApiMisuse, Error};
use crate::msgs::{
    ClientExtensionsInput, Message, MessagePayload, ServerExtensionsInput, TransportParameters,
};
use crate::server::{ChooseConfig, ClientHello, ServerConfig, ServerSide, ServerState};
use crate::suites::SupportedCipherSuite;
use crate::sync::Arc;
use crate::tls13::Tls13CipherSuite;
use crate::tls13::key_schedule::{
    hkdf_expand_label, hkdf_expand_label_aead_key, hkdf_expand_label_block,
};

/// A QUIC client or server connection.
pub trait Connection: fmt::Debug + Deref<Target = ConnectionOutputs> {
    /// Return the TLS-encoded transport parameters for the session's peer.
    ///
    /// While the transport parameters are technically available prior to the
    /// completion of the handshake, they cannot be fully trusted until the
    /// handshake completes, and reliance on them should be minimized.
    /// However, any tampering with the parameters will cause the handshake
    /// to fail.
    fn quic_transport_parameters(&self) -> Option<&[u8]>;

    /// Compute the keys for encrypting/decrypting 0-RTT packets, if available
    fn zero_rtt_keys(&self) -> Option<DirectionalKeys>;

    /// Consume unencrypted TLS handshake data.
    ///
    /// Handshake data obtained from separate encryption levels should be supplied in separate calls.
    ///
    /// How much of the `input` buffer is consumed is recorded by a call to
    /// [`TlsInputBuffer::discard()`].  Unconsumed data should be presented again on the next call.
    fn read_hs(&mut self, input: &mut dyn TlsInputBuffer) -> Result<(), Error>;

    /// Obtain pending events that the caller should process.
    ///
    /// All pending events are returned as an iterator.
    fn events(&mut self) -> impl Iterator<Item = QuicEvent>;

    /// Returns true if the connection is currently performing the TLS handshake.
    fn is_handshaking(&self) -> bool;
}

/// A QUIC client connection.
pub struct ClientConnection {
    inner: QuicCommon<ClientSide>,
}

impl ClientConnection {
    /// Make a new QUIC ClientConnection.
    ///
    /// This differs from `ClientConnection::new()` in that it takes an extra `params` argument,
    /// which contains the TLS-encoded transport parameters to send.
    pub fn new(
        config: Arc<ClientConfig>,
        quic_version: Version,
        name: ServerName<'static>,
        params: Vec<u8>,
    ) -> Result<Self, Error> {
        let alpn_protocols = config.alpn_protocols.clone();
        Self::new_with_alpn(config, quic_version, name, params, alpn_protocols)
    }

    /// Make a new QUIC ClientConnection with custom ALPN protocols.
    pub fn new_with_alpn(
        config: Arc<ClientConfig>,
        version: Version,
        name: ServerName<'static>,
        params: Vec<u8>,
        alpn_protocols: Vec<ApplicationProtocol<'static>>,
    ) -> Result<Self, Error> {
        let suites = &config.provider().tls13_cipher_suites;
        if suites.is_empty() {
            return Err(ApiMisuse::QuicRequiresTls13Support.into());
        }

        if !suites
            .iter()
            .any(|scs| scs.quic.is_some())
        {
            return Err(ApiMisuse::NoQuicCompatibleCipherSuites.into());
        }

        let exts = ClientExtensionsInput {
            transport_parameters: Some(match version {
                Version::V1 | Version::V2 => TransportParameters::Quic(Payload::new(params)),
            }),

            ..ClientExtensionsInput::from_alpn(alpn_protocols)
        };

        let mut quic = Quic {
            version,
            ..Quic::default()
        };

        let mut tls = Vec::new();
        let inner = ConnectionCommon::for_client(
            config,
            name,
            exts,
            Some(&mut quic),
            Protocol::Quic(version),
            &mut tls,
        )?;

        // In QUIC mode, handshake output is emitted via `QuicEvent`s, not `tls`.
        debug_assert!(tls.is_empty());
        Ok(Self {
            inner: QuicCommon::new(inner, quic),
        })
    }

    /// Return the FIPS validation status of the connection.
    pub fn fips(&self) -> FipsStatus {
        self.inner.fips
    }

    /// Returns True if the server signalled it will process early data.
    ///
    /// If you sent early data and this returns false at the end of the
    /// handshake then the server will not process the data.  This
    /// is not an error, but you may wish to resend the data.
    pub fn is_early_data_accepted(&self) -> bool {
        self.inner
            .common
            .is_early_data_accepted()
    }

    /// Returns the number of TLS1.3 tickets that have been received.
    pub fn tls13_tickets_received(&self) -> u32 {
        self.inner
            .common
            .common
            .recv
            .tls13_tickets_received
    }

    /// Returns an object that can derive key material from the agreed connection secrets.
    ///
    /// See [RFC 5705][] for more details on what this is for.
    ///
    /// This function can be called at most once per connection.
    ///
    /// This function will error:
    ///
    /// - if called prior to the handshake completing; (check with
    ///   [`CommonState::is_handshaking`] first).
    /// - if called more than once per connection.
    ///
    /// [RFC 5705]: https://datatracker.ietf.org/doc/html/rfc5705
    pub fn exporter(&mut self) -> Result<KeyingMaterialExporter, Error> {
        self.inner.common.exporter()
    }
}

impl Connection for ClientConnection {
    fn quic_transport_parameters(&self) -> Option<&[u8]> {
        self.inner.quic_transport_parameters()
    }

    fn zero_rtt_keys(&self) -> Option<DirectionalKeys> {
        self.inner.zero_rtt_keys()
    }

    fn read_hs(&mut self, input: &mut dyn TlsInputBuffer) -> Result<(), Error> {
        self.inner.read_hs(input)
    }

    fn events(&mut self) -> impl Iterator<Item = QuicEvent> {
        self.inner.events()
    }

    fn is_handshaking(&self) -> bool {
        self.inner.is_handshaking()
    }
}

impl Deref for ClientConnection {
    type Target = ConnectionOutputs;

    fn deref(&self) -> &Self::Target {
        &self.inner
    }
}

impl fmt::Debug for ClientConnection {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.debug_struct("quic::ClientConnection")
            .finish_non_exhaustive()
    }
}

/// A QUIC server connection.
pub struct ServerConnection {
    inner: QuicCommon<ServerSide>,
}

impl ServerConnection {
    /// Make a new QUIC ServerConnection.
    ///
    /// This differs from `ServerConnection::new()` in that it takes an extra `params` argument,
    /// which contains the TLS-encoded transport parameters to send.
    pub fn new(
        config: Arc<ServerConfig>,
        version: Version,
        params: Vec<u8>,
    ) -> Result<Self, Error> {
        check_server_config(&config)?;
        let exts = ServerExtensionsInput {
            transport_parameters: Some(match version {
                Version::V1 | Version::V2 => TransportParameters::Quic(Payload::new(params)),
            }),
        };

        let core = ConnectionCommon::for_server(config, exts, Protocol::Quic(version))?;
        let inner = QuicCommon::new(
            core,
            Quic {
                version,
                ..Quic::default()
            },
        );
        Ok(Self { inner })
    }

    /// Return the FIPS validation status of the connection.
    pub fn fips(&self) -> FipsStatus {
        self.inner.fips
    }

    /// Retrieves the server name, if any, used to select the certificate and
    /// private key.
    ///
    /// This returns `None` until some time after the client's server name indication
    /// (SNI) extension value is processed during the handshake. It will never be
    /// `None` when the connection is ready to send or process application data,
    /// unless the client does not support SNI.
    ///
    /// This is useful for application protocols that need to enforce that the
    /// server name matches an application layer protocol hostname. For
    /// example, HTTP/1.1 servers commonly expect the `Host:` header field of
    /// every request on a connection to match the hostname in the SNI extension
    /// when the client provides the SNI extension.
    ///
    /// The server name is also used to match sessions during session resumption.
    pub fn server_name(&self) -> Option<&DnsName<'_>> {
        self.inner.common.side.server_name()
    }

    /// Set the resumption data to embed in future resumption tickets supplied to the client.
    ///
    /// Defaults to the empty byte string. Must be less than 2^15 bytes to allow room for other
    /// data. Should be called while `is_handshaking` returns true to ensure all transmitted
    /// resumption tickets are affected (otherwise an error will be returned).
    ///
    /// Integrity will be assured by rustls, but the data will be visible to the client. If secrecy
    /// from the client is desired, encrypt the data separately.
    pub fn set_resumption_data(&mut self, resumption_data: &[u8]) -> Result<(), Error> {
        assert!(resumption_data.len() < 2usize.pow(15));
        match &mut self.inner.common.state {
            Ok(st) => st.set_resumption_data(resumption_data),
            Err(e) => Err(e.clone()),
        }
    }

    /// Retrieves the resumption data supplied by the client, if any.
    ///
    /// Returns `Some` if and only if a valid resumption ticket has been received from the client.
    pub fn received_resumption_data(&self) -> Option<&[u8]> {
        self.inner
            .common
            .side
            .received_resumption_data()
    }

    /// Returns an object that can derive key material from the agreed connection secrets.
    ///
    /// See [RFC 5705][] for more details on what this is for.
    ///
    /// This function can be called at most once per connection.
    ///
    /// This function will error:
    ///
    /// - if called prior to the handshake completing; (check with
    ///   [`CommonState::is_handshaking`] first).
    /// - if called more than once per connection.
    ///
    /// [RFC 5705]: https://datatracker.ietf.org/doc/html/rfc5705
    pub fn exporter(&mut self) -> Result<KeyingMaterialExporter, Error> {
        self.inner.common.exporter()
    }
}

impl Connection for ServerConnection {
    fn quic_transport_parameters(&self) -> Option<&[u8]> {
        self.inner.quic_transport_parameters()
    }

    fn zero_rtt_keys(&self) -> Option<DirectionalKeys> {
        self.inner.zero_rtt_keys()
    }

    fn read_hs(&mut self, input: &mut dyn TlsInputBuffer) -> Result<(), Error> {
        self.inner.read_hs(input)
    }

    fn events(&mut self) -> impl Iterator<Item = QuicEvent> {
        self.inner.events()
    }

    fn is_handshaking(&self) -> bool {
        self.inner.is_handshaking()
    }
}

impl Deref for ServerConnection {
    type Target = ConnectionOutputs;

    fn deref(&self) -> &Self::Target {
        &self.inner
    }
}

impl fmt::Debug for ServerConnection {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.debug_struct("quic::ServerConnection")
            .finish_non_exhaustive()
    }
}

/// An in-progress TLS server handshake.
#[non_exhaustive]
#[derive(Debug)]
pub enum ServerHandshake {
    /// More data needs to be received to make progress.
    NeedsInput(NeedsInput),

    /// A complete `ClientHello` has been received.
    ///
    /// The handshake can be progressed by choosing a [`ServerConfig`] based on
    /// [`Accepted::client_hello()`] and providing it to [`Accepted::choose_config()`].
    Accepted(Accepted),

    /// The handshake is complete.
    Complete(ServerConnection),
}

impl ServerHandshake {
    /// Creates a new QUIC [`ServerHandshake`] via the payload of the [`ServerHandshake::NeedsInput`] variant.
    ///
    /// It is a fundamental fact of server TLS connections that the server reads first; this is reflected
    /// in the returned type.
    ///
    /// You may wrap this in the [`ServerHandshake::NeedsInput`] variant to generalise the type to a
    /// [`ServerHandshake`].
    ///
    /// The returned object should be fed data from a single potential client.
    pub fn start(version: Version) -> NeedsInput {
        NeedsInput {
            inner: QuicCommon::new(
                ConnectionCommon::for_acceptor(Protocol::Quic(version)),
                Quic {
                    version,
                    ..Quic::default()
                },
            ),
        }
    }
}

impl TryFrom<QuicCommon<ServerSide>> for ServerHandshake {
    type Error = Error;

    fn try_from(mut inner: QuicCommon<ServerSide>) -> Result<Self, Error> {
        const MISUSED: Error = Error::Unreachable("forgot to restore state");

        Ok(match mem::replace(&mut inner.common.state, Err(MISUSED))? {
            ServerState::ChooseConfig(choose_config) => Self::Accepted(Accepted {
                inner,
                choose_config,
            }),

            state if state.is_traffic() => {
                inner.common.state = Ok(state);
                Self::Complete(ServerConnection { inner })
            }

            state => {
                inner.common.state = Ok(state);
                Self::NeedsInput(NeedsInput { inner })
            }
        })
    }
}

/// More data needs to be received to make progress.
///
/// Provide the data to [`Self::process()`].
pub struct NeedsInput {
    inner: QuicCommon<ServerSide>,
}

impl NeedsInput {
    /// Progress the handshake by receiving further unencrypted TLS handshake data.
    ///
    /// The input should be ordered QUIC CRYPTO stream data for one encryption level.
    ///
    /// Handshake data obtained from separate encryption levels should be supplied in separate calls.
    ///
    /// How much of the `input` buffer is consumed is recorded by a call to
    /// [`TlsInputBuffer::discard()`].  Unconsumed data should be presented again on the next call.
    ///
    /// An error from this function is fatal to the connection, as it consumes the [`NeedsInput`]
    /// object.
    ///
    /// On success, this returns:
    ///
    /// - a [`ServerHandshake::NeedsInput`] if more data is required.
    /// - a [`ServerHandshake::Accepted`] if a whole `ClientHello` has been received,
    ///   and a choice of [`ServerConfig`] is required to continue.
    /// - a [`ServerHandshake::Complete`] if the handshake is complete.
    ///
    /// `output` has any resulting handshake messages or key changes appended to it.
    pub fn process(
        mut self,
        input: &mut dyn TlsInputBuffer,
        output: &mut Vec<QuicEvent>,
    ) -> Result<ServerHandshake, Error> {
        self.inner.read_hs(input)?;
        output.extend(self.inner.events());
        ServerHandshake::try_from(self.inner)
    }
}

impl fmt::Debug for NeedsInput {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.debug_struct("quic::NeedsInput")
            .finish_non_exhaustive()
    }
}

/// Represents that a `ClientHello` message has been received.
///
/// The handshake can be progressed by choosing a [`ServerConfig`] based on
/// [`Accepted::client_hello()`] and providing it to [`Accepted::choose_config()`].
pub struct Accepted {
    // invariant: `inner.core.state` is `Err(_)` and requires restoring
    inner: QuicCommon<ServerSide>,
    choose_config: Box<ChooseConfig>,
}

impl Accepted {
    /// Get the [`ClientHello`] for this connection.
    pub fn client_hello(&self) -> ClientHello<'_> {
        self.choose_config.client_hello()
    }

    /// Choose a [`ServerConfig`] to progress the handshake.
    ///
    /// Resolves an [`Accepted`], providing the [`ServerConfig`] that should be used for
    /// the session, and the TLS-encoded QUIC transport parameters to send.
    ///
    /// Returns an error if configuration-dependent validation of the received
    /// `ClientHello` message fails.
    ///
    /// Events are appended to `output`.
    pub fn choose_config(
        mut self,
        config: Arc<ServerConfig>,
        params: Vec<u8>,
        output: &mut Vec<QuicEvent>,
    ) -> Result<ServerHandshake, Error> {
        check_server_config(&config)?;

        let mut tls = Vec::new();
        self.inner.common.accepted(
            self.choose_config,
            ServerExtensionsInput {
                transport_parameters: Some(match self.inner.quic.version {
                    Version::V1 | Version::V2 => TransportParameters::Quic(Payload::new(params)),
                }),
            },
            Some(&mut self.inner.quic),
            config,
            &mut tls,
        )?;

        // In QUIC mode, handshake output is emitted via `QuicEvent`s, not `tls`.
        debug_assert!(tls.is_empty());
        output.extend(self.inner.events());

        ServerHandshake::try_from(self.inner)
    }
}

impl fmt::Debug for Accepted {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.debug_struct("quic::Accepted")
            .finish_non_exhaustive()
    }
}

fn check_server_config(config: &ServerConfig) -> Result<(), Error> {
    let suites = &config.provider.tls13_cipher_suites;
    if suites.is_empty() {
        return Err(ApiMisuse::QuicRequiresTls13Support.into());
    }

    if !suites
        .iter()
        .any(|scs| scs.quic.is_some())
    {
        return Err(ApiMisuse::NoQuicCompatibleCipherSuites.into());
    }

    if config.max_early_data_size != 0 && config.max_early_data_size != 0xffff_ffff {
        return Err(ApiMisuse::QuicRestrictsMaxEarlyDataSize.into());
    }

    Ok(())
}

/// QUIC events that should be handled by the caller.
#[expect(clippy::large_enum_variant)]
#[derive(Debug)]
#[non_exhaustive]
pub enum QuicEvent {
    /// These bytes should be handled as an unencrypted TLS handshake message.
    Message(Vec<u8>),

    /// The key material should be changed.
    KeyChange(KeyChange),
}

/// A shared interface for QUIC connections.
struct QuicCommon<Side: SideData> {
    common: ConnectionCommon<Side>,
    quic: Quic,
}

impl<Side: SideData> QuicCommon<Side> {
    fn new(common: ConnectionCommon<Side>, quic: Quic) -> Self {
        Self { common, quic }
    }

    fn quic_transport_parameters(&self) -> Option<&[u8]> {
        self.quic
            .params
            .as_ref()
            .map(|v| v.as_ref())
    }

    fn zero_rtt_keys(&self) -> Option<DirectionalKeys> {
        let suite = self
            .common
            .common
            .negotiated_cipher_suite()
            .and_then(|suite| match suite {
                SupportedCipherSuite::Tls13(suite) => Some(suite),
                _ => None,
            })?;

        let suite = Suite {
            inner: suite,
            quic: suite.quic?,
        };

        Some(DirectionalKeys::new(
            suite,
            self.quic.early_secret.as_ref()?,
            self.quic.version,
        ))
    }

    fn read_hs(&mut self, input: &mut dyn TlsInputBuffer) -> Result<(), Error> {
        self.common
            .common
            .recv
            .deframer
            .input_quic(input.slice_mut())?;

        let mut tls = Vec::new();
        let mut iter = MessageIter::new(input, &mut tls, Some(&mut self.quic), &mut self.common);
        let result = match iter.next() {
            Some(Ok(_)) | None => Ok(()),
            Some(Err(e)) => Err(e),
        };

        input.discard(
            self.common
                .common
                .recv
                .deframer
                .take_discard(),
        );

        result
    }

    fn events(&mut self) -> impl Iterator<Item = QuicEvent> {
        self.quic.events()
    }
}

impl<Side: SideData> Deref for QuicCommon<Side> {
    type Target = CommonState;

    fn deref(&self) -> &Self::Target {
        &self.common.common
    }
}

impl<Side: SideData> DerefMut for QuicCommon<Side> {
    fn deref_mut(&mut self) -> &mut Self::Target {
        &mut self.common.common
    }
}

#[derive(Default)]
pub(crate) struct Quic {
    pub(crate) version: Version,
    /// QUIC transport parameters received from the peer during the handshake
    pub(crate) params: Option<Vec<u8>>,
    pub(crate) events: Vec<QuicEvent>,
    pub(crate) early_secret: Option<OkmBlock>,
}

impl Quic {
    pub(crate) fn send_msg(&mut self, m: Message<'_>, _must_encrypt: bool) {
        if let MessagePayload::Alert(_) = m.payload {
            // alerts are sent out-of-band in QUIC mode
            return;
        }

        debug_assert!(
            matches!(
                m.payload,
                MessagePayload::Handshake { .. } | MessagePayload::HandshakeFlight(_)
            ),
            "QUIC uses TLS for the cryptographic handshake only"
        );
        let mut bytes = Vec::new();
        m.payload.encode(&mut bytes);
        self.events
            .push(QuicEvent::Message(bytes));
    }

    pub(crate) fn events(&mut self) -> impl Iterator<Item = QuicEvent> {
        mem::take(&mut self.events).into_iter()
    }
}

impl QuicOutput for Quic {
    fn transport_parameters(&mut self, params: Vec<u8>) {
        self.params = Some(params);
    }

    fn early_secret(&mut self, secret: Option<OkmBlock>) {
        self.early_secret = secret;
    }

    fn handshake_secrets(
        &mut self,
        client_secret: OkmBlock,
        server_secret: OkmBlock,
        suite: Suite,
        side: Side,
    ) {
        self.events
            .push(QuicEvent::KeyChange(KeyChange::Handshake {
                keys: Keys::new(&Secrets::new(
                    client_secret,
                    server_secret,
                    suite,
                    side,
                    self.version,
                )),
            }));
    }

    fn traffic_secrets(
        &mut self,
        client_secret: OkmBlock,
        server_secret: OkmBlock,
        suite: Suite,
        side: Side,
    ) {
        let mut secrets = Secrets::new(client_secret, server_secret, suite, side, self.version);
        let keys = Keys::new(&secrets);
        secrets.update();
        self.events
            .push(QuicEvent::KeyChange(KeyChange::OneRtt {
                keys,
                next: secrets,
            }));
    }

    fn send_msg(&mut self, m: Message<'_>, must_encrypt: bool) {
        self.send_msg(m, must_encrypt);
    }
}

pub(crate) trait QuicOutput {
    fn transport_parameters(&mut self, params: Vec<u8>);

    fn early_secret(&mut self, secret: Option<OkmBlock>);

    fn handshake_secrets(
        &mut self,
        client_secret: OkmBlock,
        server_secret: OkmBlock,
        suite: Suite,
        side: Side,
    );

    fn traffic_secrets(
        &mut self,
        client_secret: OkmBlock,
        server_secret: OkmBlock,
        suite: Suite,
        side: Side,
    );

    fn send_msg(&mut self, m: Message<'_>, must_encrypt: bool);
}

/// Secrets used to encrypt/decrypt traffic
#[derive(Clone)]
pub struct Secrets {
    /// Secret used to encrypt packets transmitted by the client
    pub(crate) client: OkmBlock,
    /// Secret used to encrypt packets transmitted by the server
    pub(crate) server: OkmBlock,
    /// Cipher suite used with these secrets
    suite: Suite,
    side: Side,
    version: Version,
}

impl Secrets {
    pub(crate) fn new(
        client: OkmBlock,
        server: OkmBlock,
        suite: Suite,
        side: Side,
        version: Version,
    ) -> Self {
        Self {
            client,
            server,
            suite,
            side,
            version,
        }
    }

    /// Derive the next set of packet keys
    pub fn next_packet_keys(&mut self) -> PacketKeySet {
        let keys = PacketKeySet::new(self);
        self.update();
        keys
    }

    pub(crate) fn update(&mut self) {
        self.client = hkdf_expand_label_block(
            self.suite
                .inner
                .hkdf_provider
                .expander_for_okm(&self.client)
                .as_ref(),
            self.version.key_update_label(),
            &[],
        );
        self.server = hkdf_expand_label_block(
            self.suite
                .inner
                .hkdf_provider
                .expander_for_okm(&self.server)
                .as_ref(),
            self.version.key_update_label(),
            &[],
        );
    }

    fn local_remote(&self) -> (&OkmBlock, &OkmBlock) {
        match self.side {
            Side::Client => (&self.client, &self.server),
            Side::Server => (&self.server, &self.client),
        }
    }
}

/// Keys used to communicate in a single direction
#[expect(clippy::exhaustive_structs)]
pub struct DirectionalKeys {
    /// Encrypts or decrypts a packet's headers
    pub header: Box<dyn HeaderProtectionKey>,
    /// Encrypts or decrypts the payload of a packet
    pub packet: Box<dyn PacketKey>,
}

impl DirectionalKeys {
    pub(crate) fn new(suite: Suite, secret: &OkmBlock, version: Version) -> Self {
        let builder = KeyBuilder::new(secret, version, suite.quic, suite.inner.hkdf_provider);
        Self {
            header: builder.header_protection_key(),
            packet: builder.packet_key(),
        }
    }
}

/// All AEADs we support have 16-byte tags.
const TAG_LEN: usize = 16;

/// Authentication tag from an AEAD seal operation.
pub struct Tag([u8; TAG_LEN]);

impl From<&[u8]> for Tag {
    fn from(value: &[u8]) -> Self {
        let mut array = [0u8; TAG_LEN];
        array.copy_from_slice(value);
        Self(array)
    }
}

impl AsRef<[u8]> for Tag {
    fn as_ref(&self) -> &[u8] {
        &self.0
    }
}

/// How a `Tls13CipherSuite` generates `PacketKey`s and `HeaderProtectionKey`s.
pub trait Algorithm: Send + Sync {
    /// Produce a `PacketKey` encrypter/decrypter for this suite.
    ///
    /// `suite` is the entire suite this `Algorithm` appeared in.
    /// `key` and `iv` is the key material to use.
    fn packet_key(&self, key: AeadKey, iv: Iv) -> Box<dyn PacketKey>;

    /// Produce a `HeaderProtectionKey` encrypter/decrypter for this suite.
    ///
    /// `key` is the key material, which is `aead_key_len()` bytes in length.
    fn header_protection_key(&self, key: AeadKey) -> Box<dyn HeaderProtectionKey>;

    /// The length in bytes of keys for this Algorithm.
    ///
    /// This controls the size of `AeadKey`s presented to `packet_key()` and `header_protection_key()`.
    fn aead_key_len(&self) -> usize;

    /// Whether this algorithm is FIPS-approved.
    fn fips(&self) -> FipsStatus {
        FipsStatus::Unvalidated
    }
}

/// A QUIC header protection key
pub trait HeaderProtectionKey: Send + Sync {
    /// Adds QUIC Header Protection.
    ///
    /// `sample` must contain the sample of encrypted payload; see
    /// [Header Protection Sample].
    ///
    /// `first` must reference the first byte of the header, referred to as
    /// `packet[0]` in [Header Protection Application].
    ///
    /// `packet_number` must reference the Packet Number field; this is
    /// `packet[pn_offset:pn_offset+pn_length]` in [Header Protection Application].
    ///
    /// Returns an error without modifying anything if `sample` is not
    /// the correct length (see [Header Protection Sample] and [`Self::sample_len()`]),
    /// or `packet_number` is longer than allowed (see [Packet Number Encoding and Decoding]).
    ///
    /// Otherwise, `first` and `packet_number` will have the header protection added.
    ///
    /// [Header Protection Application]: https://datatracker.ietf.org/doc/html/rfc9001#section-5.4.1
    /// [Header Protection Sample]: https://datatracker.ietf.org/doc/html/rfc9001#section-5.4.2
    /// [Packet Number Encoding and Decoding]: https://datatracker.ietf.org/doc/html/rfc9000#section-17.1
    fn encrypt_in_place(
        &self,
        sample: &[u8],
        first: &mut u8,
        packet_number: &mut [u8],
    ) -> Result<(), Error>;

    /// Removes QUIC Header Protection.
    ///
    /// `sample` must contain the sample of encrypted payload; see
    /// [Header Protection Sample].
    ///
    /// `first` must reference the first byte of the header, referred to as
    /// `packet[0]` in [Header Protection Application].
    ///
    /// `packet_number` must reference the Packet Number field; this is
    /// `packet[pn_offset:pn_offset+pn_length]` in [Header Protection Application].
    ///
    /// Returns an error without modifying anything if `sample` is not
    /// the correct length (see [Header Protection Sample] and [`Self::sample_len()`]),
    /// or `packet_number` is longer than allowed (see
    /// [Packet Number Encoding and Decoding]).
    ///
    /// Otherwise, `first` and `packet_number` will have the header protection removed.
    ///
    /// [Header Protection Application]: https://datatracker.ietf.org/doc/html/rfc9001#section-5.4.1
    /// [Header Protection Sample]: https://datatracker.ietf.org/doc/html/rfc9001#section-5.4.2
    /// [Packet Number Encoding and Decoding]: https://datatracker.ietf.org/doc/html/rfc9000#section-17.1
    fn decrypt_in_place(
        &self,
        sample: &[u8],
        first: &mut u8,
        packet_number: &mut [u8],
    ) -> Result<(), Error>;

    /// Expected sample length for the key's algorithm
    fn sample_len(&self) -> usize;
}

/// Keys to encrypt or decrypt the payload of a packet
pub trait PacketKey: Send + Sync {
    /// Encrypt a QUIC packet
    ///
    /// Takes a `packet_number` and optional `path_id`, used to derive the nonce; the packet
    /// `header`, which is used as the additional authenticated data; and the `payload`. The
    /// authentication tag is returned if encryption succeeds.
    ///
    /// Fails if and only if the payload is longer than allowed by the cipher suite's AEAD algorithm.
    ///
    /// When provided, the `path_id` is used for multipath encryption as described in
    /// <https://www.ietf.org/archive/id/draft-ietf-quic-multipath-15.html#section-2.4>.
    fn encrypt_in_place(
        &self,
        packet_number: u64,
        header: &[u8],
        payload: &mut [u8],
        path_id: Option<u32>,
    ) -> Result<Tag, Error>;

    /// Decrypt a QUIC packet
    ///
    /// Takes a `packet_number` and optional `path_id`, used to derive the nonce; the packet
    /// `header`, which is used as the additional authenticated data, and the `payload`, which
    /// includes the authentication tag.
    ///
    /// On success, returns the slice of `payload` containing the decrypted data.
    ///
    /// When provided, the `path_id` is used for multipath encryption as described in
    /// <https://www.ietf.org/archive/id/draft-ietf-quic-multipath-15.html#section-2.4>.
    fn decrypt_in_place<'a>(
        &self,
        packet_number: u64,
        header: &[u8],
        payload: &'a mut [u8],
        path_id: Option<u32>,
    ) -> Result<&'a [u8], Error>;

    /// Tag length for the underlying AEAD algorithm
    fn tag_len(&self) -> usize;

    /// Number of QUIC messages that can be safely encrypted with a single key of this type.
    ///
    /// Once a `MessageEncrypter` produced for this suite has encrypted more than
    /// `confidentiality_limit` messages, an attacker gains an advantage in distinguishing it
    /// from an ideal pseudorandom permutation (PRP).
    ///
    /// This is to be set on the assumption that messages are maximally sized --
    /// 2 ** 16. For non-QUIC TCP connections see [`CipherSuiteCommon::confidentiality_limit`][csc-limit].
    ///
    /// [csc-limit]: crate::crypto::CipherSuiteCommon::confidentiality_limit
    fn confidentiality_limit(&self) -> u64;

    /// Number of QUIC messages that can be safely decrypted with a single key of this type
    ///
    /// Once a `MessageDecrypter` produced for this suite has failed to decrypt `integrity_limit`
    /// messages, an attacker gains an advantage in forging messages.
    ///
    /// This is not relevant for TLS over TCP (which is also implemented in this crate)
    /// because a single failed decryption is fatal to the connection.
    /// However, this quantity is used by QUIC.
    fn integrity_limit(&self) -> u64;
}

/// Packet protection keys for bidirectional 1-RTT communication
#[expect(clippy::exhaustive_structs)]
pub struct PacketKeySet {
    /// Encrypts outgoing packets
    pub local: Box<dyn PacketKey>,
    /// Decrypts incoming packets
    pub remote: Box<dyn PacketKey>,
}

impl PacketKeySet {
    fn new(secrets: &Secrets) -> Self {
        let (local, remote) = secrets.local_remote();
        let (version, alg, hkdf) = (
            secrets.version,
            secrets.suite.quic,
            secrets.suite.inner.hkdf_provider,
        );

        Self {
            local: KeyBuilder::new(local, version, alg, hkdf).packet_key(),
            remote: KeyBuilder::new(remote, version, alg, hkdf).packet_key(),
        }
    }
}

/// Helper for building QUIC packet and header protection keys
pub struct KeyBuilder<'a> {
    expander: Box<dyn HkdfExpander>,
    version: Version,
    alg: &'a dyn Algorithm,
}

impl<'a> KeyBuilder<'a> {
    /// Create a new KeyBuilder
    pub fn new(
        secret: &OkmBlock,
        version: Version,
        alg: &'a dyn Algorithm,
        hkdf: &'a dyn Hkdf,
    ) -> Self {
        Self {
            expander: hkdf.expander_for_okm(secret),
            version,
            alg,
        }
    }

    /// Derive packet keys
    pub fn packet_key(&self) -> Box<dyn PacketKey> {
        let aead_key_len = self.alg.aead_key_len();
        let packet_key = hkdf_expand_label_aead_key(
            self.expander.as_ref(),
            aead_key_len,
            self.version.packet_key_label(),
            &[],
        );

        let packet_iv =
            hkdf_expand_label(self.expander.as_ref(), self.version.packet_iv_label(), &[]);
        self.alg
            .packet_key(packet_key, packet_iv)
    }

    /// Derive header protection keys
    pub fn header_protection_key(&self) -> Box<dyn HeaderProtectionKey> {
        let header_key = hkdf_expand_label_aead_key(
            self.expander.as_ref(),
            self.alg.aead_key_len(),
            self.version.header_key_label(),
            &[],
        );
        self.alg
            .header_protection_key(header_key)
    }
}

/// Produces QUIC initial keys from a TLS 1.3 ciphersuite and a QUIC key generation algorithm.
#[non_exhaustive]
#[derive(Clone, Copy)]
pub struct Suite {
    /// The TLS 1.3 ciphersuite used to derive keys.
    pub inner: &'static Tls13CipherSuite,
    /// The QUIC key generation algorithm used to derive keys.
    pub quic: &'static dyn Algorithm,
}

impl Suite {
    /// Produce a set of initial keys given the connection ID, side and version
    pub fn keys(&self, client_dst_connection_id: &[u8], side: Side, version: Version) -> Keys {
        Keys::initial(version, *self, client_dst_connection_id, side)
    }
}

impl TryFrom<&'static Tls13CipherSuite> for Suite {
    type Error = ApiMisuse;

    fn try_from(suite: &'static Tls13CipherSuite) -> Result<Self, Self::Error> {
        Ok(Self {
            inner: suite,
            quic: suite
                .quic
                .ok_or(ApiMisuse::NoQuicCompatibleCipherSuites)?,
        })
    }
}

impl fmt::Debug for Suite {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.debug_struct("Suite")
            .field("inner", &self.inner)
            .finish_non_exhaustive()
    }
}

/// Complete set of keys used to communicate with the peer
#[expect(clippy::exhaustive_structs)]
pub struct Keys {
    /// Encrypts outgoing packets
    pub local: DirectionalKeys,
    /// Decrypts incoming packets
    pub remote: DirectionalKeys,
}

impl Keys {
    /// Construct keys for use with initial packets
    pub fn initial(
        version: Version,
        suite: Suite,
        client_dst_connection_id: &[u8],
        side: Side,
    ) -> Self {
        const CLIENT_LABEL: &[u8] = b"client in";
        const SERVER_LABEL: &[u8] = b"server in";
        let salt = version.initial_salt();
        let hs_secret = suite
            .inner
            .hkdf_provider
            .extract_from_secret(Some(salt), client_dst_connection_id);

        let secrets = Secrets {
            client: hkdf_expand_label_block(hs_secret.as_ref(), CLIENT_LABEL, &[]),
            server: hkdf_expand_label_block(hs_secret.as_ref(), SERVER_LABEL, &[]),
            suite,
            side,
            version,
        };

        Self::new(&secrets)
    }

    fn new(secrets: &Secrets) -> Self {
        let (local, remote) = secrets.local_remote();
        Self {
            local: DirectionalKeys::new(secrets.suite, local, secrets.version),
            remote: DirectionalKeys::new(secrets.suite, remote, secrets.version),
        }
    }
}

/// Key material for use in QUIC packet spaces
///
/// QUIC uses 4 different sets of keys (and progressive key updates for long-running connections):
///
/// * Initial: these can be created from [`Keys::initial()`]
/// * 0-RTT keys: can be retrieved from [`Connection::zero_rtt_keys()`]
/// * Handshake: these are returned from [`Connection::events()`] after `ClientHello` and
///   `ServerHello` messages have been exchanged
/// * 1-RTT keys: these are returned from [`Connection::events()`] after the handshake is done
///
/// Once the 1-RTT keys have been exchanged, either side may initiate a key update. Progressive
/// update keys can be obtained from the [`Secrets`] returned in [`KeyChange::OneRtt`]. Note that
/// only packet keys are updated by key updates; header protection keys remain the same.
#[expect(clippy::exhaustive_enums)]
pub enum KeyChange {
    /// Keys for the handshake space
    Handshake {
        /// Header and packet keys for the handshake space
        keys: Keys,
    },
    /// Keys for 1-RTT data
    OneRtt {
        /// Header and packet keys for 1-RTT data
        keys: Keys,
        /// Secrets to derive updated keys from
        next: Secrets,
    },
}

impl fmt::Debug for KeyChange {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            Self::Handshake { .. } => f
                .debug_struct("Handshake")
                .finish_non_exhaustive(),
            Self::OneRtt { .. } => f
                .debug_struct("OneRtt")
                .finish_non_exhaustive(),
        }
    }
}

/// QUIC protocol version
///
/// Governs version-specific behavior in the TLS layer
#[non_exhaustive]
#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)]
pub enum Version {
    /// First stable RFC
    #[default]
    V1,
    /// Anti-ossification variant of V1
    V2,
}

impl Version {
    fn initial_salt(self) -> &'static [u8; 20] {
        match self {
            Self::V1 => &[
                // https://www.rfc-editor.org/rfc/rfc9001.html#name-initial-secrets
                0x38, 0x76, 0x2c, 0xf7, 0xf5, 0x59, 0x34, 0xb3, 0x4d, 0x17, 0x9a, 0xe6, 0xa4, 0xc8,
                0x0c, 0xad, 0xcc, 0xbb, 0x7f, 0x0a,
            ],
            Self::V2 => &[
                // https://tools.ietf.org/html/rfc9369.html#name-initial-salt
                0x0d, 0xed, 0xe3, 0xde, 0xf7, 0x00, 0xa6, 0xdb, 0x81, 0x93, 0x81, 0xbe, 0x6e, 0x26,
                0x9d, 0xcb, 0xf9, 0xbd, 0x2e, 0xd9,
            ],
        }
    }

    /// Key derivation label for packet keys.
    pub(crate) fn packet_key_label(&self) -> &'static [u8] {
        match self {
            Self::V1 => b"quic key",
            Self::V2 => b"quicv2 key",
        }
    }

    /// Key derivation label for packet "IV"s.
    pub(crate) fn packet_iv_label(&self) -> &'static [u8] {
        match self {
            Self::V1 => b"quic iv",
            Self::V2 => b"quicv2 iv",
        }
    }

    /// Key derivation for header keys.
    pub(crate) fn header_key_label(&self) -> &'static [u8] {
        match self {
            Self::V1 => b"quic hp",
            Self::V2 => b"quicv2 hp",
        }
    }

    fn key_update_label(&self) -> &'static [u8] {
        match self {
            Self::V1 => b"quic ku",
            Self::V2 => b"quicv2 ku",
        }
    }
}

#[cfg(all(test, any(target_arch = "aarch64", target_arch = "x86_64")))]
mod tests {
    use super::*;
    use crate::crypto::TLS13_TEST_SUITE;
    use crate::crypto::tls13::OkmBlock;
    use crate::quic::{HeaderProtectionKey, Secrets, Side, Version};

    #[test]
    fn key_update_test_vector() {
        fn equal_okm(x: &OkmBlock, y: &OkmBlock) -> bool {
            x.as_ref() == y.as_ref()
        }

        let mut secrets = Secrets {
            // Constant dummy values for reproducibility
            client: OkmBlock::new(
                &[
                    0xb8, 0x76, 0x77, 0x08, 0xf8, 0x77, 0x23, 0x58, 0xa6, 0xea, 0x9f, 0xc4, 0x3e,
                    0x4a, 0xdd, 0x2c, 0x96, 0x1b, 0x3f, 0x52, 0x87, 0xa6, 0xd1, 0x46, 0x7e, 0xe0,
                    0xae, 0xab, 0x33, 0x72, 0x4d, 0xbf,
                ][..],
            ),
            server: OkmBlock::new(
                &[
                    0x42, 0xdc, 0x97, 0x21, 0x40, 0xe0, 0xf2, 0xe3, 0x98, 0x45, 0xb7, 0x67, 0x61,
                    0x34, 0x39, 0xdc, 0x67, 0x58, 0xca, 0x43, 0x25, 0x9b, 0x87, 0x85, 0x06, 0x82,
                    0x4e, 0xb1, 0xe4, 0x38, 0xd8, 0x55,
                ][..],
            ),
            suite: Suite {
                inner: TLS13_TEST_SUITE,
                quic: &FakeAlgorithm,
            },
            side: Side::Client,
            version: Version::V1,
        };
        secrets.update();

        assert!(equal_okm(
            &secrets.client,
            &OkmBlock::new(
                &[
                    0x42, 0xca, 0xc8, 0xc9, 0x1c, 0xd5, 0xeb, 0x40, 0x68, 0x2e, 0x43, 0x2e, 0xdf,
                    0x2d, 0x2b, 0xe9, 0xf4, 0x1a, 0x52, 0xca, 0x6b, 0x22, 0xd8, 0xe6, 0xcd, 0xb1,
                    0xe8, 0xac, 0xa9, 0x6, 0x1f, 0xce
                ][..]
            )
        ));
        assert!(equal_okm(
            &secrets.server,
            &OkmBlock::new(
                &[
                    0xeb, 0x7f, 0x5e, 0x2a, 0x12, 0x3f, 0x40, 0x7d, 0xb4, 0x99, 0xe3, 0x61, 0xca,
                    0xe5, 0x90, 0xd4, 0xd9, 0x92, 0xe1, 0x4b, 0x7a, 0xce, 0x3, 0xc2, 0x44, 0xe0,
                    0x42, 0x21, 0x15, 0xb6, 0xd3, 0x8a
                ][..]
            )
        ));
    }

    struct FakeAlgorithm;

    impl Algorithm for FakeAlgorithm {
        fn packet_key(&self, _key: AeadKey, _iv: Iv) -> Box<dyn PacketKey> {
            unimplemented!()
        }

        fn header_protection_key(&self, _key: AeadKey) -> Box<dyn HeaderProtectionKey> {
            unimplemented!()
        }

        fn aead_key_len(&self) -> usize {
            16
        }
    }

    #[test]
    fn auto_traits() {
        fn assert_auto<T: Send + Sync>() {}
        assert_auto::<Box<dyn PacketKey>>();
        assert_auto::<Box<dyn HeaderProtectionKey>>();
    }
}