NMR registration helpers (#4035)

Publish helper functions to simplify MsQuic dynamic loading integration.

src/inc/msquic.h

@@ -1679,6 +1679,150 @@ typedef struct MSQUIC_NMR_DISPATCH {
     MsQuicCloseFn MsQuicClose;
 } MSQUIC_NMR_DISPATCH;
 
+//
+// Stores the internal NMR client state. It's meant to be opaque to the users.
+//
+typedef struct __MSQUIC_NMR_CLIENT {
+    NPI_CLIENT_CHARACTERISTICS NpiClientCharacteristics;
+    HANDLE NmrClientHandle;
+    NPI_MODULEID ModuleId;
+    KEVENT RegistrationCompleteEvent;
+    MSQUIC_NMR_DISPATCH* ProviderDispatch;
+    BOOLEAN Deleting;
+} __MSQUIC_NMR_CLIENT;
+
+#define QUIC_GET_DISPATCH(h) (((__MSQUIC_NMR_CLIENT*)(h))->ProviderDispatch)
+
+static
+NTSTATUS
+__MsQuicClientAttachProvider(
+    _In_ HANDLE NmrBindingHandle,
+    _In_ void *ClientContext,
+    _In_ const NPI_REGISTRATION_INSTANCE *ProviderRegistrationInstance
+    )
+{
+    UNREFERENCED_PARAMETER(ProviderRegistrationInstance);
+
+    NTSTATUS Status;
+    __MSQUIC_NMR_CLIENT* Client = (__MSQUIC_NMR_CLIENT*)ClientContext;
+    void* ProviderContext;
+
+    #pragma warning(suppress:6387) // _Param_(2) could be '0' - by design.
+    Status =
+        NmrClientAttachProvider(
+            NmrBindingHandle,
+            Client,
+            NULL,
+            &ProviderContext,
+            (const void**)&Client->ProviderDispatch);
+    KeSetEvent(&Client->RegistrationCompleteEvent, IO_NO_INCREMENT, FALSE);
+    return Status;
+}
+
+static
+NTSTATUS
+__MsQuicClientDetachProvider(
+    _In_ void *ClientBindingContext
+    )
+{
+    __MSQUIC_NMR_CLIENT* Client = (__MSQUIC_NMR_CLIENT*)ClientBindingContext;
+    if (InterlockedOr8((char*)&Client->Deleting, 1)) {
+        return STATUS_SUCCESS;
+    } else {
+        return STATUS_PENDING;
+    }
+}
+
+_IRQL_requires_max_(PASSIVE_LEVEL)
+__forceinline
+void
+MsQuicNmrClientDeregister(
+    _Inout_ HANDLE* ClientHandle
+    )
+{
+    __MSQUIC_NMR_CLIENT* Client = (__MSQUIC_NMR_CLIENT*)(*ClientHandle);
+
+    if (InterlockedOr8((char*)&Client->Deleting, 1)) {
+        //
+        // We are already in the middle of detaching the client.
+        // Complete it now.
+        //
+        NmrClientDetachProviderComplete(Client->NmrClientHandle);
+    }
+
+    if (Client->NmrClientHandle) {
+        if (NmrDeregisterClient(Client->NmrClientHandle) == STATUS_PENDING) {
+            //
+            // Wait for the deregistration to complete.
+            //
+            NmrWaitForClientDeregisterComplete(Client->NmrClientHandle);
+        }
+        Client->NmrClientHandle = NULL;
+    }
+
+    ExFreePoolWithTag(Client, 'cNQM');
+    *ClientHandle = NULL;
+}
+
+_IRQL_requires_max_(PASSIVE_LEVEL)
+__forceinline
+NTSTATUS
+MsQuicNmrClientRegister(
+    _Out_ HANDLE* ClientHandle,
+    _In_ GUID* ClientModuleId
+    )
+{
+    NPI_REGISTRATION_INSTANCE *ClientRegistrationInstance;
+    NTSTATUS Status = STATUS_SUCCESS;
+
+    __MSQUIC_NMR_CLIENT* Client =
+        (__MSQUIC_NMR_CLIENT*)ExAllocatePool2(POOL_FLAG_NON_PAGED, sizeof(*Client), 'cNQM');
+    if (Client == NULL) {
+        Status = STATUS_INSUFFICIENT_RESOURCES;
+        goto Exit;
+    }
+
+    KeInitializeEvent(&Client->RegistrationCompleteEvent, SynchronizationEvent, FALSE);
+
+    Client->ModuleId.Length = sizeof(Client->ModuleId);
+    Client->ModuleId.Type = MIT_GUID;
+    Client->ModuleId.Guid = *ClientModuleId;
+
+    Client->NpiClientCharacteristics.Length = sizeof(Client->NpiClientCharacteristics);
+    Client->NpiClientCharacteristics.ClientAttachProvider = __MsQuicClientAttachProvider;
+    Client->NpiClientCharacteristics.ClientDetachProvider = __MsQuicClientDetachProvider;
+
+    ClientRegistrationInstance = &Client->NpiClientCharacteristics.ClientRegistrationInstance;
+    ClientRegistrationInstance->Size = sizeof(*ClientRegistrationInstance);
+    ClientRegistrationInstance->Version = 0;
+    ClientRegistrationInstance->NpiId = &MSQUIC_NPI_ID;
+    ClientRegistrationInstance->ModuleId = &Client->ModuleId;
+
+    Status =
+        NmrRegisterClient(
+            &Client->NpiClientCharacteristics, Client, &Client->NmrClientHandle);
+    if (!NT_SUCCESS(Status)) {
+        goto Exit;
+    }
+
+    Status =
+        KeWaitForSingleObject(
+            &Client->RegistrationCompleteEvent, Executive, KernelMode, FALSE, NULL);
+    if (Status != STATUS_SUCCESS) {
+        Status = STATUS_UNSUCCESSFUL;
+        goto Exit;
+    }
+
+    *ClientHandle = Client;
+
+Exit:
+    if (!NT_SUCCESS(Status) && Client != NULL) {
+        MsQuicNmrClientDeregister((HANDLE*)&Client);
+    }
+
+    return Status;
+}
+
 #endif
 
 //

src/test/bin/winkernel/control.cpp

@@ -64,113 +64,7 @@ PAGEDX EVT_WDF_FILE_CLEANUP QuicTestCtlEvtFileCleanup;
 WDFDEVICE QuicTestCtlDevice = nullptr;
 QUIC_DEVICE_EXTENSION* QuicTestCtlExtension = nullptr;
 QUIC_TEST_CLIENT* QuicTestClient = nullptr;
-
-typedef struct QuicTestNmrClient {
-    NPI_CLIENT_CHARACTERISTICS NpiClientCharacteristics;
-    HANDLE NmrClientHandle;
-    NPI_MODULEID ModuleId;
-    CXPLAT_EVENT RegistrationCompleteEvent;
-    MSQUIC_NMR_DISPATCH* ProviderDispatch;
-    BOOLEAN Deleting;
-} QuicTestNmrClient;
-
-static QuicTestNmrClient NmrClient;
-
-static
-NTSTATUS
-QuicTestClientAttachProvider(
-    _In_ HANDLE NmrBindingHandle,
-    _In_ void *ClientContext,
-    _In_ const NPI_REGISTRATION_INSTANCE *ProviderRegistrationInstance
-    )
-{
-    UNREFERENCED_PARAMETER(ProviderRegistrationInstance);
-
-    NTSTATUS Status;
-    QuicTestNmrClient* Client = (QuicTestNmrClient*)ClientContext;
-    void* ProviderContext;
-
-    #pragma warning(suppress:6387) // _Param_(2) could be '0' - by design.
-    Status =
-        NmrClientAttachProvider(
-            NmrBindingHandle,
-            Client,
-            NULL,
-            &ProviderContext,
-            (const void**)&Client->ProviderDispatch);
-    if (!NT_SUCCESS(Status)) {
-        QuicTraceEvent(
-            LibraryErrorStatus,
-            "[ lib] ERROR, %u, %s.",
-            Status,
-            "NmrClientAttachProvider failed");
-    }
-    CxPlatEventSet(Client->RegistrationCompleteEvent);
-    return Status;
-}
-
-static
-NTSTATUS
-QuicTestClientDetachProvider(
-    _In_ void *ClientBindingContext
-    )
-{
-    QuicTestNmrClient* Client = (QuicTestNmrClient*)ClientBindingContext;
-    if (InterlockedFetchAndSetBoolean(&Client->Deleting)) {
-        return STATUS_SUCCESS;
-    } else {
-        return STATUS_PENDING;
-    }
-}
-
-NTSTATUS
-QuicTestRegisterNmrClient(
-    void
-    )
-{
-    NPI_REGISTRATION_INSTANCE *ClientRegistrationInstance;
-    NTSTATUS Status = STATUS_SUCCESS;
-
-    CxPlatEventInitialize(&NmrClient.RegistrationCompleteEvent, FALSE, FALSE);
-    NmrClient.ModuleId.Length = sizeof(NmrClient.ModuleId);
-    NmrClient.ModuleId.Type = MIT_GUID;
-    NmrClient.ModuleId.Guid = MSQUIC_MODULE_ID;
-
-    NmrClient.NpiClientCharacteristics.Length = sizeof(NmrClient.NpiClientCharacteristics);
-    NmrClient.NpiClientCharacteristics.ClientAttachProvider = QuicTestClientAttachProvider;
-    NmrClient.NpiClientCharacteristics.ClientDetachProvider = QuicTestClientDetachProvider;
-
-    ClientRegistrationInstance = &NmrClient.NpiClientCharacteristics.ClientRegistrationInstance;
-    ClientRegistrationInstance->Size = sizeof(*ClientRegistrationInstance);
-    ClientRegistrationInstance->Version = 0;
-    ClientRegistrationInstance->NpiId = &MSQUIC_NPI_ID;
-    ClientRegistrationInstance->ModuleId = &NmrClient.ModuleId;
-
-    Status =
-        NmrRegisterClient(
-            &NmrClient.NpiClientCharacteristics, &NmrClient, &NmrClient.NmrClientHandle);
-    if (!NT_SUCCESS(Status)) {
-        QuicTraceEvent(
-            LibraryErrorStatus,
-            "[ lib] ERROR, %u, %s.",
-            Status,
-            "NmrRegisterClient failed");
-        goto Exit;
-    }
-
-    if (!CxPlatEventWaitWithTimeout(NmrClient.RegistrationCompleteEvent, 1000)) {
-        Status = STATUS_UNSUCCESSFUL;
-        QuicTraceEvent(
-            LibraryErrorStatus,
-            "[ lib] ERROR, %u, %s.",
-            Status,
-            "client registration timed out");
-        goto Exit;
-    }
-
-Exit:
-    return Status;
-}
+HANDLE NmrClient = nullptr;
 
 _No_competing_thread_
 INITCODE
@@ -188,20 +82,23 @@ QuicTestCtlInitialize(
     WDF_IO_QUEUE_CONFIG QueueConfig;
     WDFQUEUE Queue;
 
-    Status = QuicTestRegisterNmrClient();
+    Status = MsQuicNmrClientRegister(&NmrClient, &MSQUIC_MODULE_ID);
     if (!NT_SUCCESS(Status)) {
         QuicTraceEvent(
             LibraryErrorStatus,
             "[ lib] ERROR, %u, %s.",
             Status,
-            "QuicTestRegisterNmrClient failed");
+            "MsQuicNmrClientRegister failed");
         goto Error;
     }
 
+    CXPLAT_DBG_ASSERT(
+        NmrClient != nullptr && QUIC_GET_DISPATCH(NmrClient) != nullptr);
+
     MsQuic =
         new (std::nothrow) MsQuicApi(
-            NmrClient.ProviderDispatch->MsQuicOpenVersion,
-            NmrClient.ProviderDispatch->MsQuicClose);
+            QUIC_GET_DISPATCH(NmrClient)->MsQuicOpenVersion,
+            QUIC_GET_DISPATCH(NmrClient)->MsQuicClose);
     if (!MsQuic) {
         goto Error;
     }
@@ -342,24 +239,8 @@ QuicTestCtlUninitialize(
 
     delete MsQuic;
 
-    if (InterlockedFetchAndSetBoolean(&NmrClient.Deleting)) {
-        //
-        // We are already in the middle of detaching the client.
-        // Complete it now.
-        //
-        NmrClientDetachProviderComplete(NmrClient.NmrClientHandle);
-    }
-
-    if (NmrClient.NmrClientHandle) {
-        NTSTATUS Status = NmrDeregisterClient(NmrClient.NmrClientHandle);
-        CXPLAT_FRE_ASSERTMSG(Status == STATUS_PENDING, "client deregistration failed");
-        if (Status == STATUS_PENDING) {
-            //
-            // Wait for the deregistration to complete.
-            //
-            NmrWaitForClientDeregisterComplete(NmrClient.NmrClientHandle);
-        }
-        NmrClient.NmrClientHandle = NULL;
+    if (NmrClient != nullptr) {
+        MsQuicNmrClientDeregister(&NmrClient);
     }
 
     QuicTraceLogVerbose(