[tkreuzer] 61061: {NTOSKRNL] Properly implement MmGetSessionById, MmAttachSession, MmDetachSession and MmQuitNextSession. Implement more session cleanup code and make sure a session is not killed w...

Author: tkreuzer
Date: Wed Nov 20 23:50:42 2013
New Revision: 61061

URL: http://svn.reactos.org/svn/reactos?rev=61061&view=rev
Log:
{NTOSKRNL]
Properly implement MmGetSessionById, MmAttachSession, MmDetachSession and
MmQuitNextSession. Implement more session cleanup code and make sure a session is not
killed while a process is attached to it. There used to be a bug, where removing a process
from the session list caused pool corruption, which was because we called
MiSessionRemoveProcess twice: once from PspExitThread, when the last thread exited and
another time from PspDeleteProcess, when the reaper thread deleted the process object.
This was, because the flag that said the process was in a session was not cleared
properly. This is now fixed and inserting and removing processes into the session's
list works fine. Also protect it properly with a spinlock.

Modified:
    trunk/reactos/ntoskrnl/ex/win32k.c
    trunk/reactos/ntoskrnl/mm/ARM3/miarm.h
    trunk/reactos/ntoskrnl/mm/ARM3/session.c
    trunk/reactos/ntoskrnl/mm/mminit.c

