Add Zw registry overrides

README.md

@@ -21,6 +21,9 @@ To use this repository from another project:
 7. Disable warning 4324 (structure was padded due to alignment) in the project properties, by adding 4324 to
    Configuration Properties -> C/C++ -> Advanced -> Disable Specific Warnings.
 8. Make sure you are not linking with any kernel libs, since they aren't usable by DLLs.
+9. Give usersim.lib preference over ntdll.lib as follows:
+   Under Configuration Properties -> Linker -> Input -> Ignore Specific Default Libraries, add ntdll.lib.
+   Under Configuration Properties -> Linker -> Input -> Additional Dependencies, add usersim.lib before ntdll.lib.
 
 ### Leak Detection
 

inc/usersim/io.h

@@ -23,7 +23,7 @@ extern "C"
         CustomPriorityWorkQueue = 32
     } WORK_QUEUE_TYPE;
 
-    typedef struct _DEVICE_OBJECT DEVICE_OBJECT;
+    typedef struct _DEVICE_OBJECT DEVICE_OBJECT, *PDEVICE_OBJECT;
 
     typedef struct _DRIVER_OBJECT DRIVER_OBJECT, *PDRIVER_OBJECT;
 

inc/usersim/rtl.h

@@ -153,7 +153,15 @@ extern "C"
         _Out_ PUNICODE_STRING destination_string,
         _In_opt_z_ __drv_aliasesMem PCWSTR source_string);
 
-    typedef struct _object_attributes OBJECT_ATTRIBUTES, *POBJECT_ATTRIBUTES;
+    typedef struct _OBJECT_ATTRIBUTES
+    {
+        ULONG Length;
+        HANDLE RootDirectory;
+        PUNICODE_STRING ObjectName;
+        ULONG Attributes;
+        SECURITY_DESCRIPTOR* SecurityDescriptor;
+        SECURITY_QUALITY_OF_SERVICE* SecurityQualityOfService;
+    } OBJECT_ATTRIBUTES, *POBJECT_ATTRIBUTES;
 
 // Include Rtl* implementations from ntdll.lib.
 #pragma comment(lib, "ntdll.lib")

inc/usersim/wdf.h

@@ -47,6 +47,7 @@ extern "C"
     struct _DRIVER_OBJECT
     {
         WDF_DRIVER_CONFIG config;
+        PDEVICE_OBJECT device;
     };
 
 #define WDF_DRIVER_GLOBALS_NAME_LEN (32)
@@ -130,7 +131,6 @@ extern "C"
         _In_ PWDF_FILEOBJECT_CONFIG file_object_config,
         _In_opt_ PWDF_OBJECT_ATTRIBUTES file_object_attributes);
 
-
     typedef NTSTATUS(FN_WDFDEVICE_WDM_IRP_PREPROCESS)(_In_ WDFDEVICE device, _Inout_ IRP* irp);
     typedef FN_WDFDEVICE_WDM_IRP_PREPROCESS* PFN_WDFDEVICE_WDM_IRP_PREPROCESS;
 
@@ -163,6 +163,12 @@ extern "C"
         _In_opt_ PWDF_OBJECT_ATTRIBUTES queue_attributes,
         _Out_opt_ WDFQUEUE* queue);
 
+    typedef _IRQL_requires_max_(DISPATCH_LEVEL)
+        VOID(WdfControlFinishInitializing_t)(_In_ PWDF_DRIVER_GLOBALS driver_globals, _In_ WDFDEVICE device);
+
+    typedef _IRQL_requires_max_(DISPATCH_LEVEL)
+        PDEVICE_OBJECT(WdfDeviceWdmGetDeviceObject_t)(_In_ PWDF_DRIVER_GLOBALS driver_globals, _In_ WDFDEVICE device);
+
     typedef HANDLE WDFOBJECT;
 
     typedef _IRQL_requires_max_(DISPATCH_LEVEL)
