diff options
Diffstat (limited to 'src')
| -rw-r--r-- | src/zen/cmds/run_cmd.cpp | 4 | ||||
| -rw-r--r-- | src/zen/cmds/run_cmd.h | 1 | ||||
| -rw-r--r-- | src/zentrack/include/zentrack/zentrack.h | 9 | ||||
| -rw-r--r-- | src/zentrack/zentrack.cpp | 33 | ||||
| -rw-r--r-- | src/zenutil/process.cpp | 61 | ||||
| -rw-r--r-- | src/zenutil/xmake.lua | 4 |
6 files changed, 88 insertions, 24 deletions
diff --git a/src/zen/cmds/run_cmd.cpp b/src/zen/cmds/run_cmd.cpp index f576deeaf..b82f03089 100644 --- a/src/zen/cmds/run_cmd.cpp +++ b/src/zen/cmds/run_cmd.cpp @@ -51,6 +51,7 @@ RunCommand::RunCommand() "Number of base directories to retain on rotation", cxxopts::value(m_MaxBaseDirectoryCount), "<count>"); + m_Options.add_option("", "", "track", "Enable resource tracking for child process", cxxopts::value(m_RunWithTracking), ""); } RunCommand::~RunCommand() @@ -140,6 +141,7 @@ RunCommand::Run(const ZenCliOptions& GlobalOptions, int argc, char** argv) Stopwatch Timer; CreateProcOptions ProcOptions; + ProcOptions.WithTracking = m_RunWithTracking; if (!RunDir.empty()) { @@ -153,7 +155,7 @@ RunCommand::Run(const ZenCliOptions& GlobalOptions, int argc, char** argv) Proc.Initialize(CreateProc(ExecutablePath, GlobalOptions.PassthroughCommandLine, ProcOptions)); if (!Proc.IsValid()) { - throw std::runtime_error(fmt::format("failed to launch '{}'", ExecutablePath)); + throw std::runtime_error(fmt::format("failed to launch '{}': {}", ExecutablePath, GetLastErrorAsString())); } int ExitCode = Proc.WaitExitCode(); diff --git a/src/zen/cmds/run_cmd.h b/src/zen/cmds/run_cmd.h index f6512a4e8..6899636cf 100644 --- a/src/zen/cmds/run_cmd.h +++ b/src/zen/cmds/run_cmd.h @@ -22,6 +22,7 @@ private: int m_RunTime = -1; std::string m_BaseDirectory; int m_MaxBaseDirectoryCount = 10; + bool m_RunWithTracking = false; }; } // namespace zen diff --git a/src/zentrack/include/zentrack/zentrack.h b/src/zentrack/include/zentrack/zentrack.h index 14d21ea0d..b28e29757 100644 --- a/src/zentrack/include/zentrack/zentrack.h +++ b/src/zentrack/include/zentrack/zentrack.h @@ -1,3 +1,12 @@ // Copyright Epic Games, Inc. All Rights Reserved. #pragma once + +namespace zen { + +struct __declspec(uuid("{e3949b3b-9143-43fb-9a0c-b1f89640da9f}")) DetoursPayload +{ + bool HookVirtualAlloc = false; +}; + +} // namespace zen diff --git a/src/zentrack/zentrack.cpp b/src/zentrack/zentrack.cpp index 1d8289d55..464239f09 100644 --- a/src/zentrack/zentrack.cpp +++ b/src/zentrack/zentrack.cpp @@ -1,5 +1,7 @@ // Copyright Epic Games, Inc. All Rights Reserved. +#include "zentrack/zentrack.h" + #include <zenbase/zenbase.h> #include <zencore/windows.h> @@ -9,22 +11,22 @@ #define ZEN_EXITCODE_HOOKS_FAILED 21 #define ZEN_EXITCODE_PAYLOAD_FAILED 22 +__declspec(dllexport) void dummy() // we need at least one export in order for detours to be happy +{ +} + ////////////////////////////////////////////////////////////////////////// #define DETOURED_FUNCTIONS_KERNELBASE \ DETOURED_FUNCTION(VirtualAlloc) \ - DETOURED_FUNCTION(VirtualFree) + DETOURED_FUNCTION(VirtualFree) \ + DETOURED_FUNCTION(Sleep) #define DETOURED_FUNCTIONS DETOURED_FUNCTIONS_KERNELBASE ////////////////////////////////////////////////////////////////////////// -struct __declspec(uuid("{e3949b3b-9143-43fb-9a0c-b1f89640da9f}")) DetoursPayload -{ - bool HookVirtualAlloc = false; -}; - -static const _GUID DetoursPayloadGuid = __uuidof(DetoursPayload); +static const _GUID DetoursPayloadGuid = __uuidof(zen::DetoursPayload); [[noreturn]] void TerminateCurrentProcess(uint32_t ExitCode) @@ -41,7 +43,7 @@ DETOURED_FUNCTIONS BOOL Hooked_VirtualFree(LPVOID lpAddress, SIZE_T dwSize, DWORD dwFreeType) { - BOOL Result = VirtualFree(lpAddress, dwSize, dwFreeType); + BOOL Result = Real_VirtualFree(lpAddress, dwSize, dwFreeType); return Result; } @@ -54,6 +56,12 @@ Hooked_VirtualAlloc(LPVOID lpAddress, SIZE_T dwSize, DWORD flAllocationType, DWO return Result; } +void +Hooked_Sleep(DWORD dwMilliseconds) +{ + Real_Sleep(dwMilliseconds); +} + ////////////////////////////////////////////////////////////////////////// void @@ -98,7 +106,7 @@ HookFunctions() } void -InstallHooks(DetoursPayload& Payload) +InstallHooks(zen::DetoursPayload& Payload) { ZEN_UNUSED(Payload); @@ -128,8 +136,9 @@ DllMain(HINSTANCE, DWORD dwReason, LPVOID) TerminateCurrentProcess(ZEN_EXITCODE_INIT_FAILED); // Grab state - DetoursPayload* Payload = (DetoursPayload*)DetourFindPayloadEx(DetoursPayloadGuid, nullptr); - if (!Payload) + DWORD cbData = 0; + zen::DetoursPayload* Payload = (zen::DetoursPayload*)DetourFindPayloadEx(DetoursPayloadGuid, &cbData); + if (!Payload || cbData != sizeof(zen::DetoursPayload)) TerminateCurrentProcess(ZEN_EXITCODE_PAYLOAD_FAILED); InstallHooks(*Payload); @@ -147,4 +156,4 @@ foo() DetourTransactionBegin(); return; -}
\ No newline at end of file +} diff --git a/src/zenutil/process.cpp b/src/zenutil/process.cpp index afdd9b338..052d36ccf 100644 --- a/src/zenutil/process.cpp +++ b/src/zenutil/process.cpp @@ -12,6 +12,11 @@ #include <thread> #if ZEN_PLATFORM_WINDOWS +# include <detours/detours.h> +# include <zentrack/zentrack.h> +#endif + +#if ZEN_PLATFORM_WINDOWS # include <shellapi.h> # include <Shlobj.h> # include <zencore/windows.h> @@ -341,16 +346,52 @@ CreateProcNormal(const std::filesystem::path& Executable, std::string_view Comma } } - BOOL Success = CreateProcessW(Executable.c_str(), - CommandLineZ.Data(), - ProcessAttributes, - ThreadAttributes, - InheritHandles, - CreationFlags, - Environment, - WorkingDir, - &StartupInfo, - &ProcessInfo); + BOOL Success; + + if (Options.WithTracking) + { + // We want to inject a payload, so start the process in a suspended state + Success = DetourCreateProcessWithDllExW(Executable.c_str(), + CommandLineZ.Data(), + ProcessAttributes, + ThreadAttributes, + InheritHandles, + CreationFlags | CREATE_SUSPENDED, + Environment, + WorkingDir, + &StartupInfo, + &ProcessInfo, + "zentrack.dll" /* lpDllName */, + NULL /* pfCreateProcessW */); + + if (Success) + { + zen::DetoursPayload Payload; + Payload.HookVirtualAlloc = true; + + const _GUID DetoursPayloadGuid = __uuidof(zen::DetoursPayload); + + if (!DetourCopyPayloadToProcess((HANDLE)ProcessInfo.hProcess, DetoursPayloadGuid, &Payload, sizeof(Payload))) + { + // TODO: report error + } + + ResumeThread(ProcessInfo.hThread); + } + } + else + { + Success = CreateProcessW(Executable.c_str(), + CommandLineZ.Data(), + ProcessAttributes, + ThreadAttributes, + InheritHandles, + CreationFlags, + Environment, + WorkingDir, + &StartupInfo, + &ProcessInfo); + } if (StartupInfo.dwFlags & STARTF_USESTDHANDLES) { diff --git a/src/zenutil/xmake.lua b/src/zenutil/xmake.lua index 0da6d23a6..57ee66b9b 100644 --- a/src/zenutil/xmake.lua +++ b/src/zenutil/xmake.lua @@ -6,5 +6,7 @@ target('zenutil') add_headerfiles("**.h") add_files("**.cpp") add_includedirs("include", {public=true}) - add_deps("zencore") + add_deps("zencore", "zentrack") add_packages("vcpkg::spdlog") + add_packages("vcpkg::detours") + add_syslinks("detours") |