tower_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 tower_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 tower_proxy::rewrite::{PathRewriter, Static};
73/// assert_eq!(Static("bar").rewrite("foo"), "bar");
74/// ```
75#[derive(Debug, Clone, Copy, PartialEq, Eq)]
76pub struct Static<S>(pub S);
77
78impl<S: AsRef<str>> PathRewriter for Static<S> {
79    #[inline]
80    fn rewrite<'a>(&'a mut self, _path: &'a str) -> Cow<'a, str> {
81        self.0.as_ref().into()
82    }
83}
84
85/// `ReplaceAll(old, new)` replaces all matches `old` with `new`.
86///
87/// ```
88/// # use tower_proxy::rewrite::{PathRewriter, ReplaceAll};
89/// assert_eq!(ReplaceAll("foo", "bar").rewrite("foofoo"), "barbar");
90/// ```
91#[derive(Debug, Clone, Copy, PartialEq, Eq)]
92pub struct ReplaceAll<S1, S2>(pub S1, pub S2);
93
94impl<S1: AsRef<str>, S2: AsRef<str>> PathRewriter for ReplaceAll<S1, S2> {
95    fn rewrite<'a>(&mut self, path: &'a str) -> Cow<'a, str> {
96        let old = self.0.as_ref();
97        if path.contains(old) {
98            path.replace(old, self.1.as_ref()).into()
99        } else {
100            path.into()
101        }
102    }
103}
104
105/// `ReplaceN(old, new, n)` replaces first `n` matches `old` with `new`.
106///
107/// ```
108/// # use tower_proxy::rewrite::{PathRewriter, ReplaceN};
109/// assert_eq!(ReplaceN("foo", "bar", 1).rewrite("foofoo"), "barfoo");
110/// assert_eq!(ReplaceN("foo", "bar", 3).rewrite("foofoo"), "barbar");
111/// ```
112#[derive(Debug, Clone, Copy, PartialEq, Eq)]
113pub struct ReplaceN<S1, S2>(pub S1, pub S2, pub usize);
114
115impl<S1: AsRef<str>, S2: AsRef<str>> PathRewriter for ReplaceN<S1, S2> {
116    fn rewrite<'a>(&mut self, path: &'a str) -> Cow<'a, str> {
117        let old = self.0.as_ref();
118        if path.contains(old) {
119            path.replacen(old, self.1.as_ref(), self.2).into()
120        } else {
121            path.into()
122        }
123    }
124}
125
126/// Trims a prefix if exists.
127///
128/// ```
129/// # use tower_proxy::rewrite::{PathRewriter, TrimPrefix};
130/// assert_eq!(TrimPrefix("foo").rewrite("foobarfoo"), "barfoo");
131/// assert_eq!(TrimPrefix("bar").rewrite("foobarfoo"), "foobarfoo");
132/// ```
133#[derive(Debug, Clone, Copy, PartialEq, Eq)]
134pub struct TrimPrefix<S>(pub S);
135
136impl<S: AsRef<str>> PathRewriter for TrimPrefix<S> {
137    fn rewrite<'a>(&mut self, path: &'a str) -> Cow<'a, str> {
138        if let Some(stripped) = path.strip_prefix(self.0.as_ref()) {
139            stripped.into()
140        } else {
141            path.into()
142        }
143    }
144}
145
146/// Trims a suffix if exists.
147///
148/// ```
149/// # use tower_proxy::rewrite::{PathRewriter, TrimSuffix};
150/// assert_eq!(TrimSuffix("foo").rewrite("foobarfoo"), "foobar");
151/// assert_eq!(TrimSuffix("bar").rewrite("foobarfoo"), "foobarfoo");
152/// ```
153#[derive(Debug, Clone, Copy, PartialEq, Eq)]
154pub struct TrimSuffix<S>(pub S);
155
156impl<S: AsRef<str>> PathRewriter for TrimSuffix<S> {
157    fn rewrite<'a>(&mut self, path: &'a str) -> Cow<'a, str> {
158        if let Some(stripped) = path.strip_suffix(self.0.as_ref()) {
159            stripped.into()
160        } else {
161            path.into()
162        }
163    }
164}
165
166/// Appends a prefix by plain string concatenation.
167///
168/// A prefix ending in `/` produces `//` when the path starts with `/`. Use [`AppendPathPrefix`] to join slash-safely.
169///
170/// ```
171/// # use tower_proxy::rewrite::{PathRewriter, AppendPrefix};
172/// assert_eq!(AppendPrefix("foo").rewrite("bar"), "foobar");
173/// ```
174#[derive(Debug, Clone, Copy, PartialEq, Eq)]
175pub struct AppendPrefix<S>(pub S);
176
177impl<S: AsRef<str>> PathRewriter for AppendPrefix<S> {
178    fn rewrite<'a>(&mut self, path: &'a str) -> Cow<'a, str> {
179        let prefix = self.0.as_ref();
180        let mut ret = String::with_capacity(prefix.len() + path.len());
181        ret.push_str(prefix);
182        ret.push_str(path);
183        ret.into()
184    }
185}
186
187/// Appends a prefix to a path.
188///
189/// Unlike [`AppendPrefix`], this joins slash-safely: trailing slashes are trimmed from the prefix at construction, so joining with a path that starts with `/` never produces `//`.
190///
191/// ```
192/// # use tower_proxy::rewrite::{PathRewriter, AppendPathPrefix};
193/// assert_eq!(AppendPathPrefix::new("/api").rewrite("/foo"), "/api/foo");
194/// assert_eq!(AppendPathPrefix::new("/api/").rewrite("/foo"), "/api/foo");
195/// assert_eq!(AppendPathPrefix::new("/").rewrite("/foo"), "/foo");
196/// ```
197#[derive(Debug, Clone, PartialEq, Eq)]
198pub struct AppendPathPrefix<'p>(Cow<'p, str>);
199
200impl<'p> AppendPathPrefix<'p> {
201    #[must_use]
202    pub fn new<I: Into<Cow<'p, str>>>(prefix: I) -> Self {
203        let prefix = match prefix.into() {
204            Cow::Borrowed(borrowed) => Cow::Borrowed(borrowed.trim_end_matches('/')),
205            Cow::Owned(mut owned) => {
206                let trimmed = owned.trim_end_matches('/').len();
207                owned.truncate(trimmed);
208                Cow::Owned(owned)
209            },
210        };
211
212        Self(prefix)
213    }
214}
215
216impl PathRewriter for AppendPathPrefix<'_> {
217    fn rewrite<'a>(&mut self, path: &'a str) -> Cow<'a, str> {
218        if self.0.is_empty() {
219            return path.into();
220        }
221
222        let mut ret = String::with_capacity(self.0.len() + path.len());
223        ret.push_str(&self.0);
224        ret.push_str(path);
225        ret.into()
226    }
227}
228
229/// Appends a suffix.
230///
231/// ```
232/// # use tower_proxy::rewrite::{PathRewriter, AppendSuffix};
233/// assert_eq!(AppendSuffix("foo").rewrite("bar"), "barfoo");
234/// ```
235#[derive(Debug, Clone, Copy, PartialEq, Eq)]
236pub struct AppendSuffix<S>(pub S);
237
238impl<S: AsRef<str>> PathRewriter for AppendSuffix<S> {
239    fn rewrite<'a>(&mut self, path: &'a str) -> Cow<'a, str> {
240        let suffix = self.0.as_ref();
241        let mut ret = String::with_capacity(suffix.len() + path.len());
242        ret.push_str(path);
243        ret.push_str(suffix);
244        ret.into()
245    }
246}
247
248/// `RegexAll(re, new)` replaces all matches `re` with `new`.
249///
250/// The type of `new` must implement [`Replacer`].
251/// See [`regex`] for details.
252///
253/// ```
254/// # use tower_proxy::rewrite::{PathRewriter, RegexAll};
255/// # use regex::Regex;
256/// let re = Regex::new(r"(?P<y>\d{4})/(?P<m>\d{2})").unwrap();
257/// assert_eq!(
258///     RegexAll(re, "$m-$y").rewrite("2021/10/2022/12"),
259///     "10-2021/12-2022"
260/// );
261/// ```
262#[derive(Debug, Clone)]
263pub struct RegexAll<Rep>(pub LibRegex, pub Rep);
264
265impl<Rep: Replacer> PathRewriter for RegexAll<Rep> {
266    fn rewrite<'a>(&mut self, path: &'a str) -> Cow<'a, str> {
267        self.0.replace_all(path, self.1.by_ref())
268    }
269}
270
271/// `RegexN(re, new, n)` replaces first `n` matches `re` with `new`.
272///
273/// The type of `new` must implement [`Replacer`].
274/// See [`regex`] for details.
275///
276/// ```
277/// # use tower_proxy::rewrite::{PathRewriter, RegexN};
278/// # use regex::Regex;
279/// let re = Regex::new(r"(?P<y>\d{4})/(?P<m>\d{2})").unwrap();
280/// assert_eq!(
281///     RegexN(re.clone(), "$m-$y", 1).rewrite("2021/10/2022/12"),
282///     "10-2021/2022/12"
283/// );
284/// assert_eq!(
285///     RegexN(re, "$m-$y", 3).rewrite("2021/10/2022/12"),
286///     "10-2021/12-2022"
287/// );
288/// ```
289#[derive(Debug, Clone)]
290pub struct RegexN<Rep>(pub LibRegex, pub Rep, pub usize);
291
292impl<Rep: Replacer> PathRewriter for RegexN<Rep> {
293    fn rewrite<'a>(&mut self, path: &'a str) -> Cow<'a, str> {
294        self.0.replacen(path, self.2, self.1.by_ref())
295    }
296}
297
298/// Converts the `path` by a function.
299///
300/// The type of the function must be `for<'a> FnMut(&'a str) -> String`.
301///
302/// ```
303/// # use tower_proxy::rewrite::{PathRewriter, Func};
304/// let f = |path: &str| path.len().to_string();
305/// assert_eq!(Func(f).rewrite("abc"), "3");
306/// ```
307pub struct Func<F>(pub F);
308
309impl<F> PathRewriter for Func<F>
310where
311    for<'a> F: FnMut(&'a str) -> String,
312{
313    fn rewrite<'a>(&'a mut self, path: &'a str) -> Cow<'a, str> {
314        self.0(path).into()
315    }
316}
317
318#[cfg(test)]
319mod test {
320    use pretty_assertions::assert_eq;
321
322    use super::{
323        AppendPathPrefix, AppendPrefix, AppendSuffix, Func, LibRegex, PathRewriter as _, RegexAll,
324        RegexN, ReplaceAll, ReplaceN, Static, TrimPrefix, TrimSuffix,
325    };
326
327    #[test]
328    fn rewrite_static() {
329        let path = "/foo/bar";
330        let mut rw = Static("/baz");
331        assert_eq!(rw.rewrite(path), "/baz");
332    }
333
334    #[test]
335    fn replace() {
336        let path = "/foo/bar/foo/baz/foo";
337        let mut rw = ReplaceAll("foo", "FOO");
338        assert_eq!(rw.rewrite(path), "/FOO/bar/FOO/baz/FOO");
339
340        let path = "/foo/bar/foo/baz/foo";
341        let mut rw = ReplaceAll("/foo", "");
342        assert_eq!(rw.rewrite(path), "/bar/baz");
343
344        let path = "/foo/bar/foo/baz/foo";
345        let mut rw = ReplaceN("foo", "FOO", 2);
346        assert_eq!(rw.rewrite(path), "/FOO/bar/FOO/baz/foo");
347    }
348
349    #[test]
350    fn trim() {
351        let path = "/foo/foo/bar";
352        let mut rw = TrimPrefix("/foo");
353        assert_eq!(rw.rewrite(path), "/foo/bar");
354
355        let path = "/foo/foo/bar";
356        let mut rw = TrimPrefix("foo");
357        assert_eq!(rw.rewrite(path), "/foo/foo/bar");
358
359        let path = "/bar/foo/foo";
360        let mut rw = TrimSuffix("foo");
361        assert_eq!(rw.rewrite(path), "/bar/foo/");
362
363        let path = "/bar/foo/foo";
364        let mut rw = TrimSuffix("foo/");
365        assert_eq!(rw.rewrite(path), "/bar/foo/foo");
366    }
367
368    #[test]
369    fn append() {
370        let path = "/foo/bar";
371        let mut rw = AppendPrefix("/baz");
372        assert_eq!(rw.rewrite(path), "/baz/foo/bar");
373
374        let path = "/foo/bar";
375        let mut rw = AppendSuffix("/baz");
376        assert_eq!(rw.rewrite(path), "/foo/bar/baz");
377
378        let path = "/foo/bar";
379        let mut rw = AppendPrefix("/baz".to_owned());
380        assert_eq!(rw.rewrite(path), "/baz/foo/bar");
381    }
382
383    #[test]
384    fn append_path() {
385        let mut rw = AppendPathPrefix::new("/baz");
386        assert_eq!(rw.rewrite("/foo/bar"), "/baz/foo/bar");
387
388        let mut rw = AppendPathPrefix::new("/baz/");
389        assert_eq!(rw.rewrite("/foo/bar"), "/baz/foo/bar");
390
391        let mut rw = AppendPathPrefix::new(String::from("/baz/"));
392        assert_eq!(rw.rewrite("/foo/bar"), "/baz/foo/bar");
393
394        let mut rw = AppendPathPrefix::new("/");
395        assert_eq!(rw.rewrite("/foo/bar"), "/foo/bar");
396    }
397
398    #[test]
399    fn regex() {
400        let path = "/2021/10/21/2021/12/02/2022/01/13";
401        let mut rw = RegexAll(
402            LibRegex::new(r"(?P<y>\d{4})/(?P<m>\d{2})/(?P<d>\d{2})").unwrap(),
403            "$m-$d-$y",
404        );
405        assert_eq!(rw.rewrite(path), "/10-21-2021/12-02-2021/01-13-2022");
406
407        let path = "/2021/10/21/2021/12/02/2022/01/13";
408        let mut rw = RegexN(
409            LibRegex::new(r"(?P<y>\d{4})/(?P<m>\d{2})/(?P<d>\d{2})").unwrap(),
410            "$m-$d-$y",
411            2,
412        );
413        assert_eq!(rw.rewrite(path), "/10-21-2021/12-02-2021/2022/01/13");
414    }
415
416    #[test]
417    fn owned_strings() {
418        let mut rw = AppendPrefix(String::from("/baz"));
419        let mut clone = rw.clone();
420        assert_eq!(rw.rewrite("/foo/bar"), "/baz/foo/bar");
421        assert_eq!(clone.rewrite("/foo/bar"), "/baz/foo/bar");
422    }
423
424    #[test]
425    fn func() {
426        let path = "/abcdefg";
427        let mut rw = Func(|path: &str| path.len().to_string());
428        assert_eq!(rw.rewrite(path), "8");
429    }
430}
tower_proxy/rewrite.rs

tower_proxy/
rewrite.rs
