[NTOS:MM] Implement MmAllocate/FreeMappingAddress (#7260)

Implement MmAllocateMappingAddress and MmFreeMappingAddress routines.
Based on mm-implement-mappingaddress.patch by Thomas Faber with some changes from me.
Required by Microsoft NTFS driver (from Windows Server 2003 SP2 only, the one from Windows XP SP3 does not need them) and by NDIS & TDI drivers (both from Windows XP SP3 and Windows Server 2003 SP2). Also they are called when using Dr. Web Security Space 8 filter drivers together with MS FltMgr & TDI.
Fortunately, this part (these two routines) are enough to get the drivers working in both cases, and others (partially incomplete) routines are not badly required, so they can be finished and committed later.
CORE-10147, CORE-14635, CORE-17409, CORE-19318

ntoskrnl/mm/ARM3/miarm.h

@@ -154,6 +154,14 @@ C_ASSERT(SYSTEM_PD_SIZE == PAGE_SIZE);
 #error Define these please!
 #endif
 
+//
+// Some internal SYSTEM_PTE_MISUSE bugcheck subcodes
+//
+#define PTE_MAPPING_NONE        0x100
+#define PTE_MAPPING_NOT_OWNED   0x101
+#define PTE_MAPPING_EMPTY       0x102
+#define PTE_MAPPING_RESERVED    0x103
+
 //
 // Mask for image section page protection
 //

ntoskrnl/mm/ARM3/pool.c

@@ -1552,27 +1552,137 @@ MmReturnPoolQuota(
 
 /* PUBLIC FUNCTIONS ***********************************************************/
 
-/*
- * @unimplemented
+/**
+ * @brief
+ * Reserves the specified amount of memory in system virtual address space.
+ *
+ * @param[in] NumberOfBytes
+ * Size, in bytes, of memory to reserve.
+ *
+ * @param[in] PoolTag
+ * Pool Tag identifying the buffer. Usually consists from 4 characters in reversed order.
+ *
+ * @return
+ * A pointer to the 1st memory block of the reserved buffer in case of success, NULL otherwise.
+ *
+ * @remarks Must be called at IRQL <= APC_LEVEL
  */
+_Must_inspect_result_
+_IRQL_requires_max_(APC_LEVEL)
+_Ret_maybenull_
 PVOID
 NTAPI
-MmAllocateMappingAddress(IN SIZE_T NumberOfBytes,
-                         IN ULONG PoolTag)
+MmAllocateMappingAddress(
+    _In_ SIZE_T NumberOfBytes,
+    _In_ ULONG PoolTag)
 {
-    UNIMPLEMENTED;
+    PFN_NUMBER SizeInPages;
+    PMMPTE PointerPte;
+    MMPTE TempPte;
+
+    /* How many PTEs does the caller want? */
+    SizeInPages = BYTES_TO_PAGES(NumberOfBytes);
+    if (SizeInPages == 0)
+    {
+        KeBugCheckEx(SYSTEM_PTE_MISUSE,
+                     PTE_MAPPING_NONE, /* Requested 0 mappings */
+                     SizeInPages,
+                     PoolTag,
+                     (ULONG_PTR)_ReturnAddress());
+    }
+
+    /* We need two extra PTEs to store size and pool tag in */
+    SizeInPages += 2;
+
+    /* Reserve our PTEs */
+    PointerPte = MiReserveSystemPtes(SizeInPages, SystemPteSpace);
+    if (!PointerPte)
+    {
+        /* Failed to reserve PTEs */
+        DPRINT1("Failed to reserve system PTEs\n");
         return NULL;
     }
 
-/*
- * @unimplemented
+    ASSERT(SizeInPages <= MM_EMPTY_PTE_LIST);
+    TempPte.u.Long = 0;
+    TempPte.u.List.NextEntry = SizeInPages;
+    MI_WRITE_INVALID_PTE(&PointerPte[0], TempPte);
+    TempPte.u.Long = PoolTag;
+    TempPte.u.Hard.Valid = 0;
+    MI_WRITE_INVALID_PTE(&PointerPte[1], TempPte);
+    return MiPteToAddress(PointerPte + 2);
+}
+
+/**
+ * @brief
+ * Frees previously reserved amount of memory in system virtual address space.
+ *
+ * @param[in] BaseAddress
+ * A pointer to the 1st memory block of the reserved buffer.
+ *
+ * @param[in] PoolTag
+ * Pool Tag identifying the buffer. Usually consists from 4 characters in reversed order.
+ *
+ * @return
+ * Nothing.
+ *
+ * @see MmAllocateMappingAddress
+ *
+ * @remarks Must be called at IRQL <= APC_LEVEL
  */
+_IRQL_requires_max_(APC_LEVEL)
 VOID
 NTAPI
-MmFreeMappingAddress(IN PVOID BaseAddress,
-                     IN ULONG PoolTag)
+MmFreeMappingAddress(
+    _In_ __drv_freesMem(Mem) _Post_invalid_ PVOID BaseAddress,
+    _In_ ULONG PoolTag)
 {
-    UNIMPLEMENTED;
+    PMMPTE PointerPte;
+    MMPTE TempPte;
+    PFN_NUMBER SizeInPages;
+    PFN_NUMBER i;
+
+    /* Get the first PTE we reserved */
+    PointerPte = MiAddressToPte(BaseAddress) - 2;
+
+    /* Verify that the pool tag matches */
+    TempPte.u.Long = PoolTag;
+    TempPte.u.Hard.Valid = 0;
+    if (PointerPte[1].u.Long != TempPte.u.Long)
+    {
+        KeBugCheckEx(SYSTEM_PTE_MISUSE,
+                     PTE_MAPPING_NOT_OWNED, /* Trying to free an address it does not own */
+                     (ULONG_PTR)BaseAddress,
+                     PoolTag,
+                     PointerPte[1].u.Long);
+    }
+
+    /* We must have a size */
+    SizeInPages = PointerPte[0].u.List.NextEntry;
+    if (SizeInPages < 3)
+    {
+        KeBugCheckEx(SYSTEM_PTE_MISUSE,
+                     PTE_MAPPING_EMPTY, /* Mapping apparently empty */
+                     (ULONG_PTR)BaseAddress,
+                     PoolTag,
+                     (ULONG_PTR)_ReturnAddress());
+    }
+    
+    /* Enumerate all PTEs and make sure they are empty */
+    for (i = 2; i < SizeInPages; i++)
+    {
+        if (PointerPte[i].u.Long != 0)
+        {
+            KeBugCheckEx(SYSTEM_PTE_MISUSE,
+                         PTE_MAPPING_RESERVED, /* Mapping address still reserved */
+                         (ULONG_PTR)PointerPte,
+                         PoolTag,
+                         SizeInPages - 2);
+        }
+    }
+
+    /* Release the PTEs */
+    MiReleaseSystemPtes(PointerPte, SizeInPages, SystemPteSpace);
 }
 
 /* EOF */

sdk/include/xdk/mmfuncs.h

@@ -633,7 +633,7 @@ MmAdvanceMdl(
 
 _Must_inspect_result_
 _IRQL_requires_max_(APC_LEVEL)
-_When_ (return != NULL, _Out_writes_bytes_opt_ (NumberOfBytes))
+_Ret_maybenull_
 NTKERNELAPI
 PVOID
 NTAPI
@@ -646,7 +646,7 @@ NTKERNELAPI
 VOID
 NTAPI
 MmFreeMappingAddress(
-  _In_ PVOID BaseAddress,
+  _In_ __drv_freesMem(Mem) _Post_invalid_ PVOID BaseAddress,
   _In_ ULONG PoolTag);
 
 _IRQL_requires_max_ (APC_LEVEL)