// Copyright Epic Games, Inc. All Rights Reserved. #include "UbaApplication.h" #include "UbaBinaryParser.h" #include "UbaCacheClient.h" #include "UbaCacheServer.h" #include "UbaClient.h" #include "UbaCompressedFileHeader.h" #include "UbaConfig.h" #include "UbaCoordinatorWrapper.h" #include "UbaFileAccessor.h" #include "UbaNetworkBackendTcp.h" #include "UbaPathUtils.h" #include "UbaPlatform.h" #include "UbaProtocol.h" #include "UbaRootPaths.h" #include "UbaScheduler.h" #include "UbaSessionClient.h" #include "UbaSessionServer.h" #include "UbaStorageClient.h" #include "UbaStorageServer.h" #include "UbaStorageUtils.h" #include "UbaVersion.h" #include "UbaAWS.h" #if PLATFORM_WINDOWS #include #include #pragma comment (lib, "Dbghelp.lib") #endif namespace uba { const tchar* Version = GetVersionString(); u32 DefaultCapacityGb = 20; const tchar* DefaultRootDir = []() { static tchar buf[256]; if constexpr (IsWindows) ExpandEnvironmentStringsW(TC("%ProgramData%\\Epic\\" UE_APP_NAME), buf, sizeof(buf)); else GetFullPathNameW(TC("~/" UE_APP_NAME), sizeof_array(buf), buf, nullptr); return buf; }(); u32 DefaultProcessorCount = []() { return GetLogicalProcessorCount(); }(); bool PrintHelp(const tchar* message) { LoggerWithWriter logger(g_consoleLogWriter, TC("")); if (*message) { logger.Info(TC("")); logger.Error(TC("%s"), message); } const tchar* dbgStr = TC(""); #if UBA_DEBUG dbgStr = TC(" (DEBUG)"); #endif logger.Info(TC("")); logger.Info(TC("-------------------------------------------")); logger.Info(TC(" UbaCli v%s%s"), Version, dbgStr); logger.Info(TC("-------------------------------------------")); logger.Info(TC("")); logger.Info(TC(" UbaCli.exe [options...] [arguments...]")); logger.Info(TC("")); logger.Info(TC(" CommandTypes:")); logger.Info(TC(" local Will run executable locally using detoured paths")); logger.Info(TC(" remote Will wait for available agent and then run executable remotely")); logger.Info(TC(" agent Will run executable against agent spawned in process")); logger.Info(TC(" native Will run executable in a normal way")); logger.Info(TC("")); logger.Info(TC(" Options:")); logger.Info(TC(" -dir= The directory used to store data. Defaults to \"%s\""), DefaultRootDir); logger.Info(TC(" -port=[:] The ip/name and port (default: %u) of the machine we want to help"), DefaultPort); logger.Info(TC(" -log Log all processes detouring information to file (only works with debug builds)")); logger.Info(TC(" -quiet Does not output any logging in console except errors")); logger.Info(TC(" -loop= Loop the commandline number of times. Will exit when/if it fails")); logger.Info(TC(" -workdir= Working directory")); logger.Info(TC(" -config= Config file that contains options for various systems")); logger.Info(TC(" -vfs=; Will convert virtual path to local under the hood. Can have multiple -vfs")); logger.Info(TC(" -checkcas Check so all cas entries are correct")); logger.Info(TC(" -checkfiletable Check so file table has correct cas stored")); logger.Info(TC(" -checkcloud Check if we are inside cloud and output information about cloud")); logger.Info(TC(" -deletecas Deletes the casdb")); logger.Info(TC(" -getcas Will print hash of application")); logger.Info(TC(" -listimports Will print explicit imports of binary")); logger.Info(TC(" -summary Print summary at the end of a session")); logger.Info(TC(" -nocustomalloc Disable custom allocator for processes. If you see odd crashes this can be tested")); logger.Info(TC(" -nostdout Disable stdout from process.")); logger.Info(TC(" -storeraw Disable compression of storage. This will use more storage and might improve performance")); logger.Info(TC(" -maxcpu= Max number of processes that can be started. Defaults to \"%u\" on this machine"), DefaultProcessorCount); logger.Info(TC(" -visualizer Spawn a visualizer that visualizes progress")); logger.Info(TC(" -detailedtrace Add details to the trace")); logger.Info(TC(" -traceChildProcesses Trace the child processes separately")); logger.Info(TC(" -crypto=<32chars> Will enable crypto on network client/server")); logger.Info(TC(" -coordinator= Load a UbaCoordinator.dll to instantiate a coordinator to get helpers")); logger.Info(TC(" -cache=[:] Connect to cache server. Will fetch from cache unless -populatecache is set")); logger.Info(TC(" -populatecache Populate cache server if connected to one")); logger.Info(TC(" -cachecommand= Send command to cache server. Will output result in log")); logger.Info(TC(" -writecachesummary Write cache summary file about connected cache server")); logger.Info(TC("")); logger.Info(TC(" CoordinatorOptions (if coordinator set):")); logger.Info(TC(" -uri=
Uri to coordinator")); logger.Info(TC(" -pool= Name of helper pool inside coordinator")); logger.Info(TC(" -oidc= Name of oidc")); logger.Info(TC(" -maxcores= Max number of cores that will be asked for from coordinator")); logger.Info(TC("")); logger.Info(TC(" If is a .yaml-file UbaCli creates a scheduler to execute commands from the yaml file instead")); logger.Info(TC("")); return false; } StorageServer* g_storageServer; void CtrlBreakPressed() { if (g_storageServer) { g_storageServer->SaveCasTable(true); LoggerWithWriter(g_consoleLogWriter).Info(TC("CAS table saved...")); } abort(); } #if PLATFORM_WINDOWS BOOL ConsoleHandler(DWORD signal) { if (signal == CTRL_C_EVENT) CtrlBreakPressed(); return FALSE; } #else void ConsoleHandler(int sig) { CtrlBreakPressed(); } #endif StringBuffer<> g_rootDir(DefaultRootDir); bool WrappedMain(int argc, tchar* argv[]) { using namespace uba; AddExceptionHandler(); InitMemory(); u32 storageCapacityGb = DefaultCapacityGb; StringBuffer<256> workDir; StringBuffer<128> listenIp; StringBuffer<128> cacheHost; TString crypto; TString coordinatorName; TString coordinatorPool; u32 coordinatorMaxCoreCount = 400; u16 port = DefaultPort; u16 cachePort = DefaultCachePort; u32 maxProcessCount = DefaultProcessorCount; u32 agentCount = 1; bool launchVisualizer = false; bool storeCompressed = true; bool disableCustomAllocator = false; bool quiet = false; bool checkCas = false; bool checkCas2 = false; bool checkCloud = false; bool getCas = false; bool listImports = false; bool deleteCas = false; bool enableStdOut = true; bool printSummary = false; bool detailedTrace = false; bool traceChildProcesses = false; bool populateCache = false; bool writeCacheSummary = false; bool logToFile = false; bool useHackVfs = false; TString checkFileTable; TString cacheFilterString; TString cacheCommand; TString testCompress; TString testDecompress; TString addCas; TString configFile; struct VfsEntry { TString virtualPath; TString localPath; }; Vector vfsEntries; u32 loopCount = 1; enum CommandType { CommandType_NotSet, CommandType_Local, CommandType_Remote, CommandType_Native, CommandType_Agent, CommandType_None, }; CommandType commandType = CommandType_NotSet; TString application; TString arguments; auto parseOption = [&](const StringView& name, StringBufferBase& value) { if (IsWindows && name.Equals(TCV("-visualizer"))) { launchVisualizer = true; } else if (name.Equals(TCV("-crypto"))) { if (value.IsEmpty()) value.Append(TCV("0123456789abcdef0123456789abcdef")); crypto = value.data; } else if (name.Equals(TCV("-coordinator"))) { if (value.IsEmpty()) return PrintHelp(TC("-coordinator needs a value")); coordinatorName = value.data; } else if (name.Equals(TCV("-pool"))) { if (value.IsEmpty()) return PrintHelp(TC("-pool needs a value")); coordinatorPool = value.data; } else if (name.Equals(TCV("-maxcores"))) { if (value.IsEmpty()) return PrintHelp(TC("-maxcores needs a value")); if (!value.Parse(coordinatorMaxCoreCount)) return PrintHelp(TC("Invalid value for -maxcores")); } else if (name.Equals(TCV("-workdir"))) { if (value.IsEmpty()) return PrintHelp(TC("-workdir needs a value")); if ((workDir.count = GetFullPathNameW(value.data, workDir.capacity, workDir.data, nullptr)) == 0) return PrintHelp(StringBuffer<>().Appendf(TC("-workdir has invalid path %s"), value.data).data); } else if (name.Equals(TCV("-config"))) { if (value.IsEmpty()) return PrintHelp(TC("-config needs a value")); if (!ExpandEnvironmentVariables(value, PrintHelp)) return false; configFile = value.data; } else if (name.Equals(TCV("-vfs"))) { if (value.IsEmpty()) return PrintHelp(TC("-vfs needs a value")); const tchar* semi = value.First(';'); if (!semi) return PrintHelp(TC("-vfs needs a semicolon between virtual and local path")); u32 semiPos = u32(semi - value.data); vfsEntries.push_back({StringView(value.data, semiPos).ToString(), StringView(value).Skip(semiPos + 1).ToString()}); } else if (name.Equals(TCV("-capacity"))) { if (!value.Parse(storageCapacityGb)) return PrintHelp(TC("Invalid value for -capacity")); } else if (name.Equals(TCV("-port"))) { if (const tchar* portIndex = value.First(':')) { StringBuffer<> portStr(portIndex + 1); if (!portStr.Parse(port)) return PrintHelp(TC("Invalid value for port in -port")); listenIp.Append(value.data, portIndex - value.data); } else { if (!value.Parse(port)) return PrintHelp(TC("Invalid value for -port")); } } else if (name.Equals(TCV("-log"))) { logToFile = true; } else if (name.Equals(TCV("-loop"))) { if (!value.Parse(loopCount)) return PrintHelp(TC("Invalid value for -loop")); } else if (name.Equals(TCV("-quiet"))) { quiet = true; } else if (name.Equals(TCV("-nocustomalloc"))) { disableCustomAllocator = true; } else if (name.Equals(TCV("-maxcpu"))) { if (!value.Parse(maxProcessCount)) return PrintHelp(TC("Invalid value for -maxcpu")); } else if (name.Equals(TCV("-nostdout"))) { enableStdOut = false; } else if (name.Equals(TCV("-checkcas"))) { checkCas = true; } else if (name.Equals(TCV("-checkfiletable"))) { if (value.IsEmpty()) return PrintHelp(TC("-checkfiletable needs a value")); StringBuffer<> temp; if ((temp.count = GetFullPathNameW(value.Replace('/', PathSeparator).data, temp.capacity, temp.data, nullptr)) == 0) return PrintHelp(StringBuffer<>().Appendf(TC("-checkfiletable has invalid path %s"), temp.data).data); checkFileTable = temp.data; } else if (name.Equals(TCV("-checkcas2"))) { checkCas2 = true; } else if (name.Equals(TCV("-checkcloud"))) { checkCloud = true; } else if (name.Equals(TCV("-testcompress"))) { if (value.IsEmpty()) return PrintHelp(TC("-testCompress needs a value")); testCompress = value.data; } else if (name.Equals(TCV("-testdecompress"))) { if (value.IsEmpty()) { if (testCompress.empty()) return PrintHelp(TC("-testDecompress needs a value")); value.Clear().Append(g_rootDir).EnsureEndsWithSlash().Append(TCV("castemp")).EnsureEndsWithSlash().Append(TCV("TestCompress.tmp")); } testDecompress = value.data; } else if (name.Equals(TCV("-deletecas"))) { deleteCas = true; } else if (name.Equals(TCV("-addcas"))) { addCas = value.data; } else if (name.Equals(TCV("-getcas"))) { getCas = true; } else if (name.Equals(TCV("-listimports"))) { listImports = true; } else if (name.Equals(TCV("-summary"))) { printSummary = true; } else if (name.Equals(TCV("-detailedtrace"))) { detailedTrace = true; } else if (name.Equals(TCV("-traceChildProcesses"))) { traceChildProcesses = true; } else if (name.Equals(TCV("-hackvfs"))) { useHackVfs = true; } else if (name.Equals(TCV("-cache"))) { if (value.IsEmpty()) return PrintHelp(TC("-cache needs a value")); if (const tchar* colon = value.First(':')) { value.Parse(cachePort, colon - value.data + 1); cacheHost.Append(value.data, colon - value.data); } else cacheHost.Append(value); } else if (name.Equals(TCV("-populatecache"))) { populateCache = true; } else if (name.Equals(TCV("-cachecommand"))) { if (value.IsEmpty()) return PrintHelp(TC("-cachecommand needs a value")); cacheCommand = value.data; commandType = CommandType_None; quiet = true; } else if (name.Equals(TCV("-writecachesummary"))) { writeCacheSummary = true; cacheFilterString = value.data; commandType = CommandType_None; } else if (name.Equals(TCV("-storeraw"))) { storeCompressed = false; } else if (name.Equals(TCV("-dir"))) { if (value.IsEmpty()) return PrintHelp(TC("-dir needs a value")); if ((g_rootDir.count = GetFullPathNameW(value.Replace('/', PathSeparator).data, g_rootDir.capacity, g_rootDir.data, nullptr)) == 0) return PrintHelp(StringBuffer<>().Appendf(TC("-dir has invalid path %s"), g_rootDir.data).data); } else if (name.Equals(TCV("-?"))) { return PrintHelp(TC("")); } else { StringBuffer<> msg; msg.Appendf(TC("Unknown argument '%s'"), name.data); return PrintHelp(msg.data); } return true; }; auto parseArg = [&](const tchar* arg) { StringBuffer<> name; StringBuffer<32*1024> value; if (const tchar* equals = TStrchr(arg,'=')) { name.Append(arg, equals - arg); value.Append(equals+1); } else { name.Append(arg); } if (!application.empty()) { if (!arguments.empty()) arguments += ' '; TString argTemp; bool hasSpace = TStrchr(arg, ' '); if (hasSpace) { argTemp = arg; size_t index = 0; while (true) { index = argTemp.find('\"', index); if (index == std::string::npos) break; argTemp.replace(index, 1, TC("\\\"")); index += 2; } arg = argTemp.c_str(); arguments += TC("\""); } arguments += arg; if (hasSpace) arguments += TC("\""); return true; } if (commandType != CommandType_NotSet) { application = arg; } else if (name.Equals(TCV("local"))) { commandType = CommandType_Local; } else if (name.Equals(TCV("remote"))) { commandType = CommandType_Remote; } else if (name.Equals(TCV("native"))) { commandType = CommandType_Native; } else if (name.Equals(TCV("agent"))) { commandType = CommandType_Agent; } else { return parseOption(name, value); } return true; }; for (int i=1; i!=argc; ++i) if (!parseArg(argv[i])) return false; auto addOption = [&](const tchar* name, const tchar* value) { StringBuffer<512> v(value); return parseOption(ToView(name), v); };(void)addOption; if (useHackVfs) { #if PLATFORM_WINDOWS addOption(TC("-vfs"), TC("Z:/UEVFS/FortniteGame;E:\\dev\\fn\\FortniteGame")); addOption(TC("-vfs"), TC("Z:/UEVFS/QAGame;E:\\dev\\fn\\QAGame")); addOption(TC("-vfs"), TC("Z:/UEVFS/Root;E:\\dev\\fn")); //addOption(TC("-vfs"), TC("Z:/UEVFS/Clang;c:\\sdk\\AutoSDK\\HostWin64\\Win64\\LLVM\\18.1.8")); //addOption(TC("-vfs"), TC("Z:/UEVFS/MSVC;c:\\sdk\\AutoSDK\\HostWin64\\Win64\\VS2022\\14.38.33130")); //addOption(TC("-vfs"), TC("Z:/UEVFS/WinSDK;C:\\Program Files (x86)\\Windows Kits\\10")); addOption(TC("-vfs"), TC("Z:/UEVFS/Clang;\\\\localhost\\c$\\sdk\\AutoSDK\\HostWin64\\Win64\\LLVM\\18.1.8")); addOption(TC("-vfs"), TC("Z:/UEVFS/MSVC;\\\\localhost\\c$\\sdk\\AutoSDK\\HostWin64\\Win64\\VS2022\\14.38.33130")); addOption(TC("-vfs"), TC("Z:/UEVFS/WinSDK;C:\\Program Files (x86)\\Windows Kits\\10")); addOption(TC("-vfs"), TC("Z:/UEVFS/SuperLuminal;C:\\Program Files\\Superluminal\\Performance\\API")); //addOption(TC("-vfs"), TC("Z:/UEVFS/Root;E:\\dev\\fn")); //addOption(TC("-vfs"), TC("Z:/UEVFS/Compiler;c:\\sdk\\AutoSDK\\HostWin64\\Win64\\LLVM\\18.1.8")); //addOption(TC("-vfs"), TC("Z:/UEVFS/Toolchain;c:\\sdk\\AutoSDK\\HostWin64\\Win64\\VS2022\\14.38.33130")); //addOption(TC("-vfs"), TC("Z:/UEVFS/WinSDK;c:\\sdk\\AutoSDK\\HostWin64\\Win64\\Windows Kits\\10.0.22621.0")); #else addOption(TC("-vfs"), TC("/UEVFS/Root;/home/honk/fn")); addOption(TC("-vfs"), TC("/UEVFS/LinuxSDK;/home/honk/AutoSDK/HostLinux/Linux_x64/v23_clang-18.1.0-rockylinux8/x86_64-unknown-linux-gnu")); #endif } FilteredLogWriter logWriter(g_consoleLogWriter, quiet ? LogEntryType_Warning : LogEntryType_Detail); LoggerWithWriter logger(logWriter, TC("")); Config config; if (!configFile.empty()) config.LoadFromFile(logger, configFile.c_str()); if constexpr (!IsArmBinary) if (IsRunningArm()) logger.Warning(TC(" Running x64 binary on arm64 system. Use arm binaries instead")); bool exit = false; if (deleteCas) { StorageImpl(StorageCreateInfo(g_rootDir.data, logWriter)).DeleteAllCas(); for (u32 i=0; i!=agentCount; ++i) { StringBuffer<> clientRootDir; clientRootDir.Append(g_rootDir).Append("Agent").AppendValue(i); StorageImpl(StorageCreateInfo(clientRootDir.data, logWriter)).DeleteAllCas(); } exit = true; } if (!addCas.empty()) { StorageCreateInfo storageInfo(g_rootDir.data, logWriter); storageInfo.casCapacityBytes = 0; storageInfo.storeCompressed = storeCompressed; StorageImpl storage(storageInfo); CasKey casKey; if (!storage.StoreCasFile(casKey, addCas.c_str(), CasKeyZero, false)) return false; exit = true; } if (checkCas) { StorageCreateInfo storageInfo(g_rootDir.data, logWriter); storageInfo.casCapacityBytes = 0; storageInfo.storeCompressed = storeCompressed; StorageImpl storage(storageInfo); if (!storage.CheckCasContent(DefaultProcessorCount)) return false; exit = true; } if (!checkFileTable.empty()) { StorageCreateInfo storageInfo(g_rootDir.data, logWriter); storageInfo.casCapacityBytes = 0; storageInfo.storeCompressed = storeCompressed; StorageImpl storage(storageInfo); if (!storage.LoadCasTable()) return false; if (!storage.CheckFileTable(checkFileTable.data(), DefaultProcessorCount)) return false; exit = true; } if (checkCas2) // Creates a storage server and storage client and transfer _all_ cas files over network { NetworkBackendTcp networkBackend(logWriter); NetworkServerCreateInfo nsci(logWriter); bool ctorSuccess = true; NetworkServer server(ctorSuccess, nsci); StorageServerCreateInfo storageInfo(server, g_rootDir.data, logWriter); storageInfo.casCapacityBytes = 0; storageInfo.storeCompressed = storeCompressed; StorageServer storageServer(storageInfo); StringBuffer<> rootDir2(g_rootDir.data); rootDir2.Append("_CHECKCAS2"); DeleteAllFiles(logger, rootDir2.data); Client client; auto g = MakeGuard([&]() { server.DisconnectClients(); }); if (!server.StartListen(networkBackend, 1347, TC("127.0.0.1"))) return false; ClientInitInfo cii { logWriter, networkBackend, rootDir2.data, TC("127.0.0.1"), 1347, TC("foo") }; cii.createSession = false; cii.addDirSuffix = false; if (!client.Init(cii)) return false; bool success = true; WorkManagerImpl workManager(DefaultProcessorCount, TC("UbaWrk/ChkCas2")); storageServer.TraverseAllCasFiles([&](const CasKey& casKey, u64 size) { workManager.AddWork([&, casKey](const WorkContext&) { Storage::RetrieveResult res; storageServer.EnsureCasFile(casKey, TC("Dummy")); if (!client.storageClient->RetrieveCasFile(res, AsCompressed(casKey, false), TC(""))) success = false; if (!client.storageClient->RetrieveCasFile(res, casKey, TC(""))) success = false; #if 0 StorageStats storageStats; FileFetcher fetcher { client.storageClient->m_bufferSlots, storageStats }; bool destinationIsCompressed = false; if (!fetcher.RetrieveFile(logger, *client.networkClient, casKey, TC("e:\\temp\\foo"), destinationIsCompressed)) success = false; #endif }, 1, TC("CheckCas2")); }); workManager.FlushWork(); if (!success) return false; exit = true; } #if UBA_USE_CLOUD if (checkCloud) { DirectoryCache dirCache; dirCache.CreateDirectory(logger, g_rootDir.data); Cloud cloud; StringBuffer<> info; if (cloud.QueryInformation(logger, info, g_rootDir.data)) { logger.Info(TC("We are inside cloud%s (%s)"), info.data, cloud.GetAvailabilityZone()); StringBuffer<> reason; u64 terminateTime; if (cloud.IsTerminating(logger, reason, terminateTime)) logger.Info(TC(".. and are being terminated: %s"), reason.data); } else logger.Info(TC("Seems like we are not running inside cloud.")); exit = true; } #endif u64 testCompressOriginalSize = 0; if (!testCompress.empty()) { WorkManagerImpl workManager(DefaultProcessorCount, TC("UbaWrk/TstComp")); FileAccessor fa(logger, testCompress.c_str()); if (!fa.OpenMemoryRead()) return logger.Error(TC("Failed to open file %s"), testCompress.c_str()); u64 fileSize = fa.GetSize(); u8* mem = fa.GetData(); testCompressOriginalSize = fileSize; StorageCreateInfo storageInfo(g_rootDir.data, logWriter); storageInfo.casCapacityBytes = 0; storageInfo.storeCompressed = storeCompressed; storageInfo.workManager = &workManager; StorageImpl storage(storageInfo); Storage::WriteResult res; CompressedFileHeader header { CalculateCasKey(mem, fileSize, true, &workManager, testCompress.c_str()) }; StringBuffer<> dest; dest.Append(storage.GetTempPath()).Append(TCV("TestCompress.tmp")); if (!storage.WriteCompressed(res, TC("MemoryMap"), InvalidFileHandle, mem, fileSize, dest.data, &header, sizeof(header), 0)) return false; if (testDecompress.empty()) return true; logger.Info(TC("Compressing %s successful (Written to %s)"), testCompress.c_str(), dest.data); exit = true; } if (!testDecompress.empty()) { WorkManagerImpl workManager(DefaultProcessorCount, TC("UbaWrk/TstDecm")); FileAccessor fa(logger, testDecompress.c_str()); if (!fa.OpenMemoryRead()) return logger.Error(TC("Failed to open file %s"), testDecompress.c_str()); u64 fileSize = fa.GetSize(); u8* mem = fa.GetData(); if (fileSize < 16) return logger.Error(TC("File %s is too small to be compressed. Requires at least 16 bytes"), testDecompress.c_str()); StorageCreateInfo storageInfo(g_rootDir.data, logWriter); storageInfo.casCapacityBytes = 0; storageInfo.storeCompressed = storeCompressed; storageInfo.workManager = &workManager; StorageImpl storage(storageInfo); BinaryReader reader(mem, 0, fileSize); auto& h = *(CompressedFileHeader*)mem; if (h.IsValid()) reader.Skip(sizeof(CompressedFileHeader)); u64 decompressedSize = reader.ReadU64(); if (testCompressOriginalSize && decompressedSize != testCompressOriginalSize) return logger.Error(TC("Compressed file %s has wrong decompressed size. (Is it compressed?)"), testDecompress.c_str()); StringBuffer<> dest; dest.Append(storage.GetTempPath()).Append(TCV("TestDecompress.tmp")); FileAccessor faDest(logger, dest.data); if (!faDest.CreateMemoryWrite(false, DefaultAttributes(), decompressedSize)) return false; u8* destMem = faDest.GetData(); OO_SINTa decoredMemSize = OodleLZDecoder_MemorySizeNeeded(OodleLZ_Compressor_Kraken); void* decoderMem = malloc(decoredMemSize); auto mg = MakeGuard([decoderMem]() { free(decoderMem); }); while (reader.GetLeft()) { u32 compressedBlockSize = reader.ReadU32(); u32 decompressedBlockSize = reader.ReadU32(); OO_SINTa decompLen = OodleLZ_Decompress(reader.GetPositionData(), (OO_SINTa)compressedBlockSize, destMem, (OO_SINTa)decompressedBlockSize, OodleLZ_FuzzSafe_Yes, OodleLZ_CheckCRC_No, OodleLZ_Verbosity_None, NULL, 0, NULL, NULL, decoderMem, decoredMemSize); if (decompLen != decompressedBlockSize) return logger.Error(TC("Failed to decompress %s (CompressedSize: %llu DecompressedSize: %llu ReadPos: %llu CompressedBlock: %u DecompressedBlock: %u)"), testDecompress.c_str(), fileSize, decompressedSize, reader.GetPosition(), compressedBlockSize, decompressedBlockSize); destMem += decompressedBlockSize; reader.Skip(compressedBlockSize); } if (!faDest.Close()) return false; logger.Info(TC("Decompressing %s successful (Written to %s)"), testDecompress.c_str(), dest.data); exit = true; } if (exit) return true; if (commandType == CommandType_NotSet) { const tchar* errorMsg = argc == 1 ? TC("") : TC("\nERROR: First argument must be command type. Options are 'local,remote or native'"); StringBuffer<> msg; return PrintHelp(errorMsg); } StringBuffer<512> currentDir; GetCurrentDirectoryW(currentDir); if (commandType != CommandType_None) { if (application.empty()) return PrintHelp(TC("No executable provided")); if (!IsAbsolutePath(application.c_str())) { StringBuffer<> fullApplicationName; if (!SearchPathForFile(logger, fullApplicationName, application.c_str(), currentDir, {})) return logger.Error(TC("Failed to find full path to %s"), application.c_str()); application = fullApplicationName.data; } if (getCas) { FileAccessor fa(logger, application.c_str()); if (!fa.OpenMemoryRead()) return logger.Error(TC("Failed to open file %s"), application.c_str()); u64 fileSize = fa.GetSize(); u8* data = fa.GetData(); bool is64Bit = false; bool isArm64 = false; bool isX64 = false; bool isDotnet = false; CasKey key = CalculateCasKey(data, fileSize, false, nullptr, application.c_str()); CasKey uncompressedKey; if (fileSize > sizeof(CompressedFileHeader)) { auto& hdr = *(CompressedFileHeader*)data; if (hdr.IsValid()) uncompressedKey = hdr.casKey; } if (data[0] == 'M' && data[1] == 'Z') { u32 offset = *(u32*)(data + 0x3c); u32* signaturePos = (u32*)(data + offset); is64Bit = *signaturePos == 0x00004550; if (is64Bit) { u16 machine = *(u16*)(signaturePos + 1); isX64 = machine == 0x8664; isArm64 = machine == 0xaa64; if (fileSize > offset + 0x18 + 0x70 + 4) isDotnet = *(u32*)(data + offset + 0x18 + 0x70); } } logger.Info(TC("%s"), application.c_str()); logger.Info(TC(" Is64Bit: %s"), (is64Bit ? TC("true") : TC("false"))); logger.Info(TC(" Arch: %s"), (isX64 ? TC("x64") : (isArm64 ? TC("arm64") : (isDotnet ? TC(".net") : TC("unknown"))))); logger.Info(TC(" Size: %llu"), fileSize); logger.Info(TC(" CasKey: %s"), CasKeyString(key).str); if (uncompressedKey != CasKeyZero) logger.Info(TC(" CasKey (uncompressed): %s"), CasKeyString(uncompressedKey).str); return true; } if (listImports) { StringBuffer<> error; bool printImports = true; BinaryInfo info; if (!ParseBinary(application, StringView(application).GetPath(), info, [&](const tchar* import, bool isKnown, const char* const* loaderPaths) { if (printImports) { if (loaderPaths && *loaderPaths) { logger.Info(TC("LoaderPaths:")); for (auto it=loaderPaths;*it; ++it) if (**it) logger.Info(TC(" %s"), *it); } printImports = false; logger.Info(TC("Imports:")); } logger.Info(TC(" %s"), import); }, error)) return logger.Error(TC("%s"), error.data); #if PLATFORM_MAC logger.Info(TC("MinOsVersion: %u.%u.%u"), (info.minVersion >> 16) & 0xffff, (info.minVersion >> 8) & 0xff, info.minVersion & 0xff); #endif return true; } } const tchar* dbgStr = TC(""); #if UBA_DEBUG dbgStr = TC(" (DEBUG)"); #endif logger.Info(TC("UbaCli v%s%s (Rootdir: \"%s\", StoreCapacity: %uGb)\n"), Version, dbgStr, g_rootDir.data, storageCapacityGb); u64 storageCapacity = u64(storageCapacityGb)*1000*1000*1000; if (workDir.IsEmpty()) workDir.Append(currentDir); // TODO: Change workdir to make it full #if UBA_DEBUG logToFile = true; #endif StringBuffer<> logFile; if (logToFile) { logFile.count = GetFullPathNameW(g_rootDir.data, logFile.capacity, logFile.data, nullptr); logFile.EnsureEndsWithSlash().Append(TCV("DebugLog.log")); logger.Info(TC("Logging to file: %s"), logFile.data); } #if PLATFORM_WINDOWS SetConsoleCtrlHandler(ConsoleHandler, TRUE); #else signal(SIGINT, ConsoleHandler); signal(SIGTERM, ConsoleHandler); #endif NetworkBackendTcp networkBackend(logWriter); NetworkServerCreateInfo nsci(logWriter); nsci.Apply(config); //nsci.workerCount = 4; bool ctorSuccess = true; NetworkServer& networkServer = *new NetworkServer(ctorSuccess, nsci); auto destroyServer = MakeGuard([&]() { delete &networkServer; }); if (!ctorSuccess) return false; if (!crypto.empty()) { u8 crypto128Data[16]; if (!CryptoFromString(crypto128Data, 16, crypto.c_str())) return logger.Error(TC("Failed to parse crypto key %s"), crypto.c_str()); networkServer.RegisterCryptoKey(crypto128Data); logger.Info(TC("Using crypto key %s for connections"), crypto.c_str()); } bool isRemote = commandType == CommandType_Remote || commandType == CommandType_Agent; bool useScheduler = StringView(application).EndsWith(TCV(".yaml")); StorageServerCreateInfo storageInfo(networkServer, g_rootDir.data, logWriter); storageInfo.casCapacityBytes = storageCapacity; storageInfo.storeCompressed = storeCompressed; storageInfo.Apply(config); StorageServer& storageServer = *new StorageServer(storageInfo); auto destroyStorage = MakeGuard([&]() { delete &storageServer; }); SessionServerCreateInfo info(storageServer, networkServer, logWriter); info.useUniqueId = useScheduler; info.traceEnabled = cacheCommand.empty() || writeCacheSummary; info.detailedTrace = detailedTrace; info.traceChildProcesses = traceChildProcesses; info.launchVisualizer = launchVisualizer; info.disableCustomAllocator = disableCustomAllocator; //info.shouldWriteToDisk = shouldWriteToDisk; info.rootDir = g_rootDir.data; //info.traceName.Append(TCV("TESTTRACE")); //info.storeIntermediateFilesCompressed = false; info.readIntermediateFilesCompressed = true; //info.extractObjFilesSymbols = true; #if UBA_DEBUG_LOG_ENABLED info.remoteLogEnabled = true; #endif //info.remoteTraceEnabled = true; info.deleteSessionsOlderThanSeconds = 1; info.Apply(config); SessionServer& sessionServer = *new SessionServer(info); auto destroySession = MakeGuard([&]() { delete &sessionServer; }); CacheClient* cacheClient = nullptr; auto ccg = MakeGuard([&]() { if (!cacheClient) return; auto& nc = cacheClient->GetClient(); nc.Disconnect(); delete cacheClient; delete &nc; }); auto CreateCacheClient = [&]() { auto nc = new NetworkClient(ctorSuccess, {logWriter}); cacheClient = new CacheClient({logWriter, storageServer, *nc, sessionServer}); }; if (cacheHost.count) { CreateCacheClient(); if (!cacheClient->GetClient().Connect(networkBackend, cacheHost.data, cachePort)) return logger.Error(TC("Failed to connect to cache server")); if (!storageServer.LoadCasTable(true)) return false; if (!cacheCommand.empty()) { LoggerWithWriter consoleLogger(g_consoleLogWriter); const tchar* additionalInfo = nullptr; return cacheClient->ExecuteCommand(consoleLogger, cacheCommand.data(), nullptr, additionalInfo); } if (writeCacheSummary) { StringBuffer<> tempFile(sessionServer.GetTempPath()); Guid guid; CreateGuid(guid); tempFile.Append(GuidToString(guid).str).Append(TCV(".txt")); if (!cacheClient->ExecuteCommand(logger, TC("content"), tempFile.data, cacheFilterString.data())) return false; logger.Info(TC("Cache status summary written to %s"), tempFile.data); #if PLATFORM_WINDOWS ShellExecuteW(NULL, L"open", tempFile.data, NULL, NULL, SW_SHOW); #endif return true; } } // Remove empty spaces and line feeds etc at the end.. just to solve annoying copy paste command lines and accidentally getting line feed while (!arguments.empty()) { tchar lastChar = arguments[arguments.size()-1]; if (lastChar != '\n' && lastChar != '\r' && lastChar != '\t' && lastChar != ' ') break; arguments.resize(arguments.size() - 1); } // Vfs testing RootsHandle rootsHandle = 0; if (!vfsEntries.empty()) { StackBinaryWriter<8*1024> writer; for (VfsEntry& entry : vfsEntries) { writer.WriteByte(0); writer.WriteString(entry.virtualPath); writer.WriteString(entry.localPath); } rootsHandle = sessionServer.RegisterRoots(writer.GetData(), writer.GetPosition()); sessionServer.DevirtualizeString(application, rootsHandle, true); } if (isRemote || useScheduler) { if (!storageServer.m_casTableLoaded) if (!storageServer.LoadCasTable(true)) return false; if (!networkServer.StartListen(networkBackend, port, listenIp.data)) return false; } auto stopServer = MakeGuard([&]() { networkServer.DisconnectClients(); }); auto stopListen = MakeGuard([&]() { networkBackend.StopListen(); }); auto RunLocal = [&](const TString& app, const TString& arg, bool enableDetour) { u64 start = GetTime(); ProcessStartInfo pinfo; pinfo.description = app.c_str(); pinfo.application = app.c_str(); pinfo.arguments = arg.c_str(); pinfo.workingDir = workDir.data; pinfo.rootsHandle = rootsHandle; u32 bucketId = 1337; if (cacheClient) { CacheResult cacheResult; cacheClient->FetchFromCache(cacheResult, RootPaths(), bucketId, pinfo); if (cacheResult.hit) { logger.Info(TC("Cached run took %s"), TimeToText(GetTime() - start).str); return true; } } pinfo.logFile = logFile.data; pinfo.logLineUserData = &logger; if (enableStdOut) pinfo.logLineFunc = [](void* userData, const tchar* line, u32 length, LogEntryType type) { ((Logger*)userData)->Log(type, line, length); }; if (populateCache) pinfo.trackInputs = true; logger.Info(TC("Running %s %s"), app.c_str(), arg.c_str()); ProcessHandle process = sessionServer.RunProcess(pinfo, false, enableDetour); if (process.GetExitCode() != 0) return logger.Error(TC("Error exit code: %u"), process.GetExitCode()); logger.Info(TC("%s run took %s"), (enableDetour ? TC("Detoured") : TC("Native")), TimeToText(GetTime() - start).str); if (populateCache) { return logger.Error(TC("Populating cache not implemented... todo")); //RootPaths rootPaths; //cacheClient->WriteToCache(rootPaths, 0, pinfo, nullptr, 1, nullptr, 0, nullptr, 0); } return true; }; auto RunRemote = [&](const TString& app, const TString& arg) { u64 start = GetTime(); ProcessStartInfo pinfo; pinfo.description = app.c_str(); pinfo.application = app.c_str(); pinfo.arguments = arg.c_str(); pinfo.workingDir = workDir.data; pinfo.logFile = logFile.data; pinfo.logLineUserData = &logger; pinfo.rootsHandle = rootsHandle; if (enableStdOut) pinfo.logLineFunc = [](void* userData, const tchar* line, u32 length, LogEntryType type) { ((Logger*)userData)->Log(type, line, length); }; logger.Info(TC("Running %s %s"), app.c_str(), arg.c_str()); ProcessHandle process = sessionServer.RunProcessRemote(pinfo); process.WaitForExit(~0u); if (process.GetExitCode() != 0) return logger.Error(TC("Error exit code: %u"), process.GetExitCode()); u64 time = GetTime() - start; logger.Info(TC("Remote run took %s"), TimeToText(time).str); return true; }; const tchar* clientZone = TC("DummyZone"); auto RunWithClient = [&](const Function& func, u32 clientCount) { Vector clients; clients.resize(clientCount); u32 clientIndex = 0; for (auto& c : clients) { u32 maxProcessor = Min(maxProcessCount/u32(clients.size()), 32u); ClientInitInfo cii { logWriter, networkBackend, g_rootDir.data, TC("127.0.0.1"), port, clientZone, maxProcessor, clientIndex++}; if (!c.Init(cii)) return false; } return func(); }; auto RunAgent = [&](const TString& app, const TString& arg) { return RunWithClient([&]() { return RunRemote(app, arg); }, 1); }; CoordinatorWrapper coordinator; auto RunScheduler = [&](const tchar* yamlFile) { /* bool ctorSuccess; NetworkBackendMemory nbm(logWriter); StringBuffer<> cacheRootDir(g_rootDir); cacheRootDir.Append("CacheServer"); NetworkServer cacheNetworkServer(ctorSuccess); StorageServerCreateInfo storageInfo2(cacheNetworkServer, cacheRootDir.data, logWriter); storageInfo2.writeReceivedCasFilesToDisk = true; StorageServer cacheStorageServer(storageInfo2); CacheServer cacheServer(logWriter, cacheRootDir.data, cacheNetworkServer, cacheStorageServer); auto csg = MakeGuard([&]() { cacheServer.Save(); }); NetworkClient cacheNetworkClient(ctorSuccess); CacheClient cacheClient(logWriter, storageServer, cacheNetworkClient, sessionServer); if (false) { cacheServer.Load(); cacheNetworkServer.StartListen(nbm); cacheNetworkClient.Connect(nbm, TC("127.0.0.1")); } auto g = MakeGuard([&]() { cacheNetworkClient.Disconnect(); cacheNetworkServer.DisconnectClients(); }); */ auto g = MakeGuard([&]() { if (cacheClient) cacheClient->GetClient().Disconnect(); }); CacheClient* cacheClients[] = { cacheClient }; SchedulerCreateInfo info(sessionServer); info.forceRemote = isRemote; info.forceNative = commandType == CommandType_Native; info.maxLocalProcessors = maxProcessCount; info.cacheClients = cacheClients; info.cacheClientCount = cacheClient ? 1 : 0; info.writeToCache = populateCache; Scheduler scheduler(info); if (!scheduler.EnqueueFromFile(yamlFile, [&](EnqueueProcessInfo& epi) { if (rootsHandle) const_cast(epi.info).rootsHandle = rootsHandle; })) return false; u32 queued, activeLocal, activeRemote, outFinished; scheduler.GetStats(queued, activeLocal, activeRemote, outFinished); bool success = true; Atomic counter; static Event finished(true); scheduler.SetProcessFinishedCallback([&](const ProcessHandle& ph) { auto& si = ph.GetStartInfo(); const tchar* desc = si.description; if (ph.GetExitCode() != 0 && ph.GetExitCode() != ProcessCancelExitCode) { logger.Error(TC("%s - Error exit code: %u (%s %s)"), desc, ph.GetExitCode(), si.application, si.arguments); success = false; } u32 c = ++counter; logger.BeginScope(); StringBuffer<128> extra; if (ph.IsRemote()) extra.Append(TCV(" [RemoteExecutor: ")).Append(ph.GetExecutingHost()).Append(']'); else if (ph.GetExecutionType() == ProcessExecutionType_Native) extra.Append(TCV(" (Not detoured)")); else if (ph.GetExecutionType() == ProcessExecutionType_FromCache) extra.Append(TCV(" (From cache)")); logger.Info(TC("[%u/%u] %s%s"), c, queued, desc, extra.data); for (auto& line : ph.GetLogLines()) if (line.text != desc && !StartsWith(line.text.c_str(), TC(" Creating library"))) logger.Log(line.type, line.text.c_str(), u32(line.text.size())); logger.EndScope(); if (c == queued) finished.Set(); }); auto RunQueue = [&]() { logger.Info(TC("Running Scheduler with %u processes"), queued); u64 start = GetTime(); scheduler.Start(); if (!finished.IsSet()) return false; u64 time = GetTime() - start; logger.Info(TC("Scheduler run took %s"), TimeToText(time).str); logger.Info(TC("")); stopServer.Execute(); return success; }; if (commandType == CommandType_Agent) { u32 clientCount = maxProcessCount == 1 ? 1 : agentCount; return RunWithClient([&]() { return RunQueue(); }, clientCount); } else return RunQueue(); }; if (!coordinatorName.empty()) { StringBuffer<512> coordinatorWorkDir(g_rootDir); coordinatorWorkDir.EnsureEndsWithSlash().Append(coordinatorName); StringBuffer<512> binariesDir; if (!GetDirectoryOfCurrentModule(logger, binariesDir)) return false; CoordinatorCreateInfo cinfo; cinfo.workDir = coordinatorWorkDir.data; cinfo.binariesDir = binariesDir.data; // TODO: This is very horde specific.. maybe all these parameters should be a string or something cinfo.pool = coordinatorPool.c_str(); cinfo.maxCoreCount = coordinatorMaxCoreCount; cinfo.logging = true; if (!coordinator.Create(logger, coordinatorName.c_str(), cinfo, networkBackend, networkServer)) return false; } auto cg = MakeGuard([&]() { coordinator.Destroy(); }); #if PLATFORM_WINDOWS // Annoying that link.exe/lld-link.exe needs path to windows folder.. if (!useScheduler) { StringBuffer<512> sdkbin; //sdkbin.count = GetEnvironmentVariable(TC("WindowsSdkVerBinPath"), sdkbin.data, sdkbin.capacity); sdkbin.Append(TCV(";C:\\Program Files (x86)\\Windows Kits\\10\\bin\\10.0.22621.0\\x64")); if (sdkbin.count) { StringBuffer<4096> temp; temp.count = GetEnvironmentVariableW(TC("PATH"), temp.data, temp.capacity); temp.Append(sdkbin); SetEnvironmentVariableW(TC("PATH"), temp.data); } } #endif for (u32 i=0; i!=loopCount; ++i) { bool success = false; if (useScheduler) { success = RunScheduler(application.c_str()); } else { switch (commandType) { case CommandType_Native: success = RunLocal(application, arguments, false); break; case CommandType_Local: success = RunLocal(application, arguments, true); break; case CommandType_Remote: success = RunRemote(application, arguments); break; case CommandType_Agent: success = RunAgent(application, arguments); } } if (!success) return false; if (false) networkServer.DisconnectClients(); } logger.BeginScope(); if (printSummary) { sessionServer.PrintSummary(logger); storageServer.PrintSummary(logger); networkServer.PrintSummary(logger); KernelStats::GetGlobal().Print(logger, true); PrintContentionSummary(logger); } logger.EndScope(); return true; } } #if PLATFORM_WINDOWS int wmain(int argc, wchar_t* argv[]) { int res; //while (true) { res = uba::WrappedMain(argc, argv) ? 0 : -1; } Sleep(1); // Here to be able to put a breakpoint just before exit :-) return res; } #else int main(int argc, char* argv[]) { return uba::WrappedMain(argc, argv) ? 0 : -1; } #endif