Skip to main content

trillium_proxy/
upstream.rs

1//! Upstream selectors
2use std::fmt::Debug;
3use trillium::Conn;
4use url::Url;
5
6#[cfg(feature = "upstream-connection-counting")]
7mod connection_counting;
8#[cfg(feature = "upstream-random")]
9mod random;
10mod round_robin;
11
12#[cfg(feature = "upstream-connection-counting")]
13pub use connection_counting::ConnectionCounting;
14#[cfg(feature = "upstream-random")]
15pub use random::RandomSelector;
16pub use round_robin::RoundRobin;
17
18/// a trait for selecting the correct upstream
19pub trait UpstreamSelector: Debug + Send + Sync + 'static {
20    /// does what it says on the label
21    fn determine_upstream(&self, conn: &mut Conn) -> Option<Url>;
22
23    /// turn self into a `Box<dyn UpstreamSelector>`
24    fn boxed(self) -> Box<dyn UpstreamSelector>
25    where
26        Self: Sized,
27    {
28        Box::new(self.into_upstream())
29    }
30}
31
32impl UpstreamSelector for Box<dyn UpstreamSelector> {
33    fn determine_upstream(&self, conn: &mut Conn) -> Option<Url> {
34        UpstreamSelector::determine_upstream(&**self, conn)
35    }
36
37    fn boxed(self) -> Box<dyn UpstreamSelector> {
38        self
39    }
40}
41
42/// represents something that can be used as an upstream selector
43///
44/// This primarily exists so &str can be used as a synonym for `Url`.
45/// All `UpstreamSelector`s also are `IntoUpstreamSelector`
46pub trait IntoUpstreamSelector {
47    /// the type that Self will be transformed into
48    type UpstreamSelector: UpstreamSelector;
49    /// transform self into the upstream selector
50    fn into_upstream(self) -> Self::UpstreamSelector;
51}
52
53impl<U: UpstreamSelector> IntoUpstreamSelector for U {
54    type UpstreamSelector = Self;
55    fn into_upstream(self) -> Self {
56        self
57    }
58}
59
60impl IntoUpstreamSelector for &str {
61    type UpstreamSelector = Url;
62    fn into_upstream(self) -> Url {
63        let url = match Url::try_from(self) {
64            Ok(url) => url,
65            Err(_) => panic!("could not convert proxy target into a url"),
66        };
67
68        assert!(!url.cannot_be_a_base(), "{url} cannot be a base");
69        url
70    }
71}
72
73impl IntoUpstreamSelector for String {
74    type UpstreamSelector = Url;
75    fn into_upstream(self) -> Url {
76        (&*self).into_upstream()
77    }
78}
79
80#[derive(Debug, Copy, Clone)]
81/// an upstream selector for forward proxy
82pub struct ForwardProxy;
83impl UpstreamSelector for ForwardProxy {
84    fn determine_upstream(&self, conn: &mut Conn) -> Option<Url> {
85        conn.inner().path_and_query().parse().ok()
86    }
87}
88
89impl UpstreamSelector for Url {
90    fn determine_upstream(&self, conn: &mut Conn) -> Option<Url> {
91        self.join(conn.inner().path_and_query().trim_start_matches('/'))
92            .ok()
93    }
94}
95
96impl<F> UpstreamSelector for F
97where
98    F: Fn(&mut Conn) -> Option<Url> + Debug + Send + Sync + 'static,
99{
100    fn determine_upstream(&self, conn: &mut Conn) -> Option<Url> {
101        self(conn)
102    }
103}