Skip to main content

aster_forge_middleware/axum/
cors.rs

1//! Dynamic CORS using the same policy/validation/header engine as Actix.
2pub use crate::shared::cors::{
3    CorsAllowedOrigins, CorsMiddlewareError, CorsMiddlewareErrorKind, RuntimeCorsPolicy,
4};
5use crate::shared::cors::{CorsDecision, CorsSettings};
6use axum::{
7    extract::{Request, State},
8    middleware::Next,
9    response::{IntoResponse, Response},
10};
11use http::{HeaderValue, StatusCode, header};
12use std::sync::Arc;
13
14type PolicyResolver = dyn Fn(&Request) -> Result<RuntimeCorsPolicy, Box<Response>> + Send + Sync;
15type ExemptPredicate = dyn Fn(&str) -> bool + Send + Sync;
16type ErrorMapper = dyn Fn(CorsMiddlewareError) -> Response + Send + Sync;
17type OriginResolver = dyn Fn(&Request) -> Result<Option<String>, Box<Response>> + Send + Sync;
18
19/// State for `from_fn_with_state(config, runtime_cors)`. No CORS policy is enabled implicitly.
20#[derive(Clone)]
21pub struct RuntimeCorsConfig {
22    settings: CorsSettings,
23    policy: Arc<PolicyResolver>,
24    exempt_path: Arc<ExemptPredicate>,
25    map_error: Arc<ErrorMapper>,
26    request_origin: Arc<OriginResolver>,
27}
28
29impl RuntimeCorsConfig {
30    /// Supplies dynamic product policy, exempt paths, and classified error mapping.
31    pub fn new<P, X, M>(policy: P, exempt_path: X, map_error: M) -> Self
32    where
33        P: Fn(&Request) -> Result<RuntimeCorsPolicy, Box<Response>> + Send + Sync + 'static,
34        X: Fn(&str) -> bool + Send + Sync + 'static,
35        M: Fn(CorsMiddlewareError) -> Response + Send + Sync + 'static,
36    {
37        Self {
38            settings: CorsSettings::default(),
39            policy: Arc::new(policy),
40            exempt_path: Arc::new(exempt_path),
41            map_error: Arc::new(map_error),
42            request_origin: Arc::new(|_| Ok(None)),
43        }
44    }
45
46    /// Provides a normalized, product-trusted request origin for same-origin bypass.
47    /// The default `None` disables bypass; raw Host/Forwarded headers are never trusted implicitly.
48    #[must_use]
49    pub fn request_origin<R>(mut self, resolver: R) -> Self
50    where
51        R: Fn(&Request) -> Result<Option<String>, Box<Response>> + Send + Sync + 'static,
52    {
53        self.request_origin = Arc::new(resolver);
54        self
55    }
56
57    /// Sets preflight methods.
58    #[must_use]
59    pub fn allowed_methods(mut self, methods: impl IntoIterator<Item = &'static str>) -> Self {
60        self.settings.allowed_methods = methods.into_iter().collect();
61        self
62    }
63    /// Sets allowed request headers; matching ignores name case.
64    #[must_use]
65    pub fn allowed_headers(mut self, headers: impl IntoIterator<Item = &'static str>) -> Self {
66        self.settings.allowed_headers = headers.into_iter().collect();
67        self
68    }
69    /// Sets response headers exposed to browser JavaScript.
70    #[must_use]
71    pub fn exposed_headers(mut self, headers: impl IntoIterator<Item = &'static str>) -> Self {
72        self.settings.exposed_headers = headers.into_iter().collect();
73        self
74    }
75    /// Adds origin schemes to parsing; full origins must still match the policy.
76    #[must_use]
77    pub fn additional_origin_schemes(
78        mut self,
79        schemes: impl IntoIterator<Item = &'static str>,
80    ) -> Self {
81        self.settings.additional_origin_schemes = schemes.into_iter().collect();
82        self
83    }
84}
85
86/// Enforces a freshly resolved policy per cross-origin request, including preflight.
87pub async fn runtime_cors(
88    State(config): State<RuntimeCorsConfig>,
89    request: Request,
90    next: Next,
91) -> Response {
92    if (config.exempt_path)(request.uri().path()) || !request.headers().contains_key(header::ORIGIN)
93    {
94        return next.run(request).await;
95    }
96    let policy = match (config.policy)(&request) {
97        Ok(policy) => policy,
98        Err(response) => return *response,
99    };
100    if !policy.enforces_requests() {
101        return next.run(request).await;
102    }
103    let origin = match (config.request_origin)(&request) {
104        Ok(origin) => origin,
105        Err(response) => return *response,
106    };
107    let decision = crate::shared::cors::evaluate(
108        request.method().as_str(),
109        request.headers(),
110        origin.as_deref(),
111        &policy,
112        &config.settings,
113    );
114    match decision {
115        Err(error) => (config.map_error)(error),
116        Ok(CorsDecision::Pass) => next.run(request).await,
117        Ok(CorsDecision::Reject) => {
118            let mut response = StatusCode::FORBIDDEN.into_response();
119            response.headers_mut().insert(
120                header::VARY,
121                HeaderValue::from_static(
122                    "Access-Control-Request-Headers, Access-Control-Request-Method, Origin",
123                ),
124            );
125            response
126        }
127        Ok(CorsDecision::Preflight(headers)) => {
128            let mut response = StatusCode::NO_CONTENT.into_response();
129            *response.headers_mut() = headers;
130            response
131        }
132        Ok(CorsDecision::Actual(origin)) => {
133            let mut response = next.run(request).await;
134            let result =
135                crate::shared::cors::apply_origin_headers(response.headers_mut(), &policy, &origin)
136                    .and_then(|()| {
137                        crate::shared::cors::apply_actual_headers(
138                            response.headers_mut(),
139                            &config.settings,
140                        )
141                    });
142            match result {
143                Ok(()) => response,
144                Err(error) => (config.map_error)(error),
145            }
146        }
147    }
148}
149
150#[cfg(test)]
151mod tests {
152    use super::{
153        CorsAllowedOrigins, CorsMiddlewareErrorKind, RuntimeCorsConfig, RuntimeCorsPolicy,
154        runtime_cors,
155    };
156    use axum::{
157        Router,
158        body::{Body, to_bytes},
159        middleware,
160        response::{IntoResponse, Response},
161        routing::any,
162    };
163    use http::{HeaderValue, Request, StatusCode, header};
164    use std::sync::{Arc, Mutex};
165    use tower::ServiceExt;
166
167    fn policy(origins: CorsAllowedOrigins, credentials: bool) -> RuntimeCorsPolicy {
168        RuntimeCorsPolicy {
169            enabled: true,
170            allowed_origins: origins,
171            allow_credentials: credentials,
172            max_age_secs: 600,
173        }
174    }
175    fn config(policy: RuntimeCorsPolicy) -> RuntimeCorsConfig {
176        RuntimeCorsConfig::new(
177            move |_| Ok(policy.clone()),
178            |path| path == "/public",
179            |error| {
180                let status = match error.kind() {
181                    CorsMiddlewareErrorKind::InvalidRequest => StatusCode::BAD_REQUEST,
182                    CorsMiddlewareErrorKind::InvalidResponse => StatusCode::INTERNAL_SERVER_ERROR,
183                };
184                (status, format!("{:?}", error.kind())).into_response()
185            },
186        )
187        .allowed_methods(["GET", "POST", "OPTIONS"])
188        .allowed_headers(["Content-Type", "X-CSRF-Token"])
189        .exposed_headers(["x-request-id", "content-length"])
190    }
191    fn listed() -> RuntimeCorsConfig {
192        config(policy(
193            CorsAllowedOrigins::List(vec!["https://panel.example.com".into()]),
194            true,
195        ))
196    }
197    fn app(config: RuntimeCorsConfig) -> Router {
198        Router::new()
199            .fallback(any(|| async { "called" }))
200            .layer(middleware::from_fn_with_state(config, runtime_cors))
201    }
202    fn request(method: &str, origin: &str) -> http::request::Builder {
203        Request::builder().method(method).header("origin", origin)
204    }
205    async fn send(router: Router, builder: http::request::Builder) -> Response {
206        router
207            .oneshot(builder.body(Body::empty()).unwrap())
208            .await
209            .unwrap()
210    }
211    async fn body(response: Response) -> String {
212        String::from_utf8(to_bytes(response.into_body(), 4096).await.unwrap().to_vec()).unwrap()
213    }
214
215    #[tokio::test]
216    async fn preflight_allows_case_insensitive_header_names_and_does_not_call_handler() {
217        let response = send(
218            app(listed()),
219            request("OPTIONS", "https://panel.example.com")
220                .header(header::ACCESS_CONTROL_REQUEST_METHOD, "POST")
221                .header(
222                    header::ACCESS_CONTROL_REQUEST_HEADERS,
223                    "content-type, X-Csrf-Token",
224                ),
225        )
226        .await;
227        assert_eq!(response.status(), StatusCode::NO_CONTENT);
228        assert_eq!(
229            response.headers()[header::ACCESS_CONTROL_ALLOW_ORIGIN],
230            "https://panel.example.com"
231        );
232        assert_eq!(
233            response.headers()[header::ACCESS_CONTROL_ALLOW_CREDENTIALS],
234            "true"
235        );
236        assert_eq!(response.headers()[header::ACCESS_CONTROL_MAX_AGE], "600");
237        assert_eq!(
238            response.headers()[header::ACCESS_CONTROL_ALLOW_METHODS],
239            "GET, POST, OPTIONS"
240        );
241        assert_eq!(
242            response.headers()[header::ACCESS_CONTROL_ALLOW_HEADERS],
243            "Content-Type, X-CSRF-Token"
244        );
245        for name in [
246            "Origin",
247            "Access-Control-Request-Method",
248            "Access-Control-Request-Headers",
249        ] {
250            assert!(
251                response.headers()[header::VARY]
252                    .to_str()
253                    .unwrap()
254                    .contains(name)
255            );
256        }
257        assert_eq!(body(response).await, "");
258    }
259
260    #[tokio::test]
261    async fn preflight_and_actual_requests_deny_disallowed_origins_methods_or_headers() {
262        for builder in [
263            request("POST", "https://evil.example.com"),
264            request("OPTIONS", "https://evil.example.com")
265                .header(header::ACCESS_CONTROL_REQUEST_METHOD, "POST"),
266            request("OPTIONS", "https://panel.example.com")
267                .header(header::ACCESS_CONTROL_REQUEST_METHOD, "DELETE"),
268            request("OPTIONS", "https://panel.example.com")
269                .header(header::ACCESS_CONTROL_REQUEST_METHOD, "POST")
270                .header(header::ACCESS_CONTROL_REQUEST_HEADERS, "x-secret"),
271        ] {
272            let response = send(app(listed()), builder).await;
273            assert_eq!(response.status(), StatusCode::FORBIDDEN);
274            assert!(response.headers().contains_key(header::VARY));
275            assert!(
276                !response
277                    .headers()
278                    .contains_key(header::ACCESS_CONTROL_ALLOW_ORIGIN)
279            );
280        }
281    }
282
283    #[tokio::test]
284    async fn actual_requests_reflect_credentials_or_use_wildcard_without_credentials() {
285        for credentials in [false, true] {
286            let response = send(
287                app(config(policy(CorsAllowedOrigins::Any, credentials))),
288                request("GET", "https://panel.example.com"),
289            )
290            .await;
291            assert_eq!(response.status(), StatusCode::OK);
292            assert_eq!(
293                response.headers()[header::ACCESS_CONTROL_ALLOW_ORIGIN],
294                if credentials {
295                    "https://panel.example.com"
296                } else {
297                    "*"
298                }
299            );
300            assert_eq!(
301                response
302                    .headers()
303                    .contains_key(header::ACCESS_CONTROL_ALLOW_CREDENTIALS),
304                credentials
305            );
306            assert_eq!(
307                response.headers()[header::ACCESS_CONTROL_EXPOSE_HEADERS],
308                "x-request-id, content-length"
309            );
310            assert_eq!(body(response).await, "called");
311        }
312    }
313
314    #[tokio::test]
315    async fn inactive_policy_absent_origin_and_exempt_path_pass_through_without_parsing() {
316        for mut policy in [
317            policy(CorsAllowedOrigins::None, false),
318            policy(CorsAllowedOrigins::Any, false),
319        ] {
320            if policy.allowed_origins == CorsAllowedOrigins::Any {
321                policy.enabled = false;
322            }
323            let response = send(
324                app(config(policy)),
325                Request::builder().header("origin", HeaderValue::from_bytes(&[0xff]).unwrap()),
326            )
327            .await;
328            assert_eq!(response.status(), StatusCode::OK);
329            assert!(
330                !response
331                    .headers()
332                    .contains_key(header::ACCESS_CONTROL_ALLOW_ORIGIN)
333            );
334        }
335        for builder in [Request::builder(), request("GET", "null").uri("/public")] {
336            let response = send(app(listed()), builder).await;
337            assert_eq!(response.status(), StatusCode::OK);
338            assert!(
339                !response
340                    .headers()
341                    .contains_key(header::ACCESS_CONTROL_ALLOW_ORIGIN)
342            );
343        }
344    }
345
346    #[tokio::test]
347    async fn runtime_policy_changes_take_effect_on_the_same_router() {
348        let current = Arc::new(Mutex::new(policy(
349            CorsAllowedOrigins::List(vec!["https://one.example.com".into()]),
350            false,
351        )));
352        let state = Arc::clone(&current);
353        let config = RuntimeCorsConfig::new(
354            move |_| Ok(state.lock().unwrap().clone()),
355            |_| false,
356            |_| StatusCode::BAD_REQUEST.into_response(),
357        );
358        let router = app(config);
359        assert_eq!(
360            send(router.clone(), request("GET", "https://one.example.com"))
361                .await
362                .status(),
363            StatusCode::OK
364        );
365        *current.lock().unwrap() = policy(
366            CorsAllowedOrigins::List(vec!["https://two.example.com".into()]),
367            false,
368        );
369        assert_eq!(
370            send(router.clone(), request("GET", "https://one.example.com"))
371                .await
372                .status(),
373            StatusCode::FORBIDDEN
374        );
375        assert_eq!(
376            send(router, request("GET", "https://two.example.com"))
377                .await
378                .headers()[header::ACCESS_CONTROL_ALLOW_ORIGIN],
379            "https://two.example.com"
380        );
381    }
382
383    #[tokio::test]
384    async fn additional_scheme_requires_an_exact_full_origin_match() {
385        let config = config(policy(
386            CorsAllowedOrigins::List(vec!["chrome-extension://first-extension".into()]),
387            true,
388        ))
389        .additional_origin_schemes(["chrome-extension"]);
390        let router = app(config);
391        assert_eq!(
392            send(
393                router.clone(),
394                request("GET", "chrome-extension://first-extension")
395            )
396            .await
397            .status(),
398            StatusCode::OK
399        );
400        assert_eq!(
401            send(router, request("GET", "chrome-extension://other-extension"))
402                .await
403                .status(),
404            StatusCode::FORBIDDEN
405        );
406    }
407
408    #[tokio::test]
409    async fn same_origin_bypass_requires_a_product_trusted_origin_resolver() {
410        let response = send(
411            app(listed()),
412            request("GET", "https://evil.example.com")
413                .header("host", "evil.example.com")
414                .header("x-forwarded-proto", "https"),
415        )
416        .await;
417        assert_eq!(response.status(), StatusCode::FORBIDDEN);
418        let response = send(
419            app(listed().request_origin(|_| Ok(Some("https://api.example.com".into())))),
420            request("GET", "https://api.example.com"),
421        )
422        .await;
423        assert_eq!(response.status(), StatusCode::OK);
424        assert!(
425            !response
426                .headers()
427                .contains_key(header::ACCESS_CONTROL_ALLOW_ORIGIN)
428        );
429    }
430
431    #[tokio::test]
432    async fn vary_merges_all_values_case_insensitively_and_preserves_wildcard() {
433        for wildcard in [false, true] {
434            let router = Router::new()
435                .route(
436                    "/",
437                    any(move || async move {
438                        let mut response = "called".into_response();
439                        response.headers_mut().append(
440                            header::VARY,
441                            HeaderValue::from_static("Accept-Encoding, origin"),
442                        );
443                        response.headers_mut().append(
444                            header::VARY,
445                            HeaderValue::from_static(if wildcard {
446                                "*"
447                            } else {
448                                "Accept-Language"
449                            }),
450                        );
451                        response.headers_mut().insert(
452                            header::ACCESS_CONTROL_ALLOW_ORIGIN,
453                            HeaderValue::from_static("https://product.example.com"),
454                        );
455                        response
456                    }),
457                )
458                .layer(middleware::from_fn_with_state(listed(), runtime_cors));
459            let response = send(router, request("GET", "https://panel.example.com")).await;
460            let values = response
461                .headers()
462                .get_all(header::VARY)
463                .iter()
464                .map(|v| v.to_str().unwrap())
465                .collect::<Vec<_>>()
466                .join(", ");
467            assert!(values.contains("Accept-Encoding"));
468            if wildcard {
469                assert!(values.contains('*'));
470            } else {
471                assert!(values.contains("Accept-Language"));
472                assert_eq!(values.to_ascii_lowercase().matches("origin").count(), 1);
473            }
474            assert_eq!(
475                response.headers()[header::ACCESS_CONTROL_ALLOW_ORIGIN],
476                "https://product.example.com"
477            );
478        }
479    }
480
481    #[tokio::test]
482    async fn request_and_response_errors_keep_distinct_product_mappings() {
483        for builder in [
484            Request::builder().header("origin", HeaderValue::from_bytes(&[0xff]).unwrap()),
485            request("GET", "null"),
486            request("OPTIONS", "https://panel.example.com")
487                .header(header::ACCESS_CONTROL_REQUEST_METHOD, "POST")
488                .header(header::ACCESS_CONTROL_REQUEST_HEADERS, "bad name"),
489        ] {
490            let response = send(app(listed()), builder).await;
491            assert_eq!(response.status(), StatusCode::BAD_REQUEST);
492            assert_eq!(body(response).await, "InvalidRequest");
493        }
494        let response = send(
495            app(listed().exposed_headers(["bad\nheader"])),
496            request("GET", "https://panel.example.com"),
497        )
498        .await;
499        assert_eq!(response.status(), StatusCode::INTERNAL_SERVER_ERROR);
500        assert_eq!(body(response).await, "InvalidResponse");
501        let router = Router::new()
502            .fallback(any(|| async {
503                let mut response = "called".into_response();
504                response
505                    .headers_mut()
506                    .insert(header::VARY, HeaderValue::from_bytes(&[0xff]).unwrap());
507                response
508            }))
509            .layer(middleware::from_fn_with_state(listed(), runtime_cors));
510        let response = send(router, request("GET", "https://panel.example.com")).await;
511        assert_eq!(response.status(), StatusCode::INTERNAL_SERVER_ERROR);
512        assert_eq!(body(response).await, "InvalidResponse");
513    }
514
515    #[tokio::test]
516    async fn product_policy_resolver_errors_are_returned_unchanged() {
517        let config = RuntimeCorsConfig::new(
518            |_| {
519                Err(Box::new(
520                    (StatusCode::SERVICE_UNAVAILABLE, "policy unavailable").into_response(),
521                ))
522            },
523            |_| false,
524            |_| StatusCode::BAD_REQUEST.into_response(),
525        );
526        let response = send(app(config), request("GET", "https://panel.example.com")).await;
527        assert_eq!(response.status(), StatusCode::SERVICE_UNAVAILABLE);
528        assert_eq!(body(response).await, "policy unavailable");
529    }
530
531    #[tokio::test]
532    async fn product_origin_resolver_failure_is_returned_unchanged() {
533        let config = listed().request_origin(|_| {
534            Err(Box::new(
535                (StatusCode::SERVICE_UNAVAILABLE, "origin unavailable").into_response(),
536            ))
537        });
538        let response = send(app(config), request("GET", "https://panel.example.com")).await;
539        assert_eq!(response.status(), StatusCode::SERVICE_UNAVAILABLE);
540        assert_eq!(body(response).await, "origin unavailable");
541    }
542
543    #[tokio::test]
544    async fn invalid_generated_preflight_headers_are_response_errors() {
545        let response = send(
546            app(listed().allowed_headers(["content-type", "bad\nheader"])),
547            request("OPTIONS", "https://panel.example.com")
548                .header(header::ACCESS_CONTROL_REQUEST_METHOD, "POST")
549                .header(header::ACCESS_CONTROL_REQUEST_HEADERS, "content-type"),
550        )
551        .await;
552        assert_eq!(response.status(), StatusCode::INTERNAL_SERVER_ERROR);
553        assert_eq!(body(response).await, "InvalidResponse");
554    }
555}