Modified: trunk/reactos/ntoskrnl/ex/win32k.c
URL:
http://svn.reactos.org/svn/reactos/trunk/reactos/ntoskrnl/ex/win32k.c?rev=6…
==============================================================================
--- trunk/reactos/ntoskrnl/ex/win32k.c	[iso-8859-1] (original)
+++ trunk/reactos/ntoskrnl/ex/win32k.c	[iso-8859-1] Wed Nov 20 23:50:42 2013
@@ -72,7 +72,7 @@
         (PsGetCurrentProcessSessionId() != Win32ObjectHeader->SessionId))
     {
         /* Get the session from the objects session Id */
-        DPRINT1("SessionId == %d\n", Win32ObjectHeader->SessionId);
+        DPRINT("SessionId == %d\n", Win32ObjectHeader->SessionId);
         SessionEntry = MmGetSessionById(Win32ObjectHeader->SessionId);
         if (SessionEntry == NULL)
         {

Modified: trunk/reactos/ntoskrnl/mm/ARM3/miarm.h
URL:
http://svn.reactos.org/svn/reactos/trunk/reactos/ntoskrnl/mm/ARM3/miarm.h?r…
==============================================================================
--- trunk/reactos/ntoskrnl/mm/ARM3/miarm.h	[iso-8859-1] (original)
+++ trunk/reactos/ntoskrnl/mm/ARM3/miarm.h	[iso-8859-1] Wed Nov 20 23:50:42 2013
@@ -18,7 +18,7 @@
 
 #define MI_MIN_INIT_PAGED_POOLSIZE              (32 * _1MB)
 
-#define MI_SESSION_VIEW_SIZE                    (20 * _1MB)
+#define MI_SESSION_VIEW_SIZE                    (48 * _1MB)
 #define MI_SESSION_POOL_SIZE                    (16 * _1MB)
 #define MI_SESSION_IMAGE_SIZE                   (8 * _1MB)
 #define MI_SESSION_WORKING_SET_SIZE             (4 * _1MB)
@@ -716,6 +716,7 @@
 extern SIZE_T MmSystemLockPagesCount;
 extern ULONG_PTR MmSubsectionBase;
 extern LARGE_INTEGER MmCriticalSectionTimeout;
+extern LIST_ENTRY MmWorkingSetExpansionHead;
 
 BOOLEAN
 FORCEINLINE
@@ -1732,6 +1733,12 @@
 
 VOID
 NTAPI
+MiInitializeSessionWsSupport(
+    VOID
+);
+
+VOID
+NTAPI
 MiInitializeSessionIds(
     VOID
 );

Modified: trunk/reactos/ntoskrnl/mm/ARM3/session.c
URL:
http://svn.reactos.org/svn/reactos/trunk/reactos/ntoskrnl/mm/ARM3/session.c…
==============================================================================
--- trunk/reactos/ntoskrnl/mm/ARM3/session.c	[iso-8859-1] (original)
+++ trunk/reactos/ntoskrnl/mm/ARM3/session.c	[iso-8859-1] Wed Nov 20 23:50:42 2013
@@ -4,6 +4,7 @@
  * FILE:            ntoskrnl/mm/ARM3/session.c
  * PURPOSE:         Session support routines
  * PROGRAMMERS:     ReactOS Portable Systems Group
+ *                  Timo Kreuzer (timo.kreuzer(a)reactos.org)
  */
 
 /* INCLUDES *******************************************************************/
@@ -25,11 +26,46 @@
 PRTL_BITMAP MiSessionIdBitmap;
 volatile LONG MiSessionLeaderExists;
 
-// HACK: we support only one process. The creator is CSRSS and that lives!
-PEPROCESS Session0CreatorProcess;
+LIST_ENTRY MiSessionWsList;
+LIST_ENTRY MmWorkingSetExpansionHead;
+
+KSPIN_LOCK MmExpansionLock;
+PETHREAD MiExpansionLockOwner;
 
 
 /* PRIVATE FUNCTIONS **********************************************************/
+
+FORCEINLINE
+KIRQL
+MiAcquireExpansionLock()
+{
+    KIRQL OldIrql;
+
+    ASSERT(KeGetCurrentIrql() <= APC_LEVEL);
+    KeAcquireSpinLock(&MmExpansionLock, &OldIrql);
+    ASSERT(MiExpansionLockOwner == NULL);
+    MiExpansionLockOwner = PsGetCurrentThread();
+    return OldIrql;
+}
+
+FORCEINLINE
+VOID
+MiReleaseExpansionLock(KIRQL OldIrql)
+{
+    ASSERT(MiExpansionLockOwner == PsGetCurrentThread());
+    MiExpansionLockOwner = NULL;
+    KeReleaseSpinLock(&MmExpansionLock, OldIrql);
+    ASSERT(KeGetCurrentIrql() <= APC_LEVEL);
+}
+
+VOID
+NTAPI
+MiInitializeSessionWsSupport(VOID)
+{
+    /* Initialize the list heads */
+    InitializeListHead(&MiSessionWsList);
+    InitializeListHead(&MmWorkingSetExpansionHead);
+}
 
 LCID
 NTAPI
@@ -225,9 +261,132 @@
 
 VOID
 NTAPI
+MiDereferenceSessionFinal(VOID)
+{
+    PMM_SESSION_SPACE SessionGlobal;
+    KIRQL OldIrql;
+
+    /* Get the pointer to the global session address */
+    SessionGlobal = MmSessionSpace->GlobalVirtualAddress;
+
+    /* Acquire the expansion lock */
+    OldIrql = MiAcquireExpansionLock();
+
+    /* Set delete pending flag, so that processes can no longer attach to this
+       session and the last process that detaches sets the AttachEvent */
+    ASSERT(SessionGlobal->u.Flags.DeletePending == 0);
+    SessionGlobal->u.Flags.DeletePending = 1;
+
+    /* Check if we have any attached processes */
+    if (SessionGlobal->AttachCount)
+    {
+        /* Initialize the event (it's not in use yet!) */
+        KeInitializeEvent(&SessionGlobal->AttachEvent, NotificationEvent, FALSE);
+
+        /* Release the expansion lock for the wait */
+        MiReleaseExpansionLock(OldIrql);
+
+        /* Wait for the event to be set due to the last process detach */
+        KeWaitForSingleObject(&SessionGlobal->AttachEvent, WrVirtualMemory, 0, 0,
0);
+
+        /* Reacquire the expansion lock */
+        OldIrql = MiAcquireExpansionLock();
+
+        /* Makes sure we still have the delete flag and no attached processes */
+        ASSERT(MmSessionSpace->u.Flags.DeletePending == 1);
+        ASSERT(MmSessionSpace->AttachCount == 0);
+    }
+
+    /* Check if the session is in the workingset expansion list */
+    if (SessionGlobal->Vm.WorkingSetExpansionLinks.Flink != NULL)
+    {
+        /* Remove the session from the list and zero the list entry */
+        RemoveEntryList(&SessionGlobal->Vm.WorkingSetExpansionLinks);
+        SessionGlobal->Vm.WorkingSetExpansionLinks.Flink = 0;
+    }
+
+    /* Check if the session is in the workingset list */
+    if (SessionGlobal->WsListEntry.Flink)
+    {
+        /* Remove the session from the list and zero the list entry */
+        RemoveEntryList(&SessionGlobal->WsListEntry);
+        SessionGlobal->WsListEntry.Flink = NULL;
+    }
+
+    /* Release the expansion lock */
+    MiReleaseExpansionLock(OldIrql);
+
+    /* Check for a win32k unload routine */
+    if (SessionGlobal->Win32KDriverUnload)
+    {
+        /* Call it */
+        SessionGlobal->Win32KDriverUnload(NULL);
+    }
+}
+
+
+VOID
+NTAPI
+MiDereferenceSession(VOID)
+{
+    PMM_SESSION_SPACE SessionGlobal;
+    PEPROCESS Process;
+    ULONG ReferenceCount, SessionId;
+
+    /* Sanity checks */
+    ASSERT(PsGetCurrentProcess()->ProcessInSession ||
+           ((MmSessionSpace->u.Flags.Initialized == 0) &&
+            (PsGetCurrentProcess()->Vm.Flags.SessionLeader == 1) &&
+            (MmSessionSpace->ReferenceCount == 1)));
+
+    /* The session bit must be set */
+    SessionId = MmSessionSpace->SessionId;
+    ASSERT(RtlCheckBit(MiSessionIdBitmap, SessionId));
+
+    /* Get the current process */
+    Process = PsGetCurrentProcess();
+
+    /* Decrement the process count */
+    InterlockedDecrement(&MmSessionSpace->ResidentProcessCount);
+
+    /* Decrement the reference count and check if was the last reference */
+    ReferenceCount = InterlockedDecrement(&MmSessionSpace->ReferenceCount);
+    if (ReferenceCount == 0)
+    {
+        /* No more references left, kill the session completely */
+        MiDereferenceSessionFinal();
+    }
+
+    /* Check if tis is the session leader or the last process in the session */
+    if ((Process->Vm.Flags.SessionLeader) || (ReferenceCount == 0))
+    {
+        /* Get the global session address before we kill the session mapping */
+        SessionGlobal = MmSessionSpace->GlobalVirtualAddress;
+
+        /* Delete all session PDEs and flush the TB */
+        RtlZeroMemory(MiAddressToPde(MmSessionBase),
+                      BYTES_TO_PAGES(MmSessionSize) * sizeof(MMPDE));
+        KeFlushEntireTb(FALSE, FALSE);
+
+        /* Is this the session leader? */
+        if (Process->Vm.Flags.SessionLeader)
+        {
+            /* Clean up the references here. */
+            ASSERT(Process->Session == NULL);
+            MiReleaseProcessReferenceToSessionDataPage(SessionGlobal);
+        }
+    }
+
+    /* Reset the current process' session flag */
+    RtlInterlockedClearBits(&Process->Flags, PSF_PROCESS_IN_SESSION_BIT);
+}
+
+VOID
+NTAPI
 MiSessionRemoveProcess(VOID)
 {
     PEPROCESS CurrentProcess = PsGetCurrentProcess();
+    KIRQL OldIrql;
 
     /* If the process isn't already in a session, or if it's the leader... */
     if (!(CurrentProcess->Flags & PSF_PROCESS_IN_SESSION_BIT) ||
@@ -240,10 +399,17 @@
     /* Sanity check */
     ASSERT(MmIsAddressValid(MmSessionSpace) == TRUE);
 
-    /* Remove the process from the list ,and dereference the session */
-    // DO NOT ENABLE THIS UNLESS YOU FIXED THE NP POOL CORRUPTION THAT IT CAUSES!!!
-    //RemoveEntryList(&CurrentProcess->SessionProcessLinks);
-    //MiDereferenceSession();
+    /* Acquire the expansion lock while touching the session */
+    OldIrql = MiAcquireExpansionLock();
+
+    /* Remove the process from the list */
+    RemoveEntryList(&CurrentProcess->SessionProcessLinks);
+
+    /* Release the lock again */
+    MiReleaseExpansionLock(OldIrql);
+
+    /* Dereference the session */
+    MiDereferenceSession();
 }
 
 VOID
@@ -251,6 +417,7 @@
 MiSessionAddProcess(IN PEPROCESS NewProcess)
 {
     PMM_SESSION_SPACE SessionGlobal;
+    KIRQL OldIrql;
 
     /* The current process must already be in a session */
     if (!(PsGetCurrentProcess()->Flags & PSF_PROCESS_IN_SESSION_BIT)) return;
@@ -270,9 +437,14 @@
     ASSERT(NewProcess->Session == NULL);
     NewProcess->Session = SessionGlobal;
 
+    /* Acquire the expansion lock while touching the session */
+    OldIrql = MiAcquireExpansionLock();
+
     /* Insert it into the process list */
-    // DO NOT ENABLE THIS UNLESS YOU FIXED THE NP POOL CORRUPTION THAT IT CAUSES!!!
-    //InsertTailList(&SessionGlobal->ProcessList,
&NewProcess->SessionProcessLinks);
+    InsertTailList(&SessionGlobal->ProcessList,
&NewProcess->SessionProcessLinks);
+
+    /* Release the lock again */
+    MiReleaseExpansionLock(OldIrql);
 
     /* Set the flag */
     PspSetProcessFlag(NewProcess, PSF_PROCESS_IN_SESSION_BIT);
@@ -389,11 +561,21 @@
     WorkingSetList->HashTableSize = 0;
     WorkingSetList->Wsle = MmSessionSpace->Wsle;
 
-    /* FIXME: Handle list insertions */
+    /* Acquire the expansion lock while touching the session */
+    OldIrql = MiAcquireExpansionLock();
+
+    /* Handle list insertions */
     ASSERT(SessionGlobal->WsListEntry.Flink == NULL);
     ASSERT(SessionGlobal->WsListEntry.Blink == NULL);
+    InsertTailList(&MiSessionWsList, &SessionGlobal->WsListEntry);
+
     ASSERT(SessionGlobal->Vm.WorkingSetExpansionLinks.Flink == NULL);
     ASSERT(SessionGlobal->Vm.WorkingSetExpansionLinks.Blink == NULL);
+    InsertTailList(&MmWorkingSetExpansionHead,
+                   &SessionGlobal->Vm.WorkingSetExpansionLinks);
+
+    /* Release the lock again */
+    MiReleaseExpansionLock(OldIrql);
 
     /* All done, return */
     //MmUnlockPageableImageSection(ExPageLockHandle);
@@ -610,10 +792,6 @@
     ASSERT(SessionGlobal->ProcessReferenceToSession == 0);
     SessionGlobal->ProcessReferenceToSession = 1;
 
-    // HACK: we only support 1 session and save the creator process
-    NT_ASSERT(Session0CreatorProcess == NULL);
-    Session0CreatorProcess = PsGetCurrentProcess();
-
     /* We're done */
     InterlockedIncrement(&MmSessionDataPages);
     return STATUS_SUCCESS;
@@ -716,16 +894,63 @@
     _Out_ PKAPC_STATE ApcState)
 {
     PEPROCESS EntryProcess;
+    PMM_SESSION_SPACE EntrySession, CurrentSession;
+    PEPROCESS CurrentProcess;
+    KIRQL OldIrql;
 
     /* The parameter is the actual process! */
     EntryProcess = SessionEntry;
     NT_ASSERT(EntryProcess != NULL);
 
-    /* HACK: for now we only support 1 session! */
-    NT_ASSERT(((PMM_SESSION_SPACE)EntryProcess->Session)->SessionId == 1);
-
-    /* Very simple for now: just attach to the process we have */
+    /* Sanity checks */
+    ASSERT(KeGetCurrentIrql() <= APC_LEVEL);
+    ASSERT(EntryProcess->Vm.Flags.SessionLeader == 0);
+
+    /* Get the session from the process that was passed in */
+    EntrySession = EntryProcess->Session;
+    ASSERT(EntrySession != NULL);
+
+    /* Get the current process and it's session */
+    CurrentProcess = PsGetCurrentProcess();
+    CurrentSession = CurrentProcess->Session;
+
+    /* Acquire the expansion lock while touching the session */
+    OldIrql = MiAcquireExpansionLock();
+
+    /* Check if the session is about to be deleted */
+    if (EntrySession->u.Flags.DeletePending)
+    {
+        /* We cannot attach to it, so unlock and fail */
+        MiReleaseExpansionLock(OldIrql);
+        return STATUS_PROCESS_IS_TERMINATING;
+    }
+
+    /* Count the number of attaches */
+    EntrySession->AttachCount++;
+
+    /* we can release the lock again */
+    MiReleaseExpansionLock(OldIrql);
+
+    /* Check if we are not the session leader and we are in a session */
+    if (!CurrentProcess->Vm.Flags.SessionLeader && (CurrentSession != NULL))
+    {
+        /* Are we already in the right session? */
+        if (CurrentSession == EntrySession)
+        {
+            /* We are, so "attach" to the current process */
+            EntryProcess = CurrentProcess;
+        }
+        else
+        {
+            /* We are not, the session id should better not match! */
+            ASSERT(CurrentSession->SessionId != EntrySession->SessionId);
+        }
+    }
+
+    /* Now attach to the process that we have */
     KeStackAttachProcess(&EntryProcess->Pcb, ApcState);
+
+    /* Success! */
     return STATUS_SUCCESS;
 }
 
@@ -737,16 +962,42 @@
     _In_ PKAPC_STATE ApcState)
 {
     PEPROCESS EntryProcess;
+    PMM_SESSION_SPACE EntrySession;
+    KIRQL OldIrql;
+    BOOLEAN DeletePending;
 
     /* The parameter is the actual process! */
     EntryProcess = SessionEntry;
     NT_ASSERT(EntryProcess != NULL);
 
-    /* HACK: for now we only support 1 session! */
-    NT_ASSERT(((PMM_SESSION_SPACE)EntryProcess->Session)->SessionId == 0);
-
-    /* Very simple for now: just detach */
+    /* Sanity checks */
+    ASSERT(KeGetCurrentIrql() <= APC_LEVEL);
+    ASSERT(EntryProcess->Vm.Flags.SessionLeader == 0);
+
+    /* Get the session from the process that was passed in */
+    EntrySession = EntryProcess->Session;
+    ASSERT(EntrySession != NULL);
+
+    /* Acquire the expansion lock while touching the session */
+    OldIrql = MiAcquireExpansionLock();
+
+    /* Make sure we have at least one attach and decrement the count */
+    ASSERT(EntrySession->AttachCount >= 1);
+    EntrySession->AttachCount--;
+
+    /* Remember if a delete is pending and we were the last one attached */
+    DeletePending = EntrySession->u.Flags.DeletePending &&
+                    (EntrySession->AttachCount == 0);
+
+    /* Release the lock again */
+    MiReleaseExpansionLock(OldIrql);
+
+    /* Detach from the process */
     KeUnstackDetachProcess(ApcState);
+
+    /* Check if we need to set the attach event */
+    if (DeletePending)
+        KeSetEvent(&EntrySession->AttachEvent, IO_NO_INCREMENT, FALSE);
 }
 
 VOID
@@ -760,11 +1011,13 @@
     EntryProcess = SessionEntry;
     NT_ASSERT(EntryProcess != NULL);
 
-    /* HACK: for now we only support 1 session! */
-    NT_ASSERT(((PMM_SESSION_SPACE)EntryProcess->Session)->SessionId == 0);
-
-    /* Get rid of the reference we got */
-    ObDereferenceObject(SessionEntry);
+    /* Sanity checks */
+    ASSERT(KeGetCurrentIrql () <= APC_LEVEL);
+    ASSERT(EntryProcess->Vm.Flags.SessionLeader == 0);
+    ASSERT(EntryProcess->Session != NULL);
+
+    /* Get rid of the reference we took */
+    ObDereferenceObject(EntryProcess);
 }
 
 PVOID
