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::{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.allowed_origins.iter().any(|o| o == "*") {
252 layer = layer.allow_origin(Any);
253 } else {
254 let origins: Vec<HeaderValue> = config
255 .allowed_origins
256 .iter()
257 .filter_map(|o| o.parse().ok())
258 .collect();
259 if !origins.is_empty() {
260 layer = layer.allow_origin(origins);
261 }
262 }
263 let methods: Vec<Method> = config
264 .allowed_methods
265 .iter()
266 .filter_map(|m| Method::from_bytes(m.as_bytes()).ok())
267 .collect();
268 if !methods.is_empty() {
269 layer = layer.allow_methods(methods);
270 }
271 layer
272}
273
274impl Default for MiddlewareStack {
275 fn default() -> Self {
276 Self::new()
277 }
278}