From 5d96ba9217416e955641b7282acc62a0ed9ca277 Mon Sep 17 00:00:00 2001
From: Oleg Dubinskiy <oleg.dubinskij30@gmail.com>
Date: Thu, 8 Aug 2024 21:02:35 +0200
Subject: [PATCH] [NTOS:MM] Implement MmProbeAndLockProcessPages (#7221)

Implement undocumented MmProbeAndLockProcessPages routine. Based on mm-implement-mappingaddress.patch by Thomas Faber from CORE-10147, with some improvements from me.
It's badly required by FltMgr.sys driver from Windows XP/Server 2003 and closely used by a lot of apps those are depending on this driver (e. g., Avast Free Antivirus several versions, Avira Antivir Personal 8.2 etc. etc.).
Fixes several asserts from MDL support routines when the 3rd-party minifilter drivers are loading FltMgr.
CORE-14157
---
 ntoskrnl/mm/ARM3/mdlsup.c | 56 +++++++++++++++++++++++++++++++++------
 1 file changed, 48 insertions(+), 8 deletions(-)

diff --git a/ntoskrnl/mm/ARM3/mdlsup.c b/ntoskrnl/mm/ARM3/mdlsup.c
index cc6d2ffb4ec..76b06f0c72e 100644
--- a/ntoskrnl/mm/ARM3/mdlsup.c
+++ b/ntoskrnl/mm/ARM3/mdlsup.c
@@ -1671,19 +1671,59 @@ MmProtectMdlSystemAddress(IN PMDL MemoryDescriptorList,
     return STATUS_NOT_IMPLEMENTED;
 }
 
-/*
- * @unimplemented
+/**
+ * @brief
+ * Probes and locks virtual pages in memory for the specified process.
+ *
+ * @param[in,out] MemoryDescriptorList
+ * Memory Descriptor List (MDL) containing the buffer to be probed and locked.
+ *
+ * @param[in] Process
+ * The process for which the buffer should be probed and locked.
+ *
+ * @param[in] AccessMode
+ * Access mode for probing the pages. Can be KernelMode or UserMode.
+ *
+ * @param[in] LockOperation
+ * The type of the probing and locking operation. Can be IoReadAccess, IoWriteAccess or IoModifyAccess.
+ *
+ * @return
+ * Nothing.
+ *
+ * @see MmProbeAndLockPages
+ *
+ * @remarks Must be called at IRQL <= APC_LEVEL
  */
+_IRQL_requires_max_(APC_LEVEL)
 VOID
 NTAPI
-MmProbeAndLockProcessPages(IN OUT PMDL MemoryDescriptorList,
-                           IN PEPROCESS Process,
-                           IN KPROCESSOR_MODE AccessMode,
-                           IN LOCK_OPERATION Operation)
+MmProbeAndLockProcessPages(
+    _Inout_ PMDL MemoryDescriptorList,
+    _In_ PEPROCESS Process,
+    _In_ KPROCESSOR_MODE AccessMode,
+    _In_ LOCK_OPERATION Operation)
 {
-    UNIMPLEMENTED;
-}
+    KAPC_STATE ApcState;
+    BOOLEAN IsAttached = FALSE;
 
+    if (Process != PsGetCurrentProcess())
+    {
+        KeStackAttachProcess(&Process->Pcb, &ApcState);
+        IsAttached = TRUE;
+    }
+
+    /* Protect in try/finally to ensure we detach even if MmProbeAndLockPages() throws an exception */
+    _SEH2_TRY
+    {
+        MmProbeAndLockPages(MemoryDescriptorList, AccessMode, Operation);
+    }
+    _SEH2_FINALLY
+    {
+        if (IsAttached)
+            KeUnstackDetachProcess(&ApcState);
+    }
+    _SEH2_END;
+}
 
 /*
  * @unimplemented