[NTOSKRNL]
[reactos.git] / reactos / ntoskrnl / mm / ARM3 / procsup.c
index 3afd8c4..ae5d496 100644 (file)
 #define NDEBUG
 #include <debug.h>
 
-#line 15 "ARMĀ³::PROCSUP"
 #define MODULE_INVOLVED_IN_ARM3
 #include "../ARM3/miarm.h"
 
 /* GLOBALS ********************************************************************/
 
 ULONG MmProcessColorSeed = 0x12345678;
+PMMWSL MmWorkingSetList;
+ULONG MmMaximumDeadKernelStacks = 5;
+SLIST_HEADER MmDeadStackSListHead;
 
 /* PRIVATE FUNCTIONS **********************************************************/
 
 VOID
 NTAPI
-MiRosTakeOverPebTebRanges(IN PEPROCESS Process)
+MiRosTakeOverSharedUserPage(IN PEPROCESS Process)
 {
     NTSTATUS Status;
     PMEMORY_AREA MemoryArea;
     PHYSICAL_ADDRESS BoundaryAddressMultiple;
-    PVOID AllocatedBase = (PVOID)USER_SHARED_DATA;
+    PVOID AllocatedBase = (PVOID)MM_SHARED_USER_DATA_VA;
     BoundaryAddressMultiple.QuadPart = 0;
 
     Status = MmCreateMemoryArea(&Process->Vm,
                                 MEMORY_AREA_OWNED_BY_ARM3,
                                 &AllocatedBase,
-                                ((ULONG_PTR)MM_HIGHEST_USER_ADDRESS - 1) -
-                                (ULONG_PTR)USER_SHARED_DATA,
+                                PAGE_SIZE,
                                 PAGE_READWRITE,
                                 &MemoryArea,
                                 TRUE,
@@ -59,11 +60,11 @@ MiCreatePebOrTeb(IN PEPROCESS Process,
     LARGE_INTEGER CurrentTime;
     TABLE_SEARCH_RESULT Result = TableFoundNode;
     PMMADDRESS_NODE Parent;
-    
+
     /* Allocate a VAD */
     Vad = ExAllocatePoolWithTag(NonPagedPool, sizeof(MMVAD_LONG), 'ldaV');
     if (!Vad) return STATUS_NO_MEMORY;
-    
+
     /* Setup the primary flags with the size, and make it commited, private, RW */
     Vad->u.LongFlags = 0;
     Vad->u.VadFlags.CommitCharge = BYTES_TO_PAGES(Size);
@@ -71,13 +72,13 @@ MiCreatePebOrTeb(IN PEPROCESS Process,
     Vad->u.VadFlags.PrivateMemory = TRUE;
     Vad->u.VadFlags.Protection = MM_READWRITE;
     Vad->u.VadFlags.NoChange = TRUE;
-    
+
     /* Setup the secondary flags to make it a secured, writable, long VAD */
     Vad->u2.LongFlags2 = 0;
     Vad->u2.VadFlags2.OneSecured = TRUE;
     Vad->u2.VadFlags2.LongVad = TRUE;
     Vad->u2.VadFlags2.ReadOnly = FALSE;
-    
+
     /* Lock the process address space */
     KeAcquireGuardedMutex(&Process->AddressCreationLock);
 
@@ -117,25 +118,29 @@ MiCreatePebOrTeb(IN PEPROCESS Process,
                                                  Base,
                                                  &Parent);
         /* Bail out, if still nothing free was found */
-        if (Result == TableFoundNode) return STATUS_NO_MEMORY;
+        if (Result == TableFoundNode)
+        {
+            ExFreePoolWithTag(Vad, 'ldaV');
+            return STATUS_NO_MEMORY;
+        }
     }
-    
+
     /* Validate that it came from the VAD ranges */
     ASSERT(*Base >= (ULONG_PTR)MI_LOWEST_VAD_ADDRESS);
-    
+
     /* Build the rest of the VAD now */
     Vad->StartingVpn = (*Base) >> PAGE_SHIFT;
     Vad->EndingVpn = ((*Base) + Size - 1) >> PAGE_SHIFT;
     Vad->u3.Secured.StartVpn = *Base;
     Vad->u3.Secured.EndVpn = (Vad->EndingVpn << PAGE_SHIFT) | (PAGE_SIZE - 1);
     Vad->u1.Parent = NULL;
-    
+
     /* FIXME: Should setup VAD bitmap */
     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);
     Process->VadRoot.NodeHint = Vad;
@@ -146,7 +151,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);
@@ -165,20 +170,20 @@ MmDeleteTeb(IN PEPROCESS Process,
     PMMVAD Vad;
     PMM_AVL_TABLE VadTree = &Process->VadRoot;
     DPRINT("Deleting TEB: %p in %16s\n", Teb, Process->ImageFileName);
-    
+
     /* TEB is one page */
     TebEnd = (ULONG_PTR)Teb + ROUND_TO_PAGES(sizeof(TEB)) - 1;
-    
+
     /* Attach to the process */
     KeAttachProcess(&Process->Pcb);
-    
+
     /* Lock the process address space */
     KeAcquireGuardedMutex(&Process->AddressCreationLock);
-    
+
     /* Find the VAD, make sure it's a TEB VAD */
     Vad = MiLocateAddress(Teb);
     DPRINT("Removing node for VAD: %lx %lx\n", Vad->StartingVpn, Vad->EndingVpn);
-    ASSERT(Vad != NULL);    
+    ASSERT(Vad != NULL);
     if (Vad->StartingVpn != ((ULONG_PTR)Teb >> PAGE_SHIFT))
     {
         /* Bug in the AVL code? */
@@ -190,25 +195,29 @@ MmDeleteTeb(IN PEPROCESS Process,
         ASSERT((Vad->StartingVpn == ((ULONG_PTR)Teb >> PAGE_SHIFT) &&
                (Vad->EndingVpn == (TebEnd >> PAGE_SHIFT))));
         ASSERT(Vad->u.VadFlags.NoChange == TRUE);
+        ASSERT(Vad->u2.VadFlags2.OneSecured == TRUE);
         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);
         MiRemoveNode((PMMADDRESS_NODE)Vad, VadTree);
-    
+
+        /* Delete the pages */
+        MiDeleteVirtualAddresses((ULONG_PTR)Teb, TebEnd, NULL);
+
         /* Release the working set */
-        MiUnlockProcessWorkingSet(Process, Thread);
-    
+        MiUnlockProcessWorkingSetUnsafe(Process, Thread);
+
         /* Remove the VAD */
         ExFreePool(Vad);
     }
 
     /* Release the address space lock */
     KeReleaseGuardedMutex(&Process->AddressCreationLock);
-    
+
     /* Detach */
     KeDetachProcess();
 }
@@ -219,26 +228,41 @@ MmDeleteKernelStack(IN PVOID StackBase,
                     IN BOOLEAN GuiStack)
 {
     PMMPTE PointerPte;
-    PFN_NUMBER StackPages, PageFrameNumber;//, PageTableFrameNumber;
-    PMMPFN Pfn1;//, Pfn2;
+    PFN_NUMBER PageFrameNumber, PageTableFrameNumber;
+    PFN_COUNT StackPages;
+    PMMPFN Pfn1, Pfn2;
     ULONG i;
     KIRQL OldIrql;
-    
+
     //
     // This should be the guard page, so decrement by one
     //
     PointerPte = MiAddressToPte(StackBase);
     PointerPte--;
-    
+
+    //
+    // If this is a small stack, just push the stack onto the dead stack S-LIST
+    //
+    if (!GuiStack)
+    {
+        if (ExQueryDepthSList(&MmDeadStackSListHead) < MmMaximumDeadKernelStacks)
+        {
+            Pfn1 = MiGetPfnEntry(PointerPte->u.Hard.PageFrameNumber);
+            InterlockedPushEntrySList(&MmDeadStackSListHead,
+                                      (PSLIST_ENTRY)&Pfn1->u1.NextStackPfn);
+            return;
+        }
+    }
+
     //
     // Calculate pages used
     //
     StackPages = BYTES_TO_PAGES(GuiStack ?
                                 KERNEL_LARGE_STACK_SIZE : KERNEL_STACK_SIZE);
-    
+
     /* Acquire the PFN lock */
     OldIrql = KeAcquireQueuedSpinLock(LockQueuePfnLock);
-                            
+
     //
     // Loop them
     //
@@ -252,32 +276,32 @@ MmDeleteKernelStack(IN PVOID StackBase,
             /* Get the PTE's page */
             PageFrameNumber = PFN_FROM_PTE(PointerPte);
             Pfn1 = MiGetPfnEntry(PageFrameNumber);
-#if 0 // ARM3 might not own the page table, so don't take this risk. Leak it instead!
+
             /* Now get the page of the page table mapping it */
             PageTableFrameNumber = Pfn1->u4.PteFrame;
             Pfn2 = MiGetPfnEntry(PageTableFrameNumber);
-            
+
             /* Remove a shared reference, since the page is going away */
             MiDecrementShareCount(Pfn2, PageTableFrameNumber);
-#endif
+
             /* Set the special pending delete marker */
             MI_SET_PFN_DELETED(Pfn1);
-            
+
             /* And now delete the actual stack page */
             MiDecrementShareCount(Pfn1, PageFrameNumber);
         }
-        
+
         //
         // Next one
         //
         PointerPte--;
     }
-    
+
     //
     // We should be at the guard page now
     //
     ASSERT(PointerPte->u.Hard.Valid == 0);
-    
+
     /* Release the PFN lock */
     KeReleaseQueuedSpinLock(LockQueuePfnLock, OldIrql);
 
@@ -292,14 +316,15 @@ NTAPI
 MmCreateKernelStack(IN BOOLEAN GuiStack,
                     IN UCHAR Node)
 {
-    PFN_NUMBER StackPtes, StackPages;
+    PFN_COUNT StackPtes, StackPages;
     PMMPTE PointerPte, StackPte;
     PVOID BaseAddress;
     MMPTE TempPte, InvalidPte;
     KIRQL OldIrql;
     PFN_NUMBER PageFrameIndex;
     ULONG i;
-    
+    PMMPFN Pfn1;
+
     //
     // Calculate pages needed
     //
@@ -310,47 +335,62 @@ MmCreateKernelStack(IN BOOLEAN GuiStack,
         //
         StackPtes = BYTES_TO_PAGES(KERNEL_LARGE_STACK_SIZE);
         StackPages = BYTES_TO_PAGES(KERNEL_LARGE_STACK_COMMIT);
-        
+
     }
     else
     {
+        //
+        // If the dead stack S-LIST has a stack on it, use it instead of allocating
+        // new system PTEs for this stack
+        //
+        if (ExQueryDepthSList(&MmDeadStackSListHead))
+        {
+            Pfn1 = (PMMPFN)InterlockedPopEntrySList(&MmDeadStackSListHead);
+            if (Pfn1)
+            {
+                PointerPte = Pfn1->PteAddress;
+                BaseAddress = MiPteToAddress(++PointerPte);
+                return BaseAddress;
+            }
+        }
+
         //
         // We'll allocate 12K and that's it
         //
         StackPtes = BYTES_TO_PAGES(KERNEL_STACK_SIZE);
         StackPages = StackPtes;
     }
-    
+
     //
     // Reserve stack pages, plus a guard page
     //
     StackPte = MiReserveSystemPtes(StackPtes + 1, SystemPteSpace);
     if (!StackPte) return NULL;
-    
+
     //
     // Get the stack address
     //
     BaseAddress = MiPteToAddress(StackPte + StackPtes + 1);
-    
+
     //
     // Select the right PTE address where we actually start committing pages
     //
     PointerPte = StackPte;
     if (GuiStack) PointerPte += BYTES_TO_PAGES(KERNEL_LARGE_STACK_SIZE -
                                                KERNEL_LARGE_STACK_COMMIT);
-    
+
 
     /* Setup the temporary invalid PTE */
     MI_MAKE_SOFTWARE_PTE(&InvalidPte, MM_NOACCESS);
 
     /* Setup the template stack PTE */
     MI_MAKE_HARDWARE_PTE_KERNEL(&TempPte, PointerPte + 1, MM_READWRITE, 0);
-    
+
     //
     // Acquire the PFN DB lock
     //
     OldIrql = KeAcquireQueuedSpinLock(LockQueuePfnLock);
-    
+
     //
     // Loop each stack page
     //
@@ -360,14 +400,16 @@ MmCreateKernelStack(IN BOOLEAN GuiStack,
         // Next PTE
         //
         PointerPte++;
-        
+
         /* Get a page and write the current invalid PTE */
+        MI_SET_USAGE(MI_USAGE_KERNEL_STACK);
+        MI_SET_PROCESS2(PsGetCurrentProcess()->ImageFileName);
         PageFrameIndex = MiRemoveAnyPage(MI_GET_NEXT_COLOR());
         MI_WRITE_INVALID_PTE(PointerPte, InvalidPte);
 
         /* Initialize the PFN entry for this page */
         MiInitializePfn(PageFrameIndex, PointerPte, 1);
-        
+
         /* Write the valid PTE */
         TempPte.u.Hard.PageFrameNumber = PageFrameIndex;
         MI_WRITE_VALID_PTE(PointerPte, TempPte);
@@ -377,7 +419,7 @@ MmCreateKernelStack(IN BOOLEAN GuiStack,
     // Release the PFN lock
     //
     KeReleaseQueuedSpinLock(LockQueuePfnLock, OldIrql);
-    
+
     //
     // Return the stack address
     //
@@ -394,25 +436,25 @@ MmGrowKernelStackEx(IN PVOID StackPointer,
     KIRQL OldIrql;
     MMPTE TempPte, InvalidPte;
     PFN_NUMBER PageFrameIndex;
-    
+
     //
     // Make sure the stack did not overflow
     //
     ASSERT(((ULONG_PTR)Thread->StackBase - (ULONG_PTR)Thread->StackLimit) <=
            (KERNEL_LARGE_STACK_SIZE + PAGE_SIZE));
-    
+
     //
     // Get the current stack limit
     //
     LimitPte = MiAddressToPte(Thread->StackLimit);
     ASSERT(LimitPte->u.Hard.Valid == 1);
-    
+
     //
     // Get the new one and make sure this isn't a retarded request
     //
     NewLimitPte = MiAddressToPte((PVOID)((ULONG_PTR)StackPointer - GrowSize));
     if (NewLimitPte == LimitPte) return STATUS_SUCCESS;
-    
+
     //
     // Now make sure you're not going past the reserved space
     //
@@ -426,15 +468,15 @@ MmGrowKernelStackEx(IN PVOID StackPointer,
         DPRINT1("Thread wants too much stack\n");
         return STATUS_STACK_OVERFLOW;
     }
-    
+
     //
     // Calculate the number of new pages
     //
     LimitPte--;
-    
+
     /* Setup the temporary invalid PTE */
     MI_MAKE_SOFTWARE_PTE(&InvalidPte, MM_NOACCESS);
-    
+
     //
     // Acquire the PFN DB lock
     //
@@ -446,24 +488,26 @@ MmGrowKernelStackEx(IN PVOID StackPointer,
     while (LimitPte >= NewLimitPte)
     {
         /* Get a page and write the current invalid PTE */
+        MI_SET_USAGE(MI_USAGE_KERNEL_STACK_EXPANSION);
+        MI_SET_PROCESS2(PsGetCurrentProcess()->ImageFileName);
         PageFrameIndex = MiRemoveAnyPage(MI_GET_NEXT_COLOR());
         MI_WRITE_INVALID_PTE(LimitPte, InvalidPte);
 
         /* Initialize the PFN entry for this page */
         MiInitializePfn(PageFrameIndex, LimitPte, 1);
-        
+
         /* Setup the template stack PTE */
         MI_MAKE_HARDWARE_PTE_KERNEL(&TempPte, LimitPte, MM_READWRITE, PageFrameIndex);
-        
+
         /* Write the valid PTE */
         MI_WRITE_VALID_PTE(LimitPte--, TempPte);
     }
-    
+
     //
     // Release the PFN lock
     //
     KeReleaseQueuedSpinLock(LockQueuePfnLock, OldIrql);
-    
+
     //
     // Set the new limit
     //
@@ -487,7 +531,7 @@ MmSetMemoryPriorityProcess(IN PEPROCESS Process,
                            IN UCHAR MemoryPriority)
 {
     UCHAR OldPriority;
-    
+
     //
     // Check if we have less then 16MB of Physical Memory
     //
@@ -499,13 +543,13 @@ MmSetMemoryPriorityProcess(IN PEPROCESS Process,
         //
         MemoryPriority = MEMORY_PRIORITY_BACKGROUND;
     }
-    
+
     //
     // Save the old priority and update it
     //
     OldPriority = (UCHAR)Process->Vm.Flags.MemoryPriority;
     Process->Vm.Flags.MemoryPriority = MemoryPriority;
-    
+
     //
     // Return the old priority
     //
@@ -518,12 +562,12 @@ MmGetSessionLocaleId(VOID)
 {
     PEPROCESS Process;
     PAGED_CODE();
-    
+
     //
     // Get the current process
     //
     Process = PsGetCurrentProcess();
-    
+
     //
     // Check if it's the Session Leader
     //
@@ -537,12 +581,10 @@ MmGetSessionLocaleId(VOID)
             //
             // Get the Locale ID
             //
-#if ROS_HAS_SESSIONS
             return ((PMM_SESSION_SPACE)Process->Session)->LocaleId;
-#endif
         }
     }
-    
+
     //
     // Not a session leader, return the default
     //
@@ -566,17 +608,11 @@ MmCreatePeb(IN PEPROCESS Process,
     KAFFINITY ProcessAffinityMask = 0;
     SectionOffset.QuadPart = (ULONGLONG)0;
     *BasePeb = NULL;
-    
+
     //
     // Attach to Process
     //
     KeAttachProcess(&Process->Pcb);
-       
-    //
-    // Allocate the PEB
-    //
-    Status = MiCreatePebOrTeb(Process, sizeof(PEB), (PULONG_PTR)&Peb);
-    ASSERT(NT_SUCCESS(Status));
 
     //
     // Map NLS Tables
@@ -591,8 +627,26 @@ MmCreatePeb(IN PEPROCESS Process,
                                 ViewShare,
                                 MEM_TOP_DOWN,
                                 PAGE_READONLY);
-    if (!NT_SUCCESS(Status)) return Status;
-    
+    DPRINT("NLS Tables at: %p\n", TableBase);
+    if (!NT_SUCCESS(Status))
+    {
+        /* Cleanup and exit */
+        KeDetachProcess();
+        return Status;
+    }
+
+    //
+    // Allocate the PEB
+    //
+    Status = MiCreatePebOrTeb(Process, sizeof(PEB), (PULONG_PTR)&Peb);
+    DPRINT("PEB at: %p\n", Peb);
+    if (!NT_SUCCESS(Status))
+    {
+        /* Cleanup and exit */
+        KeDetachProcess();
+        return Status;
+    }
+
     //
     // Use SEH in case we can't load the PEB
     //
@@ -602,7 +656,7 @@ MmCreatePeb(IN PEPROCESS Process,
         // Initialize the PEB
         //
         RtlZeroMemory(Peb, sizeof(PEB));
-        
+
         //
         // Set up data
         //
@@ -610,14 +664,14 @@ MmCreatePeb(IN PEPROCESS Process,
         Peb->InheritedAddressSpace = InitialPeb->InheritedAddressSpace;
         Peb->Mutant = InitialPeb->Mutant;
         Peb->ImageUsesLargePages = InitialPeb->ImageUsesLargePages;
-        
+
         //
         // NLS
         //
         Peb->AnsiCodePageData = (PCHAR)TableBase + ExpAnsiCodePageDataOffset;
         Peb->OemCodePageData = (PCHAR)TableBase + ExpOemCodePageDataOffset;
         Peb->UnicodeCaseTableData = (PCHAR)TableBase + ExpUnicodeCaseTableDataOffset;
-        
+
         //
         // Default Version Data (could get changed below)
         //
@@ -626,12 +680,12 @@ MmCreatePeb(IN PEPROCESS Process,
         Peb->OSBuildNumber = (USHORT)(NtBuildNumber & 0x3FFF);
         Peb->OSPlatformId = 2; /* VER_PLATFORM_WIN32_NT */
         Peb->OSCSDVersion = (USHORT)CmNtCSDVersion;
-        
+
         //
         // Heap and Debug Data
         //
         Peb->NumberOfProcessors = KeNumberProcessors;
-        Peb->BeingDebugged = (BOOLEAN)(Process->DebugPort != NULL ? TRUE : FALSE);
+        Peb->BeingDebugged = (BOOLEAN)(Process->DebugPort != NULL);
         Peb->NtGlobalFlag = NtGlobalFlag;
         /*Peb->HeapSegmentReserve = MmHeapSegmentReserve;
          Peb->HeapSegmentCommit = MmHeapSegmentCommit;
@@ -642,11 +696,11 @@ MmCreatePeb(IN PEPROCESS Process,
          */
         Peb->MaximumNumberOfHeaps = (PAGE_SIZE - sizeof(PEB)) / sizeof(PVOID);
         Peb->ProcessHeaps = (PVOID*)(Peb + 1);
-        
+
         //
         // Session ID
         //
-        if (Process->Session) Peb->SessionId = 0; // MmGetSessionId(Process);
+        MmGetSessionId(Process);
     }
     _SEH2_EXCEPT(EXCEPTION_EXECUTE_HANDLER)
     {
@@ -657,7 +711,7 @@ MmCreatePeb(IN PEPROCESS Process,
         _SEH2_YIELD(return _SEH2_GetExceptionCode());
     }
     _SEH2_END;
-    
+
     //
     // Use SEH in case we can't load the image
     //
@@ -678,7 +732,7 @@ MmCreatePeb(IN PEPROCESS Process,
         _SEH2_YIELD(return STATUS_INVALID_IMAGE_PROTECT);
     }
     _SEH2_END;
-    
+
     //
     // Parse the headers
     //
@@ -705,7 +759,7 @@ MmCreatePeb(IN PEPROCESS Process,
                              sizeof(IMAGE_LOAD_CONFIG_DIRECTORY),
                              sizeof(ULONG));
             }
-            
+
             //
             // Write subsystem data
             //
@@ -725,36 +779,22 @@ MmCreatePeb(IN PEPROCESS Process,
                 Peb->OSMinorVersion = (NtHeaders->OptionalHeader.Win32VersionValue >> 8) & 0xFF;
                 Peb->OSBuildNumber = (NtHeaders->OptionalHeader.Win32VersionValue >> 16) & 0x3FFF;
                 Peb->OSPlatformId = (NtHeaders->OptionalHeader.Win32VersionValue >> 30) ^ 2;
-            }
-            
-            //
-            // Process the image config data overrides if specfied
-            //
-            if (ImageConfigData != NULL)
-            {
-                //
-                // Process CSD version override
-                //
-                if (ImageConfigData->CSDVersion)
+
+                /* Process CSD version override */
+                if ((ImageConfigData) && (ImageConfigData->CSDVersion))
                 {
-                    //
-                    // Set new data
-                    //
+                    /* Take the value from the image configuration directory */
                     Peb->OSCSDVersion = ImageConfigData->CSDVersion;
                 }
-                
-                //
-                // Process affinity mask ovverride
-                //
-                if (ImageConfigData->ProcessAffinityMask)
-                {
-                    //
-                    // Set new data
-                    //
-                    ProcessAffinityMask = ImageConfigData->ProcessAffinityMask;
-                }
             }
-            
+
+            /* Process optional process affinity mask override */
+            if ((ImageConfigData) && (ImageConfigData->ProcessAffinityMask))
+            {
+                /* Take the value from the image configuration directory */
+                ProcessAffinityMask = ImageConfigData->ProcessAffinityMask;
+            }
+
             //
             // Check if this is a UP image
             if (Characteristics & IMAGE_FILE_UP_SYSTEM_ONLY)
@@ -762,6 +802,7 @@ MmCreatePeb(IN PEPROCESS Process,
                 //
                 // Force it to use CPU 0
                 //
+                /* FIXME: this should use the MmRotatingUniprocessorNumber */
                 Peb->ImageProcessAffinityMask = 0;
             }
             else
@@ -782,7 +823,7 @@ MmCreatePeb(IN PEPROCESS Process,
         }
         _SEH2_END;
     }
-    
+
     //
     // Detach from the Process
     //
@@ -801,18 +842,18 @@ MmCreateTeb(IN PEPROCESS Process,
     PTEB Teb;
     NTSTATUS Status = STATUS_SUCCESS;
     *BaseTeb = NULL;
-    
+
     //
     // Attach to Target
     //
     KeAttachProcess(&Process->Pcb);
-    
+
     //
     // Allocate the TEB
     //
     Status = MiCreatePebOrTeb(Process, sizeof(TEB), (PULONG_PTR)&Teb);
     ASSERT(NT_SUCCESS(Status));
-    
+
     //
     // Use SEH in case we can't load the TEB
     //
@@ -822,18 +863,22 @@ MmCreateTeb(IN PEPROCESS Process,
         // Initialize the PEB
         //
         RtlZeroMemory(Teb, sizeof(TEB));
-        
+
         //
         // Set TIB Data
         //
+#ifdef _M_AMD64
+        Teb->NtTib.ExceptionList = NULL;
+#else
         Teb->NtTib.ExceptionList = EXCEPTION_CHAIN_END;
+#endif
         Teb->NtTib.Self = (PNT_TIB)Teb;
-        
+
         //
         // Identify this as an OS/2 V3.0 ("Cruiser") TIB
         //
         Teb->NtTib.Version = 30 << 8;
-        
+
         //
         // Set TEB Data
         //
@@ -841,7 +886,7 @@ MmCreateTeb(IN PEPROCESS Process,
         Teb->RealClientId = *ClientId;
         Teb->ProcessEnvironmentBlock = Process->Peb;
         Teb->CurrentLocale = PsDefaultThreadLocaleId;
-        
+
         //
         // Check if we have a grandparent TEB
         //
@@ -887,6 +932,40 @@ MmCreateTeb(IN PEPROCESS Process,
     return Status;
 }
 
+VOID
+NTAPI
+MiInitializeWorkingSetList(IN PEPROCESS CurrentProcess)
+{
+    PMMPFN Pfn1;
+    PMMPTE sysPte;
+    MMPTE tempPte;
+
+    /* Setup some bogus list data */
+    MmWorkingSetList->LastEntry = CurrentProcess->Vm.MinimumWorkingSetSize;
+    MmWorkingSetList->HashTable = NULL;
+    MmWorkingSetList->HashTableSize = 0;
+    MmWorkingSetList->NumberOfImageWaiters = 0;
+    MmWorkingSetList->Wsle = (PVOID)0xDEADBABE;
+    MmWorkingSetList->VadBitMapHint = 1;
+    MmWorkingSetList->HashTableStart = (PVOID)0xBADAB00B;
+    MmWorkingSetList->HighestPermittedHashAddress = (PVOID)0xCAFEBABE;
+    MmWorkingSetList->FirstFree = 1;
+    MmWorkingSetList->FirstDynamic = 2;
+    MmWorkingSetList->NextSlot = 3;
+    MmWorkingSetList->LastInitializedWsle = 4;
+
+    /* The rule is that the owner process is always in the FLINK of the PDE's PFN entry */
+    Pfn1 = MiGetPfnEntry(CurrentProcess->Pcb.DirectoryTableBase[0] >> PAGE_SHIFT);
+    ASSERT(Pfn1->u4.PteFrame == MiGetPfnEntryIndex(Pfn1));
+    Pfn1->u1.Event = (PKEVENT)CurrentProcess;
+
+    /* Map the process working set in kernel space */
+    sysPte = MiReserveSystemPtes(1, SystemPteSpace);
+    MI_MAKE_HARDWARE_PTE_KERNEL(&tempPte, sysPte, MM_READWRITE, CurrentProcess->WorkingSetPage);
+    MI_WRITE_VALID_PTE(sysPte, tempPte);
+    CurrentProcess->Vm.VmWorkingSetList = MiPteToAddress(sysPte);
+}
+
 NTSTATUS
 NTAPI
 MmInitializeProcessAddressSpace(IN PEPROCESS Process,
@@ -907,14 +986,15 @@ MmInitializeProcessAddressSpace(IN PEPROCESS Process,
     PWCHAR Source;
     PCHAR Destination;
     USHORT Length = 0;
-    
+    MMPTE TempPte;
+
     /* We should have a PDE */
     ASSERT(Process->Pcb.DirectoryTableBase[0] != 0);
     ASSERT(Process->PdeUpdateNeeded == FALSE);
 
     /* Attach to the process */
     KeAttachProcess(&Process->Pcb);
-    
+
     /* The address space should now been in phase 1 or 0 */
     ASSERT(Process->AddressSpaceInitialized <= 1);
     Process->AddressSpaceInitialized = 2;
@@ -931,20 +1011,46 @@ MmInitializeProcessAddressSpace(IN PEPROCESS Process,
     OldIrql = KeAcquireQueuedSpinLock(LockQueuePfnLock);
 
     /* Setup the PFN for the PDE base of this process */
+#ifdef _M_AMD64
+    PointerPte = MiAddressToPte(PXE_BASE);
+#else
     PointerPte = MiAddressToPte(PDE_BASE);
+#endif
     PageFrameNumber = PFN_FROM_PTE(PointerPte);
+    ASSERT(Process->Pcb.DirectoryTableBase[0] == PageFrameNumber * PAGE_SIZE);
     MiInitializePfn(PageFrameNumber, PointerPte, TRUE);
 
     /* Do the same for hyperspace */
+#ifdef _M_AMD64
+    PointerPde = MiAddressToPxe((PVOID)HYPER_SPACE);
+#else
     PointerPde = MiAddressToPde(HYPER_SPACE);
+#endif
     PageFrameNumber = PFN_FROM_PTE(PointerPde);
-    MiInitializePfn(PageFrameNumber, PointerPde, TRUE);
+    //ASSERT(Process->Pcb.DirectoryTableBase[0] == PageFrameNumber * PAGE_SIZE); // we're not lucky
+    MiInitializePfn(PageFrameNumber, (PMMPTE)PointerPde, TRUE);
+
+    /* Setup the PFN for the PTE for the working set */
+    PointerPte = MiAddressToPte(MI_WORKING_SET_LIST);
+    MI_MAKE_HARDWARE_PTE(&TempPte, PointerPte, MM_READWRITE, 0);
+    ASSERT(PointerPte->u.Long != 0);
+    PageFrameNumber = PFN_FROM_PTE(PointerPte);
+    MI_WRITE_INVALID_PTE(PointerPte, DemandZeroPte);
+    MiInitializePfn(PageFrameNumber, PointerPte, TRUE);
+    TempPte.u.Hard.PageFrameNumber = PageFrameNumber;
+    MI_WRITE_VALID_PTE(PointerPte, TempPte);
+
+    /* Now initialize the working set list */
+    MiInitializeWorkingSetList(Process);
+
+    /* Sanity check */
+    ASSERT(Process->PhysicalVadRoot == NULL);
 
     /* Release PFN lock */
     KeReleaseQueuedSpinLock(LockQueuePfnLock, OldIrql);
 
     /* Lock the VAD, ARM3-owned ranges away */
-    MiRosTakeOverPebTebRanges(Process);
+    MiRosTakeOverSharedUserPage(Process);
 
     /* Check if there's a Section Object */
     if (SectionObject)
@@ -1008,7 +1114,7 @@ MmInitializeProcessAddressSpace(IN PEPROCESS Process,
         /* Save the pointer */
         Process->SectionBaseAddress = ImageBase;
     }
-    
+
     /* Be nice and detach */
     KeDetachProcess();
 
@@ -1018,6 +1124,7 @@ MmInitializeProcessAddressSpace(IN PEPROCESS Process,
 
 NTSTATUS
 NTAPI
+INIT_FUNCTION
 MmInitializeHandBuiltProcess(IN PEPROCESS Process,
                              IN PULONG_PTR DirectoryTableBase)
 {
@@ -1032,6 +1139,9 @@ MmInitializeHandBuiltProcess(IN PEPROCESS Process,
     ASSERT(Process->VadRoot.NumberGenericTableElements == 0);
     Process->VadRoot.BalancedRoot.u1.Parent = &Process->VadRoot.BalancedRoot;
 
+    /* Use idle process Working set */
+    Process->Vm.VmWorkingSetList = PsGetCurrentProcess()->Vm.VmWorkingSetList;
+
     /* Done */
     Process->HasAddressSpace = TRUE;//??
     return STATUS_SUCCESS;
@@ -1039,10 +1149,11 @@ MmInitializeHandBuiltProcess(IN PEPROCESS Process,
 
 NTSTATUS
 NTAPI
+INIT_FUNCTION
 MmInitializeHandBuiltProcess2(IN PEPROCESS Process)
 {
-    /* Lock the VAD, ARM3-owned ranges away */                            
-    MiRosTakeOverPebTebRanges(Process);
+    /* Lock the VAD, ARM3-owned ranges away */
+    MiRosTakeOverSharedUserPage(Process);
     return STATUS_SUCCESS;
 }
 
@@ -1055,47 +1166,65 @@ MmCreateProcessAddressSpace(IN ULONG MinWs,
                             OUT PULONG_PTR DirectoryTableBase)
 {
     KIRQL OldIrql;
-    PFN_NUMBER PdeIndex, HyperIndex;
+    PFN_NUMBER PdeIndex, HyperIndex, WsListIndex;
     PMMPTE PointerPte;
     MMPTE TempPte, PdePte;
     ULONG PdeOffset;
-    PMMPTE SystemTable;
+    PMMPTE SystemTable, HyperTable;
     ULONG Color;
+    PMMPFN Pfn1;
 
     /* Choose a process color */
-    Process->NextPageColor = RtlRandom(&MmProcessColorSeed);
-    
+    Process->NextPageColor = (USHORT)RtlRandom(&MmProcessColorSeed);
+
     /* Setup the hyperspace lock */
     KeInitializeSpinLock(&Process->HyperSpaceLock);
 
     /* Lock PFN database */
     OldIrql = KeAcquireQueuedSpinLock(LockQueuePfnLock);
-    
+
     /* Get a zero page for the PDE, if possible */
     Color = MI_GET_NEXT_PROCESS_COLOR(Process);
+    MI_SET_USAGE(MI_USAGE_PAGE_DIRECTORY);
     PdeIndex = MiRemoveZeroPageSafe(Color);
     if (!PdeIndex)
     {
         /* No zero pages, grab a free one */
         PdeIndex = MiRemoveAnyPage(Color);
-        
+
         /* Zero it outside the PFN lock */
         KeReleaseQueuedSpinLock(LockQueuePfnLock, OldIrql);
         MiZeroPhysicalPage(PdeIndex);
         OldIrql = KeAcquireQueuedSpinLock(LockQueuePfnLock);
     }
-    
+
     /* Get a zero page for hyperspace, if possible */
+    MI_SET_USAGE(MI_USAGE_PAGE_DIRECTORY);
     Color = MI_GET_NEXT_PROCESS_COLOR(Process);
     HyperIndex = MiRemoveZeroPageSafe(Color);
     if (!HyperIndex)
     {
         /* No zero pages, grab a free one */
         HyperIndex = MiRemoveAnyPage(Color);
-        
+
         /* Zero it outside the PFN lock */
         KeReleaseQueuedSpinLock(LockQueuePfnLock, OldIrql);
         MiZeroPhysicalPage(HyperIndex);
+        OldIrql = KeAcquireQueuedSpinLock(LockQueuePfnLock);
+    }
+
+    /* Get a zero page for the woring set list, if possible */
+    MI_SET_USAGE(MI_USAGE_PAGE_TABLE);
+    Color = MI_GET_NEXT_PROCESS_COLOR(Process);
+    WsListIndex = MiRemoveZeroPageSafe(Color);
+    if (!WsListIndex)
+    {
+        /* No zero pages, grab a free one */
+        WsListIndex = MiRemoveAnyPage(Color);
+
+        /* Zero it outside the PFN lock */
+        KeReleaseQueuedSpinLock(LockQueuePfnLock, OldIrql);
+        MiZeroPhysicalPage(WsListIndex);
     }
     else
     {
@@ -1108,12 +1237,45 @@ MmCreateProcessAddressSpace(IN ULONG MinWs,
     Process->AddressSpaceInitialized = 1;
 
     /* Set the base directory pointers */
+    Process->WorkingSetPage = WsListIndex;
     DirectoryTableBase[0] = PdeIndex << PAGE_SHIFT;
     DirectoryTableBase[1] = HyperIndex << PAGE_SHIFT;
 
     /* Make sure we don't already have a page directory setup */
     ASSERT(Process->Pcb.DirectoryTableBase[0] == 0);
 
+    /* Get a PTE to map hyperspace */
+    PointerPte = MiReserveSystemPtes(1, SystemPteSpace);
+    ASSERT(PointerPte != NULL);
+
+    /* Build it */
+    MI_MAKE_HARDWARE_PTE_KERNEL(&PdePte,
+                                PointerPte,
+                                MM_READWRITE,
+                                HyperIndex);
+
+    /* Set it dirty and map it */
+    MI_MAKE_DIRTY_PAGE(&PdePte);
+    MI_WRITE_VALID_PTE(PointerPte, PdePte);
+
+    /* Now get hyperspace's page table */
+    HyperTable = MiPteToAddress(PointerPte);
+
+    /* Now write the PTE/PDE entry for the working set list index itself */
+    TempPte = ValidKernelPte;
+    TempPte.u.Hard.PageFrameNumber = WsListIndex;
+    /* Hyperspace is local */
+    MI_MAKE_LOCAL_PAGE(&TempPte);
+    PdeOffset = MiAddressToPteOffset(MmWorkingSetList);
+    HyperTable[PdeOffset] = TempPte;
+
+    /* Let go of the system PTE */
+    MiReleaseSystemPtes(PointerPte, 1, SystemPteSpace);
+
+    /* Save the PTE address of the page directory itself */
+    Pfn1 = MiGetPfnEntry(PdeIndex);
+    Pfn1->PteAddress = (PMMPTE)PDE_BASE;
+
     /* Insert us into the Mm process list */
     InsertTailList(&MmProcessList, &Process->MmProcessLinks);
 
@@ -1128,7 +1290,7 @@ MmCreateProcessAddressSpace(IN ULONG MinWs,
                                 PdeIndex);
 
     /* Set it dirty and map it */
-    PdePte.u.Hard.Dirty = TRUE;
+    MI_MAKE_DIRTY_PAGE(&PdePte);
     MI_WRITE_VALID_PTE(PointerPte, PdePte);
 
     /* Now get the page directory (which we'll double map, so call it a page table */
@@ -1157,6 +1319,9 @@ MmCreateProcessAddressSpace(IN ULONG MinWs,
 
     /* Let go of the system PTE */
     MiReleaseSystemPtes(PointerPte, 1, SystemPteSpace);
+
+    /* Add the process to the session */
+    MiSessionAddProcess(Process);
     return TRUE;
 }
 #endif
@@ -1168,10 +1333,21 @@ MmCleanProcessAddressSpace(IN PEPROCESS Process)
     PMMVAD Vad;
     PMM_AVL_TABLE VadTree;
     PETHREAD Thread = PsGetCurrentThread();
-    
+
+    /* Only support this */
+    ASSERT(Process->AddressSpaceInitialized == 2);
+
+    /* Remove from the session */
+    MiSessionRemoveProcess();
+
     /* 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;
     while (VadTree->NumberGenericTableElements)
@@ -1180,24 +1356,32 @@ 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);
         MiRemoveNode((PMMADDRESS_NODE)Vad, VadTree);
 
-        /* Only PEB/TEB VADs supported for now */
-        ASSERT(Vad->u.VadFlags.PrivateMemory == 1);
+        /* Only regular VADs supported for now */
         ASSERT(Vad->u.VadFlags.VadType == VadNone);
-        
-        /* Delete the addresses */
-        MiDeleteVirtualAddresses(Vad->StartingVpn << PAGE_SHIFT,
-                                 (Vad->EndingVpn << PAGE_SHIFT) | (PAGE_SIZE - 1),
-                                 Vad);
-        
-        /* Release the working set */
-        MiUnlockProcessWorkingSet(Process, Thread);
-        
+
+        /* Check if this is a section VAD */
+        if (!(Vad->u.VadFlags.PrivateMemory) && (Vad->ControlArea))
+        {
+            /* Remove the view */
+            MiRemoveMappedView(Process, Vad);
+        }
+        else
+        {
+            /* Delete the addresses */
+            MiDeleteVirtualAddresses(Vad->StartingVpn << PAGE_SHIFT,
+                                     (Vad->EndingVpn << PAGE_SHIFT) | (PAGE_SIZE - 1),
+                                     Vad);
+
+            /* Release the working set */
+            MiUnlockProcessWorkingSetUnsafe(Process, Thread);
+        }
+
         /* Skip ARM3 fake VADs, they'll be freed by MmDeleteProcessAddresSpace */
         if (Vad->u.VadFlags.Spare == 1)
         {
@@ -1205,15 +1389,717 @@ MmCleanProcessAddressSpace(IN PEPROCESS Process)
             Vad->u.VadFlags.Spare = 2;
             continue;
         }
-        
+
         /* Free the VAD memory */
         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);
 }
 
+VOID
+NTAPI
+MmDeleteProcessAddressSpace2(IN PEPROCESS Process)
+{
+    PMMPFN Pfn1, Pfn2;
+    KIRQL OldIrql;
+    PFN_NUMBER PageFrameIndex;
+
+    //ASSERT(Process->CommitCharge == 0);
+
+    /* Acquire the PFN lock */
+    OldIrql = KeAcquireQueuedSpinLock(LockQueuePfnLock);
+
+    /* Check for fully initialized process */
+    if (Process->AddressSpaceInitialized == 2)
+    {
+        /* Map the working set page and its page table */
+        Pfn1 = MiGetPfnEntry(Process->WorkingSetPage);
+        Pfn2 = MiGetPfnEntry(Pfn1->u4.PteFrame);
+
+        /* Nuke it */
+        MI_SET_PFN_DELETED(Pfn1);
+        MiDecrementShareCount(Pfn2, Pfn1->u4.PteFrame);
+        MiDecrementShareCount(Pfn1, Process->WorkingSetPage);
+        ASSERT((Pfn1->u3.e2.ReferenceCount == 0) || (Pfn1->u3.e1.WriteInProgress));
+        MiReleaseSystemPtes(MiAddressToPte(Process->Vm.VmWorkingSetList), 1, SystemPteSpace);
+
+        /* Now map hyperspace and its page table */
+        PageFrameIndex = Process->Pcb.DirectoryTableBase[1] >> PAGE_SHIFT;
+        Pfn1 = MiGetPfnEntry(PageFrameIndex);
+        Pfn2 = MiGetPfnEntry(Pfn1->u4.PteFrame);
+
+        /* Nuke it */
+        MI_SET_PFN_DELETED(Pfn1);
+        MiDecrementShareCount(Pfn2, Pfn1->u4.PteFrame);
+        MiDecrementShareCount(Pfn1, PageFrameIndex);
+        ASSERT((Pfn1->u3.e2.ReferenceCount == 0) || (Pfn1->u3.e1.WriteInProgress));
+
+        /* Finally, nuke the PDE itself */
+        PageFrameIndex = Process->Pcb.DirectoryTableBase[0] >> PAGE_SHIFT;
+        Pfn1 = MiGetPfnEntry(PageFrameIndex);
+        MI_SET_PFN_DELETED(Pfn1);
+        MiDecrementShareCount(Pfn1, PageFrameIndex);
+        MiDecrementShareCount(Pfn1, PageFrameIndex);
+
+        /* Page table is now dead. Bye bye... */
+        ASSERT((Pfn1->u3.e2.ReferenceCount == 0) || (Pfn1->u3.e1.WriteInProgress));
+    }
+    else
+    {
+        /* A partly-initialized process should never exit through here */
+        ASSERT(FALSE);
+    }
+
+    /* Release the PFN lock */
+    KeReleaseQueuedSpinLock(LockQueuePfnLock, OldIrql);
+
+    /* Drop a reference on the session */
+    if (Process->Session) MiReleaseProcessReferenceToSessionDataPage(Process->Session);
+
+    /* Clear out the PDE pages */
+    Process->Pcb.DirectoryTableBase[0] = 0;
+    Process->Pcb.DirectoryTableBase[1] = 0;
+}
+
+/* SESSION CODE TO MOVE TO SESSION.C ******************************************/
+
+PMM_SESSION_SPACE MmSessionSpace;
+PFN_NUMBER MiSessionDataPages, MiSessionTagPages, MiSessionTagSizePages;
+PFN_NUMBER MiSessionBigPoolPages, MiSessionCreateCharge;
+KGUARDED_MUTEX MiSessionIdMutex;
+LONG MmSessionDataPages;
+PRTL_BITMAP MiSessionIdBitmap;
+volatile LONG MiSessionLeaderExists;
+
+VOID
+NTAPI
+MiInitializeSessionIds(VOID)
+{
+    ULONG Size, BitmapSize;
+    PFN_NUMBER TotalPages;
+
+    /* Setup the total number of data pages needed for the structure */
+    TotalPages = MI_SESSION_DATA_PAGES_MAXIMUM;
+    MiSessionDataPages = ROUND_TO_PAGES(sizeof(MM_SESSION_SPACE)) >> PAGE_SHIFT;
+    ASSERT(MiSessionDataPages <= MI_SESSION_DATA_PAGES_MAXIMUM - 3);
+    TotalPages -= MiSessionDataPages;
+
+    /* Setup the number of pages needed for session pool tags */
+    MiSessionTagSizePages = 2;
+    MiSessionBigPoolPages = 1;
+    MiSessionTagPages = MiSessionTagSizePages + MiSessionBigPoolPages;
+    ASSERT(MiSessionTagPages <= TotalPages);
+    ASSERT(MiSessionTagPages < MI_SESSION_TAG_PAGES_MAXIMUM);
+
+    /* Total pages needed for a session (FIXME: Probably different on PAE/x64) */
+    MiSessionCreateCharge = 1 + MiSessionDataPages + MiSessionTagPages;
+
+    /* Initialize the lock */
+    KeInitializeGuardedMutex(&MiSessionIdMutex);
+
+    /* Allocate the bitmap */
+    Size = MI_INITIAL_SESSION_IDS;
+    BitmapSize = ((Size + 31) / 32) * sizeof(ULONG);
+    MiSessionIdBitmap = ExAllocatePoolWithTag(PagedPool,
+                                              sizeof(RTL_BITMAP) + BitmapSize,
+                                              '  mM');
+    if (MiSessionIdBitmap)
+    {
+        /* Free all the bits */
+        RtlInitializeBitMap(MiSessionIdBitmap,
+                            (PVOID)(MiSessionIdBitmap + 1),
+                            Size);
+        RtlClearAllBits(MiSessionIdBitmap);
+    }
+    else
+    {
+        /* Die if we couldn't allocate the bitmap */
+        KeBugCheckEx(INSTALL_MORE_MEMORY,
+                     MmNumberOfPhysicalPages,
+                     MmLowestPhysicalPage,
+                     MmHighestPhysicalPage,
+                     0x200);
+    }
+}
+
+VOID
+NTAPI
+MiSessionLeader(IN PEPROCESS Process)
+{
+    KIRQL OldIrql;
+
+    /* Set the flag while under the expansion lock */
+    OldIrql = KeAcquireQueuedSpinLock(LockQueueExpansionLock);
+    Process->Vm.Flags.SessionLeader = TRUE;
+    KeReleaseQueuedSpinLock(LockQueueExpansionLock, OldIrql);
+}
+
+ULONG
+NTAPI
+MmGetSessionId(IN PEPROCESS Process)
+{
+    PMM_SESSION_SPACE SessionGlobal;
+
+    /* The session leader is always session zero */
+    if (Process->Vm.Flags.SessionLeader == 1) return 0;
+
+    /* Otherwise, get the session global, and read the session ID from it */
+    SessionGlobal = (PMM_SESSION_SPACE)Process->Session;
+    if (!SessionGlobal) return 0;
+    return SessionGlobal->SessionId;
+}
+
+VOID
+NTAPI
+MiReleaseProcessReferenceToSessionDataPage(IN PMM_SESSION_SPACE SessionGlobal)
+{
+    ULONG i, SessionId;
+    PMMPTE PointerPte;
+    PFN_NUMBER PageFrameIndex[MI_SESSION_DATA_PAGES_MAXIMUM];
+    PMMPFN Pfn1;
+    KIRQL OldIrql;
+
+    /* Is there more than just this reference? If so, bail out */
+    if (InterlockedDecrement(&SessionGlobal->ProcessReferenceToSession)) return;
+
+    /* Get the session ID */
+    SessionId = SessionGlobal->SessionId;
+    DPRINT1("Last process in sessino %d going down!!!\n", SessionId);
+
+    /* Free the session page tables */
+    ExFreePool(SessionGlobal->PageTables);
+    ASSERT(!MI_IS_PHYSICAL_ADDRESS(SessionGlobal));
+
+    /* Capture the data page PFNs */
+    PointerPte = MiAddressToPte(SessionGlobal);
+    for (i = 0; i < MiSessionDataPages; i++)
+    {
+        PageFrameIndex[i] = PFN_FROM_PTE(PointerPte + i);
+    }
+
+    /* Release them */
+    MiReleaseSystemPtes(PointerPte, MiSessionDataPages, SystemPteSpace);
+
+    /* Mark them as deleted */
+    for (i = 0; i < MiSessionDataPages; i++)
+    {
+        Pfn1 = MI_PFN_ELEMENT(PageFrameIndex[i]);
+        MI_SET_PFN_DELETED(Pfn1);
+    }
+
+    /* Loop every data page and drop a reference count */
+    OldIrql = KeAcquireQueuedSpinLock(LockQueuePfnLock);
+    for (i = 0; i < MiSessionDataPages; i++)
+    {
+        /* Sanity check that the page is correct, then decrement it */
+        Pfn1 = MI_PFN_ELEMENT(PageFrameIndex[i]);
+        ASSERT(Pfn1->u2.ShareCount == 1);
+        ASSERT(Pfn1->u3.e2.ReferenceCount == 1);
+        MiDecrementShareCount(Pfn1, PageFrameIndex[i]);
+    }
+
+    /* Done playing with pages, release the lock */
+    KeReleaseQueuedSpinLock(LockQueuePfnLock, OldIrql);
+
+    /* Decrement the number of data pages */
+    InterlockedDecrement(&MmSessionDataPages);
+
+    /* Free this session ID from the session bitmap */
+    KeAcquireGuardedMutex(&MiSessionIdMutex);
+    ASSERT(RtlCheckBit(MiSessionIdBitmap, SessionId));
+    RtlClearBit(MiSessionIdBitmap, SessionId);
+    KeReleaseGuardedMutex(&MiSessionIdMutex);
+}
+
+VOID
+NTAPI
+MiSessionRemoveProcess(VOID)
+{
+    PEPROCESS CurrentProcess = PsGetCurrentProcess();
+
+    /* If the process isn't already in a session, or if it's the leader... */
+    if (!(CurrentProcess->Flags & PSF_PROCESS_IN_SESSION_BIT) ||
+        (CurrentProcess->Vm.Flags.SessionLeader))
+    {
+        /* Then there's nothing to do */
+        return;
+    }
+
+    /* Sanity check */
+    ASSERT(MmIsAddressValid(MmSessionSpace) == TRUE);
+
+    /* Remove the process from the list ,and dereference the session */
+    RemoveEntryList(&CurrentProcess->SessionProcessLinks);
+    //MiDereferenceSession();
+}
+
+VOID
+NTAPI
+MiSessionAddProcess(IN PEPROCESS NewProcess)
+{
+    PMM_SESSION_SPACE SessionGlobal;
+
+    /* The current process must already be in a session */
+    if (!(PsGetCurrentProcess()->Flags & PSF_PROCESS_IN_SESSION_BIT)) return;
+
+    /* Sanity check */
+    ASSERT(MmIsAddressValid(MmSessionSpace) == TRUE);
+
+    /* Get the global session */
+    SessionGlobal = MmSessionSpace->GlobalVirtualAddress;
+
+    /* Increment counters */
+    InterlockedIncrement((PLONG)&SessionGlobal->ReferenceCount);
+    InterlockedIncrement(&SessionGlobal->ResidentProcessCount);
+    InterlockedIncrement(&SessionGlobal->ProcessReferenceToSession);
+
+    /* Set the session pointer */
+    ASSERT(NewProcess->Session == NULL);
+    NewProcess->Session = SessionGlobal;
+
+    /* Insert it into the process list */
+    InsertTailList(&SessionGlobal->ProcessList, &NewProcess->SessionProcessLinks);
+
+    /* Set the flag */
+    PspSetProcessFlag(NewProcess, PSF_PROCESS_IN_SESSION_BIT);
+}
+
+NTSTATUS
+NTAPI
+MiSessionInitializeWorkingSetList(VOID)
+{
+    KIRQL OldIrql;
+    PMMPTE PointerPte, PointerPde;
+    MMPTE TempPte;
+    ULONG Color, Index;
+    PFN_NUMBER PageFrameIndex;
+    PMM_SESSION_SPACE SessionGlobal;
+    BOOLEAN AllocatedPageTable;
+    PMMWSL WorkingSetList;
+
+    /* Get pointers to session global and the session working set list */
+    SessionGlobal = MmSessionSpace->GlobalVirtualAddress;
+    WorkingSetList = (PMMWSL)MiSessionSpaceWs;
+
+    /* Fill out the two pointers */
+    MmSessionSpace->Vm.VmWorkingSetList = WorkingSetList;
+    MmSessionSpace->Wsle = (PMMWSLE)WorkingSetList->UsedPageTableEntries;
+
+    /* Get the PDE for the working set, and check if it's already allocated */
+    PointerPde = MiAddressToPde(WorkingSetList);
+    if (PointerPde->u.Hard.Valid == 1)
+    {
+        /* Nope, we'll have to do it */
+        ASSERT(PointerPde->u.Hard.Global == 0);
+        AllocatedPageTable = FALSE;
+    }
+    else
+    {
+        /* Yep, that makes our job easier */
+        AllocatedPageTable = TRUE;
+    }
+
+    /* Get the PTE for the working set */
+    PointerPte = MiAddressToPte(WorkingSetList);
+
+    /* Initialize the working set lock, and lock the PFN database */
+    ExInitializePushLock(&SessionGlobal->Vm.WorkingSetMutex);
+    //MmLockPageableSectionByHandle(ExPageLockHandle);
+    OldIrql = KeAcquireQueuedSpinLock(LockQueuePfnLock);
+
+    /* Check if we need a page table */
+    if (AllocatedPageTable == TRUE)
+    {
+        /* Get a zeroed colored zero page */
+        Color = MI_GET_NEXT_COLOR();
+        PageFrameIndex = MiRemoveZeroPageSafe(Color);
+        if (!PageFrameIndex)
+        {
+            /* No zero pages, grab a free one */
+            PageFrameIndex = MiRemoveAnyPage(Color);
+
+            /* Zero it outside the PFN lock */
+            KeReleaseQueuedSpinLock(LockQueuePfnLock, OldIrql);
+            MiZeroPhysicalPage(PageFrameIndex);
+            OldIrql = KeAcquireQueuedSpinLock(LockQueuePfnLock);
+        }
+
+        /* Write a valid PDE for it */
+        TempPte.u.Long = ValidKernelPdeLocal.u.Long;
+        TempPte.u.Hard.PageFrameNumber = PageFrameIndex;
+        MI_WRITE_VALID_PTE(PointerPde, TempPte);
+
+        /* Add this into the list */
+        Index = ((ULONG_PTR)WorkingSetList - (ULONG_PTR)MmSessionBase) >> 22;
+        MmSessionSpace->PageTables[Index] = TempPte;
+
+        /* Initialize the page directory page, and now zero the working set list itself */
+        MiInitializePfnForOtherProcess(PageFrameIndex,
+                                       PointerPde,
+                                       MmSessionSpace->SessionPageDirectoryIndex);
+        KeZeroPages(PointerPte, PAGE_SIZE);
+    }
+
+    /* Get a zeroed colored zero page */
+    Color = MI_GET_NEXT_COLOR();
+    PageFrameIndex = MiRemoveZeroPageSafe(Color);
+    if (!PageFrameIndex)
+    {
+        /* No zero pages, grab a free one */
+        PageFrameIndex = MiRemoveAnyPage(Color);
+
+        /* Zero it outside the PFN lock */
+        KeReleaseQueuedSpinLock(LockQueuePfnLock, OldIrql);
+        MiZeroPhysicalPage(PageFrameIndex);
+        OldIrql = KeAcquireQueuedSpinLock(LockQueuePfnLock);
+    }
+
+    /* Write a valid PTE for it */
+    TempPte.u.Long = ValidKernelPteLocal.u.Long;
+    TempPte.u.Hard.Dirty = TRUE;
+    TempPte.u.Hard.PageFrameNumber = PageFrameIndex;
+
+    /* Initialize the working set list page */
+    MiInitializePfnAndMakePteValid(PageFrameIndex, PointerPte, TempPte);
+
+    /* Now we can release the PFN database lock */
+    KeReleaseQueuedSpinLock(LockQueuePfnLock, OldIrql);
+
+    /* Fill out the working set structure */
+    MmSessionSpace->Vm.Flags.SessionSpace = 1;
+    MmSessionSpace->Vm.MinimumWorkingSetSize = 20;
+    MmSessionSpace->Vm.MaximumWorkingSetSize = 384;
+    WorkingSetList->LastEntry = 20;
+    WorkingSetList->HashTable = NULL;
+    WorkingSetList->HashTableSize = 0;
+    WorkingSetList->Wsle = MmSessionSpace->Wsle;
+
+    /* FIXME: Handle list insertions */
+    ASSERT(SessionGlobal->WsListEntry.Flink == NULL);
+    ASSERT(SessionGlobal->WsListEntry.Blink == NULL);
+    ASSERT(SessionGlobal->Vm.WorkingSetExpansionLinks.Flink == NULL);
+    ASSERT(SessionGlobal->Vm.WorkingSetExpansionLinks.Blink == NULL);
+
+    /* All done, return */
+    //MmUnlockPageableImageSection(ExPageLockHandle);
+    return STATUS_SUCCESS;
+}
+
+NTSTATUS
+NTAPI
+MiSessionCreateInternal(OUT PULONG SessionId)
+{
+    PEPROCESS Process = PsGetCurrentProcess();
+    ULONG NewFlags, Flags, Size, i, Color;
+    KIRQL OldIrql;
+    PMMPTE PointerPte, PageTables, SessionPte;
+    PMMPDE PointerPde;
+    PMM_SESSION_SPACE SessionGlobal;
+    MMPTE TempPte;
+    NTSTATUS Status;
+    BOOLEAN Result;
+    PFN_NUMBER SessionPageDirIndex;
+    PFN_NUMBER TagPage[MI_SESSION_TAG_PAGES_MAXIMUM];
+    PFN_NUMBER DataPage[MI_SESSION_DATA_PAGES_MAXIMUM];
+
+    /* This should not exist yet */
+    ASSERT(MmIsAddressValid(MmSessionSpace) == FALSE);
+
+    /* Loop so we can set the session-is-creating flag */
+    Flags = Process->Flags;
+    while (TRUE)
+    {
+        /* Check if it's already set */
+        if (Flags & PSF_SESSION_CREATION_UNDERWAY_BIT)
+        {
+            /* Bail out */
+            DPRINT1("Lost session race\n");
+            return STATUS_ALREADY_COMMITTED;
+        }
+
+        /* Now try to set it */
+        NewFlags = InterlockedCompareExchange((PLONG)&Process->Flags,
+                                              Flags | PSF_SESSION_CREATION_UNDERWAY_BIT,
+                                              Flags);
+        if (NewFlags == Flags) break;
+
+        /* It changed, try again */
+        Flags = NewFlags;
+    }
+
+    /* Now we should own the flag */
+    ASSERT(Process->Flags & PSF_SESSION_CREATION_UNDERWAY_BIT);
+
+    /*
+     * Session space covers everything from 0xA0000000 to 0xC0000000.
+     * Allocate enough page tables to describe the entire region
+     */
+    Size = (0x20000000 / PDE_MAPPED_VA) * sizeof(MMPTE);
+    PageTables = ExAllocatePoolWithTag(NonPagedPool, Size, 'tHmM');
+    ASSERT(PageTables != NULL);
+    RtlZeroMemory(PageTables, Size);
+
+    /* Lock the session ID creation mutex */
+    KeAcquireGuardedMutex(&MiSessionIdMutex);
+
+    /* Allocate a new Session ID */
+    *SessionId = RtlFindClearBitsAndSet(MiSessionIdBitmap, 1, 0);
+    if (*SessionId == 0xFFFFFFFF)
+    {
+        /* We ran out of session IDs, we should expand */
+        DPRINT1("Too many sessions created. Expansion not yet supported\n");
+        ExFreePoolWithTag(PageTables, 'tHmM');
+        return STATUS_NO_MEMORY;
+    }
+
+    /* Unlock the session ID creation mutex */
+    KeReleaseGuardedMutex(&MiSessionIdMutex);
+
+    /* Reserve the global PTEs */
+    SessionPte = MiReserveSystemPtes(MiSessionDataPages, SystemPteSpace);
+    ASSERT(SessionPte != NULL);
+
+    /* Acquire the PFN lock while we set everything up */
+    OldIrql = KeAcquireQueuedSpinLock(LockQueuePfnLock);
+
+    /* Loop the global PTEs */
+    TempPte.u.Long = ValidKernelPte.u.Long;
+    for (i = 0; i < MiSessionDataPages; i++)
+    {
+        /* Get a zeroed colored zero page */
+        Color = MI_GET_NEXT_COLOR();
+        DataPage[i] = MiRemoveZeroPageSafe(Color);
+        if (!DataPage[i])
+        {
+            /* No zero pages, grab a free one */
+            DataPage[i] = MiRemoveAnyPage(Color);
+
+            /* Zero it outside the PFN lock */
+            KeReleaseQueuedSpinLock(LockQueuePfnLock, OldIrql);
+            MiZeroPhysicalPage(DataPage[i]);
+            OldIrql = KeAcquireQueuedSpinLock(LockQueuePfnLock);
+        }
+
+        /* Fill the PTE out */
+        TempPte.u.Hard.PageFrameNumber = DataPage[i];
+        MI_WRITE_VALID_PTE(SessionPte + i, TempPte);
+    }
+
+    /* Set the pointer to global space */
+    SessionGlobal = MiPteToAddress(SessionPte);
+
+    /* Get a zeroed colored zero page */
+    Color = MI_GET_NEXT_COLOR();
+    SessionPageDirIndex = MiRemoveZeroPageSafe(Color);
+    if (!SessionPageDirIndex)
+    {
+        /* No zero pages, grab a free one */
+        SessionPageDirIndex = MiRemoveAnyPage(Color);
+
+        /* Zero it outside the PFN lock */
+        KeReleaseQueuedSpinLock(LockQueuePfnLock, OldIrql);
+        MiZeroPhysicalPage(SessionPageDirIndex);
+        OldIrql = KeAcquireQueuedSpinLock(LockQueuePfnLock);
+    }
+
+    /* Fill the PTE out */
+    TempPte.u.Long = ValidKernelPdeLocal.u.Long;
+    TempPte.u.Hard.PageFrameNumber = SessionPageDirIndex;
+
+    /* Setup, allocate, fill out the MmSessionSpace PTE */
+    PointerPde = MiAddressToPde(MmSessionSpace);
+    ASSERT(PointerPde->u.Long == 0);
+    MI_WRITE_VALID_PTE(PointerPde, TempPte);
+    MiInitializePfnForOtherProcess(SessionPageDirIndex,
+                                   PointerPde,
+                                   SessionPageDirIndex);
+    ASSERT(MI_PFN_ELEMENT(SessionPageDirIndex)->u1.WsIndex == 0);
+
+     /* Loop all the local PTEs for it */
+    TempPte.u.Long = ValidKernelPteLocal.u.Long;
+    PointerPte = MiAddressToPte(MmSessionSpace);
+    for (i = 0; i < MiSessionDataPages; i++)
+    {
+        /* And fill them out */
+        TempPte.u.Hard.PageFrameNumber = DataPage[i];
+        MiInitializePfnAndMakePteValid(DataPage[i], PointerPte + i, TempPte);
+        ASSERT(MI_PFN_ELEMENT(DataPage[i])->u1.WsIndex == 0);
+    }
+
+     /* Finally loop all of the session pool tag pages */
+    for (i = 0; i < MiSessionTagPages; i++)
+    {
+        /* Grab a zeroed colored page */
+        Color = MI_GET_NEXT_COLOR();
+        TagPage[i] = MiRemoveZeroPageSafe(Color);
+        if (!TagPage[i])
+        {
+            /* No zero pages, grab a free one */
+            TagPage[i] = MiRemoveAnyPage(Color);
+
+            /* Zero it outside the PFN lock */
+            KeReleaseQueuedSpinLock(LockQueuePfnLock, OldIrql);
+            MiZeroPhysicalPage(TagPage[i]);
+            OldIrql = KeAcquireQueuedSpinLock(LockQueuePfnLock);
+        }
+
+        /* Fill the PTE out */
+        TempPte.u.Hard.PageFrameNumber = TagPage[i];
+        MiInitializePfnAndMakePteValid(TagPage[i],
+                                       PointerPte + MiSessionDataPages + i,
+                                       TempPte);
+    }
+
+    /* PTEs have been setup, release the PFN lock */
+    KeReleaseQueuedSpinLock(LockQueuePfnLock, OldIrql);
+
+    /* Fill out the session space structure now */
+    MmSessionSpace->GlobalVirtualAddress = SessionGlobal;
+    MmSessionSpace->ReferenceCount = 1;
+    MmSessionSpace->ResidentProcessCount = 1;
+    MmSessionSpace->u.LongFlags = 0;
+    MmSessionSpace->SessionId = *SessionId;
+    MmSessionSpace->LocaleId = PsDefaultSystemLocaleId;
+    MmSessionSpace->SessionPageDirectoryIndex = SessionPageDirIndex;
+    MmSessionSpace->Color = Color;
+    MmSessionSpace->NonPageablePages = MiSessionCreateCharge;
+    MmSessionSpace->CommittedPages = MiSessionCreateCharge;
+    MmSessionSpace->PageTables = PageTables;
+    MmSessionSpace->PageTables[PointerPde - MiAddressToPde(MmSessionBase)] = *PointerPde;
+    InitializeListHead(&MmSessionSpace->ImageList);
+    DPRINT1("Session %d is ready to go: 0x%p 0x%p, %lx 0x%p\n",
+            *SessionId, MmSessionSpace, SessionGlobal, SessionPageDirIndex, PageTables);
+
+    /* Initialize session pool */
+    //Status = MiInitializeSessionPool();
+    Status = STATUS_SUCCESS;
+    ASSERT(NT_SUCCESS(Status) == TRUE);
+
+    /* Initialize system space */
+    Result = MiInitializeSystemSpaceMap(&SessionGlobal->Session);
+    ASSERT(Result == TRUE);
+
+    /* Initialize the process list, make sure the workign set list is empty */
+    ASSERT(SessionGlobal->WsListEntry.Flink == NULL);
+    ASSERT(SessionGlobal->WsListEntry.Blink == NULL);
+    InitializeListHead(&SessionGlobal->ProcessList);
+
+    /* We're done, clear the flag */
+    ASSERT(Process->Flags & PSF_SESSION_CREATION_UNDERWAY_BIT);
+    PspClearProcessFlag(Process, PSF_SESSION_CREATION_UNDERWAY_BIT);
+
+    /* Insert the process into the session  */
+    ASSERT(Process->Session == NULL);
+    ASSERT(SessionGlobal->ProcessReferenceToSession == 0);
+    SessionGlobal->ProcessReferenceToSession = 1;
+
+    /* We're done */
+    InterlockedIncrement(&MmSessionDataPages);
+    return STATUS_SUCCESS;
+}
+
+NTSTATUS
+NTAPI
+MmSessionCreate(OUT PULONG SessionId)
+{
+    PEPROCESS Process = PsGetCurrentProcess();
+    ULONG SessionLeaderExists;
+    NTSTATUS Status;
+
+    /* Fail if the process is already in a session */
+    if (Process->Flags & PSF_PROCESS_IN_SESSION_BIT)
+    {
+        DPRINT1("Process already in session\n");
+        return STATUS_ALREADY_COMMITTED;
+    }
+
+    /* Check if the process is already the session leader */
+    if (!Process->Vm.Flags.SessionLeader)
+    {
+        /* Atomically set it as the leader */
+        SessionLeaderExists = InterlockedCompareExchange(&MiSessionLeaderExists, 1, 0);
+        if (SessionLeaderExists)
+        {
+            DPRINT1("Session leader race\n");
+            return STATUS_INVALID_SYSTEM_SERVICE;
+        }
+
+        /* Do the work required to upgrade him */
+        MiSessionLeader(Process);
+    }
+
+    /* Create the session */
+    KeEnterCriticalRegion();
+    Status = MiSessionCreateInternal(SessionId);
+    if (!NT_SUCCESS(Status))
+    {
+        KeLeaveCriticalRegion();
+        return Status;
+    }
+
+    /* Set up the session working set */
+    Status = MiSessionInitializeWorkingSetList();
+    if (!NT_SUCCESS(Status))
+    {
+        /* Fail */
+        //MiDereferenceSession();
+        ASSERT(FALSE);
+        KeLeaveCriticalRegion();
+        return Status;
+    }
+
+    /* All done */
+    KeLeaveCriticalRegion();
+
+    /* Set and assert the flags, and return */
+    MmSessionSpace->u.Flags.Initialized = 1;
+    PspSetProcessFlag(Process, PSF_PROCESS_IN_SESSION_BIT);
+    ASSERT(MiSessionLeaderExists == 1);
+    return Status;
+}
+
+NTSTATUS
+NTAPI
+MmSessionDelete(IN ULONG SessionId)
+{
+    PEPROCESS Process = PsGetCurrentProcess();
+
+    /* Process must be in a session */
+    if (!(Process->Flags & PSF_PROCESS_IN_SESSION_BIT))
+    {
+        DPRINT1("Not in a session!\n");
+        return STATUS_UNABLE_TO_FREE_VM;
+    }
+
+    /* It must be the session leader */
+    if (!Process->Vm.Flags.SessionLeader)
+    {
+        DPRINT1("Not a session leader!\n");
+        return STATUS_UNABLE_TO_FREE_VM;
+    }
+
+    /* Remove one reference count */
+    KeEnterCriticalRegion();
+    /* FIXME: Do it */
+    KeLeaveCriticalRegion();
+
+    /* All done */
+    return STATUS_SUCCESS;
+}
+
 /* SYSTEM CALLS ***************************************************************/
 
 NTSTATUS