#![allow(non_snake_case)]
use super::RiGatewayRequest;
use crate::core::RiResult;
use async_trait::async_trait;
use std::sync::Arc;
#[async_trait]
pub trait RiMiddleware: Send + Sync {
async fn execute(&self, request: &mut RiGatewayRequest) -> RiResult<()>;
fn name(&self) -> &'static str;
}
pub struct RiMiddlewareChain {
middlewares: Vec<Arc<dyn RiMiddleware>>,
}
impl Default for RiMiddlewareChain {
fn default() -> Self {
Self::new()
}
}
impl RiMiddlewareChain {
pub fn new() -> Self {
Self {
middlewares: Vec::new(),
}
}
pub fn add(&mut self, middleware: Arc<dyn RiMiddleware>) {
self.middlewares.push(middleware);
}
pub async fn execute(&self, request: &mut RiGatewayRequest) -> RiResult<()> {
for middleware in &self.middlewares {
let start = std::time::Instant::now();
middleware.execute(request).await?;
let duration = start.elapsed();
let _duration_ms = duration.as_secs_f64() * 1000.0;
}
Ok(())
}
pub fn clear(&mut self) {
self.middlewares.clear();
}
pub fn len(&self) -> usize {
self.middlewares.len()
}
pub fn is_empty(&self) -> bool {
self.middlewares.is_empty()
}
}
pub struct RiAuthMiddleware {
auth_header: String,
}
impl RiAuthMiddleware {
pub fn new(auth_header: String) -> Self {
Self { auth_header }
}
}
#[async_trait]
impl RiMiddleware for RiAuthMiddleware {
async fn execute(&self, request: &mut RiGatewayRequest) -> RiResult<()> {
if let Some(auth_header) = request.headers.get(&self.auth_header) {
if let Some(token) = auth_header.strip_prefix("Bearer ") {
if token.is_empty() {
return Err(crate::core::RiError::Other("Empty bearer token".to_string()));
}
} else {
return Err(crate::core::RiError::Other("Invalid authorization header format".to_string()));
}
} else {
}
Ok(())
}
fn name(&self) -> &'static str {
"AuthMiddleware"
}
}
pub struct RiCorsMiddleware {
allowed_origins: Vec<String>,
#[allow(dead_code)]
allowed_methods: Vec<String>,
#[allow(dead_code)]
allowed_headers: Vec<String>,
}
impl RiCorsMiddleware {
pub fn new(
allowed_origins: Vec<String>,
allowed_methods: Vec<String>,
allowed_headers: Vec<String>,
) -> Self {
Self {
allowed_origins,
allowed_methods,
allowed_headers,
}
}
fn is_origin_allowed(&self, origin: &str) -> bool {
if self.allowed_origins.contains(&"*".to_string()) {
return false;
}
self.allowed_origins.iter().any(|allowed| allowed == origin)
}
fn is_wildcard_allowed(&self) -> bool {
self.allowed_origins.contains(&"*".to_string())
}
}
#[async_trait]
impl RiMiddleware for RiCorsMiddleware {
async fn execute(&self, request: &mut RiGatewayRequest) -> RiResult<()> {
if let Some(origin) = request.headers.get("origin") {
if !self.is_wildcard_allowed() && !self.is_origin_allowed(origin) {
return Err(crate::core::RiError::Other("Origin not allowed".to_string()));
}
}
Ok(())
}
fn name(&self) -> &'static str {
"CorsMiddleware"
}
}
pub struct RiLoggingMiddleware {
#[allow(dead_code)]
log_level: String,
}
impl RiLoggingMiddleware {
pub fn new(log_level: String) -> Self {
Self { log_level }
}
}
#[async_trait]
impl RiMiddleware for RiLoggingMiddleware {
async fn execute(&self, request: &mut RiGatewayRequest) -> RiResult<()> {
log::info!("[{}] {} {} from {}",
chrono::Utc::now().format("%Y-%m-%d %H:%M:%S"),
request.method,
request.path,
request.remote_addr
);
Ok(())
}
fn name(&self) -> &'static str {
"LoggingMiddleware"
}
}
pub struct RiRequestIdMiddleware;
impl Default for RiRequestIdMiddleware {
fn default() -> Self {
Self::new()
}
}
impl RiRequestIdMiddleware {
pub fn new() -> Self {
Self
}
}
#[async_trait]
impl RiMiddleware for RiRequestIdMiddleware {
async fn execute(&self, _request: &mut RiGatewayRequest) -> RiResult<()> {
Ok(())
}
fn name(&self) -> &'static str {
"RequestIdMiddleware"
}
}
pub struct RiRateLimitMiddleware {
rate_limiter: Arc<crate::gateway::RiRateLimiter>,
}
impl RiRateLimitMiddleware {
pub fn new(rate_limiter: Arc<crate::gateway::RiRateLimiter>) -> Self {
Self {
rate_limiter,
}
}
}
#[async_trait]
impl RiMiddleware for RiRateLimitMiddleware {
async fn execute(&self, request: &mut RiGatewayRequest) -> RiResult<()> {
if !self.rate_limiter.check_request(request).await {
return Err(crate::core::RiError::Other("Rate limit exceeded".to_string()));
}
Ok(())
}
fn name(&self) -> &'static str {
"RateLimitMiddleware"
}
}