aboutsummaryrefslogtreecommitdiff
path: root/src
diff options
context:
space:
mode:
Diffstat (limited to 'src')
-rw-r--r--src/zen/cmds/run_cmd.cpp4
-rw-r--r--src/zen/cmds/run_cmd.h1
-rw-r--r--src/zentrack/include/zentrack/zentrack.h9
-rw-r--r--src/zentrack/zentrack.cpp33
-rw-r--r--src/zenutil/process.cpp61
-rw-r--r--src/zenutil/xmake.lua4
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")