1use crate::config::{ProxyConfig, RateLimitKey as RouteRateLimitKey, Route, ServerConfig};
6use crate::health::HealthChecker;
7use crate::ratelimit::{per_second, RateLimitKey, RateLimiter};
8use crate::vendored::backend::BackendAppManager;
9use crate::vendored::conn::{ConnActivity, IdleIo, InFlight};
10use crate::vendored::forwarder::ForwarderClient;
11use crate::vendored::handler::MessageHandler;
12use crate::vendored::hyper_ext::timer::Timeout;
13use crate::vendored::hyper_ext::{IncomingBodyExt, ProxyBody};
14use crate::vendored::types::PathName;
15use http_body_util::BodyExt;
16use hyper::body::Incoming;
17use hyper::server::conn::http2;
18use hyper::{Request, Response};
19use hyper_util::rt::{TokioExecutor, TokioIo, TokioTimer};
20use hyper_util::server::conn::auto;
21use std::convert::Infallible;
22use std::future::Future;
23use std::net::SocketAddr;
24use std::pin::Pin;
25use std::sync::atomic::{AtomicBool, Ordering};
26use std::sync::{Arc, OnceLock};
27use std::time::Duration;
28use tokio::net::{TcpListener, TcpStream};
29use tokio::sync::{RwLock, Semaphore};
30use tokio_util::sync::CancellationToken;
31use tracing::{debug, error, info, warn};
32use uuid::Uuid;
33
34pub struct ProxyService {
36 backend_manager: Arc<BackendAppManager>,
38 forwarder: Arc<ForwarderClient>,
40 rate_limiter: Arc<RateLimiter>,
42 config: Arc<RwLock<ProxyConfig>>,
44 connection_slots: OnceLock<Arc<Semaphore>>,
47 cleanup_started: AtomicBool,
49}
50
51#[derive(Clone, Debug)]
54struct ListenerSettings {
55 header_read_timeout: Duration,
56 idle_timeout: Duration,
57 http2_keepalive_interval: Option<Duration>,
58 http2_keepalive_timeout: Duration,
59}
60
61impl ListenerSettings {
62 fn from_server(server: &ServerConfig) -> Self {
63 let at_least_a_second = |secs: u64| Duration::from_secs(secs.max(1));
67 Self {
68 header_read_timeout: at_least_a_second(server.header_read_timeout_secs),
69 idle_timeout: at_least_a_second(server.idle_timeout_secs),
70 http2_keepalive_interval: match server.http2_keepalive_interval_secs {
71 0 => None,
72 secs => Some(Duration::from_secs(secs)),
73 },
74 http2_keepalive_timeout: at_least_a_second(server.http2_keepalive_timeout_secs),
75 }
76 }
77
78 fn auto_builder(&self) -> auto::Builder<TokioExecutor> {
84 let mut builder = auto::Builder::new(TokioExecutor::new());
85 builder
86 .http1()
87 .timer(TokioTimer::new())
88 .header_read_timeout(self.header_read_timeout);
89 builder
92 .http2()
93 .timer(TokioTimer::new())
94 .enable_connect_protocol()
95 .keep_alive_interval(self.http2_keepalive_interval)
96 .keep_alive_timeout(self.http2_keepalive_timeout);
97 builder
98 }
99
100 fn http2_builder(&self) -> http2::Builder<TokioExecutor> {
102 let mut builder = http2::Builder::new(TokioExecutor::new());
103 builder
104 .timer(TokioTimer::new())
105 .enable_connect_protocol()
106 .keep_alive_interval(self.http2_keepalive_interval)
107 .keep_alive_timeout(self.http2_keepalive_timeout);
108 builder
109 }
110}
111
112#[derive(Clone)]
119struct ConnectionService {
120 service: Arc<ProxyService>,
121 activity: Arc<ConnActivity>,
122 client_addr: SocketAddr,
123 proto: &'static str,
124}
125
126impl hyper::service::Service<Request<Incoming>> for ConnectionService {
127 type Response = Response<ProxyBody>;
128 type Error = Infallible;
129 type Future = Pin<Box<dyn Future<Output = Result<Self::Response, Self::Error>> + Send>>;
130
131 fn call(&self, mut request: Request<Incoming>) -> Self::Future {
132 let in_flight = InFlight::new(&self.activity);
133 request.extensions_mut().insert(in_flight.clone());
136 let service = self.service.clone();
137 let client_addr = self.client_addr;
138 let proto = self.proto;
139 Box::pin(async move {
140 let response = service.handle_request(request, client_addr, proto).await?;
141 Ok(response.map(|body| {
142 body.map_frame(move |frame| {
143 let _held_until_the_body_is_dropped = &in_flight;
144 frame
145 })
146 .boxed()
147 }))
148 })
149 }
150}
151
152fn route_scope(route: &Route) -> Uuid {
157 route.id.unwrap_or_else(|| {
158 Uuid::new_v5(
159 &Uuid::NAMESPACE_OID,
160 format!("postrust-proxy route {}", route.name).as_bytes(),
161 )
162 })
163}
164
165impl ProxyService {
166 pub fn new(
171 config: Arc<RwLock<ProxyConfig>>,
172 health_checker: Arc<HealthChecker>,
173 rate_limiter: Arc<RateLimiter>,
174 ) -> Self {
175 let backend_manager =
176 Arc::new(BackendAppManager::new().with_health_checker(health_checker));
177 let forwarder = Arc::new(ForwarderClient::default());
178
179 Self {
180 backend_manager,
181 forwarder,
182 rate_limiter,
183 config,
184 connection_slots: OnceLock::new(),
185 cleanup_started: AtomicBool::new(false),
186 }
187 }
188
189 pub async fn load_config(&self) {
191 let config = self.config.read().await;
192
193 for upstream in &config.upstreams {
195 self.backend_manager.register_upstream(upstream.clone());
196 }
197
198 for route in &config.routes {
200 if let Some(upstream) = config.upstreams.iter().find(|u| u.name == route.upstream) {
202 use crate::vendored::types::ServerName;
203 let host = route.match_.host.as_deref().unwrap_or("*");
204 let path = route.match_.path.as_deref().unwrap_or("/");
205 let host = ServerName::new(host);
206 let path = PathName::new(path);
207 self.backend_manager
208 .register_route(host, path, upstream.resolved_id());
209 }
210 }
211
212 info!(
213 "Loaded {} routes and {} upstreams",
214 config.routes.len(),
215 config.upstreams.len()
216 );
217 }
218
219 pub async fn serve_http(
221 self: Arc<Self>,
222 addr: SocketAddr,
223 cancel_token: CancellationToken,
224 ) -> std::io::Result<()> {
225 let listener = TcpListener::bind(addr).await?;
226 info!("HTTP proxy listening on {}", addr);
227 self.serve_http_listener(listener, cancel_token).await
228 }
229
230 pub async fn serve_http_listener(
233 self: Arc<Self>,
234 listener: TcpListener,
235 cancel_token: CancellationToken,
236 ) -> std::io::Result<()> {
237 self.run_listener(
238 listener,
239 cancel_token,
240 "HTTP proxy",
241 |service, io, activity, client_addr, settings| async move {
242 let conn = ConnectionService {
243 service,
244 activity,
245 client_addr,
246 proto: "http",
247 };
248 if let Err(err) = settings
252 .auto_builder()
253 .serve_connection_with_upgrades(TokioIo::new(io), conn)
254 .await
255 {
256 debug!("Connection error: {}", err);
257 }
258 },
259 )
260 .await
261 }
262
263 pub async fn serve_https(
273 self: Arc<Self>,
274 addr: SocketAddr,
275 tls_config: Arc<tokio_rustls::rustls::ServerConfig>,
276 cancel_token: CancellationToken,
277 ) -> std::io::Result<()> {
278 let listener = TcpListener::bind(addr).await?;
279 info!("HTTPS proxy listening on {} (ALPN: h2, http/1.1)", addr);
280 self.serve_https_listener(listener, tls_config, cancel_token)
281 .await
282 }
283
284 pub async fn serve_https_listener(
286 self: Arc<Self>,
287 listener: TcpListener,
288 tls_config: Arc<tokio_rustls::rustls::ServerConfig>,
289 cancel_token: CancellationToken,
290 ) -> std::io::Result<()> {
291 let acceptor = tokio_rustls::TlsAcceptor::from(tls_config);
292 self.run_listener(
293 listener,
294 cancel_token,
295 "HTTPS proxy",
296 move |service, io, activity, client_addr, settings| {
297 let acceptor = acceptor.clone();
298 async move {
299 let tls_stream = match acceptor.accept(io).await {
303 Ok(s) => s,
304 Err(e) => {
305 debug!("TLS handshake failed: {}", e);
306 return;
307 }
308 };
309
310 let is_h2 = tls_stream
311 .get_ref()
312 .1
313 .alpn_protocol()
314 .map(|p| p == b"h2")
315 .unwrap_or(false);
316
317 let conn = ConnectionService {
318 service,
319 activity,
320 client_addr,
321 proto: "https",
322 };
323
324 let io = TokioIo::new(tls_stream);
325 let result = if is_h2 {
326 settings
327 .http2_builder()
328 .serve_connection(io, conn)
329 .await
330 .map_err(|e| {
331 Box::<dyn std::error::Error + Send + Sync>::from(e.to_string())
332 })
333 } else {
334 settings
338 .auto_builder()
339 .serve_connection_with_upgrades(io, conn)
340 .await
341 .map_err(|e| {
342 Box::<dyn std::error::Error + Send + Sync>::from(e.to_string())
343 })
344 };
345
346 if let Err(err) = result {
347 debug!("HTTPS connection error: {}", err);
348 }
349 }
350 },
351 )
352 .await
353 }
354
355 pub async fn serve_h2c(
366 self: Arc<Self>,
367 addr: SocketAddr,
368 cancel_token: CancellationToken,
369 ) -> std::io::Result<()> {
370 let listener = TcpListener::bind(addr).await?;
371 info!("HTTP/2-only proxy listening on {}", addr);
372 self.serve_h2c_listener(listener, cancel_token).await
373 }
374
375 pub async fn serve_h2c_listener(
377 self: Arc<Self>,
378 listener: TcpListener,
379 cancel_token: CancellationToken,
380 ) -> std::io::Result<()> {
381 self.run_listener(
382 listener,
383 cancel_token,
384 "HTTP/2-only proxy",
385 |service, io, activity, client_addr, settings| async move {
386 let conn = ConnectionService {
387 service,
388 activity,
389 client_addr,
390 proto: "http",
391 };
392 if let Err(err) = settings
393 .http2_builder()
394 .serve_connection(TokioIo::new(io), conn)
395 .await
396 {
397 debug!("HTTP/2 connection error: {}", err);
398 }
399 },
400 )
401 .await
402 }
403
404 async fn run_listener<F, Fut>(
413 self: Arc<Self>,
414 listener: TcpListener,
415 cancel_token: CancellationToken,
416 what: &'static str,
417 serve: F,
418 ) -> std::io::Result<()>
419 where
420 F: Fn(Arc<Self>, IdleIo<TcpStream>, Arc<ConnActivity>, SocketAddr, ListenerSettings) -> Fut,
421 Fut: Future<Output = ()> + Send + 'static,
422 {
423 let (settings, max_connections) = {
424 let config = self.config.read().await;
425 (
426 ListenerSettings::from_server(&config.server),
427 config.server.max_connections,
428 )
429 };
430 let slots = self
431 .connection_slots
432 .get_or_init(|| {
433 Arc::new(Semaphore::new(
434 max_connections.clamp(1, Semaphore::MAX_PERMITS),
435 ))
436 })
437 .clone();
438 self.start_background_tasks(&cancel_token);
439
440 loop {
441 let slot = tokio::select! {
442 _ = cancel_token.cancelled() => break,
443 slot = slots.clone().acquire_owned() => match slot {
444 Ok(slot) => slot,
445 Err(_) => break,
447 },
448 };
449
450 let accepted = tokio::select! {
451 _ = cancel_token.cancelled() => break,
452 accepted = listener.accept() => accepted,
453 };
454
455 match accepted {
456 Ok((stream, client_addr)) => {
457 if let Err(e) = stream.set_nodelay(true) {
460 debug!("Could not set TCP_NODELAY on the client socket: {}", e);
461 }
462 let activity = ConnActivity::new(Some(slot));
463 let io = IdleIo::new(stream, activity.clone(), settings.idle_timeout);
464 tokio::spawn(serve(
465 self.clone(),
466 io,
467 activity,
468 client_addr,
469 settings.clone(),
470 ));
471 }
472 Err(e) => {
473 error!("Accept error: {}", e);
477 drop(slot);
478 tokio::select! {
479 _ = cancel_token.cancelled() => break,
480 _ = tokio::time::sleep(Duration::from_millis(100)) => {}
481 }
482 }
483 }
484 }
485
486 info!("{} stopped", what);
487 Ok(())
488 }
489
490 fn start_background_tasks(&self, cancel_token: &CancellationToken) {
496 if !self.cleanup_started.swap(true, Ordering::AcqRel) {
497 tokio::spawn(
498 self.rate_limiter
499 .clone()
500 .start_cleanup(cancel_token.clone()),
501 );
502 }
503 }
504
505 async fn handle_request(
507 &self,
508 request: Request<Incoming>,
509 client_addr: SocketAddr,
510 proto: &str,
511 ) -> Result<Response<ProxyBody>, Infallible> {
512 let uri = request.uri().clone();
513 let method = request.method().clone();
514
515 let host = MessageHandler::request_host(&request);
519 let host = host.as_str();
520
521 let path = uri.path();
522
523 debug!("{} {} {} from {}", method, host, path, client_addr);
524
525 if let Some(reason) = MessageHandler::unsafe_path_reason(path) {
527 debug!("Refusing {} from {}: {}", path, client_addr, reason);
528 return Ok(MessageHandler::bad_request(reason));
529 }
530
531 let (route, upstream, server) = {
533 let config = self.config.read().await;
534 let server = (
535 config.server.allowed_upgrade_protocols.clone(),
536 config.server.trust_forwarded_headers,
537 Duration::from_secs(config.server.tunnel_idle_timeout_secs.max(1)),
538 );
539 let matched = config
545 .routes
546 .iter()
547 .filter(|r| r.enabled)
548 .filter(|r| {
549 r.match_
550 .matches(host, path, method.as_str(), request.headers())
551 })
552 .max_by_key(|r| {
553 let path_len = r.match_.path.as_ref().map(|p| p.len()).unwrap_or(0);
554 (r.priority, path_len)
555 })
556 .cloned();
557
558 match matched {
559 Some(r) => {
560 let upstream = config
561 .upstreams
562 .iter()
563 .find(|u| u.name == r.upstream)
564 .map(|u| (u.resolved_id(), u.enabled));
565 (Some(r), upstream, server)
566 }
567 None => (None, None, server),
568 }
569 };
570 let (allowed_upgrades, trust_forwarded, tunnel_idle) = server;
571
572 let route = match route {
573 Some(r) => r,
574 None => {
575 debug!("No route found for {} {}", host, path);
576 return Ok(MessageHandler::not_found());
577 }
578 };
579
580 let upstream_id = match upstream {
581 Some((id, true)) => id,
582 Some((_, false)) => {
583 debug!(
584 "Upstream '{}' for route {} is disabled",
585 route.upstream, route.name
586 );
587 return Ok(MessageHandler::service_unavailable("Upstream disabled"));
588 }
589 None => {
590 warn!(
591 "Upstream '{}' not found for route {}",
592 route.upstream, route.name
593 );
594 return Ok(MessageHandler::service_unavailable("Upstream not found"));
595 }
596 };
597
598 if let Some(ref rate_limit) = route.rate_limit {
602 let scope = route_scope(&route);
603 let by_ip = || RateLimitKey::client_ip(client_addr.ip());
604 let key = match &rate_limit.key {
605 RouteRateLimitKey::ClientIp => by_ip(),
606 RouteRateLimitKey::Route => RateLimitKey::Route(scope),
607 RouteRateLimitKey::Header(name) => {
611 let values: Vec<&[u8]> = request
612 .headers()
613 .get_all(name.as_str())
614 .iter()
615 .map(|v| v.as_bytes())
616 .collect();
617 if values.is_empty() {
618 by_ip()
619 } else {
620 RateLimitKey::header_value(&values.join(&b',')[..])
621 }
622 }
623 };
624 let rate = per_second(rate_limit.requests, rate_limit.window_secs);
625 if !self
626 .rate_limiter
627 .check_with_config(scope, key, rate, rate_limit.requests)
628 {
629 return Ok(MessageHandler::too_many_requests());
630 }
631 }
632
633 let backend = match self.backend_manager.select_backend(upstream_id, None) {
635 Some(b) => b,
636 None => {
637 warn!("No enabled backend for route {}", route.name);
638 return Ok(MessageHandler::service_unavailable("No backend available"));
639 }
640 };
641
642 let (parts, body) = request.into_parts();
644 let mut forwarded_request = Request::from_parts(parts, body.boxed_body());
645
646 let upgrade_protocol = MessageHandler::upgrade_protocol(&forwarded_request).filter(|p| {
652 let allowed = MessageHandler::upgrade_allowed(p, &allowed_upgrades);
653 if !allowed {
654 debug!("Declining an upgrade to {:?} from {}", p, client_addr);
655 }
656 allowed
657 });
658 let h2_protocol = MessageHandler::h2_connect_protocol(&forwarded_request);
663 if let Some(ref protocol) = h2_protocol {
664 if !MessageHandler::upgrade_allowed(protocol, &allowed_upgrades) {
665 debug!(
666 "Refusing extended CONNECT for {:?} from {}",
667 protocol, client_addr
668 );
669 return Ok(MessageHandler::not_implemented(
670 "Extended CONNECT protocol not supported",
671 ));
672 }
673 }
674
675 MessageHandler::strip_hop_by_hop_headers(&mut forwarded_request);
678
679 MessageHandler::add_forwarding_headers(
681 &mut forwarded_request,
682 client_addr,
683 proto,
684 trust_forwarded,
685 );
686
687 if route.strip_path {
689 if let Some(ref prefix) = route.match_.path {
690 MessageHandler::strip_path_prefix(&mut forwarded_request, prefix);
691 }
692 }
693
694 MessageHandler::apply_route_headers(&mut forwarded_request, &route);
696
697 MessageHandler::rewrite_host_header(&mut forwarded_request, &backend.address);
699
700 if let Some(protocol) = h2_protocol {
702 debug!(
703 "Extended CONNECT to {} via {} (RFC 8441)",
704 protocol, backend.address
705 );
706 return match self
707 .forwarder
708 .forward_h2_websocket(&backend, forwarded_request, &protocol, tunnel_idle)
709 .await
710 {
711 Ok(response) => Ok(response),
712 Err(e) => {
713 error!("Extended CONNECT error to {}: {}", backend.address, e);
714 Ok(MessageHandler::bad_gateway(&format!(
715 "Extended CONNECT error: {}",
716 e
717 )))
718 }
719 };
720 }
721
722 if let Some(protocol) = upgrade_protocol {
724 MessageHandler::restore_upgrade_headers(&mut forwarded_request, &protocol);
725 debug!(
726 "Upgrading connection to {} via {}",
727 protocol, backend.address
728 );
729 return match self
730 .forwarder
731 .forward_upgrade(&backend, forwarded_request, tunnel_idle)
732 .await
733 {
734 Ok(response) => Ok(response),
735 Err(e) => {
736 error!("Upgrade error to {}: {}", backend.address, e);
737 Ok(MessageHandler::bad_gateway(&format!(
738 "Upgrade error: {}",
739 e
740 )))
741 }
742 };
743 }
744
745 forwarded_request.extensions_mut().remove::<InFlight>();
759 let timeout = Duration::from_secs(u64::from(route.timeout_secs));
760 let forwarded = Timeout::new(self.forwarder.forward(&backend, forwarded_request), timeout);
761
762 match forwarded.await {
763 Ok(Ok(response)) => {
764 let (parts, body) = response.into_parts();
765 Ok(Response::from_parts(parts, body.boxed_body()))
766 }
767 Ok(Err(e)) => {
768 error!("Forward error to {}: {}", backend.address, e);
769 Ok(MessageHandler::bad_gateway(&format!(
770 "Backend error: {}",
771 e
772 )))
773 }
774 Err(_) => {
775 error!(
776 "Upstream {} did not answer within {}s (route {})",
777 backend.address, route.timeout_secs, route.name
778 );
779 Ok(MessageHandler::gateway_timeout(&format!(
780 "Upstream did not answer within {}s",
781 route.timeout_secs
782 )))
783 }
784 }
785 }
786}