Basic NMR integration for MsQuic

scripts/clog.inputs

@@ -75,6 +75,7 @@
 ../src/bin/winuser/dllmain.c
 ../src/bin/winkernel/driver.c
 ../src/bin/winkernel/msquicpcw.c
+../src/bin/winkernel/nmrprovider.c
 ../src/bin/static/empty.c
 ../src/core/operation.h
 ../src/core/stream.h

src/bin/winkernel/driver.c

@@ -40,6 +40,20 @@ MsQuicPcwCleanup(
     void
     );
 
+_No_competing_thread_
+_IRQL_requires_max_(PASSIVE_LEVEL)
+NTSTATUS
+MsQuicRegisterNmrProvider(
+    void
+    );
+
+_No_competing_thread_
+_IRQL_requires_max_(PASSIVE_LEVEL)
+void
+MsQuicDeregisterNmrProvider(
+    void
+    );
+
 INITCODE DRIVER_INITIALIZE DriverEntry;
 PAGEDX EVT_WDF_DRIVER_UNLOAD EvtDriverUnload;
 
@@ -118,6 +132,11 @@ Return Value:
         goto Error;
     }
 
+    Status = MsQuicRegisterNmrProvider();
+    if (!NT_SUCCESS(Status)) {
+        goto Error;
+    }
+
 Error:
 
     if (!NT_SUCCESS(Status)) {
@@ -151,6 +170,8 @@ Routine Description:
     UNREFERENCED_PARAMETER(Driver);
 
     PAGED_CODE();
+
+    MsQuicDeregisterNmrProvider();
     MsQuicPcwCleanup();
     MsQuicLibraryUnload();
 }

src/bin/winkernel/msquic.kernel.vcxproj

@@ -25,6 +25,7 @@
   <ItemGroup>
     <ClCompile Include="driver.c" />
     <ClCompile Include="msquicpcw.c" />
+    <ClCompile Include="nmrprovider.c" />
   </ItemGroup>
   <ItemGroup>
     <ClInclude Include="resource.h" />

src/bin/winkernel/msquicpriv.kernel.vcxproj

@@ -25,6 +25,7 @@
   <ItemGroup>
     <ClCompile Include="driver.c" />
     <ClCompile Include="msquicpcw.c" />
+    <ClCompile Include="nmrprovider.c" />
   </ItemGroup>
   <ItemGroup>
     <ClInclude Include="resource.h" />

src/bin/winkernel/nmrprovider.c

@@ -0,0 +1,109 @@
+#include "quic_platform.h"
+#include <wdm.h>
+#include "msquic.h"
+#include "quic_trace.h"
+#ifdef QUIC_CLOG
+#include "nmrprovider.c.clog.h"
+#endif
+
+typedef struct MSQUIC_NMR_PROVIDER {
+    NPI_PROVIDER_CHARACTERISTICS NpiProviderCharacteristics;
+    HANDLE NmrProviderHandle;
+    NPI_MODULEID ModuleId;
+} MSQUIC_NMR_PROVIDER;
+
+MSQUIC_NMR_PROVIDER NmrProvider;
+
+MSQUIC_NMR_DISPATCH MsQuicNmrDispatch = {
+    .Version = 0,
+    .Reserved = 0,
+    .MsQuicOpenVersion = MsQuicOpenVersion,
+    .MsQuicClose = MsQuicClose,
+};
+
+NTSTATUS
+MsQuicNmrProviderAttachClient(
+    _In_ HANDLE NmrBindingHandle,
+    _In_ VOID *ProviderContext,
+    _In_ CONST NPI_REGISTRATION_INSTANCE *ClientRegistrationInstance,
+    _In_ VOID *ClientBindingContext,
+    _In_ CONST VOID *ClientDispatch,
+    _Out_ VOID **ProviderBindingContext,
+    _Out_ CONST VOID **ProviderDispatch
+    )
+{
+    UNREFERENCED_PARAMETER(NmrBindingHandle);
+    UNREFERENCED_PARAMETER(ProviderContext);
+    UNREFERENCED_PARAMETER(ClientRegistrationInstance);
+    UNREFERENCED_PARAMETER(ClientBindingContext);
+    UNREFERENCED_PARAMETER(ClientDispatch);
+
+    *ProviderBindingContext = NULL;
+    *ProviderDispatch = &MsQuicNmrDispatch;
+    return STATUS_SUCCESS;
+}
+
+NTSTATUS
+MsQuicNmrProviderDetachClient(
+    _In_ VOID *ProviderBindingContext
+    )
+{
+    UNREFERENCED_PARAMETER(ProviderBindingContext);
+    return STATUS_SUCCESS;
+}
+
+_No_competing_thread_
+_IRQL_requires_max_(PASSIVE_LEVEL)
+NTSTATUS
+MsQuicRegisterNmrProvider(
+    void
+    )
+{
+    NTSTATUS Status = STATUS_SUCCESS;
+    NPI_REGISTRATION_INSTANCE *ProviderRegistrationInstance;
+
+    NmrProvider.ModuleId.Length = sizeof(NmrProvider.ModuleId);
+    NmrProvider.ModuleId.Type = MIT_GUID;
+    NmrProvider.ModuleId.Guid = MSQUIC_MODULE_ID;
+
+    NmrProvider.NpiProviderCharacteristics.Version = 0;
+    NmrProvider.NpiProviderCharacteristics.Length = sizeof(NmrProvider.NpiProviderCharacteristics);
+    NmrProvider.NpiProviderCharacteristics.ProviderAttachClient = MsQuicNmrProviderAttachClient;
+    NmrProvider.NpiProviderCharacteristics.ProviderDetachClient = MsQuicNmrProviderDetachClient;
+
+    ProviderRegistrationInstance = &NmrProvider.NpiProviderCharacteristics.ProviderRegistrationInstance;
+    ProviderRegistrationInstance->Version = 0;
+    ProviderRegistrationInstance->Size = sizeof(*ProviderRegistrationInstance);
+    ProviderRegistrationInstance->NpiId = &MSQUIC_NPI_ID;
+    ProviderRegistrationInstance->ModuleId = &NmrProvider.ModuleId;
+
+    Status =
+        NmrRegisterProvider(
+            &NmrProvider.NpiProviderCharacteristics,
+            &NmrProvider,
+            &NmrProvider.NmrProviderHandle);
+    if (!NT_SUCCESS(Status)) {
+        QuicTraceEvent(
+            LibraryErrorStatus,
+            "[ lib] ERROR, %u, %s.",
+            Status,
+            "NmrRegisterProvider");
+        goto Exit;
+    }
+
+Exit:
+    return Status;
+}
+
+void
+MsQuicDeregisterNmrProvider(
+    void
+    )
+{
+    if (NmrProvider.NmrProviderHandle != NULL) {
+        NTSTATUS Status = NmrDeregisterProvider(NmrProvider.NmrProviderHandle);
+        CXPLAT_FRE_ASSERTMSG(Status == STATUS_PENDING, "deregistration failed");
+        NmrWaitForProviderDeregisterComplete(NmrProvider.NmrProviderHandle);
+        NmrProvider.NmrProviderHandle = NULL;
+    }
+}

