feat(s2n-quic-transport): discard client initial keys when we have handshake keys

quic/s2n-quic-transport/src/connection/connection_impl.rs

@@ -785,6 +785,7 @@ impl<Config: endpoint::Config> connection::Trait for ConnectionImpl<Config> {
             error,
             self.path_manager.active_path_mut(),
             active_path_id,
+            timestamp,
             &mut publisher,
         );
     }
@@ -1339,6 +1340,7 @@ impl<Config: endpoint::Config> connection::Trait for ConnectionImpl<Config> {
                 self.space_manager.discard_initial(
                     self.path_manager.active_path_mut(),
                     path_id,
+                    datagram.timestamp,
                     &mut publisher,
                 );
             }

quic/s2n-quic-transport/src/connection/transmission.rs

@@ -274,19 +274,6 @@ impl<'a, 'sub, Config: endpoint::Config> tx::Message for ConnectionTransmission<
                     encoder,
                 ) {
                     Ok((outcome, encoder)) => {
-                        //= https://www.rfc-editor.org/rfc/rfc9001#section-4.9.1
-                        //# a client MUST discard Initial keys when it first sends a
-                        //# Handshake packet
-
-                        if Config::ENDPOINT_TYPE.is_client() {
-                            let path = &mut self.context.path_manager[self.context.path_id];
-                            space_manager.discard_initial(
-                                path,
-                                self.context.path_id,
-                                self.context.publisher,
-                            );
-                        }
-
                         *self.context.outcome += outcome;
                         encoder
                     }
@@ -308,13 +295,24 @@ impl<'a, 'sub, Config: endpoint::Config> tx::Message for ConnectionTransmission<
                     }
                 };
 
