Perform a key update early in a connection's lifetime

quinn-proto/src/connection/mod.rs

@@ -142,6 +142,8 @@ pub struct Connection {
     /// Set if 0-RTT is supported, then cleared when no longer needed.
     zero_rtt_crypto: Option<ZeroRttCrypto>,
     key_phase: bool,
+    /// How many packets are in the current key phase. Used only for `Data` space.
+    key_phase_size: u64,
     /// Transport parameters set by the peer
     peer_params: TransportParameters,
     /// Source ConnectionId of the first packet received from the peer
@@ -284,6 +286,13 @@ impl Connection {
             zero_rtt_enabled: false,
             zero_rtt_crypto: None,
             key_phase: false,
+            // A small initial key phase size ensures peers that don't handle key updates correctly
+            // fail sooner rather than later. It's okay for both peers to do this, as the first one
+            // to perform an update will reset the other's key phase size in `update_keys`, and a
+            // simultaneous key update by both is just like a regular key update with a really fast
+            // response. Inspired by quic-go's similar behavior of performing the first key update
+            // at the 100th short-header packet.
+            key_phase_size: rng.gen_range(10..1000),
             peer_params: TransportParameters::default(),
             orig_rem_cid: rem_cid,
             initial_dst_cid: init_cid,
@@ -3137,13 +3146,18 @@ impl Connection {
     }
 
     fn update_keys(&mut self, end_packet: Option<(u64, Instant)>, remote: bool) {
+        trace!("executing key update");
         // Generate keys for the key phase after the one we're switching to, store them in
         // `next_crypto`, make the contents of `next_crypto` current, and move the current keys into
         // `prev_crypto`.
         let new = self
             .crypto
             .next_1rtt_keys()
             .expect("only called for `Data` packets");
+        self.key_phase_size = new
+            .local
+            .confidentiality_limit()
+            .saturating_sub(KEY_UPDATE_MARGIN);
         let old = mem::replace(
             &mut self.spaces[SpaceId::Data]
                 .crypto
@@ -3477,6 +3491,11 @@ const MIN_PACKET_SPACE: usize = 40;
 /// that numbers around 10 are a good compromise.
 const MAX_TRANSMIT_SEGMENTS: usize = 10;
 
+/// Perform key updates this many packets before the AEAD confidentiality limit.
+///
+/// Chosen arbitrarily, intended to be large enough to prevent spurious connection loss.
+const KEY_UPDATE_MARGIN: u64 = 10000;
+
 struct ZeroRttCrypto {
     header: Box<dyn HeaderKey>,
     packet: Box<dyn PacketKey>,

quinn-proto/src/connection/packet_builder.rs

@@ -40,33 +40,37 @@ impl PacketBuilder {
         version: u32,
     ) -> Option<Self> {
         // Initiate key update if we're approaching the confidentiality limit
-        let confidentiality_limit = conn.spaces[space_id]
-            .crypto
-            .as_ref()
-            .map_or_else(
-                || &conn.zero_rtt_crypto.as_ref().unwrap().packet,
-                |keys| &keys.packet.local,
-            )
-            .confidentiality_limit();
         let sent_with_keys = conn.spaces[space_id].sent_with_keys;
         if space_id == SpaceId::Data {
-            if sent_with_keys.saturating_add(KEY_UPDATE_MARGIN) >= confidentiality_limit {
+            if sent_with_keys >= conn.key_phase_size {
                 conn.initiate_key_update();
             }
-        } else if sent_with_keys.saturating_add(1) == confidentiality_limit {
-            // We still have time to attempt a graceful close
-            conn.close_inner(
-                now,
-                Close::Connection(frame::ConnectionClose {
-                    error_code: TransportErrorCode::AEAD_LIMIT_REACHED,
-                    frame_type: None,
-                    reason: Bytes::from_static(b"confidentiality limit reached"),
-                }),
-            )
-        } else if sent_with_keys > confidentiality_limit {
-            // Confidentiality limited violated and there's nothing we can do
-            conn.kill(TransportError::AEAD_LIMIT_REACHED("confidentiality limit reached").into());
-            return None;
+        } else {
+            let confidentiality_limit = conn.spaces[space_id]
+                .crypto
+                .as_ref()
+                .map_or_else(
+                    || &conn.zero_rtt_crypto.as_ref().unwrap().packet,
+                    |keys| &keys.packet.local,
+                )
+                .confidentiality_limit();
+            if sent_with_keys.saturating_add(1) == confidentiality_limit {
+                // We still have time to attempt a graceful close
+                conn.close_inner(
+                    now,
+                    Close::Connection(frame::ConnectionClose {
+                        error_code: TransportErrorCode::AEAD_LIMIT_REACHED,
+                        frame_type: None,
+                        reason: Bytes::from_static(b"confidentiality limit reached"),
+                    }),
+                )
+            } else if sent_with_keys > confidentiality_limit {
+                // Confidentiality limited violated and there's nothing we can do
+                conn.kill(
+                    TransportError::AEAD_LIMIT_REACHED("confidentiality limit reached").into(),
+                );
+                return None;
+            }
         }
 
         let space = &mut conn.spaces[space_id];
@@ -247,8 +251,3 @@ impl PacketBuilder {
         (buffer.len() - encode_start, pad)
     }
 }
-
-/// Perform key updates this many packets before the AEAD confidentiality limit.
-///
-/// Chosen arbitrarily, intended to be large enough to prevent spurious connection loss.
-const KEY_UPDATE_MARGIN: u64 = 10000;