src/generated/linux/nmrprovider.c.clog.h

@@ -0,0 +1,49 @@
+#ifndef CLOG_DO_NOT_INCLUDE_HEADER
+#include <clog.h>
+#endif
+#undef TRACEPOINT_PROVIDER
+#define TRACEPOINT_PROVIDER CLOG_NMRPROVIDER_C
+#undef TRACEPOINT_PROBE_DYNAMIC_LINKAGE
+#define  TRACEPOINT_PROBE_DYNAMIC_LINKAGE
+#undef TRACEPOINT_INCLUDE
+#define TRACEPOINT_INCLUDE "nmrprovider.c.clog.h.lttng.h"
+#if !defined(DEF_CLOG_NMRPROVIDER_C) || defined(TRACEPOINT_HEADER_MULTI_READ)
+#define DEF_CLOG_NMRPROVIDER_C
+#include <lttng/tracepoint.h>
+#define __int64 __int64_t
+#include "nmrprovider.c.clog.h.lttng.h"
+#endif
+#include <lttng/tracepoint-event.h>
+#ifndef _clog_MACRO_QuicTraceEvent
+#define _clog_MACRO_QuicTraceEvent  1
+#define QuicTraceEvent(a, ...) _clog_CAT(_clog_ARGN_SELECTOR(__VA_ARGS__), _clog_CAT(_,a(#a, __VA_ARGS__)))
+#endif
+#ifdef __cplusplus
+extern "C" {
+#endif
+/*----------------------------------------------------------
+// Decoder Ring for LibraryErrorStatus
+// [ lib] ERROR, %u, %s.
+// QuicTraceEvent(
+            LibraryErrorStatus,
+            "[ lib] ERROR, %u, %s.",
+            Status,
+            "NmrRegisterProvider");
+// arg2 = arg2 = Status = arg2
+// arg3 = arg3 = "NmrRegisterProvider" = arg3
+----------------------------------------------------------*/
+#ifndef _clog_4_ARGS_TRACE_LibraryErrorStatus
+#define _clog_4_ARGS_TRACE_LibraryErrorStatus(uniqueId, encoded_arg_string, arg2, arg3)\
+tracepoint(CLOG_NMRPROVIDER_C, LibraryErrorStatus , arg2, arg3);\
+
+#endif
+
+
+
+
+#ifdef __cplusplus
+}
+#endif
+#ifdef CLOG_INLINE_IMPLEMENTATION
+#include "quic.clog_nmrprovider.c.clog.h.c"
+#endif

src/generated/linux/nmrprovider.c.clog.h.lttng.h

