s2n-quic 1.86.0

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

//! Default provider for Address Validation tokens
//!
//! The default provider will randomly generate a 256 bit key. This key will be used to sign and
//! verify tokens. The key can be rotated at a duration set by the user.
//!
//! The default provider does not support tokens delivered in a NEW_TOKEN frame.

use core::{mem::size_of, time::Duration};
use s2n_codec::{DecoderBuffer, DecoderBufferMut};
use s2n_quic_core::{
    connection, event::api::SocketAddress, random, time::Timestamp, token::Source,
};
use s2n_quic_crypto::{constant_time, digest, hmac};
use std::collections::HashSet;
use zerocopy::{FromBytes, IntoBytes, Unaligned};
use zeroize::Zeroizing;

//= https://www.rfc-editor.org/rfc/rfc9000#section-8.1.4
//# To protect against such attacks, servers MUST ensure that
//# replay of tokens is prevented or limited.
/// Tracks retry tokens that have already been validated so that replays can be
/// detected.
///
/// A retry token is only accepted while the [`BaseKey`] that signed it remains
/// valid (at most two key rotation periods, ~2 seconds by default); once the
/// key rotates the token fails HMAC validation regardless. The filter therefore
/// only needs to remember tokens seen within that short window, and it is
/// cleared wholesale when the key rotates.
///
/// Because the window is short and bounded, an exact set of token HMACs is used
/// rather than a probabilistic filter. This avoids false positives entirely: a
/// probabilistic filter can spuriously flag a legitimate first-use token as a
/// replay, causing that client's handshake to fail. The tradeoff is memory per
/// entry, which is bounded by [`Self::MAX_ENTRIES`].
#[derive(Debug, Default)]
struct DuplicateFilter {
    seen: HashSet<[u8; 32]>,
}

impl DuplicateFilter {
    /// The maximum number of tokens to track per key.
    ///
    /// Each entry is the 32-byte token HMAC. At this cap the underlying table
    /// grows to 32Ki buckets, bounding worst-case memory to ~1 MiB per key
    /// (~2 MiB across both keys). The set is allocated lazily and grows with the
    /// live token count, so typical usage is far below this bound; the cap only
    /// limits the worst case under high retry volume.
    ///
    /// When the set is full, a token that is not already recorded is rejected
    /// rather than accepted without tracking. Accepting an untracked token would
    /// mean a later replay of it could not be detected, so this fails closed to
    /// preserve the anti-replay guarantee.
    ///
    /// Reaching this cap requires sustaining on the order of 8k validated retry
    /// tokens per second across a key rotation period, far above expected rates
    /// and further bounded by any endpoint rate limiter used to trigger Retry,
    /// so this rejection path is not expected to be hit in practice.
    const MAX_ENTRIES: usize = 16 * 1024;

    /// Records the token as seen, returning `true` if it was newly recorded.
    ///
    /// Returns `false` if the token was already recorded, indicating a replay.
    /// Also returns `false` if [`Self::MAX_ENTRIES`] tokens are already being
    /// tracked and the token could not be recorded, since a later replay of it
    /// could not be detected. In either case the caller must reject the token.
    fn insert(&mut self, token: &Token) -> bool {
        if self.seen.len() < Self::MAX_ENTRIES {
            self.seen.insert(token.hmac)
        } else {
            false
        }
    }
}

struct BaseKey {
    active_duration: Duration,

    // HMAC key for signing and verifying
    key: Option<(Timestamp, hmac::Key)>,

    // Tracks previously validated tokens to detect replays. An empty filter
    // performs no heap allocation; the backing storage is allocated on the
    // first validated token and released when the key rotates.
    duplicate_filter: DuplicateFilter,
}

impl BaseKey {
    pub fn new(active_duration: Duration) -> Self {
        Self {
            active_duration,
            key: None,
            duplicate_filter: DuplicateFilter::default(),
        }
    }

    pub fn hasher(&mut self, random: &mut dyn random::Generator) -> Option<hmac::Context> {
        let key = self.poll_key(random)?;
        Some(hmac::Context::with_key(&key))
    }

