From f714751342e996697e136dbf027cd12393fad493 Mon Sep 17 00:00:00 2001
From: Dan Engelbrecht <dan.engelbrecht@epicgames.com>
Date: Tue, 10 Dec 2024 09:09:51 +0100
Subject: improved payload validation in HttpClient (#259)

* improved payload validation in HttpClient
* separate error messages for FromCompressed and Decompress
* refactor so we can do retry if decompression of block fails
---
 src/zenhttp/httpclient.cpp | 78 +++++++++++++++++++++++++++++++++++-----------
 1 file changed, 60 insertions(+), 18 deletions(-)

(limited to 'src/zenhttp/httpclient.cpp')

diff --git a/src/zenhttp/httpclient.cpp b/src/zenhttp/httpclient.cpp
index 9af909fcf..d8ce25304 100644
--- a/src/zenhttp/httpclient.cpp
+++ b/src/zenhttp/httpclient.cpp
@@ -375,37 +375,75 @@ ShouldRetry(const cpr::Response& Response)
 static bool
 ValidatePayload(cpr::Response& Response, std::unique_ptr<detail::TempPayloadFile>& PayloadFile)
 {
-	if (IsHttpSuccessCode(Response.status_code))
+	IoBuffer ResponseBuffer = (Response.text.empty() && PayloadFile) ? PayloadFile->BorrowIoBuffer()
+																	 : IoBuffer(IoBuffer::Wrap, Response.text.data(), Response.text.size());
+
+	if (auto ContentLength = Response.header.find("Content-Length"); ContentLength != Response.header.end())
+	{
+		std::optional<uint64_t> ExpectedContentSize = ParseInt<uint64_t>(ContentLength->second);
+		if (!ExpectedContentSize.has_value())
+		{
+			Response.error =
+				cpr::Error(/*CURLE_READ_ERROR*/ 26, fmt::format("Can not parse Content-Length header. Value: '{}'", ContentLength->second));
+			return false;
+		}
+		if (ExpectedContentSize.value() != ResponseBuffer.GetSize())
+		{
+			Response.error = cpr::Error(
+				/*CURLE_READ_ERROR*/ 26,
+				fmt::format("Payload size {} does not match Content-Length {}", ResponseBuffer.GetSize(), ContentLength->second));
+			return false;
+		}
+	}
+
+	if (auto JupiterHash = Response.header.find("X-Jupiter-IoHash"); JupiterHash != Response.header.end())
 	{
-		if (auto ContentType = Response.header.find("Content-Type");
-			ContentType != Response.header.end() && ContentType->second == "application/x-ue-comp")
+		IoHash ExpectedPayloadHash;
+		if (IoHash::TryParse(JupiterHash->second, ExpectedPayloadHash))
+		{
+			IoHash PayloadHash = IoHash::HashBuffer(ResponseBuffer);
+			if (PayloadHash != ExpectedPayloadHash)
+			{
+				Response.error = cpr::Error(/*CURLE_READ_ERROR*/ 26,
+											fmt::format("Payload hash {} does not match X-Jupiter-IoHash {}",
+														PayloadHash.ToHexString(),
+														ExpectedPayloadHash.ToHexString()));
+				return false;
+			}
+		}
+	}
+
+	if (auto ContentType = Response.header.find("Content-Type"); ContentType != Response.header.end())
+	{
+		if (ContentType->second == "application/x-ue-comp")
 		{
-			IoBuffer ResponseBuffer = (Response.text.empty() && PayloadFile)
-										  ? PayloadFile->BorrowIoBuffer()
-										  : IoBuffer(IoBuffer::Wrap, Response.text.data(), Response.text.size());
 			IoHash	 RawHash;
 			uint64_t RawSize;
-			if (!CompressedBuffer::ValidateCompressedHeader(ResponseBuffer, RawHash, RawSize))
+			if (CompressedBuffer::ValidateCompressedHeader(ResponseBuffer, RawHash, RawSize))
+			{
+				return true;
+			}
+			else
 			{
 				Response.error = cpr::Error(/*CURLE_READ_ERROR*/ 26, "Compressed binary failed validation");
 				return false;
 			}
 		}
-		else if (auto JupiterHash = Response.header.find("X-Jupiter-IoHash"); JupiterHash != Response.header.end())
+		if (ContentType->second == "application/x-ue-cb")
 		{
-			IoHash ExpectedPayloadHash;
-			if (IoHash::TryParse(JupiterHash->second, ExpectedPayloadHash))
+			if (CbValidateError Error = ValidateCompactBinary(ResponseBuffer.GetView(), CbValidateMode::Default);
+				Error == CbValidateError::None)
 			{
-				IoBuffer ResponseBuffer(IoBuffer::Wrap, Response.text.data(), Response.text.size());
-				IoHash	 PayloadHash = IoHash::HashBuffer(ResponseBuffer);
-				if (PayloadHash != ExpectedPayloadHash)
-				{
-					Response.error = cpr::Error(/*CURLE_READ_ERROR*/ 26, "Compressed binary failed validation");
-					return false;
-				}
+				return true;
+			}
+			else
+			{
+				Response.error = cpr::Error(/*CURLE_READ_ERROR*/ 26, fmt::format("Compact binary failed validation: {}", ToString(Error)));
+				return false;
 			}
 		}
 	}
+
 	return true;
 }
 
@@ -433,7 +471,11 @@ DoWithRetry(std::function<cpr::Response()>&& Func, std::unique_ptr<detail::TempP
 	{
 		if (!ShouldRetry(Result))
 		{
-			if (Result.error || ValidatePayload(Result, PayloadFile))
+			if (Result.error || !IsHttpSuccessCode(Result.status_code))
+			{
+				break;
+			}
+			if (ValidatePayload(Result, PayloadFile))
 			{
 				break;
 			}
-- 
cgit v1.2.3