Files
UnrealEngine/Engine/Source/Programs/Unsync/Private/UnsyncCmdQuery.cpp
2025-05-18 13:04:45 +08:00

665 lines
18 KiB
C++

// Copyright Epic Games, Inc. All Rights Reserved.
#include "UnsyncCmdQuery.h"
#include "UnsyncFile.h"
#include "UnsyncHttp.h"
#include "UnsyncProxy.h"
#include "UnsyncThread.h"
#include "UnsyncUtil.h"
#include "UnsyncAuth.h"
#include "UnsyncPool.h"
#include "UnsyncScheduler.h"
#include <float.h>
#include <algorithm>
#include <regex>
#include <json11.hpp>
#include <fmt/format.h>
namespace unsync {
using FMirrorInfoResult = TResult<std::vector<FMirrorInfo>>;
// Runs a basic HTTP request against the remote server and returns the time it took to get the response, -1 if connection could not be
// established.
static double
RunHttpPing(std::string_view Address, uint16 Port)
{
FTlsClientSettings TlsSettings;
TlsSettings.Subject = Address;
ETlsRequirement TlsRequirement = ETlsRequirement::None;
if (Port == 443)
{
TlsRequirement = ETlsRequirement::Preferred;
}
FHttpConnection Connection(Address, Port, TlsRequirement, TlsSettings);
FHttpRequest Request;
Request.Url = "/api/v1/ping";
Request.Method = EHttpMethod::GET;
// Don't time the connection handshake, only query time itself
if (!Connection.Open())
{
return -1;
}
FTimePoint TimeBegin = TimePointNow();
FHttpResponse PingResponse = HttpRequest(Connection, Request);
FTimePoint TimeEnd = TimePointNow();
if (PingResponse.Success())
{
return DurationSec(TimeBegin, TimeEnd);
}
else
{
return -1;
}
}
FMirrorInfoResult
RunQueryMirrors(const FRemoteDesc& RemoteDesc)
{
const char* Url = "/api/v1/mirrors";
FHttpResponse Response = HttpRequest(RemoteDesc, EHttpMethod::GET, Url);
if (!Response.Success())
{
return HttpError(Url, Response.Code);
}
using namespace json11;
std::string JsonString = std::string(Response.AsStringView());
std::string JsonErrorString;
Json JsonObject = Json::parse(JsonString, JsonErrorString);
if (!JsonErrorString.empty())
{
return AppError(std::string("JSON parse error while getting server mirrors: ") + JsonErrorString);
}
std::vector<FMirrorInfo> Result;
for (const auto& Elem : JsonObject.array_items())
{
FMirrorInfo Info;
for (const auto& Field : Elem.object_items())
{
if (Field.first == "name")
{
Info.Name = Field.second.string_value();
}
else if (Field.first == "description")
{
Info.Description = Field.second.string_value();
}
else if (Field.first == "address")
{
Info.Address = Field.second.string_value();
}
else if (Field.first == "port")
{
int PortValue = Field.second.int_value();
if (PortValue > 0 && PortValue < 65536)
{
Info.Port = uint16(PortValue);
}
else
{
UNSYNC_WARNING(L"Unexpected port value: %d", PortValue);
Info.Port = 0;
}
}
}
Result.push_back(Info);
}
return ResultOk(Result);
}
int32
CmdQueryMirrors(const FCmdQueryOptions& Options)
{
FMirrorInfoResult MirrorsResult = RunQueryMirrors(Options.Remote);
if (MirrorsResult.IsError())
{
LogError(MirrorsResult.GetError(), L"Failed to get mirror list from the server");
return 1;
}
std::vector<FMirrorInfo> Mirrors = MirrorsResult.GetData();
ParallelForEach(Mirrors, [](FMirrorInfo& Mirror) { Mirror.Ping = RunHttpPing(Mirror.Address, Mirror.Port); });
std::sort(Mirrors.begin(), Mirrors.end(), [](const FMirrorInfo& InA, const FMirrorInfo& InB) {
double A = InA.Ping > 0 ? InA.Ping : FLT_MAX;
double B = InB.Ping > 0 ? InB.Ping : FLT_MAX;
return A < B;
});
LogPrintf(ELogLevel::MachineReadable, L"[\n");
for (size_t I = 0; I < Mirrors.size(); ++I)
{
const FMirrorInfo& Mirror = Mirrors[I];
int32 PingMs = (Mirror.Ping == 0) ? 0 : std::max(1, int32(Mirror.Ping * 1000.0));
LogPrintf(ELogLevel::MachineReadable,
L" {\"address\":\"%hs\", \"port\":%d, \"ok\":%hs, \"ping\":%d, \"name\":\"%hs\", \"description\":\"%hs\"}%hs\n",
StringEscape(Mirror.Address).c_str(),
Mirror.Port,
Mirror.Ping > 0 ? "true" : "false",
PingMs,
StringEscape(Mirror.Name).c_str(),
StringEscape(Mirror.Description).c_str(),
I + 1 == Mirrors.size() ? "" : ",");
}
LogPrintf(ELogLevel::MachineReadable, L"]\n");
return 0;
}
int32
CmdQueryList(const FCmdQueryOptions& Options)
{
FHttpConnection Connection = FHttpConnection::CreateDefaultHttps(Options.Remote);
TResult<ProxyQuery::FHelloResponse> HelloResponse = ProxyQuery::Hello(Options.Remote.Protocol, Connection);
if (HelloResponse.IsError())
{
UNSYNC_ERROR("Failed establish a handshake with server '%hs'", Options.Remote.Host.Address.c_str());
LogError(HelloResponse.GetError());
return -1;
}
FAuthDesc AuthDesc = FAuthDesc::FromHelloResponse(*HelloResponse);
if (Options.Args.empty())
{
UNSYNC_ERROR(L"Path argument is required");
return -1;
}
TResult<ProxyQuery::FDirectoryListing> ListingResult = ProxyQuery::ListDirectory(Options.Remote.Protocol, Connection, &AuthDesc, Options.Args[0]);
if (ListingResult.IsError())
{
LogError(ListingResult.GetError(), L"Failed to list remote directory");
return -1;
}
std::string ListingJson = ListingResult->ToJson();
LogPrintf(ELogLevel::MachineReadable, L"%hs\n", ListingJson.c_str());
return 0;
}
int32
CmdQuerySearch(const FCmdQueryOptions& Options)
{
using namespace ProxyQuery;
if (Options.Args.empty())
{
UNSYNC_ERROR(L"Path argument is required");
return -1;
}
auto CreateConnection = [Remote = Options.Remote]
{
FTlsClientSettings TlsSettings = Remote.GetTlsClientSettings();
return new FHttpConnection(Remote.Host.Address, Remote.Host.Port, Remote.TlsRequirement, TlsSettings);
};
TObjectPool<FHttpConnection> ConnectionPool(CreateConnection);
TResult<ProxyQuery::FHelloResponse> HelloResponse = [&ConnectionPool, &Options]
{
std::unique_ptr<FHttpConnection> Connection = ConnectionPool.Acquire();
TResult<ProxyQuery::FHelloResponse> Result = ProxyQuery::Hello(Options.Remote.Protocol, *Connection);
ConnectionPool.Release(std::move(Connection));
return Result;
}();
if (HelloResponse.IsError())
{
UNSYNC_ERROR("Failed establish a handshake with server '%hs'", Options.Remote.Host.Address.c_str());
LogError(HelloResponse.GetError());
return -1;
}
FAuthDesc AuthDesc = FAuthDesc::FromHelloResponse(*HelloResponse);
const std::string& RootPath = Options.Args[0];
std::vector<std::regex> SubdirPatterns;
UNSYNC_LOG("Searching '%hs'", RootPath.c_str());
if (Options.Args.size() > 1)
{
UNSYNC_LOG_INDENT;
UNSYNC_LOG("Subdirectory patterns:", RootPath.c_str());
UNSYNC_LOG_INDENT;
for (size_t i = 1; i < Options.Args.size(); ++i)
{
UNSYNC_LOG("%hs", Options.Args[i].c_str());
std::regex Pattern = std::regex(Options.Args[i], std::regex_constants::icase);
SubdirPatterns.push_back(Pattern);
}
}
struct FEntry
{
std::string Path;
uint32 Depth = 0;
};
std::vector<FEntry> PendingDirectories;
{
FEntry RootEntry;
RootEntry.Path = RootPath;
RootEntry.Depth = 0;
PendingDirectories.push_back(RootEntry);
}
struct FResultEntry : FEntry
{
FDirectoryListingEntry DirEntry;
};
struct FTaskContext
{
std::mutex Mutex;
std::vector<FResultEntry> FoundEntries;
std::unordered_set<std::string> VisitedDirectories;
bool bParentThreadVerbose = false;
int32 ParentThreadIndent = 0;
};
FTaskGroup Tasks = GScheduler->CreateTaskGroup();
FTaskContext Context;
Context.bParentThreadVerbose = GLogVerbose;
Context.ParentThreadIndent = GLogIndent;
std::function<void(std::string, int32, const FDirectoryListing*)> ExploreDirectory =
[&Context, &AuthDesc, &ConnectionPool, &ExploreDirectory, &SubdirPatterns, &Tasks, &Options](
std::string Path,
int32 CurrentDepth,
const FDirectoryListing* DirectoryListingPtr)
{
FLogVerbosityScope VerboseScope(Context.bParentThreadVerbose);
FLogIndentScope IndentScope(Context.ParentThreadIndent, true);
FDirectoryListing RemoteDirectoryListing;
if (!DirectoryListingPtr)
{
UNSYNC_VERBOSE2(L"Listing '%hs'", Path.c_str());
GScheduler->NetworkSemaphore.Acquire(false);
std::unique_ptr<FHttpConnection> Connection = ConnectionPool.Acquire();
TResult<FDirectoryListing> DirectoryListingResult =
ProxyQuery::ListDirectory(Options.Remote.Protocol, *Connection, &AuthDesc, Path);
ConnectionPool.Release(std::move(Connection));
GScheduler->NetworkSemaphore.Release();
if (DirectoryListingResult.IsError())
{
LogError(DirectoryListingResult.GetError(), L"Failed to list remote directory");
return;
}
std::swap(RemoteDirectoryListing, DirectoryListingResult.GetData());
DirectoryListingPtr = &RemoteDirectoryListing;
}
const FDirectoryListing& DirectoryListing = *DirectoryListingPtr;
for (const FDirectoryListingEntry& DirEntry : DirectoryListing.Entries)
{
std::string DirEntryName = DirEntry.Name;
// Only include one subdirectory level
const size_t SeparatorPos = DirEntryName.find(PATH_SEPARATOR);
if (SeparatorPos != std::string::npos)
{
DirEntryName = DirEntryName.substr(0, PATH_SEPARATOR);
}
FEntry NextEntry;
NextEntry.Path = Path + PATH_SEPARATOR + DirEntryName;
NextEntry.Depth = CurrentDepth + 1;
// Include only leaf directory entries in the final output
if (SeparatorPos == std::string::npos)
{
std::lock_guard<std::mutex> LockGuard(Context.Mutex);
FResultEntry ResultEntry;
ResultEntry.Path = NextEntry.Path;
ResultEntry.Depth = NextEntry.Depth;
ResultEntry.DirEntry = DirEntry;
Context.FoundEntries.push_back(ResultEntry);
}
if (NextEntry.Depth > SubdirPatterns.size())
{
continue;
}
if (DirEntry.bDirectory && std::regex_match(DirEntryName, SubdirPatterns[CurrentDepth], std::regex_constants::match_any))
{
UNSYNC_VERBOSE2(L"Matched: '%hs'", DirEntryName.c_str());
// Directory listing may already include some child sub-directories, so we can skip some network requests
FDirectoryListing SubDirectoryListing;
std::string RequiredPrefix = DirEntryName + PATH_SEPARATOR;
for (const FDirectoryListingEntry& DirEntry2 : DirectoryListing.Entries)
{
if (DirEntry2.Name.starts_with(RequiredPrefix))
{
FDirectoryListingEntry SubDirEntry = DirEntry2;
SubDirEntry.Name = SubDirEntry.Name.substr(RequiredPrefix.length());
SubDirectoryListing.Entries.push_back(SubDirEntry);
}
}
{
// Only visit each sub-directory once
std::lock_guard<std::mutex> LockGuard(Context.Mutex);
if (!Context.VisitedDirectories.insert(NextEntry.Path).second)
{
continue;
}
}
if (SubDirectoryListing.Entries.empty())
{
Tasks.run([ExploreDirectory, NextEntry]() { ExploreDirectory(NextEntry.Path, NextEntry.Depth, nullptr); });
}
else
{
ExploreDirectory(NextEntry.Path, NextEntry.Depth, &SubDirectoryListing);
}
}
}
};
ExploreDirectory(RootPath, 0, nullptr);
Tasks.wait();
std::vector<FResultEntry>& ResultEntries = Context.FoundEntries;
std::sort(ResultEntries.begin(), ResultEntries.end(), [](const FResultEntry& A, const FResultEntry& B) { return A.Path < B.Path; });
LogPrintf(ELogLevel::MachineReadable, L"{\n");
LogPrintf(ELogLevel::MachineReadable, L" \"root\": \"%hs\",\n", StringEscape(RootPath).c_str());
LogPrintf(ELogLevel::MachineReadable, L" \"entries\": [\n");
for (size_t i = 0; i < ResultEntries.size(); ++i)
{
const FResultEntry& ResultEntry = ResultEntries[i];
const char* TrailingComma = i + 1 == ResultEntries.size() ? "" : ",";
std::string_view RelativePath = std::string_view(ResultEntry.Path).substr(RootPath.length() + 1);
LogPrintf(ELogLevel::MachineReadable,
L" { \"path\": \"%hs\", \"is_directory\": %hs, \"mtime\": %llu, \"size\": %llu }%hs\n",
StringEscape(RelativePath).c_str(),
ResultEntry.DirEntry.bDirectory ? "true" : "false",
llu(ResultEntry.DirEntry.Mtime),
llu(ResultEntry.DirEntry.Size),
TrailingComma);
}
LogPrintf(ELogLevel::MachineReadable, L" ]\n");
LogPrintf(ELogLevel::MachineReadable, L"}\n");
return 0;
}
int32
CmdQueryFile(const FCmdQueryOptions& Options)
{
if (Options.Args.empty())
{
UNSYNC_ERROR(L"Path argument is required");
return -1;
}
// TODO: use a global connection pool and use it for ProxyQuery::DownloadFile too
FHttpConnection HelloConnection = FHttpConnection::CreateDefaultHttps(Options.Remote);
TResult<ProxyQuery::FHelloResponse> HelloResponse = ProxyQuery::Hello(Options.Remote.Protocol, HelloConnection);
if (HelloResponse.IsError())
{
UNSYNC_ERROR("Failed establish a handshake with server '%hs'", Options.Remote.Host.Address.c_str());
LogError(HelloResponse.GetError());
return -1;
}
FAuthDesc AuthDesc = FAuthDesc::FromHelloResponse(*HelloResponse);
HelloConnection.Close();
TResult<FAuthToken> AuthToken = Authenticate(AuthDesc);
if (!AuthToken.IsOk())
{
LogError(AuthToken.GetError(), L"Failed to authenticate");
return -1;
}
UNSYNC_LOG(L"Downloading file: '%hs'", Options.Args[0].c_str());
UNSYNC_LOG_INDENT;
std::unique_ptr<FNativeFile> ResultWriter;
FPath OutputPath = Options.OutputPath;
if (OutputPath.empty())
{
FPath RequestPath = Options.Args[0];
if (RequestPath.has_filename())
{
OutputPath = RequestPath.filename();
OutputPath = GetAbsoluteNormalPath(OutputPath);
}
else
{
UNSYNC_ERROR(
L"Output could not be derived from the request string. "
L"Use `-o <filename>` command line argument to specify it explicitly.");
return 1;
}
}
UNSYNC_LOG(L"Output file: '%ls'", OutputPath.wstring().c_str());
auto OutputCallback = [&ResultWriter, &Options, OutputPath](uint64 Size) -> FIOWriter&
{
UNSYNC_LOG(L"Size: %llu bytes (%.3f MB)", llu(Size), SizeMb(Size));
ResultWriter = std::make_unique<FNativeFile>(OutputPath, EFileMode::CreateWriteOnly, Size);
return *ResultWriter;
};
FHttpConnection Connection = FHttpConnection::CreateDefaultHttps(Options.Remote);
TResult<> Response = ProxyQuery::DownloadFile(Connection, &AuthDesc, Options.Args[0], OutputCallback);
if (Response.IsOk())
{
UNSYNC_LOG(L"Output written to file '%ls'", OutputPath.wstring().c_str());
}
else
{
LogError(Response.GetError(), L"Failed to download file");
}
return 0;
}
int32
CmdQueryHttpGet(const FCmdQueryOptions& Options)
{
FHttpConnection HttpConnection = FHttpConnection::CreateDefaultHttps(Options.Remote);
std::string BearerToken;
if (Options.Remote.bAuthenticationRequired)
{
TResult<ProxyQuery::FHelloResponse> HelloResponse = ProxyQuery::Hello(Options.Remote.Protocol, HttpConnection);
if (!HelloResponse.IsOk())
{
LogError(HelloResponse.GetError(), L"Failed to query basic server information");
return -1;
}
FAuthDesc AuthDesc = FAuthDesc::FromHelloResponse(*HelloResponse);
TResult<FAuthToken> AuthToken = Authenticate(AuthDesc);
if (!AuthToken.IsOk())
{
LogError(AuthToken.GetError(), L"Failed to authenticate with the server");
return -1;
}
BearerToken = AuthToken->Access;
}
// TODO: RequestPath should include leading slash
std::string RequestUrl = fmt::format("/{}", Options.Remote.RequestPath);
FHttpRequest Request;
Request.Method = EHttpMethod::GET;
Request.BearerToken = BearerToken;
Request.Url = RequestUrl;
FHttpResponse Response = HttpRequest(HttpConnection, Request);
if (!Response.Success())
{
LogError(HttpError(std::move(RequestUrl), Response.Code));
return -1;
}
FPath OutputPath = Options.OutputPath;
if (OutputPath.empty())
{
if (Response.Buffer.Empty())
{
return 0;
}
else if (Response.ContentType == EHttpContentType::Application_Json || Response.ContentType == EHttpContentType::Text_Plain ||
Response.ContentType == EHttpContentType::Text_Html)
{
Response.Buffer.PushBack('\n');
Response.Buffer.PushBack(0);
LogPrintf(ELogLevel::MachineReadable, L"%hs", (const char*)Response.Buffer.Data());
return 0;
}
else
{
UNSYNC_ERROR(L"Unexpected response content type. Only plain text or json are supported. Use `-o <filename>` command line argument to write response body to a file.");
return -1;
}
}
else
{
UNSYNC_LOG(L"Output file: '%ls'", OutputPath.wstring().c_str());
OutputPath = GetAbsoluteNormalPath(OutputPath);
if (EnsureDirectoryExists(OutputPath.parent_path()))
{
if (WriteBufferToFile(OutputPath, Response.Buffer))
{
UNSYNC_VERBOSE(L"Wrote bytes: %llu", llu(Response.Buffer.Size()));
return 0;
}
else
{
UNSYNC_ERROR(L"Failed to write output file '%ls'", OutputPath.wstring().c_str());
}
}
else
{
UNSYNC_ERROR(L"Failed to create output directory '%ls'", OutputPath.parent_path().wstring().c_str());
}
return -1;
}
}
int32
CmdQuery(const FCmdQueryOptions& Options)
{
if (!Options.Remote.IsValid())
{
UNSYNC_ERROR(L"Server address is not specified or is invalid");
return 1;
}
if (Options.Query == "mirrors")
{
return CmdQueryMirrors(Options);
}
else if (Options.Query == "list")
{
return CmdQueryList(Options);
}
else if (Options.Query == "search" || Options.Query == "explore")
{
return CmdQuerySearch(Options);
}
else if (Options.Query == "file")
{
return CmdQueryFile(Options);
}
if (Options.Query == "http-get")
{
return CmdQueryHttpGet(Options);
}
else
{
UNSYNC_ERROR(L"Unknown query command. Allowed options: mirrors, list, search, file, http-get");
return 1;
}
}
TResult<FMirrorInfo>
FindClosestMirror(const FRemoteDesc& Remote)
{
FMirrorInfoResult MirrorsResult = RunQueryMirrors(Remote);
if (MirrorsResult.IsError())
{
return FError(MirrorsResult.GetError());
}
std::vector<FMirrorInfo> Mirrors = MirrorsResult.GetData();
ParallelForEach(Mirrors, [](FMirrorInfo& Mirror) { Mirror.Ping = RunHttpPing(Mirror.Address, Mirror.Port); });
std::sort(Mirrors.begin(), Mirrors.end(), [](const FMirrorInfo& InA, const FMirrorInfo& InB) {
double A = InA.Ping > 0 ? InA.Ping : FLT_MAX;
double B = InB.Ping > 0 ? InB.Ping : FLT_MAX;
return A < B;
});
for (const FMirrorInfo& Mirror : Mirrors)
{
if (Mirror.Ping > 0)
{
return ResultOk(Mirror);
}
}
return AppError("No reachable mirror found");
}
} // namespace unsync