Added support for defining test/non-test server environments

1 files changed, 117 insertions, 35 deletions

diff --git a/zenutil/zenserverprocess.cpp b/zenutil/zenserverprocess.cpp
index f58a9f166..9d38d3939 100644
--- a/zenutil/zenserverprocess.cpp
+++ b/zenutil/zenserverprocess.cpp
@@ -6,6 +6,7 @@
 #include <zencore/fmtutils.h>
 #include <zencore/string.h>
 
+#include <shellapi.h>
 #include <spdlog/spdlog.h>
 
 //////////////////////////////////////////////////////////////////////////
@@ -21,7 +22,17 @@ ZenServerEnvironment::~ZenServerEnvironment()
 }
 
 void
-ZenServerEnvironment::Initialize(std::filesystem::path ProgramBaseDir, std::filesystem::path TestBaseDir)
+ZenServerEnvironment::Initialize(std::filesystem::path ProgramBaseDir)
+{
+	m_ProgramBaseDir = ProgramBaseDir;
+
+	spdlog::debug("Program base dir is '{}'", ProgramBaseDir);
+
+	m_IsInitialized = true;
+}
+
+void
+ZenServerEnvironment::InitializeForTest(std::filesystem::path ProgramBaseDir, std::filesystem::path TestBaseDir)
 {
 	m_ProgramBaseDir = ProgramBaseDir;
 	m_TestBaseDir	 = TestBaseDir;
@@ -30,7 +41,8 @@ ZenServerEnvironment::Initialize(std::filesystem::path ProgramBaseDir, std::file
 	spdlog::info("Cleaning test base dir '{}'", TestBaseDir);
 	zen::DeleteDirectories(TestBaseDir.c_str());
 
-	m_IsInitialized = true;
+	m_IsTestInstance = true;
+	m_IsInitialized	 = true;
 }
 
 std::filesystem::path
@@ -51,7 +63,7 @@ ZenServerEnvironment::CreateNewTestDir()
 }
 
 std::filesystem::path
-ZenServerEnvironment::RootDir(std::string_view Path)
+ZenServerEnvironment::GetTestRootDir(std::string_view Path)
 {
 	std::filesystem::path Root = m_ProgramBaseDir.parent_path().parent_path();
 
@@ -112,12 +124,21 @@ ZenServerInstance::SpawnServer(int BasePort)
 	zen::ExtendableWideStringBuilder<128> CommandLine;
 	CommandLine << "\"";
 	CommandLine.Append(Executable.c_str());
-	CommandLine << "\" --test --owner-pid ";
-	CommandLine << MyPid;
-	CommandLine << " ";
-	CommandLine << "--port " << BasePort;
+	CommandLine << "\"";
+
+	const bool IsTest = m_Env.IsTestEnvironment();
+
+	if (IsTest)
+	{
+		CommandLine << " --test --owner-pid " << MyPid << " --log-id " << LogId;
+	}
+
 	CommandLine << " --child-id " << ChildEventName;
-	CommandLine << " --log-id " << LogId;
+
+	if (BasePort)
+	{
+		CommandLine << " --port " << BasePort;
+	}
 
 	if (!m_TestDir.empty())
 	{
@@ -135,39 +156,100 @@ ZenServerInstance::SpawnServer(int BasePort)
 	spdlog::debug("Spawning server");
 
 	PROCESS_INFORMATION ProcessInfo{};
-	STARTUPINFO			Sinfo{.cb = sizeof(STARTUPINFO)};
-
-	DWORD				  CreationFlags		= 0;  // CREATE_NEW_CONSOLE;
-	const bool			  InheritHandles	= false;
-	void*				  Environment		= nullptr;
-	LPSECURITY_ATTRIBUTES ProcessAttributes = nullptr;
-	LPSECURITY_ATTRIBUTES ThreadAttributes	= nullptr;
-
-	BOOL Success = CreateProcessW(Executable.c_str(),
-								  (LPWSTR)CommandLine.c_str(),
-								  ProcessAttributes,
-								  ThreadAttributes,
-								  InheritHandles,
-								  CreationFlags,
-								  Environment,
-								  CurrentDirectory.c_str(),
-								  &Sinfo,
-								  &ProcessInfo);
-
-	if (Success == FALSE)
+	STARTUPINFO			StartupInfo{.cb = sizeof(STARTUPINFO)};
+
+	DWORD CreationFlags = 0;
+
+	if (!IsTest)
+	{
+		CreationFlags |= CREATE_NEW_CONSOLE;
+	}
+
+	HANDLE hProcess = NULL;
+
 	{
-		std::error_code err(::GetLastError(), std::system_category());
+		const bool			  InheritHandles	= false;
+		void*				  Environment		= nullptr;
+		LPSECURITY_ATTRIBUTES ProcessAttributes = nullptr;
+		LPSECURITY_ATTRIBUTES ThreadAttributes	= nullptr;
+
+		BOOL Success = CreateProcessW(Executable.c_str(),
+									  (LPWSTR)CommandLine.c_str(),
+									  ProcessAttributes,
+									  ThreadAttributes,
+									  InheritHandles,
+									  CreationFlags,
+									  Environment,
+									  CurrentDirectory.c_str(),
+									  &StartupInfo,
+									  &ProcessInfo);
+
+		if (Success)
+		{
+			hProcess = ProcessInfo.hProcess;
+			CloseHandle(ProcessInfo.hThread);
+		}
+		else
+		{
+			DWORD WinError = ::GetLastError();
+
+			if (WinError == ERROR_ELEVATION_REQUIRED)
+			{
+				// Try launching elevated process
 
-		spdlog::error("Server spawn failed: {}", err.message());
+				spdlog::debug("Regular spawn failed - spawning elevated server");
 
-		throw std::system_error(err, "failed to create server process");
+				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 = CommandLine.c_str();
+
+				if (::ShellExecuteEx(&ShellExecuteInfo))
+				{
+					WinError = NO_ERROR;
+
+					hProcess = ShellExecuteInfo.hProcess;
+				}
+			}
+
+			if (WinError != NO_ERROR)
+			{
+				std::error_code err(WinError, std::system_category());
+
+				spdlog::error("Server spawn failed: {}", err.message());
+
+				throw std::system_error(err, "failed to create server process");
+			}
+		}
 	}
 
 	spdlog::debug("Server spawned OK");
 
-	CloseHandle(ProcessInfo.hThread);
+	if (IsTest)
+	{
+		m_Process.Initialize(hProcess);
+		m_ShutdownEvent = std::move(ChildShutdownEvent);
+	}
+	else
+	{
+		CloseHandle(hProcess);
+	}
+
+	m_ReadyEvent = std::move(ChildEvent);
+}
 
-	m_Process.Initialize(ProcessInfo.hProcess);
-	m_ReadyEvent	= std::move(ChildEvent);
-	m_ShutdownEvent = std::move(ChildShutdownEvent);
+void
+ZenServerInstance::WaitUntilReady()
+{
+	m_ReadyEvent.Wait();
+}
+
+bool
+ZenServerInstance::WaitUntilReady(int Timeout)
+{
+	return m_ReadyEvent.Wait(Timeout);
 }