From 11087af9af01c61e02d5cfe5a9016f78c9cf5cba Mon Sep 17 00:00:00 2001
From: Benjamin Saunders <ben.e.saunders@gmail.com>
Date: Sat, 30 Mar 2024 11:42:25 -0700
Subject: [PATCH] Use hashed connection IDs by default

---
 quinn-proto/src/config.rs    |  8 +++-----
 quinn-proto/src/tests/mod.rs | 26 +++++++++++++++++---------
 2 files changed, 20 insertions(+), 14 deletions(-)

diff --git a/quinn-proto/src/config.rs b/quinn-proto/src/config.rs
index 7b06977a6..24d416ec1 100644
--- a/quinn-proto/src/config.rs
+++ b/quinn-proto/src/config.rs
@@ -6,7 +6,7 @@ use thiserror::Error;
 use rand::RngCore;
 
 use crate::{
-    cid_generator::{ConnectionIdGenerator, RandomConnectionIdGenerator},
+    cid_generator::{ConnectionIdGenerator, HashedConnectionIdGenerator},
     congestion,
     crypto::{self, HandshakeTokenKey, HmacKey},
     VarInt, VarIntBoundsExceeded, DEFAULT_SUPPORTED_VERSIONS, INITIAL_MTU, MAX_UDP_PAYLOAD,
@@ -620,7 +620,7 @@ impl EndpointConfig {
     /// Create a default config with a particular `reset_key`
     pub fn new(reset_key: Arc<dyn HmacKey>) -> Self {
         let cid_factory: fn() -> Box<dyn ConnectionIdGenerator> =
-            || Box::<RandomConnectionIdGenerator>::default();
+            || Box::<HashedConnectionIdGenerator>::default();
         Self {
             reset_key,
             max_udp_payload_size: (1500u32 - 28).into(), // Ethernet MTU minus IP + UDP headers
@@ -638,9 +638,7 @@ impl EndpointConfig {
     /// connections involving that  `Endpoint`. A custom CID generator allows applications to embed
     /// information in local connection IDs, e.g. to support stateless packet-level load balancers.
     ///
-    /// `EndpointConfig::new()` applies a default random CID generator factory. This functions
-    /// accepts any customized CID generator to reset CID generator factory that implements
-    /// the `ConnectionIdGenerator` trait.
+    /// Defaults to [`HashedConnectionIdGenerator`].
     pub fn cid_generator<F: Fn() -> Box<dyn ConnectionIdGenerator> + Send + Sync + 'static>(
         &mut self,
         factory: F,
diff --git a/quinn-proto/src/tests/mod.rs b/quinn-proto/src/tests/mod.rs
index b4f8726e3..e631003f0 100644
--- a/quinn-proto/src/tests/mod.rs
+++ b/quinn-proto/src/tests/mod.rs
@@ -173,12 +173,15 @@ fn stateless_retry() {
 #[test]
 fn server_stateless_reset() {
     let _guard = subscribe();
-    let mut reset_key = vec![0; 64];
+    let mut key_material = vec![0; 64];
     let mut rng = rand::thread_rng();
-    rng.fill_bytes(&mut reset_key);
-    let reset_key = hmac::Key::new(hmac::HMAC_SHA256, &reset_key);
+    rng.fill_bytes(&mut key_material);
+    let reset_key = hmac::Key::new(hmac::HMAC_SHA256, &key_material);
+    rng.fill_bytes(&mut key_material);
 
-    let endpoint_config = Arc::new(EndpointConfig::new(Arc::new(reset_key)));
+    let mut endpoint_config = EndpointConfig::new(Arc::new(reset_key));
+    endpoint_config.cid_generator(move || Box::new(HashedConnectionIdGenerator::from_key(0)));
+    let endpoint_config = Arc::new(endpoint_config);
 
     let mut pair = Pair::new(endpoint_config.clone(), server_config());
     let (client_ch, _) = pair.connect();
@@ -200,12 +203,15 @@ fn server_stateless_reset() {
 #[test]
 fn client_stateless_reset() {
     let _guard = subscribe();
-    let mut reset_key = vec![0; 64];
+    let mut key_material = vec![0; 64];
     let mut rng = rand::thread_rng();
-    rng.fill_bytes(&mut reset_key);
-    let reset_key = hmac::Key::new(hmac::HMAC_SHA256, &reset_key);
+    rng.fill_bytes(&mut key_material);
+    let reset_key = hmac::Key::new(hmac::HMAC_SHA256, &key_material);
+    rng.fill_bytes(&mut key_material);
 
-    let endpoint_config = Arc::new(EndpointConfig::new(Arc::new(reset_key)));
+    let mut endpoint_config = EndpointConfig::new(Arc::new(reset_key));
+    endpoint_config.cid_generator(move || Box::new(HashedConnectionIdGenerator::from_key(0)));
+    let endpoint_config = Arc::new(endpoint_config);
 
     let mut pair = Pair::new(endpoint_config.clone(), server_config());
     let (_, server_ch) = pair.connect();
@@ -232,7 +238,9 @@ fn client_stateless_reset() {
 fn stateless_reset_limit() {
     let _guard = subscribe();
     let remote = SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), 42);
-    let endpoint_config = Arc::new(EndpointConfig::default());
+    let mut endpoint_config = EndpointConfig::default();
+    endpoint_config.cid_generator(move || Box::new(RandomConnectionIdGenerator::new(8)));
+    let endpoint_config = Arc::new(endpoint_config);
     let mut endpoint = Endpoint::new(
         endpoint_config.clone(),
         Some(Arc::new(server_config())),