builds allow redirect option (#379)

* add --allow-redirect to zen builds upload/download

2 files changed, 57 insertions, 22 deletions

diff --git a/src/zenutil/jupiter/jupiterbuildstorage.cpp b/src/zenutil/jupiter/jupiterbuildstorage.cpp
index 6ada72e1e..53cad78da 100644
--- a/src/zenutil/jupiter/jupiterbuildstorage.cpp
+++ b/src/zenutil/jupiter/jupiterbuildstorage.cpp
@@ -25,8 +25,9 @@ public:
 						Statistics&					 Stats,
 						std::string_view			 Namespace,
 						std::string_view			 Bucket,
+						bool						 AllowRedirect,
 						const std::filesystem::path& TempFolderPath)
-	: m_Session(InLog, InHttpClient)
+	: m_Session(InLog, InHttpClient, AllowRedirect)
 	, m_Stats(Stats)
 	, m_Namespace(Namespace)
 	, m_Bucket(Bucket)
@@ -462,11 +463,12 @@ CreateJupiterBuildStorage(LoggerRef					   InLog,
 						  BuildStorage::Statistics&	   Stats,
 						  std::string_view			   Namespace,
 						  std::string_view			   Bucket,
+						  bool						   AllowRedirect,
 						  const std::filesystem::path& TempFolderPath)
 {
 	ZEN_TRACE_CPU("CreateJupiterBuildStorage");
 
-	return std::make_unique<JupiterBuildStorage>(InLog, InHttpClient, Stats, Namespace, Bucket, TempFolderPath);
+	return std::make_unique<JupiterBuildStorage>(InLog, InHttpClient, Stats, Namespace, Bucket, AllowRedirect, TempFolderPath);
 }
 
 }  // namespace zen
diff --git a/src/zenutil/jupiter/jupitersession.cpp b/src/zenutil/jupiter/jupitersession.cpp
index d3076d36b..de138f994 100644
--- a/src/zenutil/jupiter/jupitersession.cpp
+++ b/src/zenutil/jupiter/jupitersession.cpp
@@ -5,6 +5,7 @@
 #include <zencore/compactbinary.h>
 #include <zencore/compactbinarybuilder.h>
 #include <zencore/compositebuffer.h>
+#include <zencore/compress.h>
 #include <zencore/fmtutils.h>
 #include <zencore/trace.h>
 
@@ -48,7 +49,10 @@ namespace detail {
 	}
 }  // namespace detail
 
-JupiterSession::JupiterSession(LoggerRef InLog, HttpClient& InHttpClient) : m_Log(InLog), m_HttpClient(InHttpClient)
+JupiterSession::JupiterSession(LoggerRef InLog, HttpClient& InHttpClient, bool AllowRedirect)
+: m_Log(InLog)
+, m_HttpClient(InHttpClient)
+, m_AllowRedirect(AllowRedirect)
 {
 }
 
