gatel_core/hoops/
replace.rs

1use http::header::{CONTENT_LENGTH, CONTENT_TYPE};
2use salvo::{Depot, FlowCtrl, Request, Response, async_trait};
3use tracing::debug;
4
5/// Response body text-replacement middleware.
6///
7/// After the inner chain produces a response this middleware collects the body,
8/// converts it to a UTF-8 string, applies each `(search, replacement)` rule in
9/// order, updates the `Content-Length` header to reflect the new body size, and
10/// returns the modified response.
11///
12/// Replacements are only applied when the response `Content-Type` matches one of
13/// the configured MIME types (default: `["text/html"]`).
14pub struct ReplaceHoop {
15    rules: Vec<(String, String)>,
16    once: bool,
17    content_types: Vec<String>,
18}
19
20impl ReplaceHoop {
21    pub fn new(rules: Vec<(String, String)>, once: bool) -> Self {
22        Self {
23            rules,
24            once,
25            content_types: vec!["text/html".to_string()],
26        }
27    }
28
29    /// Return true if the response Content-Type header matches any of the
30    /// configured MIME types.
31    fn content_type_matches(&self, headers: &http::HeaderMap) -> bool {
32        let ct = headers
33            .get(CONTENT_TYPE)
34            .and_then(|v| v.to_str().ok())
35            .unwrap_or("");
36        self.content_types
37            .iter()
38            .any(|allowed| ct.contains(allowed.as_str()))
39    }
40}
41
42#[async_trait]
43impl salvo::Handler for ReplaceHoop {
44    async fn handle(
45        &self,
46        req: &mut Request,
47        depot: &mut Depot,
48        res: &mut Response,
49        ctrl: &mut FlowCtrl,
50    ) {
51        ctrl.call_next(req, depot, res).await;
52
53        // Only modify responses with matching content types.
54        if !self.content_type_matches(res.headers()) {
55            return;
56        }
57
58        // Take the body and collect it.
59        let body = res.take_body();
60        let body_bytes = match super::compress::collect_res_body_bytes(body).await {
61            Ok(b) => b,
62            Err(_) => return,
63        };
64
65        // Convert to string; skip replacement on non-UTF-8 bodies.
66        let text = match std::str::from_utf8(&body_bytes) {
67            Ok(s) => s.to_string(),
68            Err(_) => {
69                debug!("response body is not valid UTF-8; skipping replace middleware");
70                res.body(body_bytes);
71                return;
72            }
73        };
74
75        // Apply each rule.
76        let mut output = text;
77        for (search, replacement) in &self.rules {
78            output = if self.once {
79                output.replacen(search.as_str(), replacement.as_str(), 1)
80            } else {
81                output.replace(search.as_str(), replacement.as_str())
82            };
83            debug!(
84                search = search.as_str(),
85                replacement = replacement.as_str(),
86                once = self.once,
87                "applied body replacement rule"
88            );
89        }
90
91        // Rebuild with the updated body and Content-Length.
92        let new_bytes = output.into_bytes();
93        res.headers_mut()
94            .insert(CONTENT_LENGTH, http::HeaderValue::from(new_bytes.len()));
95        res.body(new_bytes);
96    }
97}
gatel_core/hoops/replace.rs

gatel_core/hoops/
replace.rs