@@ -0,0 +1,23 @@
+
+
+
+/*----------------------------------------------------------
+// Decoder Ring for LibraryErrorStatus
+// [ lib] ERROR, %u, %s.
+// QuicTraceEvent(
+            LibraryErrorStatus,
+            "[ lib] ERROR, %u, %s.",
+            Status,
+            "NmrRegisterProvider");
+// arg2 = arg2 = Status = arg2
+// arg3 = arg3 = "NmrRegisterProvider" = arg3
+----------------------------------------------------------*/
+TRACEPOINT_EVENT(CLOG_NMRPROVIDER_C, LibraryErrorStatus,
+    TP_ARGS(
+        unsigned int, arg2,
+        const char *, arg3), 
+    TP_FIELDS(
+        ctf_integer(unsigned int, arg2, arg2)
+        ctf_string(arg3, arg3)
+    )
+)

src/generated/linux/quic.clog_nmrprovider.c.clog.h.c

@@ -0,0 +1,7 @@
+#include <clog.h>
+#ifdef BUILDING_TRACEPOINT_PROVIDER
+#define TRACEPOINT_CREATE_PROBES
+#else
+#define TRACEPOINT_DEFINE
+#endif
+#include "nmrprovider.c.clog.h"

src/inc/msquic.h

@@ -1646,6 +1646,41 @@ MsQuicClose(
 
 #endif
 
+_IRQL_requires_max_(PASSIVE_LEVEL)
+_Check_return_
+typedef
+QUIC_STATUS
+(QUIC_API *MsQuicOpenVersionFn)(
+    _In_ uint32_t Version,
+    _Out_ _Pre_defensive_ const void** QuicApi
+    );
+
+_IRQL_requires_max_(PASSIVE_LEVEL)
+typedef
+void
+(QUIC_API *MsQuicCloseFn)(
+    _In_ _Pre_defensive_ const void* QuicApi
+    );
+
+#ifdef _KERNEL_MODE
+
+static const GUID MSQUIC_NPI_ID = {
+    0xC43138E3, 0xCD13, 0x4CB1, { 0x9C, 0xAE, 0xE0, 0x05, 0xC8, 0x55, 0x7A, 0xBA }
+}; // C43138E3-CD13-4CB1-9CAE-E005C8557ABA
+
+static const GUID MSQUIC_MODULE_ID = {
+    0x698F7C72, 0xC2E6, 0x49CD, { 0x8C, 0x39, 0x98, 0x85, 0x1D, 0x50, 0x19, 0x01 }
+}; // 698F7C72-C2E6-49CD-8C39-98851D501901
+
+typedef struct MSQUIC_NMR_DISPATCH {
+    USHORT Version;
+    USHORT Reserved;
+    MsQuicOpenVersionFn MsQuicOpenVersion;
+    MsQuicCloseFn MsQuicClose;
+} MSQUIC_NMR_DISPATCH;
+
+#endif
+
 //
 // Version specific helpers that wrap MsQuicOpenVersion.
 //

src/inc/msquic.hpp

@@ -296,18 +296,25 @@ class UniquePtrArray {
 };
 
 class MsQuicApi : public QUIC_API_TABLE {
-    const QUIC_API_TABLE* ApiTable {nullptr};
-    QUIC_STATUS InitStatus;
+    const void* ApiTable {nullptr};
+    QUIC_STATUS InitStatus {QUIC_STATUS_INVALID_STATE};
+    MsQuicOpenVersionFn OpenFn {nullptr};
+    MsQuicCloseFn CloseFn {nullptr};
 public:
-    MsQuicApi() noexcept {
-        if (QUIC_SUCCEEDED(InitStatus = MsQuicOpen2(&ApiTable))) {
+    MsQuicApi(
+        MsQuicOpenVersionFn _OpenFn = MsQuicOpenVersion,
+        MsQuicCloseFn _CloseFn = MsQuicClose) noexcept {
+
+        CloseFn = _CloseFn;
+
+        if (QUIC_SUCCEEDED(InitStatus = _OpenFn(QUIC_API_VERSION_2, &ApiTable))) {
             QUIC_API_TABLE* thisTable = this;
-            memcpy(thisTable, ApiTable, sizeof(*ApiTable));
+            memcpy(thisTable, ApiTable, sizeof(QUIC_API_TABLE));
         }
     }
     ~MsQuicApi() noexcept {
         if (QUIC_SUCCEEDED(InitStatus)) {
-            MsQuicClose(ApiTable);
+            CloseFn(ApiTable);
             ApiTable = nullptr;
             QUIC_API_TABLE* thisTable = this;
             memset(thisTable, 0, sizeof(*thisTable));

src/inc/msquic_winkernel.h

@@ -24,6 +24,7 @@
 #include <minwindef.h>
 #include <ntstatus.h>
 #include <basetsd.h>
+#include <netioddk.h>
 
 typedef INT8 int8_t;
 typedef INT16 int16_t;