Use keyed connection IDs by default

quinn-rs · Mar 30, 2024 · c66a24a · c66a24a
1 parent de0ad91
commit c66a24a
Show file tree

Hide file tree

Showing 3 changed files with 35 additions and 24 deletions.
diff --git 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, KeyedConnectionIdGenerator},
     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::<KeyedConnectionIdGenerator>::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 [`KeyedConnectionIdGenerator`].
     pub fn cid_generator<F: Fn() -> Box<dyn ConnectionIdGenerator> + Send + Sync + 'static>(
         &mut self,
         factory: F,

diff --git a/quinn-proto/src/endpoint.rs b/quinn-proto/src/endpoint.rs
@@ -180,16 +180,6 @@ impl Endpoint {
             }
         };
 
-        if !first_decode.is_initial()
-            && self
-                .local_cid_generator
-                .validate(first_decode.dst_cid())
-                .is_err()
-        {
-            debug!("dropping packet with invalid CID");
-            return None;
-        }
-
         //
         // Handle packet on existing connection, if any
         //
@@ -278,6 +268,17 @@ impl Endpoint {
         // If we got this far, we're a server receiving a seemingly valid packet for an unknown
         // connection. Send a stateless reset if possible.
         //
+
+        if !first_decode.is_initial()
+            && self
+                .local_cid_generator
+                .validate(first_decode.dst_cid())
+                .is_err()
+        {
+            debug!("dropping packet with invalid CID");
+            return None;
+        }
+
         if !dst_cid.is_empty() {
             return self
                 .stateless_reset(now, datagram_len, addresses, dst_cid, buf)

diff --git a/quinn-proto/src/tests/mod.rs b/quinn-proto/src/tests/mod.rs
@@ -178,12 +178,17 @@ 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 cid_key = Arc::new(hmac::Key::new(hmac::HMAC_SHA256, &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(KeyedConnectionIdGenerator::from_key(cid_key.clone())));
+    let endpoint_config = Arc::new(endpoint_config);
 
     let mut pair = Pair::new(endpoint_config.clone(), server_config());
     let (client_ch, _) = pair.connect();
@@ -205,12 +210,17 @@ 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 cid_key = Arc::new(hmac::Key::new(hmac::HMAC_SHA256, &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(KeyedConnectionIdGenerator::from_key(cid_key.clone())));
+    let endpoint_config = Arc::new(endpoint_config);
 
     let mut pair = Pair::new(endpoint_config.clone(), server_config());
     let (_, server_ch) = pair.connect();
@@ -237,7 +247,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())),