Files
2025-05-18 13:04:45 +08:00

170 lines
3.6 KiB
C++

// Copyright Epic Games, Inc. All Rights Reserved.
#include "UbaBinaryReaderWriter.h"
#include "UbaDetoursShared.h"
#include "UbaEvent.h"
#include "UbaProtocol.h"
#include "UbaProcessStats.h"
#include <dlfcn.h>
#if PLATFORM_LINUX
#include <sys/prctl.h>
#endif
using namespace uba;
namespace uba
{
int g_comFd = -1;
int g_sessionPid = -1;
SharedEvent* g_cancelEvent;
SharedEvent* g_readEvent;
SharedEvent* g_writeEvent;
u8* g_messageMappingMem;
extern bool g_isCancelled;
void PreInit(const char* logFile);
void Init();
void Deinit();
void AddExceptionHandler();
}
static void __attribute__((constructor(102))) PreInitCtor()
{
using namespace uba;
SuppressDetourScope s;
#if PLATFORM_LINUX
prctl(PR_SET_PDEATHSIG, SIGHUP, 0, 0, 0); // We want the process to die if the parent die
#endif
InitSharedVariables();
g_runningRemote = getenv("UBA_REMOTE") != nullptr;
const char* logFile = getenv("UBA_LOGFILE");
PreInit(logFile);
unsetenv("UBA_LOGFILE");
const char* comId = getenv("UBA_COMID");
//printf("Starting up %s... (comid: %s) %u\n", __progname, (comId && *comId) ? comId : "NOTSET", getpid());
if (!comId || !*comId)
return;
StringBuffer<256> comIdName;
const char* plusIndex = strchr(comId, '+');
comIdName.Append(comId, u64(plusIndex - comId));
u64 comIdUid;
if (!comIdName.Parse(comIdUid))
UBA_ASSERT(false);
u32 comIdOffset = strtoul(plusIndex + 1, nullptr, 10);
unsetenv("UBA_COMID");
comIdName.Clear();
GetMappingHandleName(comIdName, comIdUid);
g_comFd = shm_open(comIdName.data, O_RDWR, S_IRUSR | S_IWUSR);
if (g_comFd == -1)
{
printf("UbaDetours: Failed to open shared mem: %s\n", comIdName.data);
return;
}
u8* rptr = (u8*)mmap(NULL, CommunicationMemSize, PROT_READ | PROT_WRITE, MAP_SHARED, g_comFd, s64(comIdOffset));
if (rptr == MAP_FAILED)
{
printf("UbaDetours: Failed to mmap fd: %u\n", g_comFd);
return;
}
g_cancelEvent = ((SharedEvent*)rptr);
g_readEvent = ((SharedEvent*)rptr) + 1;
g_writeEvent = ((SharedEvent*)rptr) + 2;
g_messageMappingMem = rptr + sizeof(SharedEvent) * 3;
AddExceptionHandler();
g_sessionPid = atoi(getenv("UBA_SESSION_PROCESS"));
g_rulesIndex = strtoul(getenv("UBA_RULES"), nullptr, 10);
g_rules = GetApplicationRules()[g_rulesIndex].rules;
unsetenv("UBA_RULES");
const char* realCwd = getenv("UBA_CWD");
UBA_ASSERT(realCwd);
chdir(realCwd);
unsetenv("UBA_CWD");
Init();
}
void CloseCom()
{
if (g_comFd == -1)
return;
void* mem = g_cancelEvent;
g_cancelEvent = nullptr;
g_readEvent = nullptr;
g_writeEvent = nullptr;
g_messageMappingMem = nullptr;
munmap(mem, CommunicationMemSize);
close(g_comFd);
g_comFd = -1;
}
static void __attribute__((destructor(102))) InitDtor()
{
using namespace uba;
Deinit();
CloseCom();
//printf("Exiting... %s %u \n", __progname, getpid());
}
namespace uba
{
constexpr u32 WritableMemSize = CommunicationMemSize - sizeof(SharedEvent) * 3;
BinaryWriter::BinaryWriter()
{
m_begin = g_messageMappingMem;
m_pos = m_begin;
m_end = m_begin + WritableMemSize;
}
void BinaryWriter::Flush(bool waitOnResponse)
{
g_writeEvent->Set();
if (!waitOnResponse)
return;
TimerScope ts(g_stats.waitOnResponse);
do
{
if (g_readEvent->IsSet(1000))
break;
// check if session process is gone
if (kill(g_sessionPid, 0) == -1 && errno == ESRCH)
{
g_isCancelled = true;
exit(1337);
}
if (g_cancelEvent->IsSet(0))
{
g_isCancelled = true;
exit(1339);
}
} while (true);
}
BinaryReader::BinaryReader()
{
m_begin = g_messageMappingMem;
m_pos = m_begin;
m_end = m_begin + WritableMemSize;
}
}