use mtuconfirmcomplete in test

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

@@ -15,7 +15,10 @@ use s2n_quic_core::{
     crypto::tls,
     dc::testing::MockDcEndpoint,
     event::{
-        api::{DatagramDropReason, DcState, EndpointDatagramDropped, EndpointMeta, Subject},
+        api::{
+            ConnectionMeta, DatagramDropReason, DcState, EndpointDatagramDropped, EndpointMeta,
+            MtuUpdated, Subject,
+        },
         Timestamp,
     },
     frame::ConnectionClose,
@@ -315,7 +318,10 @@ fn self_test<S: ServerProviders, C: ClientProviders>(
     test(model, |handle| {
         let mut server = server
             .with_io(handle.builder().build()?)?
-            .with_event((dc::ConfirmComplete, (tracing_events(), server_subscriber)))?
+            .with_event((
+                (dc::ConfirmComplete, dc::MtuConfirmComplete),
+                (tracing_events(), server_subscriber),
+            ))?
             .with_random(Random::with_seed(456))?
             .start()?;
 
@@ -335,14 +341,18 @@ fn self_test<S: ServerProviders, C: ClientProviders>(
                         }
                     } else {
                         assert!(result.is_ok());
+                        assert!(dc::MtuConfirmComplete::wait_ready(&mut conn).await.is_ok());
                     }
                 }
             }
         });
 
         let client = client
             .with_io(handle.builder().build().unwrap())?
-            .with_event((dc::ConfirmComplete, (tracing_events(), client_subscriber)))?
+            .with_event((
+                (dc::ConfirmComplete, dc::MtuConfirmComplete),
+                (tracing_events(), client_subscriber),
+            ))?
             .with_random(Random::with_seed(456))?
             .start()?;
 
@@ -373,9 +383,10 @@ fn self_test<S: ServerProviders, C: ClientProviders>(
                             .unwrap()
                             .clone();
                         assert_dc_complete(&client_events);
-                        // wait briefly so the ack for the `DC_STATELESS_RESET_TOKENS` frame from the server is sent
-                        // before the client closes the connection. This is only necessary to confirm the `dc::State`
-                        // on the server moves to `DcState::Complete`
+
+                        assert!(dc::MtuConfirmComplete::wait_ready(&mut conn).await.is_ok());
+
+                        // wait briefly for MTU probing to complete on the server
                         delay(Duration::from_millis(100)).await;
                     }
                 }
@@ -435,6 +446,12 @@ fn self_test<S: ServerProviders, C: ClientProviders>(
             .duration_since_start()
     );
 
+    let client_mtu_events = client_events.mtu_updated_events.lock().unwrap().clone();
+    let server_mtu_events = server_events.mtu_updated_events.lock().unwrap().clone();
+
+    assert_mtu_probing_completed(&client_mtu_events);
+    assert_mtu_probing_completed(&server_mtu_events);
+
     Ok((client_events, server_events))
 }
 
@@ -452,6 +469,14 @@ fn assert_dc_complete(events: &[DcStateChangedEvent]) {
     assert!(matches!(events[2].state, DcState::Complete { .. }));
 }
 
+fn assert_mtu_probing_completed(events: &[MtuUpdatedEvent]) {
+    assert!(events.len() > 0);
+    let last_event = events.last().unwrap();
+    assert!(last_event.search_complete);
+    // 1472 = default MaxMtu (1500) - headers
+    assert_eq!(1472, last_event.mtu);
+}
+
 fn convert_io_result(io_result: std::io::Result<()>) -> Option<connection::Error> {
     io_result
         .err()?
@@ -468,9 +493,16 @@ struct DcStateChangedEvent {
     state: DcState,
 }
 
+#[derive(Clone)]
+struct MtuUpdatedEvent {
+    mtu: u16,
+    search_complete: bool,
+}
+
 #[derive(Clone, Default)]
 struct DcRecorder {
     pub dc_state_changed_events: Arc<Mutex<Vec<DcStateChangedEvent>>>,
+    pub mtu_updated_events: Arc<Mutex<Vec<MtuUpdatedEvent>>>,
     pub endpoint_datagram_dropped_events: Arc<Mutex<Vec<EndpointDatagramDropped>>>,
 }
 impl DcRecorder {
@@ -510,6 +542,22 @@ impl events::Subscriber for DcRecorder {
         store(event, &mut buffer);
     }
 
+    fn on_mtu_updated(
+        &mut self,
+        context: &mut Self::ConnectionContext,
+        _meta: &ConnectionMeta,
+        event: &MtuUpdated,
+    ) {
+        let store = |event: &events::MtuUpdated, storage: &mut Vec<MtuUpdatedEvent>| {
+            storage.push(MtuUpdatedEvent {
+                mtu: event.mtu,
+                search_complete: event.search_complete,
+            });
+        };
+        let mut buffer = context.mtu_updated_events.lock().unwrap();
+        store(event, &mut buffer);
+    }
+
     fn on_endpoint_datagram_dropped(
         &mut self,
         _meta: &EndpointMeta,