doido_controller/
stack.rs1use 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::{AllowHeaders, AllowMethods, AllowOrigin, Any, CorsLayer},
14};
15
16type 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 api_only: bool,
27 before: Vec<RouterTransform>,
28 after: Vec<RouterTransform>,
29}
30
31impl MiddlewareStack {
32 pub fn new() -> Self {
33 Self {
34 cors: false,
35 cors_config: None,
36 allowed_hosts: Vec::new(),
37 csrf: false,
38 force_ssl: false,
39 api_only: false,
40 before: Vec::new(),
41 after: Vec::new(),
42 }
43 }
44
45 pub fn with_api_only(mut self, api_only: bool) -> Self {
52 self.api_only = api_only;
53 self
54 }
55
56 pub fn with_csrf(mut self) -> Self {
59 self.csrf = true;
60 self
61 }
62
63 pub fn with_force_ssl(mut self) -> Self {
67 self.force_ssl = true;
68 self
69 }
70
71 pub fn insert_before(
75 mut self,
76 transform: impl FnOnce(Router) -> Router + Send + 'static,
77 ) -> Self {
78 self.before.push(Box::new(transform));
79 self
80 }
81
82 pub fn insert_after(
85 mut self,
86 transform: impl FnOnce(Router) -> Router + Send + 'static,
87 ) -> Self {
88 self.after.push(Box::new(transform));
89 self
90 }
91
92 pub fn with_cors(mut self) -> Self {
95 self.cors = true;
96 self
97 }
98
99 pub fn with_cors_config(mut self, config: CorsConfig) -> Self {
102 self.cors_config = Some(config);
103 self
104 }
105
106 pub fn with_allowed_hosts(mut self, hosts: Vec<String>) -> Self {
110 self.allowed_hosts = hosts;
111 self
112 }
113
114 pub fn apply(self, router: Router) -> Router {
115 let mut r = router;
117 for transform in self.before {
118 r = transform(r);
119 }
120 r = r
124 .layer(CatchPanicLayer::new())
125 .layer(from_fn(crate::logging::log_requests));
126 match &self.cors_config {
127 Some(config) if config.enabled => r = r.layer(build_cors(config)),
128 _ if self.cors => r = r.layer(CorsLayer::permissive()),
129 _ => {}
130 }
131 if !self.allowed_hosts.is_empty() {
132 r = r.layer(from_fn_with_state(Arc::new(self.allowed_hosts), host_guard));
133 }
134 if self.csrf && !self.api_only {
137 r = r.layer(from_fn(csrf_guard));
138 }
139 if self.force_ssl {
140 r = r.layer(from_fn(force_ssl_guard));
141 }
142 for transform in self.after {
144 r = transform(r);
145 }
146 r
147 }
148}
149
150async fn force_ssl_guard(request: Request, next: Next) -> Response {
152 let forwarded_https = request
153 .headers()
154 .get("x-forwarded-proto")
155 .and_then(|v| v.to_str().ok())
156 .map(|p| p.eq_ignore_ascii_case("https"))
157 .unwrap_or(false);
158 let scheme_https = request.uri().scheme_str() == Some("https");
159 if forwarded_https || scheme_https {
160 return next.run(request).await;
161 }
162
163 let host = request
164 .headers()
165 .get(header::HOST)
166 .and_then(|v| v.to_str().ok())
167 .unwrap_or("");
168 let path = request
169 .uri()
170 .path_and_query()
171 .map(|pq| pq.as_str())
172 .unwrap_or("/");
173 let location = format!("https://{host}{path}");
174
175 match HeaderValue::from_str(&location) {
176 Ok(value) => {
177 let mut response = Response::builder()
178 .status(StatusCode::MOVED_PERMANENTLY)
179 .body(Body::empty())
180 .expect("valid 301 response");
181 response.headers_mut().insert(header::LOCATION, value);
182 response
183 }
184 Err(_) => Response::builder()
185 .status(StatusCode::BAD_REQUEST)
186 .body(Body::empty())
187 .expect("valid 400 response"),
188 }
189}
190
191async fn csrf_guard(request: Request, next: Next) -> Response {
195 let method = request.method();
196 let is_safe = matches!(
197 *method,
198 Method::GET | Method::HEAD | Method::OPTIONS | Method::TRACE
199 );
200 if is_safe {
201 return next.run(request).await;
202 }
203
204 let cookie_token = request
205 .headers()
206 .get(header::COOKIE)
207 .and_then(|c| c.to_str().ok())
208 .and_then(crate::csrf::token_from_cookie_header);
209 let header_token = request
210 .headers()
211 .get("x-csrf-token")
212 .and_then(|h| h.to_str().ok())
213 .map(str::to_string);
214
215 match (cookie_token, header_token) {
216 (Some(cookie), Some(header)) if crate::csrf::tokens_match(&cookie, &header) => {
217 next.run(request).await
218 }
219 _ => Response::builder()
220 .status(StatusCode::FORBIDDEN)
221 .body(Body::from("CSRF token mismatch"))
222 .expect("valid 403 response"),
223 }
224}
225
226async fn host_guard(
228 State(allowed): State<Arc<Vec<String>>>,
229 request: Request,
230 next: Next,
231) -> Response {
232 let host = request
233 .headers()
234 .get(header::HOST)
235 .and_then(|h| h.to_str().ok())
236 .map(|h| h.split(':').next().unwrap_or(h).to_string());
237 match host {
238 Some(h) if allowed.iter().any(|a| a == &h) => next.run(request).await,
239 _ => Response::builder()
240 .status(StatusCode::FORBIDDEN)
241 .body(Body::from("Forbidden host"))
242 .expect("valid 403 response"),
243 }
244}
245
246fn build_cors(config: &CorsConfig) -> CorsLayer {
250 let mut layer = CorsLayer::new();
251 if config.allow_credentials {
252 layer = layer.allow_credentials(true);
253 }
254 if let Some(origin) = configured_allow_origin(config) {
255 layer = layer.allow_origin(origin);
256 }
257 if let Some(methods) = configured_allow_methods(config) {
258 layer = layer.allow_methods(methods);
259 }
260 if let Some(headers) = configured_allow_headers(config) {
261 layer = layer.allow_headers(headers);
262 }
263 layer
264}
265
266fn configured_allow_origin(config: &CorsConfig) -> Option<AllowOrigin> {
267 if config.allowed_origins.is_empty() {
268 return None;
269 }
270 if config.allowed_origins.iter().any(|o| o == "*") {
271 return Some(if config.allow_credentials {
272 AllowOrigin::mirror_request()
273 } else {
274 Any.into()
275 });
276 }
277 let origins: Vec<HeaderValue> = config
278 .allowed_origins
279 .iter()
280 .filter_map(|o| o.parse().ok())
281 .collect();
282 (!origins.is_empty()).then(|| AllowOrigin::list(origins))
283}
284
285fn configured_allow_methods(config: &CorsConfig) -> Option<AllowMethods> {
286 if config.allowed_methods.is_empty() {
287 return None;
288 }
289 if config.allowed_methods.iter().any(|m| m == "*") {
290 return Some(if config.allow_credentials {
291 AllowMethods::mirror_request()
292 } else {
293 Any.into()
294 });
295 }
296 let methods: Vec<Method> = config
297 .allowed_methods
298 .iter()
299 .filter_map(|m| Method::from_bytes(m.as_bytes()).ok())
300 .collect();
301 (!methods.is_empty()).then(|| AllowMethods::list(methods))
302}
303
304fn configured_allow_headers(config: &CorsConfig) -> Option<AllowHeaders> {
305 if config.allowed_headers.is_empty() {
306 return None;
307 }
308 if config.allowed_headers.iter().any(|h| h == "*") {
309 return Some(if config.allow_credentials {
310 AllowHeaders::mirror_request()
311 } else {
312 Any.into()
313 });
314 }
315 let headers: Vec<http::HeaderName> = config
316 .allowed_headers
317 .iter()
318 .filter_map(|h| h.parse().ok())
319 .collect();
320 (!headers.is_empty()).then(|| AllowHeaders::list(headers))
321}
322
323impl Default for MiddlewareStack {
324 fn default() -> Self {
325 Self::new()
326 }
327}