    fn poll_key(&mut self, random: &mut dyn random::Generator) -> Option<hmac::Key> {
        let now = s2n_quic_platform::time::now();

        //= https://www.rfc-editor.org/rfc/rfc9000#section-21.3
        //# Servers SHOULD provide mitigations for this attack by limiting the
        //# usage and lifetime of address validation tokens; see Section 8.1.3.
        if let Some((expires_at, key)) = self.key.as_ref() {
            if expires_at > &now {
                // key is still valid
                return Some(key.clone());
            }
        }

        let expires_at = now.checked_add(self.active_duration)?;

        let mut key_material = Zeroizing::new([0; digest::SHA256_OUTPUT_LEN]);
        random.private_random_fill(&mut key_material[..]);
        let key = hmac::Key::new(hmac::HMAC_SHA256, key_material.as_ref());

        // Replace the duplicate filter along with the key material. Tokens
        // signed by the previous key will no longer validate, so there is no
        // need to track them for replay detection.
        //
        // Replacing with a fresh (allocation-free) filter releases the old
        // allocation rather than clearing it in place. Retry token issuance is
        // expected to be bursty and transient, and clearing would pin the
        // burst's high-water-mark allocation for the lifetime of the endpoint.
        // Replacement returns memory to zero between bursts; the set re-grows
        // lazily on the next validated token, and that re-growth cost is only
        // paid if another burst actually arrives.
        self.duplicate_filter = DuplicateFilter::default();

        self.key = Some((expires_at, key));

        self.key.as_ref().map(|key| key.1.clone())
    }
}

const DEFAULT_KEY_ROTATION_PERIOD: Duration = Duration::from_millis(1000);

#[derive(Debug)]
pub struct Provider {
    //= https://www.rfc-editor.org/rfc/rfc9000#section-8.1.3
    //# Thus, a token SHOULD have an
    //# expiration time, which could be either an explicit expiration time or
    //# an issued timestamp that can be used to dynamically calculate the
    //# expiration time.
    /// To fulfill this SHOULD, we rotate the key periodically. This allows
    /// customers to control the token lifetime without adding bytes to the token itself.
    key_rotation_period: Duration,
}

impl Default for Provider {
    fn default() -> Self {
        Self {
            key_rotation_period: DEFAULT_KEY_ROTATION_PERIOD,
        }
    }
}

impl super::Provider for Provider {
    type Format = Format;
    type Error = core::convert::Infallible;

    fn start(self) -> Result<Self::Format, Self::Error> {
        // The keys must remain valid for two rotation periods or they will regenerate their
        // material and validation will fail.
        let format = Format {
            key_rotation_period: self.key_rotation_period,
            current_key_rotates_at: s2n_quic_platform::time::now(),
            current_key: 0,
            keys: [
                BaseKey::new(self.key_rotation_period * 2),
                BaseKey::new(self.key_rotation_period * 2),
            ],
        };

        Ok(format)
    }
}

pub struct Format {
    //= https://www.rfc-editor.org/rfc/rfc9000#section-8.1.4
    //= type=exception
    //= reason=We use a duplicate filter to prevent tokens from being used more than once.
    //# Servers are encouraged to allow tokens to be used only
    //# once, if possible; tokens MAY include additional information about
    //# clients to further narrow applicability or reuse.
    /// Key validity period
    key_rotation_period: Duration,

    //= https://www.rfc-editor.org/rfc/rfc9000#section-8.1.4
    //# Servers SHOULD ensure that
    //# tokens sent in Retry packets are only accepted for a short time.
    /// Timestamp to rotate current key
    current_key_rotates_at: s2n_quic_core::time::Timestamp,

    /// Which key is used to sign
    current_key: u8,

    /// Key used to sign keys
    keys: [BaseKey; 2],
}

impl Format {
    fn current_key(&mut self) -> u8 {
        let now = s2n_quic_platform::time::now();
        if now > self.current_key_rotates_at {
            self.current_key ^= 1;
            self.current_key_rotates_at = now + self.key_rotation_period;

            // The duplicate filter is not touched here: this only flips which
            // key signs new tokens, while both keys remain valid for validation
            // during the overlap. Each key's filter is dropped when that key's
            // material regenerates, in `BaseKey::poll_key`.
        }
        self.current_key
    }

