Allow server to use NEW_TOKEN frames

Whenever a path becomes validated, the server sends the client NEW_TOKEN
frames. These may cause an Incoming to be validated.

- Adds dependency on `fastbloom`
- Converts TokenInner to enum with Retry and Validation variants
- Adds relevant configuration to ServerConfig
- Incoming now has `may_retry`
- Adds `TokenLog` object to server to mitigate token reuse

Cargo.toml

@@ -25,6 +25,7 @@ crc = "3"
 directories-next = "2"
 futures-io = "0.3.19"
 getrandom = { version = "0.2", default-features = false }
+fastbloom = "0.8"
 hdrhistogram = { version = "7.2", default-features = false }
 hex-literal = "0.4"
 lazy_static = "1"

quinn-proto/Cargo.toml

@@ -11,7 +11,7 @@ categories.workspace = true
 workspace = ".."
 
 [features]
-default = ["rustls-ring", "log"]
+default = ["rustls-ring", "log", "fastbloom"]
 aws-lc-rs = ["dep:aws-lc-rs", "aws-lc-rs?/aws-lc-sys", "aws-lc-rs?/prebuilt-nasm"]
 aws-lc-rs-fips = ["aws-lc-rs", "aws-lc-rs?/fips"]
 # For backwards compatibility, `rustls` forwards to `rustls-ring`
@@ -34,6 +34,7 @@ rustls-log = ["rustls?/logging"]
 arbitrary = { workspace = true, optional = true }
 aws-lc-rs = { workspace = true, optional = true }
 bytes = { workspace = true }
+fastbloom = { workspace = true, optional = true }
 rustc-hash = { workspace = true }
 rand = { workspace = true }
 ring = { workspace = true, optional = true }

quinn-proto/src/bloom_token_log.rs

