[NTOS:CM] Improve the capture of user-mode parameters.

- Improve the capture of OBJECT_ATTRIBUTES parameters that are passed
  (by pointer) to the Cm* helper functions, and the capture of
  UNICODE_STRINGs.

- Correctly differentiate user-mode vs. kernel-mode root directory handles
  (in OBJECT_ATTRIBUTES): note that most of the Cm* APIs expect their
  parameters to be kernel-mode (pointers, handles...).

CORE-13448

ntoskrnl/config/ntapi.c

@@ -17,6 +17,222 @@ BOOLEAN CmBootAcceptFirstTime = TRUE;
 BOOLEAN CmFirstTime = TRUE;
 extern ULONG InitSafeBootMode;
 
+
+/* PRIVATE FUNCTIONS *********************************************************/
+
+/*
+ * Adapted from ntoskrnl/include/internal/ob_x.h:ObpReleaseObjectCreateInformation()
+ */
+VOID
+ReleaseCapturedObjectAttributes(
+    _In_ POBJECT_ATTRIBUTES CapturedObjectAttributes,
+    _In_ KPROCESSOR_MODE AccessMode)
+{
+    /* Check if we have a security descriptor */
+    if (CapturedObjectAttributes->SecurityDescriptor)
+    {
+        /* Release it */
+        SeReleaseSecurityDescriptor(CapturedObjectAttributes->SecurityDescriptor,
+                                    AccessMode,
+                                    TRUE);
+        CapturedObjectAttributes->SecurityDescriptor = NULL;
+    }
+
+    /* Check if we have an object name */
+    if (CapturedObjectAttributes->ObjectName)
+    {
+        /* Release it */
+        ReleaseCapturedUnicodeString(CapturedObjectAttributes->ObjectName, AccessMode);
+    }
+}
+
+/*
+ * Adapted from ntoskrnl/ob/oblife.c:ObpCaptureObjectCreateInformation()
+ */
+NTSTATUS
+ProbeAndCaptureObjectAttributes(
+    _Out_ POBJECT_ATTRIBUTES CapturedObjectAttributes,
+    _Out_ PUNICODE_STRING ObjectName,
+    _In_ KPROCESSOR_MODE AccessMode,
+    _In_ POBJECT_ATTRIBUTES ObjectAttributes,
+    _In_ BOOLEAN CaptureSecurity)
+{
+    NTSTATUS Status = STATUS_SUCCESS;
+    PSECURITY_DESCRIPTOR SecurityDescriptor;
+    // PSECURITY_QUALITY_OF_SERVICE SecurityQos;
+    PUNICODE_STRING LocalObjectName = NULL;
+
+    /* Zero out the Capture Data */
+    RtlZeroMemory(CapturedObjectAttributes, sizeof(*CapturedObjectAttributes));
+
+    /* SEH everything here for protection */
+    _SEH2_TRY
+    {
+        /* Check if we got attributes */
+        if (ObjectAttributes)
+        {
+            /* Check if we're in user mode */
+            if (AccessMode != KernelMode)
+            {
+                /* Probe the attributes */
+                ProbeForRead(ObjectAttributes,
+                             sizeof(OBJECT_ATTRIBUTES),
+                             sizeof(ULONG));
+            }
+
+            /* Validate the Size and Attributes */
+            if ((ObjectAttributes->Length != sizeof(OBJECT_ATTRIBUTES)) ||
+                (ObjectAttributes->Attributes & ~OBJ_VALID_KERNEL_ATTRIBUTES))  // Understood as all the possible valid attributes
+            {
+                /* Invalid combination, fail */
+                _SEH2_YIELD(return STATUS_INVALID_PARAMETER);
+            }
+
+            /* Set some Create Info and do not allow user-mode kernel handles */
+            CapturedObjectAttributes->Length = sizeof(OBJECT_ATTRIBUTES);
+            CapturedObjectAttributes->RootDirectory = ObjectAttributes->RootDirectory;
+            CapturedObjectAttributes->Attributes = ObpValidateAttributes(ObjectAttributes->Attributes, AccessMode);
+            LocalObjectName = ObjectAttributes->ObjectName;
+            SecurityDescriptor = ObjectAttributes->SecurityDescriptor;
+            // SecurityQos = ObjectAttributes->SecurityQualityOfService;
+
+            /* Check if we have a security descriptor */
+            if (CaptureSecurity && SecurityDescriptor)
+            {
+                /*
+                 * Capture it.
+                 * Note: This has an implicit memory barrier due
+                 * to the function call, so cleanup is safe here.
+                 */
+                Status = SeCaptureSecurityDescriptor(SecurityDescriptor,
+                                                     AccessMode,
+                                                     NonPagedPool,
+                                                     TRUE,
+                                                     &CapturedObjectAttributes->
+                                                        SecurityDescriptor);
+                if (!NT_SUCCESS(Status))
+                {
+                    /* Capture failed, quit */
+                    CapturedObjectAttributes->SecurityDescriptor = NULL;
+                    _SEH2_YIELD(return Status);
+                }
+            }
+            else
+            {
+                CapturedObjectAttributes->SecurityDescriptor = NULL;
+            }
+
+#if 0
+// We don't use the QoS!
+
+            /* Check if we have QoS */
+            if (SecurityQos)
+            {
+                /* Check if we came from user mode */
+                if (AccessMode != KernelMode)
+                {
+                    /* Validate the QoS */
+                    ProbeForRead(SecurityQos,
+                                 sizeof(SECURITY_QUALITY_OF_SERVICE),
+                                 sizeof(ULONG));
+                }
+
+                /* Save Info */
+                CapturedObjectAttributes->SecurityQualityOfService = *SecurityQos;
+                CapturedObjectAttributes->SecurityQos =
+                    &CapturedObjectAttributes->SecurityQualityOfService;
+            }
+#else
+            CapturedObjectAttributes->SecurityQualityOfService = NULL;
+#endif
+        }
+        else
+        {
+            /* We don't have a name */
+            LocalObjectName = NULL;
+        }
+    }
+    _SEH2_EXCEPT(ExSystemExceptionFilter())
+    {
+        /* Cleanup and return the exception code */
+        ReleaseCapturedObjectAttributes(CapturedObjectAttributes, AccessMode);
+        _SEH2_YIELD(return _SEH2_GetExceptionCode());
+    }
+    _SEH2_END;
+
+    /* Now check if the Object Attributes had an Object Name */
+    if (LocalObjectName)
+    {
+        Status = ProbeAndCaptureUnicodeString(ObjectName, AccessMode, LocalObjectName);
+    }
+    else
+    {
+        /* Clear the string */
+        RtlInitEmptyUnicodeString(ObjectName, NULL, 0);
+
+        /* It cannot have specified a Root Directory */
+        if (CapturedObjectAttributes->RootDirectory)
+        {
+            Status = STATUS_OBJECT_NAME_INVALID;
+        }
+    }
+
+    /* Set the caputured object attributes name pointer to the one the user gave to us */
+    CapturedObjectAttributes->ObjectName = ObjectName;
+
+    /* Cleanup if we failed */
+    if (!NT_SUCCESS(Status))
+    {
+        ReleaseCapturedObjectAttributes(CapturedObjectAttributes, AccessMode);
+    }
+
+    /* Return status to caller */
+    return Status;
+}
+
+static
+NTSTATUS
+CmpConvertHandleToKernelHandle(
+    _In_ HANDLE SourceHandle,
+    _In_opt_ POBJECT_TYPE ObjectType,
+    _In_ ACCESS_MASK DesiredAccess,
+    _In_ KPROCESSOR_MODE AccessMode,
+    _Out_ PHANDLE KernelHandle)
+{
+    NTSTATUS Status;
+    PVOID Object;
+
+    *KernelHandle = NULL;
+
+    /* NULL handle is valid */
+    if (SourceHandle == NULL)
+        return STATUS_SUCCESS;
+
+    /* Get the object pointer */
+    Status = ObReferenceObjectByHandle(SourceHandle,
+                                       DesiredAccess,
+                                       ObjectType,
+                                       AccessMode,
+                                       &Object,
+                                       NULL);
+    if (!NT_SUCCESS(Status))
+        return Status;
+
+    /* Create a kernel handle from the pointer */
+    Status = ObOpenObjectByPointer(Object,
+                                   OBJ_KERNEL_HANDLE,
+                                   NULL,
+                                   DesiredAccess,
+                                   ObjectType,
+                                   KernelMode,
+                                   KernelHandle);
+
+    /* Dereference the object */
+    ObDereferenceObject(Object);
+    return Status;
+}
+
+
 /* FUNCTIONS *****************************************************************/
 
 NTSTATUS
