Skip to main content

ferron_common/http_proxy/
builder.rs

1use std::sync::atomic::AtomicUsize;
2use std::sync::Arc;
3use std::time::Duration;
4use std::{collections::HashMap, net::IpAddr};
5
6use hickory_resolver::config::{NameServerConfig, ResolverConfig};
7use hickory_resolver::net::runtime::TokioRuntimeProvider;
8use hyper::header::HeaderName;
9use tokio::sync::RwLock;
10
11use super::{Connections, LoadBalancerAlgorithm, LoadBalancerAlgorithmInner, ProxyHeader, ProxyToKey, ReverseProxy};
12use crate::{
13  http_proxy::{SrvUpstreamData, Upstream, UpstreamInner},
14  util::TtlCache,
15};
16
17/// Builder for configuring and constructing a [`ReverseProxy`].
18pub struct ReverseProxyBuilder<'a> {
19  pub(super) connections: &'a mut Connections,
20  #[allow(clippy::type_complexity)]
21  pub(super) upstreams: Vec<(Upstream, Option<usize>, Option<Duration>)>,
22  pub(super) lb_algorithm: LoadBalancerAlgorithm,
23  pub(super) lb_health_check_window: Duration,
24  pub(super) lb_health_check_max_fails: u64,
25  pub(super) lb_health_check: bool,
26  pub(super) lb_retry_connection: bool,
27  pub(super) proxy_no_verification: bool,
28  pub(super) proxy_intercept_errors: bool,
29  pub(super) proxy_http2_only: bool,
30  pub(super) proxy_http2: bool,
31  pub(super) proxy_keepalive: bool,
32  pub(super) proxy_proxy_header: Option<ProxyHeader>,
33  pub(super) proxy_request_header: Vec<(HeaderName, String)>,
34  pub(super) proxy_request_header_replace: Vec<(HeaderName, String)>,
35  pub(super) proxy_request_header_remove: Vec<HeaderName>,
36  pub(super) rewrite_host: bool,
37}
38
39impl<'a> ReverseProxyBuilder<'a> {
40  /// Adds an upstream backend target.
41  ///
42  /// `proxy_to` is the backend URL (for example `http://127.0.0.1:8080`).
43  /// `proxy_unix` can be used to target a Unix socket path.
44  /// `local_limit` controls per-upstream connection limit.
45  /// `keepalive_idle_timeout` sets pooled connection idle timeout.
46  pub fn upstream(
47    mut self,
48    proxy_to: String,
49    proxy_unix: Option<String>,
50    local_limit: Option<usize>,
51    keepalive_idle_timeout: Option<Duration>,
52  ) -> Self {
53    self.upstreams.push((
54      Upstream::Static(UpstreamInner { proxy_to, proxy_unix }),
55      local_limit,
56      keepalive_idle_timeout,
57    ));
58    self
59  }
60
61  /// Adds a dynamic (SRV-based) upstream backend target.
62  ///
63  /// `to` is the backend URL (for example `http://_http._tcp.example.com`).
64  /// `local_limit` controls per-upstream connection limit.
65  /// `keepalive_idle_timeout` sets pooled connection idle timeout.
66  pub fn upstream_srv(
67    mut self,
68    to: String,
69    local_limit: Option<usize>,
70    keepalive_idle_timeout: Option<Duration>,
71    secondary_runtime_handle: tokio::runtime::Handle,
72    dns_servers: Vec<IpAddr>,
73  ) -> Self {
74    let dns_resolver = secondary_runtime_handle.block_on(async {
75      if !dns_servers.is_empty() {
76        hickory_resolver::Resolver::builder_with_config(
77          ResolverConfig::from_parts(
78            None,
79            vec![],
80            dns_servers.iter().map(|ip| NameServerConfig::udp(*ip)).collect(),
81          ),
82          TokioRuntimeProvider::default(),
83        )
84        .build()
85      } else {
86        hickory_resolver::Resolver::builder_tokio()
87          .unwrap_or(hickory_resolver::Resolver::builder_with_config(
88            ResolverConfig::default(),
89            TokioRuntimeProvider::default(),
90          ))
91          .build()
92      }
93    });
94    self.upstreams.push((
95      Upstream::Srv(SrvUpstreamData {
96        to,
97        secondary_runtime_handle,
98        dns_resolver: dns_resolver.ok().map(Arc::new),
99      }),
100      local_limit,
101      keepalive_idle_timeout,
102    ));
103    self
104  }
105
106  /// Sets load balancing algorithm.
107  pub fn lb_algorithm(mut self, algorithm: LoadBalancerAlgorithm) -> Self {
108    self.lb_algorithm = algorithm;
109    self
110  }
111
112  /// Sets health-check TTL window for failed backend counters.
113  pub fn lb_health_check_window(mut self, window: Duration) -> Self {
114    self.lb_health_check_window = window;
115    self
116  }
117
118  /// Sets maximum consecutive failed checks before a backend is considered unhealthy.
119  pub fn lb_health_check_max_fails(mut self, max_fails: u64) -> Self {
120    self.lb_health_check_max_fails = max_fails;
121    self
122  }
123
124  /// Enables or disables backend health checks.
125  pub fn lb_health_check(mut self, enable: bool) -> Self {
126    self.lb_health_check = enable;
127    self
128  }
129
130  /// Disables certificate verification for upstream TLS connections.
131  pub fn proxy_no_verification(mut self, no_verification: bool) -> Self {
132    self.proxy_no_verification = no_verification;
133    self
134  }
135
136  /// Intercepts upstream errors and converts them to proxy-generated responses.
137  pub fn proxy_intercept_errors(mut self, intercept_errors: bool) -> Self {
138    self.proxy_intercept_errors = intercept_errors;
139    self
140  }
141
142  /// Enables retrying a different backend when connection setup fails.
143  pub fn lb_retry_connection(mut self, retry: bool) -> Self {
144    self.lb_retry_connection = retry;
145    self
146  }
147
148  /// Forces HTTP/2-only upstream connections.
149  pub fn proxy_http2_only(mut self, http2_only: bool) -> Self {
150    self.proxy_http2_only = http2_only;
151    self
152  }
153
154  /// Enables HTTP/2 support for upstream connections.
155  pub fn proxy_http2(mut self, http2: bool) -> Self {
156    self.proxy_http2 = http2;
157    self
158  }
159
160  /// Enables connection pooling and keepalive reuse.
161  pub fn proxy_keepalive(mut self, keepalive: bool) -> Self {
162    self.proxy_keepalive = keepalive;
163    self
164  }
165
166  /// Sets PROXY protocol header mode for upstream connections.
167  pub fn proxy_proxy_header(mut self, proxy_header: Option<ProxyHeader>) -> Self {
168    self.proxy_proxy_header = proxy_header;
169    self
170  }
171
172  /// Adds a request header to upstream requests.
173  pub fn proxy_request_header(mut self, header_name: HeaderName, header_value: String) -> Self {
174    self.proxy_request_header.push((header_name, header_value));
175    self
176  }
177
178  /// Replaces a request header on upstream requests.
179  pub fn proxy_request_header_replace(mut self, header_name: HeaderName, header_value: String) -> Self {
180    self.proxy_request_header_replace.push((header_name, header_value));
181    self
182  }
183
184  /// Removes a request header from upstream requests.
185  pub fn proxy_request_header_remove(mut self, header_name: HeaderName) -> Self {
186    self.proxy_request_header_remove.push(header_name);
187    self
188  }
189
190  /// Enables or disables `Host` header rewriting for non-HTTPS upstream requests.
191  pub fn rewrite_host(mut self, rewrite_host: bool) -> Self {
192    self.rewrite_host = rewrite_host;
193    self
194  }
195
196  /// Builds a [`ReverseProxy`] from the configured options.
197  pub fn build(mut self) -> ReverseProxy {
198    let connections = self.connections.connections.clone();
199    #[cfg(unix)]
200    let unix_connections = self.connections.unix_connections.clone();
201
202    let proxy_to = self
203      .upstreams
204      .drain(..)
205      .map(|(upstream, local_limit, keepalive_idle_timeout)| {
206        let is_unix_socket = match &upstream {
207          Upstream::Static(inner) => Some(inner.proxy_unix.is_some()),
208          Upstream::Srv(_) => Some(false), // SRV records lead to A/AAAA lookups, so they cannot be Unix sockets
209        };
210        (
211          upstream,
212          is_unix_socket.and_then(|is_unix_socket| {
213            apply_local_limit(
214              local_limit,
215              is_unix_socket,
216              &connections,
217              #[cfg(unix)]
218              &unix_connections,
219            )
220          }),
221          keepalive_idle_timeout,
222        )
223      })
224      .collect::<Vec<ProxyToKey>>();
225
226    let proxy_to = Arc::new(proxy_to);
227    let load_balancer_algorithm = if let Some(algorithm) = self
228      .connections
229      .load_balancer_cache
230      .get(&(self.lb_algorithm, proxy_to.clone()))
231    {
232      algorithm.clone()
233    } else {
234      let new_algorithm = Arc::new(build_load_balancer_algorithm(self.lb_algorithm));
235      self
236        .connections
237        .load_balancer_cache
238        .insert((self.lb_algorithm, proxy_to.clone()), new_algorithm.clone());
239      new_algorithm
240    };
241    let failed_backends = if let Some(failed) = self.connections.failed_backend_cache.get(&(
242      self.lb_health_check_window,
243      self.lb_health_check_max_fails,
244      proxy_to.clone(),
245    )) {
246      failed.clone()
247    } else {
248      let new_failed = Arc::new(RwLock::new(TtlCache::new(self.lb_health_check_window)));
249      self.connections.failed_backend_cache.insert(
250        (
251          self.lb_health_check_window,
252          self.lb_health_check_max_fails,
253          proxy_to.clone(),
254        ),
255        new_failed.clone(),
256      );
257      new_failed
258    };
259    ReverseProxy {
260      failed_backends,
261      load_balancer_algorithm,
262      proxy_to,
263      health_check_max_fails: self.lb_health_check_max_fails,
264      enable_health_check: self.lb_health_check,
265      disable_certificate_verification: self.proxy_no_verification,
266      proxy_intercept_errors: self.proxy_intercept_errors,
267      retry_connection: self.lb_retry_connection,
268      proxy_http2_only: self.proxy_http2_only,
269      proxy_http2: self.proxy_http2,
270      proxy_keepalive: self.proxy_keepalive,
271      proxy_header: self.proxy_proxy_header,
272      headers_to_add: Arc::new(self.proxy_request_header.drain(..).collect()),
273      headers_to_replace: Arc::new(self.proxy_request_header_replace.drain(..).collect()),
274      headers_to_remove: Arc::new(self.proxy_request_header_remove.drain(..).collect()),
275      rewrite_host: self.rewrite_host,
276      connections,
277      #[cfg(unix)]
278      unix_connections,
279    }
280  }
281}
282
283fn build_load_balancer_algorithm(algorithm: LoadBalancerAlgorithm) -> LoadBalancerAlgorithmInner {
284  match algorithm {
285    LoadBalancerAlgorithm::TwoRandomChoices => {
286      LoadBalancerAlgorithmInner::TwoRandomChoices(Arc::new(RwLock::new(HashMap::new())))
287    }
288    LoadBalancerAlgorithm::LeastConnections => {
289      LoadBalancerAlgorithmInner::LeastConnections(Arc::new(RwLock::new(HashMap::new())))
290    }
291    LoadBalancerAlgorithm::RoundRobin => LoadBalancerAlgorithmInner::RoundRobin(Arc::new(AtomicUsize::new(0))),
292    LoadBalancerAlgorithm::Random => LoadBalancerAlgorithmInner::Random,
293  }
294}
295
296fn apply_local_limit(
297  local_limit: Option<usize>,
298  is_unix_socket: bool,
299  connections: &super::ConnectionPool,
300  #[cfg(unix)] unix_connections: &super::ConnectionPool,
301) -> Option<usize> {
302  #[allow(clippy::bind_instead_of_map)]
303  local_limit.and_then(|limit| {
304    if is_unix_socket {
305      #[cfg(unix)]
306      {
307        Some(unix_connections.set_local_limit(limit))
308      }
309      #[cfg(not(unix))]
310      {
311        None
312      }
313    } else {
314      Some(connections.set_local_limit(limit))
315    }
316  })
317}