[NTOS:MM] Add ASSERTs for VAD table locking

This is required to satisfy VAD locking rules.

ntoskrnl/include/internal/mm.h

@@ -1690,6 +1690,12 @@ FORCEINLINE
 VOID
 MmLockAddressSpace(PMMSUPPORT AddressSpace)
 {
+    ASSERT(!PsGetCurrentThread()->OwnsProcessWorkingSetExclusive);
+    ASSERT(!PsGetCurrentThread()->OwnsProcessWorkingSetShared);
+    ASSERT(!PsGetCurrentThread()->OwnsSystemWorkingSetExclusive);
+    ASSERT(!PsGetCurrentThread()->OwnsSystemWorkingSetShared);
+    ASSERT(!PsGetCurrentThread()->OwnsSessionWorkingSetExclusive);
+    ASSERT(!PsGetCurrentThread()->OwnsSessionWorkingSetShared);
     KeAcquireGuardedMutex(&CONTAINING_RECORD(AddressSpace, EPROCESS, Vm)->AddressCreationLock);
 }
 

ntoskrnl/mm/ARM3/vadnode.c

@@ -43,6 +43,74 @@ CHAR MmReadWrite[32] =
 
 /* FUNCTIONS ******************************************************************/
 
+extern MM_AVL_TABLE MiRosKernelVadRoot;
+
+#if DBG
+
+static
+VOID
+MiDbgAssertIsLockedForRead(_In_ PMM_AVL_TABLE Table)
+{
+    if (Table == &MmSectionBasedRoot)
+    {
+        /* Need to hold MmSectionBasedMutex */
+        ASSERT(MmSectionBasedMutex.Owner == KeGetCurrentThread());
+    }
+    else if (Table == &MiRosKernelVadRoot)
+    {
+        /* Need to hold either the system working-set lock or
+           the idle process' AddressCreationLock */
+        ASSERT(PsGetCurrentThread()->OwnsSystemWorkingSetExclusive ||
+               PsGetCurrentThread()->OwnsSystemWorkingSetShared ||
+               (PsIdleProcess->AddressCreationLock.Owner == KeGetCurrentThread()));
+    }
+    else
+    {
+        /* Need to hold either the process working-set lock or
+           the current process' AddressCreationLock */
+        PEPROCESS Process = CONTAINING_RECORD(Table, EPROCESS, VadRoot);
+        ASSERT(MI_WS_OWNER(Process) ||
+               (Process->AddressCreationLock.Owner == KeGetCurrentThread()));
+    }
+}
+
+static
+VOID
+MiDbgAssertIsLockedForWrite(_In_ PMM_AVL_TABLE Table)
+{
+    if (Table == &MmSectionBasedRoot)
+    {
+        /* Need to hold MmSectionBasedMutex */
+        ASSERT(MmSectionBasedMutex.Owner == KeGetCurrentThread());
+    }
+    else if (Table == &MiRosKernelVadRoot)
+    {
+        /* Need to hold both the system working-set lock exclusive and
+           the idle process' AddressCreationLock */
+        ASSERT(PsGetCurrentThread()->OwnsSystemWorkingSetExclusive);
+        ASSERT(PsIdleProcess->AddressCreationLock.Owner == KeGetCurrentThread());
+    }
+    else
+    {
+        /* Need to hold both the process working-set lock exclusive and
+           the current process' AddressCreationLock */
+        PEPROCESS Process = CONTAINING_RECORD(Table, EPROCESS, VadRoot);
+        ASSERT(Process == PsGetCurrentProcess());
+        ASSERT(PsGetCurrentThread()->OwnsProcessWorkingSetExclusive);
+        ASSERT(Process->AddressCreationLock.Owner == KeGetCurrentThread());
+    }
+}
+
+#define ASSERT_LOCKED_FOR_READ(Table) MiDbgAssertIsLockedForRead(Table)
+#define ASSERT_LOCKED_FOR_WRITE(Table) MiDbgAssertIsLockedForWrite(Table)
+
+#else // DBG
+
+#define ASSERT_LOCKED_FOR_READ(Table)
+#define ASSERT_LOCKED_FOR_WRITE(Table)
+
+#endif // DBG
+
 PMMVAD
 NTAPI
 MiLocateAddress(IN PVOID VirtualAddress)
