tweak(core/sbac): cleanup code, return on read failing, add ReadBitsSafe

code/components/citizen-resources-core/src/StateBagComponent.cpp

@@ -601,19 +601,30 @@ void StateBagComponentImpl::HandlePacket(int source, std::string_view dataRaw, s
 	{
 		uint16_t length = 0;
 
-		buffer.Read<uint16_t>(16, &length);
+		auto emptyVec = []
+		{
+			return std::vector<char>{};
+		};
+
+		if (!buffer.Read<uint16_t>(16, &length))
+		{
+			return emptyVec();
+		}
 
 		// validate input
-		if (length <= 0 || length > std::numeric_limits<uint16_t>::max())
+		if (length <= 0)
 		{
-			return std::vector<char>{};
+			return emptyVec();
 		}
 
-		std::vector<char> rbuffer(length - 1);
-		buffer.ReadBits(rbuffer.data(), rbuffer.size() * 8);
+		auto char_buffer = buffer.ReadBitsSafe<char>((length - 1) * 8);
+		if (!char_buffer)
+		{
+			return emptyVec();
+		}
 		buffer.Read<uint8_t>(8);
 
-		return std::move(rbuffer);
+		return std::move(*char_buffer);
 	};
 
 	// read id
@@ -632,23 +643,13 @@ void StateBagComponentImpl::HandlePacket(int source, std::string_view dataRaw, s
 		return;
 	}
 
-	// if m_curBit is greater then m_maxBit we will overflow the dataLength, which would lead to an allocation of an
-	// extremely large buffer, which would fail and crash the server.
-	if (buffer.IsAtEnd())
-	{
-		return;
-	}
-
-	// read data
-	size_t dataLength = (buffer.GetLength() * 8) - buffer.GetCurrentBit();
-
-	if (dataLength == 0)
+	auto optional_data = buffer.ReadBitsSafe<uint8_t>(buffer.GetLength() * 8 - buffer.GetCurrentBit());
+	if (!optional_data)
 	{
 		return;
 	}
 
-	std::vector<uint8_t> data(dataLength / 8);
-	buffer.ReadBits(data.data(), dataLength);
+	auto data = *optional_data;
 
 	// handle data
 	auto bagName = std::string_view{

code/shared/state/RlMessageBuffer.h

@@ -311,7 +311,6 @@ class MessageBuffer
 		return true;
 	}
 
-
 	// copied IDA code, please improve!
 	inline bool ReadBits(void* data, int length)
 	{
@@ -332,6 +331,31 @@ class MessageBuffer
 		return rv;
 	}
 
+	template<typename T>
+	inline std::optional<std::vector<T>> ReadBitsSafe(int length)
+	{
+		static_assert(std::is_integral_v<T>, "ReadBitsSafe wants an int value");
+
+		if (length <= 0)
+		{
+			return std::nullopt;
+		}
+
+		if ((m_curBit + length) > m_maxBit)
+		{
+			return std::nullopt;
+		}
+
+		std::vector<T> data(length);
+
+		// CopyBits only seems to ever return true, so no need for a check here
+		auto rv = CopyBits(data.data(), m_data.data(), length, 0, m_curBit);
+
+		m_curBit += length;
+
+		return data;
+	}
+
 	// copied IDA code, please improve!
 	inline bool WriteBits(const void* data, int length)
 	{