1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
use axum::{extract::State, http::HeaderValue, middleware::Next, response::Response};
use axum_csp::*;

use crate::web::GoatState;

pub async fn cspheaders(
    State(state): State<GoatState>,
    req: axum::extract::Request,
    next: Next,
) -> Response {
    let uri: String = req.uri().path().to_string();
    let url_matcher: Option<CspUrlMatcher> = state.read().await.csp_matchers.iter().find_map(|c| {
        if c.matcher.is_match(&uri) {
            Some(c.to_owned())
        } else {
            None
        }
    });

    // wait for the middleware to come back
    let mut response = next.run(req).await;

    // if we found one, woot
    if let Some(rule) = url_matcher {
        let headers = response.headers_mut();
        if rule.matcher.is_match(&uri) {
            let header: HeaderValue = rule.into();
            headers.insert("Content-Security-Policy", header);
        }
    } else {
        log::debug!("didn't match uri");
    }

    response
}