Index: ntoskrnl/config/ntapi.c
===================================================================
--- ntoskrnl/config/ntapi.c	(revision 75167)
+++ ntoskrnl/config/ntapi.c	(working copy)
@@ -17,6 +17,233 @@
 BOOLEAN CmFirstTime = TRUE;
 extern ULONG InitSafeBootMode;
 
+/* PRIVATE FUNCTIONS *********************************************************/
+
+// static __inline
+// FORCEINLINE
+VOID
+ReleaseCapturedObjectAttributes(
+    IN POBJECT_ATTRIBUTES CapturedObjectAttributes,
+    IN KPROCESSOR_MODE AccessMode);
+
+/*
+ * Adapted from ntoskrnl/ob/oblife.c:ObpCaptureObjectCreateInformation()
+ */
+// static __inline
+// FORCEINLINE
+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;
+}
+
+/*
+ * Adapted from ntoskrnl/include/internal/ob_x.h:ObpReleaseObjectCreateInformation()
+ */
+// static __inline
+// FORCEINLINE
+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);
+    }
+}
+
+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
@@ -158,7 +388,7 @@
     /* Just let the object manager handle this */
     Status = ObOpenObjectByName(ObjectAttributes,
                                 CmpKeyObjectType,
-                                ExGetPreviousMode(),
+                                PreviousMode,
                                 NULL,
                                 DesiredAccess,
                                 &ParseContext,
@@ -267,7 +500,7 @@
     Status = ObReferenceObjectByHandle(KeyHandle,
                                        KEY_ENUMERATE_SUB_KEYS,
                                        CmpKeyObjectType,
-                                       ExGetPreviousMode(),
+                                       PreviousMode,
                                        (PVOID*)&KeyObject,
                                        NULL);
     if (!NT_SUCCESS(Status)) return Status;
@@ -352,7 +589,7 @@
     Status = ObReferenceObjectByHandle(KeyHandle,
                                        KEY_QUERY_VALUE,
                                        CmpKeyObjectType,
-                                       ExGetPreviousMode(),
+                                       PreviousMode,
                                        (PVOID*)&KeyObject,
                                        NULL);
     if (!NT_SUCCESS(Status)) return Status;
@@ -443,7 +682,7 @@
         Status = ObReferenceObjectByHandle(KeyHandle,
                                            0,
                                            CmpKeyObjectType,
-                                           ExGetPreviousMode(),
+                                           PreviousMode,
                                            (PVOID*)&KeyObject,
                                            &HandleInfo);
         if (NT_SUCCESS(Status))
@@ -463,7 +702,7 @@
         Status = ObReferenceObjectByHandle(KeyHandle,
                                            KEY_QUERY_VALUE,
                                            CmpKeyObjectType,
-                                           ExGetPreviousMode(),
+                                           PreviousMode,
                                            (PVOID*)&KeyObject,
                                            NULL);
     }
@@ -527,21 +766,34 @@
                 IN ULONG Length,
                 OUT PULONG ResultLength)
 {
     KPROCESSOR_MODE PreviousMode = ExGetPreviousMode();
     NTSTATUS Status;
     PCM_KEY_BODY KeyObject;
     REG_QUERY_VALUE_KEY_INFORMATION QueryValueKeyInfo;
     REG_POST_OPERATION_INFORMATION PostOperationInfo;
-    UNICODE_STRING ValueNameCopy = *ValueName;
+    UNICODE_STRING ValueNameCopy;
+
     PAGED_CODE();
+
     DPRINT("NtQueryValueKey() KH 0x%p, VN '%wZ', KVIC %d, Length %lu\n",
         KeyHandle, ValueName, KeyValueInformationClass, Length);
 
+    /* Reject classes we don't know about */
+    if ((KeyValueInformationClass != KeyValueBasicInformation)       &&
+        (KeyValueInformationClass != KeyValueFullInformation)        &&
+        (KeyValueInformationClass != KeyValuePartialInformation)     &&
+        (KeyValueInformationClass != KeyValueFullInformationAlign64) &&
+        (KeyValueInformationClass != KeyValuePartialInformationAlign64))
+    {
+        /* Fail */
+        return STATUS_INVALID_PARAMETER;
+    }
+
     /* Verify that the handle is valid and is a registry key */
     Status = ObReferenceObjectByHandle(KeyHandle,
                                        KEY_QUERY_VALUE,
                                        CmpKeyObjectType,
-                                       ExGetPreviousMode(),
+                                       PreviousMode,
                                        (PVOID*)&KeyObject,
                                        NULL);
     if (!NT_SUCCESS(Status)) return Status;
@@ -564,6 +816,18 @@
         _SEH2_END;
     }
 
