Added ZenServerState implementation, used to track and enumerate live Zen server instances

1 files changed, 248 insertions, 0 deletions

diff --git a/zenutil/zenserverprocess.cpp b/zenutil/zenserverprocess.cpp
index 9d38d3939..8427e612e 100644
--- a/zenutil/zenserverprocess.cpp
+++ b/zenutil/zenserverprocess.cpp
@@ -2,12 +2,260 @@
 
 #include "zenserverprocess.h"
 
+#include <zencore/except.h>
 #include <zencore/filesystem.h>
 #include <zencore/fmtutils.h>
 #include <zencore/string.h>
 
+#include <atlbase.h>
 #include <shellapi.h>
 #include <spdlog/spdlog.h>
+#include <source_location>
+
+#include <zencore/windows.h>
+
+//////////////////////////////////////////////////////////////////////////
+
+namespace zenutil {
+
+class SecurityAttributes
+{
+public:
+	inline SECURITY_ATTRIBUTES* Attributes() { return &m_Attributes; }
+
+protected:
+	SECURITY_ATTRIBUTES m_Attributes{};
+	SECURITY_DESCRIPTOR m_Sd{};
+};
+
+// Security attributes which allows any user access
+
+class AnyUserSecurityAttributes : public SecurityAttributes
+{
+public:
+	AnyUserSecurityAttributes()
+	{
+		m_Attributes.nLength		= sizeof m_Attributes;
+		m_Attributes.bInheritHandle = false;  // Disable inheritance
+
+		const BOOL success = InitializeSecurityDescriptor(&m_Sd, SECURITY_DESCRIPTOR_REVISION);
+
+		if (success)
+		{
+			const BOOL bSetOk = SetSecurityDescriptorDacl(&m_Sd, TRUE, (PACL)NULL, FALSE);
+
+			if (bSetOk)
+			{
+				m_Attributes.lpSecurityDescriptor = &m_Sd;
+			}
+			else
+			{
+				zen::ThrowLastError("SetSecurityDescriptorDacl failed", std::source_location::current());
+			}
+		}
+	}
+};
+
+}  // namespace zenutil
+
+//////////////////////////////////////////////////////////////////////////
+
+ZenServerState::ZenServerState()
+{
+}
+
+ZenServerState::~ZenServerState()
+{
+	if (m_OurEntry)
+	{
+		// Clean up our entry now that we're leaving
+
+		m_OurEntry->Reset();
+		m_OurEntry = nullptr;
+	}
+
+	if (m_Data)
+	{
+		UnmapViewOfFile(m_Data);
+		m_Data = nullptr;
+	}
+
+	if (m_hMapFile)
+	{
+		CloseHandle(m_hMapFile);
+	}
+}
+
+void
+ZenServerState::Initialize()
+{
+	// TODO: there's a small chance of a race here, this logic could be tightened up with a mutex to
+	//		 ensure only a single process at a time creates the mapping
+
+	if (HANDLE hMap = OpenFileMapping(FILE_MAP_ALL_ACCESS, FALSE, L"Global\\ZenMap"))
+	{
+		m_hMapFile = hMap;
+	}
+	else
+	{
+		// Security attributes to enable any user to access state
+		zenutil::AnyUserSecurityAttributes Attrs;
+
+		hMap = CreateFileMapping(INVALID_HANDLE_VALUE,						// use paging file
+								 Attrs.Attributes(),						// allow anyone to access
+								 PAGE_READWRITE,							// read/write access
+								 0,											// maximum object size (high-order DWORD)
+								 m_MaxEntryCount * sizeof(ZenServerEntry),	// maximum object size (low-order DWORD)
+								 L"Global\\ZenMap");						// name of mapping object
+
+		if (hMap == NULL)
+		{
+			zen::ThrowLastError("Could not open or create file mapping object for Zen server state");
+		}
+
+		m_hMapFile = hMap;
+	}
+
+	void* pBuf = MapViewOfFile(m_hMapFile,			 // handle to map object
+							   FILE_MAP_ALL_ACCESS,	 // read/write permission
+							   0,					 // offset high
+							   0,					 // offset low
+							   m_MaxEntryCount * sizeof(ZenServerEntry));
+
+	if (pBuf == NULL)
+	{
+		zen::ThrowLastError("Could not map view of Zen server state");
+	}
+
+	m_Data = reinterpret_cast<ZenServerEntry*>(pBuf);
+}
+
+bool
+ZenServerState::InitializeReadOnly()
+{
+	if (HANDLE hMap = OpenFileMapping(FILE_MAP_ALL_ACCESS, FALSE, L"Global\\ZenMap"))
+	{
+		m_hMapFile = hMap;
+	}
+	else
+	{
+		return false;
+	}
+
+	void* pBuf = MapViewOfFile(m_hMapFile,	   // handle to map object
+							   FILE_MAP_READ,  // read permission
+							   0,			   // offset high
+							   0,			   // offset low
+							   m_MaxEntryCount * sizeof(ZenServerEntry));
+
+	if (pBuf == NULL)
+	{
+		zen::ThrowLastError("Could not map view of Zen server state");
+	}
+
+	m_Data = reinterpret_cast<ZenServerEntry*>(pBuf);
+
+	return true;
+}
+
+ZenServerState::ZenServerEntry*
+ZenServerState::Lookup(int ListenPort)
+{
+	for (int i = 0; i < m_MaxEntryCount; ++i)
+	{
+		if (m_Data[i].ListenPort == ListenPort)
+		{
+			return &m_Data[i];
+		}
+	}
+
+	return nullptr;
+}
+
+ZenServerState::ZenServerEntry*
+ZenServerState::Register(int ListenPort)
+{
+	if (m_Data == nullptr)
+	{
+		return nullptr;
+	}
+
+	// Allocate an entry
+
+	int Pid = zen::GetCurrentProcessId();
+
+	for (int i = 0; i < m_MaxEntryCount; ++i)
+	{
+		ZenServerEntry& Entry = m_Data[i];
+
+		if (Entry.ListenPort.load(std::memory_order::memory_order_relaxed) == 0)
+		{
+			uint16_t Expected = 0;
+			if (Entry.ListenPort.compare_exchange_strong(Expected, ListenPort))
+			{
+				// Successfully allocated entry
+
+				m_OurEntry = &Entry;
+
+				Entry.Pid = Pid;
+
+				return &Entry;
+			}
+		}
+	}
+
+	return nullptr;
+}
+
+void
+ZenServerState::Sweep()
+{
+	if (m_Data == nullptr)
+	{
+		return;
+	}
+
+	for (int i = 0; i < m_MaxEntryCount; ++i)
+	{
+		ZenServerEntry& Entry = m_Data[i];
+
+		if (Entry.ListenPort)
+		{
+			if (zen::IsProcessRunning(Entry.Pid) == false)
+			{
+				spdlog::debug("Sweep - pid {} not running, reclaiming entry (port {})", Entry.Pid, Entry.ListenPort);
+
+				Entry.Reset();
+			}
+		}
+	}
+}
+
+void
+ZenServerState::Snapshot(std::function<void(const ZenServerEntry&)>&& Callback)
+{
+	if (m_Data == nullptr)
+	{
+		return;
+	}
+
+	for (int i = 0; i < m_MaxEntryCount; ++i)
+	{
+		ZenServerEntry& Entry = m_Data[i];
+
+		if (Entry.ListenPort)
+		{
+			Callback(Entry);
+		}
+	}
+}
+
+void
+ZenServerState::ZenServerEntry::Reset()
+{
+	Pid		   = 0;
+	ListenPort = 0;
+}
 
 //////////////////////////////////////////////////////////////////////////