Some improvements

.gitignore

@@ -26,6 +26,9 @@
 *.a
 *.lib
 
+# Auto-fetched at build time (see Example.vcxproj FetchInfoDb target)
+info.db
+
 # Executables
 *.exe
 *.out

.gitmodules

@@ -0,0 +1,3 @@
+[submodule "extern/MemProcFS"]
+	path = extern/MemProcFS
+	url = https://github.com/ufrisk/MemProcFS.git

DMALibrary/DMALibrary.vcxproj

@@ -30,26 +30,26 @@
   <PropertyGroup Condition="'$(Configuration)|$(Platform)'=='Debug|Win32'" Label="Configuration">
     <ConfigurationType>StaticLibrary</ConfigurationType>
     <UseDebugLibraries>true</UseDebugLibraries>
-    <PlatformToolset>v143</PlatformToolset>
+    <PlatformToolset>v145</PlatformToolset>
     <CharacterSet>Unicode</CharacterSet>
   </PropertyGroup>
   <PropertyGroup Condition="'$(Configuration)|$(Platform)'=='Release|Win32'" Label="Configuration">
     <ConfigurationType>StaticLibrary</ConfigurationType>
     <UseDebugLibraries>false</UseDebugLibraries>
-    <PlatformToolset>v143</PlatformToolset>
+    <PlatformToolset>v145</PlatformToolset>
     <WholeProgramOptimization>true</WholeProgramOptimization>
     <CharacterSet>Unicode</CharacterSet>
   </PropertyGroup>
   <PropertyGroup Condition="'$(Configuration)|$(Platform)'=='Debug|x64'" Label="Configuration">
     <ConfigurationType>StaticLibrary</ConfigurationType>
     <UseDebugLibraries>true</UseDebugLibraries>
-    <PlatformToolset>v143</PlatformToolset>
+    <PlatformToolset>v145</PlatformToolset>
     <CharacterSet>Unicode</CharacterSet>
   </PropertyGroup>
   <PropertyGroup Condition="'$(Configuration)|$(Platform)'=='Release|x64'" Label="Configuration">
     <ConfigurationType>StaticLibrary</ConfigurationType>
     <UseDebugLibraries>false</UseDebugLibraries>
-    <PlatformToolset>v143</PlatformToolset>
+    <PlatformToolset>v145</PlatformToolset>
     <WholeProgramOptimization>true</WholeProgramOptimization>
     <CharacterSet>Unicode</CharacterSet>
   </PropertyGroup>
@@ -74,6 +74,9 @@
   <PropertyGroup Condition="'$(Configuration)|$(Platform)'=='Release|x64'">
     <OutDir>$(SolutionDir)lib\</OutDir>
   </PropertyGroup>
+  <PropertyGroup Condition="'$(Configuration)|$(Platform)'=='Debug|x64'">
+    <OutDir>$(SolutionDir)lib\Debug\</OutDir>
+  </PropertyGroup>
   <ItemDefinitionGroup Condition="'$(Configuration)|$(Platform)'=='Debug|Win32'">
     <ClCompile>
       <WarningLevel>Level3</WarningLevel>
@@ -82,6 +85,7 @@
       <ConformanceMode>true</ConformanceMode>
       <PrecompiledHeader>Use</PrecompiledHeader>
       <PrecompiledHeaderFile>pch.h</PrecompiledHeaderFile>
+      <AdditionalIncludeDirectories>$(SolutionDir)extern\MemProcFS\includes;$(ProjectDir);%(AdditionalIncludeDirectories)</AdditionalIncludeDirectories>
     </ClCompile>
     <Link>
       <SubSystem>
@@ -99,6 +103,7 @@
       <ConformanceMode>true</ConformanceMode>
       <PrecompiledHeader>Use</PrecompiledHeader>
       <PrecompiledHeaderFile>pch.h</PrecompiledHeaderFile>
+      <AdditionalIncludeDirectories>$(SolutionDir)extern\MemProcFS\includes;$(ProjectDir);%(AdditionalIncludeDirectories)</AdditionalIncludeDirectories>
     </ClCompile>
     <Link>
       <SubSystem>
@@ -116,13 +121,17 @@
       <ConformanceMode>true</ConformanceMode>
       <PrecompiledHeader>Use</PrecompiledHeader>
       <PrecompiledHeaderFile>pch.h</PrecompiledHeaderFile>