+    //
+    // FIXME: Probe and certainly capture 'ValueName' !
+    //
+#if 0
+    /* Capture the string */
+    Status = ProbeAndCaptureUnicodeString(&ValueNameCopy, PreviousMode, ValueName);
+    if (!NT_SUCCESS(Status))
+        goto end;
+#else
+    ValueNameCopy = *ValueName;
+#endif
+
     /* Make sure the name is aligned properly */
     if ((ValueNameCopy.Length & (sizeof(WCHAR) - 1)))
     {
@@ -638,6 +902,7 @@
     /* Probe and copy the data */
     if ((PreviousMode != KernelMode) && (DataSize != 0))
     {
+        // FIXME: Not captured by Win2k3!!
         PVOID DataCopy = ExAllocatePoolWithTag(PagedPool, DataSize, TAG_CM);
         if (!DataCopy)
             return STATUS_INSUFFICIENT_RESOURCES;
@@ -672,7 +937,7 @@
     Status = ObReferenceObjectByHandle(KeyHandle,
                                        KEY_SET_VALUE,
                                        CmpKeyObjectType,
-                                       ExGetPreviousMode(),
+                                       PreviousMode,
                                        (PVOID*)&KeyObject,
                                        NULL);
     if (!NT_SUCCESS(Status))
@@ -723,12 +988,12 @@
                                Type,
                                Data,
                                DataSize);
+
+        /* Do the post-callback */
+        PostOperationInfo.Status = Status;
+        CmiCallRegisteredCallbacks(RegNtPostSetValueKey, &PostOperationInfo);
     }
 
-    /* Do the post-callback */
-    PostOperationInfo.Status = Status;
-    CmiCallRegisteredCallbacks(RegNtPostSetValueKey, &PostOperationInfo);
-
 end:
     /* Dereference and return status */
     if (KeyObject)
@@ -749,7 +1014,8 @@
     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 */
@@ -757,17 +1023,21 @@
                                        KEY_SET_VALUE,
                                        CmpKeyObjectType,
                                        PreviousMode,
-                                       (PVOID *)&KeyObject,
+                                       (PVOID*)&KeyObject,
                                        NULL);
     if (!NT_SUCCESS(Status)) return Status;
 
-    /* Don't touch read-only keys */
-    if (KeyObject->KeyControlBlock->ExtFlags & CM_KCB_READ_ONLY_KEY)
-    {
-        /* Fail */
-        ObDereferenceObject(KeyObject);
-        return STATUS_ACCESS_DENIED;
-    }
+    //
+    // FIXME: Probe and certainly capture 'ValueName' !
+    //
+#if 0
+    /* Capture the string */
+    Status = ProbeAndCaptureUnicodeString(&ValueNameCopy, PreviousMode, ValueName);
+    if (!NT_SUCCESS(Status))
+        goto end;
+#else
+    ValueNameCopy = *ValueName;
+#endif
 
     /* Make sure the name is aligned properly */
     if ((ValueNameCopy.Length & (sizeof(WCHAR) - 1)))
@@ -777,6 +1047,14 @@
         return STATUS_INVALID_PARAMETER;
     }
 
+    /* Don't touch read-only keys */
+    if (KeyObject->KeyControlBlock->ExtFlags & CM_KCB_READ_ONLY_KEY)
+    {
+        /* Fail */
+        ObDereferenceObject(KeyObject);
+        return STATUS_ACCESS_DENIED;
+    }
+
     /* Do the callback */
     DeleteValueKeyInfo.Object = (PVOID)KeyObject;
     DeleteValueKeyInfo.ValueName = ValueName;
@@ -871,16 +1150,20 @@
 {
     NTSTATUS Status;
     KPROCESSOR_MODE PreviousMode = ExGetPreviousMode();
+    OBJECT_ATTRIBUTES CapturedTargetKey;
+    OBJECT_ATTRIBUTES CapturedSourceFile;
+    UNICODE_STRING TargetKeyName, SourceFileName;
     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))
     {
-        /* Fail */
         DPRINT1("Restore Privilege missing!\n");
         return STATUS_PRIVILEGE_NOT_HELD;
     }