+                //= https://www.rfc-editor.org/rfc/rfc9001#section-4.9.1
+                //# a client MUST discard Initial keys when it first sends a
+                //# Handshake packet
+                if Config::ENDPOINT_TYPE.is_client() {
+                    space_manager.discard_initial(
+                        &mut self.context.path_manager[self.context.path_id],
+                        self.context.path_id,
+                        self.context.timestamp,
+                        self.context.publisher,
+                    );
+                }
+
                 //= https://www.rfc-editor.org/rfc/rfc9001#section-4.9.2
                 //# An endpoint MUST discard its handshake keys when the TLS handshake is
                 //# confirmed (Section 4.1.2).
                 if space_manager.is_handshake_confirmed() {
-                    let path = &mut self.context.path_manager[self.context.path_id];
                     space_manager.discard_handshake(
-                        path,
+                        &mut self.context.path_manager[self.context.path_id],
                         self.context.path_id,
                         self.context.publisher,
                     );

quic/s2n-quic-transport/src/space/handshake.rs

@@ -274,8 +274,7 @@ impl<Config: endpoint::Config> HandshakeSpace<Config> {
         //# datagram unblocks it, even if none of the packets in the datagram are
         //# successfully processed.  In such a case, the PTO timer will need to
         //# be re-armed.
-        self.recovery_manager
-            .update_pto_timer(path, timestamp, is_handshake_confirmed);
+        self.update_pto_timer(path, timestamp, is_handshake_confirmed);
     }
 
     /// Called when the connection timer expired
@@ -309,6 +308,17 @@ impl<Config: endpoint::Config> HandshakeSpace<Config> {
             .on_packet_number_space_discarded(path, path_id, publisher);
     }
 
+    /// Updates the probe timeout timer
+    pub fn update_pto_timer(
+        &mut self,
+        path: &Path<Config>,
+        now: Timestamp,
+        is_handshake_confirmed: bool,
+    ) {
+        self.recovery_manager
+            .update_pto_timer(path, now, is_handshake_confirmed);
+    }
+
     pub fn requires_probe(&self) -> bool {
         self.recovery_manager.requires_probe()
     }

quic/s2n-quic-transport/src/space/mod.rs

@@ -89,52 +89,16 @@ impl<Config: endpoint::Config> fmt::Debug for PacketSpaceManager<Config> {
 }
 
 macro_rules! packet_space_api {
-    ($ty:ty, $field:ident, $get_mut:ident $(, $discard:ident $(, $assert_discard:ident)? )?) => {
+    ($ty:ty, $field:ident, $get_mut:ident) => {
         #[allow(dead_code)]
         pub fn $field(&self) -> Option<&$ty> {
-            self.$field
-                .as_ref()
-                .map(Box::as_ref)
+            self.$field.as_ref().map(Box::as_ref)
         }
 
         pub fn $get_mut(&mut self) -> Option<(&mut $ty, &mut HandshakeStatus)> {
-            let space = self.$field
-                .as_mut()
-                .map(Box::as_mut)?;
+            let space = self.$field.as_mut().map(Box::as_mut)?;
             Some((space, &mut self.handshake_status))
         }
-
-        $(
-            pub fn $discard<Pub: event::ConnectionPublisher>(
-                &mut self,
-                path: &mut Path<Config>,
-                path_id: path::Id,
-                publisher: &mut Pub,
-            ) {
-                //= https://www.rfc-editor.org/rfc/rfc9002#section-6.2.2
-                //# When Initial or Handshake keys are discarded, the PTO and loss
-                //# detection timers MUST be reset, because discarding keys indicates
-                //# forward progress and the loss detection timer might have been set for
-                //# a now discarded packet number space.
-                if let Some(mut space) = self.$field.take() {
-                    // reset the PTO backoff value as part of resetting the PTO timer.
-                    path.reset_pto_backoff();
-
-                    space.on_discard(path, path_id, publisher);
-                }
-
-                //= https://www.rfc-editor.org/rfc/rfc9001#section-4.9.1
-                //# Endpoints MUST NOT send
-                //# Initial packets after this point.
-                // By discarding a space, we are no longer capable of sending packets with those
-                // keys.
-
-                debug_assert!(self.$field.is_none(), "space should have been discarded");
-                $(
-                    debug_assert!(self.$assert_discard.is_none(), "space should have been discarded");
-                )?
-            }
-        )?
     };
 }
 
@@ -174,15 +138,9 @@ impl<Config: endpoint::Config> PacketSpaceManager<Config> {
         }
     }
 
-    packet_space_api!(InitialSpace<Config>, initial, initial_mut, discard_initial);
+    packet_space_api!(InitialSpace<Config>, initial, initial_mut);
 
-    packet_space_api!(
-        HandshakeSpace<Config>,
-        handshake,
-        handshake_mut,
-        discard_handshake,
-        initial
-    );
+    packet_space_api!(HandshakeSpace<Config>, handshake, handshake_mut);
 
     packet_space_api!(ApplicationSpace<Config>, application, application_mut);
 
@@ -193,6 +151,87 @@ impl<Config: endpoint::Config> PacketSpaceManager<Config> {
         self.zero_rtt_crypto.as_ref().map(Box::as_ref)
     }
 
+    /// Discard the initial packet space
+    pub fn discard_initial<Pub: event::ConnectionPublisher>(
+        &mut self,
+        path: &mut Path<Config>,
+        path_id: path::Id,
+        now: Timestamp,
+        publisher: &mut Pub,
+    ) {
+        //= https://www.rfc-editor.org/rfc/rfc9002#section-6.2.2
+        //# When Initial or Handshake keys are discarded, the PTO and loss
+        //# detection timers MUST be reset, because discarding keys indicates
+        //# forward progress and the loss detection timer might have been set for
+        //# a now discarded packet number space.
+        if let Some(mut space) = self.initial.take() {
+            path.reset_pto_backoff();
+            space.on_discard(path, path_id, publisher);
+
+            if let Some((handshake, handshake_status)) = self.handshake_mut() {
+                //= https://www.rfc-editor.org/rfc/rfc9002#section-6.2.1
+                //# A sender SHOULD restart its PTO timer every time an ack-eliciting
+                //# packet is sent or acknowledged, or when Initial or Handshake keys are
+                //# discarded (Section 4.9 of [QUIC-TLS]).
+
+                //= https://www.rfc-editor.org/rfc/rfc9002#section-6.2.2.1
+                //# Since the server could be blocked until more datagrams are received
+                //# from the client, it is the client's responsibility to send packets to
+                //# unblock the server until it is certain that the server has finished
+                //# its address validation
+
+                // Arm the PTO timer on the handshake space so the handshake can make progress
+                // even if no handshake packets have been transmitted or received yet
+                handshake.update_pto_timer(path, now, handshake_status.is_confirmed());
+            }
+        }
+
+        //= https://www.rfc-editor.org/rfc/rfc9001#section-4.9.1
+        //# Endpoints MUST NOT send
+        //# Initial packets after this point.
+        // By discarding a space, we are no longer capable of sending packets with those
+        // keys.
+
+        debug_assert!(
+            self.initial.is_none(),
+            "initial space should have been discarded"
+        );
+    }
+
+    /// Discard the handshake packet space
+    pub fn discard_handshake<Pub: event::ConnectionPublisher>(
+        &mut self,
+        path: &mut Path<Config>,
+        path_id: path::Id,
+        publisher: &mut Pub,
+    ) {
+        //= https://www.rfc-editor.org/rfc/rfc9002#section-6.2.2
+        //# When Initial or Handshake keys are discarded, the PTO and loss
+        //# detection timers MUST be reset, because discarding keys indicates
+        //# forward progress and the loss detection timer might have been set for
+        //# a now discarded packet number space.
+        if let Some(mut space) = self.handshake.take() {
+            path.reset_pto_backoff();
+            space.on_discard(path, path_id, publisher);
+            // Dropping handshake will clear the PTO timer for the handshake space.
+            // The PTO timer for the application space is reset when the
+            // handshake is confirmed.
+
+            //= https://www.rfc-editor.org/rfc/rfc9002#section-6.2.1
+            //# An endpoint MUST NOT set its PTO timer for the Application Data
+            //# packet number space until the handshake is confirmed.
+        }
+
+        debug_assert!(
+            self.initial.is_none(),
+            "initial space should have been discarded"
+        );
+        debug_assert!(
+            self.handshake.is_none(),
+            "handshake space should have been discarded"
+        );
+    }
+
     pub fn discard_zero_rtt_crypto(&mut self) {
         self.zero_rtt_crypto = None;
     }
@@ -450,11 +489,12 @@ impl<Config: endpoint::Config> PacketSpaceManager<Config> {
         error: connection::Error,
         path: &mut path::Path<Config>,
         path_id: path::Id,
+        now: Timestamp,
         publisher: &mut Pub,
     ) {
         self.session_info = None;
         self.retry_cid = None;
-        self.discard_initial(path, path_id, publisher);
+        self.discard_initial(path, path_id, now, publisher);
         self.discard_handshake(path, path_id, publisher);
         self.discard_zero_rtt_crypto();
 

quic/s2n-quic/src/tests.rs

@@ -31,6 +31,7 @@ mod blackhole;
 mod interceptor;
 mod mtu;
 mod no_tls;
+mod pto;
 mod self_test;
 
 // TODO: https://github.com/aws/s2n-quic/issues/1726

quic/s2n-quic/src/tests/pto.rs

@@ -0,0 +1,94 @@
+// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+// SPDX-License-Identifier: Apache-2.0
+
+use super::*;
+use s2n_codec::EncoderBuffer;
+use s2n_quic_core::{
+    event::api::{PacketHeader, Subject},
+    packet::interceptor::{Interceptor, Packet},
+};
+
+/// This test ensures the PTO timer in the Handshake space is armed even
+/// when the client does not otherwise receive or send any handshake
+/// packets
+#[test]
+fn handshake_pto_timer_is_armed() {
+    let model = Model::default();
+    let pto_subscriber = recorder::Pto::new();
+    let packet_sent_subscriber = recorder::PacketSent::new();
+    let pto_events = pto_subscriber.events();
+    let packet_sent_events = packet_sent_subscriber.events();
+
+    test(model.clone(), |handle| {
+        let mut server = Server::builder()
+            .with_io(handle.builder().build()?)?
+            .with_tls(SERVER_CERTS)?
+            .with_packet_interceptor(DropHandshakeTx)?
+            .start()?;
+
+        let addr = server.local_addr()?;
+        spawn(async move {
+            // We would expect this connection to time out since the server
+            // is not able to send any handshake packets
+            assert!(server.accept().await.is_none());
+        });
+
+        let client = Client::builder()
+            .with_io(handle.builder().build().unwrap())?
+            .with_tls(certificates::CERT_PEM)?
+            .with_event((pto_subscriber, packet_sent_subscriber))?
+            .start()?;
+
+        primary::spawn(async move {
+            let connect = Connect::new(addr).with_server_name("localhost");
+            // We would expect this connection to time out since the server
+            // is not able to send any handshake packets
+            assert!(client.connect(connect).await.is_err());
+        });
+
+        Ok(addr)
+    })
+    .unwrap();
+
+    let pto_events = pto_events.lock().unwrap();
+    let pto_count = *pto_events.iter().max().unwrap_or(&0) as usize;
+
+    // Assert that the client sent some PTOs
+    assert!(pto_count > 0);
+
+    let packet_sent_events = packet_sent_events.lock().unwrap();
+    let initial_packets_sent = packet_sent_events
+        .iter()
+        .filter(|&packet_sent| matches!(packet_sent.packet_header, PacketHeader::Initial { .. }))
+        .count();
+    let handshake_packets_sent = packet_sent_events
+        .iter()
+        .filter(|&packet_sent| matches!(packet_sent.packet_header, PacketHeader::Handshake { .. }))
+        .count();
+
+    // Assert that only 2 initial packets were sent (the Initial[ClientHello] and the Initial[ACK])
+    assert_eq!(2, initial_packets_sent);
+
+    // Assert that all handshake packets that were sent were due to the PTO timer firing.
+    // The first PTO that fires will send a single packet, since there are no packets
+    // in flight. Subsequent PTOs will send two packets.
+    let expected_handshake_packet_count = pto_count * 2 - 1;
+    assert_eq!(expected_handshake_packet_count, handshake_packets_sent);
+}
+
+/// Drops all outgoing handshake packets
+struct DropHandshakeTx;
+
+impl Interceptor for DropHandshakeTx {
+    #[inline]
+    fn intercept_tx_payload(
+        &mut self,
+        _subject: &Subject,
+        packet: &Packet,
+        payload: &mut EncoderBuffer,
+    ) {
+        if packet.number.space().is_handshake() {
+            payload.set_position(0);
+        }
+    }
+}