-      <AdditionalIncludeDirectories>$(ProjectDir)</AdditionalIncludeDirectories>
+      <AdditionalIncludeDirectories>$(SolutionDir)extern\MemProcFS\includes;$(ProjectDir)</AdditionalIncludeDirectories>
+      <LanguageStandard>stdcpp20</LanguageStandard>
     </ClCompile>
     <Link>
       <SubSystem>
       </SubSystem>
       <GenerateDebugInformation>true</GenerateDebugInformation>
     </Link>
+    <Lib>
+      <AdditionalDependencies>$(SolutionDir)extern\MemProcFS\includes\lib64\leechcore.lib;$(SolutionDir)extern\MemProcFS\includes\lib64\vmm.lib;%(AdditionalDependencies)</AdditionalDependencies>
+    </Lib>
   </ItemDefinitionGroup>
   <ItemDefinitionGroup Condition="'$(Configuration)|$(Platform)'=='Release|x64'">
     <ClCompile>
@@ -135,7 +144,7 @@
       <PrecompiledHeader>Use</PrecompiledHeader>
       <PrecompiledHeaderFile>pch.h</PrecompiledHeaderFile>
       <LanguageStandard>stdcpp20</LanguageStandard>
-      <AdditionalIncludeDirectories>$(SolutionDir);$(ProjectDir);%(AdditionalIncludeDirectories)</AdditionalIncludeDirectories>
+      <AdditionalIncludeDirectories>$(SolutionDir)extern\MemProcFS\includes;$(SolutionDir);$(ProjectDir);%(AdditionalIncludeDirectories)</AdditionalIncludeDirectories>
       <MultiProcessorCompilation>true</MultiProcessorCompilation>
       <LanguageStandard_C>Default</LanguageStandard_C>
     </ClCompile>
@@ -147,7 +156,7 @@
       <GenerateDebugInformation>true</GenerateDebugInformation>
     </Link>
     <Lib>
-      <AdditionalDependencies>$(ProjectDir)libs\leechcore.lib;$(ProjectDir)libs\vmm.lib;%(AdditionalDependencies)</AdditionalDependencies>
+      <AdditionalDependencies>$(SolutionDir)extern\MemProcFS\includes\lib64\leechcore.lib;$(SolutionDir)extern\MemProcFS\includes\lib64\vmm.lib;%(AdditionalDependencies)</AdditionalDependencies>
     </Lib>
   </ItemDefinitionGroup>
   <ItemGroup>

DMALibrary/Memory/InputManager.cpp

@@ -134,13 +134,12 @@ bool c_keys::InitKeyboard()
 				return false;
 			}
 
-			uintptr_t gafAsyncKeyState;
-			if (!VMMDLL_PdbSymbolAddress(mem.vHandle, str, const_cast<LPSTR>("gafAsyncKeyState"), &gafAsyncKeyState))
+			if (!VMMDLL_PdbSymbolAddress(mem.vHandle, str, const_cast<LPSTR>("gafAsyncKeyState"), &gafAsyncKeyStateExport))
 			{
 				LOG("failed to find gafAsyncKeyState\n");
 				return false;
 			}
-			LOG("found gafAsyncKeyState at: 0x%p\n", gafAsyncKeyState);
+			LOG("found gafAsyncKeyState at: 0x%p\n", gafAsyncKeyStateExport);
 		}
 		if (gafAsyncKeyStateExport > 0x7FFFFFFFFFFF)
 			return true;

DMALibrary/Memory/Memory.cpp

@@ -4,6 +4,24 @@
 #include <thread>
 #include <iostream>
 
+namespace
+{
+	bool IsValidScatterHandle(VMMDLL_SCATTER_HANDLE handle)
+	{
+		return handle && handle != INVALID_HANDLE_VALUE;
+	}
+
+	bool IsValidRemoteAddress(uint64_t address)
+	{
+		return address && address != static_cast<uint64_t>(-1);
+	}
+
+	bool IsValidLocalBuffer(void* buffer)
+	{
+		return buffer && buffer != INVALID_HANDLE_VALUE;
+	}
+}
+
 Memory::Memory()
 {
 	LOG("loading libraries...\n");
@@ -16,7 +34,7 @@ Memory::Memory()
 		LOG("vmm: %p\n", modules.VMM);
 		LOG("ftd: %p\n", modules.FTD3XX);
 		LOG("leech: %p\n", modules.LEECHCORE);
-		THROW("[!] Could not load a library\n");
+		printf("[!] Could not load a library\n");
 	}
 
 	this->key = std::make_shared<c_keys>();
