Skip to main content

ullun/
middleware.rs

1//! Middleware types and traits
2
3use crate::error::{Error, Result};
4use crate::request::Request;
5use crate::response::Response;
6use async_trait::async_trait;
7use std::sync::Arc;
8
9/// Middleware trait
10#[async_trait]
11pub trait Middleware: Send + Sync + 'static {
12    async fn handle(&self, req: Request, next: Next) -> Result<Response>;
13}
14
15/// Next middleware in chain
16#[derive(Clone)]
17pub struct Next {
18    middlewares: Arc<Vec<Arc<dyn Middleware>>>,
19    index: usize,
20}
21
22impl Next {
23    pub fn new(middlewares: Arc<Vec<Arc<dyn Middleware>>>) -> Self {
24        Self {
25            middlewares,
26            index: 0,
27        }
28    }
29
30    pub async fn run(mut self, req: Request) -> Result<Response> {
31        if self.index < self.middlewares.len() {
32            let middleware = self.middlewares[self.index].clone();
33            self.index += 1;
34            middleware.handle(req, self).await
35        } else {
36            // End of chain without handler match
37            Err(Error::not_found("Route not found"))
38        }
39    }
40}
41
42/// Implement Middleware for async functions
43#[async_trait]
44impl<F, Fut> Middleware for F
45where
46    F: Fn(Request, Next) -> Fut + Send + Sync + 'static,
47    Fut: std::future::Future<Output = Result<Response>> + Send + 'static,
48{
49    async fn handle(&self, req: Request, next: Next) -> Result<Response> {
50        (self)(req, next).await
51    }
52}
53
54/// Logger middleware
55pub async fn logger(req: Request, next: Next) -> Result<Response> {
56    let method = req.method.clone();
57    let uri = req.uri.clone();
58    let start = std::time::Instant::now();
59
60    let response = next.run(req).await;
61
62    let duration = start.elapsed();
63    println!("{} {} - {:?}", method, uri, duration);
64
65    response
66}
67
68/// CORS middleware with preflight support
69pub fn cors(
70    allowed_origins: Vec<String>,
71) -> impl Fn(
72    Request,
73    Next,
74) -> std::pin::Pin<Box<dyn std::future::Future<Output = Result<Response>> + Send>>
75       + Clone {
76    move |req: Request, next: Next| {
77        let origins = allowed_origins.clone();
78        Box::pin(async move {
79            // Handle preflight OPTIONS requests
80            if req.method == http::Method::OPTIONS {
81                let mut response = Response::new(http::StatusCode::NO_CONTENT);
82
83                if !origins.is_empty() {
84                    let origin = if origins.contains(&"*".to_string()) {
85                        "*".to_string()
86                    } else {
87                        // In production, should check Origin header against allowed list
88                        origins[0].clone()
89                    };
90                    response = response.with_header("Access-Control-Allow-Origin", origin);
91                    response = response.with_header(
92                        "Access-Control-Allow-Methods",
93                        "GET, POST, PUT, DELETE, PATCH, OPTIONS",
94                    );
95                    response = response.with_header(
96                        "Access-Control-Allow-Headers",
97                        "Content-Type, Authorization",
98                    );
99                    response = response.with_header("Access-Control-Max-Age", "86400");
100                }
101
102                return Ok(response);
103            }
104
105            // Regular request - add CORS headers to response
106            let mut response = next.run(req).await?;
107
108            if !origins.is_empty() {
109                let origin = if origins.contains(&"*".to_string()) {
110                    "*".to_string()
111                } else {
112                    origins[0].clone()
113                };
114                response = response.with_header("Access-Control-Allow-Origin", origin);
115                response = response.with_header("Access-Control-Allow-Credentials", "true");
116            }
117
118            Ok(response)
119        })
120    }
121}