Merge branch 'main' into lm/restrict-content-type

1 files changed, 329 insertions, 0 deletions

diff --git a/src/zencore/process.cpp b/src/zencore/process.cpp
index 56849a10d..f657869dc 100644
--- a/src/zencore/process.cpp
+++ b/src/zencore/process.cpp
@@ -9,6 +9,7 @@
 #include <zencore/string.h>
 #include <zencore/testing.h>
 #include <zencore/timer.h>
+#include <zencore/trace.h>
 
 #include <thread>
 
@@ -490,6 +491,8 @@ CreateProcNormal(const std::filesystem::path& Executable, std::string_view Comma
 	LPSECURITY_ATTRIBUTES ProcessAttributes = nullptr;
 	LPSECURITY_ATTRIBUTES ThreadAttributes	= nullptr;
 
+	const bool AssignToJob = Options.AssignToJob && Options.AssignToJob->IsValid();
+
 	DWORD CreationFlags = 0;
 	if (Options.Flags & CreateProcOptions::Flag_NewConsole)
 	{
@@ -503,6 +506,10 @@ CreateProcNormal(const std::filesystem::path& Executable, std::string_view Comma
 	{
 		CreationFlags |= CREATE_NEW_PROCESS_GROUP;
 	}
+	if (AssignToJob)
+	{
+		CreationFlags |= CREATE_SUSPENDED;
+	}
 
 	const wchar_t* WorkingDir = nullptr;
 	if (Options.WorkingDirectory != nullptr)
@@ -571,6 +578,15 @@ CreateProcNormal(const std::filesystem::path& Executable, std::string_view Comma
 		return nullptr;
 	}
 
+	if (AssignToJob)
+	{
+		if (!Options.AssignToJob->AssignProcess(ProcessInfo.hProcess))
+		{
+			ZEN_WARN("Failed to assign newly created process to job object");
+		}
+		ResumeThread(ProcessInfo.hThread);
+	}
+
 	CloseHandle(ProcessInfo.hThread);
 	return ProcessInfo.hProcess;
 }
@@ -644,6 +660,8 @@ CreateProcUnelevated(const std::filesystem::path& Executable, std::string_view C
 	};
 	PROCESS_INFORMATION ProcessInfo = {};
 
+	const bool AssignToJob = Options.AssignToJob && Options.AssignToJob->IsValid();
+
 	if (Options.Flags & CreateProcOptions::Flag_NewConsole)
 	{
 		CreateProcFlags |= CREATE_NEW_CONSOLE;
@@ -652,6 +670,10 @@ CreateProcUnelevated(const std::filesystem::path& Executable, std::string_view C
 	{
 		CreateProcFlags |= CREATE_NO_WINDOW;
 	}
+	if (AssignToJob)
+	{
+		CreateProcFlags |= CREATE_SUSPENDED;
+	}
 
 	ExtendableWideStringBuilder<256> CommandLineZ;
 	CommandLineZ << CommandLine;
@@ -679,6 +701,15 @@ CreateProcUnelevated(const std::filesystem::path& Executable, std::string_view C
 		return nullptr;
 	}
 
+	if (AssignToJob)
+	{
+		if (!Options.AssignToJob->AssignProcess(ProcessInfo.hProcess))
+		{
+			ZEN_WARN("Failed to assign newly created process to job object");
+		}
+		ResumeThread(ProcessInfo.hThread);
+	}
+
 	CloseHandle(ProcessInfo.hThread);
 	return ProcessInfo.hProcess;
 }
@@ -715,6 +746,8 @@ CreateProcElevated(const std::filesystem::path& Executable, std::string_view Com
 CreateProcResult
 CreateProc(const std::filesystem::path& Executable, std::string_view CommandLine, const CreateProcOptions& Options)
 {
+	ZEN_TRACE_CPU("CreateProc");
+
 #if ZEN_PLATFORM_WINDOWS
 	if (Options.Flags & CreateProcOptions::Flag_Unelevated)
 	{
@@ -746,6 +779,17 @@ CreateProc(const std::filesystem::path& Executable, std::string_view CommandLine
 			ZEN_UNUSED(Result);
 		}
 
+		if (!Options.StdoutFile.empty())
+		{
+			int Fd = open(Options.StdoutFile.c_str(), O_WRONLY | O_CREAT | O_TRUNC, 0644);
+			if (Fd >= 0)
+			{
+				dup2(Fd, STDOUT_FILENO);
+				dup2(Fd, STDERR_FILENO);
+				close(Fd);
+			}
+		}
+
 		if (execv(Executable.c_str(), ArgV.data()) < 0)
 		{
 			ThrowLastError("Failed to exec() a new process image");
@@ -845,6 +889,65 @@ ProcessMonitor::IsActive() const
 
 //////////////////////////////////////////////////////////////////////////
 
+#if ZEN_PLATFORM_WINDOWS
+JobObject::JobObject() = default;
+
+JobObject::~JobObject()
+{
+	if (m_JobHandle)
+	{
+		CloseHandle(m_JobHandle);
+		m_JobHandle = nullptr;
+	}
+}
+
+void
+JobObject::Initialize()
+{
+	ZEN_ASSERT(m_JobHandle == nullptr, "JobObject already initialized");
+
+	m_JobHandle = CreateJobObjectW(nullptr, nullptr);
+	if (!m_JobHandle)
+	{
+		ZEN_WARN("Failed to create job object: {}", zen::GetLastError());
+		return;
+	}
+
+	JOBOBJECT_EXTENDED_LIMIT_INFORMATION LimitInfo = {};
+	LimitInfo.BasicLimitInformation.LimitFlags	   = JOB_OBJECT_LIMIT_KILL_ON_JOB_CLOSE;
+
+	if (!SetInformationJobObject(m_JobHandle, JobObjectExtendedLimitInformation, &LimitInfo, sizeof(LimitInfo)))
+	{
+		ZEN_WARN("Failed to set job object limits: {}", zen::GetLastError());
+		CloseHandle(m_JobHandle);
+		m_JobHandle = nullptr;
+	}
+}
+
+bool
+JobObject::AssignProcess(void* ProcessHandle)
+{
+	ZEN_ASSERT(m_JobHandle != nullptr, "JobObject not initialized");
+	ZEN_ASSERT(ProcessHandle != nullptr, "ProcessHandle is null");
+
+	if (!AssignProcessToJobObject(m_JobHandle, ProcessHandle))
+	{
+		ZEN_WARN("Failed to assign process to job object: {}", zen::GetLastError());
+		return false;
+	}
+
+	return true;
+}
+
+bool
+JobObject::IsValid() const
+{
+	return m_JobHandle != nullptr;
+}
+#endif	// ZEN_PLATFORM_WINDOWS
+
+//////////////////////////////////////////////////////////////////////////
+
 bool
 IsProcessRunning(int pid, std::error_code& OutEc)
 {
@@ -1001,6 +1104,232 @@ GetProcessExecutablePath(int Pid, std::error_code& OutEc)
 #endif	// ZEN_PLATFORM_LINUX
 }
 
+std::string
+GetProcessCommandLine(int Pid, std::error_code& OutEc)
+{
+#if ZEN_PLATFORM_WINDOWS
+	HANDLE hProcess = OpenProcess(PROCESS_QUERY_INFORMATION, FALSE, static_cast<DWORD>(Pid));
+	if (!hProcess)
+	{
+		OutEc = MakeErrorCodeFromLastError();
+		return {};
+	}
+	auto _ = MakeGuard([hProcess] { CloseHandle(hProcess); });
+
+	// NtQueryInformationProcess is an undocumented NT API; load it dynamically.
+	// Info class 60 = ProcessCommandLine, available since Windows 8.1.
+	using PFN_NtQIP = LONG(WINAPI*)(HANDLE, UINT, PVOID, ULONG, PULONG);
+	static const PFN_NtQIP s_NtQIP =
+		reinterpret_cast<PFN_NtQIP>(GetProcAddress(GetModuleHandleW(L"ntdll.dll"), "NtQueryInformationProcess"));
+	if (!s_NtQIP)
+	{
+		return {};
+	}
+
+	constexpr UINT ProcessCommandLineClass	= 60;
+	constexpr LONG StatusInfoLengthMismatch = static_cast<LONG>(0xC0000004L);
+
+	ULONG ReturnLength = 0;
+	LONG  Status	   = s_NtQIP(hProcess, ProcessCommandLineClass, nullptr, 0, &ReturnLength);
+	if (Status != StatusInfoLengthMismatch || ReturnLength == 0)
+	{
+		return {};
+	}
+
+	std::vector<char> Buf(ReturnLength);
+	Status = s_NtQIP(hProcess, ProcessCommandLineClass, Buf.data(), ReturnLength, &ReturnLength);
+	if (Status < 0)
+	{
+		OutEc = MakeErrorCodeFromLastError();
+		return {};
+	}
+
+	// Output: UNICODE_STRING header immediately followed by the UTF-16 string data.
+	// The UNICODE_STRING.Buffer field points into our Buf.
+	struct LocalUnicodeString
+	{
+		USHORT Length;
+		USHORT MaximumLength;
+		WCHAR* Buffer;
+	};
+	if (ReturnLength < sizeof(LocalUnicodeString))
+	{
+		return {};
+	}
+	const auto* Us = reinterpret_cast<const LocalUnicodeString*>(Buf.data());
+	if (Us->Length == 0 || Us->Buffer == nullptr)
+	{
+		return {};
+	}
+
+	// Skip argv[0]: may be a quoted path ("C:\...\exe.exe") or a bare path
+	const WCHAR* p	 = Us->Buffer;
+	const WCHAR* End = Us->Buffer + Us->Length / sizeof(WCHAR);
+	if (p < End && *p == L'"')
+	{
+		++p;
+		while (p < End && *p != L'"')
+		{
+			++p;
+		}
+		if (p < End)
+		{
+			++p;  // skip closing quote
+		}
+	}
+	else
+	{
+		while (p < End && *p != L' ')
+		{
+			++p;
+		}
+	}
+	while (p < End && *p == L' ')
+	{
+		++p;
+	}
+	if (p >= End)
+	{
+		return {};
+	}
+
+	int Utf8Size = WideCharToMultiByte(CP_UTF8, 0, p, static_cast<int>(End - p), nullptr, 0, nullptr, nullptr);
+	if (Utf8Size <= 0)
+	{
+		OutEc = MakeErrorCodeFromLastError();
+		return {};
+	}
+	std::string Result(Utf8Size, '\0');
+	WideCharToMultiByte(CP_UTF8, 0, p, static_cast<int>(End - p), Result.data(), Utf8Size, nullptr, nullptr);
+	return Result;
+
+#elif ZEN_PLATFORM_LINUX
+	std::string CmdlinePath = fmt::format("/proc/{}/cmdline", Pid);
+	FILE*		F			= fopen(CmdlinePath.c_str(), "rb");
+	if (!F)
+	{
+		OutEc = MakeErrorCodeFromLastError();
+		return {};
+	}
+	auto FGuard = MakeGuard([F] { fclose(F); });
+
+	// /proc/{pid}/cmdline contains null-separated argv entries; read it all
+	std::string Raw;
+	char		Chunk[4096];
+	size_t		BytesRead;
+	while ((BytesRead = fread(Chunk, 1, sizeof(Chunk), F)) > 0)
+	{
+		Raw.append(Chunk, BytesRead);
+	}
+	if (Raw.empty())
+	{
+		return {};
+	}
+
+	// Skip argv[0] (first null-terminated entry)
+	const char* p	= Raw.data();
+	const char* End = Raw.data() + Raw.size();
+	while (p < End && *p != '\0')
+	{
+		++p;
+	}
+	if (p < End)
+	{
+		++p;  // skip null terminator of argv[0]
+	}
+
+	// Build result: remaining entries joined by spaces (inter-arg nulls → spaces)
+	std::string Result;
+	Result.reserve(static_cast<size_t>(End - p));
+	for (const char* q = p; q < End; ++q)
+	{
+		Result += (*q == '\0') ? ' ' : *q;
+	}
+	while (!Result.empty() && Result.back() == ' ')
+	{
+		Result.pop_back();
+	}
+	return Result;
+
+#elif ZEN_PLATFORM_MAC
+	int	   Mib[3]  = {CTL_KERN, KERN_PROCARGS2, Pid};
+	size_t BufSize = 0;
+	if (sysctl(Mib, 3, nullptr, &BufSize, nullptr, 0) != 0 || BufSize == 0)
+	{
+		OutEc = MakeErrorCodeFromLastError();
+		return {};
+	}
+
+	std::vector<char> Buf(BufSize);
+	if (sysctl(Mib, 3, Buf.data(), &BufSize, nullptr, 0) != 0)
+	{
+		OutEc = MakeErrorCodeFromLastError();
+		return {};
+	}
+
+	// Layout: [int argc][exec_path\0][null padding][argv[0]\0][argv[1]\0]...[envp\0]...
+	if (BufSize < sizeof(int))
+	{
+		return {};
+	}
+	int Argc = 0;
+	memcpy(&Argc, Buf.data(), sizeof(int));
+	if (Argc <= 1)
+	{
+		return {};
+	}
+
+	const char* p	= Buf.data() + sizeof(int);
+	const char* End = Buf.data() + BufSize;
+
+	// Skip exec_path and any trailing null padding that follows it
+	while (p < End && *p != '\0')
+	{
+		++p;
+	}
+	while (p < End && *p == '\0')
+	{
+		++p;
+	}
+
+	// Skip argv[0]
+	while (p < End && *p != '\0')
+	{
+		++p;
+	}
+	if (p < End)
+	{
+		++p;
+	}
+
+	// Collect argv[1..Argc-1]
+	std::string Result;
+	for (int i = 1; i < Argc && p < End; ++i)
+	{
+		if (i > 1)
+		{
+			Result += ' ';
+		}
+		const char* ArgStart = p;
+		while (p < End && *p != '\0')
+		{
+			++p;
+		}
+		Result.append(ArgStart, p);
+		if (p < End)
+		{
+			++p;
+		}
+	}
+	return Result;
+
+#else
+	ZEN_UNUSED(Pid);
+	ZEN_UNUSED(OutEc);
+	return {};
+#endif
+}
+
 std::error_code
 FindProcess(const std::filesystem::path& ExecutableImage, ProcessHandle& OutHandle, bool IncludeSelf)
 {