@@ -689,6 +707,11 @@ uint64_t Memory::FindSignature(const char* signature, uint64_t range_start, uint
 
 bool Memory::Write(uintptr_t address, void* buffer, size_t size) const
 {
+	if (!IsValidRemoteAddress(address) || !IsValidLocalBuffer(buffer))
+	{
+		LOG("[!] Failed to write Memory at 0x%p\n", address);
+		return false;
+	}
 	if (!VMMDLL_MemWrite(this->vHandle, current_process.PID, address, static_cast<PBYTE>(buffer), size))
 	{
 		LOG("[!] Failed to write Memory at 0x%p\n", address);
@@ -699,6 +722,11 @@ bool Memory::Write(uintptr_t address, void* buffer, size_t size) const
 
 bool Memory::Write(uintptr_t address, void* buffer, size_t size, int pid) const
 {
+	if (!IsValidRemoteAddress(address) || !IsValidLocalBuffer(buffer))
+	{
+		LOG("[!] Failed to write Memory at 0x%p\n", address);
+		return false;
+	}
 	if (!VMMDLL_MemWrite(this->vHandle, pid, address, static_cast<PBYTE>(buffer), size))
 	{
 		LOG("[!] Failed to write Memory at 0x%p\n", address);
@@ -709,6 +737,11 @@ bool Memory::Write(uintptr_t address, void* buffer, size_t size, int pid) const
 
 bool Memory::Read(uintptr_t address, void* buffer, size_t size) const
 {
+	if (!IsValidRemoteAddress(address) || !IsValidLocalBuffer(buffer))
+	{
+		LOG("[!] Failed to read Memory at 0x%p\n", address);
+		return false;
+	}
 	DWORD read_size = 0;
 	if (!VMMDLL_MemReadEx(this->vHandle, current_process.PID, address, static_cast<PBYTE>(buffer), size, &read_size, VMMDLL_FLAG_NOCACHE))
 	{
@@ -721,6 +754,11 @@ bool Memory::Read(uintptr_t address, void* buffer, size_t size) const
 
 bool Memory::Read(uintptr_t address, void* buffer, size_t size, int pid) const
 {
+	if (!IsValidRemoteAddress(address) || !IsValidLocalBuffer(buffer))
+	{
+		LOG("[!] Failed to read Memory at 0x%p\n", address);
+		return false;
+	}
 	DWORD read_size = 0;
 	if (!VMMDLL_MemReadEx(this->vHandle, pid, address, static_cast<PBYTE>(buffer), size, &read_size, VMMDLL_FLAG_NOCACHE))
 	{
@@ -730,11 +768,37 @@ bool Memory::Read(uintptr_t address, void* buffer, size_t size, int pid) const
 	return (read_size == size);
 }
 
+std::shared_ptr<Memory::ScatterPending> Memory::GetOrCreateScatterPending(VMMDLL_SCATTER_HANDLE handle) const
+{
+	if (!IsValidScatterHandle(handle))
+		return nullptr;
+
+	std::lock_guard<std::mutex> lock(scatter_pending_mutex);
+	auto& pending = scatter_pending[handle];
+	if (!pending)
+		pending = std::make_shared<ScatterPending>();
+	return pending;
+}
+
+std::shared_ptr<Memory::ScatterPending> Memory::FindScatterPending(VMMDLL_SCATTER_HANDLE handle) const
+{
+	if (!IsValidScatterHandle(handle))
+		return nullptr;
+
+	std::lock_guard<std::mutex> lock(scatter_pending_mutex);
+	const auto it = scatter_pending.find(handle);
+	if (it == scatter_pending.end())
+		return nullptr;
+	return it->second;
+}
+
 VMMDLL_SCATTER_HANDLE Memory::CreateScatterHandle() const
 {
 	const VMMDLL_SCATTER_HANDLE ScatterHandle = VMMDLL_Scatter_Initialize(this->vHandle, current_process.PID, VMMDLL_FLAG_NOCACHE);
 	if (!ScatterHandle)
 		LOG("[!] Failed to create scatter handle\n");
+	else
+		GetOrCreateScatterPending(ScatterHandle);
 	return ScatterHandle;
 }
 
@@ -743,32 +807,67 @@ VMMDLL_SCATTER_HANDLE Memory::CreateScatterHandle(int pid) const
 	const VMMDLL_SCATTER_HANDLE ScatterHandle = VMMDLL_Scatter_Initialize(this->vHandle, pid, VMMDLL_FLAG_NOCACHE);
 	if (!ScatterHandle)
 		LOG("[!] Failed to create scatter handle\n");
+	else
+		GetOrCreateScatterPending(ScatterHandle);
 	return ScatterHandle;
 }
 
 void Memory::CloseScatterHandle(VMMDLL_SCATTER_HANDLE handle)
 {
+	if (!IsValidScatterHandle(handle))
+		return;
+
+	{
+		std::lock_guard<std::mutex> lock(scatter_pending_mutex);
+		scatter_pending.erase(handle);
+	}
 	VMMDLL_Scatter_CloseHandle(handle);
 }
 
-void Memory::AddScatterReadRequest(VMMDLL_SCATTER_HANDLE handle, uint64_t address, void* buffer, size_t size)
+bool Memory::AddScatterReadRequest(VMMDLL_SCATTER_HANDLE handle, uint64_t address, void* buffer, size_t size)
 {
+	if (!IsValidScatterHandle(handle) || !IsValidRemoteAddress(address) || !IsValidLocalBuffer(buffer))
+	{
+		LOG("[!] Failed to prepare scatter read at 0x%p\n", address);
+		return false;
+	}
 	if (!VMMDLL_Scatter_PrepareEx(handle, address, size, static_cast<PBYTE>(buffer), NULL))
 	{
 		LOG("[!] Failed to prepare scatter read at 0x%p\n", address);
+		return false;
 	}
+	const auto pending = GetOrCreateScatterPending(handle);
+	if (!pending)
+		return false;
+	pending->reads.fetch_add(1, std::memory_order_relaxed);
+	return true;
 }
 
-void Memory::AddScatterWriteRequest(VMMDLL_SCATTER_HANDLE handle, uint64_t address, void* buffer, size_t size)
+bool Memory::AddScatterWriteRequest(VMMDLL_SCATTER_HANDLE handle, uint64_t address, void* buffer, size_t size)
 {
+	if (!IsValidScatterHandle(handle) || !IsValidRemoteAddress(address) || !IsValidLocalBuffer(buffer))
+	{
+		LOG("[!] Failed to prepare scatter write at 0x%p\n", address);
+		return false;
+	}
 	if (!VMMDLL_Scatter_PrepareWrite(handle, address, static_cast<PBYTE>(buffer), size))
 	{
 		LOG("[!] Failed to prepare scatter write at 0x%p\n", address);
+		return false;
 	}
+	const auto pending = GetOrCreateScatterPending(handle);
+	if (!pending)
+		return false;
+	pending->writes.fetch_add(1, std::memory_order_relaxed);
+	return true;
 }
 
 void Memory::ExecuteReadScatter(VMMDLL_SCATTER_HANDLE handle, int pid)
 {
+	const auto pending = FindScatterPending(handle);
+	if (!pending || pending->reads.load(std::memory_order_relaxed) == 0)
+		return;
+
 	if (pid == 0)
 		pid = current_process.PID;
 
@@ -781,20 +880,30 @@ void Memory::ExecuteReadScatter(VMMDLL_SCATTER_HANDLE handle, int pid)
 	{
 		LOG("[-] Failed to clear Scatter\n");
 	}
+
+	pending->reads.store(0, std::memory_order_relaxed);
+	pending->writes.store(0, std::memory_order_relaxed);
 }
 
 void Memory::ExecuteWriteScatter(VMMDLL_SCATTER_HANDLE handle, int pid)
 {
+	const auto pending = FindScatterPending(handle);
+	if (!pending || pending->writes.load(std::memory_order_relaxed) == 0)
+		return;
+
 	if (pid == 0)
 		pid = current_process.PID;
 
 	if (!VMMDLL_Scatter_Execute(handle))
 	{
-		LOG("[-] Failed to Execute Scatter Read\n");
+		LOG("[-] Failed to Execute Scatter Write\n");
 	}
 	//Clear after using it
 	if (!VMMDLL_Scatter_Clear(handle, pid, VMMDLL_FLAG_NOCACHE))
 	{
 		LOG("[-] Failed to clear Scatter\n");
 	}
+
+	pending->reads.store(0, std::memory_order_relaxed);
+	pending->writes.store(0, std::memory_order_relaxed);
 }