From a93d6d0cb03037a73e2cc4ea27bef1516660340f Mon Sep 17 00:00:00 2001
From: LAN Xingcan <thynson@live.com>
Date: Sat, 15 Jun 2024 12:11:51 +0800
Subject: [PATCH] proto: make initial destination cid configurable

---
 quinn-proto/src/config.rs   | 28 ++++++++++++++++++++++++++--
 quinn-proto/src/endpoint.rs |  4 ++--
 2 files changed, 28 insertions(+), 4 deletions(-)

diff --git a/quinn-proto/src/config.rs b/quinn-proto/src/config.rs
index c75fbe00f..94d829f96 100644
--- a/quinn-proto/src/config.rs
+++ b/quinn-proto/src/config.rs
@@ -20,7 +20,9 @@ use crate::{
     cid_generator::{ConnectionIdGenerator, HashedConnectionIdGenerator},
     congestion,
     crypto::{self, HandshakeTokenKey, HmacKey},
-    VarInt, VarIntBoundsExceeded, DEFAULT_SUPPORTED_VERSIONS, INITIAL_MTU, MAX_UDP_PAYLOAD,
+    shared::ConnectionId,
+    RandomConnectionIdGenerator, VarInt, VarIntBoundsExceeded, DEFAULT_SUPPORTED_VERSIONS,
+    INITIAL_MTU, MAX_CID_SIZE, MAX_UDP_PAYLOAD,
 };
 
 /// Parameters governing the core QUIC state machine
@@ -963,6 +965,9 @@ pub struct ClientConfig {
     /// Cryptographic configuration to use
     pub(crate) crypto: Arc<dyn crypto::ClientConfig>,
 
+    /// Provider that populates the destination connection ID of Initial Packets
+    pub(crate) initial_dst_cid_provider: Arc<dyn Fn() -> ConnectionId + Send + Sync>,
+
     /// QUIC protocol version to use
     pub(crate) version: u32,
 }
@@ -973,10 +978,29 @@ impl ClientConfig {
         Self {
             transport: Default::default(),
             crypto,
+            initial_dst_cid_provider: Arc::new(|| {
+                RandomConnectionIdGenerator::new(MAX_CID_SIZE).generate_cid()
+            }),
             version: 1,
         }
     }
 
+    /// Configure how to populate the destination CID of the initial packet when attempting to
+    /// establish a new connection.
+    ///
+    /// By default, it's populated with random bytes with reasonable length, so unless you have
+    /// a good reason, you do not need to change it.
+    ///
+    /// When prefer to override the default, please note that the generated connection ID MUST be
+    /// at least 8 bytes long and unpredictable, as per section 7.2 of RFC 9000.
+    pub fn initial_dst_cid_provider(
+        &mut self,
+        initial_dst_cid_provider: Arc<dyn Fn() -> ConnectionId + Send + Sync>,
+    ) -> &mut Self {
+        self.initial_dst_cid_provider = initial_dst_cid_provider;
+        self
+    }
+
     /// Set a custom [`TransportConfig`]
     pub fn transport_config(&mut self, transport: Arc<TransportConfig>) -> &mut Self {
         self.transport = transport;
@@ -1016,7 +1040,7 @@ impl fmt::Debug for ClientConfig {
             .field("transport", &self.transport)
             .field("crypto", &"ClientConfig { elided }")
             .field("version", &self.version)
-            .finish()
+            .finish_non_exhaustive()
     }
 }
 
diff --git a/quinn-proto/src/endpoint.rs b/quinn-proto/src/endpoint.rs
index e4df5787d..22744c1a3 100644
--- a/quinn-proto/src/endpoint.rs
+++ b/quinn-proto/src/endpoint.rs
@@ -16,7 +16,7 @@ use thiserror::Error;
 use tracing::{debug, error, trace, warn};
 
 use crate::{
-    cid_generator::{ConnectionIdGenerator, RandomConnectionIdGenerator},
+    cid_generator::ConnectionIdGenerator,
     coding::BufMutExt,
     config::{ClientConfig, EndpointConfig, ServerConfig},
     connection::{Connection, ConnectionError},
@@ -395,7 +395,7 @@ impl Endpoint {
             return Err(ConnectError::UnsupportedVersion);
         }
 
-        let remote_id = RandomConnectionIdGenerator::new(MAX_CID_SIZE).generate_cid();
+        let remote_id = (config.initial_dst_cid_provider)();
         trace!(initial_dcid = %remote_id);
 
         let ch = ConnectionHandle(self.connections.vacant_key());