From 8fd89d68fe2ad6c39fd7cb272954865b697a52d3 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             | 77 ++++++++++++++++++-
 1 file changed, 76 insertions(+), 1 deletion(-)

diff --git a/code/components/citizen-server-impl/src/state/ServerGameState.cpp b/code/components/citizen-server-impl/src/state/ServerGameState.cpp
index 08fcf322d7..15f8f66e0b 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,81 @@ 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() };
+			
+		constexpr double kStateBagRateLimit = 75.f;
+		constexpr double kStateBagRateLimitBurst = 125.f;
+		static auto stateBagRateLimiter = stateBagRateLimiterStore.GetRateLimiter("stateBag", fx::RateLimiterDefaults{ kStateBagRateLimit, kStateBagRateLimitBurst });
+			
+		constexpr double kStateBagRateFloodLimit = 150.f;
+		constexpr double kStateBagRateFloodLimitBurst = 175.f;
+		static auto stateBagRateFloodLimiter = stateBagRateLimiterStore.GetRateLimiter("stateBagFlood", fx::RateLimiterDefaults{ kStateBagRateFloodLimit, kStateBagRateFloodLimitBurst });
+
+		constexpr double kStateBagSizeLimit = 128 * 1024.0;
+		constexpr double kStateBagSizeLimitBurst = 256 * 1024.0;
+		static auto stateBagSizeRateLimiter = stateBagRateLimiterStore.GetRateLimiter("stateBagSize", fx::RateLimiterDefaults{ kStateBagSizeLimit, kStateBagSizeLimitBurst });
+
+		static fx::KeyedRateLimiter<uint32_t, true> logLimiter { 1.0, 1.0 };
+
+		const uint32_t netId = client->GetNetId();
+
+		const bool hitRateLimit = !stateBagRateLimiter->Consume(netId);
+		const bool hitFloodRateLimit = !stateBagRateFloodLimiter->Consume(netId);
+
+		if (hitRateLimit)
+		{
+			const std::string& clientName = client->GetName();
+			auto printStateWarning = [&clientName, netId](const std::string& logChannel, const std::string_view logReason, const std::string_view rateLimiter, double rateLimit, double burstRateLimit)
+			{
+				console::Printf(logChannel, logReason, clientName, netId);
+				console::Printf(logChannel, "If you believe this to be a mistake please increase your rateLimiter_%s_rate and rateLimiter_%s_burst. ", rateLimiter, rateLimiter);
+				console::Printf(logChannel, "You can do this with `set rateLimiter_%s_rate [new value]`. The default rate limit is %0.0f and burst limit is %0.0f\n", rateLimiter, rateLimit, burstRateLimit);
+				console::Printf(logChannel, "You can disable this warning with `con_addChannelFilter %s drop` if you think you have this properly set up.\n", logChannel);
+			};
+
+			if (hitFloodRateLimit)
+			{
+				if (!client->IsDropping())
+				{
+					printStateWarning("sbag-client-flood",
+						"Client %s %d got dropped for sending too many state bag value updates.\n",
+						"stateBagFlood",
+						 kStateBagRateFloodLimit, kStateBagRateFloodLimitBurst);
+				}
+
+				instance->GetComponent<fx::GameServer>()->DropClient(client, "Reliable state bag packet overflow.");
+				return;
+			}
+
+			// only log here if we didn't log before, logging should only happen once in every 15 seconds.
+			if (logLimiter.Consume(netId))
+			{
+				printStateWarning("sbag-update-dropped",
+					"Client %s %d sent too many state bag updates and had their updates dropped.\n",
+					"stateBag",
+					kStateBagRateLimit, kStateBagRateLimitBurst);
+			}
+
+			return;
+		}
+
+		uint32_t dataLength = buffer.GetRemainingBytes();
+		if (!stateBagSizeRateLimiter->Consume(netId, double(dataLength)))
+		{
+			if (!client->IsDropping())
+			{
+				const std::string& logChannel = "sbag-size-kick";
+				console::Printf(logChannel, "Client %s %d got dropped for sending too large of a state bag update.\n", client->GetName(), netId);
+				console::Printf(logChannel, "If you believe this to be a mistake please increase your rateLimiter_stateBagSize_rate. ");
+				console::Printf(logChannel, "You can do this with `set rateLimiter_stateBagSize_rate [new value]`. The default size rate limit is %0.0f.\n", kStateBagSizeLimit);
+				console::Printf(logChannel, "You can disable this warning with `con_addChannelFilter %s drop` if you think you have this properly set up.\n", logChannel);
+			}
+
+			instance->GetComponent<fx::GameServer>()->DropClient(client, "Reliable state bag packet overflow.");
+			return;
+		}
 		
 		uint32_t slotId = client->GetSlotId();
 		if (slotId != -1)