Merge branch 'de/zen-service-command' of https://github.ol.epicgames.net/ue-foundation/zen into de/zen-service-command

1 files changed, 168 insertions, 69 deletions

diff --git a/src/zenhttp/httpclient.cpp b/src/zenhttp/httpclient.cpp
index 211a5d05c..30711a432 100644
--- a/src/zenhttp/httpclient.cpp
+++ b/src/zenhttp/httpclient.cpp
@@ -414,6 +414,11 @@ ValidatePayload(cpr::Response& Response, std::unique_ptr<detail::TempPayloadFile
 		}
 	}
 
+	if (Response.status_code == (long)HttpResponseCode::PartialContent)
+	{
+		return true;
+	}
+
 	if (auto ContentType = Response.header.find("Content-Type"); ContentType != Response.header.end())
 	{
 		if (ContentType->second == "application/x-ue-comp")
@@ -449,12 +454,27 @@ ValidatePayload(cpr::Response& Response, std::unique_ptr<detail::TempPayloadFile
 }
 
 static cpr::Response
-DoWithRetry(std::string_view SessionId, std::function<cpr::Response()>&& Func, uint8_t RetryCount)
+DoWithRetry(
+	std::string_view							 SessionId,
+	std::function<cpr::Response()>&&			 Func,
+	uint8_t										 RetryCount,
+	std::function<bool(cpr::Response& Result)>&& Validate = [](cpr::Response&) { return true; })
 {
 	uint8_t		  Attempt = 0;
 	cpr::Response Result  = Func();
-	while (Attempt < RetryCount && ShouldRetry(Result))
+	while (Attempt < RetryCount)
 	{
+		if (!ShouldRetry(Result))
+		{
+			if (Result.error || !IsHttpSuccessCode(Result.status_code))
+			{
+				break;
+			}
+			if (Validate(Result))
+			{
+				break;
+			}
+		}
 		Sleep(100 * (Attempt + 1));
 		Attempt++;
 		ZEN_INFO("{} Attempt {}/{}", CommonResponse(SessionId, std::move(Result)).ErrorMessage("Retry"), Attempt, RetryCount + 1);
@@ -881,7 +901,11 @@ HttpClient::Get(std::string_view Url, const KeyValueMap& AdditionalHeader, const
 					m_Impl->AllocSession(m_BaseUri, Url, m_ConnectionSettings, AdditionalHeader, Parameters, m_SessionId, GetAccessToken());
 				return Sess.Get();
 			},
-			m_ConnectionSettings.RetryCount));
+			m_ConnectionSettings.RetryCount,
+			[](cpr::Response& Result) {
+				std::unique_ptr<detail::TempPayloadFile> NoTempFile;
+				return ValidatePayload(Result, NoTempFile);
+			}));
 }
 
 HttpClient::Response
@@ -1124,6 +1148,30 @@ HttpClient::Download(std::string_view Url, const std::filesystem::path& TempFold
 				   return true;
 			   };
 
