diff options
Diffstat (limited to 'src/zenutil')
| -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 | 99 | ||||
| -rw-r--r-- | src/zenutil/include/zenutil/zenserverprocess.h | 2 | ||||
| -rw-r--r-- | src/zenutil/process.cpp | 806 | ||||
| -rw-r--r-- | src/zenutil/xmake.lua | 7 | ||||
| -rw-r--r-- | src/zenutil/zenutil.cpp | 2 |
8 files changed, 1017 insertions, 2 deletions
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/zenutil/include/zenutil/process.h b/src/zenutil/include/zenutil/process.h new file mode 100644 index 000000000..429ab113a --- /dev/null +++ b/src/zenutil/include/zenutil/process.h @@ -0,0 +1,99 @@ +// Copyright Epic Games, Inc. All Rights Reserved. + +#pragma once + +#include <zencore/thread.h> +#include <zencore/zencore.h> + +#include <filesystem> + +namespace zen { + +/** Basic process abstraction + */ +class ProcessHandle +{ +public: + ZENCORE_API ProcessHandle(); + + ProcessHandle(const ProcessHandle&) = delete; + ProcessHandle& operator=(const ProcessHandle&) = delete; + + ZENCORE_API ~ProcessHandle(); + + ZENCORE_API void Initialize(int Pid); + ZENCORE_API void Initialize(void* ProcessHandle); /// Initialize with an existing handle - takes ownership of the handle + ZENCORE_API [[nodiscard]] bool IsRunning() const; + ZENCORE_API [[nodiscard]] bool IsValid() const; + ZENCORE_API bool Wait(int TimeoutMs = -1); + ZENCORE_API int WaitExitCode(); + ZENCORE_API void Terminate(int ExitCode); + ZENCORE_API void Reset(); + [[nodiscard]] inline int Pid() const { return m_Pid; } + +private: + void* m_ProcessHandle = nullptr; + int m_Pid = 0; +#if ZEN_PLATFORM_LINUX || ZEN_PLATFORM_MAC + int m_ExitCode = -1; +#endif +}; + +/** Basic process creation + */ +struct CreateProcOptions +{ + enum + { + Flag_NewConsole = 1 << 0, + Flag_Elevated = 1 << 1, + Flag_Unelevated = 1 << 2, + }; + + const std::filesystem::path* WorkingDirectory = nullptr; + uint32_t Flags = 0; + std::filesystem::path StdoutFile; + bool WithTracking = false; +}; + +#if ZEN_PLATFORM_WINDOWS +using CreateProcResult = void*; // handle to the process +#else +using CreateProcResult = int32_t; // pid +#endif + +ZENCORE_API CreateProcResult CreateProc(const std::filesystem::path& Executable, + std::string_view CommandLine, // should also include arg[0] (executable name) + const CreateProcOptions& Options = {}); + +/** Process monitor - monitors a list of running processes via polling + + Intended to be used to monitor a set of "sponsor" processes, where + we need to determine when none of them remain alive + + */ + +class ProcessMonitor +{ +public: + ProcessMonitor(); + ~ProcessMonitor(); + + ZENCORE_API bool IsRunning(); + ZENCORE_API void AddPid(int Pid); + ZENCORE_API bool IsActive() const; + +private: + using HandleType = void*; + + mutable RwLock m_Lock; + std::vector<HandleType> m_ProcessHandles; +}; + +ZENCORE_API bool IsProcessRunning(int pid); +ZENCORE_API int GetCurrentProcessId(); +int GetProcessId(CreateProcResult ProcId); + +void process_forcelink(); // internal + +} // namespace zen 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/zenutil/process.cpp b/src/zenutil/process.cpp new file mode 100644 index 000000000..052d36ccf --- /dev/null +++ b/src/zenutil/process.cpp @@ -0,0 +1,806 @@ +// Copyright Epic Games, Inc. All Rights Reserved. + +#include <zenutil/process.h> + +#include <zencore/except.h> +#include <zencore/filesystem.h> +#include <zencore/fmtutils.h> +#include <zencore/scopeguard.h> +#include <zencore/string.h> +#include <zencore/testing.h> + +#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> +#else +# include <fcntl.h> +# include <pthread.h> +# include <signal.h> +# include <sys/file.h> +# include <sys/sem.h> +# include <sys/stat.h> +# include <sys/syscall.h> +# include <sys/wait.h> +# include <time.h> +# include <unistd.h> +#endif + +ZEN_THIRD_PARTY_INCLUDES_START +#include <fmt/format.h> +ZEN_THIRD_PARTY_INCLUDES_END + +namespace zen { + +#if ZEN_PLATFORM_LINUX +const bool bNoZombieChildren = []() { + // When a child process exits it is put into a zombie state until the parent + // collects its result. This doesn't fit the Windows-like model that Zen uses + // where there is a less strict familial model and no zombification. Ignoring + // SIGCHLD signals removes the need to call wait() on zombies. Another option + // would be for the child to call setsid() but that would detatch the child + // from the terminal. + struct sigaction Action = {}; + sigemptyset(&Action.sa_mask); + Action.sa_handler = SIG_IGN; + sigaction(SIGCHLD, &Action, nullptr); + return true; +}(); +#endif + +ProcessHandle::ProcessHandle() = default; + +#if ZEN_PLATFORM_WINDOWS +void +ProcessHandle::Initialize(void* ProcessHandle) +{ + ZEN_ASSERT(m_ProcessHandle == nullptr); + + if (ProcessHandle == INVALID_HANDLE_VALUE) + { + ProcessHandle = nullptr; + } + + // TODO: perform some debug verification here to verify it's a valid handle? + m_ProcessHandle = ProcessHandle; + m_Pid = GetProcessId(m_ProcessHandle); +} +#endif // ZEN_PLATFORM_WINDOWS + +ProcessHandle::~ProcessHandle() +{ + Reset(); +} + +void +ProcessHandle::Initialize(int Pid) +{ + ZEN_ASSERT(m_ProcessHandle == nullptr); + +#if ZEN_PLATFORM_WINDOWS + m_ProcessHandle = OpenProcess(PROCESS_QUERY_INFORMATION | SYNCHRONIZE, FALSE, Pid); +#elif ZEN_PLATFORM_LINUX || ZEN_PLATFORM_MAC + if (Pid > 0) + { + m_ProcessHandle = (void*)(intptr_t(Pid)); + } +#endif + + if (!m_ProcessHandle) + { + ThrowLastError(fmt::format("ProcessHandle::Initialize(pid: {}) failed", Pid)); + } + + m_Pid = Pid; +} + +bool +ProcessHandle::IsRunning() const +{ + bool bActive = false; + +#if ZEN_PLATFORM_WINDOWS + DWORD ExitCode = 0; + GetExitCodeProcess(m_ProcessHandle, &ExitCode); + bActive = (ExitCode == STILL_ACTIVE); +#elif ZEN_PLATFORM_LINUX || ZEN_PLATFORM_MAC + bActive = (kill(pid_t(m_Pid), 0) == 0); +#endif + + return bActive; +} + +bool +ProcessHandle::IsValid() const +{ + return (m_ProcessHandle != nullptr); +} + +void +ProcessHandle::Terminate(int ExitCode) +{ + if (!IsRunning()) + { + return; + } + + bool bSuccess = false; + +#if ZEN_PLATFORM_WINDOWS + TerminateProcess(m_ProcessHandle, ExitCode); + DWORD WaitResult = WaitForSingleObject(m_ProcessHandle, INFINITE); + bSuccess = (WaitResult != WAIT_OBJECT_0); +#elif ZEN_PLATFORM_LINUX || ZEN_PLATFORM_MAC + ZEN_UNUSED(ExitCode); + bSuccess = (kill(m_Pid, SIGKILL) == 0); +#endif + + if (!bSuccess) + { + // What might go wrong here, and what is meaningful to act on? + } +} + +void +ProcessHandle::Reset() +{ + if (IsValid()) + { +#if ZEN_PLATFORM_WINDOWS + CloseHandle(m_ProcessHandle); +#endif + m_ProcessHandle = nullptr; + m_Pid = 0; + } +} + +bool +ProcessHandle::Wait(int TimeoutMs) +{ + using namespace std::literals; + +#if ZEN_PLATFORM_WINDOWS + const DWORD Timeout = (TimeoutMs < 0) ? INFINITE : TimeoutMs; + + const DWORD WaitResult = WaitForSingleObject(m_ProcessHandle, Timeout); + + switch (WaitResult) + { + case WAIT_OBJECT_0: + return true; + + case WAIT_TIMEOUT: + return false; + + case WAIT_FAILED: + break; + } +#elif ZEN_PLATFORM_LINUX || ZEN_PLATFORM_MAC + const int SleepMs = 20; + timespec SleepTime = {0, SleepMs * 1000 * 1000}; + for (int SleepedTimeMS = 0;; SleepedTimeMS += SleepMs) + { + int WaitState = 0; + waitpid(m_Pid, &WaitState, WNOHANG | WCONTINUED | WUNTRACED); + + if (WIFEXITED(WaitState)) + { + m_ExitCode = WEXITSTATUS(WaitState); + } + + if (kill(m_Pid, 0) < 0) + { + int32_t LastError = zen::GetLastError(); + if (LastError == ESRCH) + { + return true; + } + ThrowSystemError(static_cast<uint32_t>(LastError), "Process::Wait kill failed"sv); + } + + if (TimeoutMs >= 0 && SleepedTimeMS >= TimeoutMs) + { + return false; + } + + nanosleep(&SleepTime, nullptr); + } +#endif + + // What might go wrong here, and what is meaningful to act on? + ThrowLastError("Process::Wait failed"sv); +} + +int +ProcessHandle::WaitExitCode() +{ + Wait(-1); + +#if ZEN_PLATFORM_WINDOWS + DWORD ExitCode = 0; + GetExitCodeProcess(m_ProcessHandle, &ExitCode); + + ZEN_ASSERT(ExitCode != STILL_ACTIVE); + + return ExitCode; +#elif ZEN_PLATFORM_LINUX || ZEN_PLATFORM_MAC + return m_ExitCode; +#else + ZEN_NOT_IMPLEMENTED(); + + return 0; +#endif +} + +////////////////////////////////////////////////////////////////////////// + +#if !ZEN_PLATFORM_WINDOWS || ZEN_WITH_TESTS +static void +BuildArgV(std::vector<char*>& Out, char* CommandLine) +{ + char* Cursor = CommandLine; + while (true) + { + // Skip leading whitespace + for (; *Cursor == ' '; ++Cursor) + ; + + // Check for nullp terminator + if (*Cursor == '\0') + { + break; + } + + Out.push_back(Cursor); + + // Extract word + int QuoteCount = 0; + do + { + QuoteCount += (*Cursor == '\"'); + if (*Cursor == ' ' && !(QuoteCount & 1)) + { + break; + } + ++Cursor; + } while (*Cursor != '\0'); + + if (*Cursor == '\0') + { + break; + } + + *Cursor = '\0'; + ++Cursor; + } +} +#endif // !WINDOWS || TESTS + +#if ZEN_PLATFORM_WINDOWS +static CreateProcResult +CreateProcNormal(const std::filesystem::path& Executable, std::string_view CommandLine, const CreateProcOptions& Options) +{ + PROCESS_INFORMATION ProcessInfo{}; + STARTUPINFO StartupInfo{.cb = sizeof(STARTUPINFO)}; + + bool InheritHandles = false; + void* Environment = nullptr; + LPSECURITY_ATTRIBUTES ProcessAttributes = nullptr; + LPSECURITY_ATTRIBUTES ThreadAttributes = nullptr; + + DWORD CreationFlags = 0; + if (Options.Flags & CreateProcOptions::Flag_NewConsole) + { + CreationFlags |= CREATE_NEW_CONSOLE; + } + + const wchar_t* WorkingDir = nullptr; + if (Options.WorkingDirectory != nullptr) + { + WorkingDir = Options.WorkingDirectory->c_str(); + } + + ExtendableWideStringBuilder<256> CommandLineZ; + CommandLineZ << CommandLine; + + if (!Options.StdoutFile.empty()) + { + SECURITY_ATTRIBUTES sa; + sa.nLength = sizeof sa; + sa.lpSecurityDescriptor = nullptr; + sa.bInheritHandle = TRUE; + + StartupInfo.hStdInput = nullptr; + StartupInfo.hStdOutput = CreateFileW(Options.StdoutFile.c_str(), + GENERIC_READ | GENERIC_WRITE, + FILE_SHARE_READ, + &sa, + CREATE_ALWAYS, + FILE_ATTRIBUTE_NORMAL, + nullptr); + + const BOOL Success = DuplicateHandle(GetCurrentProcess(), + StartupInfo.hStdOutput, + GetCurrentProcess(), + &StartupInfo.hStdError, + 0, + TRUE, + DUPLICATE_SAME_ACCESS); + + if (Success) + { + StartupInfo.dwFlags |= STARTF_USESTDHANDLES; + InheritHandles = true; + } + else + { + CloseHandle(StartupInfo.hStdOutput); + StartupInfo.hStdOutput = 0; + } + } + + 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) + { + CloseHandle(StartupInfo.hStdError); + CloseHandle(StartupInfo.hStdOutput); + } + + if (!Success) + { + return nullptr; + } + + CloseHandle(ProcessInfo.hThread); + return ProcessInfo.hProcess; +} + +static CreateProcResult +CreateProcUnelevated(const std::filesystem::path& Executable, std::string_view CommandLine, const CreateProcOptions& Options) +{ + /* Launches a binary with the shell as its parent. The shell (such as + Explorer) should be an unelevated process. */ + + // No sense in using this route if we are not elevated in the first place + if (IsUserAnAdmin() == FALSE) + { + return CreateProcNormal(Executable, CommandLine, Options); + } + + // Get the users' shell process and open it for process creation + HWND ShellWnd = GetShellWindow(); + if (ShellWnd == nullptr) + { + return nullptr; + } + + DWORD ShellPid; + GetWindowThreadProcessId(ShellWnd, &ShellPid); + + HANDLE Process = OpenProcess(PROCESS_CREATE_PROCESS, FALSE, ShellPid); + if (Process == nullptr) + { + return nullptr; + } + auto $0 = MakeGuard([&] { CloseHandle(Process); }); + + // Creating a process as a child of another process is done by setting a + // thread-attribute list on the startup info passed to CreateProcess() + SIZE_T AttrListSize; + InitializeProcThreadAttributeList(nullptr, 1, 0, &AttrListSize); + + auto AttrList = (PPROC_THREAD_ATTRIBUTE_LIST)malloc(AttrListSize); + auto $1 = MakeGuard([&] { free(AttrList); }); + + if (!InitializeProcThreadAttributeList(AttrList, 1, 0, &AttrListSize)) + { + return nullptr; + } + + BOOL bOk = + UpdateProcThreadAttribute(AttrList, 0, PROC_THREAD_ATTRIBUTE_PARENT_PROCESS, (HANDLE*)&Process, sizeof(Process), nullptr, nullptr); + if (!bOk) + { + return nullptr; + } + + // By this point we know we are an elevated process. It is not allowed to + // create a process as a child of another unelevated process that share our + // elevated console window if we have one. So we'll need to create a new one. + uint32_t CreateProcFlags = EXTENDED_STARTUPINFO_PRESENT; + if (GetConsoleWindow() != nullptr) + { + CreateProcFlags |= CREATE_NEW_CONSOLE; + } + else + { + CreateProcFlags |= DETACHED_PROCESS; + } + + // Everything is set up now so we can proceed and launch the process + STARTUPINFOEXW StartupInfo = { + .StartupInfo = {.cb = sizeof(STARTUPINFOEXW)}, + .lpAttributeList = AttrList, + }; + PROCESS_INFORMATION ProcessInfo = {}; + + if (Options.Flags & CreateProcOptions::Flag_NewConsole) + { + CreateProcFlags |= CREATE_NEW_CONSOLE; + } + + ExtendableWideStringBuilder<256> CommandLineZ; + CommandLineZ << CommandLine; + + ExtendableWideStringBuilder<256> CurrentDirZ; + LPCWSTR WorkingDirectoryPtr = nullptr; + if (Options.WorkingDirectory) + { + CurrentDirZ << Options.WorkingDirectory->native(); + WorkingDirectoryPtr = CurrentDirZ.c_str(); + } + + bOk = CreateProcessW(Executable.c_str(), + CommandLineZ.Data(), + nullptr, + nullptr, + FALSE, + CreateProcFlags, + nullptr, + WorkingDirectoryPtr, + &StartupInfo.StartupInfo, + &ProcessInfo); + if (bOk == FALSE) + { + return nullptr; + } + + CloseHandle(ProcessInfo.hThread); + return ProcessInfo.hProcess; +} + +static CreateProcResult +CreateProcElevated(const std::filesystem::path& Executable, std::string_view CommandLine, const CreateProcOptions& Options) +{ + ExtendableWideStringBuilder<256> CommandLineZ; + CommandLineZ << CommandLine; + + SHELLEXECUTEINFO ShellExecuteInfo; + ZeroMemory(&ShellExecuteInfo, sizeof(ShellExecuteInfo)); + ShellExecuteInfo.cbSize = sizeof(ShellExecuteInfo); + ShellExecuteInfo.fMask = SEE_MASK_UNICODE | SEE_MASK_NOCLOSEPROCESS; + ShellExecuteInfo.lpFile = Executable.c_str(); + ShellExecuteInfo.lpVerb = TEXT("runas"); + ShellExecuteInfo.nShow = SW_SHOW; + ShellExecuteInfo.lpParameters = CommandLineZ.c_str(); + + if (Options.WorkingDirectory != nullptr) + { + ShellExecuteInfo.lpDirectory = Options.WorkingDirectory->c_str(); + } + + if (::ShellExecuteEx(&ShellExecuteInfo)) + { + return ShellExecuteInfo.hProcess; + } + + return nullptr; +} +#endif // ZEN_PLATFORM_WINDOWS + +CreateProcResult +CreateProc(const std::filesystem::path& Executable, std::string_view CommandLine, const CreateProcOptions& Options) +{ +#if ZEN_PLATFORM_WINDOWS + if (Options.Flags & CreateProcOptions::Flag_Unelevated) + { + return CreateProcUnelevated(Executable, CommandLine, Options); + } + + if (Options.Flags & CreateProcOptions::Flag_Elevated) + { + return CreateProcElevated(Executable, CommandLine, Options); + } + + return CreateProcNormal(Executable, CommandLine, Options); +#else + std::vector<char*> ArgV; + std::string CommandLineZ(CommandLine); + BuildArgV(ArgV, CommandLineZ.data()); + ArgV.push_back(nullptr); + + int ChildPid = fork(); + if (ChildPid < 0) + { + ThrowLastError("Failed to fork a new child process"); + } + else if (ChildPid == 0) + { + if (Options.WorkingDirectory != nullptr) + { + int Result = chdir(Options.WorkingDirectory->c_str()); + ZEN_UNUSED(Result); + } + + if (execv(Executable.c_str(), ArgV.data()) < 0) + { + ThrowLastError("Failed to exec() a new process image"); + } + } + + return ChildPid; +#endif +} + +////////////////////////////////////////////////////////////////////////// + +ProcessMonitor::ProcessMonitor() +{ +} + +ProcessMonitor::~ProcessMonitor() +{ + RwLock::ExclusiveLockScope _(m_Lock); + + for (HandleType& Proc : m_ProcessHandles) + { +#if ZEN_PLATFORM_WINDOWS + CloseHandle(Proc); +#endif + Proc = 0; + } +} + +bool +ProcessMonitor::IsRunning() +{ + RwLock::ExclusiveLockScope _(m_Lock); + + bool FoundOne = false; + + for (HandleType& Proc : m_ProcessHandles) + { + bool ProcIsActive; + +#if ZEN_PLATFORM_WINDOWS + DWORD ExitCode = 0; + GetExitCodeProcess(Proc, &ExitCode); + + ProcIsActive = (ExitCode == STILL_ACTIVE); + if (!ProcIsActive) + { + CloseHandle(Proc); + } +#else + int Pid = int(intptr_t(Proc)); + ProcIsActive = IsProcessRunning(Pid); +#endif + + if (!ProcIsActive) + { + Proc = 0; + } + + // Still alive + FoundOne |= ProcIsActive; + } + + std::erase_if(m_ProcessHandles, [](HandleType Handle) { return Handle == 0; }); + + return FoundOne; +} + +void +ProcessMonitor::AddPid(int Pid) +{ + HandleType ProcessHandle; + +#if ZEN_PLATFORM_WINDOWS + ProcessHandle = OpenProcess(PROCESS_QUERY_INFORMATION | SYNCHRONIZE, FALSE, Pid); +#else + ProcessHandle = HandleType(intptr_t(Pid)); +#endif + + if (ProcessHandle) + { + RwLock::ExclusiveLockScope _(m_Lock); + m_ProcessHandles.push_back(ProcessHandle); + } +} + +bool +ProcessMonitor::IsActive() const +{ + RwLock::SharedLockScope _(m_Lock); + return m_ProcessHandles.empty() == false; +} + +////////////////////////////////////////////////////////////////////////// + +bool +IsProcessRunning(int pid) +{ + // This function is arguably not super useful, a pid can be re-used + // by the OS so holding on to a pid and polling it over some time + // period will not necessarily tell you what you probably want to know. + +#if ZEN_PLATFORM_WINDOWS + HANDLE hProc = OpenProcess(PROCESS_QUERY_LIMITED_INFORMATION, FALSE, pid); + + if (!hProc) + { + DWORD Error = zen::GetLastError(); + + if (Error == ERROR_INVALID_PARAMETER) + { + return false; + } + + ThrowSystemError(Error, fmt::format("failed to open process with pid {}", pid)); + } + + bool bStillActive = true; + DWORD ExitCode = 0; + if (0 != GetExitCodeProcess(hProc, &ExitCode)) + { + bStillActive = ExitCode == STILL_ACTIVE; + } + else + { + ZEN_WARN("Unable to get exit code from handle for process '{}', treating the process as active", pid); + } + + CloseHandle(hProc); + + return bStillActive; +#elif ZEN_PLATFORM_LINUX || ZEN_PLATFORM_MAC + return (kill(pid_t(pid), 0) == 0); +#endif +} + +int +GetCurrentProcessId() +{ +#if ZEN_PLATFORM_WINDOWS + return ::GetCurrentProcessId(); +#else + return int(getpid()); +#endif +} + +int +GetProcessId(CreateProcResult ProcId) +{ +#if ZEN_PLATFORM_WINDOWS + return static_cast<int>(::GetProcessId(ProcId)); +#else + return ProcId; +#endif +} + +#if ZEN_WITH_TESTS + +void +process_forcelink() +{ +} + +TEST_SUITE_BEGIN("core.process"); + +TEST_CASE("Process") +{ + int Pid = GetCurrentProcessId(); + CHECK(Pid > 0); + CHECK(IsProcessRunning(Pid)); +} + +TEST_CASE("BuildArgV") +{ + const char* Words[] = {"one", "two", "three", "four", "five"}; + struct + { + int WordCount; + const char* Input; + } Cases[] = { + {0, ""}, + {0, " "}, + {1, "one"}, + {1, " one"}, + {1, "one "}, + {2, "one two"}, + {2, " one two"}, + {2, "one two "}, + {2, " one two"}, + {2, "one two "}, + {2, "one two "}, + {3, "one two three"}, + {3, "\"one\" two \"three\""}, + {5, "one two three four five"}, + }; + + for (const auto& Case : Cases) + { + std::vector<char*> OutArgs; + StringBuilder<64> Mutable; + Mutable << Case.Input; + BuildArgV(OutArgs, Mutable.Data()); + + CHECK_EQ(OutArgs.size(), Case.WordCount); + + for (int i = 0, n = int(OutArgs.size()); i < n; ++i) + { + const char* Truth = Words[i]; + size_t TruthLen = strlen(Truth); + + const char* Candidate = OutArgs[i]; + bool bQuoted = (Candidate[0] == '\"'); + Candidate += bQuoted; + + CHECK(strncmp(Truth, Candidate, TruthLen) == 0); + + if (bQuoted) + { + CHECK_EQ(Candidate[TruthLen], '\"'); + } + } + } +} + +TEST_SUITE_END(/* core.process */); + +#endif + +} // namespace zen 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(); } |