#![allow(non_snake_case)]
use crate::core::{RiModule, RiServiceContext};
use log;
use serde::{Deserialize, Serialize};
use std::collections::HashMap as FxHashMap;
use std::sync::Arc;
use tokio::sync::RwLock;
pub mod middleware;
pub mod routing;
pub mod radix_tree;
pub mod circuit_breaker;
pub mod load_balancer;
pub mod rate_limiter;
pub mod server;
pub use routing::{RiRoute, RiRouter, RiRouteHandler};
pub use radix_tree::{RiRadixTree, RadixNode, PathSegment, SegmentType, RouteMatch};
pub use middleware::{RiMiddleware, RiMiddlewareChain};
pub use load_balancer::{RiLoadBalancer, RiLoadBalancerStrategy, RiBackendServer, RiLoadBalancerServerStats};
pub use rate_limiter::{RiRateLimiter, RiRateLimitConfig, RiRateLimitStats, RiSlidingWindowRateLimiter};
pub use circuit_breaker::{RiCircuitBreaker, RiCircuitBreakerConfig, RiCircuitBreakerState, RiCircuitBreakerMetrics};
#[cfg(feature = "gateway")]
pub use server::{RiGatewayServer, load_tls_config};
#[cfg_attr(feature = "pyo3", pyo3::prelude::pyclass(get_all, set_all))]
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct RiGatewayConfig {
pub listen_address: String,
pub listen_port: u16,
pub max_connections: usize,
pub request_timeout_seconds: u64,
pub enable_rate_limiting: bool,
pub enable_circuit_breaker: bool,
pub enable_load_balancing: bool,
pub cors_enabled: bool,
pub cors_origins: Vec<String>,
pub cors_methods: Vec<String>,
pub cors_headers: Vec<String>,
pub enable_logging: bool,
pub log_level: String,
}
#[cfg(feature = "pyo3")]
#[pyo3::prelude::pymethods]
impl RiGatewayConfig {
#[new]
fn py_new() -> Self {
Self::default()
}
#[staticmethod]
fn py_new_with_address(listen_address: String, listen_port: u16) -> Self {
Self {
listen_address,
listen_port,
..Self::default()
}
}
}
impl RiGatewayConfig {
pub fn new() -> Self {
Self::default()
}
}
impl Default for RiGatewayConfig {
fn default() -> Self {
Self {
listen_address: "0.0.0.0".to_string(),
listen_port: 8080,
max_connections: 10000,
request_timeout_seconds: 30,
enable_rate_limiting: true,
enable_circuit_breaker: true,
enable_load_balancing: true,
cors_enabled: true,
cors_origins: vec![],
cors_methods: vec!["GET".to_string(), "POST".to_string(), "PUT".to_string(), "DELETE".to_string(), "OPTIONS".to_string()],
cors_headers: vec!["Content-Type".to_string(), "Authorization".to_string(), "X-Requested-With".to_string()],
enable_logging: true,
log_level: "info".to_string(),
}
}
}
#[derive(Debug, Clone)]
#[cfg_attr(feature = "pyo3", pyo3::prelude::pyclass)]
pub struct RiGatewayRequest {
pub id: String,
pub method: String,
pub path: String,
pub headers: FxHashMap<String, String>,
pub query_params: FxHashMap<String, String>,
pub body: Option<Vec<u8>>,
pub remote_addr: String,
pub timestamp: std::time::Instant,
}
impl RiGatewayRequest {
pub fn new(
method: String,
path: String,
headers: FxHashMap<String, String>,
query_params: FxHashMap<String, String>,
body: Option<Vec<u8>>,
remote_addr: String,
) -> Self {
Self {
id: uuid::Uuid::new_v4().to_string(),
method,
path,
headers,
query_params,
body,
remote_addr,
timestamp: std::time::Instant::now(),
}
}
}
#[derive(Debug, Clone)]
#[cfg_attr(feature = "pyo3", pyo3::prelude::pyclass)]
pub struct RiGatewayResponse {
pub status_code: u16,
pub headers: FxHashMap<String, String>,
pub body: Vec<u8>,
pub request_id: String,
}
impl RiGatewayResponse {
pub fn new(status_code: u16, body: Vec<u8>, request_id: String) -> Self {
let mut headers = FxHashMap::default();
headers.insert("Content-Type".to_string(), "application/json".to_string());
headers.insert("X-Request-ID".to_string(), request_id.clone());
Self {
status_code,
headers,
body,
request_id,
}
}
pub fn with_header(mut self, key: String, value: String) -> Self {
self.headers.insert(key, value);
self
}
pub fn json<T: serde::Serialize>(status_code: u16, data: &T, request_id: String) -> crate::core::RiResult<Self> {
let body = serde_json::to_vec(data)?;
Ok(Self::new(status_code, body, request_id))
}
pub fn error(status_code: u16, message: String, request_id: String) -> Self {
let error_body = serde_json::json!({
"error": message,
"request_id": request_id
});
let body = serde_json::to_vec(&error_body).unwrap_or_else(|_| b"{}".to_vec());
Self::new(status_code, body, request_id)
}
}
#[cfg_attr(feature = "pyo3", pyo3::prelude::pyclass)]
pub struct RiGateway {
config: RwLock<RiGatewayConfig>,
router: Arc<RiRouter>,
middleware_chain: Arc<RiMiddlewareChain>,
rate_limiter: Option<Arc<RiRateLimiter>>,
circuit_breaker: Option<Arc<RiCircuitBreaker>>,
}
impl Default for RiGateway {
fn default() -> Self {
Self::new()
}
}
impl RiGateway {
pub fn new() -> Self {
let config = RiGatewayConfig::default();
let router = Arc::new(RiRouter::new());
let middleware_chain = Arc::new(RiMiddlewareChain::new());
let rate_limiter = if config.enable_rate_limiting {
Some(Arc::new(RiRateLimiter::new(RiRateLimitConfig::default())))
} else {
None
};
let circuit_breaker = if config.enable_circuit_breaker {
Some(Arc::new(RiCircuitBreaker::new(RiCircuitBreakerConfig::default())))
} else {
None
};
Self {
config: RwLock::new(config),
router,
middleware_chain,
rate_limiter,
circuit_breaker,
}
}
pub fn router(&self) -> Arc<RiRouter> {
self.router.clone()
}
pub fn middleware_chain(&self) -> Arc<RiMiddlewareChain> {
self.middleware_chain.clone()
}
pub async fn handle_request(&self, request: RiGatewayRequest) -> RiGatewayResponse {
let request_id = request.id.clone();
if let Some(rate_limiter) = &self.rate_limiter {
if !rate_limiter.check_request(&request).await {
return RiGatewayResponse::new(429, "Rate limit exceeded".to_string().into_bytes(), request_id);
}
}
if let Some(circuit_breaker) = &self.circuit_breaker {
if !circuit_breaker.allow_request() {
return RiGatewayResponse::new(503, "Service temporarily unavailable".to_string().into_bytes(), request_id);
}
}
let mut request = request;
match self.middleware_chain.execute(&mut request).await {
Ok(()) => {
match self.router.route(&request).await {
Ok(route_handler) => {
match route_handler(request).await {
Ok(response) => response,
Err(e) => {
log::error!("[Ri.Gateway] Internal server error for request {}: {}", request_id, e);
RiGatewayResponse::new(500, "Internal Server Error".to_string().into_bytes(), request_id)
}
}
},
Err(e) => {
log::warn!("[Ri.Gateway] Route not found for request {}: {}", request_id, e);
RiGatewayResponse::new(404, "Not Found".to_string().into_bytes(), request_id)
}
}
},
Err(e) => {
log::warn!("[Ri.Gateway] Middleware error for request {}: {}", request_id, e);
RiGatewayResponse::new(403, "Forbidden".to_string().into_bytes(), request_id)
}
}
}
}
#[async_trait::async_trait]
impl RiModule for RiGateway {
fn name(&self) -> &str {
"Ri.Gateway"
}
async fn init(&mut self, ctx: &mut RiServiceContext) -> crate::core::RiResult<()> {
let logger = ctx.logger();
logger.info("Ri.Gateway", "Initializing API gateway module")?;
let config = self.config.read().await;
logger.info(
"Ri.Gateway",
format!("Gateway will listen on {}:{}", config.listen_address, config.listen_port)
)?;
logger.info("Ri.Gateway", "API gateway module initialized successfully")?;
Ok(())
}
async fn after_shutdown(&mut self, _ctx: &mut RiServiceContext) -> crate::core::RiResult<()> {
log::info!("Cleaning up Ri Gateway Module");
if let Some(rate_limiter) = &self.rate_limiter {
rate_limiter.clear_all_buckets();
log::info!("Rate limiter cleanup completed");
}
if let Some(circuit_breaker) = &self.circuit_breaker {
circuit_breaker.reset();
log::info!("Circuit breaker reset completed");
}
self.router.clear_routes();
log::info!("Router cleanup completed");
log::info!("Ri Gateway Module cleanup completed");
Ok(())
}
}
#[cfg(feature = "pyo3")]
#[pyo3::prelude::pymethods]
impl RiGateway {
#[new]
fn py_new() -> Self {
Self::new()
}
}