s2n-quic-dc 0.85.0

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

use crate::{
    credentials::{Credentials, Id},
    event,
    packet::{secret_control as control, Packet},
    path::secret::{
        open,
        schedule::{Ciphersuite, ExportSecret},
        seal, stateless_reset,
    },
    psk::io::HandshakeReason,
    stream::TransportFeatures,
};
use core::fmt;
use s2n_quic_core::{dc, time, varint::VarInt};
use std::{net::SocketAddr, sync::Arc};
use tokio::task::JoinHandle;

mod cleaner;
mod disk;
mod entry;
pub mod handshake;
mod peer;
mod rehandshake;
mod size_of;
mod state;
mod status;
mod store;

#[cfg(any(test, feature = "testing"))]
pub mod testing;

#[cfg(test)]
mod event_tests;

pub use disk::{deserialize, DiskEntry, Entries, Serializer, SerializerBuilder};
pub use entry::Entry;
use state::StateBuilderError;
use store::Store;

pub(crate) use cleaner::Epoch;
pub use entry::{
    ApplicationData, ApplicationDataError, ApplicationPair, Bidirectional, ControlPair,
};
pub use handshake::HandshakingPath;
pub use peer::Peer;

pub(crate) use size_of::SizeOf;
pub(crate) use status::Dedup;

// FIXME: Most of this comment is not true today, we're expecting to implement the details
// contained here. This is presented as a roadmap.
/// This map caches path secrets derived from handshakes.
///
/// The cache is configurable on two axes:
///
/// * Maximum size (in megabytes)
/// * Maximum per-peer/secret derivation per-second rate (in derived secrets, e.g., accepted/opened streams)
///
/// Each entry in the cache will take around 550 bytes plus 15 bits per derived secret at the
/// maximum rate (corresponding to no false positives in replay prevention for 15 seconds).
#[derive(Clone)]
pub struct Map {
    store: Arc<dyn Store>,
}

impl fmt::Debug for Map {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        f.debug_struct("Map")
            .field("secrets_len", &self.secrets_len())
            .field("peers_len", &self.peers_len())
            .field("secrets_capacity", &self.secrets_capacity())
            .finish_non_exhaustive()
    }
}

/// A builder for [`Map`].
pub struct Builder<C, S>
where
    C: 'static + time::Clock + Sync + Send,
    S: event::Subscriber,
{
    inner: state::StateBuilder<C, S>,
}

impl<C, S> Builder<C, S>
where
    C: 'static + time::Clock + Sync + Send,
    S: event::Subscriber,
{
    pub fn with_signer(mut self, signer: stateless_reset::Signer) -> Self {
        self.inner = self.inner.with_signer(signer);
        self
    }

    pub fn with_capacity(mut self, capacity: usize) -> Self {
        self.inner = self.inner.with_capacity(capacity);
        self
    }

    pub fn with_evict_on_unknown_path_secret(mut self, should_evict: bool) -> Self {
        self.inner = self.inner.with_evict_on_unknown_path_secret(should_evict);
        self
    }

    pub fn with_clock<C2: 'static + time::Clock + Sync + Send>(self, clock: C2) -> Builder<C2, S> {
        Builder {
            inner: self.inner.with_clock(clock),
        }
    }

    pub fn with_subscriber<S2: event::Subscriber>(self, subscriber: S2) -> Builder<C, S2> {
        Builder {
            inner: self.inner.with_subscriber(subscriber),
        }
    }

    /// Configures on-disk serialization of the map.
    pub fn with_serializer(mut self, serializer: Serializer) -> Self {
        self.inner = self.inner.with_serializer(serializer);
        self
    }

    /// Builds the [`Map`].
    pub fn build(self) -> Result<Map, MapBuilderError> {
        Ok(Map {
            store: self.inner.build().map_err(MapBuilderError)?,
        })
    }
}

#[derive(Debug, thiserror::Error)]
#[error(transparent)]
pub struct MapBuilderError(StateBuilderError);

impl Map {
    /// Begins configuring a [`Map`].
    ///
    /// The signer, capacity, clock, and subscriber are all required; [`Builder::build`] fails if
    /// any are missing.
    pub fn builder() -> Builder<time::StdClock, crate::event::tracing::Subscriber> {
        Builder {
            inner: state::State::builder(),
        }
    }