+			   uint64_t RequestedContentLength = (uint64_t)-1;
+			   if (auto RangeIt = AdditionalHeader.Entries.find("Range"); RangeIt != AdditionalHeader.Entries.end())
+			   {
+				   if (RangeIt->second.starts_with("bytes"))
+				   {
+					   size_t RangeStartPos = RangeIt->second.find('=', 5);
+					   if (RangeStartPos != std::string::npos)
+					   {
+						   RangeStartPos++;
+						   size_t RangeSplitPos = RangeIt->second.find('-', RangeStartPos);
+						   if (RangeSplitPos != std::string::npos)
+						   {
+							   std::optional<size_t> RequestedRangeStart =
+								   ParseInt<size_t>(RangeIt->second.substr(RangeStartPos, RangeSplitPos - RangeStartPos));
+							   std::optional<size_t> RequestedRangeEnd = ParseInt<size_t>(RangeIt->second.substr(RangeStartPos + 1));
+							   if (RequestedRangeStart.has_value() && RequestedRangeEnd.has_value())
+							   {
+								   RequestedContentLength = RequestedRangeEnd.value() - 1;
+							   }
+						   }
+					   }
+				   }
+			   }
+
 			   cpr::Response Response;
 			   {
 				   std::vector<std::pair<std::string, std::string>> ReceivedHeaders;
@@ -1131,10 +1179,10 @@ HttpClient::Download(std::string_view Url, const std::filesystem::path& TempFold
 					   std::pair<std::string, std::string> Header = GetHeader(header);
 					   if (Header.first == "Content-Length"sv)
 					   {
-						   std::optional<size_t> ContentSize = ParseInt<size_t>(Header.second);
-						   if (ContentSize.has_value())
+						   std::optional<size_t> ContentLength = ParseInt<size_t>(Header.second);
+						   if (ContentLength.has_value())
 						   {
-							   if (ContentSize.value() > 1024 * 1024)
+							   if (ContentLength.value() > 1024 * 1024)
 							   {
 								   PayloadFile = std::make_unique<detail::TempPayloadFile>();
 								   std::error_code Ec = PayloadFile->Open(TempFolderPath);
@@ -1148,7 +1196,7 @@ HttpClient::Download(std::string_view Url, const std::filesystem::path& TempFold
 							   }
 							   else
 							   {
-								   PayloadString.reserve(ContentSize.value());
+								   PayloadString.reserve(ContentLength.value());
 							   }
 						   }
 					   }
@@ -1194,85 +1242,90 @@ HttpClient::Download(std::string_view Url, const std::filesystem::path& TempFold
 					   auto It = Response.header.find("Content-Length");
 					   if (It != Response.header.end())
 					   {
-						   std::optional<int64_t> ContentLength = ParseInt<int64_t>(It->second);
-						   if (ContentLength)
-						   {
-							   std::vector<std::pair<std::string, std::string>> ReceivedHeaders;
+						   std::vector<std::pair<std::string, std::string>> ReceivedHeaders;
 
-							   auto HeaderCallback = [&](std::string header, intptr_t) {
-								   std::pair<std::string, std::string> Header = GetHeader(header);
-								   if (!Header.first.empty())
-								   {
-									   ReceivedHeaders.emplace_back(std::move(Header));
-								   }
+						   auto HeaderCallback = [&](std::string header, intptr_t) {
+							   std::pair<std::string, std::string> Header = GetHeader(header);
+							   if (!Header.first.empty())
+							   {
+								   ReceivedHeaders.emplace_back(std::move(Header));
+							   }
 
-								   if (Header.first == "Content-Range"sv)
+							   if (Header.first == "Content-Range"sv)
+							   {
+								   if (Header.second.starts_with("bytes "sv))
 								   {
-									   if (Header.second.starts_with("bytes "sv))
+									   size_t RangeStartEnd = Header.second.find('-', 6);
+									   if (RangeStartEnd != std::string::npos)
 									   {
-										   size_t RangeStartEnd = Header.second.find('-', 6);
-										   if (RangeStartEnd != std::string::npos)
+										   const auto Start = ParseInt<uint64_t>(Header.second.substr(6, RangeStartEnd - 6));
+										   if (Start)
 										   {
-											   const auto Start = ParseInt<uint64_t>(Header.second.substr(6, RangeStartEnd - 6));
-											   if (Start)
+											   uint64_t DownloadedSize = PayloadFile ? PayloadFile->GetSize() : PayloadString.length();
+											   if (Start.value() == DownloadedSize)
 											   {
-												   uint64_t DownloadedSize = PayloadFile ? PayloadFile->GetSize() : PayloadString.length();
-												   if (Start.value() == DownloadedSize)
-												   {
-													   return 1;
-												   }
-												   else if (Start.value() > DownloadedSize)
-												   {
-													   return 0;
-												   }
-												   if (PayloadFile)
-												   {
-													   PayloadFile->ResetWritePos(Start.value());
-												   }
-												   else
-												   {
-													   PayloadString = PayloadString.substr(0, Start.value());
-												   }
 												   return 1;
 											   }
+											   else if (Start.value() > DownloadedSize)
+											   {
+												   return 0;
+											   }
+											   if (PayloadFile)
+											   {
+												   PayloadFile->ResetWritePos(Start.value());
+											   }
+											   else
+											   {
+												   PayloadString = PayloadString.substr(0, Start.value());
+											   }
+											   return 1;
 										   }
 									   }
-									   return 0;
 								   }
-								   return 1;
-							   };
+								   return 0;
+							   }
+							   return 1;
+						   };
 
-							   KeyValueMap HeadersWithRange(AdditionalHeader);
-							   do
-							   {
-								   uint64_t DownloadedSize = PayloadFile ? PayloadFile->GetSize() : PayloadString.length();
+						   KeyValueMap HeadersWithRange(AdditionalHeader);
+						   do
+						   {
+							   uint64_t DownloadedSize = PayloadFile ? PayloadFile->GetSize() : PayloadString.length();
 
-								   std::string Range = fmt::format("bytes={}-{}", DownloadedSize, ContentLength.value());
-								   if (auto RangeIt = HeadersWithRange.Entries.find("Range"); RangeIt != HeadersWithRange.Entries.end())
+							   uint64_t ContentLength = RequestedContentLength;
+							   if (ContentLength == uint64_t(-1))
+							   {
+								   if (auto ParsedContentLength = ParseInt<int64_t>(It->second); ParsedContentLength.has_value())
 								   {
-									   if (RangeIt->second == Range)
-									   {
-										   // If we didn't make any progress, abort
-										   break;
-									   }
+									   ContentLength = ParsedContentLength.value();
 								   }
-								   HeadersWithRange.Entries.insert_or_assign("Range", Range);
-
-								   Impl::Session Sess = m_Impl->AllocSession(m_BaseUri,
-																			 Url,
-																			 m_ConnectionSettings,
-																			 HeadersWithRange,
-																			 {},
-																			 m_SessionId,
-																			 GetAccessToken());
-								   Response = Sess.Download(cpr::WriteCallback{DownloadCallback}, cpr::HeaderCallback{HeaderCallback});
-								   for (const std::pair<std::string, std::string>& H : ReceivedHeaders)
+							   }
+
+							   std::string Range = fmt::format("bytes={}-{}", DownloadedSize, DownloadedSize + ContentLength - 1);
+							   if (auto RangeIt = HeadersWithRange.Entries.find("Range"); RangeIt != HeadersWithRange.Entries.end())
+							   {
+								   if (RangeIt->second == Range)
 								   {
-									   Response.header.insert_or_assign(H.first, H.second);
+									   // If we didn't make any progress, abort
+									   break;
 								   }
-								   ReceivedHeaders.clear();
-							   } while (ShouldResume(Response));
-						   }
+							   }
+							   HeadersWithRange.Entries.insert_or_assign("Range", Range);
+
+							   Impl::Session Sess = m_Impl->AllocSession(m_BaseUri,
+																		 Url,
+																		 m_ConnectionSettings,
+																		 HeadersWithRange,
+																		 {},
+																		 m_SessionId,
+																		 GetAccessToken());
+							   Response = Sess.Download(cpr::WriteCallback{DownloadCallback}, cpr::HeaderCallback{HeaderCallback});
+							   for (const std::pair<std::string, std::string>& H : ReceivedHeaders)
+							   {
+								   Response.header.insert_or_assign(H.first, H.second);
+							   }
+							   ReceivedHeaders.clear();
+						   } while (ShouldResume(Response));
 					   }
 				   }
 			   }
@@ -1398,6 +1451,52 @@ HttpClient::Response::ThrowError(std::string_view ErrorPrefix)
 
 #if ZEN_WITH_TESTS
 
+TEST_CASE("responseformat")
+{
+	using namespace std::literals;
+
+	SUBCASE("identity")
+	{
+		BodyLogFormatter _{"abcd"};
+		CHECK_EQ(_.GetText(), "abcd"sv);
+	}
+
+	SUBCASE("very long")
+	{
+		std::string_view LongView =
+			"abcdefghijklmnopqrstuvwxyzabcdefghijklmnopqrstuvwxyzabcdefghijklmnopqrstuvwxyzabcdefghijklmnopqrstuvwxyz"
+			"abcdefghijklmnopqrstuvwxyzabcdefghijklmnopqrstuvwxyzabcdefghijklmnopqrstuvwxyzabcdefghijklmnopqrstuvwxyz"
+			"abcdefghijklmnopqrstuvwxyzabcdefghijklmnopqrstuvwxyzabcdefghijklmnopqrstuvwxyzabcdefghijklmnopqrstuvwxyz"
+			"abcdefghijklmnopqrstuvwxyzabcdefghijklmnopqrstuvwxyzabcdefghijklmnopqrstuvwxyzabcdefghijklmnopqrstuvwxyz"
+			"abcdefghijklmnopqrstuvwxyzabcdefghijklmnopqrstuvwxyzabcdefghijklmnopqrstuvwxyzabcdefghijklmnopqrstuvwxyz"
+			"abcdefghijklmnopqrstuvwxyzabcdefghijklmnopqrstuvwxyzabcdefghijklmnopqrstuvwxyzabcdefghijklmnopqrstuvwxyz"
+			"abcdefghijklmnopqrstuvwxyzabcdefghijklmnopqrstuvwxyzabcdefghijklmnopqrstuvwxyzabcdefghijklmnopqrstuvwxyz"
+			"abcdefghijklmnopqrstuvwxyzabcdefghijklmnopqrstuvwxyzabcdefghijklmnopqrstuvwxyzabcdefghijklmnopqrstuvwxyz"
+			"abcdefghijklmnopqrstuvwxyzabcdefghijklmnopqrstuvwxyzabcdefghijklmnopqrstuvwxyzabcdefghijklmnopqrstuvwxyz"
+			"abcdefghijklmnopqrstuvwxyzabcdefghijklmnopqrstuvwxyzabcdefghijklmnopqrstuvwxyzabcdefghijklmnopqrstuvwxyz"
+			"abcdefghijklmnopqrstuvwxyzabcdefghijklmnopqrstuvwxyzabcdefghijklmnopqrstuvwxyzabcdefghijklmnopqrstuvwxyz"
+			"abcdefghijklmnopqrstuvwxyzabcdefghijklmnopqrstuvwxyzabcdefghijklmnopqrstuvwxyzabcdefghijklmnopqrstuvwxyz"
+			"abcdefghijklmnopqrstuvwxyzabcdefghijklmnopqrstuvwxyzabcdefghijklmnopqrstuvwxyzabcdefghijklmnopqrstuvwxyz"
+			"abcdefghijklmnopqrstuvwxyzabcdefghijklmnopqrstuvwxyzabcdefghijklmnopqrstuvwxyzabcdefghijklmnopqrstuvwxyz"
+			"abcdefghijklmnopqrstuvwxyzabcdefghijklmnopqrstuvwxyzabcdefghijklmnopqrstuvwxyzabcdefghijklmnopqrstuvwxyz"
+			"abcdefghijklmnopqrstuvwxyzabcdefghijklmnopqrstuvwxyzabcdefghijklmnopqrstuvwxyzabcdefghijklmnopqrstuvwxyz";
+
+		BodyLogFormatter _{LongView};
+
+		CHECK(_.GetText().size() < LongView.size());
+		CHECK(_.GetText().starts_with("[truncated"sv));
+	}
+
+	SUBCASE("invalid text")
+	{
+		std::string_view BadText = "totobaba\xff\xfe";
+
+		BodyLogFormatter _{BadText};
+
+		CHECK_EQ(_.GetText(), "totobaba");
+	}
+}
+
 TEST_CASE("httpclient")
 {
 	using namespace std::literals;