@@ -888,6 +1171,77 @@
     /* 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;
+    }
+
+    DbgBreakPoint(); // For testing purposes only!!
+
+    /* 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 handle is a kernel handle */
+    Status = CmpConvertHandleToKernelHandle(CapturedTargetKey.RootDirectory,
+                                            CmpKeyObjectType,
+                                            KEY_READ,
+                                            PreviousMode,
+                                            &CapturedTargetKey.RootDirectory);
+    if (!NT_SUCCESS(Status))
+        goto Cleanup;
+    CapturedTargetKey.Attributes |= OBJ_KERNEL_HANDLE;
+
+    /* Make sure the source file handle is a kernel handle */
+    Status = CmpConvertHandleToKernelHandle(CapturedSourceFile.RootDirectory,
+                                            IoFileObjectType,
+                                            FILE_TRAVERSE,
+                                            PreviousMode,
+                                            &CapturedSourceFile.RootDirectory);
+    if (!NT_SUCCESS(Status))
+        goto Cleanup;
+    CapturedSourceFile.Attributes |= OBJ_KERNEL_HANDLE;
+
     /* Check if we have a trust class */
     if (TrustClassKey)
     {
@@ -896,16 +1250,31 @@
                                            0,
                                            CmpKeyObjectType,
                                            PreviousMode,
-                                           (PVOID *)&KeyBody,
+                                           (PVOID*)&KeyBody,
                                            NULL);
     }
 
     /* 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 (CapturedSourceFile.RootDirectory)
+        ObCloseHandle(CapturedSourceFile.RootDirectory, KernelMode);
+    if (CapturedTargetKey.RootDirectory)
+        ObCloseHandle(CapturedTargetKey.RootDirectory, KernelMode);
+
+    /* Release the captured object attributes */
+    ReleaseCapturedObjectAttributes(&CapturedSourceFile, PreviousMode);
+    ReleaseCapturedObjectAttributes(&CapturedTargetKey, PreviousMode);
+
+Quit:
     /* Bring back APCs */
     KeLeaveCriticalRegion();
 
@@ -1128,6 +1499,7 @@
     /* Get the processor mode */
     PreviousMode = KeGetPreviousMode();
 
+    /* Check for user-mode caller */
     if (PreviousMode != KernelMode)
     {
         /* Prepare to probe parameters */
@@ -1149,6 +1521,10 @@
         _SEH2_END;
     }
 
+    //
+    // FIXME: Shouldn't it be better to get this private handle in kernel mode?
+    //
+
     /* Open a handle to the key */
     Status = ObOpenObjectByName(TargetKey,
                                 CmpKeyObjectType,
@@ -1164,7 +1540,7 @@
                                            KEY_READ,
                                            CmpKeyObjectType,
                                            PreviousMode,
-                                           (PVOID *)&KeyBody,
+                                           (PVOID*)&KeyBody,
                                            NULL);
 
         /* Close the handle */
@@ -1275,12 +1651,37 @@
         return STATUS_INVALID_PARAMETER;
     }
 
-    /* Check for the SeBackupPrivilege */
+    /* Validate privilege */
     if (!SeSinglePrivilegeCheck(SeBackupPrivilege, PreviousMode))
     {
+        DPRINT1("Backup Privilege missing!\n");
         return STATUS_PRIVILEGE_NOT_HELD;
     }
 
