[PATCH 1/2] Use MmGetSystemRoutineAddress to test for IoOpenDriverRegistryKey

[Owen Smith <owen.smith@xxxxxxxxx>]

Server 2016 does not define the function IoOpenDriverRegistryKey, use
MmGetSystemRoutineAddress to dynamically find the function so that a
single binary can be used on Server 2016 (and Win10-1607) and Server 2025.

Signed-off-by: Owen Smith <owen.smith@xxxxxxxxx>
---
 src/common/registry.c | 47 ++++++++++++++++++++++++++++++++-----------
 1 file changed, 35 insertions(+), 12 deletions(-)

diff --git a/src/common/registry.c b/src/common/registry.c
index 211c177..fc68941 100644
--- a/src/common/registry.c
+++ b/src/common/registry.c
@@ -41,6 +41,10 @@
 static PDRIVER_OBJECT   RegistryDriverObject;
 static UNICODE_STRING   RegistryPath;
 
+typedef NTSTATUS(*IOOPENDRIVERREGISTRYKEY)(PDRIVER_OBJECT, DRIVER_REGKEY_TYPE, 
ACCESS_MASK, ULONG, PHANDLE);
+
+static IOOPENDRIVERREGISTRYKEY __IoOpenDriverRegistryKey;
+
 static FORCEINLINE PVOID
 __RegistryAllocate(
     IN  ULONG   Length
@@ -63,6 +67,8 @@ RegistryInitialize(
     IN  PUNICODE_STRING Path
     )
 {
+    UNICODE_STRING      Unicode;
+    PVOID               Func;
     NTSTATUS            status;
 
     ASSERT3P(RegistryPath.Buffer, ==, NULL);
@@ -74,6 +80,13 @@ RegistryInitialize(
     ASSERT3P(RegistryDriverObject, ==, NULL);
     RegistryDriverObject = DriverObject;
 
+    ASSERT3P(__IoOpenDriverRegistryKey, ==, NULL);
+    RtlInitUnicodeString(&Unicode, L"IoOpenDriverRegistryKey");
+
+    Func = MmGetSystemRoutineAddress(&Unicode);
+    if (Func != NULL)
+        __IoOpenDriverRegistryKey = (IOOPENDRIVERREGISTRYKEY)Func;
+
     return STATUS_SUCCESS;
 
 fail1:
@@ -87,6 +100,8 @@ RegistryTeardown(
     VOID
     )
 {
+    __IoOpenDriverRegistryKey = NULL;
+
     RegistryDriverObject = NULL;
 
     RtlFreeUnicodeString(&RegistryPath);
@@ -100,38 +115,46 @@ RegistryOpenParametersKey(
     OUT PHANDLE         Key
     )
 {
-#ifdef VERIFIER_REG_ISOLATION
-    return IoOpenDriverRegistryKey(RegistryDriverObject,
-                                   DriverRegKeyParameters,
-                                   DesiredAccess,
-                                   0,
-                                   Key);
-#else
     HANDLE              ServiceKey;
     NTSTATUS            status;
 
+    if (__IoOpenDriverRegistryKey != NULL) {
+        status = __IoOpenDriverRegistryKey(RegistryDriverObject,
+                                           DriverRegKeyParameters,
+                                           DesiredAccess,
+                                           0,
+                                           Key);
+        if (!NT_SUCCESS(status))
+            goto fail1;
+
+        goto done;
+    }
+
     status = RegistryOpenKey(NULL, &RegistryPath, DesiredAccess, &ServiceKey);
     if (!NT_SUCCESS(status))
-        goto fail1;
+        goto fail2;
 
     status = RegistryOpenSubKey(ServiceKey, "Parameters", DesiredAccess, Key);
     if (!NT_SUCCESS(status))
-        goto fail2;
+        goto fail3;
 
     RegistryCloseKey(ServiceKey);
 
+done:
     return STATUS_SUCCESS;
 
-fail2:
-    Error("fail2\n");
+fail3:
+    Error("fail3\n");
 
     RegistryCloseKey(ServiceKey);
 
+fail2:
+    Error("fail2\n");
+
 fail1:
     Error("fail1 %08x\n", status);
 
     return status;
-#endif
 }
 
 NTSTATUS
-- 
2.44.0.windows.1