From 0b7bfaf44b0416615b1efe101ae4808ab889e5ec Mon Sep 17 00:00:00 2001
From: Dillon Skaggs <dillon@dillonskaggs.dev>
Date: Wed, 17 Jan 2024 13:21:43 -0600
Subject: [PATCH] fix(state/statebag): filter frequent & large data from client

* if we accept these willy-nilly, clients could possibly OOM the server, or do a one-to-many make-shift DoS on larger a server
* the size limits here are kept high to keep compatibility with people who "misuse" state bags
* sane defaults for servers that don't send large data via state bags would be `set rateLimiter_stateBagSize_rate 8196` and `set rateLimiter_stateBagSize_burst 16364`
---
 .../src/state/ServerGameState.cpp             | 44 ++++++++++++++++++-
 1 file changed, 42 insertions(+), 2 deletions(-)

diff --git a/code/components/citizen-server-impl/src/state/ServerGameState.cpp b/code/components/citizen-server-impl/src/state/ServerGameState.cpp
index 08fcf322d7..1bcefc0ef2 100644
--- a/code/components/citizen-server-impl/src/state/ServerGameState.cpp
+++ b/code/components/citizen-server-impl/src/state/ServerGameState.cpp
@@ -30,6 +30,8 @@
 #include <boost/range/adaptors.hpp>
 #include <boost/math/constants/constants.hpp>
 
+#include <KeyedRateLimiter.h>
+
 #include <OneSyncVars.h>
 #include <DebugAlias.h>
 
@@ -4250,8 +4252,46 @@ void ServerGameState::AttachToObject(fx::ServerInstanceBase* instance)
 	sbac->SetGameInterface(this);
 
 	instance->GetComponent<fx::GameServer>()->GetComponent<fx::HandlerMapComponent>()->Add(HashRageString("msgStateBag"), 
-		{ fx::ThreadIdx::Sync, [this](const fx::ClientSharedPtr& client, net::Buffer& buffer)
-	{
+		{ fx::ThreadIdx::Sync, [this, instance](const fx::ClientSharedPtr& client, net::Buffer& buffer)
+	{
+		static fx::RateLimiterStore<uint32_t, false> stateBagRateLimiterStore{ instance->GetComponent<console::Context>().GetRef() };
+		static auto stateBagRateLimiter = stateBagRateLimiterStore.GetRateLimiter("stateBag", fx::RateLimiterDefaults{ 50.f, 100.f });
+		static auto stateBagFloodRateLimiter = stateBagRateLimiterStore.GetRateLimiter("stateBagFlood", fx::RateLimiterDefaults{ 50.f, 150.f });
+		/*
+		 * You can use rateLimiter_stateBagSize_rate to increase the regular rate limit, increasing the burst shouldn't be done as
+		 * state bag packets *should never be* above this limit any way (refer to StateBagImpl::SendKeyValue "dataBuffer"
+		 * 
+		 * It might be better to have these set to lower defaults, though in this case maintaining compatibility is of higher priority.
+		 */
+		static auto stateBagSizeRateLimiter = stateBagRateLimiterStore.GetRateLimiter("stateBagSize", fx::RateLimiterDefaults{ 64 * 1024.0, 128 * 1024.0 });
+		
+		const uint32_t netId = client->GetNetId();
+
+		if (!stateBagRateLimiter->Consume(netId))
+		{
+			if (!stateBagFloodRateLimiter->Consume(netId))
+			{
+				gscomms_execute_callback_on_main_thread([client, instance]()
+				{
+					instance->GetComponent<fx::GameServer>()->DropClient(client, "Reliable state bag packet overflow.");
+				});
+			}
+
+			return;
+		}
+
+		uint32_t dataLength = buffer.GetRemainingBytes();
+		if (!stateBagSizeRateLimiter->Consume(netId, double(dataLength)))
+		{
+			gscomms_execute_callback_on_main_thread([client, instance]()
+			{
+				// if this happens, try increasing rateLimiter_stateBagSize_rate and rateLimiter_stateBagSize_burst, though
+				// the client shouldn't be sending large packets to the server
+				instance->GetComponent<fx::GameServer>()->DropClient(client, "Reliable state bag packet overflow.");
+			});
+
+			return;
+		}
 		
 		uint32_t slotId = client->GetSlotId();
 		if (slotId != -1)