huiali/rust-skills/rust-middleware

use actix_web::{
    dev::{forward_ready, Service, ServiceRequest, ServiceResponse, Transform},
    Error, HttpMessage,
};
use futures::future::{ready, LocalBoxFuture, Ready};
use std::rc::Rc;
use uuid::Uuid;
 
#[derive(Debug, Clone)]
pub struct RequestId(pub String);
 
pub struct RequestTracking;
 
impl<S, B> Transform<S, ServiceRequest> for RequestTracking
where
    S: Service<ServiceRequest, Response = ServiceResponse<B>, Error = Error> + 'static,
    S::Future: 'static,
    B: 'static,
{
    type Response = ServiceResponse<B>;
    type Error = Error;
    type InitError = ();
    type Transform = RequestTrackingMiddleware<S>;
    type Future = Ready<Result<Self::Transform, Self::InitError>>;
 
    fn new_transform(&self, service: S) -> Self::Future {
        ready(Ok(RequestTrackingMiddleware {
            service: Rc::new(service),
        }))
    }
}
 
pub struct RequestTrackingMiddleware<S> {
    service: Rc<S>,
}
 
impl<S, B> Service<ServiceRequest> for RequestTrackingMiddleware<S>
where
    S: Service<ServiceRequest, Response = ServiceResponse<B>, Error = Error> + 'static,
    S::Future: 'static,
    B: 'static,
{
    type Response = ServiceResponse<B>;
    type Error = Error;
    type Future = LocalBoxFuture<'static, Result<Self::Response, Self::Error>>;
 
    forward_ready!(service);
 
    fn call(&self, req: ServiceRequest) -> Self::Future {
        let service = Rc::clone(&self.service);
 
        Box::pin(async move {
            // Extract or generate request ID
            let request_id = req
                .headers()
                .get("X-Request-ID")
                .or_else(|| req.headers().get("Request-ID"))
                .and_then(|v| v.to_str().ok())
                .map(|s| s.to_string())
                .unwrap_or_else(|| Uuid::new_v4().to_string());
 
            // Store in request extensions
            req.extensions_mut().insert(RequestId(request_id.clone()));
 
            let start_time = std::time::Instant::now();
            log::info!("Request: {} {} - ID: {}", req.method(), req.path(), request_id);
 
            // Call next middleware
            let mut res = service.call(req).await?;
 
            let duration = start_time.elapsed();
            log::info!(
                "Response: {} - Duration: {:?} - ID: {}",
                res.status().as_u16(),
                duration,
                request_id
            );
 
            // Add request ID to response headers
            res.headers_mut().insert(
                actix_web::http::header::HeaderName::from_static("x-request-id"),
                actix_web::http::header::HeaderValue::from_str(&request_id).unwrap(),
            );
 
            Ok(res)
        })
    }
}
 
// Usage
use actix_web::{web, App, HttpServer};
 
#[actix_web::main]
async fn main() -> std::io::Result<()> {
    HttpServer::new(|| {
        App::new()
            .wrap(RequestTracking)  // Add middleware
            .route("/api", web::get().to(handler))
    })
    .bind("127.0.0.1:8080")?
    .run()
    .await
}

use actix_cors::Cors;
use actix_web::http::header;
 
pub struct CorsBuilder;
 
impl CorsBuilder {
    /// Production CORS with strict origin validation
    pub fn production(allowed_origins: &[&str]) -> Cors {
        Cors::default()
            .allowed_origin_fn(move |origin, _req_head| {
                let origin_str = origin.to_str().unwrap_or("");
                let allowed = allowed_origins.iter().any(|&o| o == origin_str);
                log::debug!("CORS: origin={} allowed={}", origin_str, allowed);
                allowed
            })
            .allowed_methods(vec!["OPTIONS", "HEAD", "GET", "POST", "PUT", "PATCH", "DELETE"])
            .allowed_headers(vec![
                header::AUTHORIZATION,
                header::ACCEPT,
                header::CONTENT_TYPE,
            ])
            .expose_headers(vec!["x-request-id"])
            .supports_credentials()
            .max_age(3600)
    }
 
