// Copyright Epic Games, Inc. All Rights Reserved. #pragma once #include "UbaNetwork.h" #include "UbaSession.h" namespace uba { class NetworkClient; struct SessionClientCreateInfo : SessionCreateInfo { SessionClientCreateInfo(Storage& s, NetworkClient& c, LogWriter& writer = g_consoleLogWriter) : SessionCreateInfo(s, writer), client(c) {} void Apply(const Config& config); NetworkClient& client; StringBuffer<128> name; u32 maxProcessCount = 1; u32 defaultPriorityClass = ProcessPriority_BelowNormal; u32 maxIdleSeconds = ~0u; u32 pingTimeoutSecondsPrintCallstacks = 0; u32 osVersion = 0; u8 memWaitLoadPercent = 80; // When memory usage goes above this percent, no new processes will be spawned until back below u8 memKillLoadPercent = 90; // When memory usage goes above this percent, newest processes will be killed to bring it back below bool dedicated = false; // If true, server will not disconnect client when starting to run out of work. bool disableCustomAllocator = false; bool useBinariesAsVersion = false; bool killRandom = false; bool useStorage = true; bool downloadDetoursLib = true; bool useDependencyCrawler = false; // Dependency crawler will go wide and try to prefetch all dependencies needed by the process Function processFinished; }; class SessionClient final : public Session { public: SessionClient(const SessionClientCreateInfo& info); ~SessionClient(); bool Start(); void Stop(bool wait = true); bool Wait(u32 milliseconds = 0xFFFFFFFF, Event* wakeupEvent = nullptr); void SendSummary(const Function& extraInfo); void SetIsTerminating(const tchar* reason = TC("Terminating"), u64 delayMs = 0); // Session stores pointer directly. Can't be temporary void SetMaxProcessCount(u32 count); void SetAllowSpawn(bool allow); u64 GetBestPing(); u32 GetMaxProcessCount(); bool Exists(const StringView& path, u32& outAttributes); private: bool RetrieveCasFile(CasKey& outNewKey, u64& outSize, const CasKey& casKey, const tchar* hint, bool storeUncompressed, bool allowProxy = true); virtual bool PrepareProcess(ProcessImpl& process, bool isChild, StringBufferBase& outRealApplication, const tchar*& outRealWorkingDir) override; virtual void* GetProcessEnvironmentVariables() override; virtual bool ProcessThreadStart(ProcessImpl& process) override; virtual bool CreateFileForRead(CreateFileResponse& out, TrackWorkScope& tws, const StringView& fileName, const StringKey& fileNameKey, ProcessImpl& process, const ApplicationRules& rules) override; virtual bool DeleteFile(DeleteFileResponse& out, const DeleteFileMessage& msg) override; virtual bool CopyFile(CopyFileResponse& out, const CopyFileMessage& msg) override; virtual bool MoveFile(MoveFileResponse& out, const MoveFileMessage& msg) override; virtual bool Chmod(ChmodResponse& out, const ChmodMessage& msg) override; virtual bool CreateDirectory(CreateDirectoryResponse& out, const CreateDirectoryMessage& msg) override; virtual bool RemoveDirectory(RemoveDirectoryResponse& out, const RemoveDirectoryMessage& msg) override; virtual bool GetFullFileName(GetFullFileNameResponse& out, const GetFullFileNameMessage& msg) override; virtual bool GetLongPathName(GetLongPathNameResponse& out, const GetLongPathNameMessage& msg) override; virtual bool GetListDirectoryInfo(ListDirectoryResponse& out, const StringView& dirName, const StringKey& dirKey) override; virtual bool WriteFilesToDisk(ProcessImpl& process, WrittenFile** files, u32 fileCount) override; virtual bool AllocFailed(Process& process, const tchar* allocType, u32 error) override; virtual void PrintSessionStats(Logger& logger) override; virtual bool GetNextProcess(Process& process, bool& outNewProcess, NextProcessInfo& outNextProcess, u32 prevExitCode, BinaryReader& statsReader) override; virtual bool CustomMessage(Process& process, BinaryReader& reader, BinaryWriter& writer) override; virtual bool SHGetKnownFolderPath(Process& process, BinaryReader& reader, BinaryWriter& writer) override; virtual bool HostRun(BinaryReader& reader, BinaryWriter& writer) override; virtual bool GetSymbols(const tchar* application, bool isArm, BinaryReader& reader, BinaryWriter& writer); virtual bool FlushWrittenFiles(ProcessImpl& process) override; virtual bool UpdateEnvironment(ProcessImpl& process, const StringView& reason, bool resetStats) override; virtual bool LogLine(ProcessImpl& process, const tchar* line, LogEntryType logType) override; virtual void TraceSessionUpdate() override; struct InternalProcessStartInfo; struct ModuleInfo; bool GetCasKeyForFile(CasKey& out, u32 processId, const StringView& fileName, const StringKey& fileNameKey); bool ReadModules(List& outModules, u32 processId, const tchar* application); bool EnsureBinaryFile(StringBufferBase& out, StringBufferBase& outVirtual, u32 processId, StringView fileName, const StringKey& fileNameKey, StringView applicationDir, StringView workingDir, const u8* loaderPaths, u32 loaderPathsSize); bool WriteBinFile(StringBufferBase& out, const StringView& binaryName, const CasKey& casKey, const KeyToString& applicationDir, u32 fileAttributes); bool SendFiles(ProcessImpl& process, Timer& sendFiles); bool SendFile(WrittenFile& file, u32 processId, bool keepMappingInMemory, bool compressed); bool SendUpdateDirectoryTable(StackBinaryReader& reader); // Note, reader is sent in to save stack space. bool UpdateDirectoryTableFromServer(StackBinaryReader& reader); bool SendUpdateNameToHashTable(StackBinaryReader& reader); bool UpdateNameToHashTableFromServer(StackBinaryReader& reader); void Connect(); void BuildEnvironmentVariables(BinaryReader& reader); bool SendProcessAvailable(Vector& out, float availableWeight); void SendReturnProcess(u32 processId, const tchar* reason); bool SendProcessInputs(ProcessImpl& process); bool SendProcessFinished(ProcessImpl& process, u32 exitCode); void SendPing(u64 memAvail, u64 memTotal); void SendNotification(const StringView& text); bool SendRootsHandle(RootsHandle rootsHandle); void SendLogFileToServer(ProcessImpl& pi); void GetLogFileName(StringBufferBase& out, const tchar* logFile, const tchar* arguments, u32 processId); u32 WriteLogLines(BinaryWriter& writer, ProcessImpl& process); bool ParseDirectoryTable(); bool EntryExists(const StringView& path, u32& outTableOffset); void ThreadCreateProcessLoop(); NetworkClient& m_client; static constexpr u8 ServiceId = SessionServiceId; StringBuffer<128> m_name; StringBuffer m_processWorkingDir; u32 m_sessionId = 0; u32 m_uiLanguage = 0; u32 m_defaultPriorityClass; u32 m_maxIdleSeconds = ~0u; u32 m_osVersion = 0; u32 m_killRandomIndex = ~0u; u32 m_killRandomCounter = 0; u8 m_memWaitLoadPercent; u8 m_memKillLoadPercent; bool m_disableCustomAllocator; bool m_useBinariesAsVersion; bool m_dedicated = false; bool m_useStorage = true; bool m_downloadDetoursLib = true; bool m_shouldSendLogToServer = false; bool m_shouldSendTraceToServer = false; bool m_remoteExecutionEnabled = true; bool m_useDependencyCrawler = false; Atomic m_terminationReason; Atomic m_terminationTime; Atomic m_maxProcessCount; Atomic m_cpuUsage; struct ApplicationEnvironment { Futex lock; TString virtualApplication; TString realApplication; }; Futex m_handledApplicationEnvironmentsLock; UnorderedMap m_handledApplicationEnvironments; Futex m_binFileLock; UnorderedMap m_writtenBinFiles; Futex m_nameToNameLookupLock; struct NameRec { TString name; TString virtualName; Futex lock; bool handled = false; }; UnorderedMap m_nameToNameLookup; struct HashRec { CasKey key; u64 serverTime = 0; Futex lock; }; UnorderedMap m_nameToHashLookup; Futex m_nameToHashLookupLock; ReaderWriterLock m_nameToHashMemLock; Futex m_directoryTableLock; u32 m_directoryTableMemPos = 0; bool m_directoryTableError = false; struct ActiveUpdateDirectoryEntry; ActiveUpdateDirectoryEntry* m_firstEmptyWait = nullptr; ActiveUpdateDirectoryEntry* m_firstReadWait = nullptr; Event m_waitToSendEvent; Thread m_loopThread; Atomic m_loop; Atomic m_sendPing; Atomic m_allowSpawn; Function m_processFinished; SessionSummaryStats m_stats; Atomic m_bestPing; Atomic m_lastPing; u64 m_lastPingSendTime = 0; u32 m_pingTimeoutSecondsPrintCallstacks = 0; Atomic m_memAvail; Atomic m_memTotal; UnorderedMap> m_hostRunCache; Futex m_hostRunCacheLock; u32 m_dirtableParsedPosition = 0; Futex m_dirVisitedLock; struct DirVisitedEntry { Futex lock; bool handled = false; }; UnorderedMap m_dirVisited; }; }