fix: Allow turning off greasing and connection migration (#2100)

* fix: Allow turning off greasing * Add `disable_migration` * fmt * Update neqo-transport/src/connection/tests/mod.rs Co-authored-by: Max Inden <[email protected]> Signed-off-by: Lars Eggert <[email protected]> * Update neqo-transport/src/connection/tests/mod.rs Co-authored-by: Max Inden <[email protected]> Signed-off-by: Lars Eggert <[email protected]> * clippy --------- Signed-off-by: Lars Eggert <[email protected]> Co-authored-by: Max Inden <[email protected]>
mozilla · Sep 13, 2024 · c1e7dc3 · c1e7dc3
1 parent 4aa0319
commit c1e7dc3
Show file tree

Hide file tree

Showing 2 changed files with 34 additions and 2 deletions.
diff --git a/neqo-transport/src/connection/params.rs b/neqo-transport/src/connection/params.rs
@@ -44,6 +44,7 @@ pub enum PreferredAddressConfig {
 /// `ConnectionParameters` use for setting intitial value for QUIC parameters.
 /// This collects configuration like initial limits, protocol version, and
 /// congestion control algorithm.
+#[allow(clippy::struct_excessive_bools)] // We need that many, sorry.
 #[derive(Debug, Clone)]
 pub struct ConnectionParameters {
     versions: VersionConfig,
@@ -78,6 +79,7 @@ pub struct ConnectionParameters {
     incoming_datagram_queue: usize,
     fast_pto: u8,
     grease: bool,
+    disable_migration: bool,
     pacing: bool,
     /// Whether the connection performs PLPMTUD.
     pmtud: bool,
@@ -102,6 +104,7 @@ impl Default for ConnectionParameters {
             incoming_datagram_queue: MAX_QUEUED_DATAGRAMS_DEFAULT,
             fast_pto: FAST_PTO_SCALE,
             grease: true,
+            disable_migration: false,
             pacing: true,
             pmtud: false,
         }
@@ -336,6 +339,12 @@ impl ConnectionParameters {
         self
     }
 
+    #[must_use]
+    pub const fn disable_migration(mut self, disable_migration: bool) -> Self {
+        self.disable_migration = disable_migration;
+        self
+    }
+
     #[must_use]
     pub const fn pacing_enabled(&self) -> bool {
         self.pacing
@@ -373,8 +382,12 @@ impl ConnectionParameters {
             tparams::ACTIVE_CONNECTION_ID_LIMIT,
             u64::try_from(LOCAL_ACTIVE_CID_LIMIT)?,
         );
-        tps.local.set_empty(tparams::DISABLE_MIGRATION);
-        tps.local.set_empty(tparams::GREASE_QUIC_BIT);
+        if self.disable_migration {
+            tps.local.set_empty(tparams::DISABLE_MIGRATION);
+        }
+        if self.grease {
+            tps.local.set_empty(tparams::GREASE_QUIC_BIT);
+        }
         tps.local.set_integer(
             tparams::MAX_ACK_DELAY,
             u64::try_from(DEFAULT_ACK_DELAY.as_millis())?,

diff --git a/neqo-transport/src/connection/tests/mod.rs b/neqo-transport/src/connection/tests/mod.rs
@@ -28,6 +28,7 @@ use crate::{
     pmtud::Pmtud,
     recovery::ACK_ONLY_SIZE_LIMIT,
     stats::{FrameStats, Stats, MAX_PTO_COUNTS},
+    tparams::{DISABLE_MIGRATION, GREASE_QUIC_BIT},
     ConnectionIdDecoder, ConnectionIdGenerator, ConnectionParameters, Error, StreamId, StreamType,
     Version,
 };
@@ -678,3 +679,21 @@ fn create_server() {
     // Server won't have a default path, so no RTT.
     assert_eq!(stats.rtt, Duration::from_secs(0));
 }
+
+#[test]
+fn tp_grease() {
+    for enable in [true, false] {
+        let client = new_client(ConnectionParameters::default().grease(enable));
+        let grease = client.tps.borrow_mut().local.get_empty(GREASE_QUIC_BIT);
+        assert_eq!(enable, grease);
+    }
+}
+
+#[test]
+fn tp_disable_migration() {
+    for disable in [true, false] {
+        let client = new_client(ConnectionParameters::default().disable_migration(disable));
+        let disable_migration = client.tps.borrow_mut().local.get_empty(DISABLE_MIGRATION);
+        assert_eq!(disable, disable_migration);
+    }
+}