Merge branch 'main' into lm/restrict-content-type

1 files changed, 210 insertions, 69 deletions

diff --git a/src/zenserver/storage/projectstore/httpprojectstore.cpp b/src/zenserver/storage/projectstore/httpprojectstore.cpp
index fe32fa15b..836d84292 100644
--- a/src/zenserver/storage/projectstore/httpprojectstore.cpp
+++ b/src/zenserver/storage/projectstore/httpprojectstore.cpp
@@ -13,7 +13,12 @@
 #include <zencore/scopeguard.h>
 #include <zencore/stream.h>
 #include <zencore/trace.h>
+#include <zenhttp/httpclientauth.h>
 #include <zenhttp/packageformat.h>
+#include <zenremotestore/builds/buildstoragecache.h>
+#include <zenremotestore/builds/buildstorageutil.h>
+#include <zenremotestore/jupiter/jupiterhost.h>
+#include <zenremotestore/operationlogoutput.h>
 #include <zenremotestore/projectstore/buildsremoteprojectstore.h>
 #include <zenremotestore/projectstore/fileremoteprojectstore.h>
 #include <zenremotestore/projectstore/jupiterremoteprojectstore.h>
@@ -244,6 +249,22 @@ namespace {
 	{
 		std::shared_ptr<RemoteProjectStore> Store;
 		std::string							Description;
+		double								LatencySec				= -1.0;
+		uint64_t							MaxRangeCountPerRequest = 1;
+
+		struct Cache
+		{
+			std::unique_ptr<HttpClient>		   Http;
+			std::unique_ptr<BuildStorageCache> Cache;
+			Oid								   BuildsId = Oid::Zero;
+			std::string						   Description;
+			double							   LatencySec			   = -1.0;
+			uint64_t						   MaxRangeCountPerRequest = 1;
+			BuildStorageCache::Statistics	   Stats;
+			bool							   Populate = false;
+		};
+
+		std::unique_ptr<Cache> OptionalCache;
 	};
 
 	CreateRemoteStoreResult CreateRemoteStore(LoggerRef					   InLog,
@@ -260,7 +281,7 @@ namespace {
 
 		using namespace std::literals;
 
-		std::shared_ptr<RemoteProjectStore> RemoteStore;
+		CreateRemoteStoreResult Result;
 
 		if (CbObjectView File = Params["file"sv].AsObjectView(); File)
 		{
@@ -285,7 +306,9 @@ namespace {
 				std::string(OptionalBaseName),
 				ForceDisableBlocks,
 				ForceEnableTempBlocks};
-			RemoteStore = CreateFileRemoteStore(Log(), Options);
+			Result.Store				   = CreateFileRemoteStore(Log(), Options);
+			Result.LatencySec			   = 0.5 / 1000.0;	// 0.5 ms
+			Result.MaxRangeCountPerRequest = 1024u;
 		}
 
 		if (CbObjectView Cloud = Params["cloud"sv].AsObjectView(); Cloud)
@@ -363,21 +386,32 @@ namespace {
 			bool ForceDisableTempBlocks = Cloud["disabletempblocks"sv].AsBool(false);
 			bool AssumeHttp2			= Cloud["assumehttp2"sv].AsBool(false);
 
-			JupiterRemoteStoreOptions Options = {
-				RemoteStoreOptions{.MaxBlockSize = MaxBlockSize, .MaxChunksPerBlock = 1000, .MaxChunkEmbedSize = MaxChunkEmbedSize},
-				Url,
-				std::string(Namespace),
-				std::string(Bucket),
-				Key,
-				BaseKey,
-				std::string(OpenIdProvider),
-				AccessToken,
-				AuthManager,
-				OidcExePath,
-				ForceDisableBlocks,
-				ForceDisableTempBlocks,
-				AssumeHttp2};
-			RemoteStore = CreateJupiterRemoteStore(Log(), Options, TempFilePath, /*Quiet*/ false, /*Unattended*/ false, /*Hidden*/ true);
+			if (JupiterEndpointTestResult TestResult = TestJupiterEndpoint(Url, AssumeHttp2, /*Verbose*/ false); TestResult.Success)
+			{
+				Result.LatencySec			   = TestResult.LatencySeconds;
+				Result.MaxRangeCountPerRequest = TestResult.MaxRangeCountPerRequest;
+
+				JupiterRemoteStoreOptions Options = {
+					RemoteStoreOptions{.MaxBlockSize = MaxBlockSize, .MaxChunksPerBlock = 1000, .MaxChunkEmbedSize = MaxChunkEmbedSize},
+					Url,
+					std::string(Namespace),
+					std::string(Bucket),
+					Key,
+					BaseKey,
+					std::string(OpenIdProvider),
+					AccessToken,
+					AuthManager,
+					OidcExePath,
+					ForceDisableBlocks,
+					ForceDisableTempBlocks,
+					AssumeHttp2};
+				Result.Store =
+					CreateJupiterRemoteStore(Log(), Options, TempFilePath, /*Quiet*/ false, /*Unattended*/ false, /*Hidden*/ true);
+			}
+			else
+			{
+				return {nullptr, fmt::format("Unable to connect to jupiter host '{}'", Url)};
+			}
 		}
 
 		if (CbObjectView Zen = Params["zen"sv].AsObjectView(); Zen)
@@ -393,12 +427,13 @@ namespace {
 			{
 				return {nullptr, "Missing oplog"};
 			}
+
 			ZenRemoteStoreOptions Options = {
 				RemoteStoreOptions{.MaxBlockSize = MaxBlockSize, .MaxChunksPerBlock = 1000, .MaxChunkEmbedSize = MaxChunkEmbedSize},
 				std::string(Url),
 				std::string(Project),
 				std::string(Oplog)};
-			RemoteStore = CreateZenRemoteStore(Log(), Options, TempFilePath);
+			Result.Store = CreateZenRemoteStore(Log(), Options, TempFilePath);
 		}
 
 		if (CbObjectView Builds = Params["builds"sv].AsObjectView(); Builds)
@@ -471,11 +506,76 @@ namespace {
 			MemoryView MetaDataSection = Builds["metadata"sv].AsBinaryView();
 			IoBuffer   MetaData(IoBuffer::Wrap, MetaDataSection.GetData(), MetaDataSection.GetSize());
 
+			auto EnsureHttps = [](const std::string& Host, std::string_view PreferredProtocol) {
+				if (!Host.empty() && Host.find("://"sv) == std::string::npos)
+				{
+					// Assume https URL
+					return fmt::format("{}://{}"sv, PreferredProtocol, Host);
+				}
+				return Host;
+			};
+
+			Host		 = EnsureHttps(Host, "https");
+			OverrideHost = EnsureHttps(OverrideHost, "https");
+			ZenHost		 = EnsureHttps(ZenHost, "http");
+
+			std::function<HttpClientAccessToken()> TokenProvider;
+			if (!OpenIdProvider.empty())
+			{
+				TokenProvider = httpclientauth::CreateFromOpenIdProvider(AuthManager, OpenIdProvider);
+			}
+			else if (!AccessToken.empty())
+			{
+				TokenProvider = httpclientauth::CreateFromStaticToken(AccessToken);
+			}
+			else if (!OidcExePath.empty())
+			{
+				if (auto TokenProviderMaybe = httpclientauth::CreateFromOidcTokenExecutable(OidcExePath,
+																							Host.empty() ? OverrideHost : Host,
+																							/*Quiet*/ false,
+																							/*Unattended*/ false,
+																							/*Hidden*/ true);
+					TokenProviderMaybe)
+				{
+					TokenProvider = TokenProviderMaybe.value();
+				}
+			}
+
+			if (!TokenProvider)
+			{
+				TokenProvider = httpclientauth::CreateFromDefaultOpenIdProvider(AuthManager);
+			}
+
+			BuildStorageResolveResult ResolveResult;
+			{
+				HttpClientSettings ClientSettings{.LogCategory		   = "httpbuildsclient",
+												  .AccessTokenProvider = TokenProvider,
+												  .AssumeHttp2		   = AssumeHttp2,
+												  .AllowResume		   = true,
+												  .RetryCount		   = 2};
+
+				std::unique_ptr<OperationLogOutput> Output(CreateStandardLogOutput(Log()));
+
+				try
+				{
+					ResolveResult = ResolveBuildStorage(*Output,
+														ClientSettings,
+														Host,
+														OverrideHost,
+														ZenHost,
+														ZenCacheResolveMode::Discovery,
+														/*Verbose*/ false);
+				}
+				catch (const std::exception& Ex)
+				{
+					return {nullptr, fmt::format("Failed resolving storage host and cache. Reason: '{}'", Ex.what())};
+				}
+			}
+			Result.LatencySec			   = ResolveResult.Cloud.LatencySec;
+			Result.MaxRangeCountPerRequest = ResolveResult.Cloud.Caps.MaxRangeCountPerRequest;
+
 			BuildsRemoteStoreOptions Options = {
 				RemoteStoreOptions{.MaxBlockSize = MaxBlockSize, .MaxChunksPerBlock = 1000, .MaxChunkEmbedSize = MaxChunkEmbedSize},
-				Host,
-				OverrideHost,
-				ZenHost,
 				std::string(Namespace),
 				std::string(Bucket),
 				BuildId,
@@ -485,25 +585,43 @@ namespace {
 				OidcExePath,
 				ForceDisableBlocks,
 				ForceDisableTempBlocks,
-				AssumeHttp2,
-				PopulateCache,
 				MetaData,
 				MaximumInMemoryDownloadSize};
-			RemoteStore = CreateJupiterBuildsRemoteStore(Log(),
-														 Options,
-														 TempFilePath,
-														 /*Quiet*/ false,
-														 /*Unattended*/ false,
-														 /*Hidden*/ true,
-														 GetTinyWorkerPool(EWorkloadType::Background));
+			Result.Store = CreateJupiterBuildsRemoteStore(Log(), ResolveResult, std::move(TokenProvider), Options, TempFilePath);
+
+			if (!ResolveResult.Cache.Address.empty())
+			{
+				Result.OptionalCache = std::make_unique<CreateRemoteStoreResult::Cache>();
+
+				HttpClientSettings CacheClientSettings{.LogCategory					= "httpcacheclient",
+													   .ConnectTimeout				= std::chrono::milliseconds{3000},
+													   .Timeout						= std::chrono::milliseconds{30000},
+													   .AssumeHttp2					= ResolveResult.Cache.AssumeHttp2,
+													   .AllowResume					= true,
+													   .RetryCount					= 0,
+													   .MaximumInMemoryDownloadSize = MaximumInMemoryDownloadSize};
+
+				Result.OptionalCache->Http		 = std::make_unique<HttpClient>(ResolveResult.Cache.Address, CacheClientSettings);
+				Result.OptionalCache->Cache		 = CreateZenBuildStorageCache(*Result.OptionalCache->Http,
+																		  Result.OptionalCache->Stats,
+																		  Namespace,
+																		  Bucket,
+																		  TempFilePath,
+																		  GetTinyWorkerPool(EWorkloadType::Background));
+				Result.OptionalCache->BuildsId	 = BuildId;
+				Result.OptionalCache->LatencySec = ResolveResult.Cache.LatencySec;
+				Result.OptionalCache->MaxRangeCountPerRequest = ResolveResult.Cache.Caps.MaxRangeCountPerRequest;
+				Result.OptionalCache->Populate				  = PopulateCache;
+				Result.OptionalCache->Description =
+					fmt::format("[zenserver] {} namespace {} bucket {}", ResolveResult.Cache.Address, Namespace, Bucket);
+			}
 		}
-
-		if (!RemoteStore)
+		if (!Result.Store)
 		{
 			return {nullptr, "Unknown remote store type"};
 		}
 
-		return {std::move(RemoteStore), ""};
+		return Result;
 	}
 
 	std::pair<HttpResponseCode, std::string> ConvertResult(const RemoteProjectStore::Result& Result)
@@ -714,8 +832,8 @@ HttpProjectService::HandleRequest(HttpServerRequest& Request)
 	}
 }
 
-void
-HttpProjectService::HandleStatsRequest(HttpServerRequest& HttpReq)
+CbObject
+HttpProjectService::CollectStats()
 {
 	ZEN_TRACE_CPU("ProjectService::Stats");
 
@@ -781,7 +899,13 @@ HttpProjectService::HandleStatsRequest(HttpServerRequest& HttpReq)
 	}
 	Cbo.EndObject();
 
-	return HttpReq.WriteResponse(HttpResponseCode::OK, Cbo.Save());
+	return Cbo.Save();
+}
+
+void
+HttpProjectService::HandleStatsRequest(HttpServerRequest& HttpReq)
+{
+	HttpReq.WriteResponse(HttpResponseCode::OK, CollectStats());
 }
 
 void
