rws_server/rewrite/
mod.rs

1//! Request and response rewriting middleware.
2//!
3//! [`RewriteLayer`] is a [`Middleware`] that transforms requests before they
4//! reach handlers and responses before they leave the server. Build one with the
5//! fluent builder API and add it to any [`crate::middleware::WithMiddleware`] stack.
6//!
7//! # Example
8//!
9//! ```rust,no_run
10//! # extern crate rws_server as rust_web_server;
11//! use rust_web_server::app::App;
12//! use rust_web_server::core::New;
13//! use rust_web_server::rewrite::RewriteLayer;
14//!
15//! let app = App::new()
16//!     .wrap(RewriteLayer::new()
17//!         .request_header_set("X-Env", "production")
18//!         .request_uri_strip_prefix("/api/v1")
19//!         .response_header_set("Cache-Control", "no-store")
20//!         .response_body_replace("http://staging.internal", "https://example.com"));
21//! ```
22//!
23//! # Regex URI rewriting (`rewrite-regex` feature)
24//!
25//! The prefix/set operations above cover fixed strings. When the rewrite
26//! depends on part of the incoming path — versioning schemes, locale
27//! prefixes, ID extraction — [`RewriteLayer::request_uri_regex_rewrite`]
28//! matches the URI against a regex and rewrites it using the match's capture
29//! groups, the same `rewrite` semantics as nginx: if the pattern matches
30//! anywhere in the URI, the **entire** URI is replaced by the expanded
31//! replacement string; otherwise the URI is left untouched.
32//!
33//! ```rust,no_run
34//! # extern crate rws_server as rust_web_server;
35//! # #[cfg(feature = "rewrite-regex")]
36//! # fn example() -> Result<(), regex::Error> {
37//! use rust_web_server::app::App;
38//! use rust_web_server::core::New;
39//! use rust_web_server::rewrite::RewriteLayer;
40//!
41//! let app = App::new()
42//!     .wrap(RewriteLayer::new()
43//!         .request_uri_regex_rewrite(r"^/api/v\d+/(.*)$", "/$1")?);
44//! # Ok(())
45//! # }
46//! ```
47//!
48//! Requires the `rewrite-regex` feature (adds the `regex` crate) — this is
49//! the one place in `rws` a third-party regex engine is worth the
50//! dependency; hand-rolling one is out of scope for this crate's "no
51//! third-party HTTP dependencies" philosophy, which doesn't extend to
52//! general-purpose text processing.
53
54#[cfg(test)]
55mod tests;
56
57use crate::application::Application;
58use crate::header::Header;
59use crate::middleware::Middleware;
60use crate::request::Request;
61use crate::response::Response;
62use crate::server::ConnectionInfo;
63
64#[cfg(feature = "rewrite-regex")]
65use regex::Regex;
66
67enum RequestRule {
68    SetHeader { name: String, value: String },
69    RemoveHeader(String),
70    SetUri(String),
71    StripUriPrefix(String),
72    AddUriPrefix(String),
73    #[cfg(feature = "rewrite-regex")]
74    RewriteUri { pattern: Regex, replacement: String },
75}
76
77enum ResponseRule {
78    SetHeader { name: String, value: String },
79    RemoveHeader(String),
80    SetStatus { code: i16, reason: String },
81    BodyReplace { from: Vec<u8>, to: Vec<u8> },
82}
83
84/// Composable request/response rewriting middleware.
85///
86/// Clones the incoming [`Request`], applies request rules, dispatches to the
87/// next handler, then applies response rules on the returned [`Response`].
88///
89/// All builder methods take `self` by value and return `Self` for chaining.
90pub struct RewriteLayer {
91    request_rules: Vec<RequestRule>,
92    response_rules: Vec<ResponseRule>,
93}
94
95impl RewriteLayer {
96    pub fn new() -> Self {
97        RewriteLayer { request_rules: Vec::new(), response_rules: Vec::new() }
98    }
99
100    /// Add or replace a request header (case-insensitive name match).
101    pub fn request_header_set(mut self, name: &str, value: &str) -> Self {
102        self.request_rules.push(RequestRule::SetHeader {
103            name: name.to_string(),
104            value: value.to_string(),
105        });
106        self
107    }
108
109    /// Remove a request header (case-insensitive).
110    pub fn request_header_remove(mut self, name: &str) -> Self {
111        self.request_rules.push(RequestRule::RemoveHeader(name.to_string()));
112        self
113    }
114
115    /// Replace the entire request URI.
116    pub fn request_uri_set(mut self, uri: &str) -> Self {
117        self.request_rules.push(RequestRule::SetUri(uri.to_string()));
118        self
119    }
120
121    /// Strip a path prefix from the request URI. No-op if the prefix is absent.
122    /// Normalizes to `"/"` if stripping leaves an empty path.
123    pub fn request_uri_strip_prefix(mut self, prefix: &str) -> Self {
124        self.request_rules.push(RequestRule::StripUriPrefix(prefix.to_string()));
125        self
126    }
127
128    /// Prepend a prefix to the request URI.
129    pub fn request_uri_add_prefix(mut self, prefix: &str) -> Self {
130        self.request_rules.push(RequestRule::AddUriPrefix(prefix.to_string()));
131        self
132    }
133
134    /// Rewrite the request URI by regex, nginx `rewrite`-directive style.
135    ///
136    /// If `pattern` matches anywhere in the URI, the **entire** URI is
137    /// replaced by `replacement` with capture-group references expanded —
138    /// `$1`, `$2`, ... for numbered groups, `${name}` for named groups
139    /// (`(?P<name>...)`), or `$0`/`${0}` for the whole match. If `pattern`
140    /// does not match, the URI is left unchanged.
141    ///
142    /// Returns `Err` if `pattern` is not a valid regex. Requires the
143    /// `rewrite-regex` feature.
144    #[cfg(feature = "rewrite-regex")]
145    pub fn request_uri_regex_rewrite(mut self, pattern: &str, replacement: &str) -> Result<Self, regex::Error> {
146        let compiled = Regex::new(pattern)?;
147        self.request_rules.push(RequestRule::RewriteUri {
148            pattern: compiled,
149            replacement: replacement.to_string(),
150        });
151        Ok(self)
152    }
153
154    /// Add or replace a response header (case-insensitive name match).
155    pub fn response_header_set(mut self, name: &str, value: &str) -> Self {
156        self.response_rules.push(ResponseRule::SetHeader {
157            name: name.to_string(),
158            value: value.to_string(),
159        });
160        self
161    }
162
163    /// Remove a response header (case-insensitive).
164    pub fn response_header_remove(mut self, name: &str) -> Self {
165        self.response_rules.push(ResponseRule::RemoveHeader(name.to_string()));
166        self
167    }
168
169    /// Override the response status code and reason phrase.
170    pub fn response_status(mut self, code: i16, reason: &str) -> Self {
171        self.response_rules.push(ResponseRule::SetStatus { code, reason: reason.to_string() });
172        self
173    }
174
175    /// Byte-level find-and-replace across all response body content ranges.
176    pub fn response_body_replace(mut self, from: &str, to: &str) -> Self {
177        self.response_rules.push(ResponseRule::BodyReplace {
178            from: from.as_bytes().to_vec(),
179            to: to.as_bytes().to_vec(),
180        });
181        self
182    }
183}
184
185impl Middleware for RewriteLayer {
186    fn handle(
187        &self,
188        request: &Request,
189        connection: &ConnectionInfo,
190        next: &dyn Application,
191    ) -> Result<Response, String> {
192        let mut req = request.clone();
193
194        for rule in &self.request_rules {
195            match rule {
196                RequestRule::SetHeader { name, value } => {
197                    req.headers.retain(|h| !h.name.eq_ignore_ascii_case(name));
198                    req.headers.push(Header { name: name.clone().into(), value: value.clone().into() });
199                }
200                RequestRule::RemoveHeader(name) => {
201                    req.headers.retain(|h| !h.name.eq_ignore_ascii_case(name));
202                }
203                RequestRule::SetUri(uri) => {
204                    req.request_uri = uri.clone();
205                }
206                RequestRule::StripUriPrefix(prefix) => {
207                    if let Some(stripped) = req.request_uri.strip_prefix(prefix.as_str()) {
208                        req.request_uri = if stripped.is_empty() || !stripped.starts_with('/') {
209                            format!("/{}", stripped)
210                        } else {
211                            stripped.to_string()
212                        };
213                    }
214                }
215                RequestRule::AddUriPrefix(prefix) => {
216                    req.request_uri = format!("{}{}", prefix, req.request_uri);
217                }
218                #[cfg(feature = "rewrite-regex")]
219                RequestRule::RewriteUri { pattern, replacement } => {
220                    if let Some(captures) = pattern.captures(&req.request_uri) {
221                        let mut expanded = String::new();
222                        captures.expand(replacement, &mut expanded);
223                        req.request_uri = expanded;
224                    }
225                }
226            }
227        }
228
229        let mut response = next.execute(&req, connection)?;
230
231        for rule in &self.response_rules {
232            match rule {
233                ResponseRule::SetHeader { name, value } => {
234                    response.headers.retain(|h| !h.name.eq_ignore_ascii_case(name));
235                    response.headers.push(Header { name: name.clone().into(), value: value.clone().into() });
236                }
237                ResponseRule::RemoveHeader(name) => {
238                    response.headers.retain(|h| !h.name.eq_ignore_ascii_case(name));
239                }
240                ResponseRule::SetStatus { code, reason } => {
241                    response.status_code = *code;
242                    response.reason_phrase = reason.clone();
243                }
244                ResponseRule::BodyReplace { from, to } => {
245                    for cr in &mut response.content_range_list {
246                        cr.body = replace_bytes(&cr.body, from, to);
247                    }
248                }
249            }
250        }
251
252        Ok(response)
253    }
254}
255
256fn replace_bytes(haystack: &[u8], needle: &[u8], replacement: &[u8]) -> Vec<u8> {
257    if needle.is_empty() {
258        return haystack.to_vec();
259    }
260    let mut result = Vec::with_capacity(haystack.len());
261    let mut i = 0;
262    while i + needle.len() <= haystack.len() {
263        if haystack[i..].starts_with(needle) {
264            result.extend_from_slice(replacement);
265            i += needle.len();
266        } else {
267            result.push(haystack[i]);
268            i += 1;
269        }
270    }
271    result.extend_from_slice(&haystack[i..]);
272    result
273}
rws_server/rewrite/mod.rs

rws_server/rewrite/
mod.rs
