doido_controller/
stack.rs1use crate::config::CorsConfig;
2use axum::{
3 body::Body,
4 extract::{Request, State},
5 middleware::{from_fn, from_fn_with_state, Next},
6 response::Response,
7 Router,
8};
9use http::{header, HeaderValue, Method, StatusCode};
10use std::sync::Arc;
11use tower_http::{
12 catch_panic::CatchPanicLayer,
13 cors::{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 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 pub fn with_csrf(mut self) -> Self {
46 self.csrf = true;
47 self
48 }
49
50 pub fn with_force_ssl(mut self) -> Self {
54 self.force_ssl = true;
55 self
56 }
57
58 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 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 pub fn with_cors(mut self) -> Self {
82 self.cors = true;
83 self
84 }
85
86 pub fn with_cors_config(mut self, config: CorsConfig) -> Self {
89 self.cors_config = Some(config);
90 self
91 }
92
93 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 let mut r = router;
104 for transform in self.before {
105 r = transform(r);
106 }
107 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 for transform in self.after {
129 r = transform(r);
130 }
131 r
132 }
133}
134
135async 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
176async 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
211async 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
231fn 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}