[NTOSKRNL]: Implement correct locking and unlocking of the working set, one of the biggest blunders in ARM3 so far.

[NTOSKRNL]: Implement MiDereferenceControlArea to avoid leaking CAs in failure cases.

svn path=/trunk/; revision=57228

reactos/ntoskrnl/include/internal/mm.h

@@ -1718,17 +1718,13 @@ VOID
 MmLockAddressSpace(PMMSUPPORT AddressSpace)
 {
     KeAcquireGuardedMutex(&CONTAINING_RECORD(AddressSpace, EPROCESS, Vm)->AddressCreationLock);
-    //ASSERT(Thread->OwnsProcessAddressSpaceExclusive == 0);
-    //Thread->OwnsProcessAddressSpaceExclusive = TRUE;
 }
 
 FORCEINLINE
 VOID
 MmUnlockAddressSpace(PMMSUPPORT AddressSpace)
 {
-    //ASSERT(Thread->OwnsProcessAddressSpaceExclusive == 1);
     KeReleaseGuardedMutex(&CONTAINING_RECORD(AddressSpace, EPROCESS, Vm)->AddressCreationLock);
-    //Thread->OwnsProcessAddressSpaceExclusive = 0;
 }
 
 FORCEINLINE

reactos/ntoskrnl/mm/ARM3/miarm.h