@@ -171,6 +177,8 @@ extern "C"
     typedef enum _WDFFUNCENUM
     {
         WdfControlDeviceInitAllocateTableIndex = 25,
+        WdfControlFinishInitializingTableIndex = 27,
+        WdfDeviceWdmGetDeviceObjectTableIndex = 31,
         WdfDeviceInitFreeTableIndex = 54,
         WdfDeviceInitSetDeviceTypeTableIndex = 66,
         WdfDeviceInitAssignNameTableIndex = 67,

inc/usersim/zw.h

@@ -0,0 +1,31 @@
+// Copyright (c) Microsoft Corporation
+// SPDX-License-Identifier: MIT
+#pragma once
+#include "usersim/rtl.h" // For UNICODE_STRING
+
+#if defined(__cplusplus)
+extern "C"
+{
+#endif
+    USERSIM_API _IRQL_requires_max_(PASSIVE_LEVEL) NTSTATUS
+    ZwCreateKey(
+        _Out_ PHANDLE key_handle,
+        _In_ ACCESS_MASK desired_access,
+        _In_ POBJECT_ATTRIBUTES object_attributes,
+        _Reserved_ ULONG title_index,
+        _In_opt_ PUNICODE_STRING class_string,
+        _In_ ULONG create_options,
+        _Out_opt_ PULONG disposition);
+
+    // The following APIs are exported by ntdll.dll but the prototypes
+    // are not in system headers, so define them here.
+
+    NTSTATUS
+    ZwDeleteKey(_In_ HANDLE key_handle);
+
+    NTSTATUS
+    ZwClose(_In_ HANDLE handle);
+
+#if defined(__cplusplus)
+}
+#endif

src/CMakeLists.txt

@@ -34,6 +34,7 @@ add_library(usersim SHARED
   tracelog.h
   utilities.h
   wdf.cpp
+  zw.cpp
   Source.def
 )
 

src/usersim.vcxproj

@@ -212,6 +212,7 @@
     <ClCompile Include="se.cpp" />
     <ClCompile Include="tracelog.c" />
     <ClCompile Include="wdf.cpp" />
+    <ClCompile Include="zw.cpp" />
   </ItemGroup>
   <ItemGroup>
     <ClInclude Include="..\inc\TraceLoggingProvider.h" />
@@ -221,6 +222,7 @@
     <ClInclude Include="..\inc\usersim\fwp_test.h" />
     <ClInclude Include="..\inc\usersim\ob.h" />
     <ClInclude Include="..\inc\usersim\wdf.h" />
+    <ClInclude Include="..\inc\usersim\zw.h" />
     <ClInclude Include="fault_injection.h" />
     <ClInclude Include="framework.h" />
     <ClInclude Include="fwp_um.h" />

src/usersim.vcxproj.filters

@@ -75,6 +75,9 @@
     <ClCompile Include="ob.cpp">
       <Filter>Source Files</Filter>
     </ClCompile>
+    <ClCompile Include="zw.cpp">
+      <Filter>Source Files</Filter>
+    </ClCompile>
   </ItemGroup>
   <ItemGroup>
     <ClInclude Include="fwp_um.h">
@@ -146,6 +149,9 @@
     <ClInclude Include="..\inc\usersim\ob.h">
       <Filter>Header Files</Filter>
     </ClInclude>
+    <ClInclude Include="..\inc\usersim\zw.h">
+      <Filter>Header Files</Filter>
+    </ClInclude>
   </ItemGroup>
   <ItemGroup>
     <None Include="Source.def">

src/wdf.cpp

@@ -12,6 +12,7 @@ WDF_DRIVER_GLOBALS g_UsersimWdfDriverGlobals = {0};
 static WdfDriverCreate_t _WdfDriverCreate;
 static WdfDeviceCreate_t _WdfDeviceCreate;
 static WdfControlDeviceInitAllocate_t _WdfControlDeviceInitAllocate;
+static WdfControlFinishInitializing_t _WdfControlFinishInitializing;
 static WdfDeviceInitFree_t _WdfDeviceInitFree;
 static WdfDeviceInitSetDeviceType_t _WdfDeviceInitSetDeviceType;
 static WdfDeviceInitSetCharacteristics_t _WdfDeviceInitSetCharacteristics;
@@ -20,6 +21,7 @@ static WdfDeviceInitSetFileObjectConfig_t _WdfDeviceInitSetFileObjectConfig;
 static WdfDeviceInitAssignWdmIrpPreprocessCallback_t _WdfDeviceInitAssignWdmIrpPreprocessCallback;
 static WdfDeviceCreateSymbolicLink_t _WdfDeviceCreateSymbolicLink;
 static WdfIoQueueCreate_t _WdfIoQueueCreate;
+static WdfDeviceWdmGetDeviceObject_t _WdfDeviceWdmGetDeviceObject;
 static WdfObjectDelete_t _WdfObjectDelete;
 
 static NTSTATUS
@@ -42,6 +44,7 @@ _WdfDriverCreate(
     }
     g_UsersimWdfDriverGlobals.Driver = driver_object;
     driver_object->config = *driver_config;
+    driver_object->device = nullptr;
     if (driver != nullptr) {
         *driver = driver_object;
     }
@@ -64,6 +67,8 @@ _WdfDeviceCreate(
         return STATUS_INSUFFICIENT_RESOURCES;
     }
 
+    DRIVER_OBJECT* driver_object = (DRIVER_OBJECT*)driver_globals->Driver;
+    driver_object->device = (PDEVICE_OBJECT)*device_init;
     *device = *device_init;
     return STATUS_SUCCESS;
 }
