From ccfc2057ab52354de82913a35d7559d1b0dd3618 Mon Sep 17 00:00:00 2001
From: BlackINT3 <blackint3@gmail.com>
Date: Tue, 25 Aug 2020 23:32:50 +0800
Subject: [PATCH] Fix memory write bugs.

---
 src/OpenArk/kernel/object/object.cpp          |  4 ++--
 .../arkdrv-api/api-memory/api-memory.cpp      | 22 ++++++++++++++-----
 .../arkdrv-api/api-object/api-object.cpp      |  2 --
 3 files changed, 19 insertions(+), 9 deletions(-)

diff --git a/src/OpenArk/kernel/object/object.cpp b/src/OpenArk/kernel/object/object.cpp
index 8efc228..cfa58a9 100644
--- a/src/OpenArk/kernel/object/object.cpp
+++ b/src/OpenArk/kernel/object/object.cpp
@@ -140,9 +140,9 @@ void KernelObject::InitObjectSectionsView()
 			prefix = L"";
 			map_name = section_name;
 		}
-		map_hd = OpenFileMappingW(FILE_MAP_READ, FALSE, map_name.c_str());
+		map_hd = OpenFileMappingW(FILE_MAP_READ|FILE_MAP_WRITE, FALSE, map_name.c_str());
 		if (map_hd) {
-			map_addr = (ULONG64)MapViewOfFileEx(map_hd, FILE_MAP_READ, 0, 0, size, NULL);
+			map_addr = (ULONG64)MapViewOfFileEx(map_hd, FILE_MAP_READ | FILE_MAP_WRITE, 0, 0, size, NULL);
 			if (!map_addr) {
 				CloseHandle(map_hd);
 				return;
diff --git a/src/OpenArkDrv/arkdrv-api/api-memory/api-memory.cpp b/src/OpenArkDrv/arkdrv-api/api-memory/api-memory.cpp
index 47aa7d1..bb374ef 100644
--- a/src/OpenArkDrv/arkdrv-api/api-memory/api-memory.cpp
+++ b/src/OpenArkDrv/arkdrv-api/api-memory/api-memory.cpp
@@ -69,8 +69,12 @@ bool MemoryRead(ULONG pid, ULONG64 addr, ULONG size, std::string &readbuf)
 	std::string data;
 	data.resize(size);
 	BOOL ret = ReadProcessMemory(phd, (PVOID)addr, (PVOID)data.data(), (SIZE_T)size, (SIZE_T*)&readlen);
+	if (!ret) {
+		ERR(L"ReadProcessMemory pid:%d, err:%d", pid, GetLastError());
+		CloseHandle(phd);
+		return FALSE;
+	}
 	CloseHandle(phd);
-	if (!ret) return FALSE;
 
 	readbuf = std::move(data);
 	return true;
@@ -105,14 +109,22 @@ bool MemoryWrite(ULONG pid, ULONG64 addr, std::string &writebuf)
 	PVOID buf = (PVOID)writebuf.data();
 	SIZE_T bufsize = (SIZE_T)writebuf.size();
 	DWORD written, oldprotect;
-	VirtualProtectEx(phd, (PVOID)addr, bufsize, PAGE_READWRITE, &oldprotect);
-	BOOL ret = WriteProcessMemory(phd, (PVOID)addr, buf, bufsize, (SIZE_T*)&written);
-	VirtualProtectEx(phd, (PVOID)addr, bufsize, oldprotect, &oldprotect);
-	CloseHandle(phd);
+	BOOL ret = FALSE;
+	ret = VirtualProtectEx(phd, (PVOID)addr, bufsize, PAGE_READWRITE, &oldprotect);
+	if (!ret) {
+		ERR(L"VirtualProtectEx pid:%d, err:%d", pid, GetLastError());
+		CloseHandle(phd);
+		return FALSE;
+	}
+	ret = WriteProcessMemory(phd, (PVOID)addr, buf, bufsize, (SIZE_T*)&written);
 	if (!ret) {
 		ERR(L"WriteProcessMemory pid:%d, err:%d", pid, GetLastError());
+		VirtualProtectEx(phd, (PVOID)addr, bufsize, oldprotect, &oldprotect);
+		CloseHandle(phd);
 		return FALSE;
 	}
+	VirtualProtectEx(phd, (PVOID)addr, bufsize, oldprotect, &oldprotect);
+	CloseHandle(phd);
 
 	return true;
 }
diff --git a/src/OpenArkDrv/arkdrv-api/api-object/api-object.cpp b/src/OpenArkDrv/arkdrv-api/api-object/api-object.cpp
index 9e75521..4fc41a2 100644
--- a/src/OpenArkDrv/arkdrv-api/api-object/api-object.cpp
+++ b/src/OpenArkDrv/arkdrv-api/api-object/api-object.cpp
@@ -269,8 +269,6 @@ bool ObjectSectionEnumR3(std::vector<ARK_OBJECT_SECTION_ITEM> &items, ULONG sess
 				item.section_size = (ULONG)mbi.RegionSize;
 				UnmapViewOfFile(mapaddr);
 				CloseHandle(maphd);
-			} else {
-				ERR(L"%s %d", map_name.c_str(), GetLastError());
 			}
 			item.session_id = session;
 			items.push_back(item);