rpc server: add rate limiting middleware (#3301)

Add RPC server rate limiting which can be utilized by the CLI
`--rpc-rate-limit <calls/per minute>`

Resolves first part of
https://github.com/paritytech/polkadot-sdk/issues/3028

//cc @PierreBesson @kogeler  you might be interested in this one

---------

Co-authored-by: James Wilson <james@jsdw.me>
Co-authored-by: Xiliang Chen <xlchen1291@gmail.com>

substrate/client/rpc-servers/src/lib.rs

@@ -22,7 +22,9 @@
 
 pub mod middleware;
 
-use std::{convert::Infallible, error::Error as StdError, net::SocketAddr, time::Duration};
+use std::{
+	convert::Infallible, error::Error as StdError, net::SocketAddr, num::NonZeroU32, time::Duration,
+};
 
 use http::header::HeaderValue;
 use hyper::{
@@ -31,10 +33,7 @@ use hyper::{
 };
 use jsonrpsee::{
 	server::{
-		middleware::{
-			http::{HostFilterLayer, ProxyGetRequestLayer},
-			rpc::RpcServiceBuilder,
-		},
+		middleware::http::{HostFilterLayer, ProxyGetRequestLayer},
 		stop_channel, ws, PingConfig, StopHandle, TowerServiceBuilder,
 	},
 	Methods, RpcModule,
@@ -43,11 +42,14 @@ use tokio::net::TcpListener;
 use tower::Service;
 use tower_http::cors::{AllowOrigin, CorsLayer};
 
-pub use jsonrpsee::core::{
-	id_providers::{RandomIntegerIdProvider, RandomStringIdProvider},
-	traits::IdProvider,
+pub use jsonrpsee::{
+	core::{
+		id_providers::{RandomIntegerIdProvider, RandomStringIdProvider},
+		traits::IdProvider,
+	},
+	server::middleware::rpc::RpcServiceBuilder,
 };
-pub use middleware::{MetricsLayer, RpcMetrics};
+pub use middleware::{MetricsLayer, RateLimitLayer, RpcMetrics};
 
 const MEGABYTE: u32 = 1024 * 1024;
 
@@ -79,12 +81,26 @@ pub struct Config<'a, M: Send + Sync + 'static> {
 	pub id_provider: Option<Box<dyn IdProvider>>,
 	/// Tokio runtime handle.
 	pub tokio_handle: tokio::runtime::Handle,
+	/// Rate limit calls per minute.
+	pub rate_limit: Option<NonZeroU32>,
+}
+
+#[derive(Debug, Clone)]
+struct PerConnection<RpcMiddleware, HttpMiddleware> {
+	methods: Methods,
+	stop_handle: StopHandle,
+	metrics: Option<RpcMetrics>,
+	tokio_handle: tokio::runtime::Handle,
+	service_builder: TowerServiceBuilder<RpcMiddleware, HttpMiddleware>,
 }
 
 /// Start RPC server listening on given address.
-pub async fn start_server<M: Send + Sync + 'static>(
+pub async fn start_server<M>(
 	config: Config<'_, M>,
-) -> Result<Server, Box<dyn StdError + Send + Sync>> {
+) -> Result<Server, Box<dyn StdError + Send + Sync>>
+where
+	M: Send + Sync,
+{
 	let Config {
 		addrs,
 		cors,
@@ -97,6 +113,7 @@ pub async fn start_server<M: Send + Sync + 'static>(
 		id_provider,
 		tokio_handle,
 		rpc_api,
+		rate_limit,
 	} = config;
 
 	let std_listener = TcpListener::bind(addrs.as_slice()).await?.into_std()?;
@@ -153,7 +170,13 @@ pub async fn start_server<M: Send + Sync + 'static>(
 				let transport_label = if is_websocket { "ws" } else { "http" };
 
 				let metrics = metrics.map(|m| MetricsLayer::new(m, transport_label));
-				let rpc_middleware = RpcServiceBuilder::new().option_layer(metrics.clone());
+				let rate_limit = rate_limit.map(|r| RateLimitLayer::per_minute(r));
+
+				// NOTE: The metrics needs to run first to include rate-limited calls in the
+				// metrics.
+				let rpc_middleware =
+					RpcServiceBuilder::new().option_layer(metrics.clone()).option_layer(rate_limit);
+
 				let mut svc =
 					service_builder.set_rpc_middleware(rpc_middleware).build(methods, stop_handle);
 
@@ -245,12 +268,3 @@ fn format_cors(maybe_cors: Option<&Vec<String>>) -> String {
 		format!("{:?}", ["*"])
 	}
 }
-
-#[derive(Clone)]
-struct PerConnection<RpcMiddleware, HttpMiddleware> {
-	methods: Methods,
-	stop_handle: StopHandle,
-	metrics: Option<RpcMetrics>,
-	tokio_handle: tokio::runtime::Handle,
-	service_builder: TowerServiceBuilder<RpcMiddleware, HttpMiddleware>,
-}

substrate/client/rpc-servers/src/middleware/mod.rs

@@ -18,6 +18,10 @@
 
 //! JSON-RPC specific middleware.
 
+/// Grafana metrics middleware.
 pub mod metrics;
+/// Rate limit middleware.
+pub mod rate_limit;
 
 pub use metrics::*;
+pub use rate_limit::*;

substrate/client/rpc-servers/src/middleware/rate_limit.rs

@@ -0,0 +1,107 @@
+// This file is part of Substrate.
+
+// Copyright (C) Parity Technologies (UK) Ltd.
+// SPDX-License-Identifier: GPL-3.0-or-later WITH Classpath-exception-2.0
+
+// This program is free software: you can redistribute it and/or modify
+// it under the terms of the GNU General Public License as published by
+// the Free Software Foundation, either version 3 of the License, or
+// (at your option) any later version.
+
+// This program is distributed in the hope that it will be useful,
+// but WITHOUT ANY WARRANTY; without even the implied warranty of
+// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
+// GNU General Public License for more details.
+
+// You should have received a copy of the GNU General Public License
+// along with this program. If not, see <https://www.gnu.org/licenses/>.
+
+//! RPC rate limiting middleware.
+
+use std::{num::NonZeroU32, sync::Arc, time::Duration};
+
+use futures::future::{BoxFuture, FutureExt};
+use governor::{
+	clock::{Clock, DefaultClock, QuantaClock},
+	middleware::NoOpMiddleware,
+	state::{InMemoryState, NotKeyed},
+	Jitter,
+};
+use jsonrpsee::{
+	server::middleware::rpc::RpcServiceT,
+	types::{ErrorObject, Id, Request},
+	MethodResponse,
+};
+
+type RateLimitInner = governor::RateLimiter<NotKeyed, InMemoryState, DefaultClock, NoOpMiddleware>;
+
+const MAX_JITTER: Duration = Duration::from_millis(50);
+const MAX_RETRIES: usize = 10;
+
+/// JSON-RPC rate limit middleware layer.
+#[derive(Debug, Clone)]
+pub struct RateLimitLayer(governor::Quota);
+
+impl RateLimitLayer {
+	/// Create new rate limit enforced per minute.
+	pub fn per_minute(n: NonZeroU32) -> Self {
+		Self(governor::Quota::per_minute(n))
+	}
+}
+
+/// JSON-RPC rate limit middleware
+pub struct RateLimit<S> {
+	service: S,
+	rate_limit: Arc<RateLimitInner>,
+	clock: QuantaClock,
+}
+
+impl<S> tower::Layer<S> for RateLimitLayer {
+	type Service = RateLimit<S>;
+
+	fn layer(&self, service: S) -> Self::Service {
+		let clock = QuantaClock::default();
+		RateLimit {
+			service,
+			rate_limit: Arc::new(RateLimitInner::direct_with_clock(self.0, &clock)),
+			clock,
+		}
+	}
+}
+
+impl<'a, S> RpcServiceT<'a> for RateLimit<S>
+where
+	S: Send + Sync + RpcServiceT<'a> + Clone + 'static,
+{
+	type Future = BoxFuture<'a, MethodResponse>;
+
+	fn call(&self, req: Request<'a>) -> Self::Future {
+		let service = self.service.clone();
+		let rate_limit = self.rate_limit.clone();
+		let clock = self.clock.clone();
+
+		async move {
+			let mut attempts = 0;
+			let jitter = Jitter::up_to(MAX_JITTER);
+
+			loop {
+				if attempts >= MAX_RETRIES {
+					break reject_too_many_calls(req.id);
+				}
+
+				if let Err(rejected) = rate_limit.check() {
+					tokio::time::sleep(jitter + rejected.wait_time_from(clock.now())).await;
+				} else {
+					break service.call(req).await;
+				}
+
+				attempts += 1;
+			}
+		}
+		.boxed()
+	}
+}
+
+fn reject_too_many_calls(id: Id) -> MethodResponse {
+	MethodResponse::error(id, ErrorObject::owned(-32999, "RPC rate limit exceeded", None::<()>))
+}