@@ -1041,6 +1041,14 @@ BOOLEAN
 MI_WS_OWNER(IN PEPROCESS Process)
 {
     /* Check if this process is the owner, and that the thread owns the WS */
+    if (PsGetCurrentThread()->OwnsProcessWorkingSetExclusive == 0)
+    {
+        DPRINT1("Thread: %p is not an owner\n", PsGetCurrentThread());
+    }
+    if (KeGetCurrentThread()->ApcState.Process != &Process->Pcb)
+    {
+        DPRINT1("Current thread %p is attached to another process %p\n", PsGetCurrentThread(), Process);
+    }
     return ((KeGetCurrentThread()->ApcState.Process == &Process->Pcb) &&
             ((PsGetCurrentThread()->OwnsProcessWorkingSetExclusive) ||
              (PsGetCurrentThread()->OwnsProcessWorkingSetShared)));
@@ -1083,6 +1091,13 @@ MiDecrementReferenceCount(
     IN PFN_NUMBER PageFrameIndex
 );
 
+FORCEINLINE
+BOOLEAN
+MI_IS_WS_UNSAFE(IN PEPROCESS Process)
+{
+    return (Process->Vm.Flags.AcquiredUnsafe == TRUE);
+}
+
 //
 // Locks the working set for the given process
 //
@@ -1099,12 +1114,57 @@ MiLockProcessWorkingSet(IN PEPROCESS Process,
     KeEnterGuardedRegion();
     ASSERT(!MM_ANY_WS_LOCK_HELD(Thread));
 
-    /* FIXME: Actually lock it (we can't because Vm is used by MAREAs) */
+    /* Lock the working set */
+    ExAcquirePushLockExclusive(&Process->Vm.WorkingSetMutex);
 
-    /* FIXME: This also can't be checked because Vm is used by MAREAs) */
+    /* Now claim that we own the lock */
-    //ASSERT(Process->Vm.Flags.AcquiredUnsafe == 0);
+    ASSERT(!MI_IS_WS_UNSAFE(Process));
+    ASSERT(Thread->OwnsProcessWorkingSetExclusive == FALSE);
+    Thread->OwnsProcessWorkingSetExclusive = TRUE;
+}
 
-    /* Okay, now we can own it exclusively */
+FORCEINLINE
+VOID
+MiLockProcessWorkingSetShared(IN PEPROCESS Process,
+                              IN PETHREAD Thread)
+{
+    /* Shouldn't already be owning the process working set */
+    ASSERT(Thread->OwnsProcessWorkingSetShared == FALSE);
+    ASSERT(Thread->OwnsProcessWorkingSetExclusive == FALSE);
+
+    /* Block APCs, make sure that still nothing is already held */
+    KeEnterGuardedRegion();
+    ASSERT(!MM_ANY_WS_LOCK_HELD(Thread));
+
+    /* Lock the working set */
+    ExAcquirePushLockShared(&Process->Vm.WorkingSetMutex);
+
+    /* Now claim that we own the lock */
+    ASSERT(!MI_IS_WS_UNSAFE(Process));
+    ASSERT(Thread->OwnsProcessWorkingSetShared == FALSE);
+    ASSERT(Thread->OwnsProcessWorkingSetExclusive == FALSE);
+    Thread->OwnsProcessWorkingSetShared = TRUE;
+}
+
+FORCEINLINE
+VOID
+MiLockProcessWorkingSetUnsafe(IN PEPROCESS Process,
+                              IN PETHREAD Thread)
+{
+    /* Shouldn't already be owning the process working set */
+    ASSERT(Thread->OwnsProcessWorkingSetExclusive == FALSE);
+
+    /* APCs must be blocked, make sure that still nothing is already held */
+    ASSERT(KeAreAllApcsDisabled() == TRUE);
+    ASSERT(!MM_ANY_WS_LOCK_HELD(Thread));
+
+    /* Lock the working set */
+    ExAcquirePushLockExclusive(&Process->Vm.WorkingSetMutex);
+
+    /* Now claim that we own the lock */
+    ASSERT(!MI_IS_WS_UNSAFE(Process));
+    Process->Vm.Flags.AcquiredUnsafe = 1;
+    ASSERT(Thread->OwnsProcessWorkingSetExclusive == FALSE);
     Thread->OwnsProcessWorkingSetExclusive = TRUE;
 }
 
@@ -1116,21 +1176,45 @@ VOID
 MiUnlockProcessWorkingSet(IN PEPROCESS Process,
                           IN PETHREAD Thread)
 {
-    /* Make sure this process really is owner, and it was a safe acquisition */
+    /* Make sure we are the owner of a safe acquisition */
     ASSERT(MI_WS_OWNER(Process));
-    /* This can't be checked because Vm is used by MAREAs) */
+    ASSERT(!MI_IS_WS_UNSAFE(Process));
-    //ASSERT(Process->Vm.Flags.AcquiredUnsafe == 0);
 
     /* The thread doesn't own it anymore */
     ASSERT(Thread->OwnsProcessWorkingSetExclusive == TRUE);
     Thread->OwnsProcessWorkingSetExclusive = FALSE;
 
-    /* FIXME: Actually release it (we can't because Vm is used by MAREAs) */
+    /* Release the lock and re-enable APCs */
-
+    ExReleasePushLockExclusive(&Process->Vm.WorkingSetMutex);
-    /* Unblock APCs */
     KeLeaveGuardedRegion();
 }
 
+//
+// Unlocks the working set for the given process
+//
+FORCEINLINE
+VOID
+MiUnlockProcessWorkingSetUnsafe(IN PEPROCESS Process,
+                                IN PETHREAD Thread)
+{
+    /* Make sure we are the owner of an unsafe acquisition */
+    ASSERT(KeGetCurrentIrql() <= APC_LEVEL);
+    ASSERT(KeAreAllApcsDisabled() == TRUE);
+    ASSERT(MI_WS_OWNER(Process));
+    ASSERT(MI_IS_WS_UNSAFE(Process));
+
+    /* No longer unsafe */
+    Process->Vm.Flags.AcquiredUnsafe = 0;
+
+    /* The thread doesn't own it anymore */
+    ASSERT(Thread->OwnsProcessWorkingSetExclusive == TRUE);
+    Thread->OwnsProcessWorkingSetExclusive = FALSE;
+
+    /* Release the lock but don't touch APC state */
+    ExReleasePushLockExclusive(&Process->Vm.WorkingSetMutex);
+    ASSERT(KeGetCurrentIrql() <= APC_LEVEL);
+}
+
 //
 // Locks the working set
 //
@@ -1148,7 +1232,8 @@ MiLockWorkingSet(IN PETHREAD Thread,
     /* Thread shouldn't already be owning something */
     ASSERT(!MM_ANY_WS_LOCK_HELD(Thread));
 
-    /* FIXME: Actually lock it (we can't because Vm is used by MAREAs) */
+    /* Lock this working set */
+    ExAcquirePushLockExclusive(&WorkingSet->WorkingSetMutex);
 
     /* Which working set is this? */
     if (WorkingSet == &MmSystemCacheWs)
@@ -1208,12 +1293,68 @@ MiUnlockWorkingSet(IN PETHREAD Thread,
         Thread->OwnsProcessWorkingSetExclusive = FALSE;
     }
 
-    /* FIXME: Actually release it (we can't because Vm is used by MAREAs) */
+    /* Release the working set lock */
+    ExReleasePushLockExclusive(&WorkingSet->WorkingSetMutex);
 
     /* Unblock APCs */
     KeLeaveGuardedRegion();
 }
 
+FORCEINLINE
+VOID
+MiUnlockProcessWorkingSetForFault(IN PEPROCESS Process,
+                                  IN PETHREAD Thread,
+                                  IN BOOLEAN Safe,
+                                  IN BOOLEAN Shared)
+{
+    ASSERT(MI_WS_OWNER(Process));
+
+    /* Check if the current owner is unsafe */
+    if (MI_IS_WS_UNSAFE(Process))
+    {
+        /* Release unsafely */
+        MiUnlockProcessWorkingSetUnsafe(Process, Thread);
+        Safe = FALSE;
+        Shared = FALSE;
+    }
+    else if (Thread->OwnsProcessWorkingSetExclusive == 1)
+    {
+        /* Owner is safe and exclusive, release normally */
+        MiUnlockProcessWorkingSet(Process, Thread);
+        Safe = TRUE;
+        Shared = FALSE;
+    }
+    else
+    {
+        /* Owner is shared (implies safe), release normally */
+        ASSERT(FALSE);
+        Safe = TRUE;
+        Shared = TRUE;
+    }
+}
+
+FORCEINLINE
+VOID
+MiLockProcessWorkingSetForFault(IN PEPROCESS Process,
+                                IN PETHREAD Thread,
+                                IN BOOLEAN Safe,
+                                IN BOOLEAN Shared)
+{
+    ASSERT(Shared == FALSE);
+
+    /* Check if this was a safe lock or not */
+    if (Safe)
+    {
+        /* Reacquire safely */
+        MiLockProcessWorkingSet(Process, Thread);
+    }
+    else
+    {
+        /* Reacquire unsafely */
+        MiLockProcessWorkingSetUnsafe(Process, Thread);
+    }
+}
+
 //
 // Returns the ProtoPTE inside a VAD for the given VPN
 //

reactos/ntoskrnl/mm/ARM3/procsup.c

@@ -135,7 +135,7 @@ MiCreatePebOrTeb(IN PEPROCESS Process,
     Status = STATUS_SUCCESS;
 
     /* Pretend as if we own the working set */
-    MiLockProcessWorkingSet(Process, Thread);
+    MiLockProcessWorkingSetUnsafe(Process, Thread);
 
     /* Insert the VAD */
     ASSERT(Vad->EndingVpn >= Vad->StartingVpn);
@@ -147,7 +147,7 @@ MiCreatePebOrTeb(IN PEPROCESS Process,
     MiInsertNode(&Process->VadRoot, (PVOID)Vad, Parent, Result);
 
     /* Release the working set */
-    MiUnlockProcessWorkingSet(Process, Thread);
+    MiUnlockProcessWorkingSetUnsafe(Process, Thread);
 
     /* Release the address space lock */
     KeReleaseGuardedMutex(&Process->AddressCreationLock);
@@ -195,7 +195,7 @@ MmDeleteTeb(IN PEPROCESS Process,
         ASSERT(Vad->u2.VadFlags2.MultipleSecured == FALSE);
 
         /* Lock the working set */
-        MiLockProcessWorkingSet(Process, Thread);
+        MiLockProcessWorkingSetUnsafe(Process, Thread);
 
         /* Remove this VAD from the tree */
         ASSERT(VadTree->NumberGenericTableElements >= 1);
@@ -205,7 +205,7 @@ MmDeleteTeb(IN PEPROCESS Process,
         MiDeleteVirtualAddresses((ULONG_PTR)Teb, TebEnd, NULL);
 
         /* Release the working set */
-        MiUnlockProcessWorkingSet(Process, Thread);
+        MiUnlockProcessWorkingSetUnsafe(Process, Thread);
 
         /* Remove the VAD */
         ExFreePool(Vad);
@@ -1338,9 +1338,11 @@ MmCleanProcessAddressSpace(IN PEPROCESS Process)
 
     /* Lock the process address space from changes */
     MmLockAddressSpace(&Process->Vm);
+    MiLockProcessWorkingSetUnsafe(Process, Thread);
 
     /* VM is deleted now */
     Process->VmDeleted = TRUE;
+    MiUnlockProcessWorkingSetUnsafe(Process, Thread);
 
     /* Enumerate the VADs */
     VadTree = &Process->VadRoot;
@@ -1350,7 +1352,7 @@ MmCleanProcessAddressSpace(IN PEPROCESS Process)
         Vad = (PMMVAD)VadTree->BalancedRoot.RightChild;
 
         /* Lock the working set */
-        MiLockProcessWorkingSet(Process, Thread);
+        MiLockProcessWorkingSetUnsafe(Process, Thread);
 
         /* Remove this VAD from the tree */
         ASSERT(VadTree->NumberGenericTableElements >= 1);
@@ -1373,7 +1375,7 @@ MmCleanProcessAddressSpace(IN PEPROCESS Process)
                                      Vad);
 
             /* Release the working set */
-            MiUnlockProcessWorkingSet(Process, Thread);
+            MiUnlockProcessWorkingSetUnsafe(Process, Thread);
         }
 
         /* Skip ARM3 fake VADs, they'll be freed by MmDeleteProcessAddresSpace */
@@ -1388,9 +1390,17 @@ MmCleanProcessAddressSpace(IN PEPROCESS Process)
         ExFreePool(Vad);
     }
 
+    /* Lock the working set */
+    MiLockProcessWorkingSetUnsafe(Process, Thread);
+    ASSERT(Process->CloneRoot == NULL);
+    ASSERT(Process->PhysicalVadRoot == NULL);
+
     /* Delete the shared user data section */
     MiDeleteVirtualAddresses(USER_SHARED_DATA, USER_SHARED_DATA, NULL);
 
+    /* Release the working set */
+    MiUnlockProcessWorkingSetUnsafe(Process, Thread);
+
     /* Release the address space */
     MmUnlockAddressSpace(&Process->Vm);
 }

reactos/ntoskrnl/mm/ARM3/section.c

@@ -84,22 +84,22 @@ CHAR MmUserProtectionToMask2[16] =
 ULONG MmCompatibleProtectionMask[8] =
 {
     PAGE_NOACCESS,
-    
+
     PAGE_NOACCESS | PAGE_READONLY | PAGE_WRITECOPY,
-    
+
     PAGE_NOACCESS | PAGE_EXECUTE,
-    
+
     PAGE_NOACCESS | PAGE_READONLY | PAGE_WRITECOPY | PAGE_EXECUTE |
     PAGE_EXECUTE_READ,
-    
+
     PAGE_NOACCESS | PAGE_READONLY | PAGE_WRITECOPY | PAGE_READWRITE,
-    
+
     PAGE_NOACCESS | PAGE_READONLY | PAGE_WRITECOPY,
-    
+
     PAGE_NOACCESS | PAGE_READONLY | PAGE_WRITECOPY | PAGE_READWRITE |
     PAGE_EXECUTE | PAGE_EXECUTE_READ | PAGE_EXECUTE_READWRITE |
     PAGE_EXECUTE_WRITECOPY,
-    
+
     PAGE_NOACCESS | PAGE_READONLY | PAGE_WRITECOPY | PAGE_EXECUTE |
     PAGE_EXECUTE_READ | PAGE_EXECUTE_WRITECOPY
 };
@@ -711,6 +711,23 @@ MiCheckControlArea(IN PCONTROL_AREA ControlArea,
     }
 }
 