    // Retry Tokens need to include the original destination connection id from the transport
    // parameters. This OCID is included in the tag.
    fn tag_retry_token(
        &mut self,
        token: &Token,
        context: &mut super::Context<'_>,
    ) -> Option<hmac::Tag> {
        let mut ctx = self.keys[token.header.key_id() as usize].hasher(context.random)?;

        //= https://www.rfc-editor.org/rfc/rfc9000#section-8.1.4
        //# Tokens
        //# sent in Retry packets SHOULD include information that allows the
        //# server to verify that the source IP address and port in client
        //# packets remain constant.
        ctx.update(&token.original_destination_connection_id);
        ctx.update(&token.nonce);
        ctx.update(context.peer_connection_id);
        match context.remote_address {
            SocketAddress::IpV4 { ip, port, .. } => {
                ctx.update(ip);
                ctx.update(&port.to_be_bytes());
            }
            SocketAddress::IpV6 { ip, port, .. } => {
                ctx.update(ip);
                ctx.update(&port.to_be_bytes());
            }
            _ => {
                // we are unable to hash the address so bail
                return None;
            }
        };

        Some(ctx.sign())
    }

    // Using the key id in the token, verify the token
    fn validate_retry_token(
        &mut self,
        context: &mut super::Context<'_>,
        token: &Token,
    ) -> Option<connection::InitialId> {
        let tag = self.tag_retry_token(token, context)?;

        // Verify the token is authentic before consulting the duplicate filter.
        //
        // The filter is backed by a HashSet whose lookups are not constant-time.
        // Gating it behind the constant-time HMAC comparison ensures those
        // variable-time operations only ever run on tokens we actually minted,
        // so a client submitting forged tokens cannot use the filter's timing as
        // an oracle (and cannot cause any filter work on unauthenticated input).
        if constant_time::verify_slices_are_equal(&token.hmac, tag.as_ref()).is_err() {
            return None;
        }

        // Record the token now that it has been validated. Only recording
        // validated tokens prevents the filter from being filled with garbage
        // tokens.
        //
        // `insert` returns `false` if the token was already recorded (a replay)
        // or if the filter is full and cannot record it. Reject in both cases:
        // accepting an untracked token would allow a later replay of it to go
        // undetected.
        if !self.keys[token.header.key_id() as usize]
            .duplicate_filter
            .insert(token)
        {
            return None;
        }

        token.original_destination_connection_id()
    }
}

impl super::Format for Format {
    const TOKEN_LEN: usize = size_of::<Token>();

    /// The default provider does not support NEW_TOKEN frame tokens
    fn generate_new_token(
        &mut self,
        _context: &mut super::Context<'_>,
        _source_connection_id: &connection::LocalId,
        _output_buffer: &mut [u8],
    ) -> Option<()> {
        //= https://www.rfc-editor.org/rfc/rfc9000#section-8.1.3
        //= type=TODO
        //= tracking-issue=418
        //# A server MAY provide clients with an address validation token during
        //# one connection that can be used on a subsequent connection.

        //= https://www.rfc-editor.org/rfc/rfc9000#section-8.1.4
        //= type=TODO
        //= tracking-issue=346
        //# Tokens sent in NEW_TOKEN frames MUST include information that allows
        //# the server to verify that the client IP address has not changed from
        //# when the token was issued.

        //= https://www.rfc-editor.org/rfc/rfc9000#section-8.1.3
        //= type=TODO
        //= tracking-issue=345
        //# A token issued with NEW_TOKEN MUST NOT include information that would
        //# allow values to be linked by an observer to the connection on which
        //# it was issued.

        //= https://www.rfc-editor.org/rfc/rfc9000#section-8.1.3
        //= type=TODO
        //= tracking-issue=387
        //# A server MUST ensure that every NEW_TOKEN frame it sends
        //# is unique across all clients, with the exception of those sent to
        //# repair losses of previously sent NEW_TOKEN frames.

        //= https://www.rfc-editor.org/rfc/rfc9000#section-8.1.3
        //= type=TODO
        //= tracking-issue=394
        //# A server MAY provide clients with an address validation token during
        //# one connection that can be used on a subsequent connection.

        None
    }

