Add memory r/w helper class.

 * add configurable wait prefs during module retrieval

Signed-off-by: Toni Uhlig <matzeton@googlemail.com>

2 files changed, 64 insertions, 6 deletions

diff --git a/memory.cpp b/memory.cpp
index 51c3e1f..366d619 100644
--- a/memory.cpp
+++ b/memory.cpp
@@ -43,6 +43,14 @@ NTSTATUS NTAPI ZwQuerySystemInformation(_In_ int SystemInformationClass,
 NTKERNELAPI PVOID NTAPI PsGetProcessWow64Process(_In_ PEPROCESS Process);
 };
 
+static int g_waitCount = 100;
+static LONGLONG g_waitTimeout = (-1LL) * 10LL * 1000LL * 250LL; // 250ms
+
+void SetLdrInitWaitPrefs(int waitCount, LONGLONG waitTimeout) {
+    g_waitCount = waitCount;
+    g_waitTimeout = waitTimeout;
+}
+
 eastl::vector<Process> GetProcesses() {
   eastl::vector<Process> result;
   ULONG memoryNeeded = 0;
@@ -161,8 +169,8 @@ eastl::vector<Module> GetModules(_In_ PEPROCESS Process, _In_ BOOLEAN isWow64) {
     }
 
     if (!ldr32->Initialized) {
-      while (!ldr32->Initialized && waitCount++ < 4) {
-        LARGE_INTEGER wait = {.QuadPart = -2500};
+      while (!ldr32->Initialized && waitCount++ < g_waitCount) {
+        LARGE_INTEGER wait = {.QuadPart = g_waitTimeout};
         KeDelayExecutionThread(KernelMode, TRUE, &wait);
       }
 
@@ -204,8 +212,8 @@ eastl::vector<Module> GetModules(_In_ PEPROCESS Process, _In_ BOOLEAN isWow64) {
     }
 
     if (!ldr->Initialized) {
-      while (!ldr->Initialized && waitCount++ < 4) {
-        LARGE_INTEGER wait = {.QuadPart = -2500};
+      while (!ldr->Initialized && waitCount++ < g_waitCount) {
+        LARGE_INTEGER wait = {.QuadPart = g_waitTimeout};
         KeDelayExecutionThread(KernelMode, TRUE, &wait);
       }
 
@@ -310,7 +318,7 @@ NTSTATUS ReadVirtualMemory(_In_ PEPROCESS pep, _In_ uint64_t sourceAddress,
   return status;
 }
 
-NTSTATUS WriteVirtualMemory(_In_ PEPROCESS pep, _In_ UCHAR *sourceAddress,
+NTSTATUS WriteVirtualMemory(_In_ PEPROCESS pep, _In_ const UCHAR *sourceAddress,
                             _In_ _Out_ uint64_t targetAddress,
                             _In_ _Out_ SIZE_T *size) {
   NTSTATUS status = STATUS_SUCCESS;
diff --git a/memory.hpp b/memory.hpp
index c22c0a4..37329dc 100644
--- a/memory.hpp
+++ b/memory.hpp
@@ -42,6 +42,7 @@ struct Page {
   uint32_t Type;
 };
 
+void SetLdrInitWaitPrefs(int waitCount, LONGLONG waitTimeout);
 eastl::vector<Process> GetProcesses();
 NTSTATUS OpenProcess(_In_ HANDLE pid, _Out_ PEPROCESS *pep, _Out_ HANDLE *process);
 NTSTATUS CloseProcess(_In_ _Out_ PEPROCESS *pep, _In_ _Out_ HANDLE *process);
@@ -56,8 +57,57 @@ NTSTATUS RestoreProtectVirtualMemory(_In_ PEPROCESS pep, _In_ uint64_t addr,
                                      _In_ SIZE_T siz, _In_ ULONG old_prot);
 NTSTATUS ReadVirtualMemory(_In_ PEPROCESS pep, _In_ uint64_t sourceAddress,
                            _Out_ UCHAR *targetAddress, _In_ _Out_ SIZE_T *size);
-NTSTATUS WriteVirtualMemory(_In_ PEPROCESS pep, _In_ UCHAR *sourceAddress,
+NTSTATUS WriteVirtualMemory(_In_ PEPROCESS pep, _In_ const UCHAR *sourceAddress,
                             _Out_ uint64_t targetAddress,
                             _In_ _Out_ SIZE_T *size);
 
+class Memory {
+public:
+    Memory(_In_ PEPROCESS& pep) : m_pep(pep) {
+        ClearLastErrorAndSize();
+    }
+    Memory(const Memory&) = delete;
+
+    void ClearLastErrorAndSize() {
+        m_last_error = STATUS_SUCCESS;
+        m_last_size = 0;
+    }
+
+    static bool IsValidAddress(uint64_t address) {
+        return address >= 0x10000 && address < 0x000F000000000000;
+    }
+
+    template<typename T>
+    T Read(uint64_t sourceAddress) {
+        if (!IsValidAddress(sourceAddress))
+            return T();
+        T value;
+        SIZE_T size = sizeof(value);
+        m_last_error = ReadVirtualMemory(m_pep, sourceAddress, reinterpret_cast<UCHAR*>(&value), &size);
+        m_last_size = size;
+        if (m_last_error == STATUS_SUCCESS && m_last_size == sizeof(T))
+            return value;
+        return T();
+    }
+
+    template<typename T>
+    bool Write(uint64_t targetAddress, const T& writeData) {
+        if (!IsValidAddress(targetAddress))
+            return false;
+        SIZE_T size = sizeof(T);
+        m_last_error = WriteVirtualMemory(m_pep, reinterpret_cast<const UCHAR*>(&writeData), targetAddress, &size);
+        m_last_size = size;
+        if (m_last_error == STATUS_SUCCESS && m_last_size == sizeof(T))
+            return true;
+        return false;
+    }
+
+    NTSTATUS LastError() { return m_last_error; }
+    SIZE_T LastSize() { return m_last_size; }
+
+private:
+    PEPROCESS& m_pep;
+    NTSTATUS m_last_error;
+    SIZE_T m_last_size;
+};
 #endif