@@ -542,12 +546,14 @@ JupiterSession::PutMultipartBuildBlob(std::string_view												  Namespace,
 									  bool& OutIsComplete) -> JupiterResult {
 			const MultipartUploadResponse::Part& Part		 = Workload->PartDescription.Parts[PartIndex];
 			IoBuffer							 PartPayload = Workload->Transmitter(Part.FirstByte, Part.LastByte - Part.FirstByte);
-			std::string MultipartUploadResponseRequestString = fmt::format("/api/v2/builds/{}/{}/{}/blobs/{}/uploadMultipart{}",
-																		   Namespace,
-																		   BucketId,
-																		   BuildId,
-																		   Hash.ToHexString(),
-																		   Part.QueryString);
+			std::string							 MultipartUploadResponseRequestString =
+				fmt::format("/api/v2/builds/{}/{}/{}/blobs/{}/uploadMultipart{}&supportsRedirect={}",
+							Namespace,
+							BucketId,
+							BuildId,
+							Hash.ToHexString(),
+							Part.QueryString,
+							m_AllowRedirect ? "true"sv : "false"sv);
 			// ZEN_INFO("PUT: {}", MultipartUploadResponseRequestString);
 			HttpClient::Response MultipartUploadResponse = m_HttpClient.Put(MultipartUploadResponseRequestString, PartPayload);
 			if (!MultipartUploadResponse.IsSuccess())
@@ -605,12 +611,13 @@ JupiterSession::PutMultipartBuildBlob(std::string_view												  Namespace,
 								IoBuffer							 RetryPartPayload =
 									Workload->Transmitter(RetryPart.FirstByte, RetryPart.LastByte - RetryPart.FirstByte - 1);
 								std::string RetryMultipartUploadResponseRequestString =
-									fmt::format("/api/v2/builds/{}/{}/{}/blobs/{}/uploadMultipart{}",
+									fmt::format("/api/v2/builds/{}/{}/{}/blobs/{}/uploadMultipart{}&supportsRedirect={}",
 												Namespace,
 												BucketId,
 												BuildId,
 												Hash.ToHexString(),
-												RetryPart.QueryString);
+												RetryPart.QueryString,
+												m_AllowRedirect ? "true"sv : "false"sv);
 
 								MultipartUploadResponse = m_HttpClient.Put(RetryMultipartUploadResponseRequestString, RetryPartPayload);
 								TotalUploadedBytes		= MultipartUploadResponse.UploadedBytes;
@@ -650,7 +657,12 @@ JupiterSession::GetMultipartBuildBlob(std::string_view												Namespace,
 									  std::function<void()>&&										OnComplete,
 									  std::vector<std::function<JupiterResult()>>&					OutWorkItems)
 {
-	std::string			 RequestUrl = fmt::format("/api/v2/builds/{}/{}/{}/blobs/{}", Namespace, BucketId, BuildId, Hash.ToHexString());
+	std::string			 RequestUrl = fmt::format("/api/v2/builds/{}/{}/{}/blobs/{}?supportsRedirect={}",
+										  Namespace,
+										  BucketId,
+										  BuildId,
+										  Hash.ToHexString(),
+										  m_AllowRedirect ? "true"sv : "false"sv);
 	HttpClient::Response Response =
 		m_HttpClient.Get(RequestUrl, HttpClient::KeyValueMap({{"Range", fmt::format("bytes={}-{}", 0, ChunkSize - 1)}}));
 	if (Response.IsSuccess())
@@ -687,11 +699,15 @@ JupiterSession::GetMultipartBuildBlob(std::string_view												Namespace,
 							uint64_t PartSize = Min(ChunkSize, TotalSize - Offset);
 							OutWorkItems.emplace_back(
 								[this, Namespace, BucketId, BuildId, Hash, TotalSize, Workload, Offset, PartSize]() -> JupiterResult {
-									std::string RequestUrl =
-										fmt::format("/api/v2/builds/{}/{}/{}/blobs/{}", Namespace, BucketId, BuildId, Hash.ToHexString());
-									HttpClient::Response Response = m_HttpClient.Get(
-										RequestUrl,
-										HttpClient::KeyValueMap({{"Range", fmt::format("bytes={}-{}", Offset, Offset + PartSize - 1)}}));
+									std::string			 RequestUrl = fmt::format("/api/v2/builds/{}/{}/{}/blobs/{}?supportsRedirect={}",
+																		  Namespace,
+																		  BucketId,
+																		  BuildId,
+																		  Hash.ToHexString(),
+																		  m_AllowRedirect ? "true"sv : "false"sv);
+									HttpClient::Response Response	= m_HttpClient.Get(
+										  RequestUrl,
+										  HttpClient::KeyValueMap({{"Range", fmt::format("bytes={}-{}", Offset, Offset + PartSize - 1)}}));
 									if (Response.IsSuccess())
 									{
 										Workload->OnReceive(Offset, Response.ResponsePayload);
@@ -734,11 +750,28 @@ JupiterSession::GetBuildBlob(std::string_view	   Namespace,
 	{
 		Headers.Entries.insert({"Range", fmt::format("bytes={}-{}", Offset, Offset + Size - 1)});
 	}
-	HttpClient::Response Response =
-		m_HttpClient.Download(fmt::format("/api/v2/builds/{}/{}/{}/blobs/{}", Namespace, BucketId, BuildId, Hash.ToHexString()),
-							  TempFolderPath,
-							  Headers);
-
+	HttpClient::Response Response = m_HttpClient.Download(fmt::format("/api/v2/builds/{}/{}/{}/blobs/{}?supportsRedirect={}",
+																	  Namespace,
+																	  BucketId,
+																	  BuildId,
+																	  Hash.ToHexString(),
+																	  m_AllowRedirect ? "true"sv : "false"sv),
+														  TempFolderPath,
+														  Headers);
+	if (Response.IsSuccess())
+	{
+		// If we get a redirect to S3 or a non-Jupiter endpoint the content type will not be correct, validate it and set it
+		if (m_AllowRedirect && (Response.ResponsePayload.GetContentType() == HttpContentType::kBinary))
+		{
+			IoHash	 ValidateRawHash;
+			uint64_t ValidateRawSize = 0;
+			ZEN_ASSERT_SLOW(CompressedBuffer::ValidateCompressedHeader(Response.ResponsePayload, ValidateRawHash, ValidateRawSize));
+			ZEN_ASSERT_SLOW(ValidateRawHash == Hash);
+			ZEN_ASSERT_SLOW(ValidateRawSize > 0);
+			ZEN_UNUSED(ValidateRawHash, ValidateRawSize);
+			Response.ResponsePayload.SetContentType(ZenContentType::kCompressedBinary);
+		}
+	}
 	return detail::ConvertResponse(Response, "JupiterSession::GetBuildBlob"sv);
 }