added zentrack module

this is meant to be used to inject memory tracking code into a child process

3 files changed, 164 insertions, 0 deletions

diff --git a/src/zentrack/include/zentrack/zentrack.h b/src/zentrack/include/zentrack/zentrack.h
new file mode 100644
index 000000000..14d21ea0d
--- /dev/null
+++ b/src/zentrack/include/zentrack/zentrack.h
@@ -0,0 +1,3 @@
+// Copyright Epic Games, Inc. All Rights Reserved.
+
+#pragma once
diff --git a/src/zentrack/xmake.lua b/src/zentrack/xmake.lua
new file mode 100644
index 000000000..94692afb2
--- /dev/null
+++ b/src/zentrack/xmake.lua
@@ -0,0 +1,11 @@
+-- Copyright Epic Games, Inc. All Rights Reserved.
+
+target('zentrack')
+    set_kind("shared")
+    set_group("libs")
+    add_headerfiles("**.h")
+    add_files("**.cpp")
+    add_includedirs("include", {public=true})
+    add_deps("zencore")
+    add_packages("vcpkg::spdlog")
+    add_links("detours")
diff --git a/src/zentrack/zentrack.cpp b/src/zentrack/zentrack.cpp
new file mode 100644
index 000000000..1d8289d55
--- /dev/null
+++ b/src/zentrack/zentrack.cpp
@@ -0,0 +1,150 @@
+// Copyright Epic Games, Inc. All Rights Reserved.
+
+#include <zenbase/zenbase.h>
+#include <zencore/windows.h>
+
+#include <detours/detours.h>
+
+#define ZEN_EXITCODE_INIT_FAILED	20
+#define ZEN_EXITCODE_HOOKS_FAILED	21
+#define ZEN_EXITCODE_PAYLOAD_FAILED 22
+
+//////////////////////////////////////////////////////////////////////////
+
+#define DETOURED_FUNCTIONS_KERNELBASE \
+	DETOURED_FUNCTION(VirtualAlloc)   \
+	DETOURED_FUNCTION(VirtualFree)
+
+#define DETOURED_FUNCTIONS DETOURED_FUNCTIONS_KERNELBASE
+
+//////////////////////////////////////////////////////////////////////////
+
+struct __declspec(uuid("{e3949b3b-9143-43fb-9a0c-b1f89640da9f}")) DetoursPayload
+{
+	bool HookVirtualAlloc = false;
+};
+
+static const _GUID DetoursPayloadGuid = __uuidof(DetoursPayload);
+
+[[noreturn]] void
+TerminateCurrentProcess(uint32_t ExitCode)
+{
+	TerminateProcess(GetCurrentProcess(), ExitCode);
+}
+
+#define DETOURED_FUNCTION(func)           \
+	using Symbol_##func = decltype(func); \
+	Symbol_##func* Real_##func;
+DETOURED_FUNCTIONS
+#undef DETOURED_FUNCTION
+
+BOOL
+Hooked_VirtualFree(LPVOID lpAddress, SIZE_T dwSize, DWORD dwFreeType)
+{
+	BOOL Result = VirtualFree(lpAddress, dwSize, dwFreeType);
+
+	return Result;
+}
+
+LPVOID
+Hooked_VirtualAlloc(LPVOID lpAddress, SIZE_T dwSize, DWORD flAllocationType, DWORD flProtect)
+{
+	void* Result = Real_VirtualAlloc(lpAddress, dwSize, flAllocationType, flProtect);
+
+	return Result;
+}
+
+//////////////////////////////////////////////////////////////////////////
+
+void
+HookFunction(void** RealFunc, void* HookFunc, const char* FunctionName)
+{
+	if (!*RealFunc)
+	{
+		return;
+	}
+
+	LONG Result = DetourAttach(RealFunc, HookFunc);
+	if (Result == NO_ERROR)
+	{
+		return;
+	}
+
+	// TODO: notify any listener about problem
+	ZEN_UNUSED(FunctionName);
+
+	ExitProcess(Result);
+}
+
+void
+HookFunctions()
+{
+	// First grab the original pre-hook function pointers
+
+#define DETOURED_FUNCTION(FuncName) Real_##FuncName = (decltype(Real_##FuncName))GetProcAddress(hModule, #FuncName);
+
+	if (HMODULE hModule = GetModuleHandleW(L"kernelbase.dll"))
+	{
+		DETOURED_FUNCTIONS_KERNELBASE
+	}
+
+#undef DETOURED_FUNCTION
+
+	// Then install our hooks
+
+#define DETOURED_FUNCTION(Func) HookFunction((PVOID*)&Real_##Func, Hooked_##Func, #Func);
+
+	DETOURED_FUNCTIONS
+}
+
+void
+InstallHooks(DetoursPayload& Payload)
+{
+	ZEN_UNUSED(Payload);
+
+	DetourTransactionBegin();
+	DetourUpdateThread(GetCurrentThread());
+	HookFunctions();
+
+	if (LONG Result = DetourTransactionCommit(); Result != NO_ERROR)
+	{
+		ExitProcess(ZEN_EXITCODE_HOOKS_FAILED);
+	}
+}
+
+//////////////////////////////////////////////////////////////////////////
+
+BOOL WINAPI
+DllMain(HINSTANCE, DWORD dwReason, LPVOID)
+{
+	if (DetourIsHelperProcess())
+		return TRUE;
+
+	if (dwReason == DLL_PROCESS_ATTACH)
+	{
+		// Restore the contents in memory import table after a process was started
+		// with DetourCreateProcessWithDllEx or DetourCreateProcessWithDlls
+		if (!DetourRestoreAfterWith())
+			TerminateCurrentProcess(ZEN_EXITCODE_INIT_FAILED);
+
+		// Grab state
+		DetoursPayload* Payload = (DetoursPayload*)DetourFindPayloadEx(DetoursPayloadGuid, nullptr);
+		if (!Payload)
+			TerminateCurrentProcess(ZEN_EXITCODE_PAYLOAD_FAILED);
+
+		InstallHooks(*Payload);
+	}
+	else if (dwReason == DLL_PROCESS_DETACH)
+	{
+	}
+
+	return TRUE;
+}
+
+void
+foo()
+{
+	DetourTransactionBegin();
+
+	return;
+}
\ No newline at end of file