@@ -52,6 +120,8 @@ MiLocateAddress(IN PVOID VirtualAddress)
     PMM_AVL_TABLE Table = &PsGetCurrentProcess()->VadRoot;
     TABLE_SEARCH_RESULT SearchResult;
 
+    ASSERT_LOCKED_FOR_READ(Table);
+
     /* Start with the the hint */
     FoundVad = (PMMVAD)Table->NodeHint;
     if (!FoundVad) return NULL;
@@ -69,6 +139,8 @@ MiLocateAddress(IN PVOID VirtualAddress)
     /* We found it, update the hint */
     ASSERT(FoundVad != NULL);
     ASSERT((Vpn >= FoundVad->StartingVpn) && (Vpn <= FoundVad->EndingVpn));
+
+    /* We allow this (atomic) update without exclusive lock, because it's a hint only */
     Table->NodeHint = FoundVad;
     return FoundVad;
 }
@@ -82,6 +154,8 @@ MiCheckForConflictingNode(IN ULONG_PTR StartVpn,
 {
     PMMADDRESS_NODE ParentNode, CurrentNode;
 
+    ASSERT_LOCKED_FOR_READ(Table);
+
     /* If the tree is empty, there is no conflict */
     if (Table->NumberGenericTableElements == 0) return TableEmptyTree;
 
@@ -132,6 +206,8 @@ MiInsertNode(IN PMM_AVL_TABLE Table,
 {
     PMMVAD_LONG Vad;
 
+    ASSERT_LOCKED_FOR_WRITE(Table);
+
     /* Insert it into the tree */
     RtlpInsertAvlTreeNode(Table, NewNode, Parent, Result);
 
@@ -186,6 +262,8 @@ MiInsertVad(IN PMMVAD Vad,
     TABLE_SEARCH_RESULT Result;
     PMMADDRESS_NODE Parent = NULL;
 
+    ASSERT_LOCKED_FOR_WRITE(VadRoot);
+
     /* Validate the VAD and set it as the current hint */
     ASSERT(Vad->EndingVpn >= Vad->StartingVpn);
     VadRoot->NodeHint = Vad;
@@ -348,6 +426,8 @@ MiInsertBasedSection(IN PSECTION Section)
     PMMADDRESS_NODE Parent = NULL;
     ASSERT(Section->Address.EndingVpn >= Section->Address.StartingVpn);
 
+    ASSERT_LOCKED_FOR_WRITE(&MmSectionBasedRoot);
+
     /* Find the parent VAD and where this child should be inserted */
     Result = RtlpFindAvlTableNodeOrParent(&MmSectionBasedRoot, (PVOID)Section->Address.StartingVpn, &Parent);
     ASSERT(Result != TableFoundNode);
@@ -362,6 +442,8 @@ MiRemoveNode(IN PMMADDRESS_NODE Node,
 {
     PMMVAD_LONG Vad;
 
+    ASSERT_LOCKED_FOR_WRITE(Table);
+
     /* Call the AVL code */
     RtlpDeleteAvlTreeNode(Table, Node);
 
@@ -509,6 +591,8 @@ MiFindEmptyAddressRangeInTree(IN SIZE_T Length,
     ULONG_PTR PageCount, AlignmentVpn, LowVpn, HighestVpn;
     ASSERT(Length != 0);
 
+    ASSERT_LOCKED_FOR_READ(Table);
+
     /* Calculate page numbers for the length, alignment, and starting address */
     PageCount = BYTES_TO_PAGES(Length);
     AlignmentVpn = Alignment >> PAGE_SHIFT;
@@ -605,6 +689,8 @@ MiFindEmptyAddressRangeDownTree(IN SIZE_T Length,
     ULONG_PTR LowVpn, HighVpn, AlignmentVpn;
     PFN_NUMBER PageCount;
 
+    ASSERT_LOCKED_FOR_READ(Table);
+
     /* Sanity checks */
     ASSERT(BoundaryAddress);
     ASSERT(BoundaryAddress <= ((ULONG_PTR)MM_HIGHEST_VAD_ADDRESS));
@@ -719,6 +805,8 @@ MiFindEmptyAddressRangeDownBasedTree(IN SIZE_T Length,
     PMMADDRESS_NODE Node, LowestNode;
     ULONG_PTR LowVpn, BestVpn;
 
+    ASSERT_LOCKED_FOR_READ(Table);
+
     /* Sanity checks */
     ASSERT(Table == &MmSectionBasedRoot);
     ASSERT(BoundaryAddress);

ntoskrnl/mm/ARM3/virtual.c

@@ -2213,6 +2213,9 @@ MiProtectVirtualMemory(IN PEPROCESS Process,
     PETHREAD Thread = PsGetCurrentThread();
     TABLE_SEARCH_RESULT Result;
 
+    /* We must be attached */
+    ASSERT(Process == PsGetCurrentProcess());
+
     /* Calculate base address for the VAD */
     StartingAddress = (ULONG_PTR)PAGE_ALIGN((*BaseAddress));
     EndingAddress = (((ULONG_PTR)*BaseAddress + *NumberOfBytesToProtect - 1) | (PAGE_SIZE - 1));
@@ -2225,18 +2228,6 @@ MiProtectVirtualMemory(IN PEPROCESS Process,
         return STATUS_INVALID_PAGE_PROTECTION;
     }
 
-    /* Check for ROS specific memory area */
-    MemoryArea = MmLocateMemoryAreaByAddress(&Process->Vm, *BaseAddress);
-    if ((MemoryArea) && (MemoryArea->Type != MEMORY_AREA_OWNED_BY_ARM3))
-    {
-        /* Evil hack */
-        return MiRosProtectVirtualMemory(Process,
-                                         BaseAddress,
-                                         NumberOfBytesToProtect,
-                                         NewAccessProtection,
-                                         OldAccessProtection);
-    }
-
     /* Lock the address space and make sure the process isn't already dead */
     AddressSpace = MmGetCurrentAddressSpace();
     MmLockAddressSpace(AddressSpace);
@@ -2247,6 +2238,19 @@ MiProtectVirtualMemory(IN PEPROCESS Process,
         goto FailPath;
     }
 
+    /* Check for ROS specific memory area */
+    MemoryArea = MmLocateMemoryAreaByAddress(&Process->Vm, *BaseAddress);
+    if ((MemoryArea) && (MemoryArea->Type != MEMORY_AREA_OWNED_BY_ARM3))
+    {
+        /* Evil hack */
+        MmUnlockAddressSpace(AddressSpace);
+        return MiRosProtectVirtualMemory(Process,
+                                         BaseAddress,
+                                         NumberOfBytesToProtect,
+                                         NewAccessProtection,
+                                         OldAccessProtection);
+    }
+
     /* Get the VAD for this address range, and make sure it exists */
     Result = MiCheckForConflictingNode(StartingAddress >> PAGE_SHIFT,
                                        EndingAddress >> PAGE_SHIFT,

ntoskrnl/mm/i386/page.c

@@ -637,9 +637,6 @@ MmCreateVirtualMappingUnsafeEx(
     PMMPTE PointerPte;
     MMPTE TempPte;
     ULONG_PTR Pte;
-#ifdef _M_AMD64
-    BOOLEAN LockReleased = FALSE;
-#endif
 
     DPRINT("MmCreateVirtualMappingUnsafe(%p, %p, %lu, %x)\n",
            Process, Address, flProtect, Page);
@@ -668,15 +665,6 @@ MmCreateVirtualMappingUnsafeEx(
         if (!MiSynchronizeSystemPde(MiAddressToPde(Address)))
             MiFillSystemPageDirectory(Address, PAGE_SIZE);
 #endif
-
-#ifdef _M_AMD64
-        /* This is a temporary hack, because we can incur a recursive page fault when accessing the PDE */
-        if (PsIdleProcess->AddressCreationLock.Owner == KeGetCurrentThread())
-        {
-            MmUnlockAddressSpace(MmGetKernelAddressSpace());
-            LockReleased = TRUE;
-        }
-#endif
     }
     else
     {
@@ -727,15 +715,6 @@ MmCreateVirtualMappingUnsafeEx(
         MiIncrementPageTableReferences(Address);
         MiUnlockProcessWorkingSetUnsafe(Process, PsGetCurrentThread());
     }
-#ifdef _M_AMD64
-    else
-    {
-        if (LockReleased)
-        {
-            MmLockAddressSpace(MmGetKernelAddressSpace());
-        }
-    }
-#endif
 
     return(STATUS_SUCCESS);
 }

ntoskrnl/mm/marea.c

@@ -300,8 +300,7 @@ MmFreeMemoryArea(
         PEPROCESS CurrentProcess = PsGetCurrentProcess();
         PEPROCESS Process = MmGetAddressSpaceOwner(AddressSpace);
 
-        if (Process != NULL &&
-                Process != CurrentProcess)
+        if ((Process != NULL) && (Process != CurrentProcess))
         {
             KeAttachProcess(&Process->Pcb);
         }
@@ -337,12 +336,6 @@ MmFreeMemoryArea(
             }
         }
 
-        if (Process != NULL &&
-                Process != CurrentProcess)
-        {
-            KeDetachProcess();
-        }
-
         //if (MemoryArea->VadNode.StartingVpn < (ULONG_PTR)MmSystemRangeStart >> PAGE_SHIFT
         if (MemoryArea->Vad)
         {
@@ -357,14 +350,23 @@ MmFreeMemoryArea(
             ASSERT(MemoryArea->VadNode.u.VadFlags.Spare != 0);
             if (((PMMVAD)MemoryArea->Vad)->u.VadFlags.Spare == 1)
             {
+                MiLockProcessWorkingSet(PsGetCurrentProcess(), PsGetCurrentThread());
                 MiRemoveNode((PMMADDRESS_NODE)&MemoryArea->VadNode, &Process->VadRoot);
+                MiUnlockProcessWorkingSet(PsGetCurrentProcess(), PsGetCurrentThread());
             }
 
             MemoryArea->Vad = NULL;
         }
         else
         {
+            MiLockWorkingSet(PsGetCurrentThread(), &MmSystemCacheWs);
             MiRemoveNode((PMMADDRESS_NODE)&MemoryArea->VadNode, &MiRosKernelVadRoot);
+            MiUnlockWorkingSet(PsGetCurrentThread(), &MmSystemCacheWs);
+        }
+
+        if ((Process != NULL) && (Process != CurrentProcess))
+        {
+            KeDetachProcess();
         }
     }
 

ntoskrnl/mm/mmfault.c

@@ -213,6 +213,7 @@ MmAccessFault(IN ULONG FaultCode,
 {
     PMEMORY_AREA MemoryArea = NULL;
     NTSTATUS Status;
+    BOOLEAN IsArm3Fault = FALSE;
 
     /* Cute little hack for ROS */
     if ((ULONG_PTR)Address >= (ULONG_PTR)MmSystemRangeStart)
@@ -227,30 +228,52 @@ MmAccessFault(IN ULONG FaultCode,
 #endif
     }
 
-    /* Handle shared user page, which doesn't have a VAD / MemoryArea */
-    if (PAGE_ALIGN(Address) == (PVOID)MM_SHARED_USER_DATA_VA)
+    /* Handle shared user page / page table, which don't have a VAD / MemoryArea */
+    if ((PAGE_ALIGN(Address) == (PVOID)MM_SHARED_USER_DATA_VA) ||
+        MI_IS_PAGE_TABLE_ADDRESS(Address))
     {
         /* This is an ARM3 fault */
-        DPRINT("ARM3 fault %p\n", MemoryArea);
+        DPRINT("ARM3 fault %p\n", Address);
         return MmArmAccessFault(FaultCode, Address, Mode, TrapInformation);
     }
 
     /* Is there a ReactOS address space yet? */
     if (MmGetKernelAddressSpace())
     {
-        /* Check if this is an ARM3 memory area */
-        MemoryArea = MmLocateMemoryAreaByAddress(MmGetKernelAddressSpace(), Address);
-        if (!(MemoryArea) && (Address <= MM_HIGHEST_USER_ADDRESS))
+        if (Address > MM_HIGHEST_USER_ADDRESS)
+        {
+            /* Check if this is an ARM3 memory area */
+            MiLockWorkingSetShared(PsGetCurrentThread(), &MmSystemCacheWs);
+            MemoryArea = MmLocateMemoryAreaByAddress(MmGetKernelAddressSpace(), Address);
+
+            if ((MemoryArea != NULL) && (MemoryArea->Type == MEMORY_AREA_OWNED_BY_ARM3))
+            {
+                IsArm3Fault = TRUE;
+            }
+
+            MiUnlockWorkingSetShared(PsGetCurrentThread(), &MmSystemCacheWs);
+        }
+        else
         {
             /* Could this be a VAD fault from user-mode? */
+            MiLockProcessWorkingSetShared(PsGetCurrentProcess(), PsGetCurrentThread());
             MemoryArea = MmLocateMemoryAreaByAddress(MmGetCurrentAddressSpace(), Address);
+
+            if ((MemoryArea != NULL) && (MemoryArea->Type == MEMORY_AREA_OWNED_BY_ARM3))
+            {
+                IsArm3Fault = TRUE;
+            }
+
+            MiUnlockProcessWorkingSetShared(PsGetCurrentProcess(), PsGetCurrentThread());
         }
     }
 
     /* Is this an ARM3 memory area, or is there no address space yet? */
-    if (((MemoryArea) && (MemoryArea->Type == MEMORY_AREA_OWNED_BY_ARM3)) ||
-            (!(MemoryArea) && ((ULONG_PTR)Address >= (ULONG_PTR)MmPagedPoolStart)) ||
-            (!MmGetKernelAddressSpace()))
+    if (IsArm3Fault ||
+        ((MemoryArea == NULL) &&
+         ((ULONG_PTR)Address >= (ULONG_PTR)MmPagedPoolStart) &&
+         ((ULONG_PTR)Address < (ULONG_PTR)MmPagedPoolEnd)) ||
+        (!MmGetKernelAddressSpace()))
     {
         /* This is an ARM3 fault */
         DPRINT("ARM3 fault %p\n", MemoryArea);

ntoskrnl/mm/mminit.c

@@ -1,4 +1,4 @@
-/*
+/*
  * PROJECT:         ReactOS Kernel
  * LICENSE:         GPL - See COPYING in the top level directory
  * FILE:            ntoskrnl/mm/mminit.c
@@ -68,6 +68,8 @@ MiInitSystemMemoryAreas(VOID)
     // Create all the static memory areas.
     //
 
+    MmLockAddressSpace(MmGetKernelAddressSpace());
+
 #ifdef _M_AMD64
     // Reserved range FFFF800000000000 - FFFFF68000000000
     MiCreateArm3StaticMemoryArea((PVOID)MI_REAL_SYSTEM_RANGE_START, PTE_BASE - MI_REAL_SYSTEM_RANGE_START, FALSE);
@@ -118,6 +120,8 @@ MiInitSystemMemoryAreas(VOID)
     // KUSER_SHARED_DATA
     MiCreateArm3StaticMemoryArea((PVOID)KI_USER_SHARED_DATA, PAGE_SIZE, FALSE);
 #endif /* _X86_ */
+
+    MmUnlockAddressSpace(MmGetKernelAddressSpace());
 }
 
 CODE_SEG("INIT")

ntoskrnl/mm/section.c

@@ -4008,6 +4008,8 @@ MmMapViewOfSection(IN PVOID SectionObject,
     PMMSUPPORT AddressSpace;
     NTSTATUS Status = STATUS_SUCCESS;
     BOOLEAN NotAtBase = FALSE;
+    BOOLEAN IsAttached = FALSE;
+    KAPC_STATE ApcState;
 
     if (MiIsRosSectionObject(SectionObject) == FALSE)
     {
@@ -4031,6 +4033,12 @@ MmMapViewOfSection(IN PVOID SectionObject,
         return STATUS_INVALID_PAGE_PROTECTION;
     }
 
+    if (PsGetCurrentProcess() != Process)
+    {
+        KeStackAttachProcess(&Process->Pcb, &ApcState);
+        IsAttached = TRUE;
+    }
+
     /* FIXME: We should keep this, but it would break code checking equality */
     Protect &= ~PAGE_NOCACHE;
 
@@ -4097,15 +4105,15 @@ MmMapViewOfSection(IN PVOID SectionObject,
             /* Fail if the user requested a fixed base address. */
             if ((*BaseAddress) != NULL)
             {
-                MmUnlockAddressSpace(AddressSpace);
-                return STATUS_CONFLICTING_ADDRESSES;
+                Status = STATUS_CONFLICTING_ADDRESSES;
+                goto Exit;
             }
             /* Otherwise find a gap to map the image. */
             ImageBase = (ULONG_PTR)MmFindGap(AddressSpace, PAGE_ROUND_UP(ImageSize), MM_VIRTMEM_GRANULARITY, FALSE);
             if (ImageBase == 0)
             {
-                MmUnlockAddressSpace(AddressSpace);
-                return STATUS_CONFLICTING_ADDRESSES;
+                Status = STATUS_CONFLICTING_ADDRESSES;
+                goto Exit;
             }
             /* Remember that we loaded image at a different base address */
             NotAtBase = TRUE;
@@ -4136,8 +4144,7 @@ MmMapViewOfSection(IN PVOID SectionObject,
                     MmUnlockSectionSegment(&SectionSegments[i]);
                 }
 
-                MmUnlockAddressSpace(AddressSpace);
-                return Status;
+                goto Exit;
             }
         }
 
@@ -4160,22 +4167,22 @@ MmMapViewOfSection(IN PVOID SectionObject,
         if ((Protect & (PAGE_READWRITE|PAGE_EXECUTE_READWRITE)) &&
                 !(Section->InitialPageProtection & (PAGE_READWRITE|PAGE_EXECUTE_READWRITE)))
         {
-            MmUnlockAddressSpace(AddressSpace);
-            return STATUS_SECTION_PROTECTION;
+            Status = STATUS_SECTION_PROTECTION;
+            goto Exit;
         }
         /* check for read access */
         if ((Protect & (PAGE_READONLY|PAGE_WRITECOPY|PAGE_EXECUTE_READ|PAGE_EXECUTE_WRITECOPY)) &&
                 !(Section->InitialPageProtection & (PAGE_READONLY|PAGE_READWRITE|PAGE_WRITECOPY|PAGE_EXECUTE_READ|PAGE_EXECUTE_READWRITE|PAGE_EXECUTE_WRITECOPY)))
         {
-            MmUnlockAddressSpace(AddressSpace);
-            return STATUS_SECTION_PROTECTION;
+            Status = STATUS_SECTION_PROTECTION;
+            goto Exit;
         }
         /* check for execute access */
         if ((Protect & (PAGE_EXECUTE|PAGE_EXECUTE_READ|PAGE_EXECUTE_READWRITE|PAGE_EXECUTE_WRITECOPY)) &&
                 !(Section->InitialPageProtection & (PAGE_EXECUTE|PAGE_EXECUTE_READ|PAGE_EXECUTE_READWRITE|PAGE_EXECUTE_WRITECOPY)))
         {
-            MmUnlockAddressSpace(AddressSpace);
-            return STATUS_SECTION_PROTECTION;
+            Status = STATUS_SECTION_PROTECTION;
+            goto Exit;
         }
 
         if (SectionOffset == NULL)
@@ -4189,8 +4196,8 @@ MmMapViewOfSection(IN PVOID SectionObject,
 
         if ((ViewOffset % PAGE_SIZE) != 0)
         {
-            MmUnlockAddressSpace(AddressSpace);
-            return STATUS_MAPPED_ALIGNMENT;
+            Status = STATUS_MAPPED_ALIGNMENT;
+            goto Exit;
         }
 
         if ((*ViewSize) == 0)
@@ -4219,18 +4226,24 @@ MmMapViewOfSection(IN PVOID SectionObject,
         MmUnlockSectionSegment(Segment);
         if (!NT_SUCCESS(Status))
         {
-            MmUnlockAddressSpace(AddressSpace);
-            return Status;
+            goto Exit;
         }
     }
 
-    MmUnlockAddressSpace(AddressSpace);
-
     if (NotAtBase)
         Status = STATUS_IMAGE_NOT_AT_BASE;
     else
         Status = STATUS_SUCCESS;
 
+Exit:
+
+    MmUnlockAddressSpace(AddressSpace);
+
+    if (IsAttached)
+    {
+        KeUnstackDetachProcess(&ApcState);
+    }
+
     return Status;
 }