+    //
+    // FIXME! Do this always in kernel mode context!!
+    //
+
+#if 0
+    /* Make sure the target key handle is a kernel handle */
+    Status = CmpConvertHandleToKernelHandle(KeyHandle,
+                                            CmpKeyObjectType,
+                                            KEY_READ,
+                                            PreviousMode,
+                                            &KeyHandle);
+    if (!NT_SUCCESS(Status))
+        return Status;
+
+    /* Make sure the target key handle is a kernel handle */
+    Status = CmpConvertHandleToKernelHandle(FileHandle,
+                                            IoFileObjectType,
+                                            FILE_WRITE_DATA,
+                                            PreviousMode,
+                                            &FileHandle);
+    if (!NT_SUCCESS(Status))
+        return Status;
+#endif
+
     /* Verify that the handle is valid and is a registry key */
     Status = ObReferenceObjectByHandle(KeyHandle,
                                        KEY_READ,
@@ -1315,12 +1716,46 @@
 
     PreviousMode = ExGetPreviousMode();
 
-    /* Check for the SeBackupPrivilege */
+    /* Validate privilege */
     if (!SeSinglePrivilegeCheck(SeBackupPrivilege, PreviousMode))
     {
+        DPRINT1("Backup Privilege missing!\n");
         return STATUS_PRIVILEGE_NOT_HELD;
     }
 
+    //
+    // FIXME! Do this always in kernel mode context!!
+    //
+
+#if 0
+    /* Make sure the target key handle is a kernel handle */
+    Status = CmpConvertHandleToKernelHandle(HighPrecedenceKeyHandle,
+                                            CmpKeyObjectType,
+                                            KEY_READ,
+                                            PreviousMode,
+                                            &HighPrecedenceKeyHandle);
+    if (!NT_SUCCESS(Status))
+        return Status;
+
+    /* Make sure the target key handle is a kernel handle */
+    Status = CmpConvertHandleToKernelHandle(LowPrecedenceKeyHandle,
+                                            CmpKeyObjectType,
+                                            KEY_READ,
+                                            PreviousMode,
+                                            &LowPrecedenceKeyHandle);
+    if (!NT_SUCCESS(Status))
+        return Status;
+
+    /* Make sure the target key handle is a kernel handle */
+    Status = CmpConvertHandleToKernelHandle(FileHandle,
+                                            IoFileObjectType,
+                                            FILE_WRITE_DATA,
+                                            PreviousMode,
+                                            &FileHandle);
+    if (!NT_SUCCESS(Status))
+        return Status;
+#endif
+
     /* Verify that the handles are valid and are registry keys */
     Status = ObReferenceObjectByHandle(HighPrecedenceKeyHandle,
                                        KEY_READ,
@@ -1379,7 +1814,7 @@
              IN ULONG Flags)
 {
     NTSTATUS Status;
-    OBJECT_ATTRIBUTES ObjectAttributes;
+    OBJECT_ATTRIBUTES CapturedTargetKey;
     UNICODE_STRING ObjectName;
     CM_PARSE_CONTEXT ParseContext = {0};
     KPROCESSOR_MODE PreviousMode = ExGetPreviousMode();
@@ -1392,7 +1827,6 @@
     /* Validate privilege */
     if (!SeSinglePrivilegeCheck(SeRestorePrivilege, PreviousMode))
     {
-        /* Fail */
         DPRINT1("Restore Privilege missing!\n");
         return STATUS_PRIVILEGE_NOT_HELD;
     }
@@ -1408,18 +1842,15 @@
                          sizeof(OBJECT_ATTRIBUTES),
                          sizeof(ULONG));
 
-            ObjectAttributes = *TargetKey;
+            CapturedTargetKey = *TargetKey;
 
             /* Probe the string */
-            ProbeForReadUnicodeString(&TargetKey->ObjectName);
-
-            ObjectName = *TargetKey->ObjectName;
-
+            ObjectName = ProbeForReadUnicodeString(TargetKey->ObjectName);
             ProbeForRead(ObjectName.Buffer,
                          ObjectName.Length,
                          sizeof(WCHAR));
 
-            ObjectAttributes.ObjectName = &ObjectName;
+            CapturedTargetKey.ObjectName = &ObjectName;
         }
         _SEH2_EXCEPT(EXCEPTION_EXECUTE_HANDLER)
         {
@@ -1431,15 +1862,26 @@
     else
     {
         /* Save the target attributes directly */
-        ObjectAttributes = *TargetKey;
+        CapturedTargetKey = *TargetKey;
     }
 
+    /* Make sure the target key handle is a kernel handle */
+    Status = CmpConvertHandleToKernelHandle(CapturedTargetKey.RootDirectory,
+                                            CmpKeyObjectType,
+                                            KEY_WRITE,
+                                            PreviousMode,
+                                            &CapturedTargetKey.RootDirectory);
+    if (!NT_SUCCESS(Status))
+        return Status;
+    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,
@@ -1446,23 +1888,26 @@
                                 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 */
-    if (!NT_SUCCESS(Status)) return Status;
+        /* Close the handle */
+        ObCloseHandle(Handle, KernelMode);
+    }
 
-    /* Reference it */
-    Status = ObReferenceObjectByHandle(Handle,
-                                       KEY_WRITE,
-                                       CmpKeyObjectType,
-                                       KernelMode,
-                                       (PVOID *)&KeyBody,
-                                       NULL);
+    /* Close the local kernel handle */
+    ObCloseHandle(CapturedTargetKey.RootDirectory, KernelMode);
 
-    /* Close the handle */
-    ZwClose(Handle);
-
     /* Return if failure encountered */
-    if (!NT_SUCCESS(Status)) return Status;
+    if (!NT_SUCCESS(Status))
+        return Status;
 
     /* Acquire the lock depending on flags */
     if (Flags == REG_FORCE_UNLOAD)
@@ -1494,26 +1939,32 @@
     {
         /* Return appropriate status */
         Status = STATUS_KEY_DELETED;
-        goto Quickie;
+        goto Quit;
     }
 
-    /* Check if it's a readonly key */
+    /* Check if it's a read-only key */
     if (KeyBody->KeyControlBlock->ExtFlags & CM_KCB_READ_ONLY_KEY)
     {
         /* Return appropriate status */
         Status = STATUS_ACCESS_DENIED;
-        goto Quickie;
+        goto Quit;
     }
 
     /* Call the internal API */
-    Status = CmUnloadKey(KeyBody->KeyControlBlock,
-                         Flags);
+    Status = CmUnloadKey(KeyBody->KeyControlBlock, 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 */
@@ -1532,7 +1983,7 @@
         CmpUnlockRegistry();
     }
 
-Quickie:
+Quit:
     /* Dereference the key */
     ObDereferenceObject(KeyBody);