Merge pull request #139 from ue-foundation/de/zen-service-command

zen service command

1 files changed, 609 insertions, 0 deletions

diff --git a/src/zen/cmds/service_cmd.cpp b/src/zen/cmds/service_cmd.cpp
new file mode 100644
index 000000000..3847c423d
--- /dev/null
+++ b/src/zen/cmds/service_cmd.cpp
@@ -0,0 +1,609 @@
+// Copyright Epic Games, Inc. All Rights Reserved.
+
+#include "service_cmd.h"
+
+#include <zencore/filesystem.h>
+#include <zencore/fmtutils.h>
+#include <zencore/logging.h>
+#include <zencore/zencore.h>
+#include <zenutil/service.h>
+#include <filesystem>
+
+#if ZEN_PLATFORM_WINDOWS
+#	include <zencore/windows.h>
+#	include <shellapi.h>
+#	include <Shlwapi.h>
+#endif
+
+#if ZEN_PLATFORM_LINUX || ZEN_PLATFORM_MAC
+#	include <unistd.h>
+#endif
+
+ZEN_THIRD_PARTY_INCLUDES_START
+#include <gsl/gsl-lite.hpp>
+ZEN_THIRD_PARTY_INCLUDES_END
+
+#include <memory>
+
+namespace zen {
+
+//////////////////////////////////////////////////////////////////////////
+
+namespace {
+
+#if ZEN_PLATFORM_WINDOWS
+	BOOL IsElevated()
+	{
+		BOOL   fRet	  = FALSE;
+		HANDLE hToken = NULL;
+		if (OpenProcessToken(GetCurrentProcess(), TOKEN_QUERY, &hToken))
+		{
+			TOKEN_ELEVATION Elevation;
+			DWORD			cbSize = sizeof(TOKEN_ELEVATION);
+			if (GetTokenInformation(hToken, TokenElevation, &Elevation, sizeof(Elevation), &cbSize))
+			{
+				fRet = Elevation.TokenIsElevated;
+			}
+		}
+		if (hToken)
+		{
+			CloseHandle(hToken);
+		}
+		return fRet;
+	}
+
+	int WinRelaunchElevated()
+	{
+		TCHAR CurrentDir[4096];
+		GetCurrentDirectory(4096, CurrentDir);
+
+		ExtendableWideStringBuilder<256> Parameters;
+		std::filesystem::path			 ExecutablePath = GetRunningExecutablePath();
+		std::wstring					 CommandLine(GetCommandLine());
+		std::wstring					 CommandArguments(PathGetArgs(CommandLine.data()));
+
+		ZEN_CONSOLE("Attempting to run '{} {}' elevated...", ExecutablePath, WideToUtf8(CommandArguments));
+
+		SHELLEXECUTEINFO shExInfo = {0};
+		shExInfo.cbSize			  = sizeof(shExInfo);
+		shExInfo.fMask			  = SEE_MASK_NOCLOSEPROCESS | SEE_MASK_NO_CONSOLE;
+		shExInfo.hwnd			  = 0;
+		shExInfo.lpVerb			  = TEXT("runas");			   // Operation to perform
+		shExInfo.lpFile			  = ExecutablePath.c_str();	   // Application to start
+		shExInfo.lpParameters	  = CommandArguments.c_str();  // Additional parameters
+		shExInfo.lpDirectory	  = CurrentDir;
+		shExInfo.nShow			  = SW_SHOW;
+		shExInfo.hInstApp		  = 0;
+
+		DWORD ReturnCode = 1;
+		if (ShellExecuteEx(&shExInfo))
+		{
+			WaitForSingleObject(shExInfo.hProcess, INFINITE);
+			GetExitCodeProcess(shExInfo.hProcess, &ReturnCode);
+			CloseHandle(shExInfo.hProcess);
+			if (ReturnCode == 0)
+			{
+				ZEN_CONSOLE("Elevated execution completed successfully.");
+			}
+			else
+			{
+				ZEN_CONSOLE("Elevated execution completed unsuccessfully, return code: '{}'.", ReturnCode);
+			}
+		}
+		else
+		{
+			ZEN_CONSOLE("Failed to run elevated, operation did not complete.");
+			ReturnCode = DWORD(-1);
+		}
+		return (int)ReturnCode;
+	}
+
+#else  // ZEN_PLATFORM_WINDOWS
+
+	bool IsElevated() { return geteuid() == 0; }
+
+#endif	// ZEN_PLATFORM_WINDOWS
+
+	int RunElevated(bool AllowElevation)
+	{
+#if ZEN_PLATFORM_WINDOWS
+		if (AllowElevation)
+		{
+			return WinRelaunchElevated();
+		}
+		else
+		{
+			ZEN_CONSOLE(
+				"This command requires elevated priviliges. Run command with elevated priviliges or add '--allow-elevation' command line "
+				"option.");
+			return 1;
+		}
+#endif	// ZEN_PLATFORM_WINDOWS
+#if ZEN_PLATFORM_LINUX || ZEN_PLATFORM_MAC
+		ZEN_UNUSED(AllowElevation);
+		ZEN_CONSOLE("This command requires elevated priviliges. Run the command with `sudo`");
+		return 1;
+#endif	// ZEN_PLATFORM_LINUX || ZEN_PLATFORM_MAC
+	}
+
+}  // namespace
+
+ServiceCommand::ServiceCommand()
+{
+	m_Options.add_option("",
+						 "v",
+						 "verb",
+						 fmt::format("Verb for service - {}, {}, {}, {}, {}. Use '--' ",
+									 m_StatusOptions.program(),
+									 m_InstallOptions.program(),
+									 m_UninstallOptions.program(),
+									 m_StartOptions.program(),
+									 m_StopOptions.program()),
+						 cxxopts::value(m_Verb),
+						 "<verb>");
+	m_Options.parse_positional({"verb"});
+	m_Options.positional_help("verb");
+
+	m_StatusOptions.add_options()("h,help", "Print help");
+	m_StatusOptions.add_option("",
+							   "n",
+							   "name",
+							   fmt::format("Service name, defaults to \"{}\"", m_ServiceName),
+							   cxxopts::value(m_ServiceName),
+							   "<name>");
+	m_StatusOptions.parse_positional({"name"});
+	m_StatusOptions.positional_help("name");
+
+	m_InstallOptions.add_options()("h,help", "Print help");
+	m_InstallOptions.add_option("", "s", "executable", "Path to server executable", cxxopts::value(m_ServerExecutable), "<path>");
+	m_InstallOptions.add_option("",
+								"n",
+								"name",
+								fmt::format("Service name, defaults to \"{}\"", m_ServiceName),
+								cxxopts::value(m_ServiceName),
+								"<name>");
+
+	m_InstallOptions.add_option("",
+								"",
+								"full",
+								fmt::format("Uninstall a running service and update service binaries before installing"),
+								cxxopts::value(m_FullInstall),
+								"<full>");
+	m_InstallOptions.add_option("",
+								"",
+								"install-path",
+								fmt::format("Path in which to install service binaries"),
+								cxxopts::value(m_InstallPath),
+								"<install-path>");
+
+	m_InstallOptions.add_option("", "u", "user", "User to run service as, defaults to current user", cxxopts::value(m_UserName), "<user>");
+#if ZEN_PLATFORM_WINDOWS
+	m_InstallOptions.add_option("",
+								"d",
+								"display-name",
+								fmt::format("Service display, defaults to \"{}\"", m_ServiceDisplayName),
+								cxxopts::value(m_ServiceDisplayName),
+								"<display-name>");
+	m_InstallOptions.add_option("",
+								"",
+								"description",
+								fmt::format("Service description", m_ServiceDescription),
+								cxxopts::value(m_ServiceDescription),
+								"<description>");
+	m_InstallOptions.add_option("",
+								"",
+								"allow-elevation",
+								fmt::format("Allow attempt to relauch command using an elevated"),
+								cxxopts::value(m_AllowElevation),
+								"<allow-elevation>");
+	m_InstallOptions.parse_positional({"executable", "name", "display-name"});
+#else
+	m_InstallOptions.parse_positional({"executable", "name"});
+#endif	// ZEN_PLATFORM_WINDOWS
+	m_InstallOptions.positional_help("executable name display-name");
+
+	m_UninstallOptions.add_options()("h,help", "Print help");
+	m_UninstallOptions.add_option("",
+								  "n",
+								  "name",
+								  fmt::format("Service name, defaults to \"{}\"", m_ServiceName),
+								  cxxopts::value(m_ServiceName),
+								  "<name>");
+	m_UninstallOptions.parse_positional({"name"});
+	m_UninstallOptions.positional_help("name");
+#if ZEN_PLATFORM_WINDOWS
+	m_UninstallOptions.add_option("",
+								  "",
+								  "allow-elevation",
+								  fmt::format("Allow attempt to relauch command using an elevated"),
+								  cxxopts::value(m_AllowElevation),
+								  "<allow-elevation>");
+#endif	// ZEN_PLATFORM_WINDOWS
+
+	m_StartOptions.add_options()("h,help", "Print help");
+	m_StartOptions.add_option("",
+							  "n",
+							  "name",
+							  fmt::format("Service name, defaults to \"{}\"", m_ServiceName),
+							  cxxopts::value(m_ServiceName),
+							  "<name>");
+	m_StartOptions.parse_positional({"name"});
+	m_StartOptions.positional_help("name");
+#if ZEN_PLATFORM_WINDOWS
+	m_StartOptions.add_option("",
+							  "",
+							  "allow-elevation",
+							  fmt::format("Allow attempt to relauch command using an elevated"),
+							  cxxopts::value(m_AllowElevation),
+							  "<allow-elevation>");
+#endif	// ZEN_PLATFORM_WINDOWS
+
+	m_StopOptions.add_options()("h,help", "Print help");
+	m_StopOptions.add_option("",
+							 "n",
+							 "name",
+							 fmt::format("Service name, defaults to \"{}\"", m_ServiceName),
+							 cxxopts::value(m_ServiceName),
+							 "<name>");
+	m_StopOptions.parse_positional({"name"});
+	m_StopOptions.positional_help("name");
+#if ZEN_PLATFORM_WINDOWS
+	m_StopOptions.add_option("",
+							 "",
+							 "allow-elevation",
+							 fmt::format("Allow attempt to relauch command using an elevated"),
+							 cxxopts::value(m_AllowElevation),
+							 "<allow-elevation>");
+#endif	// ZEN_PLATFORM_WINDOWS
+}
+
+ServiceCommand::~ServiceCommand() = default;
+
+std::string
+FmtServiceInfo(const ServiceInfo& Info, std::string_view Prefix)
+{
+	std::string Result = fmt::format(
+		"{}Status:             {}\n"
+		"{}Executable:         {}\n"
+		"{}CommandLineOptions: {}",
+		Prefix,
+		ToString(Info.Status),
+		Prefix,
+		Info.Spec.ExecutablePath,
+		Prefix,
+		Info.Spec.CommandLineOptions);
+
+#if ZEN_PLATFORM_WINDOWS
+	Result += fmt::format(
+		"\n"
+		"{}Display Name:       {}\n"
+		"{}Description:        {}",
+		Prefix,
+		Info.Spec.DisplayName,
+		Prefix,
+		Info.Spec.Description);
+#endif	// ZEN_PLATFORM_WINDOWS
+
+	return Result;
+}
+
+enum class ServiceStatusReturnCode
+{
+	Running = 0,  // successful return indicates proper installation and everything running smoothly
+	NotInstalled,
+	NotRunning,
+	UnknownError
+};
+
+int
+ServiceCommand::Run(const ZenCliOptions& GlobalOptions, int argc, char** argv)
+{
+	ZEN_UNUSED(GlobalOptions);
+
+	using namespace std::literals;
+
+	std::vector<char*> SubCommandArguments;
+	cxxopts::Options*  SubOption			 = nullptr;
+	int				   ParentCommandArgCount = GetSubCommand(m_Options, argc, argv, m_SubCommands, SubOption, SubCommandArguments);
+	if (!ParseOptions(ParentCommandArgCount, argv))
+	{
+		return 0;
+	}
+
+	if (SubOption == nullptr)
+	{
+		throw zen::OptionParseException("command verb is missing");
+	}
+
+	if (!ParseOptions(*SubOption, gsl::narrow<int>(SubCommandArguments.size()), SubCommandArguments.data()))
+	{
+		return 0;
+	}
+
+	if (SubOption == &m_StatusOptions)
+	{
+		ServiceInfo		Info;
+		std::error_code Ec = QueryInstalledService(m_ServiceName, Info);
+		if (Ec)
+		{
+			ZEN_CONSOLE("Can't get information about service '{}'. Reason: '{}'", m_ServiceName, Ec.message());
+			return gsl::narrow<int>(ServiceStatusReturnCode::UnknownError);
+		}
+		if (Info.Status == ServiceStatus::NotInstalled)
+		{
+			ZEN_CONSOLE("Service '{}' is not installed", m_ServiceName);
+			return gsl::narrow<int>(ServiceStatusReturnCode::NotInstalled);
+		}
+		else if (Info.Status != ServiceStatus::Running)
+		{
+			ZEN_CONSOLE("Service '{}' is not running", m_ServiceName);
+			return gsl::narrow<int>(ServiceStatusReturnCode::NotRunning);
+		}
+		else
+		{
+			ZEN_CONSOLE("Service '{}':\n{}", m_ServiceName, FmtServiceInfo(Info, "  "));
+		}
+	}
+
+	if (SubOption == &m_InstallOptions)
+	{
+		if (!IsElevated())
+		{
+			return RunElevated(m_AllowElevation);
+		}
+
+		ServiceInfo		Info;
+		std::error_code Ec = QueryInstalledService(m_ServiceName, Info);
+		if (!Ec && Info.Status != ServiceStatus::NotInstalled)
+		{
+			if (m_FullInstall)
+			{
+				if (Info.Status == ServiceStatus::Running)
+				{
+					Ec = StopService(m_ServiceName);
+					if (Ec)
+					{
+						ZEN_CONSOLE("Failed to stop service '{}' using '{}'. Reason: '{}'",
+									m_ServiceName,
+									m_ServerExecutable,
+									Ec.message());
+
+						return 1;
+					}
+
+					int Timeout = 30000;  // Wait up to 30 seconds for the service to fully stop before uninstalling
+					while (Timeout > 0)
+					{
+						Ec = QueryInstalledService(m_ServiceName, Info);
+						if (Ec)
+						{
+							ZEN_CONSOLE("Failed to wait for service to stop: '{}'", Ec.message());
+							return 1;
+						}
+
+						if (Info.Status == ServiceStatus::Stopped)
+						{
+							break;
+						}
+
+						Sleep(100);
+						Timeout -= 100;
+					}
+
+					if (Info.Status != ServiceStatus::Stopped)
+					{
+						ZEN_CONSOLE("Timed out waiting for service to stop");
+						return 1;
+					}
+				}
+
+				Ec = UninstallService(m_ServiceName);
+				if (Ec)
+				{
+					ZEN_CONSOLE("Failed to uninstall running service '{}' using '{}'. Reason: '{}'",
+								m_ServiceName,
+								m_ServerExecutable,
+								Ec.message());
+
+					return 1;
+				}
+			}
+			else
+			{
+				ZEN_CONSOLE("Service '{}' already installed:\n{}", m_ServiceName, FmtServiceInfo(Info, "  "));
+				return 1;
+			}
+		}
+		if (m_ServerExecutable.empty())
+		{
+			std::filesystem::path ExePath = zen::GetRunningExecutablePath();
+			ExePath.replace_filename("zenserver" + ExePath.extension().string());
+			m_ServerExecutable = ExePath;
+		}
+		m_ServerExecutable = std::filesystem::absolute(m_ServerExecutable);
+
+		if (m_FullInstall)
+		{
+			if (m_InstallPath.empty())
+			{
+				throw zen::OptionParseException("--full requires --install-path to be specified");
+			}
+
+			std::filesystem::path ExePath = zen::GetRunningExecutablePath();
+
+			std::filesystem::path FilesToCopy[] = {
+				ExePath,
+				m_ServerExecutable,
+#if ZEN_PLATFORM_WINDOWS
+				ExePath.replace_extension("pdb"),
+				m_ServerExecutable.replace_extension("pdb"),
+				ExePath
+					.replace_filename("crashpad_handler.exe")
+#endif
+#if ZEN_PLATFORM_MAC
+						ExePath.replace_filename("crashpad_handler")
+#endif
+			};
+
+			std::filesystem::path DestinationExePath	= m_InstallPath / ExePath.filename();
+			std::filesystem::path DestinationServerPath = m_InstallPath / m_ServerExecutable.filename();
+
+			if (!std::filesystem::is_directory(m_InstallPath) && !CreateDirectories(m_InstallPath))
+			{
+				ZEN_CONSOLE("Unable to create install directory '{}'", m_InstallPath);
+				return 1;
+			}
+
+			for (const std::filesystem::path& File : FilesToCopy)
+			{
+				std::filesystem::path Destination = m_InstallPath / File.filename();
+				if (!CopyFile(File, Destination, {.EnableClone = false}))
+				{
+					ZEN_CONSOLE("Failed to copy '{}' to '{}'", File, Destination);
+					return 1;
+				}
+
+				ZEN_INFO("Copied '{}' to '{}'", File, Destination);
+
+				if (File.extension() != "pdb")
+				{
+					std::filesystem::permissions(Destination, std::filesystem::perms::owner_exec, std::filesystem::perm_options::add);
+				}
+			}
+
+			m_ServerExecutable = m_InstallPath / m_ServerExecutable.filename();
+		}
+
+		Ec = InstallService(
+			m_ServiceName,
+			ServiceSpec {
+				.ExecutablePath = m_ServerExecutable, .CommandLineOptions = GlobalOptions.PassthroughCommandLine, .UserName = m_UserName
+#if ZEN_PLATFORM_WINDOWS
+					,
+				.DisplayName = m_ServiceDisplayName, .Description = m_ServiceDescription
+#endif	// ZEN_PLATFORM_WINDOWS
+			});
+		if (Ec)
+		{
+			ZEN_CONSOLE("Failed to install service '{}' using '{}' . Reason: '{}'", m_ServiceName, m_ServerExecutable, Ec.message());
+			return 1;
+		}
+		ZEN_CONSOLE("Installed service '{}' using '{}' successfully", m_ServiceName, m_ServerExecutable);
+
+		if (m_FullInstall)
+		{
+			Ec = StartService(m_ServiceName);
+			if (Ec)
+			{
+				ZEN_CONSOLE("Failed to start service '{}'. Reason: '{}'", m_ServiceName, Ec.message());
+				return 1;
+			}
+		}
+	}
+
+	if (SubOption == &m_UninstallOptions)
+	{
+		ServiceInfo		Info;
+		std::error_code Ec = QueryInstalledService(m_ServiceName, Info);
+		if (Ec)
+		{
+			ZEN_CONSOLE("Failed to inspect installed service '{}'. Reason: '{}'", m_ServiceName, Ec.message());
+			return 1;
+		}
+		if (Info.Status == ServiceStatus::NotInstalled)
+		{
+			ZEN_CONSOLE("Service '{}' is not installed", m_ServiceName);
+			return 0;
+		}
+		if (Info.Status != ServiceStatus::Stopped)
+		{
+			ZEN_CONSOLE("Service '{}' is running, stop before uninstalling", m_ServiceName);
+			return 0;
+		}
+
+		if (!IsElevated())
+		{
+			return RunElevated(m_AllowElevation);
+		}
+
+		Ec = UninstallService(m_ServiceName);
+		if (Ec)
+		{
+			ZEN_CONSOLE("Failed to uninstall service '{}'. Reason: '{}'", m_ServiceName, Ec.message());
+			return 1;
+		}
+		ZEN_CONSOLE("Uninstalled service {} successfully", m_ServiceName);
+	}
+
+	if (SubOption == &m_StartOptions)
+	{
+		ServiceInfo		Info;
+		std::error_code Ec = QueryInstalledService(m_ServiceName, Info);
+		if (Ec)
+		{
+			ZEN_CONSOLE("Failed to inspect installed service '{}'. Reason: '{}'", m_ServiceName, Ec.message());
+			return 1;
+		}
+		if (Info.Status == ServiceStatus::NotInstalled)
+		{
+			ZEN_CONSOLE("Service '{}' is not installed", m_ServiceName);
+			return 1;
+		}
+		if (Info.Status != ServiceStatus::Stopped)
+		{
+			ZEN_CONSOLE("Service '{}' is already running", m_ServiceName);
+			return 1;
+		}
+
+		if (!IsElevated())
+		{
+			return RunElevated(m_AllowElevation);
+		}
+
+		Ec = StartService(m_ServiceName);
+		if (Ec)
+		{
+			ZEN_CONSOLE("Failed to start service '{}'. Reason: '{}'", m_ServiceName, Ec.message());
+			return 1;
+		}
+		ZEN_CONSOLE("Started service '{}' successfully", m_ServiceName);
+	}
+
+	if (SubOption == &m_StopOptions)
+	{
+		ServiceInfo		Info;
+		std::error_code Ec = QueryInstalledService(m_ServiceName, Info);
+		if (Ec)
+		{
+			ZEN_CONSOLE("Failed to inspect installed service '{}'. Reason: '{}'", m_ServiceName, Ec.message());
+			return 1;
+		}
+		if (Info.Status == ServiceStatus::NotInstalled)
+		{
+			ZEN_CONSOLE("Service '{}' is not installed", m_ServiceName);
+			return 1;
+		}
+		if (Info.Status != ServiceStatus::Running)
+		{
+			ZEN_CONSOLE("Service '{}' is not running", m_ServiceName);
+			return 1;
+		}
+
+		if (!IsElevated())
+		{
+			return RunElevated(m_AllowElevation);
+		}
+
+		Ec = StopService(m_ServiceName);
+		if (Ec)
+		{
+			ZEN_CONSOLE("Failed to stop service '{}'. Reason: '{}'", m_ServiceName, Ec.message());
+			return 1;
+		}
+		ZEN_CONSOLE("Stopped service '{}' successfully", m_ServiceName);
+	}
+
+	return 0;
+}
+
+}  // namespace zen