Skip to main content

doido_controller/
stack.rs

1use crate::axum::{
2    body::Body,
3    extract::{Request, State},
4    middleware::{from_fn, from_fn_with_state, Next},
5    response::Response,
6    Router,
7};
8use crate::config::CorsConfig;
9use http::{header, HeaderValue, Method, StatusCode};
10use std::sync::Arc;
11use tower_http::{
12    catch_panic::CatchPanicLayer,
13    cors::{Any, CorsLayer},
14};
15
16/// A router transformation registered by the app to insert its own middleware
17/// relative to doido's always-on stack (Rails `config.middleware.insert_*`).
18type RouterTransform = Box<dyn FnOnce(Router) -> Router + Send>;
19
20pub struct MiddlewareStack {
21    cors: bool,
22    cors_config: Option<CorsConfig>,
23    allowed_hosts: Vec<String>,
24    csrf: bool,
25    force_ssl: bool,
26    before: Vec<RouterTransform>,
27    after: Vec<RouterTransform>,
28}
29
30impl MiddlewareStack {
31    pub fn new() -> Self {
32        Self {
33            cors: false,
34            cors_config: None,
35            allowed_hosts: Vec::new(),
36            csrf: false,
37            force_ssl: false,
38            before: Vec::new(),
39            after: Vec::new(),
40        }
41    }
42
43    /// Enable CSRF protection (double-submit cookie): state-changing requests
44    /// must send an `X-CSRF-Token` header matching the `csrf_token` cookie.
45    pub fn with_csrf(mut self) -> Self {
46        self.csrf = true;
47        self
48    }
49
50    /// Redirect insecure requests to HTTPS (Rails `force_ssl`). A request is
51    /// considered secure when its `X-Forwarded-Proto` is `https` (or the URI
52    /// scheme is `https`); otherwise it gets a 301 to the `https://` URL.
53    pub fn with_force_ssl(mut self) -> Self {
54        self.force_ssl = true;
55        self
56    }
57
58    /// Insert custom middleware **inside** the always-on layers (closer to the
59    /// router), i.e. it runs after logging/panic-recovery on the way in. The
60    /// closure receives the router and returns it with its `.layer(...)` added.
61    pub fn insert_before(
62        mut self,
63        transform: impl FnOnce(Router) -> Router + Send + 'static,
64    ) -> Self {
65        self.before.push(Box::new(transform));
66        self
67    }
68
69    /// Insert custom middleware **outside** the always-on layers (outermost), so
70    /// it wraps the whole stack.
71    pub fn insert_after(
72        mut self,
73        transform: impl FnOnce(Router) -> Router + Send + 'static,
74    ) -> Self {
75        self.after.push(Box::new(transform));
76        self
77    }
78
79    /// Enable permissive CORS (any origin/method/header). For fine-grained,
80    /// config-driven CORS use [`with_cors_config`](Self::with_cors_config).
81    pub fn with_cors(mut self) -> Self {
82        self.cors = true;
83        self
84    }
85
86    /// Enable CORS from parsed [`CorsConfig`] (spec 07 `[middleware.cors]`). A
87    /// config with `enabled: false` is ignored, keeping CORS opt-in.
88    pub fn with_cors_config(mut self, config: CorsConfig) -> Self {
89        self.cors_config = Some(config);
90        self
91    }
92
93    /// Restrict requests to an allowlist of `Host` header values (Rails
94    /// `config.hosts` / `ActionDispatch::HostAuthorization`). An empty list
95    /// permits any host; a request whose host is not listed gets a `403`.
96    pub fn with_allowed_hosts(mut self, hosts: Vec<String>) -> Self {
97        self.allowed_hosts = hosts;
98        self
99    }
100
101    pub fn apply(self, router: Router) -> Router {
102        // App-registered "before" middleware sits innermost (closest to routes).
103        let mut r = router;
104        for transform in self.before {
105            r = transform(r);
106        }
107        // Log every request and its response (method, path, status, latency)
108        // through doido's centralized logger. Added after `CatchPanicLayer` so
109        // it sits outermost and logs panic-recovered `500`s too.
110        r = r
111            .layer(CatchPanicLayer::new())
112            .layer(from_fn(crate::logging::log_requests));
113        match &self.cors_config {
114            Some(config) if config.enabled => r = r.layer(build_cors(config)),
115            _ if self.cors => r = r.layer(CorsLayer::permissive()),
116            _ => {}
117        }
118        if !self.allowed_hosts.is_empty() {
119            r = r.layer(from_fn_with_state(Arc::new(self.allowed_hosts), host_guard));
120        }
121        if self.csrf {
122            r = r.layer(from_fn(csrf_guard));
123        }
124        if self.force_ssl {
125            r = r.layer(from_fn(force_ssl_guard));
126        }
127        // App-registered "after" middleware wraps everything (outermost).
128        for transform in self.after {
129            r = transform(r);
130        }
131        r
132    }
133}
134
135/// Redirect insecure requests to their `https://` equivalent with a 301.
136async fn force_ssl_guard(request: Request, next: Next) -> Response {
137    let forwarded_https = request
138        .headers()
139        .get("x-forwarded-proto")
140        .and_then(|v| v.to_str().ok())
141        .map(|p| p.eq_ignore_ascii_case("https"))
142        .unwrap_or(false);
143    let scheme_https = request.uri().scheme_str() == Some("https");
144    if forwarded_https || scheme_https {
145        return next.run(request).await;
146    }
147
148    let host = request
149        .headers()
150        .get(header::HOST)
151        .and_then(|v| v.to_str().ok())
152        .unwrap_or("");
153    let path = request
154        .uri()
155        .path_and_query()
156        .map(|pq| pq.as_str())
157        .unwrap_or("/");
158    let location = format!("https://{host}{path}");
159
160    match HeaderValue::from_str(&location) {
161        Ok(value) => {
162            let mut response = Response::builder()
163                .status(StatusCode::MOVED_PERMANENTLY)
164                .body(Body::empty())
165                .expect("valid 301 response");
166            response.headers_mut().insert(header::LOCATION, value);
167            response
168        }
169        Err(_) => Response::builder()
170            .status(StatusCode::BAD_REQUEST)
171            .body(Body::empty())
172            .expect("valid 400 response"),
173    }
174}
175
176/// Enforce the double-submit CSRF check on state-changing methods. Safe methods
177/// (GET/HEAD/OPTIONS/TRACE) always pass; others require the `csrf_token` cookie
178/// and the `X-CSRF-Token` header to be present and equal.
179async fn csrf_guard(request: Request, next: Next) -> Response {
180    let method = request.method();
181    let is_safe = matches!(
182        *method,
183        Method::GET | Method::HEAD | Method::OPTIONS | Method::TRACE
184    );
185    if is_safe {
186        return next.run(request).await;
187    }
188
189    let cookie_token = request
190        .headers()
191        .get(header::COOKIE)
192        .and_then(|c| c.to_str().ok())
193        .and_then(crate::csrf::token_from_cookie_header);
194    let header_token = request
195        .headers()
196        .get("x-csrf-token")
197        .and_then(|h| h.to_str().ok())
198        .map(str::to_string);
199
200    match (cookie_token, header_token) {
201        (Some(cookie), Some(header)) if crate::csrf::tokens_match(&cookie, &header) => {
202            next.run(request).await
203        }
204        _ => Response::builder()
205            .status(StatusCode::FORBIDDEN)
206            .body(Body::from("CSRF token mismatch"))
207            .expect("valid 403 response"),
208    }
209}
210
211/// Reject requests whose `Host` header (port stripped) is not in the allowlist.
212async fn host_guard(
213    State(allowed): State<Arc<Vec<String>>>,
214    request: Request,
215    next: Next,
216) -> Response {
217    let host = request
218        .headers()
219        .get(header::HOST)
220        .and_then(|h| h.to_str().ok())
221        .map(|h| h.split(':').next().unwrap_or(h).to_string());
222    match host {
223        Some(h) if allowed.iter().any(|a| a == &h) => next.run(request).await,
224        _ => Response::builder()
225            .status(StatusCode::FORBIDDEN)
226            .body(Body::from("Forbidden host"))
227            .expect("valid 403 response"),
228    }
229}
230
231/// Build a [`CorsLayer`] from configuration. `"*"` in `allowed_origins` maps to
232/// "any origin"; otherwise each origin/method is parsed and unparseable entries
233/// are skipped.
234fn build_cors(config: &CorsConfig) -> CorsLayer {
235    let mut layer = CorsLayer::new();
236    if config.allowed_origins.iter().any(|o| o == "*") {
237        layer = layer.allow_origin(Any);
238    } else {
239        let origins: Vec<HeaderValue> = config
240            .allowed_origins
241            .iter()
242            .filter_map(|o| o.parse().ok())
243            .collect();
244        if !origins.is_empty() {
245            layer = layer.allow_origin(origins);
246        }
247    }
248    let methods: Vec<Method> = config
249        .allowed_methods
250        .iter()
251        .filter_map(|m| Method::from_bytes(m.as_bytes()).ok())
252        .collect();
253    if !methods.is_empty() {
254        layer = layer.allow_methods(methods);
255    }
256    layer
257}
258
259impl Default for MiddlewareStack {
260    fn default() -> Self {
261        Self::new()
262    }
263}