    //= https://www.rfc-editor.org/rfc/rfc9000#section-8.1.2
    //# Requiring the server
    //# to provide a different connection ID, along with the
    //# original_destination_connection_id transport parameter defined in
    //# Section 18.2, forces the server to demonstrate that it, or an entity
    //# it cooperates with, received the original Initial packet from the
    //# client.
    fn generate_retry_token(
        &mut self,
        context: &mut super::Context<'_>,
        original_destination_connection_id: &connection::InitialId,
        output_buffer: &mut [u8],
    ) -> Option<()> {
        let buffer = DecoderBufferMut::new(output_buffer);
        let (token, _) = buffer
            .decode::<&mut Token>()
            .expect("Provided output buffer did not match TOKEN_LEN");

        let header = Header::new(Source::RetryPacket, self.current_key());

        token.header = header;
        token.original_destination_connection_id[..original_destination_connection_id.len()]
            .copy_from_slice(original_destination_connection_id.as_bytes());
        token.odcid_len = original_destination_connection_id.len() as u8;

        // ensure the other CID bytes are zeroed out
        for b in token
            .original_destination_connection_id
            .iter_mut()
            .skip(original_destination_connection_id.len())
        {
            *b = 0;
        }

        // Populate the nonce before signing
        context.random.public_random_fill(&mut token.nonce[..]);

        let tag = self.tag_retry_token(token, context)?;

        token.hmac.copy_from_slice(tag.as_ref());

        Some(())
    }

    //= https://www.rfc-editor.org/rfc/rfc9000#section-8.1.3
    //# When a server receives an Initial packet with an address validation
    //# token, it MUST attempt to validate the token, unless it has already
    //# completed address validation.
    fn validate_token(
        &mut self,
        context: &mut super::Context<'_>,
        token: &[u8],
    ) -> Option<connection::InitialId> {
        let buffer = DecoderBuffer::new(token);
        let (token, remaining) = buffer.decode::<&Token>().ok()?;

        // Verify the provided token doesn't have any additional data
        remaining.ensure_empty().ok()?;

        if token.header.version() != TOKEN_VERSION {
            return None;
        }

        let source = token.header.token_source();

        match source {
            Source::RetryPacket => self.validate_retry_token(context, token),
            Source::NewTokenFrame => None, // Not supported in the default provider
        }
        //= https://www.rfc-editor.org/rfc/rfc9000#section-8.1.4
        //= type=TODO
        //= tracking-issue=347
        //# Tokens that are provided
        //# in NEW_TOKEN frames (Section 19.7) need to be valid for longer but
        //# SHOULD NOT be accepted multiple times.

        //= https://www.rfc-editor.org/rfc/rfc9000#section-8.1.3
        //= type=TODO
        //= tracking-issue=388
        //# Clients that want to break continuity of identity with a server can
        //# discard tokens provided using the NEW_TOKEN frame.
    }
}

#[derive(Clone, Copy, Debug, FromBytes, IntoBytes, Unaligned)]
#[repr(C)]
pub(crate) struct Header(u8);

const TOKEN_VERSION: u8 = 0x00;

const VERSION_SHIFT: u8 = 7;
const VERSION_MASK: u8 = 0x80;

const TOKEN_SOURCE_SHIFT: u8 = 6;
const TOKEN_SOURCE_MASK: u8 = 0x40;

const KEY_ID_SHIFT: u8 = 5;
const KEY_ID_MASK: u8 = 0x20;

impl Header {
    fn new(source: Source, key_id: u8) -> Header {
        let mut header: u8 = 0;
        header |= TOKEN_VERSION << VERSION_SHIFT;
        //= https://www.rfc-editor.org/rfc/rfc9000#section-8.1.3
        //# Information that
        //# allows the server to distinguish between tokens from Retry and
        //# NEW_TOKEN MAY be accessible to entities other than the server.
        header |= match source {
            Source::NewTokenFrame => 0 << TOKEN_SOURCE_SHIFT,
            Source::RetryPacket => 1 << TOKEN_SOURCE_SHIFT,
        };

        // The key_id can only be 0 or 1
        debug_assert!(key_id <= 1);
        header |= (key_id & 0x01) << KEY_ID_SHIFT;

        Header(header)
    }

