diff options
Diffstat (limited to 'src')
| -rw-r--r-- | src/zen/cmds/bench_cmd.cpp | 2 | ||||
| -rw-r--r-- | src/zen/cmds/rpcreplay_cmd.cpp | 2 | ||||
| -rw-r--r-- | src/zen/cmds/run_cmd.cpp | 6 | ||||
| -rw-r--r-- | src/zen/cmds/run_cmd.h | 1 | ||||
| -rw-r--r-- | src/zen/cmds/up_cmd.cpp | 2 | ||||
| -rw-r--r-- | src/zencore/filesystem.cpp | 3 | ||||
| -rw-r--r-- | src/zencore/zencore.cpp | 2 | ||||
| -rw-r--r-- | src/zentrack/include/zentrack/zentrack.h | 17 | ||||
| -rw-r--r-- | src/zentrack/trackertrace.cpp | 3 | ||||
| -rw-r--r-- | src/zentrack/xmake.lua | 13 | ||||
| -rw-r--r-- | src/zentrack/zentrack.cpp | 165 | ||||
| -rw-r--r-- | src/zenutil/basicfile.cpp | 70 | ||||
| -rw-r--r-- | src/zenutil/include/zenutil/basicfile.h | 1 | ||||
| -rw-r--r-- | src/zenutil/include/zenutil/mimalloc_hooks.h | 32 | ||||
| -rw-r--r-- | src/zenutil/include/zenutil/process.h (renamed from src/zencore/include/zencore/process.h) | 1 | ||||
| -rw-r--r-- | src/zenutil/include/zenutil/zenserverprocess.h | 2 | ||||
| -rw-r--r-- | src/zenutil/process.cpp (renamed from src/zencore/process.cpp) | 63 | ||||
| -rw-r--r-- | src/zenutil/xmake.lua | 7 | ||||
| -rw-r--r-- | src/zenutil/zenutil.cpp | 2 |
19 files changed, 372 insertions, 22 deletions
diff --git a/src/zen/cmds/bench_cmd.cpp b/src/zen/cmds/bench_cmd.cpp index 5c955e980..d8c9e12d6 100644 --- a/src/zen/cmds/bench_cmd.cpp +++ b/src/zen/cmds/bench_cmd.cpp @@ -7,10 +7,10 @@ #include <zencore/filesystem.h> #include <zencore/fmtutils.h> #include <zencore/logging.h> -#include <zencore/process.h> #include <zencore/string.h> #include <zencore/thread.h> #include <zencore/timer.h> +#include <zenutil/process.h> namespace zen { diff --git a/src/zen/cmds/rpcreplay_cmd.cpp b/src/zen/cmds/rpcreplay_cmd.cpp index 53f45358e..d307ef0e8 100644 --- a/src/zen/cmds/rpcreplay_cmd.cpp +++ b/src/zen/cmds/rpcreplay_cmd.cpp @@ -6,7 +6,6 @@ #include <zencore/filesystem.h> #include <zencore/fmtutils.h> #include <zencore/logging.h> -#include <zencore/process.h> #include <zencore/scopeguard.h> #include <zencore/session.h> #include <zencore/stream.h> @@ -15,6 +14,7 @@ #include <zenhttp/httpcommon.h> #include <zenhttp/httpshared.h> #include <zenutil/cache/rpcrecording.h> +#include <zenutil/process.h> ZEN_THIRD_PARTY_INCLUDES_START #include <cpr/cpr.h> diff --git a/src/zen/cmds/run_cmd.cpp b/src/zen/cmds/run_cmd.cpp index a99ba9704..b82f03089 100644 --- a/src/zen/cmds/run_cmd.cpp +++ b/src/zen/cmds/run_cmd.cpp @@ -5,9 +5,9 @@ #include <zencore/filesystem.h> #include <zencore/fmtutils.h> #include <zencore/logging.h> -#include <zencore/process.h> #include <zencore/string.h> #include <zencore/timer.h> +#include <zenutil/process.h> using namespace std::literals; @@ -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/zen/cmds/up_cmd.cpp b/src/zen/cmds/up_cmd.cpp index 837cc7edf..7a21a35b1 100644 --- a/src/zen/cmds/up_cmd.cpp +++ b/src/zen/cmds/up_cmd.cpp @@ -4,7 +4,7 @@ #include <zencore/filesystem.h> #include <zencore/logging.h> -#include <zencore/process.h> +#include <zenutil/process.h> #include <zenutil/zenserverprocess.h> #include <memory> diff --git a/src/zencore/filesystem.cpp b/src/zencore/filesystem.cpp index 29ec14e0c..36195f7c7 100644 --- a/src/zencore/filesystem.cpp +++ b/src/zencore/filesystem.cpp @@ -7,7 +7,6 @@ #include <zencore/fmtutils.h> #include <zencore/iobuffer.h> #include <zencore/logging.h> -#include <zencore/process.h> #include <zencore/stream.h> #include <zencore/string.h> #include <zencore/testing.h> @@ -1433,7 +1432,7 @@ GetRunningExecutablePath() #elif ZEN_PLATFORM_MAC char Buffer[PROC_PIDPATHINFO_MAXSIZE]; - int SelfPid = GetCurrentProcessId(); + int SelfPid = int(getpid()); if (proc_pidpath(SelfPid, Buffer, sizeof(Buffer)) <= 0) return {}; diff --git a/src/zencore/zencore.cpp b/src/zencore/zencore.cpp index d0acac608..8dd687fbd 100644 --- a/src/zencore/zencore.cpp +++ b/src/zencore/zencore.cpp @@ -25,7 +25,6 @@ #include <zencore/logging.h> #include <zencore/memory.h> #include <zencore/mpscqueue.h> -#include <zencore/process.h> #include <zencore/sha1.h> #include <zencore/stats.h> #include <zencore/stream.h> @@ -143,7 +142,6 @@ zencore_forcelinktests() zen::logging_forcelink(); zen::memory_forcelink(); zen::mpscqueue_forcelink(); - zen::process_forcelink(); zen::refcount_forcelink(); zen::sha1_forcelink(); zen::stats_forcelink(); diff --git a/src/zentrack/include/zentrack/zentrack.h b/src/zentrack/include/zentrack/zentrack.h new file mode 100644 index 000000000..6fbbeae92 --- /dev/null +++ b/src/zentrack/include/zentrack/zentrack.h @@ -0,0 +1,17 @@ +// Copyright Epic Games, Inc. All Rights Reserved. + +#pragma once + +#include <zenbase/zenbase.h> + +namespace zen { + +#if ZEN_PLATFORM_WINDOWS +struct __declspec(uuid("{e3949b3b-9143-43fb-9a0c-b1f89640da9f}")) DetoursPayload +{ + bool HookVirtualAlloc = false; +}; + +#endif + +} // namespace zen diff --git a/src/zentrack/trackertrace.cpp b/src/zentrack/trackertrace.cpp new file mode 100644 index 000000000..634b20018 --- /dev/null +++ b/src/zentrack/trackertrace.cpp @@ -0,0 +1,3 @@ +// Copyright Epic Games, Inc. All Rights Reserved. + +#include <zentrack/zentrack.h> diff --git a/src/zentrack/xmake.lua b/src/zentrack/xmake.lua new file mode 100644 index 000000000..70a83f787 --- /dev/null +++ b/src/zentrack/xmake.lua @@ -0,0 +1,13 @@ +-- 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") + if is_os("windows") then + add_links("detours") + end diff --git a/src/zentrack/zentrack.cpp b/src/zentrack/zentrack.cpp new file mode 100644 index 000000000..5b0a248f2 --- /dev/null +++ b/src/zentrack/zentrack.cpp @@ -0,0 +1,165 @@ +// Copyright Epic Games, Inc. All Rights Reserved. + +#include "zentrack/zentrack.h" + +#if ZEN_PLATFORM_WINDOWS +# 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 + +__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(Sleep) + +# define DETOURED_FUNCTIONS DETOURED_FUNCTIONS_KERNELBASE + +////////////////////////////////////////////////////////////////////////// + +static const _GUID DetoursPayloadGuid = __uuidof(zen::DetoursPayload); + +[[noreturn]] void +TerminateCurrentProcess(uint32_t ExitCode) +{ + TerminateProcess(GetCurrentProcess(), ExitCode); +} + +# define DETOURED_FUNCTION(func) \ + using Decl_##func = decltype(func); \ + Decl_##func* Real_##func; +DETOURED_FUNCTIONS +# undef DETOURED_FUNCTION + +BOOL +Hooked_VirtualFree(LPVOID lpAddress, SIZE_T dwSize, DWORD dwFreeType) +{ + BOOL Result = Real_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 +Hooked_Sleep(DWORD dwMilliseconds) +{ + Real_Sleep(dwMilliseconds); +} + +////////////////////////////////////////////////////////////////////////// + +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(zen::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 + DWORD cbData = 0; + zen::DetoursPayload* Payload = reinterpret_cast<zen::DetoursPayload*>(DetourFindPayloadEx(DetoursPayloadGuid, &cbData)); + if (!Payload || cbData != sizeof(zen::DetoursPayload)) + { + TerminateCurrentProcess(ZEN_EXITCODE_PAYLOAD_FAILED); + } + + InstallHooks(*Payload); + } + else if (dwReason == DLL_PROCESS_DETACH) + { + } + + return TRUE; +} + +void +foo() +{ + DetourTransactionBegin(); + + return; +} +#endif diff --git a/src/zenutil/basicfile.cpp b/src/zenutil/basicfile.cpp index 819d0805d..024b1e5bf 100644 --- a/src/zenutil/basicfile.cpp +++ b/src/zenutil/basicfile.cpp @@ -11,6 +11,30 @@ #if ZEN_PLATFORM_WINDOWS # include <zencore/windows.h> +extern "C" +{ +# define STATUS_SUCCESS ((NTSTATUS)0x00000000L) // ntsubauth +# define NTAPI __stdcall + + typedef DWORD NTSTATUS; + + typedef struct _IO_STATUS_BLOCK + { + union + { + NTSTATUS Status; + PVOID Pointer; + } DUMMYUNIONNAME; + + ULONG_PTR Information; + } IO_STATUS_BLOCK, *PIO_STATUS_BLOCK; + + NTSTATUS NTAPI + NtFlushBuffersFileEx(HANDLE FileHandle, ULONG Flags, PVOID Parameters, ULONG ParametersSize, PIO_STATUS_BLOCK IoStatusBlock); + + using Decl_NtFlushBuffersFileEx = decltype(NtFlushBuffersFileEx); + Decl_NtFlushBuffersFileEx* Real_NtFlushBuffersFileEx; +} #else # include <fcntl.h> # include <sys/file.h> @@ -18,11 +42,37 @@ # include <unistd.h> #endif +#if ZEN_PLATFORM_MAC +# include <fcntl.h> +#endif + #include <fmt/format.h> #include <gsl/gsl-lite.hpp> namespace zen { +#if ZEN_PLATFORM_WINDOWS + +NTSTATUS NTAPI +NtFlushBuffersFileEx(HANDLE FileHandle, ULONG Flags, PVOID Parameters, ULONG ParametersSize, PIO_STATUS_BLOCK IoStatusBlock) +{ + if (!Real_NtFlushBuffersFileEx) + { + Real_NtFlushBuffersFileEx = (Decl_NtFlushBuffersFileEx*)GetProcAddress(GetModuleHandleA("kernelbase.dll"), "NtFlushBuffersFileEx"); + } + + if (Real_NtFlushBuffersFileEx) + { + return Real_NtFlushBuffersFileEx(FileHandle, Flags, Parameters, ParametersSize, IoStatusBlock); + } + + return 0; +} + +#endif + +////////////////////////////////////////////////////////////////////////// + BasicFile::~BasicFile() { Close(); @@ -347,6 +397,26 @@ BasicFile::Flush() #endif } +void +BasicFile::FlushDataOnly() +{ +#if ZEN_PLATFORM_WINDOWS + IO_STATUS_BLOCK Iosb{}; + NTSTATUS Status = zen::NtFlushBuffersFileEx(m_FileHandle, FLUSH_FLAGS_FILE_DATA_ONLY, nullptr, 0, &Iosb); + + if (Status != STATUS_SUCCESS) + { + // warn? + } +#elif ZEN_PLATFORM_MAC + int Fd = int(uintptr_t(m_FileHandle)); + fcntl(Fd, F_FULLFSYNC); +#else + int Fd = int(uintptr_t(m_FileHandle)); + fdatasync(Fd); +#endif +} + uint64_t BasicFile::FileSize() { diff --git a/src/zenutil/include/zenutil/basicfile.h b/src/zenutil/include/zenutil/basicfile.h index f25d9f23c..42cea904a 100644 --- a/src/zenutil/include/zenutil/basicfile.h +++ b/src/zenutil/include/zenutil/basicfile.h @@ -60,6 +60,7 @@ public: void Write(const void* Data, uint64_t Size, uint64_t FileOffset); void Write(const void* Data, uint64_t Size, uint64_t FileOffset, std::error_code& Ec); void Flush(); + void FlushDataOnly(); [[nodiscard]] uint64_t FileSize(); [[nodiscard]] uint64_t FileSize(std::error_code& Ec); void SetFileSize(uint64_t FileSize); diff --git a/src/zenutil/include/zenutil/mimalloc_hooks.h b/src/zenutil/include/zenutil/mimalloc_hooks.h new file mode 100644 index 000000000..e7f354e2d --- /dev/null +++ b/src/zenutil/include/zenutil/mimalloc_hooks.h @@ -0,0 +1,32 @@ + +// Copyright Epic Games, Inc. All Rights Reserved. + +#if ZEN_USE_MIMALLOC +ZEN_THIRD_PARTY_INCLUDES_START +# include <mimalloc.h> +ZEN_THIRD_PARTY_INCLUDES_END + +# define ZEN_MIMALLOC_FUNCTIONS \ + ZEN_MI_FUNCTION(mi_malloc) \ + ZEN_MI_FUNCTION(mi_calloc) \ + ZEN_MI_FUNCTION(mi_realloc) \ + ZEN_MI_FUNCTION(mi_expand) \ + ZEN_MI_FUNCTION(mi_free) \ + ZEN_MI_FUNCTION(mi_expand) \ + ZEN_MI_FUNCTION(mi_strdup) \ + ZEN_MI_FUNCTION(mi_strndup) \ + ZEN_MI_FUNCTION(mi_realpath) \ + ZEN_MI_FUNCTION(mi_expand) \ + ZEN_MI_FUNCTION(mi_malloc_aligned) \ + ZEN_MI_FUNCTION(mi_realloc_aligned) + +struct MimallocHooks +{ +# define ZEN_MI_FUNCTION(func) \ + using Decl_##func = decltype(func); \ + Decl_##func* Pfn##func; + ZEN_MIMALLOC_FUNCTIONS +# undef ZEN_MI_FUNCTION +}; + +#endif diff --git a/src/zencore/include/zencore/process.h b/src/zenutil/include/zenutil/process.h index d90a32301..429ab113a 100644 --- a/src/zencore/include/zencore/process.h +++ b/src/zenutil/include/zenutil/process.h @@ -53,6 +53,7 @@ struct CreateProcOptions const std::filesystem::path* WorkingDirectory = nullptr; uint32_t Flags = 0; std::filesystem::path StdoutFile; + bool WithTracking = false; }; #if ZEN_PLATFORM_WINDOWS diff --git a/src/zenutil/include/zenutil/zenserverprocess.h b/src/zenutil/include/zenutil/zenserverprocess.h index 15138341c..6af48c33d 100644 --- a/src/zenutil/include/zenutil/zenserverprocess.h +++ b/src/zenutil/include/zenutil/zenserverprocess.h @@ -4,9 +4,9 @@ #include <zencore/enumflags.h> #include <zencore/logging.h> -#include <zencore/process.h> #include <zencore/thread.h> #include <zencore/uid.h> +#include <zenutil/process.h> #include <atomic> #include <filesystem> diff --git a/src/zencore/process.cpp b/src/zenutil/process.cpp index 2d0ec2de6..052d36ccf 100644 --- a/src/zencore/process.cpp +++ b/src/zenutil/process.cpp @@ -1,6 +1,6 @@ // Copyright Epic Games, Inc. All Rights Reserved. -#include <zencore/process.h> +#include <zenutil/process.h> #include <zencore/except.h> #include <zencore/filesystem.h> @@ -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..fa6ab088e 100644 --- a/src/zenutil/xmake.lua +++ b/src/zenutil/xmake.lua @@ -6,5 +6,10 @@ target('zenutil') add_headerfiles("**.h") add_files("**.cpp") add_includedirs("include", {public=true}) - add_deps("zencore") + add_deps("zencore", "zentrack") add_packages("vcpkg::spdlog") + + if is_os("windows") then + add_packages("vcpkg::detours") + add_syslinks("detours") + end diff --git a/src/zenutil/zenutil.cpp b/src/zenutil/zenutil.cpp index d9d6c83a2..eba3613f1 100644 --- a/src/zenutil/zenutil.cpp +++ b/src/zenutil/zenutil.cpp @@ -5,6 +5,7 @@ #if ZEN_WITH_TESTS # include <zenutil/basicfile.h> +# include <zenutil/process.h> # include <zenutil/cache/rpcrecording.h> namespace zen { @@ -13,6 +14,7 @@ void zenutil_forcelinktests() { basicfile_forcelink(); + process_forcelink(); cache::rpcrecord_forcelink(); } |