    /// Development CORS allowing any origin
    pub fn development() -> Cors {
        Cors::default()
            .allow_any_origin()
            .allow_any_method()
            .allow_any_header()
            .supports_credentials()
    }
 
    /// API CORS for strict origin control
    pub fn api(strict_origins: &[&str]) -> Cors {
        Cors::default()
            .allowed_origins(strict_origins)
            .allowed_methods(vec!["GET", "POST", "PUT", "DELETE", "PATCH"])
            .allowed_headers(vec![header::AUTHORIZATION, header::CONTENT_TYPE])
            .max_age(86400)
    }
}
 
// Usage
HttpServer::new(|| {
    App::new()
        .wrap(CorsBuilder::production(&["https://example.com"]))
        .service(api_routes())
})

use actix_web::{
    dev::{forward_ready, Service, ServiceRequest, ServiceResponse, Transform},
    Error, HttpResponse,
};
use futures::future::{ready, LocalBoxFuture, Ready};
use std::collections::HashMap;
use std::sync::Arc;
use std::time::{Duration, Instant};
use tokio::sync::RwLock;
 
#[derive(Debug, Clone)]
pub struct RateLimitConfig {
    pub requests_per_second: u64,
    pub window_seconds: u64,
    pub whitelist: Vec<String>,
}
 
struct SlidingWindow {
    count: u64,
    window_start: Instant,
}
 
pub struct RateLimiting {
    config: RateLimitConfig,
    counters: Arc<RwLock<HashMap<String, SlidingWindow>>>,
}
 
impl RateLimiting {
    pub fn new(config: RateLimitConfig) -> Self {
        Self {
            config,
            counters: Arc::new(RwLock::new(HashMap::new())),
        }
    }
 
    async fn is_rate_limited(&self, key: &str) -> bool {
        let mut counters = self.counters.write().await;
        let now = Instant::now();
 
        // Clean up expired windows
        counters.retain(|_, w| {
            now.duration_since(w.window_start) < Duration::from_secs(self.config.window_seconds)
        });
 
        // Check or create window for this key
        match counters.get_mut(key) {
            Some(window) => {
                if now.duration_since(window.window_start)
                    < Duration::from_secs(self.config.window_seconds)
                {
                    window.count += 1;
                    window.count > self.config.requests_per_second
                } else {
                    // Reset window
                    window.count = 1;
                    window.window_start = now;
                    false
                }
            }
            None => {
                // New window
                counters.insert(
                    key.to_string(),
                    SlidingWindow {
                        count: 1,
                        window_start: now,
                    },
                );
                false
            }
        }
    }
}
 
impl<S, B> Transform<S, ServiceRequest> for RateLimiting
where
    S: Service<ServiceRequest, Response = ServiceResponse<B>, Error = Error> + 'static,
    S::Future: 'static,
    B: 'static,
{
    type Response = ServiceResponse<B>;
    type Error = Error;
    type InitError = ();
    type Transform = RateLimitingMiddleware<S>;
    type Future = Ready<Result<Self::Transform, Self::InitError>>;
 
    fn new_transform(&self, service: S) -> Self::Future {
        ready(Ok(RateLimitingMiddleware {
            service,
            config: self.config.clone(),
            counters: self.counters.clone(),
        }))
    }
}
 
pub struct RateLimitingMiddleware<S> {
    service: S,
    config: RateLimitConfig,
    counters: Arc<RwLock<HashMap<String, SlidingWindow>>>,
}
 
