ferron_common/http_proxy/
builder.rs1use 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
17pub 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 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 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 pub fn lb_algorithm(mut self, algorithm: LoadBalancerAlgorithm) -> Self {
108 self.lb_algorithm = algorithm;
109 self
110 }
111
112 pub fn lb_health_check_window(mut self, window: Duration) -> Self {
114 self.lb_health_check_window = window;
115 self
116 }
117
118 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 pub fn lb_health_check(mut self, enable: bool) -> Self {
126 self.lb_health_check = enable;
127 self
128 }
129
130 pub fn proxy_no_verification(mut self, no_verification: bool) -> Self {
132 self.proxy_no_verification = no_verification;
133 self
134 }
135
136 pub fn proxy_intercept_errors(mut self, intercept_errors: bool) -> Self {
138 self.proxy_intercept_errors = intercept_errors;
139 self
140 }
141
142 pub fn lb_retry_connection(mut self, retry: bool) -> Self {
144 self.lb_retry_connection = retry;
145 self
146 }
147
148 pub fn proxy_http2_only(mut self, http2_only: bool) -> Self {
150 self.proxy_http2_only = http2_only;
151 self
152 }
153
154 pub fn proxy_http2(mut self, http2: bool) -> Self {
156 self.proxy_http2 = http2;
157 self
158 }
159
160 pub fn proxy_keepalive(mut self, keepalive: bool) -> Self {
162 self.proxy_keepalive = keepalive;
163 self
164 }
165
166 pub fn proxy_proxy_header(mut self, proxy_header: Option<ProxyHeader>) -> Self {
168 self.proxy_proxy_header = proxy_header;
169 self
170 }
171
172 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 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 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 pub fn rewrite_host(mut self, rewrite_host: bool) -> Self {
192 self.rewrite_host = rewrite_host;
193 self
194 }
195
196 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), };
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}