Skip to main content

gatel_core/hoops/
headers.rs

1use std::collections::HashMap;
2
3use salvo::{Depot, FlowCtrl, Request, Response, async_trait};
4use tracing::debug;
5
6use crate::config::HeadersConfig;
7use crate::placeholder::expand_placeholders;
8
9/// Request/response header manipulation middleware.
10///
11/// Supports setting, adding, and deleting headers on both the request
12/// (before forwarding) and the response (before returning to client).
13/// Header values may contain placeholders that are expanded at runtime:
14///   - `{client_ip}` — client socket address IP
15///   - `{host}` — request Host header value
16///   - `{method}` — HTTP method
17///   - `{path}` — request URI path
18///   - `{scheme}` — request URI scheme (http/https)
19pub struct HeadersHoop {
20    request_set: Vec<(String, String)>,
21    response_set: Vec<(String, String)>,
22    response_remove: Vec<String>,
23}
24
25impl HeadersHoop {
26    pub fn new(cfg: &HeadersConfig) -> Self {
27        let request_set: Vec<(String, String)> = cfg
28            .request_set
29            .iter()
30            .map(|(k, v)| (k.clone(), v.clone()))
31            .collect();
32        let response_set: Vec<(String, String)> = cfg
33            .response_set
34            .iter()
35            .map(|(k, v)| (k.clone(), v.clone()))
36            .collect();
37        let response_remove = cfg.response_remove.clone();
38        Self {
39            request_set,
40            response_set,
41            response_remove,
42        }
43    }
44}
45
46#[async_trait]
47impl salvo::Handler for HeadersHoop {
48    async fn handle(
49        &self,
50        req: &mut Request,
51        depot: &mut Depot,
52        res: &mut Response,
53        ctrl: &mut FlowCtrl,
54    ) {
55        // Build placeholder values from the request.
56        let placeholders = build_placeholders(req);
57
58        // Apply request header modifications.
59        for (name, value_template) in &self.request_set {
60            let expanded = expand_placeholders(value_template, &placeholders);
61            if let (Ok(hn), Ok(hv)) = (
62                name.parse::<http::header::HeaderName>(),
63                expanded.parse::<http::header::HeaderValue>(),
64            ) {
65                debug!(header = %name, value = %expanded, "setting request header");
66                req.headers_mut().insert(hn, hv);
67            }
68        }
69
70        // Forward to next middleware / handler.
71        ctrl.call_next(req, depot, res).await;
72
73        // Apply response header removals.
74        for name in &self.response_remove {
75            if let Ok(hn) = name.parse::<http::header::HeaderName>() {
76                res.headers_mut().remove(hn);
77            }
78        }
79
80        // Apply response header sets.
81        for (name, value_template) in &self.response_set {
82            let expanded = expand_placeholders(value_template, &placeholders);
83            if let (Ok(hn), Ok(hv)) = (
84                name.parse::<http::header::HeaderName>(),
85                expanded.parse::<http::header::HeaderValue>(),
86            ) {
87                debug!(header = %name, value = %expanded, "setting response header");
88                res.headers_mut().insert(hn, hv);
89            }
90        }
91    }
92}
93
94/// Build a map of placeholder names to their current values from the request.
95fn build_placeholders(req: &Request) -> HashMap<&'static str, String> {
96    let client_addr = super::client_addr(req);
97    let mut m = HashMap::new();
98    m.insert("client_ip", client_addr.ip().to_string());
99    m.insert(
100        "host",
101        req.headers()
102            .get(http::header::HOST)
103            .and_then(|v| v.to_str().ok())
104            .unwrap_or("")
105            .to_string(),
106    );
107    m.insert("method", req.method().to_string());
108    m.insert("path", req.uri().path().to_string());
109    m.insert(
110        "scheme",
111        req.uri().scheme_str().unwrap_or("http").to_string(),
112    );
113    m
114}