Return owned Connection on NewConnection and rely on close for cleanup

Author: LindonAliu

src/rs/lib.rs

@@ -867,9 +867,6 @@ extern "C" fn raw_conn_callback(
     event: *mut ffi::QUIC_CONNECTION_EVENT,
 ) -> QUIC_STATUS {
     let event_ref = unsafe { event.as_ref().expect("cannot get connection event") };
-    let cleanup_ctx =
-        event_ref.Type == ffi::QUIC_CONNECTION_EVENT_TYPE_QUIC_CONNECTION_EVENT_SHUTDOWN_COMPLETE;
-
     let status = match unsafe { (context as *mut Box<ConnectionCallback>).as_mut() } {
         Some(f) => {
             let event = ConnectionEvent::from(event_ref);
@@ -883,11 +880,6 @@ extern "C" fn raw_conn_callback(
         None => StatusCode::QUIC_STATUS_SUCCESS.into(),
     };
 
-    if cleanup_ctx {
-        let conn = unsafe { ConnectionRef::from_raw(connection) };
-        conn.consume_callback_ctx();
-    }
-
     status
 }
 
@@ -1058,14 +1050,12 @@ extern "C" fn raw_listener_callback(
 ) -> QUIC_STATUS {
     let listner_ref = unsafe { ListenerRef::from_raw(listener) };
     let event = ListenerEvent::from(unsafe { event.as_ref().expect("fail to get listener event") });
-    let f = unsafe {
-        (context as *mut Box<ListenerCallback>)
-            .as_ref() // allow mutation
-            .expect("cannot get ListenerCallback from ctx")
-    };
-    match f(listner_ref, event) {
-        Ok(_) => StatusCode::QUIC_STATUS_SUCCESS.into(),
-        Err(e) => e.0,
+    match unsafe { (context as *mut Box<ListenerCallback>).as_ref() } {
+        Some(f) => match f(listner_ref, event) {
+            Ok(_) => StatusCode::QUIC_STATUS_SUCCESS.into(),
+            Err(e) => e.0,
+        },
+        None => StatusCode::QUIC_STATUS_SUCCESS.into(),
     }
 }
 

src/rs/server_client_test.rs

@@ -180,6 +180,8 @@ fn test_server_client() {
             } => {
                 connection.set_callback_handler(conn_handler.clone());
                 connection.set_configuration(&config_cp)?;
+                // Keep the connection alive; will be closed on ShutdownComplete.
+                let _ = unsafe { connection.into_raw() };
             }
             crate::ListenerEvent::StopComplete {
                 app_close_in_progress: _,
@@ -345,6 +347,8 @@ fn connection_ref_callback_cleanup() {
                     Ok(())
                 });
                 connection.set_configuration(&server_config)?;
+                // Keep connection alive until ShutdownComplete closes it.
+                let _ = unsafe { connection.into_raw() };
             }
             Ok(())
         }
@@ -391,6 +395,8 @@ fn connection_ref_callback_cleanup() {
     let raw_conn = server_handle_rx
         .recv_timeout(Duration::from_secs(5))
         .expect("Server did not receive shutdown event");
+    // Close now to drop the server-side callback context.
+    unsafe { Connection::from_raw(raw_conn) };
     client_done_rx
         .recv_timeout(Duration::from_secs(5))
         .expect("Client did not complete shutdown");
@@ -405,7 +411,5 @@ fn connection_ref_callback_cleanup() {
         1,
         "ConnectionRef callback context was not cleaned up"
     );
-
-    unsafe { Connection::from_raw(raw_conn) };
     listener.stop();
 }

src/rs/types.rs

@@ -12,8 +12,7 @@ pub enum ListenerEvent<'a> {
         /// User app needs to take ownership of this new connection.
         /// User app needs to set configuration for this connection
         /// before returning from the callback.
-        /// TODO: Make this Connection type.
-        connection: crate::ConnectionRef,
+        connection: crate::Connection,
     },
     StopComplete {
         app_close_in_progress: bool,
@@ -67,7 +66,7 @@ impl<'a> From<&'a crate::ffi::QUIC_LISTENER_EVENT> for ListenerEvent<'a> {
                 let ev = unsafe { &value.__bindgen_anon_1.NEW_CONNECTION };
                 Self::NewConnection {
                     info: NewConnectionInfo::from(unsafe { ev.Info.as_ref().unwrap() }),
-                    connection: unsafe { crate::ConnectionRef::from_raw(ev.Connection) },
+                    connection: unsafe { crate::Connection::from_raw(ev.Connection) },
                 }
             }
             crate::ffi::QUIC_LISTENER_EVENT_TYPE_QUIC_LISTENER_EVENT_STOP_COMPLETE => {