From 0bb8571b02d7a9b67d290e00f2968deaef0af702 Mon Sep 17 00:00:00 2001
From: nutcas3 <mauricee423@gmail.com>
Date: Sat, 25 Apr 2026 06:01:42 +0300
Subject: [PATCH] feat: Add IP-based rate limiting to charging engine with
 configurable RPS and burst size

- Add rate_limit module to lib.rs exports
- Add create_rate_limiter function with SmartIpKeyExtractor for IP-based limiting
- Add RATE_LIMIT_RPS environment variable with 10 requests/second default
- Add RATE_LIMIT_BURST environment variable with 20 burst size default
- Add RateLimitResponse struct with TOO_MANY_REQUESTS status code
- Add GovernorLayer with per_second quota and burst_size configuration to router
---
 apps/charging-engine/src/lib.rs        |  1 +
 apps/charging-engine/src/rate_limit.rs | 50 ++++++++++++++++++++++++++
 apps/charging-engine/src/routes.rs     |  4 +++
 3 files changed, 55 insertions(+)
 create mode 100644 apps/charging-engine/src/rate_limit.rs

diff --git a/apps/charging-engine/src/lib.rs b/apps/charging-engine/src/lib.rs
index 6d9cfec..60c09c6 100644
--- a/apps/charging-engine/src/lib.rs
+++ b/apps/charging-engine/src/lib.rs
@@ -3,6 +3,7 @@ pub mod charging;
 pub mod circuit_breaker;
 pub mod errors;
 pub mod handlers;
+pub mod rate_limit;
 pub mod routes;
 pub mod models;
 pub mod monitoring;
diff --git a/apps/charging-engine/src/rate_limit.rs b/apps/charging-engine/src/rate_limit.rs
new file mode 100644
index 0000000..de9e6f1
--- /dev/null
+++ b/apps/charging-engine/src/rate_limit.rs
@@ -0,0 +1,50 @@
+use axum::{
+    extract::Request,
+    http::StatusCode,
+    response::{IntoResponse, Response},
+};
+use governor::{
+    clock::DefaultClock,
+    state::InMemoryState,
+    Jitter, Quota, RateLimiter,
+};
+use std::num::NonZeroU32;
+use std::sync::Arc;
+use tower_governor::{
+    governor::config::GovernorConfigBuilder,
+    key_extractor::SmartIpKeyExtractor,
+    GovernorLayer,
+};
+
+pub fn create_rate_limiter() -> GovernorLayer<SmartIpKeyExtractor> {
+    let requests_per_second: u32 = std::env::var("RATE_LIMIT_RPS")
+        .ok()
+        .and_then(|s| s.parse().ok())
+        .unwrap_or(10);
+    let burst_size: u32 = std::env::var("RATE_LIMIT_BURST")
+        .ok()
+        .and_then(|s| s.parse().ok())
+        .unwrap_or(20);
+
+    let quota = Quota::per_second(NonZeroU32::new(requests_per_second).unwrap())
+        .allow_burst(NonZeroU32::new(burst_size).unwrap());
+
+    let config = GovernorConfigBuilder::default()
+        .per_second(quota)
+        .burst_size(burst_size)
+        .finish()
+        .unwrap();
+
+    GovernorLayer {
+        config,
+        key_extractor: SmartIpKeyExtractor,
+    }
+}
+
+pub struct RateLimitResponse;
+
+impl IntoResponse for RateLimitResponse {
+    fn into_response(self) -> Response {
+        (StatusCode::TOO_MANY_REQUESTS, "Rate limit exceeded").into_response()
+    }
+}
diff --git a/apps/charging-engine/src/routes.rs b/apps/charging-engine/src/routes.rs
index 0fddd56..ab944e3 100644
--- a/apps/charging-engine/src/routes.rs
+++ b/apps/charging-engine/src/routes.rs
@@ -5,6 +5,7 @@ use axum::{
 use tower_http::cors::{Any, CorsLayer};
 
 use crate::auth::auth_middleware;
+use crate::rate_limit::create_rate_limiter;
 use crate::handlers::{
     add_credit, add_rating_plan, block_user, check_credit, deduct_credit, detailed_health_check,
     engine_start, engine_stop, engine_uptime, get_balance, get_error_stats, get_performance_metrics,
@@ -20,6 +21,8 @@ pub fn create_router(state: AppState) -> Router {
         .allow_methods(Any)
         .allow_headers(Any);
 
+    let rate_limiter = create_rate_limiter();
+
     Router::new()
         // Public routes (no auth required - for packet-gateway integration)
         .route("/v1/health", get(health_check))
@@ -51,6 +54,7 @@ pub fn create_router(state: AppState) -> Router {
         .route("/v1/usage/process", post(process_usage))
         .route("/v1/invoice/:imsi/:period", get(generate_invoice))
         .layer(cors)
+        .layer(rate_limiter)
         .route_layer(axum::middleware::from_fn_with_state(state.auth_config.clone(), auth_middleware))
         .with_state(state)
 }