From eab9fc71b7f796fe667ae08850ff8048d9530247 Mon Sep 17 00:00:00 2001
From: Fabian Terhorst <fabian.terhorst@gmail.com>
Date: Wed, 25 Sep 2024 18:53:40 +0200
Subject: [PATCH] fix(net): prevent component size calculation integer overflow

---
 .../components/net-base/include/ByteCounter.h | 46 +++++++++++++++++--
 code/tests/shared/TestComponents.cpp          | 25 ++++++++++
 2 files changed, 67 insertions(+), 4 deletions(-)

diff --git a/code/components/net-base/include/ByteCounter.h b/code/components/net-base/include/ByteCounter.h
index c2c806e583..e77e90ad1f 100644
--- a/code/components/net-base/include/ByteCounter.h
+++ b/code/components/net-base/include/ByteCounter.h
@@ -38,7 +38,14 @@ namespace net
 		template <typename T>
 		bool Field(T& value)
 		{
-			m_counter += sizeof(value);
+			const uint64_t size = sizeof(value);
+			if (m_counter > m_counter + size)
+			{
+				m_counter = std::numeric_limits<uint64_t>::max();
+				return true;
+			}
+
+			m_counter += size;
 			return true;
 		}
 
@@ -51,6 +58,12 @@ namespace net
 		template <typename T>
 		bool Field(T& value, const size_t size)
 		{
+			if (m_counter > m_counter + size)
+			{
+				m_counter = std::numeric_limits<uint64_t>::max();
+				return true;
+			}
+
 			m_counter += size;
 			return true;
 		}
@@ -64,8 +77,14 @@ namespace net
 		template <typename T>
 		bool Field(Span<T>& value, const size_t size)
 		{
-			m_counter += size * sizeof(T);
-		
+			const uint64_t valueSize = size * sizeof(T);
+			if (m_counter > m_counter + valueSize)
+			{
+				m_counter = std::numeric_limits<uint64_t>::max();
+				return true;
+			}
+	
+			m_counter += valueSize;
 			return true;
 		}
 	};
@@ -73,13 +92,26 @@ namespace net
 	template <>
 	inline bool ByteCounter::Field<bool>(bool& value)
 	{
-		m_counter += 1;
+		const uint64_t valueSize = 1;
+		if (m_counter > m_counter + valueSize)
+		{
+			m_counter = std::numeric_limits<uint64_t>::max();
+			return true;
+		}
+
+		m_counter += valueSize;
 		return true;
 	}
 	
 	template <>
 	inline bool ByteCounter::Field<std::string_view>(std::string_view& value, const size_t size)
 	{
+		if (m_counter > m_counter + size)
+		{
+			m_counter = std::numeric_limits<uint64_t>::max();
+			return true;
+		}
+
 		m_counter += size;
 		return true;
 	}
@@ -87,6 +119,12 @@ namespace net
 	template <>
 	inline bool ByteCounter::Field<std::string>(std::string& value, const size_t size)
 	{
+		if (m_counter > m_counter + size)
+		{
+			m_counter = std::numeric_limits<uint64_t>::max();
+			return true;
+		}
+
 		m_counter += size;
 		return true;
 	}
diff --git a/code/tests/shared/TestComponents.cpp b/code/tests/shared/TestComponents.cpp
index 9b9b21b57f..2387996b6e 100644
--- a/code/tests/shared/TestComponents.cpp
+++ b/code/tests/shared/TestComponents.cpp
@@ -195,3 +195,28 @@ TEST_CASE("Serializable component property")
 
 	REQUIRE(instanceRead.component.GetValue().value.GetValue() == random);
 }
+
+namespace SerializableComponentUnlimitedSizeTest
+{
+class Component: public net::SerializableComponent
+{
+public:
+	net::SerializableProperty<uint8_t> value1 {};
+	net::SerializableProperty<net::Span<uint8_t>, net::storage_type::StreamTail> value2 {};
+
+	template <typename T>
+	bool Process(T& stream)
+	{
+		return ProcessPropertiesInOrder<T>(
+			stream,
+			value1,
+			value2
+		);
+	}
+};
+}
+
+TEST_CASE("Serializable component size of unlimited size")
+{
+	REQUIRE(net::SerializableComponent::GetSize<SerializableComponentUnlimitedSizeTest::Component>() == UINT64_MAX);
+}