@@ -0,0 +1,216 @@
+use std::{
+    collections::HashSet,
+    f64::consts::LN_2,
+    hash::{BuildHasher, Hasher},
+    mem::{size_of, swap},
+    sync::Mutex,
+};
+
+use fastbloom::BloomFilter;
+use rustc_hash::FxBuildHasher;
+use tracing::warn;
+
+use crate::{Duration, SystemTime, TokenLog, TokenReuseError, UNIX_EPOCH};
+
+/// The token's rand needs to guarantee uniqueness because of the role it plays in the encryption
+/// of the tokens, so it is 128 bits. But since the token log can tolerate both false positives and
+/// false negatives, we trim it down to 64 bits, which would still only have a small collision rate
+/// even at significant amounts of usage, while allowing us to store twice as many in the hash set
+/// variant.
+///
+/// Token rand values are uniformly randomly generated server-side and cryptographically integrity-
+/// checked, so we don't need to employ secure hashing for this, we can simply truncate.
+fn rand_to_fingerprint(rand: u128) -> u64 {
+    (rand & 0xffffffff) as u64
+}
+
+/// `BuildHasher` of `IdentityHasher`.
+#[derive(Default)]
+struct IdentityBuildHasher;
+
+impl BuildHasher for IdentityBuildHasher {
+    type Hasher = IdentityHasher;
+
+    fn build_hasher(&self) -> Self::Hasher {
+        IdentityHasher::default()
+    }
+}
+
+/// Hasher that assumes the thing being hashes is a `u64` and is the identity operation.
+#[derive(Default)]
+struct IdentityHasher {
+    data: [u8; 8],
+    #[cfg(debug_assertions)]
+    wrote_8_byte_slice: bool,
+}
+
+impl Hasher for IdentityHasher {
+    fn write(&mut self, bytes: &[u8]) {
+        #[cfg(debug_assertions)]
+        {
+            assert!(!self.wrote_8_byte_slice);
+            assert_eq!(bytes.len(), 8);
+            self.wrote_8_byte_slice = true;
+        }
+        self.data.copy_from_slice(bytes);
+    }
+
+    fn finish(&self) -> u64 {
+        #[cfg(debug_assertions)]
+        assert!(self.wrote_8_byte_slice);
+        u64::from_ne_bytes(self.data)
+    }
+}
+
+/// Hash set of `u64` which are assumed to already be uniformly randomly distributed, and thus
+/// effectively pre-hashed.
+type IdentityHashSet = HashSet<u64, IdentityBuildHasher>;
+
+/// Bloom filter that uses `FxHasher`s.
+type FxBloomFilter = BloomFilter<512, FxBuildHasher>;
+
+/// Bloom filter-based `TokenLog`
+///
+/// Parameterizable over an approximate maximum number of bytes to allocate. Starts out by storing
+/// used tokens in a hash set. Once the hash set becomes too large, converts it to a bloom filter.
+/// This achieves a memory profile of linear growth with an upper bound.
+///
+/// Divides time into periods based on `lifetime` and stores two filters at any given moment, for
+/// each of the two periods currently non-expired tokens could expire in. As such, turns over
+/// filters as time goes on to avoid bloom filter false positive rate increasing infinitely over
+/// time.
+pub struct BloomTokenLog(Mutex<State>);
+
+/// Lockable state of [`BloomTokenLog`]
+struct State {
+    filter_max_bytes: usize,
+    k_num: u32,
+
+    // filter_1 covers tokens that expire in the period starting at
+    // UNIX_EPOCH + period_idx_1 * lifetime and extending lifetime after.
+    // filter_2 covers tokens for the next lifetime after that.
+    period_idx_1: u128,
+    filter_1: Filter,
+    filter_2: Filter,
+}
+
+/// Period filter within [`State`]
+enum Filter {
+    Set(IdentityHashSet),
+    Bloom(FxBloomFilter),
+}
+
+impl BloomTokenLog {
+    /// Construct with an approximate maximum memory usage and expected number of validation token
+    /// usages per expiration period
+    ///
+    /// Calculates the optimal bloom filter k number automatically.
+    pub fn new_expected_items(max_bytes: usize, expected_hits: u64) -> Self {
+        Self::new(max_bytes, optimal_k_num(max_bytes, expected_hits))
+    }
+
+    /// Construct with an approximate maximum memory usage and a bloom filter k number
+    ///
+    /// If choosing a custom k number, note that `BloomTokenLog` always maintains two filters
+    /// between them and divides the allocation budget of `max_bytes` evenly between them. As such,
+    /// each bloom filter will contain `max_bytes * 4` bits.
+    pub fn new(max_bytes: usize, k_num: u32) -> Self {
+        assert!(max_bytes >= 2, "BloomTokenLog max_bytes too low");
+        assert!(k_num >= 1, "BloomTokenLog k_num must be at least 1");
+
+        Self(Mutex::new(State {
+            filter_max_bytes: max_bytes / 2,
+            k_num,
+            period_idx_1: 0,
+            filter_1: Filter::Set(IdentityHashSet::default()),
+            filter_2: Filter::Set(IdentityHashSet::default()),
+        }))
+    }
+}
+
+fn optimal_k_num(num_bytes: usize, expected_hits: u64) -> u32 {
+    assert!(expected_hits > 0, "BloomTokenLog expected hits too low");
+    let num_bits = (num_bytes as u64)
+        .checked_mul(8)
+        .expect("BloomTokenLog num bytes too high");
+    (((num_bits as f64 / expected_hits as f64) * LN_2).round() as u32).max(1)
+}
+
+impl TokenLog for BloomTokenLog {
+    fn check_and_insert(
+        &self,
+        rand: u128,
+        issued: SystemTime,
+        lifetime: Duration,
+    ) -> Result<(), TokenReuseError> {
+        let mut guard = self.0.lock().unwrap();
+        let state = &mut *guard;
+        let fingerprint = rand_to_fingerprint(rand);
+
+        // calculate period index for token
+        let period_idx = (issued + lifetime)
+            .duration_since(UNIX_EPOCH)
+            .unwrap()
+            .as_nanos()
+            / lifetime.as_nanos();
+
+        // get relevant filter
+        let filter = if period_idx < state.period_idx_1 {
+            // shouldn't happen unless time travels backwards or new_token_lifetime changes
+            warn!("BloomTokenLog presented with token too far in past");
+            return Err(TokenReuseError);
+        } else if period_idx == state.period_idx_1 {
+            &mut state.filter_1
+        } else if period_idx == state.period_idx_1 + 1 {
+            &mut state.filter_2
+        } else {
+            // turn over filters
+            if period_idx == state.period_idx_1 + 2 {
+                swap(&mut state.filter_1, &mut state.filter_2);
+            } else {
+                state.filter_1 = Filter::Set(IdentityHashSet::default());
+            }
+            state.filter_2 = Filter::Set(IdentityHashSet::default());
+            state.period_idx_1 = period_idx - 1;
+
+            &mut state.filter_2
+        };
+
+        // query and insert
+        match *filter {
+            Filter::Set(ref mut hset) => {
+                if !hset.insert(fingerprint) {
+                    return Err(TokenReuseError);
+                }
+
+                if hset.capacity() * size_of::<u64>() > state.filter_max_bytes {
+                    // convert to bloom
+                    let mut bloom = BloomFilter::with_num_bits(state.filter_max_bytes * 8)
+                        .hasher(FxBuildHasher)
+                        .hashes(state.k_num);
+                    for item in hset.iter() {
+                        bloom.insert(item);
+                    }
+                    *filter = Filter::Bloom(bloom);
+                }
+            }
+            Filter::Bloom(ref mut bloom) => {
+                if bloom.insert(&fingerprint) {
+                    return Err(TokenReuseError);
+                }
+            }
+        }
+
+        Ok(())
+    }
+}
+
+const DEFAULT_MAX_BYTES: usize = 10 << 20;
+const DEFAULT_EXPECTED_HITS: u64 = 1_000_000;
+
+/// Default to 20 MiB max memory consumption and expected one million hits
+impl Default for BloomTokenLog {
+    fn default() -> Self {
+        Self::new_expected_items(DEFAULT_MAX_BYTES, DEFAULT_EXPECTED_HITS)
+    }
+}

quinn-proto/src/config.rs

