From f91ccc2f975225a278bc68238555ed0923e676ba Mon Sep 17 00:00:00 2001
From: William Yang <mscame@gmail.com>
Date: Mon, 28 Aug 2023 20:04:51 +0200
Subject: [PATCH] feat: close listeners in registration rundown (#3823)

---
 src/core/listener.c     | 23 +++++++++++++++++++++++
 src/core/listener.h     |  5 +++++
 src/core/registration.c |  9 +++++++++
 src/core/registration.h |  7 ++++++-
 4 files changed, 43 insertions(+), 1 deletion(-)

diff --git a/src/core/listener.c b/src/core/listener.c
index 9cf4c5a299..61dfdd4b56 100644
--- a/src/core/listener.c
+++ b/src/core/listener.c
@@ -75,9 +75,26 @@ MsQuicListenerOpen(
     QuicSiloAddRef(Listener->Silo);
 #endif
 
+    BOOLEAN RegistrationShuttingDown;
+
     BOOLEAN Result = CxPlatRundownAcquire(&Registration->Rundown);
     CXPLAT_DBG_ASSERT(Result); UNREFERENCED_PARAMETER(Result);
 
+    CxPlatDispatchLockAcquire(&Registration->ConnectionLock);
+    RegistrationShuttingDown = Registration->ShuttingDown;
+    if (!RegistrationShuttingDown) {
+        CxPlatListInsertTail(&Registration->Listeners, &Listener->RegistrationLink);
+    }
+    CxPlatDispatchLockRelease(&Registration->ConnectionLock);
+
+    if (RegistrationShuttingDown) {
+        CxPlatRundownRelease(&Registration->Rundown);
+        CxPlatEventUninitialize(Listener->StopEvent);
+        CXPLAT_FREE(Listener, QUIC_POOL_LISTENER);
+        Listener = NULL;
+        Status = QUIC_STATUS_INVALID_STATE;
+        goto Error;
+    }
     QuicTraceEvent(
         ListenerCreated,
         "[list][%p] Created, Registration=%p",
@@ -117,6 +134,12 @@ QuicListenerFree(
     QuicSiloRelease(Listener->Silo);
 #endif
 
+    CxPlatDispatchLockAcquire(&Listener->Registration->ConnectionLock);
+    if (!Listener->Registration->ShuttingDown) {
+        CxPlatListEntryRemove(&Listener->RegistrationLink);
+    }
+    CxPlatDispatchLockRelease(&Listener->Registration->ConnectionLock);
+
     CxPlatRefUninitialize(&Listener->RefCount);
     CxPlatEventUninitialize(Listener->StopEvent);
     CXPLAT_DBG_ASSERT(Listener->AlpnList == NULL);
diff --git a/src/core/listener.h b/src/core/listener.h
index d2c847db66..ae30967354 100644
--- a/src/core/listener.h
+++ b/src/core/listener.h
@@ -52,6 +52,11 @@ typedef struct QUIC_LISTENER {
     //
     QUIC_REGISTRATION* Registration;
 
+    //
+    // Link into the registrations's list of listeners.
+    //
+    CXPLAT_LIST_ENTRY RegistrationLink;
+
 #ifdef QUIC_SILO
     //
     // The silo.
diff --git a/src/core/registration.c b/src/core/registration.c
index c762b8e29b..b48fcd5351 100644
--- a/src/core/registration.c
+++ b/src/core/registration.c
@@ -82,6 +82,7 @@ MsQuicRegistrationOpen(
     CxPlatListInitializeHead(&Registration->Configurations);
     CxPlatDispatchLockInitialize(&Registration->ConnectionLock);
     CxPlatListInitializeHead(&Registration->Connections);
+    CxPlatListInitializeHead(&Registration->Listeners);
     CxPlatRundownInitialize(&Registration->Rundown);
     Registration->AppNameLength = (uint8_t)(AppNameLength + 1);
     if (AppNameLength != 0) {
@@ -248,6 +249,14 @@ MsQuicRegistrationShutdown(
         }
 
         CxPlatDispatchLockRelease(&Registration->ConnectionLock);
+
+        Entry = Registration->Listeners.Flink;
+        while (Entry != &Registration->Listeners) {
+            QUIC_LISTENER* Listener =
+                CXPLAT_CONTAINING_RECORD(Entry, QUIC_LISTENER, RegistrationLink);
+            Entry = Entry->Flink;
+            MsQuicListenerStop((HQUIC)Listener);
+        }
     }
 
 Exit:
diff --git a/src/core/registration.h b/src/core/registration.h
index 4aaa1555aa..5e68448af9 100644
--- a/src/core/registration.h
+++ b/src/core/registration.h
@@ -81,7 +81,7 @@ typedef struct QUIC_REGISTRATION {
     CXPLAT_LIST_ENTRY Configurations;
 
     //
-    // Protects access to the Connections list.
+    // Protects access to the Connections list and the Listeners list.
     //
     CXPLAT_DISPATCH_LOCK ConnectionLock;
 
@@ -90,6 +90,11 @@ typedef struct QUIC_REGISTRATION {
     //
     CXPLAT_LIST_ENTRY Connections;
 
+    //
+    // List of all Listeners for this registration.
+    //
+    CXPLAT_LIST_ENTRY Listeners;
+
     //
     // Rundown for all child objects.
     //