    pub fn new<C, S>(
        signer: stateless_reset::Signer,
        capacity: usize,
        should_evict_on_unknown_path_secret: bool,
        clock: C,
        subscriber: S,
    ) -> Self
    where
        C: 'static + time::Clock + Send + Sync,
        S: event::Subscriber,
    {
        #[expect(
            clippy::unwrap_used,
            reason = "build only fails if a required field is unset, we control the required field set"
        )]
        Self::builder()
            .with_clock(clock)
            .with_subscriber(subscriber)
            .with_signer(signer)
            .with_capacity(capacity)
            .with_evict_on_unknown_path_secret(should_evict_on_unknown_path_secret)
            .build()
            .unwrap()
    }

    /// The number of trusted secrets.
    pub fn secrets_len(&self) -> usize {
        self.store.secrets_len()
    }

    /// The number of trusted peers.
    ///
    /// This should be smaller than `secrets_len` (modulo momentary churn).
    pub fn peers_len(&self) -> usize {
        self.store.peers_len()
    }

    pub fn secrets_capacity(&self) -> usize {
        self.store.secrets_capacity()
    }

    pub fn drop_state(&self) {
        self.store.drop_state();
    }

    /// Serializes the current map to disk using the serializer configured at construction.
    ///
    /// This is a no-op (returning `Ok(())`) if no serializer was configured. It can be called
    /// regardless of whether background serialization is enabled, letting callers drive
    /// serialization on their own schedule.
    pub fn serialize_to_disk(&self) -> std::io::Result<()> {
        self.store.serialize_to_disk()
    }

    pub fn contains(&self, peer: &SocketAddr) -> bool {
        self.store.contains(peer)
    }

    pub fn register_request_handshake(
        &self,
        cb: Box<dyn Fn(SocketAddr, HandshakeReason) -> Option<JoinHandle<()>> + Send + Sync>,
    ) {
        self.store.register_request_handshake(cb);
    }

    /// Gets the [`Peer`] entry for the given address
    ///
    /// NOTE: This function is used to track cache hit ratios so it
    ///       should only be used for connection attempts.
    pub fn get_tracked(&self, peer: SocketAddr) -> Option<Peer> {
        let entry = self.store.get_by_addr_tracked(&peer)?;
        Some(Peer::new(&entry, self))
    }

    /// Gets the [`Peer`] entry for the given address
    ///
    /// NOTE: This function is used to track cache hit ratios so it
    ///       should only be used for connection attempts.
    pub fn get_untracked(&self, peer: SocketAddr) -> Option<Peer> {
        let entry = self.store.get_by_addr_untracked(&peer)?;
        Some(Peer::new(&entry, self))
    }

    /// Retrieve a sealer by path secret ID.
    ///
    /// Generally callers should prefer to use one of the `pair` APIs; this is primarily useful for
    /// "response" datagrams which want to be bound to the exact same shared secret.
    ///
    /// Note that unlike by-IP lookup this should typically not be done significantly after the
    /// original secret was used for decryption.
    pub fn seal_once_id(&self, id: Id) -> Option<(seal::Once, Credentials, dc::ApplicationParams)> {
        let entry = self.store.get_by_id_tracked(&id)?;
        let (sealer, credentials) = entry.uni_sealer();
        Some((sealer, credentials, entry.parameters()))
    }

    pub fn open_once(
        &self,
        credentials: &Credentials,
        queue_id: Option<VarInt>,
        control_out: &mut Vec<u8>,
    ) -> Option<open::Once> {
        let entry = self
            .store
            .pre_authentication(credentials, queue_id, control_out)?;
        let opener = entry.uni_opener(self.clone(), credentials, queue_id);
        Some(opener)
    }

    pub fn open_once_with_application_data(
        &self,
        credentials: &Credentials,
        queue_id: Option<VarInt>,
        control_out: &mut Vec<u8>,
    ) -> Option<(open::Once, Option<ApplicationData>)> {
        let entry = self
            .store
            .pre_authentication(credentials, queue_id, control_out)?;
        let application_data = entry.application_data().clone();
        let opener = entry.uni_opener(self.clone(), credentials, queue_id);
        Some((opener, application_data))
    }

    pub fn pair_for_credentials(
        &self,
        credentials: &Credentials,
        queue_id: Option<VarInt>,
        features: &TransportFeatures,
        control_out: &mut Vec<u8>,
    ) -> Option<(
        entry::Bidirectional,
        dc::ApplicationParams,
        Option<entry::ApplicationData>,
    )> {
        let entry = self
            .store
            .pre_authentication(credentials, queue_id, control_out)?;

        let params = entry.parameters();
        let keys = entry.bidi_remote(self.clone(), credentials, queue_id, features);

        let application_data = entry.application_data().clone();
        Some((keys, params, application_data))
    }

    pub fn secret_for_credentials(
        &self,
        credentials: &Credentials,
        queue_id: Option<VarInt>,
        features: &TransportFeatures,
        control_out: &mut Vec<u8>,
    ) -> Option<(
        ExportSecret,
        Ciphersuite,
        entry::Bidirectional,
        dc::ApplicationParams,
    )> {
        let entry = self
            .store
            .pre_authentication(credentials, queue_id, control_out)?;
        let params = entry.parameters();
        let keys = entry.bidi_remote(self.clone(), credentials, queue_id, features); // for dedup check
        let secret = entry.secret();

        Some((*secret.export_secret(), *secret.ciphersuite(), keys, params))
    }

    /// This can be called from anywhere to ask the map to handle a packet.
    ///
    /// For secret control packets, this will process those.
    /// For other packets, the map may collect metrics but will otherwise drop the packets.
    pub fn handle_unexpected_packet(&self, packet: &Packet, peer: &SocketAddr) {
        self.store.handle_unexpected_packet(packet, peer);
    }

    /// Emits a DcConnectionTimeout event via the subscriber
    pub fn on_dc_connection_timeout(&self, peer_address: &SocketAddr) {
        self.store.on_dc_connection_timeout(peer_address);
    }

    /// Emits a datagram encrypt event with the wire packet length
    pub(crate) fn on_datagram_encrypt(&self, packet_len: usize) {
        self.store.on_datagram_encrypt(packet_len);
    }

    /// Emits a datagram decrypt event with the wire packet length
    pub(crate) fn on_datagram_decrypt(&self, packet_len: usize) {
        self.store.on_datagram_decrypt(packet_len);
    }

    pub fn handle_control_packet(&self, packet: &control::Packet, peer: &SocketAddr) {
        match packet {
            control::Packet::StaleKey(packet) => {
                let _ = self.handle_stale_key_packet(packet, peer);
            }
            control::Packet::ReplayDetected(packet) => {
                let _ = self.handle_replay_detected_packet(packet, peer);
            }
            control::Packet::UnknownPathSecret(packet) => {
                let _ = self.handle_unknown_path_secret_packet(packet, peer);
            }
        }
    }

    /// Sends an already-encoded secret control packet in `buffer` to `dst` using the map's
    /// control socket, emitting the corresponding packet-sent metric (e.g.
    /// `UnknownPathSecretPacketSent`).
    ///
    /// `buffer` should contain a fully-encoded secret control packet, such as the one written
    /// into the `control_out` buffer by [`Map::open_once`],
    /// [`Map::open_once_with_application_data`], and related methods when the path secret is
    /// unknown. Sending is best-effort: if the map has no control socket the packet is dropped.
    pub fn send_control_packet(&self, dst: &SocketAddr, buffer: &mut [u8]) {
        self.store.send_control_packet(dst, buffer);
    }

    pub fn handle_stale_key_packet<'a>(
        &self,
        packet: &'a control::stale_key::Packet,
        peer: &SocketAddr,
    ) -> Option<&'a control::StaleKey> {
        self.store.handle_stale_key_packet(packet, peer)
    }

    pub fn handle_replay_detected_packet<'a>(
        &self,
        packet: &'a control::replay_detected::Packet,
        peer: &SocketAddr,
    ) -> Option<&'a control::ReplayDetected> {
        self.store.handle_replay_detected_packet(packet, peer)
    }

    pub fn handle_unknown_path_secret_packet<'a>(
        &self,
        packet: &'a control::unknown_path_secret::Packet,
        peer: &SocketAddr,
    ) -> Option<&'a control::UnknownPathSecret> {
        self.store.handle_unknown_path_secret_packet(packet, peer)
    }

    #[doc(hidden)]
    #[cfg(any(test, feature = "testing"))]
    #[allow(
        clippy::unwrap_used,
        reason = "test-support helper may panic to surface setup failures"
    )]
    pub fn for_test_with_peers(
        peers: Vec<(
            crate::path::secret::schedule::Ciphersuite,
            dc::Version,
            SocketAddr,
        )>,
    ) -> (Self, Vec<Id>) {
        use crate::path::secret::{receiver, schedule, sender};

        let provider = Self::new(
            stateless_reset::Signer::random(),
            peers.len() * 3,
            false,
            time::NoopClock,
            event::testing::Subscriber::no_snapshot(),
        );
        let mut secret = [0; 32];
        aws_lc_rs::rand::fill(&mut secret).unwrap();
        let mut stateless_reset = [0; control::TAG_LEN];
        aws_lc_rs::rand::fill(&mut stateless_reset).unwrap();

        let mut ids = Vec::with_capacity(peers.len());
        for (idx, (ciphersuite, version, peer)) in peers.into_iter().enumerate() {
            secret[..8].copy_from_slice(&(idx as u64).to_be_bytes()[..]);
            stateless_reset[..8].copy_from_slice(&(idx as u64).to_be_bytes()[..]);
            let secret = schedule::Secret::new(
                ciphersuite,
                version,
                s2n_quic_core::endpoint::Type::Client,
                &secret,
            );
            ids.push(*secret.id());
            let sender = sender::State::new(stateless_reset);
            let entry = Entry::new(
                peer,
                secret,
                sender,
                receiver::State::new(),
                dc::testing::TEST_APPLICATION_PARAMS,
                dc::testing::TEST_REHANDSHAKE_PERIOD,
                None,
            );
            let entry = Arc::new(entry);
            provider.store.test_insert(entry);
        }

        (provider, ids)
    }

    #[doc(hidden)]
    #[cfg(test)]
    pub fn test_stop_cleaner(&self) {
        self.store.test_stop_cleaner();
    }

    #[doc(hidden)]
    #[cfg(test)]
    pub fn reset_all_senders(&self) {
        self.store.reset_all_senders();
    }

    #[doc(hidden)]
    #[cfg(any(test, feature = "testing"))]
    pub fn test_insert(&self, peer: SocketAddr) {
        let receiver = super::receiver::State::new();
        let entry = Entry::fake(peer, Some(receiver));
        self.store.test_insert(entry);
    }

    #[cfg(any(test, feature = "testing"))]
    #[allow(
        clippy::unwrap_used,
        reason = "test-support helper may panic to surface setup failures"
    )]
    pub(crate) fn test_insert_pair(
        &self,
        local_addr: SocketAddr,
        local_params: Option<dc::ApplicationParams>,
        peer: &Self,
        peer_addr: SocketAddr,
        peer_params: Option<dc::ApplicationParams>,
    ) -> crate::credentials::Id {
        use crate::path::secret::{schedule, sender};
        use s2n_quic_core::endpoint::Type;

        let ciphersuite = schedule::Ciphersuite::AES_GCM_128_SHA256;

        let mut secret = [0; 32];
        aws_lc_rs::rand::fill(&mut secret).unwrap();

        let insert = |map: &Self,
                      peer: &Self,
                      peer_addr,
                      params: Option<dc::ApplicationParams>,
                      endpoint| {
            let secret =
                schedule::Secret::new(ciphersuite, dc::SUPPORTED_VERSIONS[0], endpoint, &secret);
            let id = *secret.id();

            let srt = peer.store.signer().sign(&id);

            let sender = sender::State::new(srt);

            let params = params.unwrap_or(dc::testing::TEST_APPLICATION_PARAMS);

            let entry = Entry::new(
                peer_addr,
                secret,
                sender,
                super::receiver::State::new(),
                params,
                dc::testing::TEST_REHANDSHAKE_PERIOD,
                None,
            );
            let entry = Arc::new(entry);
            map.store.test_insert(entry);

            id
        };

        let client_id = insert(self, peer, peer_addr, peer_params, Type::Client);
        let server_id = insert(peer, self, local_addr, local_params, Type::Server);

        assert_eq!(client_id, server_id);

        client_id
    }

    #[allow(clippy::type_complexity)]
    pub fn register_make_application_data(
        &self,
        cb: Box<
            dyn Fn(
                    &dyn s2n_quic_core::crypto::tls::TlsSession,
                ) -> Result<Option<ApplicationData>, ApplicationDataError>
                + Send
                + Sync,
        >,
    ) {
        self.store.register_make_application_data(cb);
    }
}