[PATCH] ntoskrnl.exe: Implement MmMapLockedPagesSpecifyCache & MmUnmapLockedPages and improve MmUnlockPages stub. (try 2)

[Christian Costa]

Try 2: Fix for 64-bit.
---
 dlls/ntoskrnl.exe/ntoskrnl.c        |   82 ++++++++++++++++++++++++++++++++---
 dlls/ntoskrnl.exe/ntoskrnl.exe.spec |    2 -
 2 files changed, 77 insertions(+), 7 deletions(-)

diff --git a/dlls/ntoskrnl.exe/ntoskrnl.c b/dlls/ntoskrnl.exe/ntoskrnl.c
index 3f1869b..78e9d3a 100644
--- a/dlls/ntoskrnl.exe/ntoskrnl.c
+++ b/dlls/ntoskrnl.exe/ntoskrnl.c
@@ -1509,12 +1509,48 @@ PVOID WINAPI MmMapIoSpace( PHYSICAL_ADDRESS PhysicalAddress, DWORD NumberOfBytes
 /***********************************************************************
  *           MmMapLockedPagesSpecifyCache  (NTOSKRNL.EXE.@)
  */
-PVOID WINAPI  MmMapLockedPagesSpecifyCache(PMDLX MemoryDescriptorList, KPROCESSOR_MODE AccessMode, MEMORY_CACHING_TYPE CacheType,
-                                           PVOID BaseAddress, ULONG BugCheckOnFailure, MM_PAGE_PRIORITY Priority)
+PVOID WINAPI MmMapLockedPagesSpecifyCache(PMDLX mdl, KPROCESSOR_MODE access, MEMORY_CACHING_TYPE type,
+                                          PVOID base_address, ULONG BugCheckOnFailure, MM_PAGE_PRIORITY priority)
 {
-    FIXME("(%p, %u, %u, %p, %u, %u): stub\n", MemoryDescriptorList, AccessMode, CacheType, BaseAddress, BugCheckOnFailure, Priority);
+    HANDLE process;
+    SIZE_T bytes;
 
-    return NULL;
+    TRACE("(%p, %u, %u, %p, %u, %u)\n", mdl, access, type, base_address, BugCheckOnFailure, priority);
+
+    process = OpenProcess(PROCESS_VM_READ, FALSE, (DWORD)mdl->Process);
+    if (!process)
+    {
+        ERR("Cannot open process\n");
+        return NULL;
+    }
+
+    /* Allocate a memory area aligned on pages boundary that contains the buffer described by the mdl */
+    mdl->MappedSystemVa = VirtualAlloc(NULL, mdl->ByteOffset + mdl->ByteCount, MEM_RESERVE | MEM_COMMIT, PAGE_READWRITE);
+    if (!mdl->MappedSystemVa)
+    {
+        ERR("Cannot alloc memory\n");
+        CloseHandle(process);
+        return NULL;
+    }
+
+    /* Read memory from the client process memory */
+    if (!ReadProcessMemory(process, (LPCVOID)((ULONG_PTR)(mdl->StartVa) + mdl->ByteOffset), (LPVOID)((ULONG_PTR)mdl->MappedSystemVa + mdl->ByteOffset), mdl->ByteCount, &bytes))
+    {
+        ERR("Cannot read memory (%x)\n", GetLastError());
+        CloseHandle(process);
+        VirtualFree(mdl->MappedSystemVa, 0, MEM_RELEASE);
+        mdl->MappedSystemVa = NULL;
+        return NULL;
+    }
+
+    if (bytes != mdl->ByteCount)
+        WARN("Bytes count is not correct %d instead of %d\n", (DWORD)bytes,  mdl->ByteCount);
+
+    CloseHandle(process);
+
+    TRACE("Returned mapped pages starting address %p\n", mdl->MappedSystemVa);
+
+    return mdl->MappedSystemVa;
 }
 
 
@@ -1547,11 +1583,45 @@ void WINAPI MmResetDriverPaging(PVOID AddrInSection)
 
 
 /***********************************************************************
+ *           MmUnmapLockedPages   (NTOSKRNL.EXE.@)
+ */
+void WINAPI MmUnmapLockedPages(PVOID base_address, PMDL mdl)
+{
+    HANDLE process;
+    SIZE_T bytes;
+
+    TRACE("(%p, %p)\n", base_address, mdl);
+
+    process = OpenProcess(PROCESS_VM_WRITE, FALSE, (DWORD)mdl->Process);
+    if (!process)
+    {
+        ERR("Cannot open process\n");
+        return;
+    }
+
+    /* Write back data modified by the driver to the client process memory */
+    if (!WriteProcessMemory(process, (LPVOID)((ULONG_PTR)(mdl->StartVa) + mdl->ByteOffset), (LPVOID)((ULONG_PTR)mdl->MappedSystemVa + mdl->ByteOffset), mdl->ByteCount, &bytes))
+        ERR("Cannot write memory (%x)\n", GetLastError());
+
+    if (bytes != mdl->ByteCount)
+        WARN("Bytes count is not correct %d instead of %d\n", (DWORD)bytes,  mdl->ByteCount);
+
+    CloseHandle(process);
+    VirtualFree(mdl->MappedSystemVa, 0, MEM_RELEASE);
+    mdl->MappedSystemVa = NULL;
+}
+
+
+/***********************************************************************
  *           MmUnlockPages  (NTOSKRNL.EXE.@)
  */
-void WINAPI  MmUnlockPages(PMDLX MemoryDescriptorList)
+void WINAPI MmUnlockPages(PMDLX mdl)
 {
-    FIXME("(%p): stub\n", MemoryDescriptorList);
+    FIXME("(%p): partial stub %p\n", mdl, mdl->MappedSystemVa);
+
+    /* Some drivers call MmUnlockPages with unmapping locked pages with MmUnmapLockedPages so do it here */
+    if (mdl->MappedSystemVa)
+        MmUnmapLockedPages(mdl->MappedSystemVa, mdl);
 }
 
 
diff --git a/dlls/ntoskrnl.exe/ntoskrnl.exe.spec b/dlls/ntoskrnl.exe/ntoskrnl.exe.spec
index 4bce7ec..f3f9ba4 100644
--- a/dlls/ntoskrnl.exe/ntoskrnl.exe.spec
+++ b/dlls/ntoskrnl.exe/ntoskrnl.exe.spec
@@ -719,7 +719,7 @@
 @ stub MmUnlockPagableImageSection
 @ stdcall MmUnlockPages(ptr)
 @ stdcall MmUnmapIoSpace(ptr long)
-@ stub MmUnmapLockedPages
+@ stdcall MmUnmapLockedPages(ptr ptr)
 @ stub MmUnmapReservedMapping
 @ stub MmUnmapVideoDisplay
 @ stub MmUnmapViewInSessionSpace