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
56    fn into_upstream(self) -> Self {
57        self
58    }
59}
60
61impl IntoUpstreamSelector for &str {
62    type UpstreamSelector = Url;
63
64    fn into_upstream(self) -> Url {
65        let url = match Url::try_from(self) {
66            Ok(url) => url,
67            Err(_) => panic!("could not convert proxy target into a url"),
68        };
69
70        assert!(!url.cannot_be_a_base(), "{url} cannot be a base");
71        url
72    }
73}
74
75impl IntoUpstreamSelector for String {
76    type UpstreamSelector = Url;
77
78    fn into_upstream(self) -> Url {
79        (&*self).into_upstream()
80    }
81}
82
83#[derive(Debug, Copy, Clone)]
84/// an upstream selector for forward proxy
85pub struct ForwardProxy;
86impl UpstreamSelector for ForwardProxy {
87    fn determine_upstream(&self, conn: &mut Conn) -> Option<Url> {
88        conn.path_and_query().parse().ok()
89    }
90}
91
92impl UpstreamSelector for Url {
93    fn determine_upstream(&self, conn: &mut Conn) -> Option<Url> {
94        self.join(conn.path_and_query().trim_start_matches('/'))
95            .ok()
96    }
97}
98
99impl<F> UpstreamSelector for F
100where
101    F: Fn(&mut Conn) -> Option<Url> + Debug + Send + Sync + 'static,
102{
103    fn determine_upstream(&self, conn: &mut Conn) -> Option<Url> {
104        self(conn)
105    }
106}