impl<S, B> Service<ServiceRequest> for RateLimitingMiddleware<S>
where
    S: Service<ServiceRequest, Response = ServiceResponse<B>, Error = Error> + 'static,
    S::Future: 'static,
    B: 'static,
{
    type Response = ServiceResponse<B>;
    type Error = Error;
    type Future = LocalBoxFuture<'static, Result<Self::Response, Self::Error>>;
 
    forward_ready!(service);
 
    fn call(&self, req: ServiceRequest) -> Self::Future {
        let client_ip = req
            .connection_info()
            .peer_addr()
            .map(|s| s.to_string())
            .unwrap_or_else(|| "unknown".to_string());
 
        // Skip rate limiting for whitelisted IPs
        if self.config.whitelist.contains(&client_ip) {
            return Box::pin(self.service.call(req));
        }
 
        let counters = self.counters.clone();
        let config = self.config.clone();
        let service = self.service.clone();
 
        Box::pin(async move {
            let limiter = RateLimiting {
                config,
                counters,
            };
 
            if limiter.is_rate_limited(&client_ip).await {
                let response = HttpResponse::TooManyRequests().json(serde_json::json!({
                    "error": "RATE_LIMIT_EXCEEDED",
                    "message": "Too many requests"
                }));
                return Ok(ServiceResponse::new(req.into_parts().0, response));
            }
 
            service.call(req).await
        })
    }
}
 
// Usage
let rate_limit = RateLimiting::new(RateLimitConfig {
    requests_per_second: 100,
    window_seconds: 60,
    whitelist: vec!["127.0.0.1".to_string()],
});
 
HttpServer::new(move || {
    App::new()
        .wrap(rate_limit.clone())
        .service(api_routes())
})

use axum::{
    extract::Request,
    http::StatusCode,
    middleware::Next,
    response::{IntoResponse, Response},
};
 
#[derive(Clone)]
pub struct Claims {
    pub sub: String,
    pub exp: usize,
}
 
pub async fn auth_middleware(
    mut req: Request,
    next: Next,
) -> Result<Response, StatusCode> {
    let auth_header = req
        .headers()
        .get("Authorization")
        .and_then(|h| h.to_str().ok())
        .ok_or(StatusCode::UNAUTHORIZED)?;
 
    let token = auth_header
        .strip_prefix("Bearer ")
        .ok_or(StatusCode::UNAUTHORIZED)?;
 
    // Verify JWT token
    let claims = verify_jwt(token)
        .map_err(|_| StatusCode::UNAUTHORIZED)?;
 
    // Store claims in request extensions
    req.extensions_mut().insert(claims);
 
    Ok(next.run(req).await)
}
 
fn verify_jwt(token: &str) -> Result<Claims, jsonwebtoken::errors::Error> {
    // JWT verification logic
    todo!()
}
 
// Usage with Axum
use axum::{middleware, routing::get, Router};
 
async fn protected_handler(
    Extension(claims): Extension<Claims>,
) -> impl IntoResponse {
    format!("Hello, user {}", claims.sub)
}
 
let app = Router::new()
    .route("/protected", get(protected_handler))
    .layer(middleware::from_fn(auth_middleware));

use actix_web::middleware::Compress;
 
// Actix compression middleware
HttpServer::new(|| {
    App::new()
        .wrap(Compress::default())  // Enables gzip, brotli, deflate
        .service(api_routes())
})
 
// Axum compression
use tower_http::compression::CompressionLayer;
 
let app = Router::new()
    .route("/", get(handler))
    .layer(CompressionLayer::new());  // Enables gzip, br, deflate, zstd

Request Flow:
  ↓
┌─────────────────────┐
│   Request Tracking  │  1. Generate request ID
├─────────────────────┤
│   CORS              │  2. Handle preflight
├─────────────────────┤
│   Compression       │  3. Compress response
├─────────────────────┤
│   Rate Limiting     │  4. Check rate limits
├─────────────────────┤
│   Authentication    │  5. Verify credentials
├─────────────────────┤
│   Handler           │  6. Business logic
└─────────────────────┘
  ↓
Response