reverse_proxy_service/
rewrite.rs

1//! A [`PathRewriter`] instance defines a rule to rewrite the request path.
2//!
3//! A "path" does not include a query. See [`http::uri::Uri`].
4
5use std::borrow::Cow;
6
7use http::uri::{Authority, Scheme, Uri};
8use http::Error as HttpError;
9use http::Request;
10
11use regex::{Regex as LibRegex, Replacer};
12
13/// Represents a rule to rewrite a path `/foo/bar/baz` to new one.
14///
15/// A "path" does not include a query. See [`http::uri::Uri`].
16pub trait PathRewriter {
17    fn rewrite<'a>(&'a mut self, path: &'a str) -> Cow<'a, str>;
18
19    fn rewrite_uri<B>(
20        &mut self,
21        req: &mut Request<B>,
22        scheme: &Scheme,
23        authority: &Authority,
24    ) -> Result<(), HttpError> {
25        let uri = {
26            let uri = req.uri();
27            let path = self.rewrite(uri.path());
28            if let Some(query) = uri.query() {
29                let mut p_and_q = path.into_owned();
30                p_and_q.push('?');
31                p_and_q.push_str(query);
32
33                Uri::builder()
34                    .scheme(scheme.clone())
35                    .authority(authority.clone())
36                    .path_and_query(p_and_q)
37                    .build()
38            } else {
39                Uri::builder()
40                    .scheme(scheme.clone())
41                    .authority(authority.clone())
42                    .path_and_query(&*path)
43                    .build()
44            }
45        }?;
46        *req.uri_mut() = uri;
47        Ok(())
48    }
49}
50
51/// Identity function, that is, this returns the `path` as is.
52///
53/// ```
54/// # use reverse_proxy_service::rewrite::{PathRewriter, Identity};
55/// assert_eq!(Identity.rewrite("foo"), "foo");
56/// ```
57#[derive(Debug, Clone, Copy, PartialEq, Eq)]
58pub struct Identity;
59
60impl PathRewriter for Identity {
61    #[inline]
62    fn rewrite<'a>(&mut self, path: &'a str) -> Cow<'a, str> {
63        path.into()
64    }
65}
66
67/// Returns `self.0` regardless what the `path` is.
68///
69/// ```
70/// # use reverse_proxy_service::rewrite::{PathRewriter, Static};
71/// assert_eq!(Static("bar").rewrite("foo"), "bar");
72/// ```
73#[derive(Debug, Clone, Copy, PartialEq, Eq)]
74pub struct Static<'a>(pub &'a str);
75
76impl PathRewriter for Static<'_> {
77    #[inline]
78    fn rewrite<'a>(&'a mut self, _path: &'a str) -> Cow<'a, str> {
79        self.0.into()
80    }
81}
82
83/// `ReplaceAll(old, new)` replaces all matches `old` with `new`.
84///
85/// ```
86/// # use reverse_proxy_service::rewrite::{PathRewriter, ReplaceAll};
87/// assert_eq!(ReplaceAll("foo", "bar").rewrite("foofoo"), "barbar");
88/// ```
89#[derive(Debug, Clone, Copy, PartialEq, Eq)]
90pub struct ReplaceAll<'a>(pub &'a str, pub &'a str);
91
92impl PathRewriter for ReplaceAll<'_> {
93    fn rewrite<'a>(&mut self, path: &'a str) -> Cow<'a, str> {
94        if path.contains(self.0) {
95            path.replace(self.0, self.1).into()
96        } else {
97            path.into()
98        }
99    }
100}
101
102/// `ReplaceN(old, new, n)` replaces first `n` matches `old` with `new`.
103///
104/// ```
105/// # use reverse_proxy_service::rewrite::{PathRewriter, ReplaceN};
106/// assert_eq!(ReplaceN("foo", "bar", 1).rewrite("foofoo"), "barfoo");
107/// assert_eq!(ReplaceN("foo", "bar", 3).rewrite("foofoo"), "barbar");
108/// ```
109#[derive(Debug, Clone, Copy, PartialEq, Eq)]
110pub struct ReplaceN<'a>(pub &'a str, pub &'a str, pub usize);
111
112impl PathRewriter for ReplaceN<'_> {
113    fn rewrite<'a>(&mut self, path: &'a str) -> Cow<'a, str> {
114        if path.contains(self.0) {
115            path.replacen(self.0, self.1, self.2).into()
116        } else {
117            path.into()
118        }
119    }
120}
121
122/// Trims a prefix if exists.
123///
124/// ```
125/// # use reverse_proxy_service::rewrite::{PathRewriter, TrimPrefix};
126/// assert_eq!(TrimPrefix("foo").rewrite("foobarfoo"), "barfoo");
127/// assert_eq!(TrimPrefix("bar").rewrite("foobarfoo"), "foobarfoo");
128/// ```
129#[derive(Debug, Clone, Copy, PartialEq, Eq)]
130pub struct TrimPrefix<'a>(pub &'a str);
131
132impl PathRewriter for TrimPrefix<'_> {
133    fn rewrite<'a>(&mut self, path: &'a str) -> Cow<'a, str> {
134        if let Some(stripped) = path.strip_prefix(self.0) {
135            stripped.into()
136        } else {
137            path.into()
138        }
139    }
140}
141
142/// Trims a suffix if exists.
143///
144/// ```
145/// # use reverse_proxy_service::rewrite::{PathRewriter, TrimSuffix};
146/// assert_eq!(TrimSuffix("foo").rewrite("foobarfoo"), "foobar");
147/// assert_eq!(TrimSuffix("bar").rewrite("foobarfoo"), "foobarfoo");
148/// ```
149#[derive(Debug, Clone, Copy, PartialEq, Eq)]
150pub struct TrimSuffix<'a>(pub &'a str);
151
152impl PathRewriter for TrimSuffix<'_> {
153    fn rewrite<'a>(&mut self, path: &'a str) -> Cow<'a, str> {
154        if let Some(stripped) = path.strip_suffix(self.0) {
155            stripped.into()
156        } else {
157            path.into()
158        }
159    }
160}
161
162/// Appends a prefix.
163///
164/// ```
165/// # use reverse_proxy_service::rewrite::{PathRewriter, AppendPrefix};
166/// assert_eq!(AppendPrefix("foo").rewrite("bar"), "foobar");
167/// ```
168#[derive(Debug, Clone, Copy, PartialEq, Eq)]
169pub struct AppendPrefix<'a>(pub &'a str);
170
171impl PathRewriter for AppendPrefix<'_> {
172    fn rewrite<'a>(&mut self, path: &'a str) -> Cow<'a, str> {
173        let mut ret = String::with_capacity(self.0.len() + path.len());
174        ret.push_str(self.0);
175        ret.push_str(path);
176        ret.into()
177    }
178}
179
180/// Appends a suffix.
181///
182/// ```
183/// # use reverse_proxy_service::rewrite::{PathRewriter, AppendSuffix};
184/// assert_eq!(AppendSuffix("foo").rewrite("bar"), "barfoo");
185/// ```
186#[derive(Debug, Clone, Copy, PartialEq, Eq)]
187pub struct AppendSuffix<'a>(pub &'a str);
188
189impl PathRewriter for AppendSuffix<'_> {
190    fn rewrite<'a>(&mut self, path: &'a str) -> Cow<'a, str> {
191        let mut ret = String::with_capacity(self.0.len() + path.len());
192        ret.push_str(path);
193        ret.push_str(self.0);
194        ret.into()
195    }
196}
197
198/// `RegexAll(re, new)` replaces all matches `re` with `new`.
199///
200/// The type of `new` must implement [`Replacer`].
201/// See [`regex`] for details.
202///
203/// ```
204/// # use reverse_proxy_service::rewrite::{PathRewriter, RegexAll};
205/// # use regex::Regex;
206/// let re = Regex::new(r"(?P<y>\d{4})/(?P<m>\d{2})").unwrap();
207/// assert_eq!(
208///     RegexAll(re, "$m-$y").rewrite("2021/10/2022/12"),
209///     "10-2021/12-2022"
210/// );
211/// ```
212#[derive(Debug, Clone)]
213pub struct RegexAll<Rep>(pub LibRegex, pub Rep);
214
215impl<Rep: Replacer> PathRewriter for RegexAll<Rep> {
216    fn rewrite<'a>(&mut self, path: &'a str) -> Cow<'a, str> {
217        self.0.replace_all(path, self.1.by_ref())
218    }
219}
220
221/// `RegexN(re, new, n)` replaces first `n` matches `re` with `new`.
222///
223/// The type of `new` must implement [`Replacer`].
224/// See [`regex`] for details.
225///
226/// ```
227/// # use reverse_proxy_service::rewrite::{PathRewriter, RegexN};
228/// # use regex::Regex;
229/// let re = Regex::new(r"(?P<y>\d{4})/(?P<m>\d{2})").unwrap();
230/// assert_eq!(
231///     RegexN(re.clone(), "$m-$y", 1).rewrite("2021/10/2022/12"),
232///     "10-2021/2022/12"
233/// );
234/// assert_eq!(
235///     RegexN(re, "$m-$y", 3).rewrite("2021/10/2022/12"),
236///     "10-2021/12-2022"
237/// );
238/// ```
239#[derive(Debug, Clone)]
240pub struct RegexN<Rep>(pub LibRegex, pub Rep, pub usize);
241
242impl<Rep: Replacer> PathRewriter for RegexN<Rep> {
243    fn rewrite<'a>(&mut self, path: &'a str) -> Cow<'a, str> {
244        self.0.replacen(path, self.2, self.1.by_ref())
245    }
246}
247
248/// Converts the `path` by a function.
249///
250/// The type of the function must be `for<'a> FnMut(&'a str) -> String`.
251///
252/// ```
253/// # use reverse_proxy_service::rewrite::{PathRewriter, Func};
254/// let f = |path: &str| path.len().to_string();
255/// assert_eq!(Func(f).rewrite("abc"), "3");
256/// ```
257pub struct Func<F>(pub F);
258
259impl<F> PathRewriter for Func<F>
260where
261    for<'a> F: FnMut(&'a str) -> String,
262{
263    fn rewrite<'a>(&'a mut self, path: &'a str) -> Cow<'a, str> {
264        self.0(path).into()
265    }
266}
267
268#[cfg(test)]
269mod test {
270    use super::*;
271
272    #[test]
273    fn rewrite_static() {
274        let path = "/foo/bar";
275        let mut rw = Static("/baz");
276        assert_eq!(rw.rewrite(path), "/baz");
277    }
278
279    #[test]
280    fn replace() {
281        let path = "/foo/bar/foo/baz/foo";
282        let mut rw = ReplaceAll("foo", "FOO");
283        assert_eq!(rw.rewrite(path), "/FOO/bar/FOO/baz/FOO");
284
285        let path = "/foo/bar/foo/baz/foo";
286        let mut rw = ReplaceAll("/foo", "");
287        assert_eq!(rw.rewrite(path), "/bar/baz");
288
289        let path = "/foo/bar/foo/baz/foo";
290        let mut rw = ReplaceN("foo", "FOO", 2);
291        assert_eq!(rw.rewrite(path), "/FOO/bar/FOO/baz/foo");
292    }
293
294    #[test]
295    fn trim() {
296        let path = "/foo/foo/bar";
297        let mut rw = TrimPrefix("/foo");
298        assert_eq!(rw.rewrite(path), "/foo/bar");
299
300        let path = "/foo/foo/bar";
301        let mut rw = TrimPrefix("foo");
302        assert_eq!(rw.rewrite(path), "/foo/foo/bar");
303
304        let path = "/bar/foo/foo";
305        let mut rw = TrimSuffix("foo");
306        assert_eq!(rw.rewrite(path), "/bar/foo/");
307
308        let path = "/bar/foo/foo";
309        let mut rw = TrimSuffix("foo/");
310        assert_eq!(rw.rewrite(path), "/bar/foo/foo");
311    }
312
313    #[test]
314    fn append() {
315        let path = "/foo/bar";
316        let mut rw = AppendPrefix("/baz");
317        assert_eq!(rw.rewrite(path), "/baz/foo/bar");
318
319        let path = "/foo/bar";
320        let mut rw = AppendSuffix("/baz");
321        assert_eq!(rw.rewrite(path), "/foo/bar/baz");
322    }
323
324    #[test]
325    fn regex() {
326        let path = "/2021/10/21/2021/12/02/2022/01/13";
327        let mut rw = RegexAll(
328            LibRegex::new(r"(?P<y>\d{4})/(?P<m>\d{2})/(?P<d>\d{2})").unwrap(),
329            "$m-$d-$y",
330        );
331        assert_eq!(rw.rewrite(path), "/10-21-2021/12-02-2021/01-13-2022");
332
333        let path = "/2021/10/21/2021/12/02/2022/01/13";
334        let mut rw = RegexN(
335            LibRegex::new(r"(?P<y>\d{4})/(?P<m>\d{2})/(?P<d>\d{2})").unwrap(),
336            "$m-$d-$y",
337            2,
338        );
339        assert_eq!(rw.rewrite(path), "/10-21-2021/12-02-2021/2022/01/13");
340    }
341
342    #[test]
343    fn func() {
344        let path = "/abcdefg";
345        let mut rw = Func(|path: &str| path.len().to_string());
346        assert_eq!(rw.rewrite(path), "8");
347    }
348}