@@ -243,6 +248,20 @@ _WdfIoQueueCreate(
     return STATUS_SUCCESS;
 }
 
+_IRQL_requires_max_(DISPATCH_LEVEL) VOID
+_WdfControlFinishInitializing(_In_ PWDF_DRIVER_GLOBALS driver_globals, _In_ WDFDEVICE device)
+{
+    UNREFERENCED_PARAMETER(driver_globals);
+    UNREFERENCED_PARAMETER(device);
+}
+
+_IRQL_requires_max_(DISPATCH_LEVEL) PDEVICE_OBJECT
+_WdfDeviceWdmGetDeviceObject(_In_ PWDF_DRIVER_GLOBALS driver_globals, _In_ WDFDEVICE device)
+{
+    UNREFERENCED_PARAMETER(driver_globals);
+    return (PDEVICE_OBJECT)device;
+}
+
 _IRQL_requires_max_(DISPATCH_LEVEL) VOID
 _WdfObjectDelete(_In_ PWDF_DRIVER_GLOBALS driver_globals, _In_ WDFOBJECT object)
 {
@@ -257,6 +276,8 @@ void
 usersim_initialize_wdf()
 {
     g_UsersimWdfFunctions[WdfControlDeviceInitAllocateTableIndex] = (WDFFUNC)_WdfControlDeviceInitAllocate;
+    g_UsersimWdfFunctions[WdfControlFinishInitializingTableIndex] = (WDFFUNC)_WdfControlFinishInitializing;
+    g_UsersimWdfFunctions[WdfDeviceWdmGetDeviceObjectTableIndex] = (WDFFUNC)_WdfDeviceWdmGetDeviceObject;
     g_UsersimWdfFunctions[WdfDeviceInitFreeTableIndex] = (WDFFUNC)_WdfDeviceInitFree;
     g_UsersimWdfFunctions[WdfDeviceInitSetDeviceTypeTableIndex] = (WDFFUNC)_WdfDeviceInitSetDeviceType;
     g_UsersimWdfFunctions[WdfDeviceInitAssignNameTableIndex] = (WDFFUNC)_WdfDeviceInitAssignName;

src/zw.cpp

@@ -0,0 +1,82 @@
+// Copyright (c) Microsoft Corporation
+// SPDX-License-Identifier: MIT
+//
+// This file contains Usersim overrides of some Zw* APIs
+// exposed by ntdll.dll, where the behavior of such APIs
+// differs between kernel mode and unprivileged user
+// mode, so that we can emulate kernel mode behavior to
+// unprivileged user mode test processes.
+#include "fault_injection.h"
+#include "usersim/zw.h"
+#include "utilities.h"
+#include <string>
+
+_IRQL_requires_max_(PASSIVE_LEVEL) NTSTATUS ZwCreateKey(
+    _Out_ PHANDLE key_handle,
+    _In_ ACCESS_MASK desired_access,
+    _In_ POBJECT_ATTRIBUTES object_attributes,
+    _Reserved_ ULONG title_index,
+    _In_opt_ PUNICODE_STRING class_string,
+    _In_ ULONG create_options,
+    _Out_opt_ PULONG disposition)
+{
+    if (usersim_fault_injection_inject_fault()) {
+        return STATUS_INSUFFICIENT_RESOURCES;
+    }
+    if (class_string != nullptr || title_index != 0) {
+        return STATUS_NOT_SUPPORTED;
+    }
+
+    UNREFERENCED_PARAMETER(create_options);
+
+    HKEY root_key = (HKEY)object_attributes->RootDirectory;
+    std::wstring relative_path(object_attributes->ObjectName->Buffer, object_attributes->ObjectName->Length);
+    PCWSTR hklm_path = L"\\Registry\\Machine\\";
+    PCWSTR hkcu_path = L"\\Registry\\User\\";
+    if (relative_path.starts_with(hklm_path)) {
+        root_key = HKEY_LOCAL_MACHINE;
+        relative_path = relative_path.substr(wcslen(hklm_path));
+    } else if (relative_path.starts_with(hkcu_path)) {
+        root_key = HKEY_CURRENT_USER;
+        relative_path = relative_path.substr(wcslen(hkcu_path));
+    }
+
+    std::wstring class_wstring;
+    if (class_string) {
+        class_wstring.assign(class_string->Buffer, class_string->Length);
+    }
+    ULONG result = RegCreateKeyExW(
+        root_key,
+        relative_path.c_str(),
+        0, // Reserved
+        (class_string ? (PWSTR)class_wstring.c_str() : nullptr),
+        REG_OPTION_VOLATILE,
+        desired_access,
+        nullptr,
+        (PHKEY)key_handle,
+        disposition);
+    if (result == ERROR_ACCESS_DENIED && root_key == HKEY_LOCAL_MACHINE) {
+        // Try again with HKCU, so access to \Registry\Machine will succeed
+        // like it should when called from "kernel-mode" code.
+        //
+        // Note that HKLM will fail with access denied, even when the desired
+        // access is only read access, unless the process is running as admin,
+        // since we use RegCreateKey instead of RegOpenKey.  This ensures that
+        // the same registry key will be accessed across multiple calls
+        // regardless of desired access.
+
+        root_key = HKEY_CURRENT_USER;
+        result = RegCreateKeyExW(
+            root_key,
+            relative_path.c_str(),
+            0, // Reserved
+            (class_string ? (PWSTR)class_wstring.c_str() : nullptr),
+            REG_OPTION_VOLATILE,
+            desired_access,
+            nullptr,
+            (PHKEY)key_handle,
+            disposition);
+    }
+
+    return win32_error_code_to_usersim_result(result);
+}

tests/tests.vcxproj

@@ -129,7 +129,8 @@
     <Link>
       <SubSystem>Console</SubSystem>
       <GenerateDebugInformation>true</GenerateDebugInformation>
-      <AdditionalDependencies>ntdll.lib;$(CoreLibraryDependencies);%(AdditionalDependencies)</AdditionalDependencies>
+      <AdditionalDependencies>usersim.lib;ntdll.lib;$(CoreLibraryDependencies);%(AdditionalDependencies)</AdditionalDependencies>
+      <AdditionalLibraryDirectories>$(OutDir);$(VC_LibraryPath_VC_x64_Desktop);%(Link.AdditionalLibraryDirectories)</AdditionalLibraryDirectories>
     </Link>
   </ItemDefinitionGroup>
   <ItemDefinitionGroup Condition="'$(Configuration)|$(Platform)'=='Release|x64'">
@@ -148,7 +149,8 @@
       <EnableCOMDATFolding>true</EnableCOMDATFolding>
       <OptimizeReferences>true</OptimizeReferences>
       <GenerateDebugInformation>true</GenerateDebugInformation>
-      <AdditionalDependencies>ntdll.lib;$(CoreLibraryDependencies);%(AdditionalDependencies)</AdditionalDependencies>
+      <AdditionalDependencies>usersim.lib;ntdll.lib;$(CoreLibraryDependencies);%(AdditionalDependencies)</AdditionalDependencies>
+      <AdditionalLibraryDirectories>$(OutDir);$(VC_LibraryPath_VC_x64_Desktop);%(Link.AdditionalLibraryDirectories)</AdditionalLibraryDirectories>
     </Link>
   </ItemDefinitionGroup>
   <ItemGroup>
@@ -161,6 +163,7 @@
     <ClCompile Include="rtl_test.cpp" />
     <ClCompile Include="se_test.cpp" />
     <ClCompile Include="wdf_test.cpp" />
+    <ClCompile Include="zw_test.cpp" />
   </ItemGroup>
   <ItemGroup>
     <None Include="packages.config" />

tests/tests.vcxproj.filters

@@ -42,6 +42,9 @@
     <ClCompile Include="ob_test.cpp">
       <Filter>Source Files</Filter>
     </ClCompile>
+    <ClCompile Include="zw_test.cpp">
+      <Filter>Source Files</Filter>
+    </ClCompile>
   </ItemGroup>
   <ItemGroup>
     <None Include="packages.config" />

tests/wdf_test.cpp

@@ -66,7 +66,8 @@ TEST_CASE("WdfDeviceCreate", "[wdf]")
     NTSTATUS status = WdfDriverCreate(UsersimWdfDriverGlobals, &driver_object, &registry_path, nullptr, &config, &driver);
     REQUIRE(status == STATUS_SUCCESS);
 
-    // Allocate control device object.
+    // Allocate control device object.  For usage discussion, see
+    // https://learn.microsoft.com/en-us/windows-hardware/drivers/wdf/using-control-device-objects#creating-a-control-device-object
     WdfControlDeviceInitAllocate_t* WdfControlDeviceInitAllocate =
         (WdfControlDeviceInitAllocate_t*)UsersimWdfFunctions[WdfControlDeviceInitAllocateTableIndex];
     DECLARE_CONST_UNICODE_STRING(security_descriptor, L"");
@@ -111,6 +112,11 @@ TEST_CASE("WdfDeviceCreate", "[wdf]")
     REQUIRE(status == STATUS_SUCCESS);
     REQUIRE(device != nullptr);
 
+    WdfDeviceWdmGetDeviceObject_t* WdfDeviceWdmGetDeviceObject =
+        (WdfDeviceWdmGetDeviceObject_t*)UsersimWdfFunctions[WdfDeviceWdmGetDeviceObjectTableIndex];
+    PDEVICE_OBJECT device_object = WdfDeviceWdmGetDeviceObject(UsersimWdfDriverGlobals, device);
+    REQUIRE(device_object != nullptr);
+
     // Create symbolic link for control object for user mode.
     UNICODE_STRING symbolic_device_name;
     RtlInitUnicodeString(&symbolic_device_name, L"symbolic device name");
@@ -127,6 +133,11 @@ TEST_CASE("WdfDeviceCreate", "[wdf]")
         UsersimWdfDriverGlobals, device, &io_queue_configuration, WDF_NO_OBJECT_ATTRIBUTES, WDF_NO_HANDLE);
     REQUIRE(status == STATUS_SUCCESS);
 
+    // Finish initializing the control device object.
+    WdfControlFinishInitializing_t* WdfControlFinishInitializing =
+        (WdfControlFinishInitializing_t*)UsersimWdfFunctions[WdfControlFinishInitializingTableIndex];
+    WdfControlFinishInitializing(UsersimWdfDriverGlobals, device);
+
     // Free device, which will free the control device object.
     WdfObjectDelete_t* WdfObjectDelete = (WdfObjectDelete_t*)UsersimWdfFunctions[WdfObjectDeleteTableIndex];
     WdfObjectDelete(UsersimWdfDriverGlobals, device);