use crate::server::{Request, Response};
use crate::error::Result;
use async_trait::async_trait;
use http::{HeaderName, HeaderValue, StatusCode};
use std::sync::Arc;
#[async_trait]
pub trait Middleware: Send + Sync {
async fn handle(&self, request: Request, next: Next) -> Result<Response>;
}
pub struct Next {
middleware: Vec<Arc<dyn Middleware>>,
index: usize,
handler: Option<Arc<dyn Handler>>,
}
impl Next {
pub fn new(middleware: Vec<Arc<dyn Middleware>>, handler: Arc<dyn Handler>) -> Self {
Self {
middleware,
index: 0,
handler: Some(handler),
}
}
pub async fn run(mut self, request: Request) -> Result<Response> {
if self.index < self.middleware.len() {
let middleware = self.middleware[self.index].clone();
self.index += 1;
middleware.handle(request, self).await
} else if let Some(handler) = &self.handler {
handler.handle(request).await
} else {
Ok(Response::text("Not Found").with_status(StatusCode::NOT_FOUND))
}
}
}
#[async_trait]
pub trait Handler: Send + Sync {
async fn handle(&self, request: Request) -> Result<Response>;
}
pub struct Logger;
impl Logger {
pub fn new() -> Self {
Self
}
}
#[async_trait]
impl Middleware for Logger {
async fn handle(&self, request: Request, next: Next) -> Result<Response> {
let method = request.method.clone();
let path = request.path.clone();
let start = std::time::Instant::now();
println!("[REQUEST] {} {}", method, path);
let response = next.run(request).await?;
let duration = start.elapsed();
println!(
"[RESPONSE] {} {} - {} ({:?})",
method, path, response.status(), duration
);
Ok(response)
}
}
impl Default for Logger {
fn default() -> Self {
Self::new()
}
}
pub struct Cors {
allow_origin: String,
allow_methods: Vec<String>,
allow_headers: Vec<String>,
max_age: u32,
}
impl Cors {
pub fn permissive() -> Self {
Self {
allow_origin: "*".to_string(),
allow_methods: vec![
"GET".to_string(),
"POST".to_string(),
"PUT".to_string(),
"DELETE".to_string(),
"PATCH".to_string(),
"OPTIONS".to_string(),
],
allow_headers: vec![
"Content-Type".to_string(),
"Authorization".to_string(),
"X-Requested-With".to_string(),
],
max_age: 86400,
}
}
pub fn new(allow_origin: impl Into<String>) -> Self {
Self {
allow_origin: allow_origin.into(),
allow_methods: vec!["GET".to_string(), "POST".to_string()],
allow_headers: vec!["Content-Type".to_string()],
max_age: 3600,
}
}
pub fn allow_methods(mut self, methods: Vec<String>) -> Self {
self.allow_methods = methods;
self
}
pub fn allow_headers(mut self, headers: Vec<String>) -> Self {
self.allow_headers = headers;
self
}
pub fn max_age(mut self, seconds: u32) -> Self {
self.max_age = seconds;
self
}
}
#[async_trait]
impl Middleware for Cors {
async fn handle(&self, request: Request, next: Next) -> Result<Response> {
if request.method.as_str() == "OPTIONS" {
let response = Response::text("")
.with_header(
HeaderName::from_static("access-control-allow-origin"),
HeaderValue::from_str(&self.allow_origin).unwrap(),
)
.with_header(
HeaderName::from_static("access-control-allow-methods"),
HeaderValue::from_str(&self.allow_methods.join(", ")).unwrap(),
)
.with_header(
HeaderName::from_static("access-control-allow-headers"),
HeaderValue::from_str(&self.allow_headers.join(", ")).unwrap(),
)
.with_header(
HeaderName::from_static("access-control-max-age"),
HeaderValue::from_str(&self.max_age.to_string()).unwrap(),
)
.with_status(StatusCode::NO_CONTENT);
return Ok(response);
}
let response = next.run(request).await?;
let response = response.with_header(
HeaderName::from_static("access-control-allow-origin"),
HeaderValue::from_str(&self.allow_origin).unwrap(),
);
Ok(response)
}
}
#[allow(dead_code)]
pub struct RateLimit {
max_requests: usize,
window_secs: u64,
}
impl RateLimit {
pub fn new(max_requests: usize, window_secs: u64) -> Self {
Self {
max_requests,
window_secs,
}
}
}
#[async_trait]
impl Middleware for RateLimit {
async fn handle(&self, request: Request, next: Next) -> Result<Response> {
next.run(request).await
}
}
pub struct Auth {
token: String,
}
impl Auth {
pub fn bearer(token: impl Into<String>) -> Self {
Self {
token: token.into(),
}
}
}
#[async_trait]
impl Middleware for Auth {
async fn handle(&self, request: Request, next: Next) -> Result<Response> {
if let Some(auth_header) = request.headers.get("authorization") {
if let Ok(auth_str) = auth_header.to_str() {
if auth_str.starts_with("Bearer ") {
let token = &auth_str[7..];
if token == self.token {
return next.run(request).await;
}
}
}
}
Ok(Response::text("Unauthorized")
.with_status(StatusCode::UNAUTHORIZED)
.with_header(
HeaderName::from_static("www-authenticate"),
HeaderValue::from_static("Bearer"),
))
}
}
#[cfg(test)]
mod tests {
use super::*;
use http::Method;
struct TestHandler;
#[async_trait]
impl Handler for TestHandler {
async fn handle(&self, _request: Request) -> Result<Response> {
Ok(Response::text("Hello"))
}
}
#[tokio::test]
async fn test_logger_middleware() {
let logger = Logger::new();
let handler = Arc::new(TestHandler);
let request = Request {
method: Method::GET,
path: "/test".to_string(),
headers: http::HeaderMap::new(),
body: bytes::Bytes::new(),
};
let next = Next::new(vec![], handler);
let response = logger.handle(request, next).await.unwrap();
assert_eq!(response.status(), StatusCode::OK);
}
#[tokio::test]
async fn test_cors_middleware() {
let cors = Cors::permissive();
let handler = Arc::new(TestHandler);
let request = Request {
method: Method::GET,
path: "/test".to_string(),
headers: http::HeaderMap::new(),
body: bytes::Bytes::new(),
};
let next = Next::new(vec![], handler);
let response = cors.handle(request, next).await.unwrap();
assert!(response.headers().contains_key("access-control-allow-origin"));
}
#[tokio::test]
async fn test_cors_preflight() {
let cors = Cors::permissive();
let handler = Arc::new(TestHandler);
let request = Request {
method: Method::OPTIONS,
path: "/test".to_string(),
headers: http::HeaderMap::new(),
body: bytes::Bytes::new(),
};
let next = Next::new(vec![], handler);
let response = cors.handle(request, next).await.unwrap();
assert_eq!(response.status(), StatusCode::NO_CONTENT);
assert!(response.headers().contains_key("access-control-allow-methods"));
}
#[tokio::test]
async fn test_auth_middleware_success() {
let auth = Auth::bearer("secret-token");
let handler = Arc::new(TestHandler);
let mut headers = http::HeaderMap::new();
headers.insert(
"authorization",
http::HeaderValue::from_static("Bearer secret-token"),
);
let request = Request {
method: Method::GET,
path: "/protected".to_string(),
headers,
body: bytes::Bytes::new(),
};
let next = Next::new(vec![], handler);
let response = auth.handle(request, next).await.unwrap();
assert_eq!(response.status(), StatusCode::OK);
}
#[tokio::test]
async fn test_auth_middleware_failure() {
let auth = Auth::bearer("secret-token");
let handler = Arc::new(TestHandler);
let request = Request {
method: Method::GET,
path: "/protected".to_string(),
headers: http::HeaderMap::new(),
body: bytes::Bytes::new(),
};
let next = Next::new(vec![], handler);
let response = auth.handle(request, next).await.unwrap();
assert_eq!(response.status(), StatusCode::UNAUTHORIZED);
}
}