diff --git a/config/src/rpc_config.rs b/config/src/rpc_config.rs index efef55d3b4..8fcddfba44 100644 --- a/config/src/rpc_config.rs +++ b/config/src/rpc_config.rs @@ -336,6 +336,15 @@ pub struct ApiQuotaConfiguration { value_parser = parse_key_val::, )] pub custom_user_api_quota: Option>, + + #[serde(skip_serializing_if = "Option::is_none")] + #[clap( + name = "jsonrpc-ratelimit-ip-whitelist", + long, + help = "IPs that bypass rate limiting (requires --http-trust-forwarded-ip-headers to identify callers); eg: 1.2.3.4,5.6.7.8", + use_value_delimiter = true + )] + pub ip_whitelist: Option>, } impl ApiQuotaConfiguration { @@ -365,6 +374,10 @@ impl ApiQuotaConfiguration { self.custom_user_api_quota.clone().unwrap_or_default() } + pub fn ip_whitelist(&self) -> Vec { + self.ip_whitelist.clone().unwrap_or_default() + } + pub fn merge(&mut self, o: &Self) -> Result<()> { if o.default_global_api_quota.is_some() { self.default_global_api_quota = o.default_global_api_quota.clone(); @@ -379,6 +392,9 @@ impl ApiQuotaConfiguration { if o.custom_user_api_quota.is_some() { self.custom_user_api_quota = o.custom_user_api_quota.clone(); } + if o.ip_whitelist.is_some() { + self.ip_whitelist = o.ip_whitelist.clone(); + } Ok(()) } } diff --git a/rpc/server/src/metadata_middleware.rs b/rpc/server/src/metadata_middleware.rs index 8e8ed9799d..64e5548338 100644 --- a/rpc/server/src/metadata_middleware.rs +++ b/rpc/server/src/metadata_middleware.rs @@ -2,6 +2,7 @@ // SPDX-License-Identifier: Apache-2.0 use jsonrpsee::server::HttpRequest; +use starcoin_logger::prelude::*; use starcoin_rpc_api::metadata::Metadata; use std::future::Future; use std::net::IpAddr; @@ -63,6 +64,10 @@ where .trust_forwarded_ip_headers .then(|| extract_user_from_request(&request, &self.ip_headers)) .flatten(); + debug!( + "HTTP RPC metadata: user={:?}, trust_forwarded={}", + user, self.trust_forwarded_ip_headers + ); request.extensions_mut().insert(Metadata { user }); let fut = self.service.call(request); diff --git a/rpc/server/src/rate_limit_middleware.rs b/rpc/server/src/rate_limit_middleware.rs index 5499646af5..77a3c174ca 100644 --- a/rpc/server/src/rate_limit_middleware.rs +++ b/rpc/server/src/rate_limit_middleware.rs @@ -11,7 +11,9 @@ use jsonrpsee::{ MethodResponse, }; use starcoin_config::{ApiQuotaConfig, ApiQuotaConfiguration, QuotaDuration}; +use starcoin_logger::prelude::*; use starcoin_rpc_api::metadata::Metadata; +use std::collections::HashSet; use std::future::Future; use std::sync::Arc; use tower::Layer; @@ -34,6 +36,7 @@ impl From for QuotaWrapper { #[derive(Clone, Debug)] pub struct JsonApiRateLimitLayer { limiters: Arc>, + ip_whitelist: Arc>, } impl JsonApiRateLimitLayer { @@ -52,8 +55,17 @@ impl JsonApiRateLimitLayer { .map(|(k, v)| (k, Into::::into(v).0)) .collect(), ); + let ip_whitelist: HashSet = quotas.ip_whitelist().into_iter().collect(); + info!( + "RPC rate limit middleware initialized with {} whitelisted IP(s)", + ip_whitelist.len() + ); + if !ip_whitelist.is_empty() { + trace!("Whitelisted IPs: {:?}", ip_whitelist); + } Self { limiters: Arc::new(limiters), + ip_whitelist: Arc::new(ip_whitelist), } } } @@ -65,6 +77,7 @@ impl Layer for JsonApiRateLimitLayer { JsonApiRateLimitMiddleware { service, limiters: self.limiters.clone(), + ip_whitelist: self.ip_whitelist.clone(), } } } @@ -73,6 +86,7 @@ impl Layer for JsonApiRateLimitLayer { pub struct JsonApiRateLimitMiddleware { service: S, limiters: Arc>, + ip_whitelist: Arc>, } fn user_from_extensions(extensions: &Extensions) -> Option { @@ -116,12 +130,22 @@ where let user = user_from_extensions(request.extensions()); let service = self.service.clone(); let limiters = self.limiters.clone(); + let ip_whitelist = self.ip_whitelist.clone(); async move { + if let Some(ref ip) = user { + if ip_whitelist.contains(ip) { + debug!("Whitelisted IP bypassing rate limit for method={}", method); + return service.call(request).await; + } + } match limiters.check(&method, user.as_ref()) { Ok(_) => service.call(request).await, - Err(e) => MethodResponse::error(request.id(), rate_limit_error(e)) - .with_extensions(request.extensions), + Err(e) => { + warn!("Rate limited: method={}, reason={}", method, e); + MethodResponse::error(request.id(), rate_limit_error(e)) + .with_extensions(request.extensions) + } } } } @@ -134,11 +158,22 @@ where let user = user_from_extensions(notification.extensions()); let service = self.service.clone(); let limiters = self.limiters.clone(); + let ip_whitelist = self.ip_whitelist.clone(); async move { + if let Some(ref ip) = user { + if ip_whitelist.contains(ip) { + debug!( + "Whitelisted IP bypassing rate limit for notification={}", + method + ); + return service.notification(notification).await; + } + } match limiters.check(&method, user.as_ref()) { Ok(_) => service.notification(notification).await, Err(e) => { + warn!("Rate limited: notification={}, reason={}", method, e); S::NotificationResponse::from_rate_limited(notification, rate_limit_error(e)) } } @@ -152,20 +187,44 @@ where Ok(BatchEntry::Call(req)) => { let method = req.method_name().to_owned(); let user = user_from_extensions(req.extensions()); - match self.limiters.check(&method, user.as_ref()) { - Ok(_) => entries.push(Ok(BatchEntry::Call(req))), - Err(e) => { - entries.push(Err(BatchEntryErr::new(req.id(), rate_limit_error(e)))) + let whitelisted = user + .as_ref() + .is_some_and(|ip| self.ip_whitelist.contains(ip)); + if whitelisted { + debug!( + "Whitelisted IP bypassing rate limit for batch call={}", + method + ); + entries.push(Ok(BatchEntry::Call(req))); + } else { + match self.limiters.check(&method, user.as_ref()) { + Ok(_) => entries.push(Ok(BatchEntry::Call(req))), + Err(e) => { + warn!("Rate limited: batch call={}, reason={}", method, e); + entries.push(Err(BatchEntryErr::new(req.id(), rate_limit_error(e)))) + } } } } Ok(BatchEntry::Notification(n)) => { let method = n.method_name().to_owned(); let user = user_from_extensions(n.extensions()); - match self.limiters.check(&method, user.as_ref()) { - Ok(_) => entries.push(Ok(BatchEntry::Notification(n))), - Err(e) => { - entries.push(Err(BatchEntryErr::new(Id::Null, rate_limit_error(e)))) + let whitelisted = user + .as_ref() + .is_some_and(|ip| self.ip_whitelist.contains(ip)); + if whitelisted { + debug!( + "Whitelisted IP bypassing rate limit for batch notification={}", + method + ); + entries.push(Ok(BatchEntry::Notification(n))); + } else { + match self.limiters.check(&method, user.as_ref()) { + Ok(_) => entries.push(Ok(BatchEntry::Notification(n))), + Err(e) => { + warn!("Rate limited: batch notification={}, reason={}", method, e); + entries.push(Err(BatchEntryErr::new(Id::Null, rate_limit_error(e)))) + } } } } @@ -244,6 +303,30 @@ mod tests { JsonApiRateLimitLayer::from_config(quotas).layer(service) } + fn test_middleware_with_whitelist( + method: &str, + whitelist: Vec, + ) -> JsonApiRateLimitMiddleware { + let service = ObserveService::default(); + let quotas = ApiQuotaConfiguration { + custom_global_api_quota: Some(vec![( + method.to_owned(), + ApiQuotaConfig::from_str("1/s").expect("valid quota"), + )]), + ip_whitelist: Some(whitelist), + ..Default::default() + }; + JsonApiRateLimitLayer::from_config(quotas).layer(service) + } + + fn request_with_user<'a>(method: &'a str, user: &str) -> Request<'a> { + let mut req = Request::borrowed(method, None, Id::Number(1)); + req.extensions_mut().insert(Metadata { + user: Some(user.to_string()), + }); + req + } + #[test] fn notification_should_be_rate_limited() { let middleware = test_middleware_for_method("state.get"); @@ -276,4 +359,92 @@ mod tests { let _ = futures::executor::block_on(middleware.batch(second_batch)); assert_eq!(middleware.service.batch_errors.load(Ordering::Relaxed), 1); } + + #[test] + fn whitelisted_ip_bypasses_rate_limit() { + let middleware = test_middleware_with_whitelist("state.get", vec!["10.0.0.1".to_string()]); + + // First call consumes the quota + let req1 = request_with_user("state.get", "10.0.0.1"); + let rsp1 = futures::executor::block_on(middleware.call(req1)); + assert_ne!(rsp1.as_error_code(), Some(-10000)); + + // Second call should still succeed because 10.0.0.1 is whitelisted + let req2 = request_with_user("state.get", "10.0.0.1"); + let rsp2 = futures::executor::block_on(middleware.call(req2)); + assert_ne!(rsp2.as_error_code(), Some(-10000)); + } + + #[test] + fn non_whitelisted_ip_still_rate_limited() { + let middleware = test_middleware_with_whitelist("state.get", vec!["10.0.0.1".to_string()]); + + // First call from non-whitelisted IP + let req1 = request_with_user("state.get", "192.168.1.1"); + let rsp1 = futures::executor::block_on(middleware.call(req1)); + assert_ne!(rsp1.as_error_code(), Some(-10000)); + + // Second call should be rate limited + let req2 = request_with_user("state.get", "192.168.1.1"); + let rsp2 = futures::executor::block_on(middleware.call(req2)); + assert_eq!(rsp2.as_error_code(), Some(-10000)); + } + + #[test] + fn whitelisted_ip_bypasses_notification_rate_limit() { + let middleware = test_middleware_with_whitelist("state.get", vec!["10.0.0.1".to_string()]); + + let mut n1 = Notification::new(Cow::Borrowed("state.get"), None); + n1.extensions_mut().insert(Metadata { + user: Some("10.0.0.1".to_string()), + }); + let mut n2 = Notification::new(Cow::Borrowed("state.get"), None); + n2.extensions_mut().insert(Metadata { + user: Some("10.0.0.1".to_string()), + }); + + let rsp1 = futures::executor::block_on(middleware.notification(n1)); + let rsp2 = futures::executor::block_on(middleware.notification(n2)); + + // Both should succeed (whitelisted) + assert!(rsp1.is_notification()); + assert!(rsp2.is_notification()); + assert_eq!(middleware.service.notifications.load(Ordering::Relaxed), 2); + } + + #[test] + fn whitelisted_ip_bypasses_batch_rate_limit() { + let middleware = test_middleware_with_whitelist("state.get", vec!["10.0.0.1".to_string()]); + + // First batch consumes the quota for non-whitelisted callers + let batch1 = Batch::from(vec![Ok(BatchEntry::Call({ + let mut req = Request::borrowed("state.get", None, Id::Number(1)); + req.extensions_mut().insert(Metadata { + user: Some("10.0.0.1".to_string()), + }); + req + }))]); + let _ = futures::executor::block_on(middleware.batch(batch1)); + assert_eq!(middleware.service.batch_errors.load(Ordering::Relaxed), 0); + + // Second batch should still succeed because 10.0.0.1 is whitelisted + let batch2 = Batch::from(vec![ + Ok(BatchEntry::Call({ + let mut req = Request::borrowed("state.get", None, Id::Number(2)); + req.extensions_mut().insert(Metadata { + user: Some("10.0.0.1".to_string()), + }); + req + })), + Ok(BatchEntry::Notification({ + let mut n = Notification::new(Cow::Borrowed("state.get"), None); + n.extensions_mut().insert(Metadata { + user: Some("10.0.0.1".to_string()), + }); + n + })), + ]); + let _ = futures::executor::block_on(middleware.batch(batch2)); + assert_eq!(middleware.service.batch_errors.load(Ordering::Relaxed), 0); + } }