@@ -2373,15 +2497,19 @@ HttpProjectService::HandleOplogSaveRequest(HttpRouterRequest& Req)
 		tsl::robin_set<IoHash, IoHash::Hasher> Attachments;
 
 		auto HasAttachment = [this](const IoHash& RawHash) { return m_CidStore.ContainsChunk(RawHash); };
-		auto OnNeedBlock   = [&AttachmentsLock, &Attachments](const IoHash& BlockHash, const std::vector<IoHash>&& ChunkHashes) {
+		auto OnNeedBlock   = [&AttachmentsLock, &Attachments](ThinChunkBlockDescription&& ThinBlockDescription,
+															  std::vector<uint32_t>&&	  NeededChunkIndexes) {
 			  RwLock::ExclusiveLockScope _(AttachmentsLock);
-			  if (BlockHash != IoHash::Zero)
+			  if (ThinBlockDescription.BlockHash != IoHash::Zero)
 			  {
-				  Attachments.insert(BlockHash);
+				  Attachments.insert(ThinBlockDescription.BlockHash);
 			  }
 			  else
 			  {
-				  Attachments.insert(ChunkHashes.begin(), ChunkHashes.end());
+				  for (uint32_t ChunkIndex : NeededChunkIndexes)
+				  {
+					  Attachments.insert(ThinBlockDescription.ChunkRawHashes[ChunkIndex]);
+				  }
 			  }
 		};
 		auto OnNeedAttachment = [&AttachmentsLock, &Attachments](const IoHash& RawHash) {
@@ -2687,36 +2815,39 @@ HttpProjectService::HandleRpcRequest(HttpRouterRequest& Req)
 				bool		 CleanOplog				  = Params["clean"].AsBool(false);
 				bool		 BoostWorkerCount		  = Params["boostworkercount"].AsBool(false);
 				bool		 BoostWorkerMemory		  = Params["boostworkermemory"sv].AsBool(false);
-
-				CreateRemoteStoreResult RemoteStoreResult = CreateRemoteStore(Log(),
-																			  Params,
-																			  m_AuthMgr,
-																			  MaxBlockSize,
-																			  MaxChunkEmbedSize,
-																			  GetMaxMemoryBufferSize(MaxBlockSize, BoostWorkerMemory),
-																			  Oplog->TempPath());
-
-				if (RemoteStoreResult.Store == nullptr)
+				EPartialBlockRequestMode PartialBlockRequestMode =
+					PartialBlockRequestModeFromString(Params["partialblockrequestmode"sv].AsString("true"));
+
+				std::shared_ptr<CreateRemoteStoreResult> RemoteStoreResult =
+					std::make_shared<CreateRemoteStoreResult>(CreateRemoteStore(Log(),
+																				Params,
+																				m_AuthMgr,
+																				MaxBlockSize,
+																				MaxChunkEmbedSize,
+																				GetMaxMemoryBufferSize(MaxBlockSize, BoostWorkerMemory),
+																				Oplog->TempPath()));
+
+				if (RemoteStoreResult->Store == nullptr)
 				{
-					return HttpReq.WriteResponse(HttpResponseCode::BadRequest, HttpContentType::kText, RemoteStoreResult.Description);
+					return HttpReq.WriteResponse(HttpResponseCode::BadRequest, HttpContentType::kText, RemoteStoreResult->Description);
 				}
-				std::shared_ptr<RemoteProjectStore> RemoteStore = std::move(RemoteStoreResult.Store);
-				RemoteProjectStore::RemoteStoreInfo StoreInfo	= RemoteStore->GetInfo();
 
 				JobId JobId = m_JobQueue.QueueJob(
 					fmt::format("Import oplog '{}/{}'", Project->Identifier, Oplog->OplogId()),
 					[this,
-					 ChunkStore		   = &m_CidStore,
-					 ActualRemoteStore = std::move(RemoteStore),
+					 RemoteStoreResult = std::move(RemoteStoreResult),
 					 Oplog,
 					 Force,
 					 IgnoreMissingAttachments,
 					 CleanOplog,
+					 PartialBlockRequestMode,
 					 BoostWorkerCount](JobContext& Context) {
-						Context.ReportMessage(fmt::format("Loading oplog '{}/{}' from {}",
-														  Oplog->GetOuterProjectIdentifier(),
-														  Oplog->OplogId(),
-														  ActualRemoteStore->GetInfo().Description));
+						Context.ReportMessage(
+							fmt::format("Loading oplog '{}/{}'\n    Host: {}\n    Cache: {}",
+										Oplog->GetOuterProjectIdentifier(),
+										Oplog->OplogId(),
+										RemoteStoreResult->Store->GetInfo().Description,
+										RemoteStoreResult->OptionalCache ? RemoteStoreResult->OptionalCache->Description : "<none>"));
 
 						Ref<TransferThreadWorkers> Workers = GetThreadWorkers(BoostWorkerCount, /*SingleThreaded*/ false);
 
@@ -2724,16 +2855,26 @@ HttpProjectService::HandleRpcRequest(HttpRouterRequest& Req)
 						WorkerThreadPool& NetworkWorkerPool = Workers->GetNetworkPool();
 
 						Context.ReportMessage(fmt::format("{}", Workers->GetWorkersInfo()));
-
-						RemoteProjectStore::Result Result	= LoadOplog(m_CidStore,
-																		*ActualRemoteStore,
-																		*Oplog,
-																		NetworkWorkerPool,
-																		WorkerPool,
-																		Force,
-																		IgnoreMissingAttachments,
-																		CleanOplog,
-																		&Context);
+						RemoteProjectStore::Result Result	= LoadOplog(LoadOplogContext{
+							  .ChunkStore	 = m_CidStore,
+							  .RemoteStore	 = *RemoteStoreResult->Store,
+							  .OptionalCache = RemoteStoreResult->OptionalCache ? RemoteStoreResult->OptionalCache->Cache.get() : nullptr,
+							  .CacheBuildId	 = RemoteStoreResult->OptionalCache ? RemoteStoreResult->OptionalCache->BuildsId : Oid::Zero,
+							  .OptionalCacheStats = RemoteStoreResult->OptionalCache ? &RemoteStoreResult->OptionalCache->Stats : nullptr,
+							  .Oplog			  = *Oplog,
+							  .NetworkWorkerPool  = NetworkWorkerPool,
+							  .WorkerPool		  = WorkerPool,
+							  .ForceDownload	  = Force,
+							  .IgnoreMissingAttachments = IgnoreMissingAttachments,
+							  .CleanOplog				= CleanOplog,
+							  .PartialBlockRequestMode	= PartialBlockRequestMode,
+							  .PopulateCache   = RemoteStoreResult->OptionalCache ? RemoteStoreResult->OptionalCache->Populate : false,
+							  .StoreLatencySec = RemoteStoreResult->LatencySec,
+							  .StoreMaxRangeCountPerRequest = RemoteStoreResult->MaxRangeCountPerRequest,
+							  .CacheLatencySec = RemoteStoreResult->OptionalCache ? RemoteStoreResult->OptionalCache->LatencySec : -1.0,
+							  .CacheMaxRangeCountPerRequest =
+								  RemoteStoreResult->OptionalCache ? RemoteStoreResult->OptionalCache->MaxRangeCountPerRequest : 0,
+							  .OptionalJobContext = &Context});
 						auto					   Response = ConvertResult(Result);
 						ZEN_INFO("LoadOplog: Status: {} '{}'", ToString(Response.first), Response.second);
 						if (!IsHttpSuccessCode(Response.first))