diff options
Diffstat (limited to 'zenutil/zenserverprocess.cpp')
| -rw-r--r-- | zenutil/zenserverprocess.cpp | 248 |
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; +} ////////////////////////////////////////////////////////////////////////// |