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

axum_proxy/
rewrite.rs