    fn version(self) -> u8 {
        (self.0 & VERSION_MASK) >> VERSION_SHIFT
    }

    fn key_id(self) -> u8 {
        (self.0 & KEY_ID_MASK) >> KEY_ID_SHIFT
    }

    //= https://www.rfc-editor.org/rfc/rfc9000#section-8.1.1
    //# A token sent in a NEW_TOKEN frame or a Retry packet MUST be
    //# constructed in a way that allows the server to identify how it was
    //# provided to a client.  These tokens are carried in the same field but
    //# require different handling from servers.
    fn token_source(self) -> Source {
        match (self.0 & TOKEN_SOURCE_MASK) >> TOKEN_SOURCE_SHIFT {
            0 => Source::NewTokenFrame,
            1 => Source::RetryPacket,
            _ => Source::NewTokenFrame,
        }
    }
}

//= https://www.rfc-editor.org/rfc/rfc9000#section-8.1.4
//#   There is no need for a single well-defined format for the token
//#   because the server that generates the token also consumes it.
#[derive(Copy, Clone, Debug, FromBytes, IntoBytes, Unaligned)]
#[repr(C)]
struct Token {
    header: Header,

    odcid_len: u8,
    original_destination_connection_id: [u8; 20],

    //= https://www.rfc-editor.org/rfc/rfc9000#section-8.1.4
    //# An address validation token MUST be difficult to guess.  Including a
    //# random value with at least 128 bits of entropy in the token would be
    //# sufficient, but this depends on the server remembering the value it
    //# sends to clients.
    nonce: [u8; 32],

    //= https://www.rfc-editor.org/rfc/rfc9000#section-8.1.4
    //# A token-based scheme allows the server to offload any state
    //# associated with validation to the client.  For this design to work,
    //# the token MUST be covered by integrity protection against
    //# modification or falsification by clients.  Without integrity
    //# protection, malicious clients could generate or guess values for
    //# tokens that would be accepted by the server.  Only the server
    //# requires access to the integrity protection key for tokens.
    hmac: [u8; 32],
}

s2n_codec::zerocopy_value_codec!(Token);

