test(proto): keep test connections on their initial keys

noq-proto/src/config/transport.rs

@@ -54,6 +54,8 @@ pub struct TransportConfig {
     pub(crate) datagram_send_buffer_size: usize,
     #[cfg(test)]
     pub(crate) deterministic_packet_numbers: bool,
+    #[cfg(test)]
+    pub(crate) initial_key_phase_size: Option<u64>,
 
     pub(crate) congestion_controller_factory: Arc<dyn congestion::ControllerFactory + Send + Sync>,
 
@@ -332,6 +334,23 @@ impl TransportConfig {
         self
     }
 
+    /// Override the size of the first 1-RTT key phase
+    ///
+    /// A fresh connection picks a short key phase at random (10..1000 packets) so that peers
+    /// which mishandle key updates fail early. That makes the first autonomous key update land
+    /// wherever the seed likes, which is a liability in tests that merely happen to run long
+    /// enough: it arms a `KeyDiscard` timer ahead of the deadline they are waiting for. Set this
+    /// to a size the test will never reach to keep the connection on its initial keys; peer
+    /// initiated and forced updates are unaffected, and the AEAD confidentiality limit still
+    /// applies.
+    ///
+    /// Defaults to `None`, i.e. the random production behavior.
+    #[cfg(test)]
+    pub(crate) fn initial_key_phase_size(&mut self, packets: Option<u64>) -> &mut Self {
+        self.initial_key_phase_size = packets;
+        self
+    }
+
     /// How to construct new `congestion::Controller`s
     ///
     /// Typically the refcounted configuration of a `congestion::Controller`,
@@ -582,6 +601,8 @@ impl Default for TransportConfig {
             datagram_send_buffer_size: 1024 * 1024,
             #[cfg(test)]
             deterministic_packet_numbers: false,
+            #[cfg(test)]
+            initial_key_phase_size: None,
 
             congestion_controller_factory: Arc::new(congestion::CubicConfig::default()),
 
@@ -631,6 +652,8 @@ impl fmt::Debug for TransportConfig {
             datagram_send_buffer_size,
             #[cfg(test)]
                 deterministic_packet_numbers: _,
+            #[cfg(test)]
+                initial_key_phase_size: _,
             congestion_controller_factory: _,
             enable_segmentation_offload,
             address_discovery_role,

noq-proto/src/connection/mod.rs

@@ -360,9 +360,13 @@ impl Connection {
             ),
         )]);
 
+        let crypto_state = CryptoState::new(crypto, init_cid, side, &mut rng);
+        #[cfg(test)]
+        let crypto_state = crypto_state.with_key_phase_size(config.initial_key_phase_size);
+
         let mut this = Self {
             endpoint_config,
-            crypto_state: CryptoState::new(crypto, init_cid, side, &mut rng),
+            crypto_state,
             handshake_cid: local_cid,
             remote_handshake_cid: remote_cid,
             local_cid_state,
@@ -6819,6 +6823,12 @@ impl Connection {
         Some(packet.payload.to_vec())
     }
 
+    /// How many 1-RTT packets may still be sent before the keys are updated
+    #[cfg(test)]
+    pub(crate) fn key_phase_size(&self) -> u64 {
+        self.crypto_state.key_phase_size
+    }
+
     /// The number of bytes of packets containing retransmittable frames that have not been
     /// acknowledged or declared lost.
     #[cfg(test)]

noq-proto/src/connection/packet_crypto.rs

@@ -137,6 +137,17 @@ impl CryptoState {
         }
     }
 
+    /// Override how many packets the first key phase lasts
+    ///
+    /// See `TransportConfig::initial_key_phase_size`.
+    #[cfg(test)]
+    pub(super) fn with_key_phase_size(mut self, packets: Option<u64>) -> Self {
+        if let Some(packets) = packets {
+            self.key_phase_size = packets;
+        }
+        self
+    }
+
     /// Removes header protection of a packet, or returns `None` if the packet was dropped.
     pub(super) fn unprotect_header(
         &self,

noq-proto/src/tests/mod.rs

@@ -1196,6 +1196,25 @@ fn streams_blocked_not_sent_under_limit() {
     );
 }
 
+/// A built pair holds a connection's initial keys unless the test says otherwise
+///
+/// Production starts a connection with a short key phase (10..1000 packets) so that peers which
+/// mishandle key updates fail early. That puts the first autonomous update at an unpredictable
+/// point, so `ConnPairBuilder` pins it out of reach; see `ConnPairBuilder::with_key_update`.
+#[test]
+fn connpair_initial_key_phase_is_pinned() {
+    let pair = ConnPair::builder().connect();
+    assert_eq!(pair.conn(Client).key_phase_size(), u64::MAX);
+    assert_eq!(pair.conn(Server).key_phase_size(), u64::MAX);
+
+    let pair = ConnPair::builder().with_key_update().connect();
+    let size = pair.conn(Client).key_phase_size();
+    assert!(
+        (10..1000).contains(&size),
+        "opted into key updates, expected a short initial phase, got {size}"
+    );
+}
+
 #[test]
 fn key_update_simple() {
     let _guard = subscribe();

noq-proto/src/tests/util.rs

@@ -521,6 +521,7 @@ pub(super) struct ConnPairBuilder {
     pub(super) client_endpoint_cfg: EndpointConfig,
     server_cfg: Option<ServerConfig>,
     routes: Option<Routing>,
+    key_updates: bool,
 }
 
 impl Default for ConnPairBuilder {
@@ -542,6 +543,7 @@ impl Default for ConnPairBuilder {
             client_endpoint_cfg: Default::default(),
             server_cfg: None,
             routes: None,
+            key_updates: false,
         }
     }
 }
@@ -566,6 +568,17 @@ impl ConnPairBuilder {
         self
     }
 
+    /// Let the connection update its keys on its own, as it pleases.
+    ///
+    /// By default a built pair pins the first key phase out of reach, so that a test cannot be
+    /// interrupted by an autonomous key update it never asked for. Tests about key updates
+    /// themselves, or about what an update does to timers and paths, should call this. Peer
+    /// initiated and forced updates (`Self::force_key_update`) work either way.
+    pub(super) fn with_key_update(mut self) -> Self {
+        self.key_updates = true;
+        self
+    }
+
     /// Enables multipath.
     pub(super) fn enable_multipath(mut self) -> Self {
         self.server_transport_cfg
@@ -619,14 +632,23 @@ impl ConnPairBuilder {
         let Self {
             latency,
             mtu,
-            server_transport_cfg,
-            client_transport_cfg,
+            mut server_transport_cfg,
+            mut client_transport_cfg,
             server_endpoint_cfg,
             client_endpoint_cfg,
             server_cfg,
             routes,
+            key_updates,
         } = self;
 
+        // Tests that are not about key updates should not be subject to one: a fresh connection
+        // picks a short key phase at random, so the first autonomous update lands wherever the
+        // seed likes. Opt back in with `Self::with_key_update`.
+        if !key_updates {
+            client_transport_cfg.initial_key_phase_size(Some(u64::MAX));
+            server_transport_cfg.initial_key_phase_size(Some(u64::MAX));
+        }
+
         let server_cfg = server_cfg.unwrap_or(ServerConfig {
             transport: Arc::new(server_transport_cfg),
             ..server_config()