[NTOS/MM] Implement Mi(Un)lockWorkingSetShared

Also fix checks when exclusively locking

ntoskrnl/mm/ARM3/miarm.h

@@ -1276,6 +1276,48 @@ MiLockWorkingSet(IN PETHREAD Thread,
     }
 }
 
+FORCEINLINE
+VOID
+MiLockWorkingSetShared(
+    _In_ PETHREAD Thread,
+    _In_ PMMSUPPORT WorkingSet)
+{
+    /* Block APCs */
+    KeEnterGuardedRegion();
+
+    /* Working set should be in global memory */
+    ASSERT(MI_IS_SESSION_ADDRESS((PVOID)WorkingSet) == FALSE);
+
+    /* Thread shouldn't already be owning something */
+    ASSERT(!MM_ANY_WS_LOCK_HELD(Thread));
+
+    /* Lock this working set */
+    ExAcquirePushLockShared(&WorkingSet->WorkingSetMutex);
+
+    /* Which working set is this? */
+    if (WorkingSet == &MmSystemCacheWs)
+    {
+        /* Own the system working set */
+        ASSERT((Thread->OwnsSystemWorkingSetExclusive == FALSE) &&
+               (Thread->OwnsSystemWorkingSetShared == FALSE));
+        Thread->OwnsSystemWorkingSetShared = TRUE;
+    }
+    else if (WorkingSet->Flags.SessionSpace)
+    {
+        /* Own the session working set */
+        ASSERT((Thread->OwnsSessionWorkingSetExclusive == FALSE) &&
+               (Thread->OwnsSessionWorkingSetShared == FALSE));
+        Thread->OwnsSessionWorkingSetShared = TRUE;
+    }
+    else
+    {
+        /* Own the process working set */
+        ASSERT((Thread->OwnsProcessWorkingSetExclusive == FALSE) &&
+               (Thread->OwnsProcessWorkingSetShared == FALSE));
+        Thread->OwnsProcessWorkingSetShared = TRUE;
+    }
+}
+
 //
 // Unlocks the working set
 //
@@ -1291,22 +1333,22 @@ MiUnlockWorkingSet(IN PETHREAD Thread,
     if (WorkingSet == &MmSystemCacheWs)
     {
         /* Release the system working set */
-        ASSERT((Thread->OwnsSystemWorkingSetExclusive == TRUE) ||
+        ASSERT((Thread->OwnsSystemWorkingSetExclusive == TRUE) &&
-               (Thread->OwnsSystemWorkingSetShared == TRUE));
+               (Thread->OwnsSystemWorkingSetShared == FALSE));
         Thread->OwnsSystemWorkingSetExclusive = FALSE;
     }
     else if (WorkingSet->Flags.SessionSpace)
     {
         /* Release the session working set */
-        ASSERT((Thread->OwnsSessionWorkingSetExclusive == TRUE) ||
+        ASSERT((Thread->OwnsSessionWorkingSetExclusive == TRUE) &&
-               (Thread->OwnsSessionWorkingSetShared == TRUE));
+               (Thread->OwnsSessionWorkingSetShared == FALSE));
-        Thread->OwnsSessionWorkingSetExclusive = 0;
+        Thread->OwnsSessionWorkingSetExclusive = FALSE;
     }
     else
     {
         /* Release the process working set */
-        ASSERT((Thread->OwnsProcessWorkingSetExclusive) ||
+        ASSERT((Thread->OwnsProcessWorkingSetExclusive == TRUE) &&
-               (Thread->OwnsProcessWorkingSetShared));
+               (Thread->OwnsProcessWorkingSetShared == FALSE));
         Thread->OwnsProcessWorkingSetExclusive = FALSE;
     }
 
@@ -1317,6 +1359,45 @@ MiUnlockWorkingSet(IN PETHREAD Thread,
     KeLeaveGuardedRegion();
 }
 
+FORCEINLINE
+VOID
+MiUnlockWorkingSetShared(
+    _In_ PETHREAD Thread,
+    _In_ PMMSUPPORT WorkingSet)
+{
+    /* Working set should be in global memory */
+    ASSERT(MI_IS_SESSION_ADDRESS((PVOID)WorkingSet) == FALSE);
+
+    /* Which working set is this? */
+    if (WorkingSet == &MmSystemCacheWs)
+    {
+        /* Release the system working set */
+        ASSERT((Thread->OwnsSystemWorkingSetExclusive == FALSE) &&
+               (Thread->OwnsSystemWorkingSetShared == TRUE));
+        Thread->OwnsSystemWorkingSetShared = FALSE;
+    }
+    else if (WorkingSet->Flags.SessionSpace)
+    {
+        /* Release the session working set */
+        ASSERT((Thread->OwnsSessionWorkingSetExclusive == FALSE) &&
+               (Thread->OwnsSessionWorkingSetShared == TRUE));
+        Thread->OwnsSessionWorkingSetShared = FALSE;
+    }
+    else
+    {
+        /* Release the process working set */
+        ASSERT((Thread->OwnsProcessWorkingSetExclusive == FALSE) &&
+               (Thread->OwnsProcessWorkingSetShared == TRUE));
+        Thread->OwnsProcessWorkingSetShared = FALSE;
+    }
+
+    /* Release the working set lock */
+    ExReleasePushLockShared(&WorkingSet->WorkingSetMutex);
+
+    /* Unblock APCs */
+    KeLeaveGuardedRegion();
+}
+
 FORCEINLINE
 VOID
 MiUnlockProcessWorkingSetForFault(IN PEPROCESS Process,