NMR helper improvement (microsoft#4045)

Percieve · Jan 9, 2024 · 878886a · 878886a
1 parent 5949785
commit 878886a
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 3 deletions.
diff --git a/src/inc/msquic.h b/src/inc/msquic.h
@@ -1774,7 +1774,8 @@ __forceinline
 NTSTATUS
 MsQuicNmrClientRegister(
     _Out_ HANDLE* ClientHandle,
-    _In_ GUID* ClientModuleId
+    _In_ GUID* ClientModuleId,
+    _In_ ULONG TimeoutMs
     )
 {
     NPI_REGISTRATION_INSTANCE *ClientRegistrationInstance;
@@ -1810,9 +1811,15 @@ MsQuicNmrClientRegister(
         goto Exit;
     }
 
+    LARGE_INTEGER Timeout;
+    Timeout.QuadPart = UInt32x32To64(TimeoutMs, 10000);
+    Timeout.QuadPart = -Timeout.QuadPart;
+
     Status =
         KeWaitForSingleObject(
-            &Client->RegistrationCompleteEvent, Executive, KernelMode, FALSE, NULL);
+            &Client->RegistrationCompleteEvent,
+            Executive, KernelMode, FALSE,
+            TimeoutMs != 0 ? &Timeout : NULL);
     if (Status != STATUS_SUCCESS) {
         Status = STATUS_UNSUCCESSFUL;
         goto Exit;

diff --git a/src/test/bin/winkernel/control.cpp b/src/test/bin/winkernel/control.cpp
@@ -82,7 +82,7 @@ QuicTestCtlInitialize(
     WDF_IO_QUEUE_CONFIG QueueConfig;
     WDFQUEUE Queue;
 
-    Status = MsQuicNmrClientRegister(&NmrClient, &MSQUIC_MODULE_ID);
+    Status = MsQuicNmrClientRegister(&NmrClient, &MSQUIC_MODULE_ID, 5000);
     if (!NT_SUCCESS(Status)) {
         QuicTraceEvent(
             LibraryErrorStatus,