fix: process packets from different sources before handshake confirmed

quic/s2n-quic-transport/src/path/manager.rs

@@ -249,44 +249,25 @@ impl<Config: endpoint::Config> Manager<Config> {
     ) -> Result<(Id, AmplificationOutcome), DatagramDropReason> {
         let valid_initial_received = self.valid_initial_received();
 
-        if let Some((id, path)) = self.path_mut(path_handle) {
-            let source_cid_changed = datagram
-                .source_connection_id
-                .is_some_and(|scid| scid != path.peer_connection_id && valid_initial_received);
-
-            if source_cid_changed {
-                //= https://www.rfc-editor.org/rfc/rfc9000#section-7.2
-                //# Once a client has received a valid Initial packet from the server, it MUST
-                //# discard any subsequent packet it receives on that connection with a
-                //# different Source Connection ID.
-
-                //= https://www.rfc-editor.org/rfc/rfc9000#section-7.2
-                //# Any further changes to the Destination Connection ID are only
-                //# permitted if the values are taken from NEW_CONNECTION_ID frames; if
-                //# subsequent Initial packets include a different Source Connection ID,
-                //# they MUST be discarded.
-
-                return Err(DatagramDropReason::InvalidSourceConnectionId);
-            }
-
-            // Update the address if it was resolved
-            //
-            // NOTE: We don't update the server address since this would cause the client to drop
-            // packets from the server.
-
-            //= https://www.rfc-editor.org/rfc/rfc9000#section-9
-            //# If a client receives packets from an unknown server address, the client MUST discard these packets.
-
+        let matched_path = if handshake_confirmed {
+            self.path_mut(path_handle)
+        } else {
             //= https://www.rfc-editor.org/rfc/rfc9000#section-9
-            //# If the peer sent the disable_active_migration transport parameter, an endpoint also MUST NOT send
-            //# packets (including probing packets; see Section 9.1) from a different local address to the address
-            //# the peer used during the handshake, unless the endpoint has acted on a preferred_address transport
-            //# parameter from the peer.
-            if Config::ENDPOINT_TYPE.is_client() {
-                path.handle.maybe_update(path_handle);
-            }
+            //# The design of QUIC relies on endpoints retaining a stable address
+            //# for the duration of the handshake.  An endpoint MUST NOT initiate
+            //# connection migration before the handshake is confirmed, as defined
+            //# in section 4.1.2 of [QUIC-TLS].
+
+            // NOTE: while we must not _initiate_ a migration before the handshake is done,
+            // it doesn't mean we can't handle the packet. So instead we pick the default path.
+            let path_id = self.active_path_id();
+            let path = self.active_path_mut();
+            Some((path_id, path))
+        };
 
-            let amplification_outcome = path.on_bytes_received(datagram.payload_len);
+        if let Some((id, path)) = matched_path {
+            let amplification_outcome =
+                path.on_datagram_received(path_handle, datagram, valid_initial_received)?;
             return Ok((id, amplification_outcome));
         }
 
@@ -297,15 +278,6 @@ impl<Config: endpoint::Config> Manager<Config> {
             return Err(DatagramDropReason::UnknownServerAddress);
         }
 
-        //= https://www.rfc-editor.org/rfc/rfc9000#section-9
-        //# The design of QUIC relies on endpoints retaining a stable address
-        //# for the duration of the handshake.  An endpoint MUST NOT initiate
-        //# connection migration before the handshake is confirmed, as defined
-        //# in section 4.1.2 of [QUIC-TLS].
-        if !handshake_confirmed {
-            return Err(DatagramDropReason::ConnectionMigrationDuringHandshake);
-        }
-
         //= https://www.rfc-editor.org/rfc/rfc9000#section-9
         //# If the peer
         //# violates this requirement, the endpoint MUST either drop the incoming

quic/s2n-quic-transport/src/path/manager/tests.rs

@@ -765,7 +765,7 @@ fn test_adding_new_path() {
 // - call on_datagram_received with new remote address bit handshake_confirmed false
 //
 // Expectation:
-// - asset on_datagram_received errors
+// - assert on_datagram_received does not error
 // - assert we have one paths
 fn do_not_add_new_path_if_handshake_not_confirmed() {
     // Setup:
@@ -811,7 +811,7 @@ fn do_not_add_new_path_if_handshake_not_confirmed() {
     );
 
     // Expectation:
-    assert!(on_datagram_result.is_err());
+    assert!(on_datagram_result.is_ok());
     assert!(manager.path(&new_addr).is_none());
     assert_eq!(manager.paths.len(), 1);
 }

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

@@ -13,8 +13,10 @@ use crate::{
 };
 use s2n_quic_core::{
     counter::{Counter, Saturating},
-    event::{self, IntoEvent},
-    frame, packet, random,
+    event::{self, builder::DatagramDropReason, IntoEvent},
+    frame,
+    inet::DatagramInfo,
+    packet, random,
     time::{timer, Timestamp},
 };
 
@@ -229,6 +231,54 @@ impl<Config: endpoint::Config> Path<Config> {
         }
     }
 
+    #[inline]
+    pub fn on_datagram_received(
+        &mut self,
+        path_handle: &Config::PathHandle,
+        datagram: &DatagramInfo,
+        valid_initial_received: bool,
+    ) -> Result<AmplificationOutcome, DatagramDropReason> {
+        let source_cid_changed = datagram
+            .source_connection_id
+            .is_some_and(|scid| scid != self.peer_connection_id && valid_initial_received);
+
+        if source_cid_changed {
+            //= https://www.rfc-editor.org/rfc/rfc9000#section-7.2
+            //# Once a client has received a valid Initial packet from the server, it MUST
+            //# discard any subsequent packet it receives on that connection with a
+            //# different Source Connection ID.
+
+            //= https://www.rfc-editor.org/rfc/rfc9000#section-7.2
+            //# Any further changes to the Destination Connection ID are only
+            //# permitted if the values are taken from NEW_CONNECTION_ID frames; if
+            //# subsequent Initial packets include a different Source Connection ID,
+            //# they MUST be discarded.
+
+            return Err(DatagramDropReason::InvalidSourceConnectionId);
+        }
+
+        // Update the address if it was resolved
+        //
+        // NOTE: We don't update the server address since this would cause the client to drop
+        // packets from the server.
+
+        //= https://www.rfc-editor.org/rfc/rfc9000#section-9
+        //# If a client receives packets from an unknown server address, the client MUST discard these packets.
+
+        //= https://www.rfc-editor.org/rfc/rfc9000#section-9
+        //# If the peer sent the disable_active_migration transport parameter, an endpoint also MUST NOT send
+        //# packets (including probing packets; see Section 9.1) from a different local address to the address
+        //# the peer used during the handshake, unless the endpoint has acted on a preferred_address transport
+        //# parameter from the peer.
+        if Config::ENDPOINT_TYPE.is_client() {
+            self.handle.maybe_update(path_handle);
+        }
+
+        let amplification_outcome = self.on_bytes_received(datagram.payload_len);
+
+        Ok(amplification_outcome)
+    }
+
     #[inline]
     pub fn on_timeout<Pub: event::ConnectionPublisher>(
         &mut self,

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

@@ -225,15 +225,13 @@ fn ip_and_port_rebind_test() {
 #[derive(Default)]
 struct RebindPortBeforeHandshakeConfirmed {
     datagram_count: usize,
-    changed_port: bool,
 }
 
 const REBIND_PORT: u16 = 55555;
 impl Interceptor for RebindPortBeforeHandshakeConfirmed {
     fn intercept_rx_remote_address(&mut self, _subject: &Subject, addr: &mut RemoteAddress) {
-        if self.datagram_count == 1 && !self.changed_port {
+        if (1..5).contains(&self.datagram_count) {
             addr.set_port(REBIND_PORT);
-            self.changed_port = true;
         }
     }
 
@@ -279,14 +277,10 @@ fn rebind_before_handshake_confirmed() {
     .unwrap();
 
     let datagram_dropped_events = datagram_dropped_events.lock().unwrap();
-
-    assert_eq!(1, datagram_dropped_events.len());
-    let event = datagram_dropped_events.first().unwrap();
-    assert!(matches!(
-        event.reason,
-        DatagramDropReason::ConnectionMigrationDuringHandshake { .. },
-    ));
-    assert_eq!(REBIND_PORT, event.remote_addr.port());
+    assert!(
+        datagram_dropped_events.is_empty(),
+        "the server should allow packets to be processed before the handshake completes"
+    );
 }
 
 // Changes the remote address to ipv4-mapped after the first packet