feat: pre-check for ip/connection rate limit

src/extensions/rate_limit/connection.rs

@@ -1,4 +1,4 @@
-use crate::{extensions::rate_limit::MethodWeights, utils::errors};
+use crate::extensions::rate_limit::MethodWeights;
 use futures::{future::BoxFuture, FutureExt};
 use governor::{DefaultDirectRateLimiter, Jitter, RateLimiter};
 use jsonrpsee::{
@@ -75,9 +75,10 @@ where
 
         async move {
             if let Some(n) = NonZeroU32::new(weight) {
-                if limiter.until_n_ready_with_jitter(n, jitter).await.is_err() {
-                    return MethodResponse::error(req.id, errors::failed("rate limit exceeded"));
-                }
+                limiter
+                    .until_n_ready_with_jitter(n, jitter)
+                    .await
+                    .expect("check_n have been done during init");
             }
             service.call(req).await
         }

src/extensions/rate_limit/ip.rs

@@ -1,4 +1,4 @@
-use crate::{extensions::rate_limit::MethodWeights, utils::errors};
+use crate::extensions::rate_limit::MethodWeights;
 use futures::{future::BoxFuture, FutureExt};
 use governor::{DefaultKeyedRateLimiter, Jitter};
 use jsonrpsee::{
@@ -86,13 +86,10 @@ where
         let weight = self.method_weights.get(req.method_name());
         async move {
             if let Some(n) = NonZeroU32::new(weight) {
-                if limiter
+                limiter
                     .until_key_n_ready_with_jitter(&ip_addr, n, jitter)
                     .await
-                    .is_err()
-                {
-                    return MethodResponse::error(req.id, errors::failed("rate limit exceeded"));
-                }
+                    .expect("check_n have been done during init");
             }
             service.call(req).await
         }

src/extensions/rate_limit/mod.rs

@@ -1,3 +1,4 @@
+use anyhow::bail;
 use governor::{DefaultKeyedRateLimiter, Jitter, Quota, RateLimiter};
 use serde::Deserialize;
 use std::num::NonZeroU32;
@@ -90,6 +91,45 @@ impl RateLimitBuilder {
             }
         }
     }
+
+    pub fn pre_check_connection(&self, method_weights: &MethodWeights) -> anyhow::Result<()> {
+        if let Some(ref rule) = self.config.connection {
+            let burst = NonZeroU32::new(rule.burst).unwrap();
+            let period = Duration::from_secs(rule.period_secs);
+            let quota = build_quota(burst, period);
+            let limiter = RateLimiter::direct(quota);
+
+            for (method, weight) in method_weights.0 {
+                if let Some(n) = NonZeroU32::new(*weight) {
+                    if limiter.check_n(n).is_err() {
+                        bail!("`{method}` weight config too big for connection rate limit: {}", n);
+                    }
+                }
+            }
+        }
+
+        Ok(())
+    }
+
+    pub fn pre_check_ip(&self, method_weights: &MethodWeights) -> anyhow::Result<()> {
+        if let Some(ref rule) = self.config.ip {
+            let burst = NonZeroU32::new(rule.burst).unwrap();
+            let period = Duration::from_secs(rule.period_secs);
+            let quota = build_quota(burst, period);
+            let limiter = RateLimiter::direct(quota);
+
+            for (method, weight) in &method_weights.0 {
+                if let Some(n) = NonZeroU32::new(*weight) {
+                    if limiter.check_n(n).is_err() {
+                        bail!("`{method}` weight config too big for ip rate limit: {}", n);
+                    }
+                }
+            }
+        }
+
+        Ok(())
+    }
+
     pub fn connection_limit(&self, method_weights: MethodWeights) -> Option<ConnectionRateLimitLayer> {
         if let Some(ref rule) = self.config.connection {
             let burst = NonZeroU32::new(rule.burst).unwrap();

src/extensions/rate_limit/weight.rs

@@ -1,25 +1,23 @@
 use crate::config::RpcMethod;
 use std::collections::BTreeMap;
+use std::sync::Arc;
 
 #[derive(Clone, Debug, Default)]
-pub struct MethodWeights(BTreeMap<String, u32>);
+pub struct MethodWeights(Arc<BTreeMap<String, u32>>);
 
 impl MethodWeights {
-    pub fn add(&mut self, method: &str, weight: u32) {
-        self.0.insert(method.to_owned(), weight);
-    }
-
     pub fn get(&self, method: &str) -> u32 {
         self.0.get(method).cloned().unwrap_or(1)
     }
 }
 
 impl MethodWeights {
     pub fn from_config(methods: &[RpcMethod]) -> Self {
-        let mut weights = MethodWeights::default();
+        let mut weights = BTreeMap::default();
         for method in methods {
-            weights.add(&method.method, method.rate_limit_weight);
+            weights.insert(method.method.to_owned(), method.rate_limit_weight);
         }
-        weights
+
+        Self(Arc::new(weights))
     }
 }

src/server.rs

@@ -51,6 +51,12 @@ pub async fn build(config: Config) -> anyhow::Result<SubwayServerHandle> {
 
     let rpc_method_weights = MethodWeights::from_config(&config.rpcs.methods);
 
+    // pre-check stage
+    if let Some(r) = &rate_limit_builder {
+        r.pre_check_ip(&rpc_method_weights)?;
+        r.pre_check_connection(&rpc_method_weights)?;
+    }
+
     let request_timeout_seconds = server_builder.config.request_timeout_seconds;
 
     let metrics = get_rpc_metrics(&extensions_registry).await;