@@ -11,14 +11,16 @@ use rustls::client::WebPkiServerVerifier;
 use rustls::pki_types::{CertificateDer, PrivateKeyDer};
 use thiserror::Error;
 
+#[cfg(feature = "fastbloom")]
+use crate::bloom_token_log::BloomTokenLog;
 #[cfg(any(feature = "rustls-aws-lc-rs", feature = "rustls-ring"))]
 use crate::crypto::rustls::{configured_provider, QuicServerConfig};
 use crate::{
     cid_generator::{ConnectionIdGenerator, HashedConnectionIdGenerator},
     congestion,
     crypto::{self, HandshakeTokenKey, HmacKey},
     shared::ConnectionId,
-    Duration, RandomConnectionIdGenerator, VarInt, VarIntBoundsExceeded,
+    Duration, RandomConnectionIdGenerator, TokenLog, VarInt, VarIntBoundsExceeded,
     DEFAULT_SUPPORTED_VERSIONS, INITIAL_MTU, MAX_CID_SIZE, MAX_UDP_PAYLOAD,
 };
 
@@ -810,6 +812,16 @@ pub struct ServerConfig {
     /// Duration after a retry token was issued for which it's considered valid
     pub(crate) retry_token_lifetime: Duration,
 
+    /// Duration after an address validation token was issued for which it's considered valid
+    pub(crate) validation_token_lifetime: Duration,
+
+    /// Responsible for limiting clients' ability to reuse tokens from NEW_TOKEN frames
+    pub(crate) validation_token_log: Option<Arc<dyn TokenLog>>,
+
+    /// Number of address validation tokens sent to a client via NEW_TOKEN frames when its path is
+    /// validated
+    pub(crate) validation_tokens_sent: u32,
+
     /// Whether to allow clients to migrate to new addresses
     ///
     /// Improves behavior for clients that move between different internet connections or suffer NAT
@@ -825,19 +837,27 @@ pub struct ServerConfig {
 }
 
 const DEFAULT_RETRY_TOKEN_LIFETIME_SECS: u64 = 15;
+const DEFAULT_VALIDATION_TOKEN_LIFETIME_SECS: u64 = 2 * 7 * 24 * 60 * 60;
 
 impl ServerConfig {
     /// Create a default config with a particular handshake token key
     pub fn new(
         crypto: Arc<dyn crypto::ServerConfig>,
         token_key: Arc<dyn HandshakeTokenKey>,
     ) -> Self {
+        #[cfg(feature = "fastbloom")]
+        let validation_token_log = Some(Arc::new(BloomTokenLog::default()) as _);
+        #[cfg(not(feature = "fastbloom"))]
+        let validation_token_log = None;
         Self {
             transport: Arc::new(TransportConfig::default()),
             crypto,
 
             token_key,
             retry_token_lifetime: Duration::from_secs(DEFAULT_RETRY_TOKEN_LIFETIME_SECS),
+            validation_token_lifetime: Duration::from_secs(DEFAULT_VALIDATION_TOKEN_LIFETIME_SECS),
+            validation_token_log,
+            validation_tokens_sent: 2,
 
             migration: true,
 
@@ -870,6 +890,38 @@ impl ServerConfig {
         self
     }
 
+    /// Duration after an address validation token was issued for which it's considered valid
+    ///
+    /// This refers only to tokens sent in NEW_TOKEN frames, in contrast to retry tokens.
+    ///
+    /// Defaults to 2 weeks.
+    pub fn validation_token_lifetime(&mut self, value: Duration) -> &mut Self {
+        self.validation_token_lifetime = value;
+        self
+    }
+
+    /// Set a custom [`TokenLog`]
+    ///
+    /// Setting this to `None` makes the server ignore all address validation tokens (that is,
+    /// tokens originating from NEW_TOKEN frames--retry tokens may still be accepted).
+    ///
+    /// Defaults to a default [`BloomTokenLog`], unless the `fastbloom` default feature is
+    /// disabled, in which case this defaults to `None`.
+    pub fn validation_token_log(&mut self, log: Option<Arc<dyn TokenLog>>) -> &mut Self {
+        self.validation_token_log = log;
+        self
+    }
+
+    /// Number of address validation tokens sent to a client when its path is validated
+    ///
+    /// This refers only to tokens sent in NEW_TOKEN frames, in contrast to retry tokens.
+    ///
+    /// Defaults to 2.
+    pub fn validation_tokens_sent(&mut self, value: u32) -> &mut Self {
+        self.validation_tokens_sent = value;
+        self
+    }
+
     /// Whether to allow clients to migrate to new addresses
     ///
     /// Improves behavior for clients that move between different internet connections or suffer NAT
@@ -987,6 +1039,8 @@ impl fmt::Debug for ServerConfig {
             // crypto not debug
             // token not debug
             .field("retry_token_lifetime", &self.retry_token_lifetime)
+            .field("validation_token_lifetime", &self.validation_token_lifetime)
+            .field("validation_tokens_sent", &self.validation_tokens_sent)
             .field("migration", &self.migration)
             .field("preferred_address_v4", &self.preferred_address_v4)
             .field("preferred_address_v6", &self.preferred_address_v6)