@@ -772,9 +1025,39 @@
 MmGetSessionById(
     _In_ ULONG SessionId)
 {
-    /* HACK: for now we only support 1 session! */
-    NT_ASSERT(SessionId == 0);
-
-    /* Just return the sessions creator process, which is csrss and still alive. */
-    return Session0CreatorProcess;
-}
+    PLIST_ENTRY ListEntry;
+    PMM_SESSION_SPACE Session;
+    PEPROCESS Process = NULL;
+    KIRQL OldIrql;
+
+    /* Acquire the expansion lock while touching the session */
+    OldIrql = MiAcquireExpansionLock();
+
+    /* Loop all entries in the session ws list */
+    ListEntry = MiSessionWsList.Flink;
+    while (ListEntry != &MiSessionWsList)
+    {
+        Session = CONTAINING_RECORD(ListEntry, MM_SESSION_SPACE, WsListEntry);
+
+        /* Check if this is the session we are looking for */
+        if (Session->SessionId == SessionId)
+        {
+            /* Check if we also have a process in the process list */
+            if (!IsListEmpty(&Session->ProcessList))
+            {
+                Process = CONTAINING_RECORD(Session->ProcessList.Flink,
+                                            EPROCESS,
+                                            SessionProcessLinks);
+
+                /* Reference the process */
+                ObReferenceObject(Process);
+                break;
+            }
+        }
+    }
+
+    /* Release the lock again */
+    MiReleaseExpansionLock(OldIrql);
+
+    return Process;
+}

Modified: trunk/reactos/ntoskrnl/mm/mminit.c
URL:
http://svn.reactos.org/svn/reactos/trunk/reactos/ntoskrnl/mm/mminit.c?rev=6…
==============================================================================
--- trunk/reactos/ntoskrnl/mm/mminit.c	[iso-8859-1] (original)
+++ trunk/reactos/ntoskrnl/mm/mminit.c	[iso-8859-1] Wed Nov 20 23:50:42 2013
@@ -450,6 +450,9 @@
                                 PageFrameNumber);
     *MmSharedUserDataPte = TempPte;
 
+    /* Initialize session working set support */
+    MiInitializeSessionWsSupport();
+
     /* Setup session IDs */
     MiInitializeSessionIds();