impl Token {
    pub fn original_destination_connection_id(&self) -> Option<connection::InitialId> {
        let dcid = self
            .original_destination_connection_id
            .get(..self.odcid_len as usize)?;
        connection::InitialId::try_from_bytes(dcid)
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use s2n_quic_core::{
        inet::SocketAddress,
        random,
        token::{Context, Format as FormatTrait, Source},
    };
    use s2n_quic_platform::time;
    use std::{net::SocketAddr, sync::Arc};
    use zerocopy::FromZeros;

    const TEST_KEY_ROTATION_PERIOD: Duration = Duration::from_millis(1000);

    fn get_test_format() -> Format {
        Format {
            key_rotation_period: TEST_KEY_ROTATION_PERIOD,
            keys: [
                BaseKey::new(TEST_KEY_ROTATION_PERIOD * 2),
                BaseKey::new(TEST_KEY_ROTATION_PERIOD * 2),
            ],
            current_key_rotates_at: time::now(),
            current_key: 0,
        }
    }

    #[test]
    fn test_header() {
        // Test all combinations of values to create a header and verify the header returns the
        // expected values.
        for source in &[Source::NewTokenFrame, Source::RetryPacket] {
            for key_id in [0, 1] {
                let header = Header::new(*source, key_id);
                // The version should always be the constant TOKEN_VERSION
                assert_eq!(header.version(), TOKEN_VERSION);
                //= https://www.rfc-editor.org/rfc/rfc9000#section-8.1.1
                //= type=test
                //# A token sent in a NEW_TOKEN frames or a Retry packet MUST be
                //# constructed in a way that allows the server to identify how it was
                //# provided to a client.

                assert_eq!(header.token_source(), *source);
                assert_eq!(header.key_id(), key_id);
            }
        }
    }

    #[test]
    fn test_valid_retry_tokens() {
        let clock = Arc::new(time::testing::MockClock::new());
        time::testing::set_local_clock(clock.clone());

        let mut format = get_test_format();
        let first_conn_id = connection::PeerId::try_from_bytes(&[2, 4, 6, 8, 10]).unwrap();
        let second_conn_id = connection::PeerId::try_from_bytes(&[1, 3, 5, 7, 9]).unwrap();
        let orig_conn_id =
            connection::InitialId::try_from_bytes(&[0, 1, 2, 3, 4, 5, 6, 7]).unwrap();
        let addr = SocketAddress::default();
        let mut first_token = [0; Format::TOKEN_LEN];
        let mut second_token = [0; Format::TOKEN_LEN];
        let mut random = random::testing::Generator(5);
        let mut context = Context::new(&addr, &first_conn_id, &mut random);

        // Generate two tokens for different connections
        format
            .generate_retry_token(&mut context, &orig_conn_id, &mut first_token)
            .unwrap();

        context = Context::new(&addr, &second_conn_id, &mut random);
        format
            .generate_retry_token(&mut context, &orig_conn_id, &mut second_token)
            .unwrap();

        clock.adjust_by(TEST_KEY_ROTATION_PERIOD);
        context = Context::new(&addr, &first_conn_id, &mut random);
        assert_eq!(
            format.validate_token(&mut context, &first_token),
            Some(orig_conn_id)
        );
        context = Context::new(&addr, &second_conn_id, &mut random);
        assert_eq!(
            format.validate_token(&mut context, &second_token),
            Some(orig_conn_id)
        );
        context = Context::new(&addr, &first_conn_id, &mut random);
        assert_eq!(format.validate_token(&mut context, &second_token), None);
    }

    #[test]
    fn test_retry_ip_port_validation() {
        //= https://www.rfc-editor.org/rfc/rfc9000#section-8.1.4
        //= type=test
        //# Tokens
        //# sent in Retry packets SHOULD include information that allows the
        //# server to verify that the source IP address and port in client
        //# packets remain constant.
        let mut format = get_test_format();
        let conn_id = connection::PeerId::try_from_bytes(&[2, 4, 6, 8, 10]).unwrap();
        let orig_conn_id =
            connection::InitialId::try_from_bytes(&[0, 1, 2, 3, 4, 5, 6, 7]).unwrap();

        let mut token = [0; Format::TOKEN_LEN];
        let ip_address = "127.0.0.1:443";
        let addr: SocketAddr = ip_address.parse().unwrap();
        let correct_address: SocketAddress = addr.into();
        let mut random = random::testing::Generator(5);
        let mut context = Context::new(&correct_address, &conn_id, &mut random);
        format
            .generate_retry_token(&mut context, &orig_conn_id, &mut token)
            .unwrap();

        let ip_address = "127.0.0.2:443";
        let addr: SocketAddr = ip_address.parse().unwrap();
        let incorrect_address: SocketAddress = addr.into();
        context = Context::new(&incorrect_address, &conn_id, &mut random);
        assert_eq!(format.validate_token(&mut context, &token), None);

        let ip_address = "127.0.0.1:444";
        let addr: SocketAddr = ip_address.parse().unwrap();
        let incorrect_port: SocketAddress = addr.into();
        context = Context::new(&incorrect_port, &conn_id, &mut random);
        assert_eq!(format.validate_token(&mut context, &token), None);

        // Verify the token is still valid after the failed attempts
        context = Context::new(&correct_address, &conn_id, &mut random);
        assert!(format.validate_token(&mut context, &token).is_some());
    }

    #[test]
    fn test_key_rotation() {
        //= https://www.rfc-editor.org/rfc/rfc9000#section-8.1.3
        //= type=test
        //# Thus, a token SHOULD have an
        //# expiration time, which could be either an explicit expiration time or
        //# an issued timestamp that can be used to dynamically calculate the
        //# expiration time.
        let clock = Arc::new(time::testing::MockClock::new());
        time::testing::set_local_clock(clock.clone());

        let mut format = get_test_format();
        let conn_id = connection::PeerId::TEST_ID;
        let orig_conn_id = connection::InitialId::TEST_ID;
        let addr = SocketAddress::default();
        let mut buf = [0; Format::TOKEN_LEN];
        let mut random = random::testing::Generator(5);
        let mut context = Context::new(&addr, &conn_id, &mut random);
        format
            .generate_retry_token(&mut context, &orig_conn_id, &mut buf)
            .unwrap();

        // Validation should succeed because the signing key is still valid, even
        // though it has been rotated from the current signing key
        clock.adjust_by(TEST_KEY_ROTATION_PERIOD);
        assert!(format.validate_token(&mut context, &buf).is_some());

        // Validation should fail because the key used for signing has been regenerated
        clock.adjust_by(TEST_KEY_ROTATION_PERIOD);
        assert!(format.validate_token(&mut context, &buf).is_none());
    }

    #[test]
    fn test_expired_retry_token() {
        //= https://www.rfc-editor.org/rfc/rfc9000#section-8.1.4
        //= type=test
        //# Servers SHOULD ensure that
        //# tokens sent in Retry packets are only accepted for a short time.
        let clock = Arc::new(time::testing::MockClock::new());
        time::testing::set_local_clock(clock.clone());

        let mut format = get_test_format();
        let conn_id = connection::PeerId::TEST_ID;
        let orig_conn_id = connection::InitialId::TEST_ID;
        let addr = SocketAddress::default();
        let mut buf = [0; Format::TOKEN_LEN];
        let mut random = random::testing::Generator(5);
        let mut context = Context::new(&addr, &conn_id, &mut random);
        format
            .generate_retry_token(&mut context, &orig_conn_id, &mut buf)
            .unwrap();

        //= https://www.rfc-editor.org/rfc/rfc9000#section-21.3
        //= type=test
        //# Servers SHOULD provide mitigations for this attack by limiting the
        //# usage and lifetime of address validation tokens; see Section 8.1.3.
        // Validation should fail because multiple rotation periods have elapsed
        clock.adjust_by(TEST_KEY_ROTATION_PERIOD * 2);
        assert!(format.validate_token(&mut context, &buf).is_none());
    }

    #[test]
    fn test_retry_validation_default_format() {
        let clock = Arc::new(time::testing::MockClock::new());
        time::testing::set_local_clock(clock);

        let mut format = get_test_format();
        let conn_id = connection::PeerId::TEST_ID;
        let odcid = connection::InitialId::try_from_bytes(&[0, 1, 2, 3, 4, 5, 6, 7]).unwrap();
        let addr = SocketAddress::default();
        let mut buf = [0; Format::TOKEN_LEN];
        let mut random = random::testing::Generator(5);
        let mut context = Context::new(&addr, &conn_id, &mut random);
        format
            .generate_retry_token(&mut context, &odcid, &mut buf)
            .unwrap();

        assert_eq!(format.validate_token(&mut context, &buf), Some(odcid));

        let wrong_conn_id = connection::PeerId::try_from_bytes(&[0, 1, 2]).unwrap();
        context = Context::new(&addr, &wrong_conn_id, &mut random);
        assert!(format.validate_token(&mut context, &buf).is_none());
    }

    #[test]
    fn test_duplicate_token_detection() {
        //= https://www.rfc-editor.org/rfc/rfc9000#section-8.1.4
        //= type=test
        //# To protect against such attacks, servers MUST ensure that
        //# replay of tokens is prevented or limited.
        let mut format = get_test_format();
        let conn_id = connection::PeerId::TEST_ID;
        let odcid = connection::InitialId::try_from_bytes(&[0, 1, 2, 3, 4, 5, 6, 7]).unwrap();
        let addr = SocketAddress::default();
        let mut buf = [0; Format::TOKEN_LEN];
        let mut random = random::testing::Generator(5);
        let mut context = Context::new(&addr, &conn_id, &mut random);
        format
            .generate_retry_token(&mut context, &odcid, &mut buf)
            .unwrap();

        assert_eq!(format.validate_token(&mut context, &buf), Some(odcid));

        // Second attempt with the same token should fail because the token is a duplicate
        assert!(format.validate_token(&mut context, &buf).is_none());
    }

    fn test_token_with_hmac(byte: u8) -> Token {
        let mut token = Token::new_zeroed();
        token.hmac = [byte; 32];
        token
    }

    #[test]
    fn duplicate_filter_tracks_seen_tokens() {
        let mut filter = DuplicateFilter::default();
        let token = test_token_with_hmac(1);

        // A new token is recorded
        assert!(filter.insert(&token));

        // Re-inserting the same token reports it as already seen
        assert!(!filter.insert(&token));

        // A different token is not affected
        let other = test_token_with_hmac(2);
        assert!(filter.insert(&other));
    }

    #[test]
    fn duplicate_filter_bounds_capacity() {
        //= https://www.rfc-editor.org/rfc/rfc9000#section-8.1.4
        //= type=test
        //# To protect against such attacks, servers MUST ensure that
        //# replay of tokens is prevented or limited.
        // Once the filter is full it stops recording new tokens to bound memory.
        // `insert` reports this by returning `false` so the caller can reject the
        // token, failing closed to preserve replay detection.
        let mut filter = DuplicateFilter::default();

        for i in 0..DuplicateFilter::MAX_ENTRIES {
            let mut token = Token::new_zeroed();
            token.hmac[..8].copy_from_slice(&(i as u64).to_be_bytes());
            assert!(filter.insert(&token));
        }
        assert_eq!(filter.seen.len(), DuplicateFilter::MAX_ENTRIES);

        // A new token beyond the cap is not recorded and insert reports failure
        let overflow = test_token_with_hmac(0xff);
        assert!(!filter.insert(&overflow));
        assert_eq!(filter.seen.len(), DuplicateFilter::MAX_ENTRIES);
        assert!(!filter.seen.contains(&overflow.hmac));

        // An already-recorded token is still tracked at capacity, so a replay
        // of it continues to be rejected
        let mut seen = Token::new_zeroed();
        seen.hmac[..8].copy_from_slice(&0u64.to_be_bytes());
        assert!(filter.seen.contains(&seen.hmac));
        assert!(!filter.insert(&seen));
    }

    #[test]
    fn test_token_modification_detection() {
        //= https://www.rfc-editor.org/rfc/rfc9000#section-8.1.4
        //= type=test
        //# For this design to work,
        //# the token MUST be covered by integrity protection against
        //# modification or falsification by clients.
        let mut format = get_test_format();
        let conn_id = connection::PeerId::try_from_bytes(&[2, 4, 6, 8, 10]).unwrap();
        let orig_conn_id =
            connection::InitialId::try_from_bytes(&[0, 1, 2, 3, 4, 5, 6, 7]).unwrap();
        let addr = SocketAddress::default();
        let mut token = [0; Format::TOKEN_LEN];

        let mut random = random::testing::Generator(5);
        let mut context = Context::new(&addr, &conn_id, &mut random);
        // Generate two tokens for different connections
        format
            .generate_retry_token(&mut context, &orig_conn_id, &mut token)
            .unwrap();

        for i in 0..Format::TOKEN_LEN {
            random = random::testing::Generator(5);
            context = Context::new(&addr, &conn_id, &mut random);
            token[i] = !token[i];
            assert!(format.validate_token(&mut context, &token).is_none());
            token[i] = !token[i];
        }
    }

    #[test]
    fn test_token_length_check() {
        let mut format = get_test_format();
        let conn_id = connection::PeerId::try_from_bytes(&[2, 4, 6, 8, 10]).unwrap();
        let addr = SocketAddress::default();

        bolero::check!().for_each(move |token| {
            let mut random = random::testing::Generator(5);
            let mut context = Context::new(&addr, &conn_id, &mut random);
            assert!(format.validate_token(&mut context, token).is_none())
        });
    }

    #[test]
    fn test_token_falsification_detection() {
        let mut format = get_test_format();
        let conn_id = connection::PeerId::try_from_bytes(&[2, 4, 6, 8, 10]).unwrap();
        let addr = SocketAddress::default();

        //= https://www.rfc-editor.org/rfc/rfc9000#section-8.1.4
        //= type=test
        //# For this design to work,
        //# the token MUST be covered by integrity protection against
        //# modification or falsification by clients.
        let generator = bolero::generator::produce::<Vec<u8>>()
            .with()
            .len(Format::TOKEN_LEN);
        bolero::check!()
            .with_generator(generator)
            .for_each(move |token| {
                let mut random = random::testing::Generator(5);
                let mut context = Context::new(&addr, &conn_id, &mut random);
                assert!(format.validate_token(&mut context, token).is_none())
            });
    }
}