Use CID generator as parser in connections

quinn-rs · May 22, 2024 · 696deb7 · 696deb7
1 parent 5f9b0e5
commit 696deb7
Show file tree

Hide file tree

Showing 3 changed files with 11 additions and 9 deletions.
diff --git a/quinn-proto/src/config.rs b/quinn-proto/src/config.rs
@@ -620,7 +620,7 @@ pub struct EndpointConfig {
     ///
     /// Create a cid generator for local cid in Endpoint struct
     pub(crate) connection_id_generator_factory:
-        Arc<dyn Fn() -> Box<dyn ConnectionIdGenerator> + Send + Sync>,
+        Arc<dyn Fn() -> Arc<dyn ConnectionIdGenerator> + Send + Sync>,
     pub(crate) supported_versions: Vec<u32>,
     pub(crate) grease_quic_bit: bool,
     /// Minimum interval between outgoing stateless reset packets
@@ -631,7 +631,7 @@ impl EndpointConfig {
     /// Create a default config with a particular `reset_key`
     pub fn new(reset_key: Arc<dyn HmacKey>) -> Self {
         let cid_factory =
-            || -> Box<dyn ConnectionIdGenerator> { Box::<HashedConnectionIdGenerator>::default() };
+            || -> Arc<dyn ConnectionIdGenerator> { Arc::<HashedConnectionIdGenerator>::default() };
         Self {
             reset_key,
             max_udp_payload_size: (1500u32 - 28).into(), // Ethernet MTU minus IP + UDP headers
@@ -650,7 +650,7 @@ impl EndpointConfig {
     /// information in local connection IDs, e.g. to support stateless packet-level load balancers.
     ///
     /// Defaults to [`HashedConnectionIdGenerator`].
-    pub fn cid_generator<F: Fn() -> Box<dyn ConnectionIdGenerator> + Send + Sync + 'static>(
+    pub fn cid_generator<F: Fn() -> Arc<dyn ConnectionIdGenerator> + Send + Sync + 'static>(
         &mut self,
         factory: F,
     ) -> &mut Self {

diff --git a/quinn-proto/src/connection/mod.rs b/quinn-proto/src/connection/mod.rs
@@ -23,8 +23,8 @@ use crate::{
     frame,
     frame::{Close, Datagram, FrameStruct},
     packet::{
-        FixedLengthConnectionIdParser, Header, InitialHeader, InitialPacket, LongType, Packet,
-        PacketNumber, PartialDecode, SpaceId,
+        Header, InitialHeader, InitialPacket, LongType, Packet, PacketNumber, PartialDecode,
+        SpaceId,
     },
     range_set::ArrayRangeSet,
     shared::{
@@ -197,6 +197,7 @@ pub struct Connection {
     retry_token: Bytes,
     /// Identifies Data-space packet numbers to skip. Not used in earlier spaces.
     packet_number_filter: PacketNumberFilter,
+    cid_gen: Arc<dyn ConnectionIdGenerator>,
 
     //
     // Queued non-retransmittable 1-RTT data
@@ -252,7 +253,7 @@ impl Connection {
         remote: SocketAddr,
         local_ip: Option<IpAddr>,
         crypto: Box<dyn crypto::Session>,
-        cid_gen: &dyn ConnectionIdGenerator,
+        cid_gen: Arc<dyn ConnectionIdGenerator>,
         now: Instant,
         version: u32,
         allow_mtud: bool,
@@ -329,6 +330,7 @@ impl Connection {
             },
             #[cfg(not(test))]
             packet_number_filter: PacketNumberFilter::new(&mut rng),
+            cid_gen,
 
             path_responses: PathResponses::default(),
             close: false,
@@ -2101,7 +2103,7 @@ impl Connection {
         while let Some(data) = remaining {
             match PartialDecode::new(
                 data,
-                &FixedLengthConnectionIdParser::new(self.local_cid_state.cid_len()),
+                &*self.cid_gen,
                 &[self.version],
                 self.endpoint_config.grease_quic_bit,
             ) {

diff --git a/quinn-proto/src/endpoint.rs b/quinn-proto/src/endpoint.rs
@@ -44,7 +44,7 @@ pub struct Endpoint {
     rng: StdRng,
     index: ConnectionIndex,
     connections: Slab<ConnectionMeta>,
-    local_cid_generator: Box<dyn ConnectionIdGenerator>,
+    local_cid_generator: Arc<dyn ConnectionIdGenerator>,
     config: Arc<EndpointConfig>,
     server_config: Option<Arc<ServerConfig>>,
     /// Whether the underlying UDP socket promises not to fragment packets
@@ -825,7 +825,7 @@ impl Endpoint {
             addresses.remote,
             addresses.local_ip,
             tls,
-            self.local_cid_generator.as_ref(),
+            self.local_cid_generator.clone(),
             now,
             version,
             self.allow_mtud,