From 4d02fd860349cc49923288c1dc5777707fbf4af9 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Piotr=20Wo=CC=81jcik?=
 <3044353+pwojcikdev@users.noreply.github.com>
Date: Sun, 5 Jan 2025 22:42:11 +0100
Subject: [PATCH] Rate limit bootstrap server

---
 nano/node/bootstrap/bootstrap_server.cpp | 25 +++++++++++++++++++-----
 nano/node/bootstrap/bootstrap_server.hpp |  3 +++
 2 files changed, 23 insertions(+), 5 deletions(-)

diff --git a/nano/node/bootstrap/bootstrap_server.cpp b/nano/node/bootstrap/bootstrap_server.cpp
index 2fd6cb7b6b..aac4e8a2e3 100644
--- a/nano/node/bootstrap/bootstrap_server.cpp
+++ b/nano/node/bootstrap/bootstrap_server.cpp
@@ -16,7 +16,8 @@ nano::bootstrap_server::bootstrap_server (bootstrap_server_config const & config
 	store{ store_a },
 	ledger{ ledger_a },
 	network_constants{ network_constants_a },
-	stats{ stats_a }
+	stats{ stats_a },
+	limiter{ config.limiter, /* allow bursts */ 3.0 } // TODO: Limiter bucket capacity should be at least equal to the batch size, currently it's not configurable
 {
 	queue.max_size_query = [this] (auto const & origin) {
 		return config.max_queue;
@@ -181,6 +182,22 @@ void nano::bootstrap_server::run ()
 	nano::unique_lock<nano::mutex> lock{ mutex };
 	while (!stopped)
 	{
+		condition.wait (lock, [this] () {
+			return stopped || !queue.empty ();
+		});
+
+		// Rate limit the processing
+		while (!stopped && !limiter.should_pass (config.batch_size))
+		{
+			stats.inc (nano::stat::type::bootstrap_server, nano::stat::detail::cooldown);
+			condition.wait_for (lock, 100ms);
+		}
+
+		if (stopped)
+		{
+			return;
+		}
+
 		if (!queue.empty ())
 		{
 			stats.inc (nano::stat::type::bootstrap_server, nano::stat::detail::loop);
@@ -190,10 +207,6 @@ void nano::bootstrap_server::run ()
 
 			lock.lock ();
 		}
-		else
-		{
-			condition.wait (lock, [this] () { return stopped || !queue.empty (); });
-		}
 	}
 }
 
@@ -429,6 +442,7 @@ nano::error nano::bootstrap_server_config::serialize (nano::tomlconfig & toml) c
 	toml.put ("max_queue", max_queue, "Maximum number of queued requests per peer. \ntype:uint64");
 	toml.put ("threads", threads, "Number of threads to process requests. \ntype:uint64");
 	toml.put ("batch_size", batch_size, "Maximum number of requests to process in a single batch. \ntype:uint64");
+	toml.put ("limiter", limiter, "Rate limit for processing requests. Use 0 for unlimited. \ntype:uint64");
 
 	return toml.get_error ();
 }
@@ -438,6 +452,7 @@ nano::error nano::bootstrap_server_config::deserialize (nano::tomlconfig & toml)
 	toml.get ("max_queue", max_queue);
 	toml.get ("threads", threads);
 	toml.get ("batch_size", batch_size);
+	toml.get ("limiter", limiter);
 
 	return toml.get_error ();
 }
diff --git a/nano/node/bootstrap/bootstrap_server.hpp b/nano/node/bootstrap/bootstrap_server.hpp
index 4a211cafb7..597ee28545 100644
--- a/nano/node/bootstrap/bootstrap_server.hpp
+++ b/nano/node/bootstrap/bootstrap_server.hpp
@@ -2,6 +2,7 @@
 
 #include <nano/lib/locks.hpp>
 #include <nano/lib/observer_set.hpp>
+#include <nano/lib/rate_limiting.hpp>
 #include <nano/node/fair_queue.hpp>
 #include <nano/node/fwd.hpp>
 #include <nano/node/messages.hpp>
@@ -24,6 +25,7 @@ class bootstrap_server_config final
 	size_t max_queue{ 16 };
 	size_t threads{ 1 };
 	size_t batch_size{ 64 };
+	size_t limiter{ 500 };
 };
 
 /**
@@ -91,6 +93,7 @@ class bootstrap_server final
 
 private:
 	nano::fair_queue<request_t, nano::no_value> queue;
+	nano::rate_limiter limiter;
 
 	std::atomic<bool> stopped{ false };
 	nano::condition_variable condition;