@@ -231,7 +447,7 @@ NtDeleteKey(IN HANDLE KeyHandle)
         CmiCallRegisteredCallbacks(RegNtPostDeleteKey, &PostOperationInfo);
     }
 
-    /* Dereference the object */
+    /* Dereference and return status */
     ObDereferenceObject(KeyObject);
     return Status;
 }
@@ -402,6 +618,7 @@ NtEnumerateValueKey(IN HANDLE KeyHandle,
         CmiCallRegisteredCallbacks(RegNtPostEnumerateValueKey, &PostOperationInfo);
     }
 
+    /* Dereference and return status */
     ObDereferenceObject(KeyObject);
     return Status;
 }
@@ -527,12 +744,12 @@ NtQueryValueKey(IN HANDLE KeyHandle,
                 IN ULONG Length,
                 OUT PULONG ResultLength)
 {
-    KPROCESSOR_MODE PreviousMode = ExGetPreviousMode();
     NTSTATUS Status;
+    KPROCESSOR_MODE PreviousMode = ExGetPreviousMode();
     PCM_KEY_BODY KeyObject;
     REG_QUERY_VALUE_KEY_INFORMATION QueryValueKeyInfo;
     REG_POST_OPERATION_INFORMATION PostOperationInfo;
-    UNICODE_STRING ValueNameCopy = *ValueName;
+    UNICODE_STRING ValueNameCopy;
 
     PAGED_CODE();
 
@@ -557,7 +774,8 @@ NtQueryValueKey(IN HANDLE KeyHandle,
                                        PreviousMode,
                                        (PVOID*)&KeyObject,
                                        NULL);
-    if (!NT_SUCCESS(Status)) return Status;
+    if (!NT_SUCCESS(Status))
+        return Status;
 
     if (PreviousMode != KernelMode)
     {
@@ -577,22 +795,25 @@ NtQueryValueKey(IN HANDLE KeyHandle,
         _SEH2_END;
     }
 
+    /* Capture the string */
+    Status = ProbeAndCaptureUnicodeString(&ValueNameCopy, PreviousMode, ValueName);
+    if (!NT_SUCCESS(Status))
+        goto Quit;
+
     /* Make sure the name is aligned properly */
     if ((ValueNameCopy.Length & (sizeof(WCHAR) - 1)))
     {
         /* It isn't, so we'll fail */
-        ObDereferenceObject(KeyObject);
+        Status = STATUS_INVALID_PARAMETER;
-        return STATUS_INVALID_PARAMETER;
+        goto Quit;
     }
-    else
+
+    /* Ignore any null characters at the end */
+    while ((ValueNameCopy.Length) &&
+           !(ValueNameCopy.Buffer[ValueNameCopy.Length / sizeof(WCHAR) - 1]))
     {
-        /* Ignore any null characters at the end */
+        /* Skip it */
-        while ((ValueNameCopy.Length) &&
+        ValueNameCopy.Length -= sizeof(WCHAR);
-               !(ValueNameCopy.Buffer[ValueNameCopy.Length / sizeof(WCHAR) - 1]))
-        {
-            /* Skip it */
-            ValueNameCopy.Length -= sizeof(WCHAR);
-        }
     }
 
     /* Setup the callback */
@@ -620,6 +841,10 @@ NtQueryValueKey(IN HANDLE KeyHandle,
         CmiCallRegisteredCallbacks(RegNtPostQueryValueKey, &PostOperationInfo);
     }
 
+Quit:
+    if (ValueNameCopy.Buffer)
+        ReleaseCapturedUnicodeString(&ValueNameCopy, PreviousMode);
+
     /* Dereference and return status */
     ObDereferenceObject(KeyObject);
     return Status;
@@ -635,16 +860,26 @@ NtSetValueKey(IN HANDLE KeyHandle,
               IN ULONG DataSize)
 {
     NTSTATUS Status = STATUS_SUCCESS;
-    PCM_KEY_BODY KeyObject = NULL;
+    KPROCESSOR_MODE PreviousMode;
+    PCM_KEY_BODY KeyObject;
     REG_SET_VALUE_KEY_INFORMATION SetValueKeyInfo;
     REG_POST_OPERATION_INFORMATION PostOperationInfo;
     UNICODE_STRING ValueNameCopy;
-    KPROCESSOR_MODE PreviousMode;
 
     PAGED_CODE();
 
     PreviousMode = ExGetPreviousMode();
 
+    /* Verify that the handle is valid and is a registry key */
+    Status = ObReferenceObjectByHandle(KeyHandle,
+                                       KEY_SET_VALUE,
+                                       CmpKeyObjectType,
+                                       PreviousMode,
+                                       (PVOID*)&KeyObject,
+                                       NULL);
+    if (!NT_SUCCESS(Status))
+        return Status;
+
     if (!DataSize)
         Data = NULL;
 
@@ -653,7 +888,11 @@ NtSetValueKey(IN HANDLE KeyHandle,
     {
         PVOID DataCopy = ExAllocatePoolWithTag(PagedPool, DataSize, TAG_CM);
         if (!DataCopy)
+        {
+            /* Dereference and return status */
+            ObDereferenceObject(KeyObject);
             return STATUS_INSUFFICIENT_RESOURCES;
+        }
         _SEH2_TRY
         {
             ProbeForRead(Data, DataSize, 1);
@@ -667,7 +906,9 @@ NtSetValueKey(IN HANDLE KeyHandle,
 
         if (!NT_SUCCESS(Status))
         {
+            /* Dereference and return status */
             ExFreePoolWithTag(DataCopy, TAG_CM);
+            ObDereferenceObject(KeyObject);
             return Status;
         }
         Data = DataCopy;
@@ -676,21 +917,11 @@ NtSetValueKey(IN HANDLE KeyHandle,
     /* Capture the string */
     Status = ProbeAndCaptureUnicodeString(&ValueNameCopy, PreviousMode, ValueName);
     if (!NT_SUCCESS(Status))
-        goto end;
+        goto Quit;
 
     DPRINT("NtSetValueKey() KH 0x%p, VN '%wZ', TI %x, T %lu, DS %lu\n",
         KeyHandle, &ValueNameCopy, TitleIndex, Type, DataSize);
 
-    /* Verify that the handle is valid and is a registry key */
-    Status = ObReferenceObjectByHandle(KeyHandle,
-                                       KEY_SET_VALUE,
-                                       CmpKeyObjectType,
-                                       PreviousMode,
-                                       (PVOID*)&KeyObject,
-                                       NULL);
-    if (!NT_SUCCESS(Status))
-        goto end;
-
     /* Make sure the name is aligned, not too long, and the data under 4GB */
     if ( (ValueNameCopy.Length > 32767) ||
          ((ValueNameCopy.Length & (sizeof(WCHAR) - 1))) ||
@@ -698,7 +929,7 @@ NtSetValueKey(IN HANDLE KeyHandle,
     {
         /* Fail */
         Status = STATUS_INVALID_PARAMETER;
-        goto end;
+        goto Quit;
     }
 
     /* Ignore any null characters at the end */
@@ -714,7 +945,7 @@ NtSetValueKey(IN HANDLE KeyHandle,
     {
         /* Fail */
         Status = STATUS_ACCESS_DENIED;
-        goto end;
+        goto Quit;
     }
 
     /* Setup callback */
@@ -742,13 +973,15 @@ NtSetValueKey(IN HANDLE KeyHandle,
         CmiCallRegisteredCallbacks(RegNtPostSetValueKey, &PostOperationInfo);
     }
 
-end:
+Quit:
-    /* Dereference and return status */
+    if (ValueNameCopy.Buffer)
-    if (KeyObject)
+        ReleaseCapturedUnicodeString(&ValueNameCopy, PreviousMode);
-        ObDereferenceObject(KeyObject);
+
-    ReleaseCapturedUnicodeString(&ValueNameCopy, PreviousMode);
     if ((PreviousMode != KernelMode) && Data)
         ExFreePoolWithTag(Data, TAG_CM);
+
+    /* Dereference and return status */
+    ObDereferenceObject(KeyObject);
     return Status;
 }
 
@@ -757,12 +990,13 @@ NTAPI
 NtDeleteValueKey(IN HANDLE KeyHandle,
                  IN PUNICODE_STRING ValueName)
 {
-    PCM_KEY_BODY KeyObject;
     NTSTATUS Status;
+    PCM_KEY_BODY KeyObject;
     REG_DELETE_VALUE_KEY_INFORMATION DeleteValueKeyInfo;
     REG_POST_OPERATION_INFORMATION PostOperationInfo;
     KPROCESSOR_MODE PreviousMode = ExGetPreviousMode();
-    UNICODE_STRING ValueNameCopy = *ValueName;
+    UNICODE_STRING ValueNameCopy;
+
     PAGED_CODE();
 
     /* Verify that the handle is valid and is a registry key */
@@ -772,22 +1006,28 @@ NtDeleteValueKey(IN HANDLE KeyHandle,
                                        PreviousMode,
                                        (PVOID*)&KeyObject,
                                        NULL);
-    if (!NT_SUCCESS(Status)) return Status;
+    if (!NT_SUCCESS(Status))
+        return Status;
 
-    /* Don't touch read-only keys */
+    /* Capture the string */
-    if (KeyObject->KeyControlBlock->ExtFlags & CM_KCB_READ_ONLY_KEY)
+    Status = ProbeAndCaptureUnicodeString(&ValueNameCopy, PreviousMode, ValueName);
-    {
+    if (!NT_SUCCESS(Status))
-        /* Fail */
+        goto Quit;
-        ObDereferenceObject(KeyObject);
-        return STATUS_ACCESS_DENIED;
-    }
 
     /* Make sure the name is aligned properly */
     if ((ValueNameCopy.Length & (sizeof(WCHAR) - 1)))
     {
         /* It isn't, so we'll fail */
-        ObDereferenceObject(KeyObject);
+        Status = STATUS_INVALID_PARAMETER;
-        return STATUS_INVALID_PARAMETER;
+        goto Quit;
+    }
+
+    /* Don't touch read-only keys */
+    if (KeyObject->KeyControlBlock->ExtFlags & CM_KCB_READ_ONLY_KEY)
+    {
+        /* Fail */
+        Status = STATUS_ACCESS_DENIED;
+        goto Quit;
     }
 
     /* Do the callback */
@@ -807,7 +1047,11 @@ NtDeleteValueKey(IN HANDLE KeyHandle,
                                    &PostOperationInfo);
     }
 
-    /* Dereference the key body */
+Quit:
+    if (ValueNameCopy.Buffer)
+        ReleaseCapturedUnicodeString(&ValueNameCopy, PreviousMode);
+
+    /* Dereference and return status */
     ObDereferenceObject(KeyObject);
     return Status;
 }
@@ -884,11 +1128,17 @@ NtLoadKeyEx(IN POBJECT_ATTRIBUTES TargetKey,
 {
     NTSTATUS Status;
     KPROCESSOR_MODE PreviousMode = ExGetPreviousMode();
+    OBJECT_ATTRIBUTES CapturedTargetKey;
+    OBJECT_ATTRIBUTES CapturedSourceFile;
+    UNICODE_STRING TargetKeyName, SourceFileName;
+    HANDLE KmTargetKeyRootDir = NULL, KmSourceFileRootDir = NULL;
     PCM_KEY_BODY KeyBody = NULL;
+
     PAGED_CODE();
 
     /* Validate flags */
-    if (Flags & ~REG_NO_LAZY_FLUSH) return STATUS_INVALID_PARAMETER;
+    if (Flags & ~REG_NO_LAZY_FLUSH)
+        return STATUS_INVALID_PARAMETER;
 
     /* Validate privilege */
     if (!SeSinglePrivilegeCheck(SeRestorePrivilege, PreviousMode))
@@ -900,6 +1150,77 @@ NtLoadKeyEx(IN POBJECT_ATTRIBUTES TargetKey,
     /* Block APCs */
     KeEnterCriticalRegion();
 
+    /* Check for user-mode caller */
+    if (PreviousMode != KernelMode)
+    {
+        /* Prepare to probe parameters */
+        _SEH2_TRY
+        {
+            /* Probe target key */
+            ProbeForRead(TargetKey,
+                         sizeof(OBJECT_ATTRIBUTES),
+                         sizeof(ULONG));
+
+            /* Probe source file */
+            ProbeForRead(SourceFile,
+                         sizeof(OBJECT_ATTRIBUTES),
+                         sizeof(ULONG));
+        }
+        _SEH2_EXCEPT(EXCEPTION_EXECUTE_HANDLER)
+        {
+            /* Return the exception code */
+            Status = _SEH2_GetExceptionCode();
+            _SEH2_YIELD(goto Quit);
+        }
+        _SEH2_END;
+    }
+
+    /* Probe and capture the target key attributes, including the security */
+    Status = ProbeAndCaptureObjectAttributes(&CapturedTargetKey,
+                                             &TargetKeyName,
+                                             PreviousMode,
+                                             TargetKey,
+                                             TRUE);
+    if (!NT_SUCCESS(Status))
+        goto Quit;
+
+    /*
+     * Probe and capture the source file attributes, but not the security.
+     * A proper security context is built by CmLoadKey().
+     */
+    Status = ProbeAndCaptureObjectAttributes(&CapturedSourceFile,
+                                             &SourceFileName,
+                                             PreviousMode,
+                                             SourceFile,
+                                             FALSE);
+    if (!NT_SUCCESS(Status))
+    {
+        ReleaseCapturedObjectAttributes(&CapturedTargetKey, PreviousMode);
+        goto Quit;
+    }
+
+    /* Make sure the target key root directory handle is a kernel handle */
+    Status = CmpConvertHandleToKernelHandle(CapturedTargetKey.RootDirectory,
+                                            CmpKeyObjectType,
+                                            KEY_READ,
+                                            PreviousMode,
+                                            &KmTargetKeyRootDir);
+    if (!NT_SUCCESS(Status))
+        goto Cleanup;
+    CapturedTargetKey.RootDirectory = KmTargetKeyRootDir;
+    CapturedTargetKey.Attributes |= OBJ_KERNEL_HANDLE;
+
+    /* Make sure the source file root directory handle is a kernel handle */
+    Status = CmpConvertHandleToKernelHandle(CapturedSourceFile.RootDirectory,
+                                            IoFileObjectType,
+                                            FILE_TRAVERSE,
+                                            PreviousMode,
+                                            &KmSourceFileRootDir);
+    if (!NT_SUCCESS(Status))
+        goto Cleanup;
+    CapturedSourceFile.RootDirectory = KmSourceFileRootDir;
+    CapturedSourceFile.Attributes |= OBJ_KERNEL_HANDLE;
+
     /* Check if we have a trust class */
     if (TrustClassKey)
     {
@@ -913,11 +1234,26 @@ NtLoadKeyEx(IN POBJECT_ATTRIBUTES TargetKey,
     }
 
     /* Call the internal API */
-    Status = CmLoadKey(TargetKey, SourceFile, Flags, KeyBody);
+    Status = CmLoadKey(&CapturedTargetKey,
+                       &CapturedSourceFile,
+                       Flags,
+                       KeyBody);
 
     /* Dereference the trust key, if any */
     if (KeyBody) ObDereferenceObject(KeyBody);
 
+Cleanup:
+    /* Close the local kernel handles */
+    if (KmSourceFileRootDir)
+        ObCloseHandle(KmSourceFileRootDir, KernelMode);
+    if (KmTargetKeyRootDir)
+        ObCloseHandle(KmTargetKeyRootDir, KernelMode);
+
+    /* Release the captured object attributes */
+    ReleaseCapturedObjectAttributes(&CapturedSourceFile, PreviousMode);
+    ReleaseCapturedObjectAttributes(&CapturedTargetKey, PreviousMode);
+
+Quit:
     /* Bring back APCs */
     KeLeaveCriticalRegion();
 
@@ -1272,6 +1608,7 @@ NtSaveKeyEx(IN HANDLE KeyHandle,
             IN ULONG Flags)
 {
     NTSTATUS Status;
+    HANDLE KmFileHandle = NULL;
     PCM_KEY_BODY KeyObject;
     KPROCESSOR_MODE PreviousMode = ExGetPreviousMode();
 
@@ -1294,6 +1631,15 @@ NtSaveKeyEx(IN HANDLE KeyHandle,
         return STATUS_PRIVILEGE_NOT_HELD;
     }
 
+    /* Make sure the target file handle is a kernel handle */
+    Status = CmpConvertHandleToKernelHandle(FileHandle,
+                                            IoFileObjectType,
+                                            FILE_WRITE_DATA,
+                                            PreviousMode,
+                                            &KmFileHandle);
+    if (!NT_SUCCESS(Status))
+        goto Quit;
+
     /* Verify that the handle is valid and is a registry key */
     Status = ObReferenceObjectByHandle(KeyHandle,
                                        KEY_READ,
@@ -1301,12 +1647,20 @@ NtSaveKeyEx(IN HANDLE KeyHandle,
                                        PreviousMode,
                                        (PVOID*)&KeyObject,
                                        NULL);
-    if (!NT_SUCCESS(Status)) return Status;
+    if (!NT_SUCCESS(Status))
+        goto Quit;
 
     /* Call the internal API */
-    Status = CmSaveKey(KeyObject->KeyControlBlock, FileHandle, Flags);
+    Status = CmSaveKey(KeyObject->KeyControlBlock, KmFileHandle, Flags);
 
+    /* Dereference the registry key */
     ObDereferenceObject(KeyObject);
+
+Quit:
+    /* Close the local kernel handle */
+    if (KmFileHandle)
+        ObCloseHandle(KmFileHandle, KernelMode);
+
     return Status;
 }
 
@@ -1316,10 +1670,11 @@ NtSaveMergedKeys(IN HANDLE HighPrecedenceKeyHandle,
                  IN HANDLE LowPrecedenceKeyHandle,
                  IN HANDLE FileHandle)
 {
+    NTSTATUS Status;
     KPROCESSOR_MODE PreviousMode;
+    HANDLE KmFileHandle = NULL;
     PCM_KEY_BODY HighPrecedenceKeyObject = NULL;
     PCM_KEY_BODY LowPrecedenceKeyObject = NULL;
-    NTSTATUS Status;
 
     PAGED_CODE();
 
@@ -1334,6 +1689,15 @@ NtSaveMergedKeys(IN HANDLE HighPrecedenceKeyHandle,
         return STATUS_PRIVILEGE_NOT_HELD;
     }
 
+    /* Make sure the target file handle is a kernel handle */
+    Status = CmpConvertHandleToKernelHandle(FileHandle,
+                                            IoFileObjectType,
+                                            FILE_WRITE_DATA,
+                                            PreviousMode,
+                                            &KmFileHandle);
+    if (!NT_SUCCESS(Status))
+        goto Quit;
+
     /* Verify that the handles are valid and are registry keys */
     Status = ObReferenceObjectByHandle(HighPrecedenceKeyHandle,
                                        KEY_READ,
@@ -1342,7 +1706,7 @@ NtSaveMergedKeys(IN HANDLE HighPrecedenceKeyHandle,
                                        (PVOID*)&HighPrecedenceKeyObject,
                                        NULL);
     if (!NT_SUCCESS(Status))
-        goto done;
+        goto Quit;
 
     Status = ObReferenceObjectByHandle(LowPrecedenceKeyHandle,
                                        KEY_READ,
@@ -1351,20 +1715,24 @@ NtSaveMergedKeys(IN HANDLE HighPrecedenceKeyHandle,
                                        (PVOID*)&LowPrecedenceKeyObject,
                                        NULL);
     if (!NT_SUCCESS(Status))
-        goto done;
+        goto Quit;
 
     /* Call the internal API */
     Status = CmSaveMergedKeys(HighPrecedenceKeyObject->KeyControlBlock,
                               LowPrecedenceKeyObject->KeyControlBlock,
-                              FileHandle);
+                              KmFileHandle);
 
-done:
+Quit:
+    /* Dereference the opened key objects */
     if (LowPrecedenceKeyObject)
         ObDereferenceObject(LowPrecedenceKeyObject);
-
     if (HighPrecedenceKeyObject)
         ObDereferenceObject(HighPrecedenceKeyObject);
 
+    /* Close the local kernel handle */
+    if (KmFileHandle)
+        ObCloseHandle(KmFileHandle, KernelMode);
+
     return Status;
 }
 
@@ -1392,8 +1760,9 @@ NtUnloadKey2(IN POBJECT_ATTRIBUTES TargetKey,
              IN ULONG Flags)
 {
     NTSTATUS Status;
-    OBJECT_ATTRIBUTES ObjectAttributes;
+    OBJECT_ATTRIBUTES CapturedTargetKey;
     UNICODE_STRING ObjectName;
+    HANDLE KmTargetKeyRootDir = NULL;
     CM_PARSE_CONTEXT ParseContext = {0};
     KPROCESSOR_MODE PreviousMode = ExGetPreviousMode();
     PCM_KEY_BODY KeyBody = NULL;
@@ -1420,18 +1789,15 @@ NtUnloadKey2(IN POBJECT_ATTRIBUTES TargetKey,
                          sizeof(OBJECT_ATTRIBUTES),
                          sizeof(ULONG));
 
-            ObjectAttributes = *TargetKey;
+            CapturedTargetKey = *TargetKey;
 
             /* Probe the string */
-            ProbeForReadUnicodeString(&TargetKey->ObjectName);
+            ObjectName = ProbeForReadUnicodeString(CapturedTargetKey.ObjectName);
-
-            ObjectName = *TargetKey->ObjectName;
-
             ProbeForRead(ObjectName.Buffer,
                          ObjectName.Length,
                          sizeof(WCHAR));
 
-            ObjectAttributes.ObjectName = &ObjectName;
+            CapturedTargetKey.ObjectName = &ObjectName;
         }
         _SEH2_EXCEPT(EXCEPTION_EXECUTE_HANDLER)
         {
@@ -1443,38 +1809,54 @@ NtUnloadKey2(IN POBJECT_ATTRIBUTES TargetKey,
     else
     {
         /* Save the target attributes directly */
-        ObjectAttributes = *TargetKey;
+        CapturedTargetKey = *TargetKey;
     }
 
+    /* Make sure the target key root directory handle is a kernel handle */
+    Status = CmpConvertHandleToKernelHandle(CapturedTargetKey.RootDirectory,
+                                            CmpKeyObjectType,
+                                            KEY_WRITE,
+                                            PreviousMode,
+                                            &KmTargetKeyRootDir);
+    if (!NT_SUCCESS(Status))
+        return Status;
+    CapturedTargetKey.RootDirectory = KmTargetKeyRootDir;
+    CapturedTargetKey.Attributes |= OBJ_KERNEL_HANDLE;
+
     /* Setup the parse context */
     ParseContext.CreateOperation = TRUE;
     ParseContext.CreateOptions = REG_OPTION_BACKUP_RESTORE;
 
     /* Do the create */
-    Status = ObOpenObjectByName(&ObjectAttributes,
+    /* Open a local handle to the key */
+    Status = ObOpenObjectByName(&CapturedTargetKey,
                                 CmpKeyObjectType,
                                 KernelMode,
                                 NULL,
                                 KEY_WRITE,
                                 &ParseContext,
                                 &Handle);
+    if (NT_SUCCESS(Status))
+    {
+        /* Reference the key object */
+        Status = ObReferenceObjectByHandle(Handle,
+                                           KEY_WRITE,
+                                           CmpKeyObjectType,
+                                           KernelMode,
+                                           (PVOID*)&KeyBody,
+                                           NULL);
 
-    /* Return if failure encountered */
+        /* Close the handle */
-    if (!NT_SUCCESS(Status)) return Status;
+        ObCloseHandle(Handle, KernelMode);
+    }
 
-    /* Reference it */
+    /* Close the local kernel handle */
-    Status = ObReferenceObjectByHandle(Handle,
+    if (KmTargetKeyRootDir)
-                                       KEY_WRITE,
+        ObCloseHandle(KmTargetKeyRootDir, KernelMode);
-                                       CmpKeyObjectType,
-                                       KernelMode,
-                                       (PVOID *)&KeyBody,
-                                       NULL);
 
-    /* Close the handle */
+    /* Return if a failure was encountered */
-    ZwClose(Handle);
+    if (!NT_SUCCESS(Status))
-
+        return Status;
-    /* Return if failure encountered */
-    if (!NT_SUCCESS(Status)) return Status;
 
     /* Acquire the lock depending on flags */
     if (Flags == REG_FORCE_UNLOAD)
@@ -1506,7 +1888,7 @@ NtUnloadKey2(IN POBJECT_ATTRIBUTES TargetKey,
     {
         /* Return appropriate status */
         Status = STATUS_KEY_DELETED;
-        goto Quickie;
+        goto Quit;
     }
 
     /* Check if it's a read-only key */
@@ -1514,21 +1896,27 @@ NtUnloadKey2(IN POBJECT_ATTRIBUTES TargetKey,
     {
         /* Return appropriate status */
         Status = STATUS_ACCESS_DENIED;
-        goto Quickie;
+        goto Quit;
     }
 
-    /* Call the internal API */
+    /* Call the internal API. Note that CmUnloadKey() unlocks the registry only on success. */
-    Status = CmUnloadKey(KeyBody->KeyControlBlock,
+    Status = CmUnloadKey(KeyBody->KeyControlBlock, Flags);
-                         Flags);
 
     /* Check if we failed, but really need to succeed */
     if ((Status == STATUS_CANNOT_DELETE) && (Flags == REG_FORCE_UNLOAD))
     {
         /* TODO: We should perform another attempt here */
-        ASSERT(FALSE);
+        _SEH2_TRY
+        {
+            DPRINT1("NtUnloadKey2(%wZ): We want to force-unload the hive but couldn't unload it: Retrying is UNIMPLEMENTED!\n", TargetKey->ObjectName);
+        }
+        _SEH2_EXCEPT(EXCEPTION_EXECUTE_HANDLER)
+        {
+        }
+        _SEH2_END;
     }
 
-    /* If CmUnloadKey failed we need to unlock registry ourselves */
+    /* If CmUnloadKey() failed we need to unlock registry ourselves */
     if (!NT_SUCCESS(Status))
     {
         if (Flags != REG_FORCE_UNLOAD)
@@ -1544,7 +1932,7 @@ NtUnloadKey2(IN POBJECT_ATTRIBUTES TargetKey,
         CmpUnlockRegistry();
     }
 
-Quickie:
+Quit:
     /* Dereference the key */
     ObDereferenceObject(KeyBody);