From 0c7da328e446518036b9420db3e75451fca522e9 Mon Sep 17 00:00:00 2001 From: Georgios Konstantopoulos Date: Fri, 24 Jul 2026 11:41:57 -0700 Subject: [PATCH] feat(proxy): bound accepted client connections --- Cargo.toml | 2 +- src/proxy/builder.rs | 23 +++++++++++++++++++++++ src/proxy/mod.rs | 21 +++++++++++++++++++-- 3 files changed, 43 insertions(+), 3 deletions(-) diff --git a/Cargo.toml b/Cargo.toml index 69fbc39..49ff457 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -34,7 +34,7 @@ rand = { version = "0.10.2", optional = true } rcgen = { version = "0.14.0", features = ["x509-parser"], optional = true } thiserror = "2.0.7" time = { version = "0.3.35", optional = true } -tokio = { version = "1.24.2", features = ["macros", "rt"] } +tokio = { version = "1.24.2", features = ["macros", "rt", "sync"] } tokio-graceful = "0.2.0" tokio-native-tls = { version = "0.3.1", optional = true } tokio-rustls = { version = "0.26.3", features = ["logging", "tls12"] } diff --git a/src/proxy/builder.rs b/src/proxy/builder.rs index 9e2bb48..2a79c8e 100644 --- a/src/proxy/builder.rs +++ b/src/proxy/builder.rs @@ -13,6 +13,7 @@ use hyper_util::{ use std::{ future::{Pending, pending}, net::SocketAddr, + num::NonZeroUsize, sync::Arc, }; use thiserror::Error; @@ -141,6 +142,7 @@ impl ProxyBuilder> { return ProxyBuilder(WantsHandlers { al: self.0.al, ca: self.0.ca, + max_concurrent_connections: None, http_connector: Err(Error::from(e)), client: None, http_handler: NoopHandler::new(), @@ -165,6 +167,7 @@ impl ProxyBuilder> { ProxyBuilder(WantsHandlers { al: self.0.al, ca: self.0.ca, + max_concurrent_connections: None, http_connector: Ok(https), client: None, http_handler: NoopHandler::new(), @@ -189,6 +192,7 @@ impl ProxyBuilder> { return ProxyBuilder(WantsHandlers { al: self.0.al, ca: self.0.ca, + max_concurrent_connections: None, http_connector: Err(Error::from(e)), client: None, http_handler: NoopHandler::new(), @@ -206,6 +210,7 @@ impl ProxyBuilder> { ProxyBuilder(WantsHandlers { al: self.0.al, ca: self.0.ca, + max_concurrent_connections: None, http_connector: Ok(https), client: None, http_handler: NoopHandler::new(), @@ -227,6 +232,7 @@ impl ProxyBuilder> { ProxyBuilder(WantsHandlers { al: self.0.al, ca: self.0.ca, + max_concurrent_connections: None, http_connector: Ok(connector), client: None, http_handler: NoopHandler::new(), @@ -242,6 +248,7 @@ impl ProxyBuilder> { pub struct WantsHandlers { al: AddrOrListener, ca: CA, + max_concurrent_connections: Option, http_connector: Result, client: Option, http_handler: H, @@ -260,6 +267,7 @@ impl ProxyBuilder> { ProxyBuilder(WantsHandlers { al: self.0.al, ca: self.0.ca, + max_concurrent_connections: self.0.max_concurrent_connections, http_connector: self.0.http_connector, client: self.0.client, http_handler, @@ -278,6 +286,7 @@ impl ProxyBuilder> { ProxyBuilder(WantsHandlers { al: self.0.al, ca: self.0.ca, + max_concurrent_connections: self.0.max_concurrent_connections, http_connector: self.0.http_connector, client: self.0.client, http_handler: self.0.http_handler, @@ -312,6 +321,18 @@ impl ProxyBuilder> { }) } + /// Set the maximum number of concurrently accepted client connections. + /// + /// The proxy waits for a connection slot before accepting another socket. + /// This bounds the file descriptors held by the process and applies + /// backpressure through the listener backlog. + pub fn with_max_concurrent_connections(self, max_concurrent_connections: NonZeroUsize) -> Self { + ProxyBuilder(WantsHandlers { + max_concurrent_connections: Some(max_concurrent_connections), + ..self.0 + }) + } + /// Set a future that when ready will gracefully shutdown the proxy server. pub fn with_graceful_shutdown + Send + 'static>( self, @@ -320,6 +341,7 @@ impl ProxyBuilder> { ProxyBuilder(WantsHandlers { al: self.0.al, ca: self.0.ca, + max_concurrent_connections: self.0.max_concurrent_connections, http_connector: self.0.http_connector, client: self.0.client, http_handler: self.0.http_handler, @@ -338,6 +360,7 @@ impl ProxyBuilder> { Ok(Proxy { al: self.0.al, ca: Arc::new(self.0.ca), + max_concurrent_connections: self.0.max_concurrent_connections, http_connector: self.0.http_connector?, client: self.0.client, http_handler: self.0.http_handler, diff --git a/src/proxy/mod.rs b/src/proxy/mod.rs index 7425400..b7331aa 100644 --- a/src/proxy/mod.rs +++ b/src/proxy/mod.rs @@ -17,8 +17,8 @@ use hyper_util::{ server::conn::auto::Builder as ServerBuilder, }; use internal::InternalProxy; -use std::{error::Error as StdError, sync::Arc}; -use tokio::net::TcpListener; +use std::{error::Error as StdError, num::NonZeroUsize, sync::Arc}; +use tokio::{net::TcpListener, sync::Semaphore}; use tokio_graceful::Shutdown; use tokio_tungstenite::Connector; use tracing::error; @@ -73,6 +73,7 @@ use tracing::error; pub struct Proxy { al: AddrOrListener, ca: Arc, + max_concurrent_connections: Option, http_connector: C, client: Option, http_handler: H, @@ -130,8 +131,23 @@ where let shutdown = Shutdown::new(self.graceful_shutdown); let guard = shutdown.guard_weak(); + let connection_permits = self + .max_concurrent_connections + .map(|maximum| Arc::new(Semaphore::new(maximum.get()))); loop { + let connection_permit = if let Some(permits) = &connection_permits { + let permit = tokio::select! { + permit = Arc::clone(permits).acquire_owned() => permit, + _ = guard.cancelled() => break, + }; + let Ok(permit) = permit else { + break; + }; + Some(permit) + } else { + None + }; tokio::select! { res = listener.accept() => { let (tcp, client_addr) = match res { @@ -150,6 +166,7 @@ where let websocket_connector = self.websocket_connector.clone(); shutdown.spawn_task_fn(move |guard| async move { + let _connection_permit = connection_permit; let conn = server.serve_connection_with_upgrades( TokioIo::new(tcp), service_fn(|req| {