+VOID
+NTAPI
+MiDereferenceControlArea(IN PCONTROL_AREA ControlArea)
+{
+    KIRQL OldIrql;
+
+    /* Lock the PFN database */
+    OldIrql = KeAcquireQueuedSpinLock(LockQueuePfnLock);
+
+    /* Drop reference counts */
+    ControlArea->NumberOfMappedViews--;
+    ControlArea->NumberOfUserReferences--;
+
+    /* Check if it's time to delete the CA. This releases the lock */
+    MiCheckControlArea(ControlArea, OldIrql);
+}
+
 VOID
 NTAPI
 MiRemoveMappedView(IN PEPROCESS CurrentProcess,
@@ -718,6 +735,7 @@ MiRemoveMappedView(IN PEPROCESS CurrentProcess,
 {
     KIRQL OldIrql;
     PCONTROL_AREA ControlArea;
+    PETHREAD CurrentThread = PsGetCurrentThread();
 
     /* Get the control area */
     ControlArea = Vad->ControlArea;
@@ -734,7 +752,7 @@ MiRemoveMappedView(IN PEPROCESS CurrentProcess,
                              Vad);
 
     /* Release the working set */
-    MiUnlockProcessWorkingSet(CurrentProcess, PsGetCurrentThread());
+    MiUnlockProcessWorkingSetUnsafe(CurrentProcess, CurrentThread);
 
     /* Lock the PFN database */
     OldIrql = KeAcquireQueuedSpinLock(LockQueuePfnLock);
@@ -760,6 +778,8 @@ MiUnmapViewOfSection(IN PEPROCESS Process,
     PVOID DbgBase = NULL;
     SIZE_T RegionSize;
     NTSTATUS Status;
+    PETHREAD CurrentThread = PsGetCurrentThread();
+    PEPROCESS CurrentProcess = PsGetCurrentProcess();
     PAGED_CODE();
 
     /* Check for Mm Region */
@@ -771,7 +791,7 @@ MiUnmapViewOfSection(IN PEPROCESS Process,
     }
 
     /* Check if we should attach to the process */
-    if (PsGetCurrentProcess() != Process)
+    if (CurrentProcess != Process)
     {
         /* The process is different, do an attach */
         KeStackAttachProcess(&Process->Pcb, &ApcState);
@@ -837,13 +857,13 @@ MiUnmapViewOfSection(IN PEPROCESS Process,
     /* FIXME: Remove VAD charges */
 
     /* Lock the working set */
-    MiLockWorkingSet(PsGetCurrentThread(), &Process->Vm);
+    MiLockProcessWorkingSetUnsafe(Process, CurrentThread);
 
     /* Remove the VAD */
     ASSERT(Process->VadRoot.NumberGenericTableElements >= 1);
     MiRemoveNode((PMMADDRESS_NODE)Vad, &Process->VadRoot);
 
-    /* Remove the PTEs for this view */
+    /* Remove the PTEs for this view, which also releases the working set lock */
     MiRemoveMappedView(Process, Vad);
 
     /* FIXME: Remove commitment */
@@ -1205,7 +1225,8 @@ MiMapViewOfDataSection(IN PCONTROL_AREA ControlArea,
                                       &Process->VadRoot))
         {
             DPRINT1("Conflict with SEC_BASED or manually based section!\n");
-            return STATUS_CONFLICTING_ADDRESSES; // FIXME: CA Leak
+            MiDereferenceControlArea(ControlArea);
+            return STATUS_CONFLICTING_ADDRESSES;
         }
     }
 
@@ -1213,7 +1234,11 @@ MiMapViewOfDataSection(IN PCONTROL_AREA ControlArea,
     /* FIXME: we are allocating a LONG VAD for ReactOS compatibility only */
     ASSERT((AllocationType & MEM_RESERVE) == 0); /* ARM3 does not support this */
     Vad = ExAllocatePoolWithTag(NonPagedPool, sizeof(MMVAD_LONG), 'ldaV');
-    if (!Vad) return STATUS_INSUFFICIENT_RESOURCES; /* FIXME: CA Leak */
+    if (!Vad)
+    {
+        MiDereferenceControlArea(ControlArea);
+        return STATUS_INSUFFICIENT_RESOURCES;
+    }
     RtlZeroMemory(Vad, sizeof(MMVAD_LONG));
     Vad->u4.Banked = (PVOID)0xDEADBABE;
 
@@ -1244,13 +1269,13 @@ MiMapViewOfDataSection(IN PCONTROL_AREA ControlArea,
     Status = STATUS_SUCCESS;
 
     /* Pretend as if we own the working set */
-    MiLockProcessWorkingSet(Process, Thread);
+    MiLockProcessWorkingSetUnsafe(Process, Thread);
 
     /* Insert the VAD */
     MiInsertVad((PMMVAD)Vad, Process);
 
     /* Release the working set */
-    MiUnlockProcessWorkingSet(Process, Thread);
+    MiUnlockProcessWorkingSetUnsafe(Process, Thread);
 
     /* Windows stores this for accounting purposes, do so as well */
     if (!Segment->u2.FirstMappedVa) Segment->u2.FirstMappedVa = (PVOID)StartAddress;
@@ -1797,7 +1822,7 @@ MiSetProtectionOnSection(IN PEPROCESS Process,
     PMMPFN Pfn1;
     ULONG ProtectionMask, QuotaCharge = 0;
     PUSHORT UsedPageTableEntries;
-    //PETHREAD Thread = PsGetCurrentThread();
+    PETHREAD Thread = PsGetCurrentThread();
     PAGED_CODE();
 
     //
@@ -1829,7 +1854,7 @@ MiSetProtectionOnSection(IN PEPROCESS Process,
     //
     // Get the PTE and PDE for the address, as well as the final PTE
     //
-    //MiLockProcessWorkingSet(Thread, Process);
+    MiLockProcessWorkingSetUnsafe(Process, Thread);
     PointerPde = MiAddressToPde(StartingAddress);
     PointerPte = MiAddressToPte(StartingAddress);
     LastPte = MiAddressToPte(EndingAddress);
@@ -1943,7 +1968,7 @@ MiSetProtectionOnSection(IN PEPROCESS Process,
     //
     // Unlock the working set and update quota charges if needed, then return
     //
-    //MiUnlockProcessWorkingSet(Thread, Process);
+    MiUnlockProcessWorkingSetUnsafe(Process, Thread);
     if ((QuotaCharge > 0) && (!DontCharge))
     {
         FoundVad->u.VadFlags.CommitCharge -= QuotaCharge;

reactos/ntoskrnl/mm/ARM3/virtual.c

@@ -205,7 +205,7 @@ MiMakeSystemAddressValid(IN PVOID PageTableVirtualAddress,
                          IN PEPROCESS CurrentProcess)
 {
     NTSTATUS Status;
-    BOOLEAN WsWasLocked = FALSE, LockChange = FALSE;
+    BOOLEAN WsShared = FALSE, WsSafe = FALSE, LockChange = FALSE;
     PETHREAD CurrentThread = PsGetCurrentThread();
 
     /* Must be a non-pool page table, since those are double-mapped already */
@@ -219,13 +219,11 @@ MiMakeSystemAddressValid(IN PVOID PageTableVirtualAddress,
     /* Check if the page table is valid */
     while (!MmIsAddressValid(PageTableVirtualAddress))
     {
-        /* Check if the WS is locked */
+        /* Release the working set lock */
-        if (CurrentThread->OwnsProcessWorkingSetExclusive)
+        MiUnlockProcessWorkingSetForFault(CurrentProcess,
-        {
+                                          CurrentThread,
-            /* Unlock the working set and remember it was locked */
+                                          WsSafe,
-            MiUnlockProcessWorkingSet(CurrentProcess, CurrentThread);
+                                          WsShared);
-            WsWasLocked = TRUE;
-        }
 
         /* Fault it in */
         Status = MmAccessFault(FALSE, PageTableVirtualAddress, KernelMode, NULL);
@@ -240,7 +238,10 @@ MiMakeSystemAddressValid(IN PVOID PageTableVirtualAddress,
         }
 
         /* Lock the working set again */
-        if (WsWasLocked) MiLockProcessWorkingSet(CurrentProcess, CurrentThread);
+        MiLockProcessWorkingSetForFault(CurrentProcess,
+                                        CurrentThread,
+                                        WsSafe,
+                                        WsShared);
 
         /* This flag will be useful later when we do better locking */
         LockChange = TRUE;
@@ -1884,8 +1885,9 @@ MiProtectVirtualMemory(IN PEPROCESS Process,
     ULONG ProtectionMask, OldProtect;
     BOOLEAN Committed;
     NTSTATUS Status = STATUS_SUCCESS;
+    PETHREAD Thread = PsGetCurrentThread();
 
-    /* Calcualte base address for the VAD */
+    /* Calculate base address for the VAD */
     StartingAddress = (ULONG_PTR)PAGE_ALIGN((*BaseAddress));
     EndingAddress = (((ULONG_PTR)*BaseAddress + *NumberOfBytesToProtect - 1) | (PAGE_SIZE - 1));
 
@@ -1896,7 +1898,7 @@ MiProtectVirtualMemory(IN PEPROCESS Process,
         DPRINT1("Invalid protection mask\n");
         return STATUS_INVALID_PAGE_PROTECTION;
     }
-    
+
     /* Check for ROS specific memory area */
     MemoryArea = MmLocateMemoryAreaByAddress(&Process->Vm, *BaseAddress);
     if ((MemoryArea) && (MemoryArea->Type == MEMORY_AREA_SECTION_VIEW))
@@ -2006,7 +2008,8 @@ MiProtectVirtualMemory(IN PEPROCESS Process,
             goto FailPath;
         }
 
-        //MiLockProcessWorkingSet(Thread, Process);
+        /* Lock the working set */
+        MiLockProcessWorkingSetUnsafe(Process, Thread);
 
         /* Check if all pages in this range are committed */
         Committed = MiIsEntireRangeCommitted(StartingAddress,
@@ -2018,7 +2021,7 @@ MiProtectVirtualMemory(IN PEPROCESS Process,
             /* Fail */
             DPRINT1("The entire range is not committed\n");
             Status = STATUS_NOT_COMMITTED;
-            //MiUnlockProcessWorkingSet(Thread, Process);
+            MiUnlockProcessWorkingSetUnsafe(Process, Thread);
             goto FailPath;
         }
 
@@ -2105,7 +2108,7 @@ MiProtectVirtualMemory(IN PEPROCESS Process,
         }
 
         /* Unlock the working set */
-        //MiUnlockProcessWorkingSet(Thread, Process);
+        MiUnlockProcessWorkingSetUnsafe(Process, Thread);
     }
 
     /* Unlock the address space */
@@ -2278,7 +2281,7 @@ MiDecommitPages(IN PVOID StartingAddress,
     PointerPde = MiAddressToPde(StartingAddress);
     PointerPte = MiAddressToPte(StartingAddress);
     if (Vad->u.VadFlags.MemCommit) CommitPte = MiAddressToPte(Vad->EndingVpn << PAGE_SHIFT);
-    MiLockWorkingSet(CurrentThread, &Process->Vm);
+    MiLockProcessWorkingSetUnsafe(Process, CurrentThread);
 
     //
     // Make the PDE valid, and now loop through each page's worth of data
@@ -2403,7 +2406,7 @@ MiDecommitPages(IN PVOID StartingAddress,
     // release the working set and return the commit reduction accounting.
     //
     if (PteCount) MiProcessValidPteList(ValidPteList, PteCount);
-    MiUnlockWorkingSet(CurrentThread, &Process->Vm);
+    MiUnlockProcessWorkingSetUnsafe(Process, CurrentThread);
     return CommitReduction;
 }
 
@@ -3992,10 +3995,10 @@ NtAllocateVirtualMemory(IN HANDLE ProcessHandle,
         //
         // Lock the working set and insert the VAD into the process VAD tree
         //
-        MiLockProcessWorkingSet(Process, CurrentThread);
+        MiLockProcessWorkingSetUnsafe(Process, CurrentThread);
         Vad->ControlArea = NULL; // For Memory-Area hack
         MiInsertVad(Vad, Process);
-        MiUnlockProcessWorkingSet(Process, CurrentThread);
+        MiUnlockProcessWorkingSetUnsafe(Process, CurrentThread);
 
         //
         // Update the virtual size of the process, and if this is now the highest
@@ -4254,7 +4257,7 @@ NtAllocateVirtualMemory(IN HANDLE ProcessHandle,
     //
     // Lock the working set while we play with user pages and page tables
     //
-    //MiLockWorkingSet(CurrentThread, AddressSpace);
+    MiLockProcessWorkingSetUnsafe(Process, CurrentThread);
 
     //
     // Make the current page table valid, and then loop each page within it
@@ -4335,7 +4338,7 @@ NtAllocateVirtualMemory(IN HANDLE ProcessHandle,
     // the target process if it was not the current process. Also dereference the
     // target process if this wasn't the case.
     //
-    //MiUnlockProcessWorkingSet(Process, CurrentThread);
+    MiUnlockProcessWorkingSetUnsafe(Process, CurrentThread);
     Status = STATUS_SUCCESS;
 FailPath:
     MmUnlockAddressSpace(AddressSpace);
@@ -4584,7 +4587,7 @@ NtFreeVirtualMemory(IN HANDLE ProcessHandle,
             //
             // Finally lock the working set and remove the VAD from the VAD tree
             //
-            MiLockWorkingSet(CurrentThread, AddressSpace);
+            MiLockProcessWorkingSetUnsafe(Process, CurrentThread);
             ASSERT(Process->VadRoot.NumberGenericTableElements >= 1);
             MiRemoveNode((PMMADDRESS_NODE)Vad, &Process->VadRoot);
         }
@@ -4614,7 +4617,7 @@ NtFreeVirtualMemory(IN HANDLE ProcessHandle,
                     // the code path above when the caller sets a zero region size
                     // and the whole VAD is destroyed
                     //
-                    MiLockWorkingSet(CurrentThread, AddressSpace);
+                    MiLockProcessWorkingSetUnsafe(Process, CurrentThread);
                     ASSERT(Process->VadRoot.NumberGenericTableElements >= 1);
                     MiRemoveNode((PMMADDRESS_NODE)Vad, &Process->VadRoot);
                 }
@@ -4653,7 +4656,7 @@ NtFreeVirtualMemory(IN HANDLE ProcessHandle,
                     // and then change the ending address of the VAD to be a bit
                     // smaller.
                     //
-                    MiLockWorkingSet(CurrentThread, AddressSpace);
+                    MiLockProcessWorkingSetUnsafe(Process, CurrentThread);
                     CommitReduction = MiCalculatePageCommitment(StartingAddress,
                                                                 EndingAddress,
                                                                 Vad,
@@ -4695,7 +4698,7 @@ NtFreeVirtualMemory(IN HANDLE ProcessHandle,
         // around with process pages.
         //
         MiDeleteVirtualAddresses(StartingAddress, EndingAddress, NULL);
-        MiUnlockWorkingSet(CurrentThread, AddressSpace);
+        MiUnlockProcessWorkingSetUnsafe(Process, CurrentThread);
         Status = STATUS_SUCCESS;
 
 FinalPath: