From 76a08af366534a97b935c8d40158c49ce74ea67b Mon Sep 17 00:00:00 2001 From: Scott Weiss Date: Mon, 9 Nov 2026 10:45:01 -0500 Subject: [PATCH 01/49] fix(controller): time out fleet API TLS handshakes The fleet listener accepts connections from anyone, since enrollment runs before a device has a client certificate. Tonic sets no TLS handshake timeout, so a client could hold a connection open forever by never sending its ClientHello. Add a fleet_limits module for the listener's limits, starting with a handshake timeout, and apply it to the fleet TLS configuration. --- crates/controller/src/application.rs | 11 +- crates/controller/src/fleet_limits.rs | 33 +++++ crates/controller/src/fleet_limits/tests.rs | 144 ++++++++++++++++++++ crates/controller/src/lib.rs | 1 + 4 files changed, 185 insertions(+), 4 deletions(-) create mode 100644 crates/controller/src/fleet_limits.rs create mode 100644 crates/controller/src/fleet_limits/tests.rs diff --git a/crates/controller/src/application.rs b/crates/controller/src/application.rs index 6620999a..cab7a7eb 100644 --- a/crates/controller/src/application.rs +++ b/crates/controller/src/application.rs @@ -2,6 +2,7 @@ use crate::{ daemon_config::{self, DaemonConfigStore}, database::Database, device_ca::DeviceCertificateIssuer, + fleet_limits::FLEET_LIMITS, gateway_jwt::{self, GatewayJwtIssuer}, management::{self, Settings, State}, oidc::OidcProvider, @@ -122,10 +123,12 @@ pub async fn run( tls.client_ca_certificate.display() ) })?; - let tls_config = ServerTlsConfig::new() - .identity(Identity::from_pem(certificate, key)) - .client_ca_root(Certificate::from_pem(client_ca)) - .client_auth_optional(true); + let tls_config = FLEET_LIMITS.tls( + ServerTlsConfig::new() + .identity(Identity::from_pem(certificate, key)) + .client_ca_root(Certificate::from_pem(client_ca)) + .client_auth_optional(true), + ); let mut server = Server::builder() .accept_http1(true) .tls_config(tls_config)?; diff --git a/crates/controller/src/fleet_limits.rs b/crates/controller/src/fleet_limits.rs new file mode 100644 index 00000000..d0c8bf82 --- /dev/null +++ b/crates/controller/src/fleet_limits.rs @@ -0,0 +1,33 @@ +//! Bounds on the device-facing fleet API. +//! +//! The fleet listener has to be reachable by every enrolled device, and client +//! certificates are optional on it because enrollment runs before a device has +//! one. These limits keep a client from holding resources on it indefinitely. + +use std::time::Duration; + +use tonic::transport::ServerTlsConfig; + +/// Limits applied to the fleet gRPC listener. +#[derive(Clone, Copy, Debug)] +pub struct FleetLimits { + /// Time a client has to finish the TLS handshake after connecting. Tonic sets + /// no limit by default, so a client could hold a connection open forever by + /// never sending its ClientHello. + pub tls_handshake: Duration, +} + +/// Production limits. +pub const FLEET_LIMITS: FleetLimits = FleetLimits { + tls_handshake: Duration::from_secs(10), +}; + +impl FleetLimits { + /// Adds the handshake timeout to the fleet listener's TLS settings. + pub fn tls(&self, config: ServerTlsConfig) -> ServerTlsConfig { + config.timeout(self.tls_handshake) + } +} + +#[cfg(test)] +mod tests; diff --git a/crates/controller/src/fleet_limits/tests.rs b/crates/controller/src/fleet_limits/tests.rs new file mode 100644 index 00000000..a5fe4179 --- /dev/null +++ b/crates/controller/src/fleet_limits/tests.rs @@ -0,0 +1,144 @@ +use std::{net::SocketAddr, pin::Pin, time::Duration}; + +use agentdesktop_proto::fleet::{ + AgentMessage, BeginEnrollmentRequest, BeginEnrollmentResponse, CompleteEnrollmentRequest, + ControllerMessage, DeviceCertificateResponse, DeviceUsageRequest, DeviceUsageResponse, + EnrollResponse, LlmGatewayCredentialRequest, LlmGatewayCredentialResponse, + RenewDeviceCertificateRequest, + fleet_agent_server::{FleetAgent, FleetAgentServer}, +}; +use futures_core::Stream; +use tokio::{ + io::AsyncReadExt, + net::{TcpListener, TcpStream}, + sync::mpsc, + time, +}; +use tokio_stream::wrappers::ReceiverStream; +use tonic::{ + Request, Response, Status, Streaming, + transport::{Identity, Server, ServerTlsConfig, server::TcpIncoming}, +}; + +use super::FleetLimits; + +/// Limits small enough for tests to observe. +fn limits() -> FleetLimits { + FleetLimits { + tls_handshake: Duration::from_millis(200), + } +} + +/// Fails the test, instead of hanging it, when the server never ends the wait. +async fn bounded(call: impl std::future::Future) -> T { + time::timeout(Duration::from_secs(5), call) + .await + .expect("the server did not enforce the limit") +} + +/// A fleet service whose calls take `delay`; `Connect` sends a message after each +/// `delay`, twice. +#[derive(Clone)] +struct Slow { + delay: Duration, +} + +type MessageStream = Pin> + Send>>; + +#[tonic::async_trait] +impl FleetAgent for Slow { + type ConnectStream = MessageStream; + + async fn begin_enrollment( + &self, + _: Request, + ) -> Result, Status> { + time::sleep(self.delay).await; + Ok(Response::new(BeginEnrollmentResponse::default())) + } + + async fn complete_enrollment( + &self, + _: Request, + ) -> Result, Status> { + Err(Status::unimplemented("test")) + } + + async fn renew_device_certificate( + &self, + _: Request, + ) -> Result, Status> { + Err(Status::unimplemented("test")) + } + + async fn get_llm_gateway_credential( + &self, + _: Request, + ) -> Result, Status> { + Err(Status::unimplemented("test")) + } + + async fn get_device_usage( + &self, + _: Request, + ) -> Result, Status> { + Err(Status::unimplemented("test")) + } + + async fn connect( + &self, + _: Request>, + ) -> Result, Status> { + let (sender, receiver) = mpsc::channel(2); + let delay = self.delay; + tokio::spawn(async move { + for _ in 0..2 { + time::sleep(delay).await; + if sender.send(Ok(ControllerMessage::default())).await.is_err() { + return; + } + } + }); + Ok(Response::new(Box::pin(ReceiverStream::new(receiver)))) + } +} + +/// Serves `service` with `limits`, over TLS when `identity` is given. +async fn serve(limits: FleetLimits, identity: Option, service: Slow) -> SocketAddr { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let address = listener.local_addr().unwrap(); + let mut server = Server::builder(); + if let Some(identity) = identity { + server = server + .tls_config(limits.tls(ServerTlsConfig::new().identity(identity))) + .unwrap(); + } + let router = server.add_service(FleetAgentServer::new(service)); + tokio::spawn(router.serve_with_incoming(TcpIncoming::from(listener))); + address +} + +/// Reads and discards everything the server sends until it closes the connection. +async fn until_closed(mut stream: TcpStream) { + let mut buffer = [0_u8; 1024]; + while matches!(stream.read(&mut buffer).await, Ok(read) if read > 0) {} +} + +#[tokio::test] +async fn a_client_that_never_finishes_the_tls_handshake_is_disconnected() { + let certified = rcgen::generate_simple_self_signed(vec!["localhost".to_owned()]).unwrap(); + let identity = Identity::from_pem(certified.cert.pem(), certified.signing_key.serialize_pem()); + let address = serve( + limits(), + Some(identity), + Slow { + delay: Duration::ZERO, + }, + ) + .await; + + // Connect and never send a ClientHello. + let stream = TcpStream::connect(address).await.unwrap(); + + bounded(until_closed(stream)).await; +} diff --git a/crates/controller/src/lib.rs b/crates/controller/src/lib.rs index 6e88b029..709fd868 100644 --- a/crates/controller/src/lib.rs +++ b/crates/controller/src/lib.rs @@ -2,6 +2,7 @@ pub mod application; pub mod daemon_config; pub mod database; pub mod device_ca; +pub mod fleet_limits; pub mod gateway_jwt; pub mod management; pub mod oidc; From 3e3092004c1084f5089f11d0f8033c75a2f4a93e Mon Sep 17 00:00:00 2001 From: Scott Weiss Date: Mon, 9 Nov 2026 10:45:39 -0500 Subject: [PATCH 02/49] fix(controller): time out fleet API requests No fleet request had a time limit, so a handler that never completed held its connection and resources indefinitely. Set a request timeout on the fleet server. Tonic applies it until the handler produces its response; a streaming RPC responds once the stream opens, so long-lived Connect streams are not cut off. Tests cover both. --- crates/controller/src/application.rs | 5 +- crates/controller/src/fleet_limits.rs | 12 +++- crates/controller/src/fleet_limits/tests.rs | 71 ++++++++++++++++++--- 3 files changed, 77 insertions(+), 11 deletions(-) diff --git a/crates/controller/src/application.rs b/crates/controller/src/application.rs index cab7a7eb..07ec8466 100644 --- a/crates/controller/src/application.rs +++ b/crates/controller/src/application.rs @@ -14,7 +14,7 @@ use agentdesktop_usage::UsageIngest; use anyhow::Context; use tonic::{ service::Routes, - transport::{Certificate, Identity, Server, ServerTlsConfig}, + transport::{Certificate, Identity, ServerTlsConfig}, }; /// Runs fleet and management listeners with application-owned extensions. @@ -129,7 +129,8 @@ pub async fn run( .client_ca_root(Certificate::from_pem(client_ca)) .client_auth_optional(true), ); - let mut server = Server::builder() + let mut server = FLEET_LIMITS + .server() .accept_http1(true) .tls_config(tls_config)?; diff --git a/crates/controller/src/fleet_limits.rs b/crates/controller/src/fleet_limits.rs index d0c8bf82..22020ad6 100644 --- a/crates/controller/src/fleet_limits.rs +++ b/crates/controller/src/fleet_limits.rs @@ -6,7 +6,7 @@ use std::time::Duration; -use tonic::transport::ServerTlsConfig; +use tonic::transport::{Server, ServerTlsConfig}; /// Limits applied to the fleet gRPC listener. #[derive(Clone, Copy, Debug)] @@ -15,11 +15,16 @@ pub struct FleetLimits { /// no limit by default, so a client could hold a connection open forever by /// never sending its ClientHello. pub tls_handshake: Duration, + /// Time a request handler has to produce its response. A streaming RPC's + /// response is produced once the stream opens, so this does not end + /// established streams such as `Connect`. + pub request: Duration, } /// Production limits. pub const FLEET_LIMITS: FleetLimits = FleetLimits { tls_handshake: Duration::from_secs(10), + request: Duration::from_secs(30), }; impl FleetLimits { @@ -27,6 +32,11 @@ impl FleetLimits { pub fn tls(&self, config: ServerTlsConfig) -> ServerTlsConfig { config.timeout(self.tls_handshake) } + + /// A server builder with the fleet listener's limits applied. + pub fn server(&self) -> Server { + Server::builder().timeout(self.request) + } } #[cfg(test)] diff --git a/crates/controller/src/fleet_limits/tests.rs b/crates/controller/src/fleet_limits/tests.rs index a5fe4179..335a1916 100644 --- a/crates/controller/src/fleet_limits/tests.rs +++ b/crates/controller/src/fleet_limits/tests.rs @@ -5,6 +5,7 @@ use agentdesktop_proto::fleet::{ ControllerMessage, DeviceCertificateResponse, DeviceUsageRequest, DeviceUsageResponse, EnrollResponse, LlmGatewayCredentialRequest, LlmGatewayCredentialResponse, RenewDeviceCertificateRequest, + fleet_agent_client::FleetAgentClient, fleet_agent_server::{FleetAgent, FleetAgentServer}, }; use futures_core::Stream; @@ -16,17 +17,19 @@ use tokio::{ }; use tokio_stream::wrappers::ReceiverStream; use tonic::{ - Request, Response, Status, Streaming, - transport::{Identity, Server, ServerTlsConfig, server::TcpIncoming}, + Code, Request, Response, Status, Streaming, + transport::{Identity, ServerTlsConfig, server::TcpIncoming}, }; -use super::FleetLimits; +use super::{FLEET_LIMITS, FleetLimits}; -/// Limits small enough for tests to observe. +/// Limits small enough for tests to observe; anything a test does not exercise +/// stays at the production value. fn limits() -> FleetLimits { - FleetLimits { - tls_handshake: Duration::from_millis(200), - } + let mut limits = FLEET_LIMITS; + limits.tls_handshake = Duration::from_millis(200); + limits.request = Duration::from_millis(200); + limits } /// Fails the test, instead of hanging it, when the server never ends the wait. @@ -107,7 +110,7 @@ impl FleetAgent for Slow { async fn serve(limits: FleetLimits, identity: Option, service: Slow) -> SocketAddr { let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); let address = listener.local_addr().unwrap(); - let mut server = Server::builder(); + let mut server = limits.server(); if let Some(identity) = identity { server = server .tls_config(limits.tls(ServerTlsConfig::new().identity(identity))) @@ -118,6 +121,16 @@ async fn serve(limits: FleetLimits, identity: Option, service: Slow) - address } +async fn client(address: SocketAddr) -> FleetAgentClient { + // The generated `connect` RPC shadows the client's `connect` constructor. + let channel = tonic::transport::Endpoint::from_shared(format!("http://{address}")) + .unwrap() + .connect() + .await + .unwrap(); + FleetAgentClient::new(channel) +} + /// Reads and discards everything the server sends until it closes the connection. async fn until_closed(mut stream: TcpStream) { let mut buffer = [0_u8; 1024]; @@ -142,3 +155,45 @@ async fn a_client_that_never_finishes_the_tls_handshake_is_disconnected() { bounded(until_closed(stream)).await; } + +#[tokio::test] +async fn a_request_that_runs_too_long_is_ended() { + let address = serve( + limits(), + None, + Slow { + delay: Duration::from_secs(30), + }, + ) + .await; + + let error = bounded( + client(address) + .await + .begin_enrollment(BeginEnrollmentRequest::default()), + ) + .await + .expect_err("a request outlived the request timeout"); + + assert_eq!(error.code(), Code::Cancelled, "{error}"); +} + +#[tokio::test] +async fn the_request_timeout_does_not_end_an_open_stream() { + // Each message arrives after more than the request timeout. + let delay = limits().request * 2; + let address = serve(limits(), None, Slow { delay }).await; + let (_outbound, receiver) = mpsc::channel::(1); + + let mut stream = bounded(client(address).await.connect(ReceiverStream::new(receiver))) + .await + .expect("open the stream") + .into_inner(); + + for message in 0..2 { + let received = bounded(stream.message()) + .await + .unwrap_or_else(|error| panic!("stream ended at message {message}: {error}")); + assert!(received.is_some(), "stream closed at message {message}"); + } +} From 16833ad8648ec37d3b90d15967752f2bf394a8aa Mon Sep 17 00:00:00 2001 From: Scott Weiss Date: Mon, 9 Nov 2026 10:46:01 -0500 Subject: [PATCH 03/49] fix(controller): close fleet connections whose peer stopped answering A device that disappears without closing its connection (sleep, network change) left the controller holding that connection and its Connect stream until the operating system gave up on it. Enable HTTP/2 keepalive PINGs and TCP keepalive on the fleet server, so connections whose peer no longer answers are closed. This is the controller-side counterpart of the agent's own keepalive. --- crates/controller/src/fleet_limits.rs | 15 +++++++++++- crates/controller/src/fleet_limits/tests.rs | 26 ++++++++++++++++++++- 2 files changed, 39 insertions(+), 2 deletions(-) diff --git a/crates/controller/src/fleet_limits.rs b/crates/controller/src/fleet_limits.rs index 22020ad6..179d0b84 100644 --- a/crates/controller/src/fleet_limits.rs +++ b/crates/controller/src/fleet_limits.rs @@ -19,12 +19,21 @@ pub struct FleetLimits { /// response is produced once the stream opens, so this does not end /// established streams such as `Connect`. pub request: Duration, + /// How often an HTTP/2 PING checks a connection, and how long an unanswered + /// PING waits before the connection is closed. + pub keepalive_interval: Duration, + pub keepalive_timeout: Duration, + /// TCP keepalive probe interval, for peers that vanish without closing. + pub tcp_keepalive: Duration, } /// Production limits. pub const FLEET_LIMITS: FleetLimits = FleetLimits { tls_handshake: Duration::from_secs(10), request: Duration::from_secs(30), + keepalive_interval: Duration::from_secs(30), + keepalive_timeout: Duration::from_secs(20), + tcp_keepalive: Duration::from_secs(60), }; impl FleetLimits { @@ -35,7 +44,11 @@ impl FleetLimits { /// A server builder with the fleet listener's limits applied. pub fn server(&self) -> Server { - Server::builder().timeout(self.request) + Server::builder() + .timeout(self.request) + .http2_keepalive_interval(Some(self.keepalive_interval)) + .http2_keepalive_timeout(Some(self.keepalive_timeout)) + .tcp_keepalive(Some(self.tcp_keepalive)) } } diff --git a/crates/controller/src/fleet_limits/tests.rs b/crates/controller/src/fleet_limits/tests.rs index 335a1916..e29b688b 100644 --- a/crates/controller/src/fleet_limits/tests.rs +++ b/crates/controller/src/fleet_limits/tests.rs @@ -10,7 +10,7 @@ use agentdesktop_proto::fleet::{ }; use futures_core::Stream; use tokio::{ - io::AsyncReadExt, + io::{AsyncReadExt, AsyncWriteExt}, net::{TcpListener, TcpStream}, sync::mpsc, time, @@ -29,6 +29,8 @@ fn limits() -> FleetLimits { let mut limits = FLEET_LIMITS; limits.tls_handshake = Duration::from_millis(200); limits.request = Duration::from_millis(200); + limits.keepalive_interval = Duration::from_millis(100); + limits.keepalive_timeout = Duration::from_millis(100); limits } @@ -197,3 +199,25 @@ async fn the_request_timeout_does_not_end_an_open_stream() { assert!(received.is_some(), "stream closed at message {message}"); } } + +#[tokio::test] +async fn a_connection_that_stops_answering_pings_is_closed() { + let address = serve( + limits(), + None, + Slow { + delay: Duration::ZERO, + }, + ) + .await; + let mut stream = TcpStream::connect(address).await.unwrap(); + + // Open an HTTP/2 connection (preface and an empty SETTINGS frame), then go + // silent: never acknowledge the server's SETTINGS or PINGs. + stream + .write_all(b"PRI * HTTP/2.0\r\n\r\nSM\r\n\r\n\x00\x00\x00\x04\x00\x00\x00\x00\x00") + .await + .unwrap(); + + bounded(until_closed(stream)).await; +} From d67560c1d2aa27621b048aea37d98a03508bb276 Mon Sep 17 00:00:00 2001 From: Scott Weiss Date: Mon, 9 Nov 2026 10:46:47 -0500 Subject: [PATCH 04/49] fix(controller): limit requests and streams per fleet connection One connection could open any number of HTTP/2 streams and run any number of requests at once, taking an unbounded share of the server. Cap concurrent requests and concurrent HTTP/2 streams per connection. Requests beyond the cap wait rather than fail, so a well-behaved client sharing a connection is only slowed down. --- crates/controller/src/fleet_limits.rs | 9 ++++ crates/controller/src/fleet_limits/tests.rs | 54 +++++++++++++++++++++ 2 files changed, 63 insertions(+) diff --git a/crates/controller/src/fleet_limits.rs b/crates/controller/src/fleet_limits.rs index 179d0b84..0bbba591 100644 --- a/crates/controller/src/fleet_limits.rs +++ b/crates/controller/src/fleet_limits.rs @@ -25,6 +25,11 @@ pub struct FleetLimits { pub keepalive_timeout: Duration, /// TCP keepalive probe interval, for peers that vanish without closing. pub tcp_keepalive: Duration, + /// Concurrent requests one connection may have in progress, and concurrent + /// HTTP/2 streams it may open. Further requests wait, so one connection + /// cannot occupy an unbounded share of the server. + pub requests_per_connection: usize, + pub streams_per_connection: u32, } /// Production limits. @@ -34,6 +39,8 @@ pub const FLEET_LIMITS: FleetLimits = FleetLimits { keepalive_interval: Duration::from_secs(30), keepalive_timeout: Duration::from_secs(20), tcp_keepalive: Duration::from_secs(60), + requests_per_connection: 32, + streams_per_connection: 64, }; impl FleetLimits { @@ -49,6 +56,8 @@ impl FleetLimits { .http2_keepalive_interval(Some(self.keepalive_interval)) .http2_keepalive_timeout(Some(self.keepalive_timeout)) .tcp_keepalive(Some(self.tcp_keepalive)) + .concurrency_limit_per_connection(self.requests_per_connection) + .max_concurrent_streams(self.streams_per_connection) } } diff --git a/crates/controller/src/fleet_limits/tests.rs b/crates/controller/src/fleet_limits/tests.rs index e29b688b..c5f24704 100644 --- a/crates/controller/src/fleet_limits/tests.rs +++ b/crates/controller/src/fleet_limits/tests.rs @@ -221,3 +221,57 @@ async fn a_connection_that_stops_answering_pings_is_closed() { bounded(until_closed(stream)).await; } + +/// How long two concurrent half-second calls take on one connection. +async fn two_calls_on_one_connection(limits: FleetLimits) -> Duration { + let delay = Duration::from_millis(500); + let address = serve(limits, None, Slow { delay }).await; + let mut client = client(address).await; + // One quick call first, so the HTTP/2 SETTINGS exchange (which carries the + // stream limit) has completed before the measured calls. + let _ = client.get_device_usage(DeviceUsageRequest::default()).await; + let (mut first, mut second) = (client.clone(), client); + let started = time::Instant::now(); + let (a, b) = bounded(async { + tokio::join!( + first.begin_enrollment(BeginEnrollmentRequest::default()), + second.begin_enrollment(BeginEnrollmentRequest::default()), + ) + }) + .await; + a.unwrap(); + b.unwrap(); + started.elapsed() +} + +fn one_at_a_time() -> FleetLimits { + let mut limits = limits(); + limits.request = Duration::from_secs(5); + limits +} + +#[tokio::test] +async fn requests_beyond_the_per_connection_limit_wait() { + let mut limits = one_at_a_time(); + limits.requests_per_connection = 1; + + let elapsed = two_calls_on_one_connection(limits).await; + + assert!( + elapsed >= Duration::from_millis(950), + "ran concurrently: {elapsed:?}" + ); +} + +#[tokio::test] +async fn streams_beyond_the_per_connection_limit_wait() { + let mut limits = one_at_a_time(); + limits.streams_per_connection = 1; + + let elapsed = two_calls_on_one_connection(limits).await; + + assert!( + elapsed >= Duration::from_millis(950), + "ran concurrently: {elapsed:?}" + ); +} From f5228c2c21d8d391cf224f12090f34ead9a58584 Mon Sep 17 00:00:00 2001 From: Scott Weiss Date: Mon, 9 Nov 2026 10:47:41 -0500 Subject: [PATCH 05/49] fix(controller): cap fleet message sizes, with a small cap for enrollment The fleet service relied on tonic's default 4 MiB decode limit for every call, including the unauthenticated enrollment RPCs, and set no limit on responses. Set the 4 MiB limit explicitly in both directions for device calls, and cap enrollment requests at 64 KiB: they carry only a hostname, a PKCE challenge and a CSR. Tonic configures limits per service, so enrollment calls are routed to their own instance of the generated server. --- crates/controller/src/application.rs | 3 +- crates/controller/src/fleet_limits.rs | 84 ++++++++++++++++++++- crates/controller/src/fleet_limits/tests.rs | 71 ++++++++++++++++- 3 files changed, 150 insertions(+), 8 deletions(-) diff --git a/crates/controller/src/application.rs b/crates/controller/src/application.rs index 07ec8466..93f1e9c7 100644 --- a/crates/controller/src/application.rs +++ b/crates/controller/src/application.rs @@ -9,7 +9,6 @@ use crate::{ service::FleetAgentService, }; use agentdesktop_core::config::ControllerConfig; -use agentdesktop_proto::fleet::fleet_agent_server::FleetAgentServer; use agentdesktop_usage::UsageIngest; use anyhow::Context; use tonic::{ @@ -139,7 +138,7 @@ pub async fn run( tracing::info!(listen = %fleet_listen, "fleet controller listening"); let fleet = async move { let mut routes = Routes::from(gateway_jwt::routes(gateway_jwks)) - .add_service(FleetAgentServer::new(service)); + .add_service(FLEET_LIMITS.fleet_service(service)); if let Some(usage_ingest) = usage_ingest { routes = routes.add_service(usage_ingest); } diff --git a/crates/controller/src/fleet_limits.rs b/crates/controller/src/fleet_limits.rs index 0bbba591..e124235d 100644 --- a/crates/controller/src/fleet_limits.rs +++ b/crates/controller/src/fleet_limits.rs @@ -4,9 +4,18 @@ //! certificates are optional on it because enrollment runs before a device has //! one. These limits keep a client from holding resources on it indefinitely. -use std::time::Duration; +use std::{ + convert::Infallible, + task::{Context, Poll}, + time::Duration, +}; -use tonic::transport::{Server, ServerTlsConfig}; +use agentdesktop_proto::fleet::fleet_agent_server::{FleetAgent, FleetAgentServer, SERVICE_NAME}; +use tonic::{ + codegen::{Service, http}, + server::NamedService, + transport::{Server, ServerTlsConfig}, +}; /// Limits applied to the fleet gRPC listener. #[derive(Clone, Copy, Debug)] @@ -30,6 +39,11 @@ pub struct FleetLimits { /// cannot occupy an unbounded share of the server. pub requests_per_connection: usize, pub streams_per_connection: u32, + /// Largest encoded message an enrolled device's calls may send or receive. + pub max_message_bytes: usize, + /// Largest enrollment request. Enrollment runs before a device is + /// authenticated, and its requests are small (a hostname and a CSR). + pub enrollment_message_bytes: usize, } /// Production limits. @@ -41,6 +55,8 @@ pub const FLEET_LIMITS: FleetLimits = FleetLimits { tcp_keepalive: Duration::from_secs(60), requests_per_connection: 32, streams_per_connection: 64, + max_message_bytes: 4 * 1024 * 1024, + enrollment_message_bytes: 64 * 1024, }; impl FleetLimits { @@ -59,6 +75,70 @@ impl FleetLimits { .concurrency_limit_per_connection(self.requests_per_connection) .max_concurrent_streams(self.streams_per_connection) } + + /// The fleet gRPC service with message size limits applied. + pub fn fleet_service(&self, service: T) -> FleetService { + let service = std::sync::Arc::new(service); + FleetService { + enrollment: FleetAgentServer::from_arc(service.clone()) + .max_decoding_message_size(self.enrollment_message_bytes) + .max_encoding_message_size(self.max_message_bytes), + devices: FleetAgentServer::from_arc(service) + .max_decoding_message_size(self.max_message_bytes) + .max_encoding_message_size(self.max_message_bytes), + } + } +} + +/// The fleet gRPC service, with a smaller request limit for the enrollment RPCs. +/// Tonic sets message limits per service, so enrollment requests go to their own +/// instance of the generated server. +pub struct FleetService { + enrollment: FleetAgentServer, + devices: FleetAgentServer, +} + +impl Clone for FleetService { + fn clone(&self) -> Self { + Self { + enrollment: self.enrollment.clone(), + devices: self.devices.clone(), + } + } +} + +fn is_enrollment(path: &str) -> bool { + path.strip_prefix('/') + .and_then(|path| path.strip_prefix(SERVICE_NAME)) + .is_some_and(|method| matches!(method, "/BeginEnrollment" | "/CompleteEnrollment")) +} + +impl Service> for FleetService +where + FleetAgentServer: Service, Error = Infallible>, +{ + type Response = as Service>>::Response; + type Error = Infallible; + type Future = as Service>>::Future; + + fn poll_ready(&mut self, context: &mut Context<'_>) -> Poll> { + match self.enrollment.poll_ready(context) { + Poll::Ready(Ok(())) => self.devices.poll_ready(context), + other => other, + } + } + + fn call(&mut self, request: http::Request) -> Self::Future { + if is_enrollment(request.uri().path()) { + self.enrollment.call(request) + } else { + self.devices.call(request) + } + } +} + +impl NamedService for FleetService { + const NAME: &'static str = SERVICE_NAME; } #[cfg(test)] diff --git a/crates/controller/src/fleet_limits/tests.rs b/crates/controller/src/fleet_limits/tests.rs index c5f24704..9078be63 100644 --- a/crates/controller/src/fleet_limits/tests.rs +++ b/crates/controller/src/fleet_limits/tests.rs @@ -4,9 +4,8 @@ use agentdesktop_proto::fleet::{ AgentMessage, BeginEnrollmentRequest, BeginEnrollmentResponse, CompleteEnrollmentRequest, ControllerMessage, DeviceCertificateResponse, DeviceUsageRequest, DeviceUsageResponse, EnrollResponse, LlmGatewayCredentialRequest, LlmGatewayCredentialResponse, - RenewDeviceCertificateRequest, - fleet_agent_client::FleetAgentClient, - fleet_agent_server::{FleetAgent, FleetAgentServer}, + RenewDeviceCertificateRequest, fleet_agent_client::FleetAgentClient, + fleet_agent_server::FleetAgent, }; use futures_core::Stream; use tokio::{ @@ -118,7 +117,7 @@ async fn serve(limits: FleetLimits, identity: Option, service: Slow) - .tls_config(limits.tls(ServerTlsConfig::new().identity(identity))) .unwrap(); } - let router = server.add_service(FleetAgentServer::new(service)); + let router = server.add_service(limits.fleet_service(service)); tokio::spawn(router.serve_with_incoming(TcpIncoming::from(listener))); address } @@ -275,3 +274,67 @@ async fn streams_beyond_the_per_connection_limit_wait() { "ran concurrently: {elapsed:?}" ); } + +fn small_messages() -> FleetLimits { + let mut limits = limits(); + limits.max_message_bytes = 1024; + limits.enrollment_message_bytes = 128; + limits +} + +#[tokio::test] +async fn an_oversized_enrollment_request_is_rejected() { + let address = serve( + small_messages(), + None, + Slow { + delay: Duration::ZERO, + }, + ) + .await; + + let error = bounded( + client(address) + .await + .begin_enrollment(BeginEnrollmentRequest { + hostname: "h".repeat(512), + ..Default::default() + }), + ) + .await + .expect_err("an oversized enrollment request was accepted"); + + assert_eq!(error.code(), Code::OutOfRange, "{error}"); +} + +#[tokio::test] +async fn device_calls_keep_the_larger_message_limit() { + let address = serve( + small_messages(), + None, + Slow { + delay: Duration::ZERO, + }, + ) + .await; + let mut client = client(address).await; + + // Over the enrollment limit but within the device limit: reaches the handler. + let within = bounded( + client.renew_device_certificate(RenewDeviceCertificateRequest { + certificate_signing_request_der: vec![0; 512], + }), + ) + .await + .expect_err("the test service rejects every renewal"); + assert_eq!(within.code(), Code::Unimplemented, "{within}"); + + let over = bounded( + client.renew_device_certificate(RenewDeviceCertificateRequest { + certificate_signing_request_der: vec![0; 4096], + }), + ) + .await + .expect_err("an oversized device request was accepted"); + assert_eq!(over.code(), Code::OutOfRange, "{over}"); +} From 99fedcb5039566ed2a9d2e4b77e6ddbe15d68e26 Mon Sep 17 00:00:00 2001 From: Scott Weiss Date: Mon, 9 Nov 2026 10:48:17 -0500 Subject: [PATCH 06/49] fix(controller): limit enrollment requests per client and in total BeginEnrollment and CompleteEnrollment run before a device has a certificate and call the identity provider, so anyone who can reach the fleet listener could flood them. Check an enrollment limiter first in both handlers: a cap on enrollment requests in progress across all clients, and a per-address token bucket. Rejections return RESOURCE_EXHAUSTED; devices already retry enrollment with backoff. The per-address budget is generous because many devices can share an address behind NAT; the in-flight cap is the main protection. --- crates/controller/src/fleet_limits.rs | 97 ++++++++++++++++- crates/controller/src/fleet_limits/tests.rs | 113 ++++++++++++++++++++ crates/controller/src/service.rs | 50 +++++++++ 3 files changed, 259 insertions(+), 1 deletion(-) diff --git a/crates/controller/src/fleet_limits.rs b/crates/controller/src/fleet_limits.rs index e124235d..0482e479 100644 --- a/crates/controller/src/fleet_limits.rs +++ b/crates/controller/src/fleet_limits.rs @@ -5,13 +5,21 @@ //! one. These limits keep a client from holding resources on it indefinitely. use std::{ + collections::HashMap, convert::Infallible, + net::{IpAddr, SocketAddr}, + sync::{Arc, Mutex}, task::{Context, Poll}, time::Duration, }; use agentdesktop_proto::fleet::fleet_agent_server::{FleetAgent, FleetAgentServer, SERVICE_NAME}; +use tokio::{ + sync::{OwnedSemaphorePermit, Semaphore}, + time::Instant, +}; use tonic::{ + Status, codegen::{Service, http}, server::NamedService, transport::{Server, ServerTlsConfig}, @@ -44,6 +52,15 @@ pub struct FleetLimits { /// Largest enrollment request. Enrollment runs before a device is /// authenticated, and its requests are small (a hostname and a CSR). pub enrollment_message_bytes: usize, + /// Enrollment requests (`BeginEnrollment`, `CompleteEnrollment`) running at + /// once across all clients. Enrollment calls the identity provider, so this + /// also bounds the load a flood can put on it. + pub enrollment_in_flight: usize, + /// Per-client-address enrollment budget: up to `enrollment_burst` requests at + /// once, refilled at one request per `enrollment_refill`. Generous, because + /// many devices can share one address behind NAT. + pub enrollment_burst: u32, + pub enrollment_refill: Duration, } /// Production limits. @@ -57,6 +74,9 @@ pub const FLEET_LIMITS: FleetLimits = FleetLimits { streams_per_connection: 64, max_message_bytes: 4 * 1024 * 1024, enrollment_message_bytes: 64 * 1024, + enrollment_in_flight: 64, + enrollment_burst: 60, + enrollment_refill: Duration::from_secs(1), }; impl FleetLimits { @@ -78,7 +98,7 @@ impl FleetLimits { /// The fleet gRPC service with message size limits applied. pub fn fleet_service(&self, service: T) -> FleetService { - let service = std::sync::Arc::new(service); + let service = Arc::new(service); FleetService { enrollment: FleetAgentServer::from_arc(service.clone()) .max_decoding_message_size(self.enrollment_message_bytes) @@ -88,6 +108,10 @@ impl FleetLimits { .max_encoding_message_size(self.max_message_bytes), } } + + pub fn enrollment(&self) -> EnrollmentLimiter { + EnrollmentLimiter::new(self) + } } /// The fleet gRPC service, with a smaller request limit for the enrollment RPCs. @@ -141,5 +165,76 @@ impl NamedService for FleetService { const NAME: &'static str = SERVICE_NAME; } +/// Admission control for the unauthenticated enrollment RPCs. +#[derive(Clone)] +pub struct EnrollmentLimiter { + in_flight: Arc, + peers: Arc>>, + burst: f64, + refill: Duration, +} + +/// Peer buckets kept before idle (full) ones are dropped. +const MAX_TRACKED_PEERS: usize = 10_000; + +#[derive(Clone, Copy)] +struct Bucket { + tokens: f64, + updated: Instant, +} + +impl EnrollmentLimiter { + fn new(limits: &FleetLimits) -> Self { + Self { + in_flight: Arc::new(Semaphore::new(limits.enrollment_in_flight)), + peers: Arc::default(), + burst: f64::from(limits.enrollment_burst), + refill: limits.enrollment_refill, + } + } + + /// Admits one enrollment request from `peer`, or rejects it with + /// `RESOURCE_EXHAUSTED`. Hold the returned permit until the request finishes. + pub fn admit(&self, peer: Option) -> Result { + if let Some(peer) = peer + && !self.take(peer.ip(), Instant::now()) + { + tracing::warn!(peer = %peer.ip(), "enrollment rate limit exceeded"); + return Err(Status::resource_exhausted( + "too many enrollment requests from this address; retry later", + )); + } + self.in_flight.clone().try_acquire_owned().map_err(|_| { + tracing::warn!("enrollment concurrency limit reached"); + Status::resource_exhausted("too many enrollment requests in progress; retry later") + }) + } + + fn take(&self, peer: IpAddr, now: Instant) -> bool { + let mut peers = self.peers.lock().expect("enrollment limiter lock poisoned"); + if peers.len() >= MAX_TRACKED_PEERS && !peers.contains_key(&peer) { + let burst = self.burst; + let refill = self.refill; + peers.retain(|_, bucket| refilled(*bucket, now, refill, burst) < burst); + } + let bucket = peers.entry(peer).or_insert(Bucket { + tokens: self.burst, + updated: now, + }); + bucket.tokens = refilled(*bucket, now, self.refill, self.burst); + bucket.updated = now; + if bucket.tokens < 1.0 { + return false; + } + bucket.tokens -= 1.0; + true + } +} + +fn refilled(bucket: Bucket, now: Instant, refill: Duration, burst: f64) -> f64 { + let elapsed = now.saturating_duration_since(bucket.updated); + (bucket.tokens + elapsed.as_secs_f64() / refill.as_secs_f64()).min(burst) +} + #[cfg(test)] mod tests; diff --git a/crates/controller/src/fleet_limits/tests.rs b/crates/controller/src/fleet_limits/tests.rs index 9078be63..412541b6 100644 --- a/crates/controller/src/fleet_limits/tests.rs +++ b/crates/controller/src/fleet_limits/tests.rs @@ -338,3 +338,116 @@ async fn device_calls_keep_the_larger_message_limit() { .expect_err("an oversized device request was accepted"); assert_eq!(over.code(), Code::OutOfRange, "{over}"); } + +mod enrollment { + use std::net::{IpAddr, Ipv4Addr, SocketAddr}; + + use tokio::time::Instant; + + use super::super::{FLEET_LIMITS, FleetLimits, MAX_TRACKED_PEERS}; + use super::Duration; + + fn limits(in_flight: usize, burst: u32) -> FleetLimits { + FleetLimits { + enrollment_in_flight: in_flight, + enrollment_burst: burst, + enrollment_refill: Duration::from_secs(10), + ..FLEET_LIMITS + } + } + + fn peer(last: u8) -> IpAddr { + IpAddr::V4(Ipv4Addr::new(192, 0, 2, last)) + } + + #[test] + fn a_peer_gets_its_burst_then_one_request_per_refill() { + let limiter = limits(64, 3).enrollment(); + let start = Instant::now(); + + for _ in 0..3 { + assert!(limiter.take(peer(1), start)); + } + assert!(!limiter.take(peer(1), start), "burst exceeded"); + assert!( + !limiter.take(peer(1), start + Duration::from_secs(9)), + "refilled early" + ); + assert!(limiter.take(peer(1), start + Duration::from_secs(10))); + assert!(!limiter.take(peer(1), start + Duration::from_secs(10))); + } + + #[test] + fn peers_have_separate_budgets() { + let limiter = limits(64, 1).enrollment(); + let now = Instant::now(); + + assert!(limiter.take(peer(1), now)); + assert!(!limiter.take(peer(1), now)); + assert!(limiter.take(peer(2), now)); + } + + #[test] + fn a_full_bucket_is_never_overfilled() { + let limiter = limits(64, 2).enrollment(); + let start = Instant::now(); + + assert!(limiter.take(peer(1), start)); + // A long idle period refills only up to the burst. + let later = start + Duration::from_secs(3600); + assert!(limiter.take(peer(1), later)); + assert!(limiter.take(peer(1), later)); + assert!(!limiter.take(peer(1), later)); + } + + #[tokio::test] + async fn requests_in_progress_are_capped_until_one_finishes() { + let limiter = limits(2, 60).enrollment(); + let address = |last| Some(SocketAddr::new(peer(last), 4000)); + + let first = limiter.admit(address(1)).expect("first request"); + let _second = limiter.admit(address(2)).expect("second request"); + let error = limiter + .admit(address(3)) + .expect_err("third request admitted"); + assert_eq!(error.code(), tonic::Code::ResourceExhausted); + + drop(first); + let _third = limiter + .admit(address(3)) + .expect("admitted after one finished"); + } + + #[tokio::test] + async fn a_rate_limited_peer_is_rejected_with_resource_exhausted() { + let limiter = limits(64, 1).enrollment(); + let address = Some(SocketAddr::new(peer(1), 4000)); + + let _first = limiter.admit(address).expect("first request"); + let error = limiter.admit(address).expect_err("second request admitted"); + + assert_eq!(error.code(), tonic::Code::ResourceExhausted); + } + + #[tokio::test] + async fn requests_without_a_peer_address_still_count_against_the_cap() { + let limiter = limits(1, 60).enrollment(); + + let _first = limiter.admit(None).expect("first request"); + assert!(limiter.admit(None).is_err()); + } + + #[test] + fn idle_peers_are_dropped_when_the_table_is_full() { + let limiter = limits(64, 1).enrollment(); + let start = Instant::now(); + for index in 0..MAX_TRACKED_PEERS { + let [a, b, c, d] = u32::try_from(index).unwrap().to_be_bytes(); + assert!(limiter.take(IpAddr::V4(Ipv4Addr::new(10 | a, b, c, d)), start)); + } + + // Once every bucket has refilled, a new peer evicts the idle ones. + assert!(limiter.take(peer(1), start + Duration::from_secs(10))); + assert_eq!(limiter.peers.lock().unwrap().len(), 1); + } +} diff --git a/crates/controller/src/service.rs b/crates/controller/src/service.rs index 6fe3570b..54e66648 100644 --- a/crates/controller/src/service.rs +++ b/crates/controller/src/service.rs @@ -22,6 +22,7 @@ use crate::{ daemon_config::DaemonConfigStore, database::{Database, DevicePrincipal}, device_ca::DeviceCertificateIssuer, + fleet_limits::{EnrollmentLimiter, FLEET_LIMITS}, gateway_jwt::GatewayJwtIssuer, oidc::{OidcPrincipal, OidcProvider}, }; @@ -34,6 +35,7 @@ pub struct FleetAgentService { gateway_jwt_issuer: Option, device_certificate_issuer: Option, usage_enabled: bool, + enrollment: EnrollmentLimiter, } impl FleetAgentService { @@ -51,9 +53,16 @@ impl FleetAgentService { gateway_jwt_issuer, device_certificate_issuer, usage_enabled: false, + enrollment: FLEET_LIMITS.enrollment(), } } + /// Replaces the enrollment admission limits (tests use small ones). + pub fn with_enrollment_limiter(mut self, enrollment: EnrollmentLimiter) -> Self { + self.enrollment = enrollment; + self + } + /// Tells devices that the controller collects gateway usage. pub fn with_usage_collection(mut self, enabled: bool) -> Self { self.usage_enabled = enabled; @@ -73,6 +82,7 @@ impl FleetAgent for FleetAgentService { &self, request: Request, ) -> Result, Status> { + let _admitted = self.enrollment.admit(request.remote_addr())?; let oidc = self .oidc .as_ref() @@ -91,6 +101,7 @@ impl FleetAgent for FleetAgentService { &self, request: Request, ) -> Result, Status> { + let _admitted = self.enrollment.admit(request.remote_addr())?; let oidc = self .oidc .as_ref() @@ -662,6 +673,45 @@ mod tests { )); } + #[tokio::test] + async fn enrollment_is_refused_before_any_work_when_the_limiter_is_full() { + use agentdesktop_proto::fleet::{ + BeginEnrollmentRequest, CompleteEnrollmentRequest, fleet_agent_server::FleetAgent, + }; + use tonic::{Code, Request}; + + use crate::fleet_limits::{FLEET_LIMITS, FleetLimits}; + + let path = std::env::temp_dir().join(format!( + "agentdesktop-enrollment-limit-{}-{}.db", + std::process::id(), + rand::random::() + )); + let database = Database::connect(&format!("sqlite://{}?mode=rwc", path.display())) + .await + .expect("connect database"); + let full = FleetLimits { + enrollment_in_flight: 0, + ..FLEET_LIMITS + }; + let service = + FleetAgentService::new(None, database, DaemonConfigStore::new(None), None, None) + .with_enrollment_limiter(full.enrollment()); + + // Without the limiter these fail later, with "OIDC enrollment is disabled". + let begin = service + .begin_enrollment(Request::new(BeginEnrollmentRequest::default())) + .await + .expect_err("begin admitted"); + assert_eq!(begin.code(), Code::ResourceExhausted, "{begin}"); + let complete = service + .complete_enrollment(Request::new(CompleteEnrollmentRequest::default())) + .await + .expect_err("complete admitted"); + assert_eq!(complete.code(), Code::ResourceExhausted, "{complete}"); + let _ = std::fs::remove_file(path); + } + #[tokio::test] async fn device_authentication_binds_oidc_issuer_and_subject() { let device_id = "7ca03414-bb20-4c80-98ef-7b0538b988ba"; From 72de9e24ed6ac3247ad10904fc58e633bb06f8e5 Mon Sep 17 00:00:00 2001 From: Scott Weiss Date: Mon, 9 Nov 2026 10:50:14 -0500 Subject: [PATCH 07/49] fix(controller): cap fleet connections in total and per client address Nothing limited how many connections the fleet listener accepted. Client certificates are optional on it, so anyone could finish the TLS handshake and hold idle connections open by answering keepalive PINGs, until the controller ran out of file descriptors. Accept fleet connections through a wrapper that closes a new connection when the total or per-address cap is reached, and releases its slots when the connection closes. The per-address cap is high because many devices can share an address behind NAT. The wrapper keeps tonic's TCP connect info, so client addresses and device certificates are unchanged, and it now sets the TCP keepalive the server builder used to. --- crates/controller/src/application.rs | 2 +- crates/controller/src/fleet_limits.rs | 180 +++++++++++++++++++- crates/controller/src/fleet_limits/tests.rs | 58 ++++++- 3 files changed, 236 insertions(+), 4 deletions(-) diff --git a/crates/controller/src/application.rs b/crates/controller/src/application.rs index 93f1e9c7..e5ffb49c 100644 --- a/crates/controller/src/application.rs +++ b/crates/controller/src/application.rs @@ -144,7 +144,7 @@ pub async fn run( } server .add_routes(routes) - .serve(fleet_listen) + .serve_with_incoming(FLEET_LIMITS.incoming(fleet_listen)?) .await .context("serve fleet gRPC API") }; diff --git a/crates/controller/src/fleet_limits.rs b/crates/controller/src/fleet_limits.rs index 0482e479..0596c466 100644 --- a/crates/controller/src/fleet_limits.rs +++ b/crates/controller/src/fleet_limits.rs @@ -7,14 +7,19 @@ use std::{ collections::HashMap, convert::Infallible, + io, net::{IpAddr, SocketAddr}, + pin::Pin, sync::{Arc, Mutex}, task::{Context, Poll}, time::Duration, }; use agentdesktop_proto::fleet::fleet_agent_server::{FleetAgent, FleetAgentServer, SERVICE_NAME}; +use futures_core::Stream; use tokio::{ + io::{AsyncRead, AsyncWrite, ReadBuf}, + net::TcpStream, sync::{OwnedSemaphorePermit, Semaphore}, time::Instant, }; @@ -22,7 +27,10 @@ use tonic::{ Status, codegen::{Service, http}, server::NamedService, - transport::{Server, ServerTlsConfig}, + transport::{ + Server, ServerTlsConfig, + server::{Connected, TcpConnectInfo, TcpIncoming}, + }, }; /// Limits applied to the fleet gRPC listener. @@ -42,6 +50,12 @@ pub struct FleetLimits { pub keepalive_timeout: Duration, /// TCP keepalive probe interval, for peers that vanish without closing. pub tcp_keepalive: Duration, + /// Open connections allowed in total and from one client address. A + /// connection over either cap is closed as soon as it is accepted. The + /// per-address cap is high because many devices can share an address behind + /// NAT, or appear to when a load balancer rewrites client addresses. + pub max_connections: usize, + pub connections_per_address: usize, /// Concurrent requests one connection may have in progress, and concurrent /// HTTP/2 streams it may open. Further requests wait, so one connection /// cannot occupy an unbounded share of the server. @@ -70,6 +84,8 @@ pub const FLEET_LIMITS: FleetLimits = FleetLimits { keepalive_interval: Duration::from_secs(30), keepalive_timeout: Duration::from_secs(20), tcp_keepalive: Duration::from_secs(60), + max_connections: 50_000, + connections_per_address: 4_096, requests_per_connection: 32, streams_per_connection: 64, max_message_bytes: 4 * 1024 * 1024, @@ -91,11 +107,29 @@ impl FleetLimits { .timeout(self.request) .http2_keepalive_interval(Some(self.keepalive_interval)) .http2_keepalive_timeout(Some(self.keepalive_timeout)) - .tcp_keepalive(Some(self.tcp_keepalive)) .concurrency_limit_per_connection(self.requests_per_connection) .max_concurrent_streams(self.streams_per_connection) } + /// Binds the fleet listener with TCP keepalive and the connection caps. + /// Serve with `Server::serve_with_incoming`. + pub fn incoming(&self, address: SocketAddr) -> io::Result> { + let incoming = TcpIncoming::bind(address)? + .with_nodelay(Some(true)) + .with_keepalive(Some(self.tcp_keepalive)); + Ok(self.limit_connections(incoming)) + } + + /// Applies the connection caps to accepted connections. + pub fn limit_connections(&self, incoming: S) -> LimitedIncoming { + LimitedIncoming { + incoming, + total: Arc::new(Semaphore::new(self.max_connections)), + per_address: Arc::default(), + per_address_limit: self.connections_per_address, + } + } + /// The fleet gRPC service with message size limits applied. pub fn fleet_service(&self, service: T) -> FleetService { let service = Arc::new(service); @@ -114,6 +148,148 @@ impl FleetLimits { } } +/// Accepted connections, closed on accept when over the connection caps. +pub struct LimitedIncoming { + incoming: S, + total: Arc, + per_address: Arc>>, + per_address_limit: usize, +} + +impl LimitedIncoming { + fn admit(&self, stream: TcpStream) -> Option { + let Ok(total) = self.total.clone().try_acquire_owned() else { + tracing::debug!("fleet connection limit reached; closing connection"); + return None; + }; + let address = stream.peer_addr().ok().map(|peer| peer.ip()); + let address = match address { + Some(address) => { + let mut counts = self + .per_address + .lock() + .expect("connection limiter lock poisoned"); + let count = counts.entry(address).or_default(); + if *count >= self.per_address_limit { + tracing::debug!(%address, "per-address fleet connection limit reached; closing connection"); + return None; + } + *count += 1; + Some(AddressSlot { + counts: self.per_address.clone(), + address, + }) + } + None => None, + }; + Some(LimitedConnection { + stream, + _total: total, + _address: address, + }) + } +} + +impl Stream for LimitedIncoming +where + S: Stream> + Unpin, +{ + type Item = io::Result; + + fn poll_next(mut self: Pin<&mut Self>, context: &mut Context<'_>) -> Poll> { + loop { + match Pin::new(&mut self.incoming).poll_next(context) { + Poll::Ready(Some(Ok(stream))) => { + // Over a cap: dropping the stream closes it; wait for the next one. + if let Some(connection) = self.admit(stream) { + return Poll::Ready(Some(Ok(connection))); + } + } + Poll::Ready(Some(Err(error))) => return Poll::Ready(Some(Err(error))), + Poll::Ready(None) => return Poll::Ready(None), + Poll::Pending => return Poll::Pending, + } + } + } +} + +/// Releases a client address's connection slot when its connection closes. +struct AddressSlot { + counts: Arc>>, + address: IpAddr, +} + +impl Drop for AddressSlot { + fn drop(&mut self) { + let mut counts = self + .counts + .lock() + .expect("connection limiter lock poisoned"); + if let Some(count) = counts.get_mut(&self.address) { + *count -= 1; + if *count == 0 { + counts.remove(&self.address); + } + } + } +} + +/// An accepted fleet connection holding its slots until it closes. +pub struct LimitedConnection { + stream: TcpStream, + _total: OwnedSemaphorePermit, + _address: Option, +} + +impl Connected for LimitedConnection { + // The plain TCP info keeps `Request::remote_addr` and `peer_certs` working. + type ConnectInfo = TcpConnectInfo; + + fn connect_info(&self) -> Self::ConnectInfo { + self.stream.connect_info() + } +} + +impl AsyncRead for LimitedConnection { + fn poll_read( + mut self: Pin<&mut Self>, + context: &mut Context<'_>, + buffer: &mut ReadBuf<'_>, + ) -> Poll> { + Pin::new(&mut self.stream).poll_read(context, buffer) + } +} + +impl AsyncWrite for LimitedConnection { + fn poll_write( + mut self: Pin<&mut Self>, + context: &mut Context<'_>, + buffer: &[u8], + ) -> Poll> { + Pin::new(&mut self.stream).poll_write(context, buffer) + } + + fn poll_write_vectored( + mut self: Pin<&mut Self>, + context: &mut Context<'_>, + buffers: &[io::IoSlice<'_>], + ) -> Poll> { + Pin::new(&mut self.stream).poll_write_vectored(context, buffers) + } + + fn is_write_vectored(&self) -> bool { + self.stream.is_write_vectored() + } + + fn poll_flush(mut self: Pin<&mut Self>, context: &mut Context<'_>) -> Poll> { + Pin::new(&mut self.stream).poll_flush(context) + } + + fn poll_shutdown(mut self: Pin<&mut Self>, context: &mut Context<'_>) -> Poll> { + Pin::new(&mut self.stream).poll_shutdown(context) + } +} + /// The fleet gRPC service, with a smaller request limit for the enrollment RPCs. /// Tonic sets message limits per service, so enrollment requests go to their own /// instance of the generated server. diff --git a/crates/controller/src/fleet_limits/tests.rs b/crates/controller/src/fleet_limits/tests.rs index 412541b6..588eb1a5 100644 --- a/crates/controller/src/fleet_limits/tests.rs +++ b/crates/controller/src/fleet_limits/tests.rs @@ -118,7 +118,7 @@ async fn serve(limits: FleetLimits, identity: Option, service: Slow) - .unwrap(); } let router = server.add_service(limits.fleet_service(service)); - tokio::spawn(router.serve_with_incoming(TcpIncoming::from(listener))); + tokio::spawn(router.serve_with_incoming(limits.limit_connections(TcpIncoming::from(listener)))); address } @@ -339,6 +339,62 @@ async fn device_calls_keep_the_larger_message_limit() { assert_eq!(over.code(), Code::OutOfRange, "{over}"); } +/// Whether the server keeps `stream` open for a while. An accepted connection +/// receives the server's HTTP/2 SETTINGS and then waits for the client; a +/// rejected one is closed (end of stream) right away. +async fn stays_open(mut stream: TcpStream) -> bool { + let mut buffer = [0_u8; 1024]; + let closed = time::timeout(Duration::from_millis(300), async { + loop { + match stream.read(&mut buffer).await { + Ok(0) | Err(_) => return, + Ok(_) => {} + } + } + }) + .await; + closed.is_err() +} + +async fn connections_over_cap(limits: FleetLimits) { + let address = serve( + limits, + None, + Slow { + delay: Duration::ZERO, + }, + ) + .await; + let first = TcpStream::connect(address).await.unwrap(); + let second = TcpStream::connect(address).await.unwrap(); + + let third = TcpStream::connect(address).await.unwrap(); + bounded(until_closed(third)).await; + + // Closing a connection frees its slot. + drop(first); + time::sleep(Duration::from_millis(100)).await; + let replacement = TcpStream::connect(address).await.unwrap(); + assert!(stays_open(replacement).await, "a freed slot was not reused"); + drop(second); +} + +#[tokio::test] +async fn connections_over_the_per_address_cap_are_closed() { + let mut limits = limits(); + limits.connections_per_address = 2; + + connections_over_cap(limits).await; +} + +#[tokio::test] +async fn connections_over_the_total_cap_are_closed() { + let mut limits = limits(); + limits.max_connections = 2; + + connections_over_cap(limits).await; +} + mod enrollment { use std::net::{IpAddr, Ipv4Addr, SocketAddr}; From b93f68c8d480378876b617181d6153303d0cf388 Mon Sep 17 00:00:00 2001 From: Scott Weiss Date: Mon, 9 Nov 2026 10:51:42 -0500 Subject: [PATCH 08/49] fix(controller): one Connect stream per device, with a message rate limit An enrolled device could open any number of Connect streams and send messages on them as fast as it liked, each one stored in the database. Keep at most one Connect stream per device: opening a new stream tells the device's previous stream to close with ABORTED. The newest stream wins, so a device reconnecting after sleep is never blocked by a stream the controller has not yet noticed is dead. Each stream also gets a message budget far above normal use (heartbeats, periodic inventory, one telemetry event per tool event); a stream over budget is closed with RESOURCE_EXHAUSTED and the device reconnects with backoff. The registry and the budget have unit tests. Wiring them into Connect is not covered end to end, since that needs a device certificate and OIDC. --- crates/controller/src/fleet_limits.rs | 119 ++++++++++++++++++-- crates/controller/src/fleet_limits/tests.rs | 57 ++++++++++ crates/controller/src/service.rs | 20 +++- 3 files changed, 187 insertions(+), 9 deletions(-) diff --git a/crates/controller/src/fleet_limits.rs b/crates/controller/src/fleet_limits.rs index 0596c466..f65c2b01 100644 --- a/crates/controller/src/fleet_limits.rs +++ b/crates/controller/src/fleet_limits.rs @@ -10,7 +10,10 @@ use std::{ io, net::{IpAddr, SocketAddr}, pin::Pin, - sync::{Arc, Mutex}, + sync::{ + Arc, Mutex, + atomic::{AtomicU64, Ordering}, + }, task::{Context, Poll}, time::Duration, }; @@ -20,7 +23,7 @@ use futures_core::Stream; use tokio::{ io::{AsyncRead, AsyncWrite, ReadBuf}, net::TcpStream, - sync::{OwnedSemaphorePermit, Semaphore}, + sync::{OwnedSemaphorePermit, Semaphore, oneshot}, time::Instant, }; use tonic::{ @@ -75,6 +78,13 @@ pub struct FleetLimits { /// many devices can share one address behind NAT. pub enrollment_burst: u32, pub enrollment_refill: Duration, + /// Messages a device may send on one `Connect` stream: up to + /// `stream_message_burst` at once, refilled at one per + /// `stream_message_refill`. Devices send a heartbeat every 30 seconds, + /// inventory periodically and one telemetry event per tool event, so this is + /// far above normal use; a stream over budget is closed. + pub stream_message_burst: u32, + pub stream_message_refill: Duration, } /// Production limits. @@ -93,6 +103,8 @@ pub const FLEET_LIMITS: FleetLimits = FleetLimits { enrollment_in_flight: 64, enrollment_burst: 60, enrollment_refill: Duration::from_secs(1), + stream_message_burst: 600, + stream_message_refill: Duration::from_millis(50), }; impl FleetLimits { @@ -146,6 +158,18 @@ impl FleetLimits { pub fn enrollment(&self) -> EnrollmentLimiter { EnrollmentLimiter::new(self) } + + /// A fresh message budget for one `Connect` stream. + pub fn stream_messages(&self) -> MessageBudget { + MessageBudget { + bucket: Bucket { + tokens: f64::from(self.stream_message_burst), + updated: Instant::now(), + }, + burst: f64::from(self.stream_message_burst), + refill: self.stream_message_refill, + } + } } /// Accepted connections, closed on accept when over the connection caps. @@ -397,13 +421,92 @@ impl EnrollmentLimiter { tokens: self.burst, updated: now, }); - bucket.tokens = refilled(*bucket, now, self.refill, self.burst); - bucket.updated = now; - if bucket.tokens < 1.0 { - return false; + take_token(bucket, now, self.refill, self.burst) + } +} + +/// Refills `bucket` for the time since its last use, then takes one token. +fn take_token(bucket: &mut Bucket, now: Instant, refill: Duration, burst: f64) -> bool { + bucket.tokens = refilled(*bucket, now, refill, burst); + bucket.updated = now; + if bucket.tokens < 1.0 { + return false; + } + bucket.tokens -= 1.0; + true +} + +/// Messages a device may send on one `Connect` stream. +pub struct MessageBudget { + bucket: Bucket, + burst: f64, + refill: Duration, +} + +impl MessageBudget { + /// Takes one message from the budget, or returns false when it is spent. + pub fn take(&mut self, now: Instant) -> bool { + take_token(&mut self.bucket, now, self.refill, self.burst) + } +} + +/// Open `Connect` streams by device. A device has at most one: opening a new +/// stream tells the device's previous stream to close. A device that reconnects +/// (after sleep, for example) is never blocked by a stream the controller has +/// not yet noticed is dead. +#[derive(Clone, Default)] +pub struct DeviceStreams { + open: Arc>>, + next: Arc, +} + +/// A device's current stream: its generation and the signal that replaces it. +type Registration = (u64, oneshot::Sender<()>); + +/// A registered `Connect` stream. `superseded` resolves when the same device +/// opens a newer stream; dropping this unregisters it. +pub struct DeviceStream { + streams: DeviceStreams, + device_id: String, + generation: u64, + pub superseded: oneshot::Receiver<()>, +} + +impl DeviceStreams { + pub fn open(&self, device_id: &str) -> DeviceStream { + let generation = self.next.fetch_add(1, Ordering::Relaxed); + let (replace, superseded) = oneshot::channel(); + let previous = self + .open + .lock() + .expect("device stream registry lock poisoned") + .insert(device_id.to_owned(), (generation, replace)); + if let Some((_, previous)) = previous { + let _ = previous.send(()); + } + DeviceStream { + streams: self.clone(), + device_id: device_id.to_owned(), + generation, + superseded, + } + } +} + +impl Drop for DeviceStream { + fn drop(&mut self) { + let mut open = self + .streams + .open + .lock() + .expect("device stream registry lock poisoned"); + // A newer stream for the same device keeps its registration. + if open + .get(&self.device_id) + .is_some_and(|(generation, _)| *generation == self.generation) + { + open.remove(&self.device_id); } - bucket.tokens -= 1.0; - true } } diff --git a/crates/controller/src/fleet_limits/tests.rs b/crates/controller/src/fleet_limits/tests.rs index 588eb1a5..d6547b32 100644 --- a/crates/controller/src/fleet_limits/tests.rs +++ b/crates/controller/src/fleet_limits/tests.rs @@ -507,3 +507,60 @@ mod enrollment { assert_eq!(limiter.peers.lock().unwrap().len(), 1); } } + +mod device_streams { + use tokio::{sync::oneshot::error::TryRecvError, time::Instant}; + + use super::super::{DeviceStreams, FLEET_LIMITS}; + use super::Duration; + + #[test] + fn a_newer_stream_from_the_same_device_replaces_the_older_one() { + let streams = DeviceStreams::default(); + let mut older = streams.open("device-a"); + let mut other_device = streams.open("device-b"); + + let mut newer = streams.open("device-a"); + + assert_eq!(older.superseded.try_recv(), Ok(())); + assert_eq!(other_device.superseded.try_recv(), Err(TryRecvError::Empty)); + assert_eq!(newer.superseded.try_recv(), Err(TryRecvError::Empty)); + } + + #[test] + fn a_replaced_stream_closing_keeps_the_newer_registration() { + let streams = DeviceStreams::default(); + let older = streams.open("device-a"); + let mut newer = streams.open("device-a"); + + drop(older); + let _newest = streams.open("device-a"); + + // The newer stream was still registered, so the newest one replaced it. + assert_eq!(newer.superseded.try_recv(), Ok(())); + } + + #[test] + fn a_closed_stream_is_unregistered() { + let streams = DeviceStreams::default(); + drop(streams.open("device-a")); + + assert!(streams.open.lock().unwrap().is_empty()); + } + + #[test] + fn a_stream_gets_its_message_burst_then_the_refill_rate() { + let mut limits = FLEET_LIMITS; + limits.stream_message_burst = 3; + limits.stream_message_refill = Duration::from_secs(1); + let mut budget = limits.stream_messages(); + let start = Instant::now(); + + for _ in 0..3 { + assert!(budget.take(start)); + } + assert!(!budget.take(start), "burst exceeded"); + assert!(budget.take(start + Duration::from_secs(1))); + assert!(!budget.take(start + Duration::from_secs(1))); + } +} diff --git a/crates/controller/src/service.rs b/crates/controller/src/service.rs index 54e66648..ccc66baa 100644 --- a/crates/controller/src/service.rs +++ b/crates/controller/src/service.rs @@ -22,7 +22,7 @@ use crate::{ daemon_config::DaemonConfigStore, database::{Database, DevicePrincipal}, device_ca::DeviceCertificateIssuer, - fleet_limits::{EnrollmentLimiter, FLEET_LIMITS}, + fleet_limits::{DeviceStreams, EnrollmentLimiter, FLEET_LIMITS}, gateway_jwt::GatewayJwtIssuer, oidc::{OidcPrincipal, OidcProvider}, }; @@ -36,6 +36,7 @@ pub struct FleetAgentService { device_certificate_issuer: Option, usage_enabled: bool, enrollment: EnrollmentLimiter, + device_streams: DeviceStreams, } impl FleetAgentService { @@ -54,6 +55,7 @@ impl FleetAgentService { device_certificate_issuer, usage_enabled: false, enrollment: FLEET_LIMITS.enrollment(), + device_streams: DeviceStreams::default(), } } @@ -297,6 +299,8 @@ impl FleetAgent for FleetAgentService { let database = self.database.clone(); let auth_sender = sender.clone(); + let mut device_stream = self.device_streams.open(&authenticated_device_id); + let mut message_budget = FLEET_LIMITS.stream_messages(); tokio::spawn(async move { let enrolled_principal = match database.device_principal(&authenticated_device_id).await { @@ -333,8 +337,22 @@ impl FleetAgent for FleetAgentService { } } } + _ = &mut device_stream.superseded => { + info!(device_id = %authenticated_device_id, "closing stream replaced by a newer stream from the same device"); + let _ = auth_sender + .send(Err(Status::aborted("replaced by a newer stream from this device"))) + .await; + break; + } message = inbound.message() => match message { Ok(Some(message)) => { + if !message_budget.take(time::Instant::now()) { + warn!(device_id = %authenticated_device_id, "closing stream after the device exceeded its message rate"); + let _ = auth_sender + .send(Err(Status::resource_exhausted("too many messages on this stream"))) + .await; + break; + } if let Some(agent_message::Message::Hello(hello)) = &message.message && hello.device_id != authenticated_device_id { From cc2f8d776ef4363c34acd6cdbe7efc8a2930110e Mon Sep 17 00:00:00 2001 From: Scott Weiss Date: Mon, 9 Nov 2026 10:52:32 -0500 Subject: [PATCH 09/49] fix(helm): keep device addresses visible to the controller The controller limits fleet connections and enrollment per client address. The chart's LoadBalancer Service used Kubernetes' default externalTrafficPolicy (Cluster), under which traffic can arrive from node addresses, turning the per-address limits into shared ones. Add service.externalTrafficPolicy, defaulting to Local, rendered only for LoadBalancer and NodePort Services. Setting it to "" restores the Kubernetes default. --- deploy/README.md | 9 +++++++-- .../helm/agentdesktop-controller/templates/service.yaml | 3 +++ deploy/helm/agentdesktop-controller/values.yaml | 6 ++++++ 3 files changed, 16 insertions(+), 2 deletions(-) diff --git a/deploy/README.md b/deploy/README.md index 6da1b48b..76b24057 100644 --- a/deploy/README.md +++ b/deploy/README.md @@ -40,8 +40,13 @@ Secret containing the complete controller configuration under the key `controller.yaml`, then set `existingConfigSecret` to its name. See `values.yaml` for all chart settings. -The Service exposes only the TLS-protected fleet API. To inspect the -loopback-only management UI: +The Service exposes only the TLS-protected fleet API. It defaults to +`externalTrafficPolicy: Local` so the controller sees each device's own address, +which its per-address connection and enrollment limits rely on. Set +`service.externalTrafficPolicy` to `""` to use the Kubernetes default instead; +the per-address limits then apply to the node addresses traffic arrives from. + +To inspect the loopback-only management UI: ```console kubectl port-forward deployment/agentdesktop 8080:8080 diff --git a/deploy/helm/agentdesktop-controller/templates/service.yaml b/deploy/helm/agentdesktop-controller/templates/service.yaml index d01a1fed..07918440 100644 --- a/deploy/helm/agentdesktop-controller/templates/service.yaml +++ b/deploy/helm/agentdesktop-controller/templates/service.yaml @@ -6,6 +6,9 @@ metadata: {{- include "agentdesktop-controller.labels" . | nindent 4 }} spec: type: {{ .Values.service.type }} + {{- if and .Values.service.externalTrafficPolicy (has .Values.service.type (list "LoadBalancer" "NodePort")) }} + externalTrafficPolicy: {{ .Values.service.externalTrafficPolicy }} + {{- end }} selector: {{- include "agentdesktop-controller.selectorLabels" . | nindent 4 }} ports: diff --git a/deploy/helm/agentdesktop-controller/values.yaml b/deploy/helm/agentdesktop-controller/values.yaml index 230c7918..3d01ee8e 100644 --- a/deploy/helm/agentdesktop-controller/values.yaml +++ b/deploy/helm/agentdesktop-controller/values.yaml @@ -38,6 +38,12 @@ tlsSecretName: agentdesktop-controller-tls service: type: LoadBalancer port: 443 + # Local keeps devices' own addresses visible to the controller, which limits + # connections and enrollment per client address. With Cluster, Kubernetes may + # replace them with node addresses, so the per-address limits act as shared + # ones. Applies only to LoadBalancer and NodePort services; set to "" to use + # the Kubernetes default. + externalTrafficPolicy: Local resources: {} From 5e0cbd8371a78ad9e64beaf18aea0cccbfeab496 Mon Sep 17 00:00:00 2001 From: Scott Weiss Date: Mon, 9 Nov 2026 11:30:25 -0500 Subject: [PATCH 10/49] fix(controller): close idle fleet connections, and evict idle ones at the cap After the TLS handshake, nothing closed a connection that never sent the HTTP/2 preface (hyper's protocol detection has no deadline), or one that sent it and answered keepalive PINGs without ever making a request. A client could hold such connections for free; enough of them filled the connection caps, after which every device reconnect was refused. Track each connection's request activity: a layer marks the connection busy from the start of a request until its response body is dropped, so an open Connect stream keeps it busy. A connection is closed if it makes no request within a short deadline after being accepted, or stays idle too long after its last request. At the total cap, the longest-idle connection is closed to admit a new one; a new connection is refused only when every open connection is busy. The connection handling moves to fleet_limits/connections.rs. Adds http-body and tower-layer, both already in the lockfile through tonic. --- Cargo.lock | 2 + Cargo.toml | 2 + crates/controller/Cargo.toml | 2 + crates/controller/src/application.rs | 5 +- crates/controller/src/fleet_limits.rs | 196 +------ .../src/fleet_limits/connections.rs | 545 ++++++++++++++++++ crates/controller/src/fleet_limits/tests.rs | 163 +++++- 7 files changed, 742 insertions(+), 173 deletions(-) create mode 100644 crates/controller/src/fleet_limits/connections.rs diff --git a/Cargo.lock b/Cargo.lock index 2a52f9ed..8e38e46b 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -126,6 +126,7 @@ dependencies = [ "base64 0.23.1", "clap", "futures-core", + "http-body", "http-body-util", "jsonwebtoken", "mime_guess", @@ -146,6 +147,7 @@ dependencies = [ "tokio-stream", "tonic", "tower", + "tower-layer", "tracing", "url", "uuid", diff --git a/Cargo.toml b/Cargo.toml index 36576999..714995dd 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -30,6 +30,7 @@ base64 = "0.23" bytes = "1" clap = { version = "4", features = ["derive"] } futures-core = "0.3" +http-body = "1" http-body-util = "0.1" humantime-serde = "1" hostname = "0.4" @@ -67,6 +68,7 @@ tokio-stream = "0.1" tonic = { version = "0.14", features = ["transport", "tls-aws-lc"] } tonic-prost = "0.14" tonic-prost-build = "0.14" +tower-layer = "0.3" tracing = "0.1" tracing-appender = "0.2" tracing-subscriber = { version = "0.3", features = ["env-filter", "json"] } diff --git a/crates/controller/Cargo.toml b/crates/controller/Cargo.toml index 5cb495a4..aaf0c508 100644 --- a/crates/controller/Cargo.toml +++ b/crates/controller/Cargo.toml @@ -18,6 +18,7 @@ axum-server = { version = "0.8", features = ["tls-rustls-no-provider"] } base64.workspace = true clap.workspace = true futures-core.workspace = true +http-body.workspace = true jsonwebtoken.workspace = true mime_guess = "2" notify.workspace = true @@ -35,6 +36,7 @@ sqlx = { workspace = true, features = ["postgres", "tls-rustls-aws-lc-rs"] } tokio.workspace = true tokio-stream.workspace = true tonic.workspace = true +tower-layer.workspace = true tracing.workspace = true time.workspace = true url.workspace = true diff --git a/crates/controller/src/application.rs b/crates/controller/src/application.rs index e5ffb49c..f37a6534 100644 --- a/crates/controller/src/application.rs +++ b/crates/controller/src/application.rs @@ -128,8 +128,9 @@ pub async fn run( .client_ca_root(Certificate::from_pem(client_ca)) .client_auth_optional(true), ); + let fleet_connections = FLEET_LIMITS.connections(); let mut server = FLEET_LIMITS - .server() + .server(&fleet_connections) .accept_http1(true) .tls_config(tls_config)?; @@ -144,7 +145,7 @@ pub async fn run( } server .add_routes(routes) - .serve_with_incoming(FLEET_LIMITS.incoming(fleet_listen)?) + .serve_with_incoming(FLEET_LIMITS.incoming(fleet_listen, &fleet_connections)?) .await .context("serve fleet gRPC API") }; diff --git a/crates/controller/src/fleet_limits.rs b/crates/controller/src/fleet_limits.rs index f65c2b01..b7f2083c 100644 --- a/crates/controller/src/fleet_limits.rs +++ b/crates/controller/src/fleet_limits.rs @@ -9,7 +9,6 @@ use std::{ convert::Infallible, io, net::{IpAddr, SocketAddr}, - pin::Pin, sync::{ Arc, Mutex, atomic::{AtomicU64, Ordering}, @@ -19,10 +18,7 @@ use std::{ }; use agentdesktop_proto::fleet::fleet_agent_server::{FleetAgent, FleetAgentServer, SERVICE_NAME}; -use futures_core::Stream; use tokio::{ - io::{AsyncRead, AsyncWrite, ReadBuf}, - net::TcpStream, sync::{OwnedSemaphorePermit, Semaphore, oneshot}, time::Instant, }; @@ -30,11 +26,13 @@ use tonic::{ Status, codegen::{Service, http}, server::NamedService, - transport::{ - Server, ServerTlsConfig, - server::{Connected, TcpConnectInfo, TcpIncoming}, - }, + transport::{Server, ServerTlsConfig, server::TcpIncoming}, }; +use tower_layer::{Identity, Stack}; + +mod connections; + +pub use connections::{Connections, LimitedConnection, LimitedIncoming, TrackRequests}; /// Limits applied to the fleet gRPC listener. #[derive(Clone, Copy, Debug)] @@ -59,6 +57,13 @@ pub struct FleetLimits { /// NAT, or appear to when a load balancer rewrites client addresses. pub max_connections: usize, pub connections_per_address: usize, + /// Time a new connection has to make its first request, and time a + /// connection may stay without any request in progress after its last one + /// finished. A connection with a request in progress (including an open + /// `Connect` stream) is never idle. At the total cap, the longest-idle + /// connection is closed to admit a new one. + pub first_request: Duration, + pub idle_connection: Duration, /// Concurrent requests one connection may have in progress, and concurrent /// HTTP/2 streams it may open. Further requests wait, so one connection /// cannot occupy an unbounded share of the server. @@ -96,6 +101,8 @@ pub const FLEET_LIMITS: FleetLimits = FleetLimits { tcp_keepalive: Duration::from_secs(60), max_connections: 50_000, connections_per_address: 4_096, + first_request: Duration::from_secs(20), + idle_connection: Duration::from_secs(120), requests_per_connection: 32, streams_per_connection: 64, max_message_bytes: 4 * 1024 * 1024, @@ -113,33 +120,34 @@ impl FleetLimits { config.timeout(self.tls_handshake) } + /// The registry that admits fleet connections and tracks their activity. + /// Pass the same one to `server` and `incoming`. + pub fn connections(&self) -> Connections { + Connections::new(self) + } + /// A server builder with the fleet listener's limits applied. - pub fn server(&self) -> Server { + pub fn server(&self, connections: &Connections) -> Server> { Server::builder() .timeout(self.request) .http2_keepalive_interval(Some(self.keepalive_interval)) .http2_keepalive_timeout(Some(self.keepalive_timeout)) .concurrency_limit_per_connection(self.requests_per_connection) .max_concurrent_streams(self.streams_per_connection) + .layer(connections.track_requests()) } - /// Binds the fleet listener with TCP keepalive and the connection caps. - /// Serve with `Server::serve_with_incoming`. - pub fn incoming(&self, address: SocketAddr) -> io::Result> { + /// Binds the fleet listener with TCP keepalive, admitting connections + /// through `connections`. Serve with `Server::serve_with_incoming`. + pub fn incoming( + &self, + address: SocketAddr, + connections: &Connections, + ) -> io::Result> { let incoming = TcpIncoming::bind(address)? .with_nodelay(Some(true)) .with_keepalive(Some(self.tcp_keepalive)); - Ok(self.limit_connections(incoming)) - } - - /// Applies the connection caps to accepted connections. - pub fn limit_connections(&self, incoming: S) -> LimitedIncoming { - LimitedIncoming { - incoming, - total: Arc::new(Semaphore::new(self.max_connections)), - per_address: Arc::default(), - per_address_limit: self.connections_per_address, - } + Ok(connections.incoming(incoming)) } /// The fleet gRPC service with message size limits applied. @@ -172,148 +180,6 @@ impl FleetLimits { } } -/// Accepted connections, closed on accept when over the connection caps. -pub struct LimitedIncoming { - incoming: S, - total: Arc, - per_address: Arc>>, - per_address_limit: usize, -} - -impl LimitedIncoming { - fn admit(&self, stream: TcpStream) -> Option { - let Ok(total) = self.total.clone().try_acquire_owned() else { - tracing::debug!("fleet connection limit reached; closing connection"); - return None; - }; - let address = stream.peer_addr().ok().map(|peer| peer.ip()); - let address = match address { - Some(address) => { - let mut counts = self - .per_address - .lock() - .expect("connection limiter lock poisoned"); - let count = counts.entry(address).or_default(); - if *count >= self.per_address_limit { - tracing::debug!(%address, "per-address fleet connection limit reached; closing connection"); - return None; - } - *count += 1; - Some(AddressSlot { - counts: self.per_address.clone(), - address, - }) - } - None => None, - }; - Some(LimitedConnection { - stream, - _total: total, - _address: address, - }) - } -} - -impl Stream for LimitedIncoming -where - S: Stream> + Unpin, -{ - type Item = io::Result; - - fn poll_next(mut self: Pin<&mut Self>, context: &mut Context<'_>) -> Poll> { - loop { - match Pin::new(&mut self.incoming).poll_next(context) { - Poll::Ready(Some(Ok(stream))) => { - // Over a cap: dropping the stream closes it; wait for the next one. - if let Some(connection) = self.admit(stream) { - return Poll::Ready(Some(Ok(connection))); - } - } - Poll::Ready(Some(Err(error))) => return Poll::Ready(Some(Err(error))), - Poll::Ready(None) => return Poll::Ready(None), - Poll::Pending => return Poll::Pending, - } - } - } -} - -/// Releases a client address's connection slot when its connection closes. -struct AddressSlot { - counts: Arc>>, - address: IpAddr, -} - -impl Drop for AddressSlot { - fn drop(&mut self) { - let mut counts = self - .counts - .lock() - .expect("connection limiter lock poisoned"); - if let Some(count) = counts.get_mut(&self.address) { - *count -= 1; - if *count == 0 { - counts.remove(&self.address); - } - } - } -} - -/// An accepted fleet connection holding its slots until it closes. -pub struct LimitedConnection { - stream: TcpStream, - _total: OwnedSemaphorePermit, - _address: Option, -} - -impl Connected for LimitedConnection { - // The plain TCP info keeps `Request::remote_addr` and `peer_certs` working. - type ConnectInfo = TcpConnectInfo; - - fn connect_info(&self) -> Self::ConnectInfo { - self.stream.connect_info() - } -} - -impl AsyncRead for LimitedConnection { - fn poll_read( - mut self: Pin<&mut Self>, - context: &mut Context<'_>, - buffer: &mut ReadBuf<'_>, - ) -> Poll> { - Pin::new(&mut self.stream).poll_read(context, buffer) - } -} - -impl AsyncWrite for LimitedConnection { - fn poll_write( - mut self: Pin<&mut Self>, - context: &mut Context<'_>, - buffer: &[u8], - ) -> Poll> { - Pin::new(&mut self.stream).poll_write(context, buffer) - } - - fn poll_write_vectored( - mut self: Pin<&mut Self>, - context: &mut Context<'_>, - buffers: &[io::IoSlice<'_>], - ) -> Poll> { - Pin::new(&mut self.stream).poll_write_vectored(context, buffers) - } - - fn is_write_vectored(&self) -> bool { - self.stream.is_write_vectored() - } - - fn poll_flush(mut self: Pin<&mut Self>, context: &mut Context<'_>) -> Poll> { - Pin::new(&mut self.stream).poll_flush(context) - } - - fn poll_shutdown(mut self: Pin<&mut Self>, context: &mut Context<'_>) -> Poll> { - Pin::new(&mut self.stream).poll_shutdown(context) - } -} - /// The fleet gRPC service, with a smaller request limit for the enrollment RPCs. /// Tonic sets message limits per service, so enrollment requests go to their own /// instance of the generated server. diff --git a/crates/controller/src/fleet_limits/connections.rs b/crates/controller/src/fleet_limits/connections.rs new file mode 100644 index 00000000..b0e50edf --- /dev/null +++ b/crates/controller/src/fleet_limits/connections.rs @@ -0,0 +1,545 @@ +//! Fleet connection admission and idle handling. +//! +//! Client certificates are optional on the fleet listener, so anyone can finish +//! the TLS handshake. Without more, a client could then hold a connection open +//! for free: by never sending the HTTP/2 preface (nothing times out before it), +//! or by sending it and answering keepalive PINGs without ever making a request. +//! Enough such connections would fill the connection caps and lock devices out. +//! +//! So each connection's request activity is tracked. A connection is closed if +//! it makes no request soon after it is accepted, or stays idle too long after +//! its last request finishes. A connection with a request in progress, including +//! an open `Connect` stream, is never idle. When the total cap is reached, the +//! longest-idle connection is closed to make room instead of refusing the new +//! one, so idle connections cannot crowd out devices that are using theirs. + +use std::{ + collections::{HashMap, VecDeque}, + future::Future, + io, + net::{IpAddr, SocketAddr}, + pin::Pin, + sync::{ + Arc, Mutex, + atomic::{AtomicBool, AtomicU64, AtomicUsize, Ordering}, + }, + task::{Context, Poll, Waker}, + time::Duration, +}; + +use futures_core::Stream; +use http_body::{Body as HttpBody, Frame, SizeHint}; +use tokio::{ + io::{AsyncRead, AsyncWrite, ReadBuf}, + net::TcpStream, + sync::oneshot, + time::{Instant, Sleep}, +}; +use tonic::{ + body::Body, + codegen::{Service, http}, + transport::server::{Connected, TcpConnectInfo, TlsConnectInfo}, +}; +use tower_layer::Layer; + +use super::FleetLimits; + +/// Open fleet connections and their request activity, shared by the listener +/// (which admits, times out and evicts connections) and the request layer (which +/// marks connections busy). +#[derive(Clone)] +pub struct Connections { + shared: Arc, +} + +struct Shared { + max_connections: usize, + connections_per_address: usize, + first_request: Duration, + idle: Duration, + state: Mutex, +} + +#[derive(Default)] +struct State { + open: HashMap>, + per_address: HashMap, + /// Connections in the order they became idle. An entry is stale once its + /// connection makes a request or closes; `Activity::idle_epoch` tells. + idle: VecDeque<(SocketAddr, u64)>, +} + +/// One connection's request activity. +struct Activity { + in_flight: AtomicUsize, + /// Advanced each time the connection becomes idle. + idle_epoch: AtomicU64, + /// Whether the connection has made any request. + used: AtomicBool, + /// When the connection was accepted, or last started or finished a request. + last_active: Mutex, + /// Closes the connection when it is evicted. + evict: Mutex>>, + /// Wakes the connection so it recomputes its idle deadline when its last + /// request finishes. + waker: Mutex>, +} + +impl Activity { + fn touch(&self) { + *self + .last_active + .lock() + .expect("connection activity lock poisoned") = Instant::now(); + } + + /// When an idle connection is due to close, or `None` while it is busy. + fn idle_deadline(&self, shared: &Shared) -> Option { + if self.in_flight.load(Ordering::Acquire) > 0 { + return None; + } + let allowed = if self.used.load(Ordering::Acquire) { + shared.idle + } else { + shared.first_request + }; + Some( + *self + .last_active + .lock() + .expect("connection activity lock poisoned") + + allowed, + ) + } +} + +impl State { + /// Closes the longest-idle connection. Returns false when every open + /// connection has a request in progress. + fn evict_idle(&mut self) -> bool { + while let Some((peer, epoch)) = self.idle.pop_front() { + let Some(activity) = self.open.get(&peer) else { + continue; + }; + if activity.in_flight.load(Ordering::Acquire) > 0 + || activity.idle_epoch.load(Ordering::Acquire) != epoch + { + continue; + } + if let Some(evict) = activity + .evict + .lock() + .expect("connection eviction lock poisoned") + .take() + { + let _ = evict.send(()); + return true; + } + } + false + } + + fn mark_idle(&mut self, peer: SocketAddr, epoch: u64) { + self.idle.push_back((peer, epoch)); + // Drop stale entries once they dominate, so the queue stays bounded by + // the number of open connections. + if self.idle.len() > 2 * self.open.len() + 64 { + let open = &self.open; + self.idle.retain(|(peer, epoch)| { + open.get(peer).is_some_and(|activity| { + activity.in_flight.load(Ordering::Acquire) == 0 + && activity.idle_epoch.load(Ordering::Acquire) == *epoch + }) + }); + } + } +} + +impl Connections { + pub(super) fn new(limits: &FleetLimits) -> Self { + Self { + shared: Arc::new(Shared { + max_connections: limits.max_connections, + connections_per_address: limits.connections_per_address, + first_request: limits.first_request, + idle: limits.idle_connection, + state: Mutex::default(), + }), + } + } + + fn state(&self) -> std::sync::MutexGuard<'_, State> { + self.shared + .state + .lock() + .expect("connection registry lock poisoned") + } + + /// Accepted connections, admitted or closed by the caps. + pub fn incoming(&self, incoming: S) -> LimitedIncoming { + LimitedIncoming { + incoming, + connections: self.clone(), + } + } + + /// A layer that marks a connection busy while a request (or stream) is open. + pub fn track_requests(&self) -> TrackRequests { + TrackRequests { + connections: self.clone(), + } + } + + fn admit(&self, stream: TcpStream) -> Option { + let peer = stream.peer_addr().ok()?; + let mut state = self.state(); + if state.open.len() >= self.shared.max_connections && !state.evict_idle() { + tracing::debug!( + "fleet connection limit reached and no connection is idle; closing connection" + ); + return None; + } + let count = state.per_address.entry(peer.ip()).or_default(); + if *count >= self.shared.connections_per_address { + tracing::debug!(address = %peer.ip(), "per-address fleet connection limit reached; closing connection"); + return None; + } + *count += 1; + let now = Instant::now(); + let (evict, evicted) = oneshot::channel(); + let activity = Arc::new(Activity { + in_flight: AtomicUsize::new(0), + idle_epoch: AtomicU64::new(0), + used: AtomicBool::new(false), + last_active: Mutex::new(now), + evict: Mutex::new(Some(evict)), + waker: Mutex::new(None), + }); + state.open.insert(peer, activity.clone()); + state.mark_idle(peer, 0); + drop(state); + Some(LimitedConnection { + stream, + peer, + activity, + connections: self.clone(), + evicted: Some(evicted), + timer: Box::pin(tokio::time::sleep_until(now + self.shared.first_request)), + closed: false, + }) + } + + fn release(&self, peer: SocketAddr, activity: &Arc) { + let mut state = self.state(); + if state + .open + .get(&peer) + .is_some_and(|open| Arc::ptr_eq(open, activity)) + { + state.open.remove(&peer); + } + if let Some(count) = state.per_address.get_mut(&peer.ip()) { + *count -= 1; + if *count == 0 { + state.per_address.remove(&peer.ip()); + } + } + } + + fn start_request(&self, peer: SocketAddr) -> Option { + let activity = self.state().open.get(&peer)?.clone(); + activity.in_flight.fetch_add(1, Ordering::AcqRel); + activity.used.store(true, Ordering::Release); + activity.touch(); + Some(RequestGuard { + connections: self.clone(), + peer, + activity, + }) + } + + #[cfg(test)] + pub(super) fn open_count(&self) -> usize { + self.state().open.len() + } +} + +/// Accepted connections, closed on accept when over the caps. +pub struct LimitedIncoming { + incoming: S, + connections: Connections, +} + +impl Stream for LimitedIncoming +where + S: Stream> + Unpin, +{ + type Item = io::Result; + + fn poll_next(mut self: Pin<&mut Self>, context: &mut Context<'_>) -> Poll> { + loop { + match Pin::new(&mut self.incoming).poll_next(context) { + Poll::Ready(Some(Ok(stream))) => { + // Over a cap: dropping the stream closes it; wait for the next one. + if let Some(connection) = self.connections.admit(stream) { + return Poll::Ready(Some(Ok(connection))); + } + } + Poll::Ready(Some(Err(error))) => return Poll::Ready(Some(Err(error))), + Poll::Ready(None) => return Poll::Ready(None), + Poll::Pending => return Poll::Pending, + } + } + } +} + +/// An accepted fleet connection. It holds its slots until it closes, and fails +/// its next read once it is evicted or has been idle too long. +pub struct LimitedConnection { + stream: TcpStream, + peer: SocketAddr, + activity: Arc, + connections: Connections, + evicted: Option>, + timer: Pin>, + closed: bool, +} + +impl LimitedConnection { + /// Returns an error once the connection should close, registering wakeups + /// for eviction and the idle deadline otherwise. + fn check_open(&mut self, context: &mut Context<'_>) -> io::Result<()> { + if self.closed { + return Err(io::Error::new( + io::ErrorKind::TimedOut, + "fleet connection closed", + )); + } + if let Some(evicted) = self.evicted.as_mut() + && Pin::new(evicted).poll(context).is_ready() + { + self.evicted = None; + self.closed = true; + tracing::debug!(peer = %self.peer, "evicted idle fleet connection to make room"); + return Err(io::Error::new( + io::ErrorKind::TimedOut, + "idle fleet connection evicted", + )); + } + { + let mut waker = self + .activity + .waker + .lock() + .expect("connection waker lock poisoned"); + if !waker + .as_ref() + .is_some_and(|waker| waker.will_wake(context.waker())) + { + *waker = Some(context.waker().clone()); + } + } + loop { + let now = Instant::now(); + let target = match self.activity.idle_deadline(&self.connections.shared) { + Some(deadline) if deadline <= now => { + self.closed = true; + tracing::debug!(peer = %self.peer, "closing idle fleet connection"); + return Err(io::Error::new( + io::ErrorKind::TimedOut, + "idle fleet connection closed", + )); + } + Some(deadline) => deadline, + // Busy: check again after an idle period. + None => now + self.connections.shared.idle, + }; + if self.timer.deadline() != target { + self.timer.as_mut().reset(target); + } + if self.timer.as_mut().poll(context).is_pending() { + return Ok(()); + } + } + } +} + +impl Drop for LimitedConnection { + fn drop(&mut self) { + self.connections.release(self.peer, &self.activity); + } +} + +impl Connected for LimitedConnection { + // The plain TCP info keeps `Request::remote_addr` and `peer_certs` working, + // and lets the request layer find this connection by its peer address. + type ConnectInfo = TcpConnectInfo; + + fn connect_info(&self) -> Self::ConnectInfo { + self.stream.connect_info() + } +} + +impl AsyncRead for LimitedConnection { + fn poll_read( + mut self: Pin<&mut Self>, + context: &mut Context<'_>, + buffer: &mut ReadBuf<'_>, + ) -> Poll> { + if let Err(error) = self.check_open(context) { + return Poll::Ready(Err(error)); + } + Pin::new(&mut self.stream).poll_read(context, buffer) + } +} + +impl AsyncWrite for LimitedConnection { + fn poll_write( + mut self: Pin<&mut Self>, + context: &mut Context<'_>, + buffer: &[u8], + ) -> Poll> { + Pin::new(&mut self.stream).poll_write(context, buffer) + } + + fn poll_write_vectored( + mut self: Pin<&mut Self>, + context: &mut Context<'_>, + buffers: &[io::IoSlice<'_>], + ) -> Poll> { + Pin::new(&mut self.stream).poll_write_vectored(context, buffers) + } + + fn is_write_vectored(&self) -> bool { + self.stream.is_write_vectored() + } + + fn poll_flush(mut self: Pin<&mut Self>, context: &mut Context<'_>) -> Poll> { + Pin::new(&mut self.stream).poll_flush(context) + } + + fn poll_shutdown(mut self: Pin<&mut Self>, context: &mut Context<'_>) -> Poll> { + Pin::new(&mut self.stream).poll_shutdown(context) + } +} + +/// Marks the connection busy until the request's response has been fully sent +/// (for `Connect`, until the stream ends). +struct RequestGuard { + connections: Connections, + peer: SocketAddr, + activity: Arc, +} + +impl Drop for RequestGuard { + fn drop(&mut self) { + self.activity.touch(); + if self.activity.in_flight.fetch_sub(1, Ordering::AcqRel) == 1 { + let epoch = self.activity.idle_epoch.fetch_add(1, Ordering::AcqRel) + 1; + self.connections.state().mark_idle(self.peer, epoch); + if let Some(waker) = self + .activity + .waker + .lock() + .expect("connection waker lock poisoned") + .take() + { + waker.wake(); + } + } + } +} + +/// Layer that ties each request to its connection's activity. +#[derive(Clone)] +pub struct TrackRequests { + connections: Connections, +} + +impl Layer for TrackRequests { + type Service = Tracked; + + fn layer(&self, inner: S) -> Self::Service { + Tracked { + inner, + connections: self.connections.clone(), + } + } +} + +#[derive(Clone)] +pub struct Tracked { + inner: S, + connections: Connections, +} + +fn remote_addr(request: &http::Request) -> Option { + let extensions = request.extensions(); + extensions + .get::() + .and_then(TcpConnectInfo::remote_addr) + .or_else(|| { + extensions + .get::>() + .and_then(|info| info.get_ref().remote_addr()) + }) +} + +type BoxFuture = Pin + Send>>; + +impl Service> for Tracked +where + S: Service, Response = http::Response>, + S::Future: Send + 'static, +{ + type Response = http::Response; + type Error = S::Error; + type Future = BoxFuture>; + + fn poll_ready(&mut self, context: &mut Context<'_>) -> Poll> { + self.inner.poll_ready(context) + } + + fn call(&mut self, request: http::Request) -> Self::Future { + let guard = remote_addr(&request).and_then(|peer| self.connections.start_request(peer)); + let response = self.inner.call(request); + Box::pin(async move { + let response = response.await?; + Ok(match guard { + Some(guard) => response.map(|body| { + Body::new(TrackedBody { + inner: body, + _guard: guard, + }) + }), + None => response, + }) + }) + } +} + +/// A response body that keeps its connection busy until it is dropped. +struct TrackedBody { + inner: Body, + _guard: RequestGuard, +} + +impl HttpBody for TrackedBody { + type Data = ::Data; + type Error = ::Error; + + fn poll_frame( + mut self: Pin<&mut Self>, + context: &mut Context<'_>, + ) -> Poll, Self::Error>>> { + Pin::new(&mut self.inner).poll_frame(context) + } + + fn is_end_stream(&self) -> bool { + self.inner.is_end_stream() + } + + fn size_hint(&self) -> SizeHint { + self.inner.size_hint() + } +} diff --git a/crates/controller/src/fleet_limits/tests.rs b/crates/controller/src/fleet_limits/tests.rs index d6547b32..50ef6657 100644 --- a/crates/controller/src/fleet_limits/tests.rs +++ b/crates/controller/src/fleet_limits/tests.rs @@ -20,7 +20,7 @@ use tonic::{ transport::{Identity, ServerTlsConfig, server::TcpIncoming}, }; -use super::{FLEET_LIMITS, FleetLimits}; +use super::{Connections, FLEET_LIMITS, FleetLimits}; /// Limits small enough for tests to observe; anything a test does not exercise /// stays at the production value. @@ -109,17 +109,27 @@ impl FleetAgent for Slow { /// Serves `service` with `limits`, over TLS when `identity` is given. async fn serve(limits: FleetLimits, identity: Option, service: Slow) -> SocketAddr { + serve_tracked(limits, identity, service).await.0 +} + +/// Like `serve`, also returning the connection registry. +async fn serve_tracked( + limits: FleetLimits, + identity: Option, + service: Slow, +) -> (SocketAddr, Connections) { let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); let address = listener.local_addr().unwrap(); - let mut server = limits.server(); + let connections = limits.connections(); + let mut server = limits.server(&connections); if let Some(identity) = identity { server = server .tls_config(limits.tls(ServerTlsConfig::new().identity(identity))) .unwrap(); } let router = server.add_service(limits.fleet_service(service)); - tokio::spawn(router.serve_with_incoming(limits.limit_connections(TcpIncoming::from(listener)))); - address + tokio::spawn(router.serve_with_incoming(connections.incoming(TcpIncoming::from(listener)))); + (address, connections) } async fn client(address: SocketAddr) -> FleetAgentClient { @@ -388,11 +398,152 @@ async fn connections_over_the_per_address_cap_are_closed() { } #[tokio::test] -async fn connections_over_the_total_cap_are_closed() { +async fn at_the_total_cap_the_longest_idle_connection_makes_room() { let mut limits = limits(); limits.max_connections = 2; + let address = serve( + limits, + None, + Slow { + delay: Duration::ZERO, + }, + ) + .await; + let oldest = TcpStream::connect(address).await.unwrap(); + time::sleep(Duration::from_millis(50)).await; + let newer = TcpStream::connect(address).await.unwrap(); + time::sleep(Duration::from_millis(50)).await; - connections_over_cap(limits).await; + let newest = TcpStream::connect(address).await.unwrap(); + + bounded(until_closed(oldest)).await; + assert!(stays_open(newest).await, "the new connection was refused"); + drop(newer); +} + +#[tokio::test] +async fn at_the_total_cap_a_new_connection_is_closed_when_every_connection_is_busy() { + let mut limits = limits(); + limits.max_connections = 1; + limits.request = Duration::from_secs(5); + let address = serve( + limits, + None, + Slow { + delay: Duration::from_millis(500), + }, + ) + .await; + let (_outbound, receiver) = mpsc::channel::(1); + let mut stream = bounded(client(address).await.connect(ReceiverStream::new(receiver))) + .await + .expect("open the stream") + .into_inner(); + + let newcomer = TcpStream::connect(address).await.unwrap(); + bounded(until_closed(newcomer)).await; + + // The busy connection was kept. + assert!(bounded(stream.message()).await.unwrap().is_some()); +} + +#[tokio::test] +async fn a_connection_that_never_sends_anything_is_closed() { + let mut limits = limits(); + limits.first_request = Duration::from_millis(200); + let address = serve( + limits, + None, + Slow { + delay: Duration::ZERO, + }, + ) + .await; + + let stream = TcpStream::connect(address).await.unwrap(); + + bounded(until_closed(stream)).await; +} + +#[tokio::test] +async fn a_connection_that_never_makes_a_request_is_closed() { + let mut limits = limits(); + limits.first_request = Duration::from_millis(200); + // Keepalive alone would also close this connection; keep it out of the way. + limits.keepalive_interval = Duration::from_secs(60); + limits.keepalive_timeout = Duration::from_secs(60); + let address = serve( + limits, + None, + Slow { + delay: Duration::ZERO, + }, + ) + .await; + let mut stream = TcpStream::connect(address).await.unwrap(); + + // A valid HTTP/2 preface and SETTINGS frame, then no request. + stream + .write_all(b"PRI * HTTP/2.0\r\n\r\nSM\r\n\r\n\x00\x00\x00\x04\x00\x00\x00\x00\x00") + .await + .unwrap(); + + bounded(until_closed(stream)).await; +} + +#[tokio::test] +async fn a_connection_idle_after_its_last_request_is_closed() { + let mut limits = limits(); + limits.first_request = Duration::from_secs(5); + limits.idle_connection = Duration::from_millis(200); + let (address, connections) = serve_tracked( + limits, + None, + Slow { + delay: Duration::ZERO, + }, + ) + .await; + let mut client = client(address).await; + bounded(client.begin_enrollment(BeginEnrollmentRequest::default())) + .await + .unwrap(); + assert_eq!(connections.open_count(), 1); + + bounded(async { + while connections.open_count() > 0 { + time::sleep(Duration::from_millis(50)).await; + } + }) + .await; +} + +#[tokio::test] +async fn an_open_stream_keeps_its_connection_from_going_idle() { + let mut limits = limits(); + limits.request = Duration::from_secs(5); + limits.idle_connection = Duration::from_millis(200); + // Messages arrive further apart than the idle limit. + let address = serve( + limits, + None, + Slow { + delay: Duration::from_millis(500), + }, + ) + .await; + let (_outbound, receiver) = mpsc::channel::(1); + let mut stream = bounded(client(address).await.connect(ReceiverStream::new(receiver))) + .await + .expect("open the stream") + .into_inner(); + + for message in 0..2 { + let received = bounded(stream.message()) + .await + .unwrap_or_else(|error| panic!("stream ended at message {message}: {error}")); + assert!(received.is_some(), "stream closed at message {message}"); + } } mod enrollment { From 41b9a93f291ab2dd7e1bb5922a7323ce1d8eb0c3 Mon Sep 17 00:00:00 2001 From: Scott Weiss Date: Mon, 9 Nov 2026 11:31:39 -0500 Subject: [PATCH 11/49] fix(controller): apply the small message cap to callers without a device certificate The 64 KiB cap covered only the enrollment RPCs. Tonic decodes a unary request before its handler runs, and the device RPCs check the client certificate in the handler, so a caller without a certificate could still send 4 MiB requests to them. Route every request without a device certificate to the small-cap service instance, not only enrollment. Tests now use mutual TLS in the production shape (device CA, client certificate optional), which also shows device certificates still reach handlers through the connection wrapper. --- crates/controller/src/fleet_limits.rs | 45 ++++--- crates/controller/src/fleet_limits/tests.rs | 136 +++++++++++++++----- 2 files changed, 136 insertions(+), 45 deletions(-) diff --git a/crates/controller/src/fleet_limits.rs b/crates/controller/src/fleet_limits.rs index b7f2083c..11e0e59a 100644 --- a/crates/controller/src/fleet_limits.rs +++ b/crates/controller/src/fleet_limits.rs @@ -26,7 +26,10 @@ use tonic::{ Status, codegen::{Service, http}, server::NamedService, - transport::{Server, ServerTlsConfig, server::TcpIncoming}, + transport::{ + Server, ServerTlsConfig, + server::{TcpConnectInfo, TcpIncoming, TlsConnectInfo}, + }, }; use tower_layer::{Identity, Stack}; @@ -71,9 +74,11 @@ pub struct FleetLimits { pub streams_per_connection: u32, /// Largest encoded message an enrolled device's calls may send or receive. pub max_message_bytes: usize, - /// Largest enrollment request. Enrollment runs before a device is - /// authenticated, and its requests are small (a hostname and a CSR). - pub enrollment_message_bytes: usize, + /// Largest request from a caller without a device certificate: enrollment, + /// which runs before a device has one, and any other call made without one + /// (rejected by its handler, but only after the request has been decoded). + /// Enrollment requests are small (a hostname and a CSR). + pub unauthenticated_message_bytes: usize, /// Enrollment requests (`BeginEnrollment`, `CompleteEnrollment`) running at /// once across all clients. Enrollment calls the identity provider, so this /// also bounds the load a flood can put on it. @@ -106,7 +111,7 @@ pub const FLEET_LIMITS: FleetLimits = FleetLimits { requests_per_connection: 32, streams_per_connection: 64, max_message_bytes: 4 * 1024 * 1024, - enrollment_message_bytes: 64 * 1024, + unauthenticated_message_bytes: 64 * 1024, enrollment_in_flight: 64, enrollment_burst: 60, enrollment_refill: Duration::from_secs(1), @@ -154,8 +159,8 @@ impl FleetLimits { pub fn fleet_service(&self, service: T) -> FleetService { let service = Arc::new(service); FleetService { - enrollment: FleetAgentServer::from_arc(service.clone()) - .max_decoding_message_size(self.enrollment_message_bytes) + unauthenticated: FleetAgentServer::from_arc(service.clone()) + .max_decoding_message_size(self.unauthenticated_message_bytes) .max_encoding_message_size(self.max_message_bytes), devices: FleetAgentServer::from_arc(service) .max_decoding_message_size(self.max_message_bytes) @@ -180,18 +185,18 @@ impl FleetLimits { } } -/// The fleet gRPC service, with a smaller request limit for the enrollment RPCs. -/// Tonic sets message limits per service, so enrollment requests go to their own -/// instance of the generated server. +/// The fleet gRPC service, with a smaller request limit for enrollment and for +/// callers without a device certificate. Tonic sets message limits per service, +/// so those requests go to their own instance of the generated server. pub struct FleetService { - enrollment: FleetAgentServer, + unauthenticated: FleetAgentServer, devices: FleetAgentServer, } impl Clone for FleetService { fn clone(&self) -> Self { Self { - enrollment: self.enrollment.clone(), + unauthenticated: self.unauthenticated.clone(), devices: self.devices.clone(), } } @@ -203,6 +208,16 @@ fn is_enrollment(path: &str) -> bool { .is_some_and(|method| matches!(method, "/BeginEnrollment" | "/CompleteEnrollment")) } +/// Whether the TLS client presented a certificate signed by the device CA +/// (the handshake rejects any other). +fn has_device_certificate(request: &http::Request) -> bool { + request + .extensions() + .get::>() + .and_then(TlsConnectInfo::peer_certs) + .is_some_and(|certificates| !certificates.is_empty()) +} + impl Service> for FleetService where FleetAgentServer: Service, Error = Infallible>, @@ -212,15 +227,15 @@ where type Future = as Service>>::Future; fn poll_ready(&mut self, context: &mut Context<'_>) -> Poll> { - match self.enrollment.poll_ready(context) { + match self.unauthenticated.poll_ready(context) { Poll::Ready(Ok(())) => self.devices.poll_ready(context), other => other, } } fn call(&mut self, request: http::Request) -> Self::Future { - if is_enrollment(request.uri().path()) { - self.enrollment.call(request) + if is_enrollment(request.uri().path()) || !has_device_certificate(&request) { + self.unauthenticated.call(request) } else { self.devices.call(request) } diff --git a/crates/controller/src/fleet_limits/tests.rs b/crates/controller/src/fleet_limits/tests.rs index 50ef6657..ace79443 100644 --- a/crates/controller/src/fleet_limits/tests.rs +++ b/crates/controller/src/fleet_limits/tests.rs @@ -109,23 +109,22 @@ impl FleetAgent for Slow { /// Serves `service` with `limits`, over TLS when `identity` is given. async fn serve(limits: FleetLimits, identity: Option, service: Slow) -> SocketAddr { - serve_tracked(limits, identity, service).await.0 + let tls = identity.map(|identity| ServerTlsConfig::new().identity(identity)); + serve_tracked(limits, tls, service).await.0 } -/// Like `serve`, also returning the connection registry. +/// Like `serve`, with any TLS configuration, also returning the connection registry. async fn serve_tracked( limits: FleetLimits, - identity: Option, + tls: Option, service: Slow, ) -> (SocketAddr, Connections) { let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); let address = listener.local_addr().unwrap(); let connections = limits.connections(); let mut server = limits.server(&connections); - if let Some(identity) = identity { - server = server - .tls_config(limits.tls(ServerTlsConfig::new().identity(identity))) - .unwrap(); + if let Some(tls) = tls { + server = server.tls_config(limits.tls(tls)).unwrap(); } let router = server.add_service(limits.fleet_service(service)); tokio::spawn(router.serve_with_incoming(connections.incoming(TcpIncoming::from(listener)))); @@ -288,7 +287,7 @@ async fn streams_beyond_the_per_connection_limit_wait() { fn small_messages() -> FleetLimits { let mut limits = limits(); limits.max_message_bytes = 1024; - limits.enrollment_message_bytes = 128; + limits.unauthenticated_message_bytes = 128; limits } @@ -317,38 +316,115 @@ async fn an_oversized_enrollment_request_is_rejected() { assert_eq!(error.code(), Code::OutOfRange, "{error}"); } -#[tokio::test] -async fn device_calls_keep_the_larger_message_limit() { - let address = serve( - small_messages(), - None, +/// A device CA with a server certificate for `localhost` and a device certificate. +struct Pki { + ca: tonic::transport::Certificate, + server: Identity, + device: Identity, +} + +fn pki() -> Pki { + use rcgen::{ + BasicConstraints, CertificateParams, CertifiedIssuer, IsCa, KeyPair, KeyUsagePurpose, + }; + + let mut ca_params = CertificateParams::new(Vec::::new()).unwrap(); + ca_params.is_ca = IsCa::Ca(BasicConstraints::Unconstrained); + ca_params.key_usages = vec![ + KeyUsagePurpose::KeyCertSign, + KeyUsagePurpose::DigitalSignature, + ]; + let ca = CertifiedIssuer::self_signed(ca_params, KeyPair::generate().unwrap()).unwrap(); + let leaf = |name: &str| { + let key = KeyPair::generate().unwrap(); + let certificate = CertificateParams::new(vec![name.to_owned()]) + .unwrap() + .signed_by(&key, &ca) + .unwrap(); + Identity::from_pem(certificate.pem(), key.serialize_pem()) + }; + Pki { + ca: tonic::transport::Certificate::from_pem(ca.pem()), + server: leaf("localhost"), + device: leaf("device"), + } +} + +/// Serves with the production TLS shape: device CA, client certificate optional. +async fn serve_mtls(limits: FleetLimits, pki: &Pki) -> SocketAddr { + let tls = ServerTlsConfig::new() + .identity(pki.server.clone()) + .client_ca_root(pki.ca.clone()) + .client_auth_optional(true); + serve_tracked( + limits, + Some(tls), Slow { delay: Duration::ZERO, }, ) - .await; - let mut client = client(address).await; - - // Over the enrollment limit but within the device limit: reaches the handler. - let within = bounded( - client.renew_device_certificate(RenewDeviceCertificateRequest { - certificate_signing_request_der: vec![0; 512], - }), - ) .await - .expect_err("the test service rejects every renewal"); + .0 +} + +async fn tls_client( + address: SocketAddr, + pki: &Pki, + device: bool, +) -> FleetAgentClient { + let mut tls = tonic::transport::ClientTlsConfig::new() + .ca_certificate(pki.ca.clone()) + .domain_name("localhost"); + if device { + tls = tls.identity(pki.device.clone()); + } + let channel = tonic::transport::Endpoint::from_shared(format!("https://{address}")) + .unwrap() + .tls_config(tls) + .unwrap() + .connect() + .await + .unwrap(); + FleetAgentClient::new(channel) +} + +fn renewal(bytes: usize) -> RenewDeviceCertificateRequest { + RenewDeviceCertificateRequest { + certificate_signing_request_der: vec![0; bytes], + } +} + +#[tokio::test] +async fn device_calls_keep_the_larger_message_limit() { + let pki = pki(); + let address = serve_mtls(small_messages(), &pki).await; + let mut client = tls_client(address, &pki, true).await; + + // Over the unauthenticated limit but within the device limit: reaches the handler. + let within = bounded(client.renew_device_certificate(renewal(512))) + .await + .expect_err("the test service rejects every renewal"); assert_eq!(within.code(), Code::Unimplemented, "{within}"); - let over = bounded( - client.renew_device_certificate(RenewDeviceCertificateRequest { - certificate_signing_request_der: vec![0; 4096], - }), - ) - .await - .expect_err("an oversized device request was accepted"); + let over = bounded(client.renew_device_certificate(renewal(4096))) + .await + .expect_err("an oversized device request was accepted"); assert_eq!(over.code(), Code::OutOfRange, "{over}"); } +#[tokio::test] +async fn callers_without_a_device_certificate_get_the_small_limit_on_every_call() { + let pki = pki(); + let address = serve_mtls(small_messages(), &pki).await; + let mut client = tls_client(address, &pki, false).await; + + let error = bounded(client.renew_device_certificate(renewal(512))) + .await + .expect_err("a large request without a certificate reached the handler"); + + assert_eq!(error.code(), Code::OutOfRange, "{error}"); +} + /// Whether the server keeps `stream` open for a while. An accepted connection /// receives the server's HTTP/2 SETTINGS and then waits for the client; a /// rejected one is closed (end of stream) right away. From 381faef5f35ffbf91da145da73f963402eb73a82 Mon Sep 17 00:00:00 2001 From: Scott Weiss Date: Mon, 9 Nov 2026 11:32:21 -0500 Subject: [PATCH 12/49] fix(controller): end the Connect response when the device's side ends A separate task forwarded daemon config pushes to each Connect response and held a sender for it. When the device half-closed its request stream or sent a mismatched hello, the inbound task exited and unregistered the stream, but the response stayed open on that other task. So a device could keep many streams per connection despite the one-stream rule, and every forwarding task (and its channel) outlived its stream until the next config change, including after ordinary reconnects. Forward config pushes from the stream's own loop, so it is the only holder of the response sender, and also stop when the device goes away. A mismatched hello now ends the stream with PERMISSION_DENIED instead of silently. Not covered end to end: Connect needs a device certificate and OIDC, which the test harness does not provide. --- crates/controller/src/service.rs | 45 ++++++++++++++++++-------------- 1 file changed, 25 insertions(+), 20 deletions(-) diff --git a/crates/controller/src/service.rs b/crates/controller/src/service.rs index ccc66baa..8d2ca1ff 100644 --- a/crates/controller/src/service.rs +++ b/crates/controller/src/service.rs @@ -278,27 +278,12 @@ impl FleetAgent for FleetAgentService { .map_err(|_| Status::unavailable("stream closed"))?; } - let update_sender = sender.clone(); - tokio::spawn(async move { - while daemon_updates.changed().await.is_ok() { - let daemon = daemon_updates.borrow().clone(); - let Some(daemon) = daemon else { - continue; - }; - if update_sender - .send(Ok(ControllerMessage { - message: Some(controller_message::Message::DaemonConfig(daemon)), - })) - .await - .is_err() - { - break; - } - } - }); - + // One task owns the response stream: it ends when the device's request + // stream ends, the device goes away, or the stream is closed for cause. + // Config pushes are forwarded from the same loop, so no other task keeps + // the response open. let database = self.database.clone(); - let auth_sender = sender.clone(); + let auth_sender = sender; let mut device_stream = self.device_streams.open(&authenticated_device_id); let mut message_budget = FLEET_LIMITS.stream_messages(); tokio::spawn(async move { @@ -314,6 +299,23 @@ impl FleetAgent for FleetAgentService { auth_interval.tick().await; loop { tokio::select! { + _ = auth_sender.closed() => break, + changed = daemon_updates.changed() => { + if changed.is_err() { + break; + } + let daemon = daemon_updates.borrow_and_update().clone(); + if let Some(daemon) = daemon + && auth_sender + .send(Ok(ControllerMessage { + message: Some(controller_message::Message::DaemonConfig(daemon)), + })) + .await + .is_err() + { + break; + } + } _ = auth_interval.tick() => { let (oidc, access_token) = &oauth_revalidation; match oidc.authenticate_access_token(access_token).await { @@ -361,6 +363,9 @@ impl FleetAgent for FleetAgentService { claimed_device_id = %hello.device_id, "device certificate and hello identity do not match" ); + let _ = auth_sender + .send(Err(Status::permission_denied("hello does not match the device certificate"))) + .await; break; } if let Err(err) = From d38e52de1d7f50df05f9e6740a65bc1982e9ff53 Mon Sep 17 00:00:00 2001 From: Scott Weiss Date: Mon, 9 Nov 2026 11:35:25 -0500 Subject: [PATCH 13/49] fix(controller): keep Connect budgets per device and limit how often a device opens Connect The message budget lived on each Connect stream, so a device could reset it by reconnecting. Each new stream also checked the access token with the identity provider before anything limited it. Keep the message budget with the device, so a reconnect keeps what is left of it, and give each device a budget for opening Connect streams (20 at once, then one a minute). The open budget is taken from the device certificate before the access token is checked. --- crates/controller/src/fleet_limits.rs | 162 ++++++++++++++------ crates/controller/src/fleet_limits/tests.rs | 64 ++++++-- crates/controller/src/service.rs | 15 +- 3 files changed, 180 insertions(+), 61 deletions(-) diff --git a/crates/controller/src/fleet_limits.rs b/crates/controller/src/fleet_limits.rs index 11e0e59a..c2d10f7e 100644 --- a/crates/controller/src/fleet_limits.rs +++ b/crates/controller/src/fleet_limits.rs @@ -88,13 +88,19 @@ pub struct FleetLimits { /// many devices can share one address behind NAT. pub enrollment_burst: u32, pub enrollment_refill: Duration, - /// Messages a device may send on one `Connect` stream: up to + /// Messages a device may send on its `Connect` streams: up to /// `stream_message_burst` at once, refilled at one per - /// `stream_message_refill`. Devices send a heartbeat every 30 seconds, - /// inventory periodically and one telemetry event per tool event, so this is - /// far above normal use; a stream over budget is closed. + /// `stream_message_refill`. The budget belongs to the device, so + /// reconnecting does not reset it. Devices send a heartbeat every 30 + /// seconds, inventory periodically and one telemetry event per tool event, + /// so this is far above normal use; a stream over budget is closed. pub stream_message_burst: u32, pub stream_message_refill: Duration, + /// How often a device may open a `Connect` stream: up to + /// `stream_open_burst` at once, refilled at one per `stream_open_refill`. + /// Each open checks the access token with the identity provider. + pub stream_open_burst: u32, + pub stream_open_refill: Duration, } /// Production limits. @@ -117,6 +123,8 @@ pub const FLEET_LIMITS: FleetLimits = FleetLimits { enrollment_refill: Duration::from_secs(1), stream_message_burst: 600, stream_message_refill: Duration::from_millis(50), + stream_open_burst: 20, + stream_open_refill: Duration::from_secs(60), }; impl FleetLimits { @@ -172,16 +180,9 @@ impl FleetLimits { EnrollmentLimiter::new(self) } - /// A fresh message budget for one `Connect` stream. - pub fn stream_messages(&self) -> MessageBudget { - MessageBudget { - bucket: Bucket { - tokens: f64::from(self.stream_message_burst), - updated: Instant::now(), - }, - burst: f64::from(self.stream_message_burst), - refill: self.stream_message_refill, - } + /// Per-device `Connect` state: one stream per device, and its budgets. + pub fn device_streams(&self) -> DeviceStreams { + DeviceStreams::new(self) } } @@ -317,32 +318,28 @@ fn take_token(bucket: &mut Bucket, now: Instant, refill: Duration, burst: f64) - true } -/// Messages a device may send on one `Connect` stream. -pub struct MessageBudget { - bucket: Bucket, - burst: f64, - refill: Duration, -} - -impl MessageBudget { - /// Takes one message from the budget, or returns false when it is spent. - pub fn take(&mut self, now: Instant) -> bool { - take_token(&mut self.bucket, now, self.refill, self.burst) - } -} - -/// Open `Connect` streams by device. A device has at most one: opening a new -/// stream tells the device's previous stream to close. A device that reconnects -/// (after sleep, for example) is never blocked by a stream the controller has -/// not yet noticed is dead. -#[derive(Clone, Default)] +/// `Connect` state by device. A device has at most one open stream: opening a +/// new stream tells the device's previous one to close, so a device that +/// reconnects (after sleep, for example) is never blocked by a stream the +/// controller has not yet noticed is dead. The device's message budget and its +/// budget for opening streams live here too, so reconnecting does not reset them. +#[derive(Clone)] pub struct DeviceStreams { - open: Arc>>, + devices: Arc>>, next: Arc, + limits: FleetLimits, +} + +struct DeviceState { + /// The open stream's generation and the signal that replaces it. + current: Option<(u64, oneshot::Sender<()>)>, + messages: Bucket, + opens: Bucket, } -/// A device's current stream: its generation and the signal that replaces it. -type Registration = (u64, oneshot::Sender<()>); +/// Device entries kept before idle ones (no open stream) are dropped. Only +/// devices with a certificate from the device CA get an entry. +const MAX_TRACKED_DEVICES: usize = 100_000; /// A registered `Connect` stream. `superseded` resolves when the same device /// opens a newer stream; dropping this unregisters it. @@ -354,14 +351,68 @@ pub struct DeviceStream { } impl DeviceStreams { + pub fn new(limits: &FleetLimits) -> Self { + Self { + devices: Arc::default(), + next: Arc::default(), + limits: *limits, + } + } + + fn with_device(&self, device_id: &str, f: impl FnOnce(&mut DeviceState) -> T) -> T { + let mut devices = self + .devices + .lock() + .expect("device stream registry lock poisoned"); + if devices.len() >= MAX_TRACKED_DEVICES && !devices.contains_key(device_id) { + devices.retain(|_, device| device.current.is_some()); + } + let now = Instant::now(); + let device = devices + .entry(device_id.to_owned()) + .or_insert_with(|| DeviceState { + current: None, + messages: Bucket { + tokens: f64::from(self.limits.stream_message_burst), + updated: now, + }, + opens: Bucket { + tokens: f64::from(self.limits.stream_open_burst), + updated: now, + }, + }); + f(device) + } + + /// Admits one more `Connect` from `device_id`, or rejects it with + /// `RESOURCE_EXHAUSTED` when the device opens streams too often. Call this + /// before any per-stream work, such as checking the access token. + pub fn admit_open(&self, device_id: &str) -> Result<(), Status> { + let limits = self.limits; + let admitted = self.with_device(device_id, |device| { + take_token( + &mut device.opens, + Instant::now(), + limits.stream_open_refill, + f64::from(limits.stream_open_burst), + ) + }); + if admitted { + Ok(()) + } else { + tracing::warn!(device_id, "device opened Connect streams too often"); + Err(Status::resource_exhausted( + "this device opened too many streams; retry later", + )) + } + } + pub fn open(&self, device_id: &str) -> DeviceStream { let generation = self.next.fetch_add(1, Ordering::Relaxed); let (replace, superseded) = oneshot::channel(); - let previous = self - .open - .lock() - .expect("device stream registry lock poisoned") - .insert(device_id.to_owned(), (generation, replace)); + let previous = self.with_device(device_id, |device| { + device.current.replace((generation, replace)) + }); if let Some((_, previous)) = previous { let _ = previous.send(()); } @@ -374,19 +425,36 @@ impl DeviceStreams { } } +impl DeviceStream { + /// Takes one message from the device's budget, or returns false when it is spent. + pub fn take_message(&self, now: Instant) -> bool { + let limits = self.streams.limits; + self.streams.with_device(&self.device_id, |device| { + take_token( + &mut device.messages, + now, + limits.stream_message_refill, + f64::from(limits.stream_message_burst), + ) + }) + } +} + impl Drop for DeviceStream { fn drop(&mut self) { - let mut open = self + let mut devices = self .streams - .open + .devices .lock() .expect("device stream registry lock poisoned"); // A newer stream for the same device keeps its registration. - if open - .get(&self.device_id) - .is_some_and(|(generation, _)| *generation == self.generation) + if let Some(device) = devices.get_mut(&self.device_id) + && device + .current + .as_ref() + .is_some_and(|(generation, _)| *generation == self.generation) { - open.remove(&self.device_id); + device.current = None; } } } diff --git a/crates/controller/src/fleet_limits/tests.rs b/crates/controller/src/fleet_limits/tests.rs index ace79443..a6452b6a 100644 --- a/crates/controller/src/fleet_limits/tests.rs +++ b/crates/controller/src/fleet_limits/tests.rs @@ -741,9 +741,13 @@ mod device_streams { use super::super::{DeviceStreams, FLEET_LIMITS}; use super::Duration; + fn streams() -> DeviceStreams { + FLEET_LIMITS.device_streams() + } + #[test] fn a_newer_stream_from_the_same_device_replaces_the_older_one() { - let streams = DeviceStreams::default(); + let streams = streams(); let mut older = streams.open("device-a"); let mut other_device = streams.open("device-b"); @@ -756,7 +760,7 @@ mod device_streams { #[test] fn a_replaced_stream_closing_keeps_the_newer_registration() { - let streams = DeviceStreams::default(); + let streams = streams(); let older = streams.open("device-a"); let mut newer = streams.open("device-a"); @@ -769,25 +773,65 @@ mod device_streams { #[test] fn a_closed_stream_is_unregistered() { - let streams = DeviceStreams::default(); + let streams = streams(); drop(streams.open("device-a")); - assert!(streams.open.lock().unwrap().is_empty()); + assert!( + streams.devices.lock().unwrap()["device-a"] + .current + .is_none(), + "the closed stream is still registered" + ); } #[test] - fn a_stream_gets_its_message_burst_then_the_refill_rate() { + fn a_device_gets_its_message_burst_then_the_refill_rate() { let mut limits = FLEET_LIMITS; limits.stream_message_burst = 3; limits.stream_message_refill = Duration::from_secs(1); - let mut budget = limits.stream_messages(); + let stream = limits.device_streams().open("device-a"); let start = Instant::now(); for _ in 0..3 { - assert!(budget.take(start)); + assert!(stream.take_message(start)); } - assert!(!budget.take(start), "burst exceeded"); - assert!(budget.take(start + Duration::from_secs(1))); - assert!(!budget.take(start + Duration::from_secs(1))); + assert!(!stream.take_message(start), "burst exceeded"); + assert!(stream.take_message(start + Duration::from_secs(1))); + assert!(!stream.take_message(start + Duration::from_secs(1))); + } + + #[test] + fn reconnecting_does_not_reset_the_message_budget() { + let mut limits = FLEET_LIMITS; + limits.stream_message_burst = 2; + limits.stream_message_refill = Duration::from_secs(3600); + let streams = limits.device_streams(); + let now = Instant::now(); + let first = streams.open("device-a"); + assert!(first.take_message(now)); + assert!(first.take_message(now)); + drop(first); + + let second = streams.open("device-a"); + + assert!(!second.take_message(now), "a new stream got a fresh budget"); + } + + #[test] + fn a_device_that_opens_streams_too_often_is_refused() { + let mut limits = FLEET_LIMITS; + limits.stream_open_burst = 2; + let streams = limits.device_streams(); + + streams.admit_open("device-a").expect("first open"); + streams.admit_open("device-a").expect("second open"); + let error = streams + .admit_open("device-a") + .expect_err("third open admitted"); + + assert_eq!(error.code(), tonic::Code::ResourceExhausted); + streams + .admit_open("device-b") + .expect("another device is unaffected"); } } diff --git a/crates/controller/src/service.rs b/crates/controller/src/service.rs index 8d2ca1ff..fffca3ff 100644 --- a/crates/controller/src/service.rs +++ b/crates/controller/src/service.rs @@ -55,7 +55,7 @@ impl FleetAgentService { device_certificate_issuer, usage_enabled: false, enrollment: FLEET_LIMITS.enrollment(), - device_streams: DeviceStreams::default(), + device_streams: FLEET_LIMITS.device_streams(), } } @@ -257,7 +257,15 @@ impl FleetAgent for FleetAgentService { &self, request: Request>, ) -> Result, Status> { - let authenticated_device_id = self.authenticate_device(&request).await?; + // Rate-limit opens per device before checking the access token, which + // calls the identity provider. + let certificate_device_id = peer_certificate_device_id(&request)?; + if let Some(device_id) = &certificate_device_id { + self.device_streams.admit_open(device_id)?; + } + let authenticated_device_id = self + .authenticate_device_identity(certificate_device_id.as_deref(), request.metadata()) + .await?; let oauth_revalidation = ( self.oidc.clone().ok_or_else(|| { Status::failed_precondition("OIDC access-token validation is disabled") @@ -285,7 +293,6 @@ impl FleetAgent for FleetAgentService { let database = self.database.clone(); let auth_sender = sender; let mut device_stream = self.device_streams.open(&authenticated_device_id); - let mut message_budget = FLEET_LIMITS.stream_messages(); tokio::spawn(async move { let enrolled_principal = match database.device_principal(&authenticated_device_id).await { @@ -348,7 +355,7 @@ impl FleetAgent for FleetAgentService { } message = inbound.message() => match message { Ok(Some(message)) => { - if !message_budget.take(time::Instant::now()) { + if !device_stream.take_message(time::Instant::now()) { warn!(device_id = %authenticated_device_id, "closing stream after the device exceeded its message rate"); let _ = auth_sender .send(Err(Status::resource_exhausted("too many messages on this stream"))) From e8d114b482992b614e4b1e4a483222dc8af69967 Mon Sep 17 00:00:00 2001 From: Scott Weiss Date: Mon, 9 Nov 2026 11:36:43 -0500 Subject: [PATCH 14/49] fix(controller): count IPv6 clients by /64 and bound the enrollment peer table Per-address limits keyed on the full IPv6 address, so one host with a /64 could open the per-address connection cap and get a fresh enrollment budget from every address it holds. Count IPv6 clients by their /64 instead, and IPv4-mapped addresses as IPv4. The enrollment peer table dropped only refilled buckets when it was full, so a client cycling through addresses kept it full and made every new peer scan the whole table. Keep it in two generations instead: when the current one fills it replaces the previous one, and a peer seen in the previous one moves back with its bucket. --- crates/controller/src/fleet_limits.rs | 73 +++++++++++++++---- .../src/fleet_limits/connections.rs | 13 ++-- crates/controller/src/fleet_limits/tests.rs | 58 ++++++++++++--- 3 files changed, 113 insertions(+), 31 deletions(-) diff --git a/crates/controller/src/fleet_limits.rs b/crates/controller/src/fleet_limits.rs index c2d10f7e..002b3a2a 100644 --- a/crates/controller/src/fleet_limits.rs +++ b/crates/controller/src/fleet_limits.rs @@ -8,7 +8,7 @@ use std::{ collections::HashMap, convert::Infallible, io, - net::{IpAddr, SocketAddr}, + net::{IpAddr, Ipv6Addr, SocketAddr}, sync::{ Arc, Mutex, atomic::{AtomicU64, Ordering}, @@ -54,8 +54,9 @@ pub struct FleetLimits { pub keepalive_timeout: Duration, /// TCP keepalive probe interval, for peers that vanish without closing. pub tcp_keepalive: Duration, - /// Open connections allowed in total and from one client address. A - /// connection over either cap is closed as soon as it is accepted. The + /// Open connections allowed in total and from one client address (an IPv4 + /// address or an IPv6 /64). A connection over either cap is closed as soon + /// as it is accepted. The /// per-address cap is high because many devices can share an address behind /// NAT, or appear to when a load balancer rewrites client addresses. pub max_connections: usize, @@ -247,17 +248,59 @@ impl NamedService for FleetService { const NAME: &'static str = SERVICE_NAME; } +/// The address a per-address limit counts against: an IPv4 address, or the /64 +/// network of an IPv6 one, since one IPv6 host usually controls a whole /64. +pub(crate) fn address_key(address: IpAddr) -> IpAddr { + match address { + IpAddr::V4(_) => address, + IpAddr::V6(v6) => match v6.to_ipv4_mapped() { + Some(v4) => IpAddr::V4(v4), + None => IpAddr::V6(Ipv6Addr::from(u128::from(v6) & (u128::MAX << 64))), + }, + } +} + /// Admission control for the unauthenticated enrollment RPCs. #[derive(Clone)] pub struct EnrollmentLimiter { in_flight: Arc, - peers: Arc>>, + peers: Arc>, burst: f64, refill: Duration, } -/// Peer buckets kept before idle (full) ones are dropped. -const MAX_TRACKED_PEERS: usize = 10_000; +/// Enrollment buckets by address, in two generations. When the current +/// generation fills up it becomes the previous one, and the old previous one is +/// dropped; a peer found in the previous generation moves back to the current +/// one. Active peers keep their buckets, and the table never holds more than +/// two generations. +#[derive(Default)] +struct PeerBuckets { + current: HashMap, + previous: HashMap, +} + +const PEERS_PER_GENERATION: usize = 10_000; + +impl PeerBuckets { + fn bucket(&mut self, peer: IpAddr, fresh: Bucket) -> &mut Bucket { + if !self.current.contains_key(&peer) { + let bucket = self.previous.remove(&peer).unwrap_or(fresh); + if self.current.len() >= PEERS_PER_GENERATION { + self.previous = std::mem::take(&mut self.current); + } + self.current.insert(peer, bucket); + } + self.current + .get_mut(&peer) + .expect("bucket was just inserted") + } + + #[cfg(test)] + fn len(&self) -> usize { + self.current.len() + self.previous.len() + } +} #[derive(Clone, Copy)] struct Bucket { @@ -281,7 +324,7 @@ impl EnrollmentLimiter { if let Some(peer) = peer && !self.take(peer.ip(), Instant::now()) { - tracing::warn!(peer = %peer.ip(), "enrollment rate limit exceeded"); + tracing::warn!(peer = %address_key(peer.ip()), "enrollment rate limit exceeded"); return Err(Status::resource_exhausted( "too many enrollment requests from this address; retry later", )); @@ -294,15 +337,13 @@ impl EnrollmentLimiter { fn take(&self, peer: IpAddr, now: Instant) -> bool { let mut peers = self.peers.lock().expect("enrollment limiter lock poisoned"); - if peers.len() >= MAX_TRACKED_PEERS && !peers.contains_key(&peer) { - let burst = self.burst; - let refill = self.refill; - peers.retain(|_, bucket| refilled(*bucket, now, refill, burst) < burst); - } - let bucket = peers.entry(peer).or_insert(Bucket { - tokens: self.burst, - updated: now, - }); + let bucket = peers.bucket( + address_key(peer), + Bucket { + tokens: self.burst, + updated: now, + }, + ); take_token(bucket, now, self.refill, self.burst) } } diff --git a/crates/controller/src/fleet_limits/connections.rs b/crates/controller/src/fleet_limits/connections.rs index b0e50edf..49ac537a 100644 --- a/crates/controller/src/fleet_limits/connections.rs +++ b/crates/controller/src/fleet_limits/connections.rs @@ -42,7 +42,7 @@ use tonic::{ }; use tower_layer::Layer; -use super::FleetLimits; +use super::{FleetLimits, address_key}; /// Open fleet connections and their request activity, shared by the listener /// (which admits, times out and evicts connections) and the request layer (which @@ -63,6 +63,7 @@ struct Shared { #[derive(Default)] struct State { open: HashMap>, + /// Open connections by `address_key`. per_address: HashMap, /// Connections in the order they became idle. An entry is stale once its /// connection makes a request or closes; `Activity::idle_epoch` tells. @@ -199,9 +200,10 @@ impl Connections { ); return None; } - let count = state.per_address.entry(peer.ip()).or_default(); + let address = address_key(peer.ip()); + let count = state.per_address.entry(address).or_default(); if *count >= self.shared.connections_per_address { - tracing::debug!(address = %peer.ip(), "per-address fleet connection limit reached; closing connection"); + tracing::debug!(%address, "per-address fleet connection limit reached; closing connection"); return None; } *count += 1; @@ -238,10 +240,11 @@ impl Connections { { state.open.remove(&peer); } - if let Some(count) = state.per_address.get_mut(&peer.ip()) { + let address = address_key(peer.ip()); + if let Some(count) = state.per_address.get_mut(&address) { *count -= 1; if *count == 0 { - state.per_address.remove(&peer.ip()); + state.per_address.remove(&address); } } } diff --git a/crates/controller/src/fleet_limits/tests.rs b/crates/controller/src/fleet_limits/tests.rs index a6452b6a..22c166f7 100644 --- a/crates/controller/src/fleet_limits/tests.rs +++ b/crates/controller/src/fleet_limits/tests.rs @@ -623,11 +623,11 @@ async fn an_open_stream_keeps_its_connection_from_going_idle() { } mod enrollment { - use std::net::{IpAddr, Ipv4Addr, SocketAddr}; + use std::net::{IpAddr, Ipv4Addr, Ipv6Addr, SocketAddr}; use tokio::time::Instant; - use super::super::{FLEET_LIMITS, FleetLimits, MAX_TRACKED_PEERS}; + use super::super::{FLEET_LIMITS, FleetLimits, PEERS_PER_GENERATION, address_key}; use super::Duration; fn limits(in_flight: usize, burst: u32) -> FleetLimits { @@ -720,18 +720,56 @@ mod enrollment { assert!(limiter.admit(None).is_err()); } + fn numbered(index: usize) -> IpAddr { + let [a, b, c, d] = u32::try_from(index).unwrap().to_be_bytes(); + IpAddr::V4(Ipv4Addr::new(10 | a, b, c, d)) + } + #[test] - fn idle_peers_are_dropped_when_the_table_is_full() { + fn the_peer_table_stays_bounded_while_addresses_keep_changing() { let limiter = limits(64, 1).enrollment(); - let start = Instant::now(); - for index in 0..MAX_TRACKED_PEERS { - let [a, b, c, d] = u32::try_from(index).unwrap().to_be_bytes(); - assert!(limiter.take(IpAddr::V4(Ipv4Addr::new(10 | a, b, c, d)), start)); + let now = Instant::now(); + + for index in 0..5 * PEERS_PER_GENERATION { + assert!(limiter.take(numbered(index), now)); } - // Once every bucket has refilled, a new peer evicts the idle ones. - assert!(limiter.take(peer(1), start + Duration::from_secs(10))); - assert_eq!(limiter.peers.lock().unwrap().len(), 1); + assert!(limiter.peers.lock().unwrap().len() <= 2 * PEERS_PER_GENERATION); + } + + #[test] + fn a_peer_seen_recently_keeps_its_spent_budget_when_the_table_turns_over() { + let limiter = limits(64, 1).enrollment(); + let now = Instant::now(); + assert!(limiter.take(peer(1), now)); + + // Fill a generation so peer 1 moves to the previous one. + for index in 0..PEERS_PER_GENERATION { + assert!(limiter.take(numbered(index), now)); + } + + assert!(!limiter.take(peer(1), now), "the spent budget was reset"); + } + + fn v6(address: &str) -> IpAddr { + address.parse::().unwrap().into() + } + + #[test] + fn addresses_in_one_ipv6_64_share_a_budget() { + let limiter = limits(64, 1).enrollment(); + let now = Instant::now(); + + assert!(limiter.take(v6("2001:db8:1:2::1"), now)); + assert!(!limiter.take(v6("2001:db8:1:2:ffff::9"), now)); + assert!(limiter.take(v6("2001:db8:1:3::1"), now), "another /64"); + } + + #[test] + fn ipv4_mapped_addresses_count_as_ipv4() { + assert_eq!(address_key(v6("::ffff:192.0.2.1")), peer(1)); + assert_eq!(address_key(peer(1)), peer(1)); + assert_eq!(address_key(v6("2001:db8::1")), v6("2001:db8::")); } } From 5124d4569764ded723d7a719792887a159aad22f Mon Sep 17 00:00:00 2001 From: Scott Weiss Date: Mon, 9 Nov 2026 11:38:17 -0500 Subject: [PATCH 15/49] fix(controller): let CompleteEnrollment wait for a slot instead of failing CompleteEnrollment shared BeginEnrollment's per-address budget and failed at once when every in-flight slot was taken. By then the user has signed in through the browser, and a refusal makes them start over; behind a busy NAT that can happen to every device on it. CompleteEnrollment no longer spends the per-address budget. It needs an enrollment that BeginEnrollment started, which already spent it, and an unknown enrollment fails at a map lookup. It now waits up to 10 seconds for an in-flight slot before it is refused. --- crates/controller/src/fleet_limits.rs | 33 +++++++++++++-- crates/controller/src/fleet_limits/tests.rs | 45 +++++++++++++++++++++ crates/controller/src/service.rs | 3 +- 3 files changed, 76 insertions(+), 5 deletions(-) diff --git a/crates/controller/src/fleet_limits.rs b/crates/controller/src/fleet_limits.rs index 002b3a2a..47e30bb0 100644 --- a/crates/controller/src/fleet_limits.rs +++ b/crates/controller/src/fleet_limits.rs @@ -84,11 +84,15 @@ pub struct FleetLimits { /// once across all clients. Enrollment calls the identity provider, so this /// also bounds the load a flood can put on it. pub enrollment_in_flight: usize, - /// Per-client-address enrollment budget: up to `enrollment_burst` requests at - /// once, refilled at one request per `enrollment_refill`. Generous, because - /// many devices can share one address behind NAT. + /// Per-client-address `BeginEnrollment` budget: up to `enrollment_burst` + /// requests at once, refilled at one request per `enrollment_refill`. + /// Generous, because many devices can share one address behind NAT. pub enrollment_burst: u32, pub enrollment_refill: Duration, + /// How long `CompleteEnrollment` waits for an in-flight slot before it is + /// refused. It waits rather than failing at once because the user has + /// already signed in, and a refusal makes them start over. + pub enrollment_completion_wait: Duration, /// Messages a device may send on its `Connect` streams: up to /// `stream_message_burst` at once, refilled at one per /// `stream_message_refill`. The budget belongs to the device, so @@ -122,6 +126,7 @@ pub const FLEET_LIMITS: FleetLimits = FleetLimits { enrollment_in_flight: 64, enrollment_burst: 60, enrollment_refill: Duration::from_secs(1), + enrollment_completion_wait: Duration::from_secs(10), stream_message_burst: 600, stream_message_refill: Duration::from_millis(50), stream_open_burst: 20, @@ -267,6 +272,7 @@ pub struct EnrollmentLimiter { peers: Arc>, burst: f64, refill: Duration, + completion_wait: Duration, } /// Enrollment buckets by address, in two generations. When the current @@ -315,10 +321,11 @@ impl EnrollmentLimiter { peers: Arc::default(), burst: f64::from(limits.enrollment_burst), refill: limits.enrollment_refill, + completion_wait: limits.enrollment_completion_wait, } } - /// Admits one enrollment request from `peer`, or rejects it with + /// Admits one `BeginEnrollment` from `peer`, or rejects it with /// `RESOURCE_EXHAUSTED`. Hold the returned permit until the request finishes. pub fn admit(&self, peer: Option) -> Result { if let Some(peer) = peer @@ -335,6 +342,24 @@ impl EnrollmentLimiter { }) } + /// Admits one `CompleteEnrollment`, waiting up to `enrollment_completion_wait` + /// for an in-flight slot. It does not spend the address's budget: completing + /// needs an enrollment that `BeginEnrollment` started (and paid for), and each + /// enrollment completes at most once, so an unknown one fails at a lookup. + pub async fn admit_completion(&self) -> Result { + match tokio::time::timeout(self.completion_wait, self.in_flight.clone().acquire_owned()) + .await + { + Ok(Ok(permit)) => Ok(permit), + _ => { + tracing::warn!("enrollment concurrency limit reached"); + Err(Status::resource_exhausted( + "too many enrollment requests in progress; retry later", + )) + } + } + } + fn take(&self, peer: IpAddr, now: Instant) -> bool { let mut peers = self.peers.lock().expect("enrollment limiter lock poisoned"); let bucket = peers.bucket( diff --git a/crates/controller/src/fleet_limits/tests.rs b/crates/controller/src/fleet_limits/tests.rs index 22c166f7..7167ffaf 100644 --- a/crates/controller/src/fleet_limits/tests.rs +++ b/crates/controller/src/fleet_limits/tests.rs @@ -712,6 +712,51 @@ mod enrollment { assert_eq!(error.code(), tonic::Code::ResourceExhausted); } + #[tokio::test] + async fn completing_an_enrollment_does_not_spend_the_address_budget() { + let limiter = limits(64, 1).enrollment(); + let address = Some(SocketAddr::new(peer(1), 4000)); + let _begin = limiter.admit(address).expect("begin"); + + let _complete = limiter.admit_completion().await.expect("complete"); + } + + fn one_slot(completion_wait: Duration) -> FleetLimits { + FleetLimits { + enrollment_completion_wait: completion_wait, + ..limits(1, 60) + } + } + + #[tokio::test] + async fn completing_an_enrollment_waits_for_a_slot() { + let limiter = one_slot(Duration::from_secs(5)).enrollment(); + let begin = limiter.admit(None).expect("begin"); + let waiting = tokio::spawn({ + let limiter = limiter.clone(); + async move { limiter.admit_completion().await.map(drop) } + }); + tokio::time::sleep(Duration::from_millis(100)).await; + assert!(!waiting.is_finished(), "completion did not wait"); + + drop(begin); + + waiting + .await + .unwrap() + .expect("completion admitted once a slot freed"); + } + + #[tokio::test] + async fn completing_an_enrollment_gives_up_after_the_wait() { + let limiter = one_slot(Duration::from_millis(50)).enrollment(); + let _begin = limiter.admit(None).expect("begin"); + + let error = limiter.admit_completion().await.expect_err("admitted"); + + assert_eq!(error.code(), tonic::Code::ResourceExhausted); + } + #[tokio::test] async fn requests_without_a_peer_address_still_count_against_the_cap() { let limiter = limits(1, 60).enrollment(); diff --git a/crates/controller/src/service.rs b/crates/controller/src/service.rs index fffca3ff..f4fef99d 100644 --- a/crates/controller/src/service.rs +++ b/crates/controller/src/service.rs @@ -103,7 +103,7 @@ impl FleetAgent for FleetAgentService { &self, request: Request, ) -> Result, Status> { - let _admitted = self.enrollment.admit(request.remote_addr())?; + let _admitted = self.enrollment.admit_completion().await?; let oidc = self .oidc .as_ref() @@ -722,6 +722,7 @@ mod tests { .expect("connect database"); let full = FleetLimits { enrollment_in_flight: 0, + enrollment_completion_wait: std::time::Duration::from_millis(10), ..FLEET_LIMITS }; let service = From 88c58cbcefc1c416bee4506617d256be30df687e Mon Sep 17 00:00:00 2001 From: Scott Weiss Date: Mon, 9 Nov 2026 11:39:05 -0500 Subject: [PATCH 16/49] fix(controller): warn when fleet limits refuse clients, at most once per interval Connections closed at the connection caps were logged at debug, so an operator had no sign that a NAT or a load balancer was pushing devices over a per-address cap. The enrollment and Connect refusals were logged at warn once per request, so a flood would flood the log too. Log every refusal at warn, but at most once per 10 seconds for each kind, with the number of events the line stands for. --- crates/controller/src/fleet_limits.rs | 85 ++++++++++++++++--- .../src/fleet_limits/connections.rs | 30 +++++-- crates/controller/src/fleet_limits/tests.rs | 12 +++ 3 files changed, 109 insertions(+), 18 deletions(-) diff --git a/crates/controller/src/fleet_limits.rs b/crates/controller/src/fleet_limits.rs index 47e30bb0..13b3c403 100644 --- a/crates/controller/src/fleet_limits.rs +++ b/crates/controller/src/fleet_limits.rs @@ -331,15 +331,21 @@ impl EnrollmentLimiter { if let Some(peer) = peer && !self.take(peer.ip(), Instant::now()) { - tracing::warn!(peer = %address_key(peer.ip()), "enrollment rate limit exceeded"); + if let Some(events) = ENROLLMENT_RATE.event(Instant::now()) { + tracing::warn!( + peer = %address_key(peer.ip()), + events, + "refused enrollment requests over the per-address rate limit" + ); + } return Err(Status::resource_exhausted( "too many enrollment requests from this address; retry later", )); } - self.in_flight.clone().try_acquire_owned().map_err(|_| { - tracing::warn!("enrollment concurrency limit reached"); - Status::resource_exhausted("too many enrollment requests in progress; retry later") - }) + self.in_flight + .clone() + .try_acquire_owned() + .map_err(|_| in_flight_refusal()) } /// Admits one `CompleteEnrollment`, waiting up to `enrollment_completion_wait` @@ -351,12 +357,7 @@ impl EnrollmentLimiter { .await { Ok(Ok(permit)) => Ok(permit), - _ => { - tracing::warn!("enrollment concurrency limit reached"); - Err(Status::resource_exhausted( - "too many enrollment requests in progress; retry later", - )) - } + _ => Err(in_flight_refusal()), } } @@ -373,6 +374,16 @@ impl EnrollmentLimiter { } } +fn in_flight_refusal() -> Status { + if let Some(events) = ENROLLMENT_IN_FLIGHT.event(Instant::now()) { + tracing::warn!( + events, + "refused enrollment requests with every in-flight slot taken" + ); + } + Status::resource_exhausted("too many enrollment requests in progress; retry later") +} + /// Refills `bucket` for the time since its last use, then takes one token. fn take_token(bucket: &mut Bucket, now: Instant, refill: Duration, burst: f64) -> bool { bucket.tokens = refilled(*bucket, now, refill, burst); @@ -466,7 +477,13 @@ impl DeviceStreams { if admitted { Ok(()) } else { - tracing::warn!(device_id, "device opened Connect streams too often"); + if let Some(events) = STREAM_OPENS.event(Instant::now()) { + tracing::warn!( + device_id, + events, + "refused Connect streams from devices that open them too often" + ); + } Err(Status::resource_exhausted( "this device opened too many streams; retry later", )) @@ -530,5 +547,49 @@ fn refilled(bucket: Bucket, now: Instant, refill: Duration, burst: f64) -> f64 { (bucket.tokens + elapsed.as_secs_f64() / refill.as_secs_f64()).min(burst) } +/// A warning that fires once per refused request or connection. A flood would +/// log a line per event, so this logs at most one per `WARNING_INTERVAL` and +/// says how many events it stands for. +pub(crate) struct ThrottledWarning { + state: Mutex, +} + +struct WarningState { + last: Option, + events: u64, +} + +const WARNING_INTERVAL: Duration = Duration::from_secs(10); + +impl ThrottledWarning { + pub(crate) const fn new() -> Self { + Self { + state: Mutex::new(WarningState { + last: None, + events: 0, + }), + } + } + + /// Counts one event. Returns the number of events since the last warning, + /// including this one, when it is time to log another. + pub(crate) fn event(&self, now: Instant) -> Option { + let mut state = self.state.lock().unwrap_or_else(|error| error.into_inner()); + state.events += 1; + if state + .last + .is_some_and(|last| now.saturating_duration_since(last) < WARNING_INTERVAL) + { + return None; + } + state.last = Some(now); + Some(std::mem::take(&mut state.events)) + } +} + +static ENROLLMENT_RATE: ThrottledWarning = ThrottledWarning::new(); +static ENROLLMENT_IN_FLIGHT: ThrottledWarning = ThrottledWarning::new(); +static STREAM_OPENS: ThrottledWarning = ThrottledWarning::new(); + #[cfg(test)] mod tests; diff --git a/crates/controller/src/fleet_limits/connections.rs b/crates/controller/src/fleet_limits/connections.rs index 49ac537a..73ff10d0 100644 --- a/crates/controller/src/fleet_limits/connections.rs +++ b/crates/controller/src/fleet_limits/connections.rs @@ -42,7 +42,11 @@ use tonic::{ }; use tower_layer::Layer; -use super::{FleetLimits, address_key}; +use super::{FleetLimits, ThrottledWarning, address_key}; + +static TOTAL_CAP: ThrottledWarning = ThrottledWarning::new(); +static PER_ADDRESS_CAP: ThrottledWarning = ThrottledWarning::new(); +static EVICTIONS: ThrottledWarning = ThrottledWarning::new(); /// Open fleet connections and their request activity, shared by the listener /// (which admits, times out and evicts connections) and the request layer (which @@ -195,15 +199,24 @@ impl Connections { let peer = stream.peer_addr().ok()?; let mut state = self.state(); if state.open.len() >= self.shared.max_connections && !state.evict_idle() { - tracing::debug!( - "fleet connection limit reached and no connection is idle; closing connection" - ); + if let Some(events) = TOTAL_CAP.event(Instant::now()) { + tracing::warn!( + events, + "closed new fleet connections at the connection limit with no connection idle" + ); + } return None; } let address = address_key(peer.ip()); let count = state.per_address.entry(address).or_default(); if *count >= self.shared.connections_per_address { - tracing::debug!(%address, "per-address fleet connection limit reached; closing connection"); + if let Some(events) = PER_ADDRESS_CAP.event(Instant::now()) { + tracing::warn!( + %address, + events, + "closed new fleet connections over the per-address connection limit" + ); + } return None; } *count += 1; @@ -323,7 +336,12 @@ impl LimitedConnection { { self.evicted = None; self.closed = true; - tracing::debug!(peer = %self.peer, "evicted idle fleet connection to make room"); + if let Some(events) = EVICTIONS.event(Instant::now()) { + tracing::warn!( + events, + "closed idle fleet connections to make room at the connection limit" + ); + } return Err(io::Error::new( io::ErrorKind::TimedOut, "idle fleet connection evicted", diff --git a/crates/controller/src/fleet_limits/tests.rs b/crates/controller/src/fleet_limits/tests.rs index 7167ffaf..848b4e64 100644 --- a/crates/controller/src/fleet_limits/tests.rs +++ b/crates/controller/src/fleet_limits/tests.rs @@ -622,6 +622,18 @@ async fn an_open_stream_keeps_its_connection_from_going_idle() { } } +#[test] +fn a_repeated_warning_is_logged_once_per_interval_with_its_count() { + let warning = super::ThrottledWarning::new(); + let start = tokio::time::Instant::now(); + + assert_eq!(warning.event(start), Some(1)); + assert_eq!(warning.event(start + Duration::from_secs(1)), None); + assert_eq!(warning.event(start + Duration::from_secs(9)), None); + assert_eq!(warning.event(start + Duration::from_secs(10)), Some(3)); + assert_eq!(warning.event(start + Duration::from_secs(11)), None); +} + mod enrollment { use std::net::{IpAddr, Ipv4Addr, Ipv6Addr, SocketAddr}; From a6e7f80fa81361efca64524bc88e9ba277248b54 Mon Sep 17 00:00:00 2001 From: Scott Weiss Date: Mon, 9 Nov 2026 11:41:38 -0500 Subject: [PATCH 17/49] feat(controller): configure the deployment-dependent fleet limits The fleet limits were constants. The ones that depend on the deployment can be wrong for it: a NAT or a load balancer that hides client addresses puts many devices behind one address, a large fleet needs more connections, and large inventories need larger messages. Add a fleetLimits section to the controller configuration for the connection caps, the device message size, the enrollment limits and the device message budget. Unset fields keep the defaults, and zero values are rejected. The Helm chart passes fleetLimits through to controller.yaml. Rename the per-device message budget fields to match. --- crates/controller/src/application.rs | 14 ++- crates/controller/src/fleet_limits.rs | 47 +++++--- crates/controller/src/fleet_limits/tests.rs | 32 ++++- crates/controller/src/service.rs | 11 +- crates/core/src/config.rs | 111 ++++++++++++++++++ .../templates/configmap.yaml | 4 + .../helm/agentdesktop-controller/values.yaml | 5 + schema/controller-config.json | 73 ++++++++++++ schema/controller-config.md | 9 ++ 9 files changed, 277 insertions(+), 29 deletions(-) diff --git a/crates/controller/src/application.rs b/crates/controller/src/application.rs index f37a6534..045ee59e 100644 --- a/crates/controller/src/application.rs +++ b/crates/controller/src/application.rs @@ -103,6 +103,7 @@ pub async fn run( let device_certificate_issuer = DeviceCertificateIssuer::from_pem(ca_certificate, &ca_key) .context("initialize device certificate issuer")?; tracing::info!("device certificate issuance enabled"); + let limits = FLEET_LIMITS.configured(&config.fleet_limits); let service = FleetAgentService::new( Some(oidc), database, @@ -110,7 +111,8 @@ pub async fn run( gateway_jwt_issuer, Some(device_certificate_issuer), ) - .with_usage_collection(usage_ingest.is_some()); + .with_usage_collection(usage_ingest.is_some()) + .with_fleet_limits(&limits); let certificate = std::fs::read(&tls.certificate) .with_context(|| format!("read TLS certificate from {}", tls.certificate.display()))?; @@ -122,14 +124,14 @@ pub async fn run( tls.client_ca_certificate.display() ) })?; - let tls_config = FLEET_LIMITS.tls( + let tls_config = limits.tls( ServerTlsConfig::new() .identity(Identity::from_pem(certificate, key)) .client_ca_root(Certificate::from_pem(client_ca)) .client_auth_optional(true), ); - let fleet_connections = FLEET_LIMITS.connections(); - let mut server = FLEET_LIMITS + let fleet_connections = limits.connections(); + let mut server = limits .server(&fleet_connections) .accept_http1(true) .tls_config(tls_config)?; @@ -139,13 +141,13 @@ pub async fn run( tracing::info!(listen = %fleet_listen, "fleet controller listening"); let fleet = async move { let mut routes = Routes::from(gateway_jwt::routes(gateway_jwks)) - .add_service(FLEET_LIMITS.fleet_service(service)); + .add_service(limits.fleet_service(service)); if let Some(usage_ingest) = usage_ingest { routes = routes.add_service(usage_ingest); } server .add_routes(routes) - .serve_with_incoming(FLEET_LIMITS.incoming(fleet_listen, &fleet_connections)?) + .serve_with_incoming(limits.incoming(fleet_listen, &fleet_connections)?) .await .context("serve fleet gRPC API") }; diff --git a/crates/controller/src/fleet_limits.rs b/crates/controller/src/fleet_limits.rs index 13b3c403..83f97867 100644 --- a/crates/controller/src/fleet_limits.rs +++ b/crates/controller/src/fleet_limits.rs @@ -17,6 +17,7 @@ use std::{ time::Duration, }; +use agentdesktop_core::config::ControllerFleetLimitsConfig; use agentdesktop_proto::fleet::fleet_agent_server::{FleetAgent, FleetAgentServer, SERVICE_NAME}; use tokio::{ sync::{OwnedSemaphorePermit, Semaphore, oneshot}, @@ -38,7 +39,7 @@ mod connections; pub use connections::{Connections, LimitedConnection, LimitedIncoming, TrackRequests}; /// Limits applied to the fleet gRPC listener. -#[derive(Clone, Copy, Debug)] +#[derive(Clone, Copy, Debug, PartialEq)] pub struct FleetLimits { /// Time a client has to finish the TLS handshake after connecting. Tonic sets /// no limit by default, so a client could hold a connection open forever by @@ -94,13 +95,13 @@ pub struct FleetLimits { /// already signed in, and a refusal makes them start over. pub enrollment_completion_wait: Duration, /// Messages a device may send on its `Connect` streams: up to - /// `stream_message_burst` at once, refilled at one per - /// `stream_message_refill`. The budget belongs to the device, so + /// `device_message_burst` at once, refilled at one per + /// `device_message_refill`. The budget belongs to the device, so /// reconnecting does not reset it. Devices send a heartbeat every 30 /// seconds, inventory periodically and one telemetry event per tool event, /// so this is far above normal use; a stream over budget is closed. - pub stream_message_burst: u32, - pub stream_message_refill: Duration, + pub device_message_burst: u32, + pub device_message_refill: Duration, /// How often a device may open a `Connect` stream: up to /// `stream_open_burst` at once, refilled at one per `stream_open_refill`. /// Each open checks the access token with the identity provider. @@ -108,7 +109,8 @@ pub struct FleetLimits { pub stream_open_refill: Duration, } -/// Production limits. +/// Default limits. The controller configuration can change the +/// deployment-dependent ones (see `configured`). pub const FLEET_LIMITS: FleetLimits = FleetLimits { tls_handshake: Duration::from_secs(10), request: Duration::from_secs(30), @@ -127,13 +129,27 @@ pub const FLEET_LIMITS: FleetLimits = FleetLimits { enrollment_burst: 60, enrollment_refill: Duration::from_secs(1), enrollment_completion_wait: Duration::from_secs(10), - stream_message_burst: 600, - stream_message_refill: Duration::from_millis(50), + device_message_burst: 600, + device_message_refill: Duration::from_millis(50), stream_open_burst: 20, stream_open_refill: Duration::from_secs(60), }; impl FleetLimits { + /// These limits, with the deployment-dependent ones from the controller + /// configuration. + pub fn configured(mut self, config: &ControllerFleetLimitsConfig) -> Self { + self.max_connections = config.max_connections; + self.connections_per_address = config.connections_per_address; + self.max_message_bytes = config.max_message_bytes; + self.enrollment_in_flight = config.enrollment_in_flight; + self.enrollment_burst = config.enrollment_burst; + self.enrollment_refill = config.enrollment_refill; + self.device_message_burst = config.device_message_burst; + self.device_message_refill = config.device_message_refill; + self + } + /// Adds the handshake timeout to the fleet listener's TLS settings. pub fn tls(&self, config: ServerTlsConfig) -> ServerTlsConfig { config.timeout(self.tls_handshake) @@ -414,8 +430,10 @@ struct DeviceState { opens: Bucket, } -/// Device entries kept before idle ones (no open stream) are dropped. Only -/// devices with a certificate from the device CA get an entry. +/// Device entries kept before idle ones (no open stream) are dropped, or twice +/// `max_connections` if that is more: each open stream holds a connection, so +/// dropping idle entries always frees at least half the table. Only devices +/// with a certificate from the device CA get an entry. const MAX_TRACKED_DEVICES: usize = 100_000; /// A registered `Connect` stream. `superseded` resolves when the same device @@ -441,7 +459,8 @@ impl DeviceStreams { .devices .lock() .expect("device stream registry lock poisoned"); - if devices.len() >= MAX_TRACKED_DEVICES && !devices.contains_key(device_id) { + let max_devices = MAX_TRACKED_DEVICES.max(self.limits.max_connections.saturating_mul(2)); + if devices.len() >= max_devices && !devices.contains_key(device_id) { devices.retain(|_, device| device.current.is_some()); } let now = Instant::now(); @@ -450,7 +469,7 @@ impl DeviceStreams { .or_insert_with(|| DeviceState { current: None, messages: Bucket { - tokens: f64::from(self.limits.stream_message_burst), + tokens: f64::from(self.limits.device_message_burst), updated: now, }, opens: Bucket { @@ -516,8 +535,8 @@ impl DeviceStream { take_token( &mut device.messages, now, - limits.stream_message_refill, - f64::from(limits.stream_message_burst), + limits.device_message_refill, + f64::from(limits.device_message_burst), ) }) } diff --git a/crates/controller/src/fleet_limits/tests.rs b/crates/controller/src/fleet_limits/tests.rs index 848b4e64..abbbb779 100644 --- a/crates/controller/src/fleet_limits/tests.rs +++ b/crates/controller/src/fleet_limits/tests.rs @@ -622,6 +622,30 @@ async fn an_open_stream_keeps_its_connection_from_going_idle() { } } +#[test] +fn the_configuration_defaults_match_the_default_limits() { + use agentdesktop_core::config::ControllerFleetLimitsConfig; + + assert_eq!( + FLEET_LIMITS.configured(&ControllerFleetLimitsConfig::default()), + FLEET_LIMITS + ); +} + +#[test] +fn configured_limits_replace_the_defaults() { + use agentdesktop_core::config::ControllerFleetLimitsConfig; + + let limits = FLEET_LIMITS.configured(&ControllerFleetLimitsConfig { + connections_per_address: 20_000, + device_message_refill: Duration::from_millis(10), + ..ControllerFleetLimitsConfig::default() + }); + + assert_eq!(limits.connections_per_address, 20_000); + assert_eq!(limits.device_message_refill, Duration::from_millis(10)); +} + #[test] fn a_repeated_warning_is_logged_once_per_interval_with_its_count() { let warning = super::ThrottledWarning::new(); @@ -882,8 +906,8 @@ mod device_streams { #[test] fn a_device_gets_its_message_burst_then_the_refill_rate() { let mut limits = FLEET_LIMITS; - limits.stream_message_burst = 3; - limits.stream_message_refill = Duration::from_secs(1); + limits.device_message_burst = 3; + limits.device_message_refill = Duration::from_secs(1); let stream = limits.device_streams().open("device-a"); let start = Instant::now(); @@ -898,8 +922,8 @@ mod device_streams { #[test] fn reconnecting_does_not_reset_the_message_budget() { let mut limits = FLEET_LIMITS; - limits.stream_message_burst = 2; - limits.stream_message_refill = Duration::from_secs(3600); + limits.device_message_burst = 2; + limits.device_message_refill = Duration::from_secs(3600); let streams = limits.device_streams(); let now = Instant::now(); let first = streams.open("device-a"); diff --git a/crates/controller/src/service.rs b/crates/controller/src/service.rs index f4fef99d..d638fcc2 100644 --- a/crates/controller/src/service.rs +++ b/crates/controller/src/service.rs @@ -22,7 +22,7 @@ use crate::{ daemon_config::DaemonConfigStore, database::{Database, DevicePrincipal}, device_ca::DeviceCertificateIssuer, - fleet_limits::{DeviceStreams, EnrollmentLimiter, FLEET_LIMITS}, + fleet_limits::{DeviceStreams, EnrollmentLimiter, FLEET_LIMITS, FleetLimits}, gateway_jwt::GatewayJwtIssuer, oidc::{OidcPrincipal, OidcProvider}, }; @@ -59,9 +59,10 @@ impl FleetAgentService { } } - /// Replaces the enrollment admission limits (tests use small ones). - pub fn with_enrollment_limiter(mut self, enrollment: EnrollmentLimiter) -> Self { - self.enrollment = enrollment; + /// Replaces the default enrollment and `Connect` limits. + pub fn with_fleet_limits(mut self, limits: &FleetLimits) -> Self { + self.enrollment = limits.enrollment(); + self.device_streams = limits.device_streams(); self } @@ -727,7 +728,7 @@ mod tests { }; let service = FleetAgentService::new(None, database, DaemonConfigStore::new(None), None, None) - .with_enrollment_limiter(full.enrollment()); + .with_fleet_limits(&full); // Without the limiter these fail later, with "OIDC enrollment is disabled". let begin = service diff --git a/crates/core/src/config.rs b/crates/core/src/config.rs index 84ad790b..e204921a 100644 --- a/crates/core/src/config.rs +++ b/crates/core/src/config.rs @@ -508,6 +508,9 @@ pub struct ControllerConfig { /// Token and cost collection from agentgateway access logs. #[serde(default, skip_serializing_if = "Option::is_none")] pub usage: Option, + /// Limits on the device-facing fleet API that depend on the deployment. + #[serde(default)] + pub fleet_limits: ControllerFleetLimitsConfig, /// TLS identities used by the device-facing fleet API. /// /// A string selects a directory containing `controller.pem`, @@ -572,6 +575,60 @@ pub enum ControllerAdminTransport { }, } +/// Limits on the device-facing fleet API that depend on the deployment: how +/// many devices connect, how many share one client address, and how much they +/// send. Unset fields keep their defaults. +#[derive(Clone, Debug, Deserialize, PartialEq, Serialize)] +#[cfg_attr(feature = "schema", derive(schemars::JsonSchema))] +#[serde(rename_all = "camelCase", default, deny_unknown_fields)] +pub struct ControllerFleetLimitsConfig { + /// Open fleet connections allowed in total. At the limit, the connection + /// idle the longest is closed to admit a new one; when none is idle, the + /// new one is closed. Defaults to `50000`. + pub max_connections: usize, + /// Open fleet connections allowed from one client address: an IPv4 + /// address, or an IPv6 /64. Raise it when more devices share an address, + /// such as behind a NAT or a load balancer that does not preserve client + /// addresses. Defaults to `4096`. + pub connections_per_address: usize, + /// Largest request a device may send, in bytes. Defaults to `4194304` + /// (4 MiB). + pub max_message_bytes: usize, + /// Enrollment requests in progress at once across all clients. Enrollment + /// calls the identity provider. Defaults to `64`. + pub enrollment_in_flight: usize, + /// Enrollments one client address may start at once; it then gets one more + /// every `enrollmentRefill`. Defaults to `60`. + pub enrollment_burst: u32, + /// Time for a client address to get back one enrollment. Defaults to `1s`. + #[serde(with = "humantime_serde")] + #[cfg_attr(feature = "schema", schemars(with = "String"))] + pub enrollment_refill: Duration, + /// Messages one device may send at once on its `Connect` stream; it then + /// gets one more every `deviceMessageRefill`. A device over its budget is + /// disconnected. Defaults to `600`. + pub device_message_burst: u32, + /// Time for a device to get back one message. Defaults to `50ms`. + #[serde(with = "humantime_serde")] + #[cfg_attr(feature = "schema", schemars(with = "String"))] + pub device_message_refill: Duration, +} + +impl Default for ControllerFleetLimitsConfig { + fn default() -> Self { + Self { + max_connections: 50_000, + connections_per_address: 4_096, + max_message_bytes: 4 * 1024 * 1024, + enrollment_in_flight: 64, + enrollment_burst: 60, + enrollment_refill: Duration::from_secs(1), + device_message_burst: 600, + device_message_refill: Duration::from_millis(50), + } + } +} + /// OpenID Connect settings used for interactive device enrollment. #[derive(Clone, Debug, Deserialize, Serialize)] #[cfg_attr(feature = "schema", derive(schemars::JsonSchema))] @@ -1415,6 +1472,24 @@ pub fn parse_controller(contents: &str) -> anyhow::Result { anyhow::bail!("gatewayJwt.lifetime must be greater than zero"); } } + let limits = &config.fleet_limits; + for (field, zero) in [ + ("maxConnections", limits.max_connections == 0), + ("connectionsPerAddress", limits.connections_per_address == 0), + ("maxMessageBytes", limits.max_message_bytes == 0), + ("enrollmentInFlight", limits.enrollment_in_flight == 0), + ("enrollmentBurst", limits.enrollment_burst == 0), + ("enrollmentRefill", limits.enrollment_refill.is_zero()), + ("deviceMessageBurst", limits.device_message_burst == 0), + ( + "deviceMessageRefill", + limits.device_message_refill.is_zero(), + ), + ] { + if zero { + anyhow::bail!("fleetLimits.{field} must be greater than zero"); + } + } Ok(config) } @@ -2564,6 +2639,42 @@ oidc: .expect("explicit insecure development configuration"); } + #[test] + fn controller_fleet_limits_default_and_parse() { + let base = r#" +tls: /etc/agentdesktop/tls +oidc: + issuer: https://idp.example.com + clientId: agentdesktop +"#; + assert_eq!( + parse_controller(base).expect("defaults").fleet_limits, + super::ControllerFleetLimitsConfig::default() + ); + + let configured = parse_controller(&format!( + "{base}fleetLimits:\n connectionsPerAddress: 20000\n enrollmentRefill: 250ms\n" + )) + .expect("configured limits") + .fleet_limits; + assert_eq!(configured.connections_per_address, 20_000); + assert_eq!( + configured.enrollment_refill, + std::time::Duration::from_millis(250) + ); + assert_eq!( + configured.max_connections, 50_000, + "unset fields keep defaults" + ); + + let error = parse_controller(&format!("{base}fleetLimits:\n enrollmentRefill: 0s\n")) + .expect_err("zero refill"); + assert!( + error.to_string().contains("fleetLimits.enrollmentRefill"), + "{error}" + ); + } + #[test] fn rejects_an_invalid_llm_gateway_url() { let error = parse_daemon( diff --git a/deploy/helm/agentdesktop-controller/templates/configmap.yaml b/deploy/helm/agentdesktop-controller/templates/configmap.yaml index 1a522b9e..8e2533ab 100644 --- a/deploy/helm/agentdesktop-controller/templates/configmap.yaml +++ b/deploy/helm/agentdesktop-controller/templates/configmap.yaml @@ -18,6 +18,10 @@ data: redirectUri: {{ .Values.oidc.redirectUri | quote }} {{- with .Values.gatewayJwt }} gatewayJwt: +{{ toYaml . | indent 6 }} + {{- end }} + {{- with .Values.fleetLimits }} + fleetLimits: {{ toYaml . | indent 6 }} {{- end }} {{- with .Values.daemonConfig }} diff --git a/deploy/helm/agentdesktop-controller/values.yaml b/deploy/helm/agentdesktop-controller/values.yaml index 3d01ee8e..4c1e4674 100644 --- a/deploy/helm/agentdesktop-controller/values.yaml +++ b/deploy/helm/agentdesktop-controller/values.yaml @@ -22,6 +22,11 @@ allowInsecureDev: false # Optional short-lived JWT issuer for authenticated LLM gateways. gatewayJwt: {} +# Limits on the device-facing fleet API, such as connectionsPerAddress or +# maxMessageBytes. Unset fields keep the controller defaults; see +# schema/controller-config.md. +fleetLimits: {} + # Daemon configuration distributed to enrolled devices. Leave empty to # disable managed configuration. daemonConfig: {} diff --git a/schema/controller-config.json b/schema/controller-config.json index 774c2998..d12a8667 100644 --- a/schema/controller-config.json +++ b/schema/controller-config.json @@ -37,6 +37,20 @@ "type": "string", "default": "sqlite://agentdesktop-controller.db?mode=rwc" }, + "fleetLimits": { + "description": "Limits on the device-facing fleet API that depend on the deployment.", + "$ref": "#/$defs/ControllerFleetLimitsConfig", + "default": { + "connectionsPerAddress": 4096, + "deviceMessageBurst": 600, + "deviceMessageRefill": "50ms", + "enrollmentBurst": 60, + "enrollmentInFlight": 64, + "enrollmentRefill": "1s", + "maxConnections": 50000, + "maxMessageBytes": 4194304 + } + }, "fleetListen": { "description": "Address on which the device-facing gRPC fleet API listens.", "type": "string", @@ -191,6 +205,65 @@ "path" ] }, + "ControllerFleetLimitsConfig": { + "description": "Limits on the device-facing fleet API that depend on the deployment: how\nmany devices connect, how many share one client address, and how much they\nsend. Unset fields keep their defaults.", + "type": "object", + "properties": { + "connectionsPerAddress": { + "description": "Open fleet connections allowed from one client address: an IPv4\naddress, or an IPv6 /64. Raise it when more devices share an address,\nsuch as behind a NAT or a load balancer that does not preserve client\naddresses. Defaults to `4096`.", + "type": "integer", + "format": "uint", + "default": 4096, + "minimum": 0 + }, + "deviceMessageBurst": { + "description": "Messages one device may send at once on its `Connect` stream; it then\ngets one more every `deviceMessageRefill`. A device over its budget is\ndisconnected. Defaults to `600`.", + "type": "integer", + "format": "uint32", + "default": 600, + "minimum": 0 + }, + "deviceMessageRefill": { + "description": "Time for a device to get back one message. Defaults to `50ms`.", + "type": "string", + "default": "50ms" + }, + "enrollmentBurst": { + "description": "Enrollments one client address may start at once; it then gets one more\nevery `enrollmentRefill`. Defaults to `60`.", + "type": "integer", + "format": "uint32", + "default": 60, + "minimum": 0 + }, + "enrollmentInFlight": { + "description": "Enrollment requests in progress at once across all clients. Enrollment\ncalls the identity provider. Defaults to `64`.", + "type": "integer", + "format": "uint", + "default": 64, + "minimum": 0 + }, + "enrollmentRefill": { + "description": "Time for a client address to get back one enrollment. Defaults to `1s`.", + "type": "string", + "default": "1s" + }, + "maxConnections": { + "description": "Open fleet connections allowed in total. At the limit, the connection\nidle the longest is closed to admit a new one; when none is idle, the\nnew one is closed. Defaults to `50000`.", + "type": "integer", + "format": "uint", + "default": 50000, + "minimum": 0 + }, + "maxMessageBytes": { + "description": "Largest request a device may send, in bytes. Defaults to `4194304`\n(4 MiB).", + "type": "integer", + "format": "uint", + "default": 4194304, + "minimum": 0 + } + }, + "additionalProperties": false + }, "ControllerGatewayJwtConfig": { "description": "Settings for issuing short-lived LLM gateway JWTs.", "type": "object", diff --git a/schema/controller-config.md b/schema/controller-config.md index 9461c745..2c363606 100644 --- a/schema/controller-config.md +++ b/schema/controller-config.md @@ -18,6 +18,15 @@ |`daemonConfig.path`|string|Path to the watched YAML configuration distributed to enrolled devices.

Relative paths are resolved from the controller configuration directory.
Valid file changes are published to connected devices automatically.| |`daemonConfig.revision`|integer|Monotonically increasing revision assigned to the daemon configuration.| |`databaseUrl`|string|SQLite or PostgreSQL URL used for controller state.| +|`fleetLimits`|object|Limits on the device-facing fleet API that depend on the deployment.| +|`fleetLimits.connectionsPerAddress`|integer|Open fleet connections allowed from one client address: an IPv4
address, or an IPv6 /64. Raise it when more devices share an address,
such as behind a NAT or a load balancer that does not preserve client
addresses. Defaults to `4096`.| +|`fleetLimits.deviceMessageBurst`|integer|Messages one device may send at once on its `Connect` stream; it then
gets one more every `deviceMessageRefill`. A device over its budget is
disconnected. Defaults to `600`.| +|`fleetLimits.deviceMessageRefill`|string|Time for a device to get back one message. Defaults to `50ms`.| +|`fleetLimits.enrollmentBurst`|integer|Enrollments one client address may start at once; it then gets one more
every `enrollmentRefill`. Defaults to `60`.| +|`fleetLimits.enrollmentInFlight`|integer|Enrollment requests in progress at once across all clients. Enrollment
calls the identity provider. Defaults to `64`.| +|`fleetLimits.enrollmentRefill`|string|Time for a client address to get back one enrollment. Defaults to `1s`.| +|`fleetLimits.maxConnections`|integer|Open fleet connections allowed in total. At the limit, the connection
idle the longest is closed to admit a new one; when none is idle, the
new one is closed. Defaults to `50000`.| +|`fleetLimits.maxMessageBytes`|integer|Largest request a device may send, in bytes. Defaults to `4194304`
(4 MiB).| |`fleetListen`|string|Address on which the device-facing gRPC fleet API listens.| |`gatewayJwt`|object|LLM gateway JWT signing settings.| |`gatewayJwt.issuer`|string|Issuer claim placed in generated JWTs.| From c4e28cfd014ab018678804a6e2d9fae53cd9b20a Mon Sep 17 00:00:00 2001 From: Scott Weiss Date: Mon, 9 Nov 2026 11:42:14 -0500 Subject: [PATCH 18/49] fix(helm): leave the external traffic policy at the Kubernetes default Defaulting the Service to externalTrafficPolicy: Local drops traffic that reaches a node without a controller pod. NodePort clients that connect to another node, and load balancers that do not health-check nodes, would lose the controller after an install or a rollout. Leave the policy unset by default, and document Local (or a higher fleetLimits.connectionsPerAddress) for when the controller sees node addresses instead of devices' own. --- deploy/README.md | 12 +++++++----- deploy/helm/agentdesktop-controller/values.yaml | 14 ++++++++------ 2 files changed, 15 insertions(+), 11 deletions(-) diff --git a/deploy/README.md b/deploy/README.md index 76b24057..26143cb8 100644 --- a/deploy/README.md +++ b/deploy/README.md @@ -40,11 +40,13 @@ Secret containing the complete controller configuration under the key `controller.yaml`, then set `existingConfigSecret` to its name. See `values.yaml` for all chart settings. -The Service exposes only the TLS-protected fleet API. It defaults to -`externalTrafficPolicy: Local` so the controller sees each device's own address, -which its per-address connection and enrollment limits rely on. Set -`service.externalTrafficPolicy` to `""` to use the Kubernetes default instead; -the per-address limits then apply to the node addresses traffic arrives from. +The Service exposes only the TLS-protected fleet API. The controller limits +connections and enrollment per client address. With the Kubernetes default +external traffic policy, it may see node addresses instead of devices' own, and +many devices then share one address's limits. Set +`service.externalTrafficPolicy` to `Local` to keep devices' addresses when your +load balancer sends traffic only to nodes that run the controller, or raise +`fleetLimits.connectionsPerAddress`. To inspect the loopback-only management UI: diff --git a/deploy/helm/agentdesktop-controller/values.yaml b/deploy/helm/agentdesktop-controller/values.yaml index 4c1e4674..6fb4f8dd 100644 --- a/deploy/helm/agentdesktop-controller/values.yaml +++ b/deploy/helm/agentdesktop-controller/values.yaml @@ -43,12 +43,14 @@ tlsSecretName: agentdesktop-controller-tls service: type: LoadBalancer port: 443 - # Local keeps devices' own addresses visible to the controller, which limits - # connections and enrollment per client address. With Cluster, Kubernetes may - # replace them with node addresses, so the per-address limits act as shared - # ones. Applies only to LoadBalancer and NodePort services; set to "" to use - # the Kubernetes default. - externalTrafficPolicy: Local + # The controller limits connections and enrollment per client address. With + # the Kubernetes default (Cluster), it may see node addresses instead of + # devices' own, so many devices share one address's limits. Set Local to keep + # devices' addresses, if your load balancer only sends traffic to nodes that + # run the controller (NodePort clients must connect to such a node). + # Otherwise raise fleetLimits.connectionsPerAddress. Applies only to + # LoadBalancer and NodePort services. + externalTrafficPolicy: "" resources: {} From ed4e95b9805f12b47674058a71bd0adcc50c5de0 Mon Sep 17 00:00:00 2001 From: Scott Weiss Date: Mon, 9 Nov 2026 11:43:25 -0500 Subject: [PATCH 19/49] fix(controller): release a device enrollment's reservation when its request is cancelled BeginEnrollment with device authorization reserves a pending-enrollment slot before it calls the identity provider, and released it only when that call returned an error. The call has no timeout of its own, so when the fleet request timeout (or a client that disconnects) cancelled the request while the provider was slow, the slot stayed taken for the full ten-minute enrollment lifetime. A slow provider could fill the table that way. Release the reservation when it is dropped, unless the device got its grant. --- crates/controller/src/oidc.rs | 89 ++++++++++++++++++++++++++++------- 1 file changed, 71 insertions(+), 18 deletions(-) diff --git a/crates/controller/src/oidc.rs b/crates/controller/src/oidc.rs index f6c8672d..eeeac882 100644 --- a/crates/controller/src/oidc.rs +++ b/crates/controller/src/oidc.rs @@ -32,6 +32,39 @@ pub struct OidcProvider { inner: Arc, } +/// A pending enrollment that [`OidcProvider::begin_device`] reserved before +/// contacting the IdP. Dropping it releases the reservation: when the IdP call +/// or its response fails, and when the request is cancelled while it waits (by +/// the fleet request timeout or a client that disconnects). Already-removed +/// (expired and pruned, or raced with a concurrent completion) is not an error. +struct Reservation { + provider: Option>, + enrollment_id: String, +} + +impl Reservation { + /// Keeps the reservation, now that the device has its grant. + fn keep(mut self) { + self.provider = None; + } +} + +impl Drop for Reservation { + fn drop(&mut self) { + let Some(provider) = self.provider.take() else { + return; + }; + let enrollment_id = std::mem::take(&mut self.enrollment_id); + if let Ok(mut pending) = provider.pending.try_lock() { + pending.remove(&enrollment_id); + } else { + tokio::spawn(async move { + provider.pending.lock().await.remove(&enrollment_id); + }); + } + } +} + struct Inner { issuer: String, client_id: String, @@ -226,21 +259,20 @@ impl OidcProvider { // afterwards, on the insert. The reservation uses ENROLLMENT_LIFETIME as a // placeholder expiry — the real, possibly shorter, device lifetime isn't known // until the IdP responds — and is narrowed below once it is, or dropped - // entirely if the IdP call or the response fails. + // entirely if the IdP call or the response fails, or the request is + // cancelled while it waits. let enrollment_id = random_secret(); self.insert_pending(enrollment_id.clone(), hostname, None, ENROLLMENT_LIFETIME) .await?; + let reservation = Reservation { + provider: Some(self.inner.clone()), + enrollment_id: enrollment_id.clone(), + }; - match self.request_device_grant(&enrollment_id, endpoint).await { - Ok((response, lifetime)) => { - self.set_pending_expiry(&enrollment_id, lifetime).await; - Ok(response) - } - Err(error) => { - self.release_pending(&enrollment_id).await; - Err(error) - } - } + let (response, lifetime) = self.request_device_grant(&enrollment_id, endpoint).await?; + self.set_pending_expiry(&enrollment_id, lifetime).await; + reservation.keep(); + Ok(response) } /// The IdP round-trip and response validation for [`begin_device`], split out so @@ -300,13 +332,6 @@ impl OidcProvider { Ok(()) } - /// Drops a reservation [`insert_pending`] made, e.g. because the IdP request or - /// response validation that was supposed to fill it in failed. Already-removed - /// (expired and pruned, or raced with a concurrent completion) is not an error. - async fn release_pending(&self, enrollment_id: &str) { - self.inner.pending.lock().await.remove(enrollment_id); - } - /// Narrows a reservation's expiry to the real device-grant lifetime once the IdP /// has responded, which can be shorter than the ENROLLMENT_LIFETIME placeholder /// [`begin_device`] reserved it with. A no-op if the entry is already gone. @@ -837,6 +862,34 @@ mod tests { server.abort(); } + #[tokio::test] + async fn begin_device_releases_the_reservation_when_the_request_is_cancelled() { + // An IdP that accepts the connection and never answers. + let listener = tokio::net::TcpListener::bind("127.0.0.1:0") + .await + .expect("bind silent IdP listener"); + let endpoint = format!( + "http://{}/device", + listener.local_addr().expect("silent IdP local addr") + ) + .parse() + .expect("parse silent IdP URL"); + let provider = provider(Some(endpoint)); + + let cancelled = tokio::time::timeout( + Duration::from_millis(200), + provider.begin_device("new-host".to_owned()), + ) + .await; + + assert!(cancelled.is_err(), "the silent IdP answered"); + assert!( + provider.inner.pending.lock().await.is_empty(), + "the cancelled request's reservation must be released" + ); + drop(listener); + } + #[tokio::test] async fn unknown_key_refreshes_and_replaces_jwks() { let cache = cache(jwks(&["old"])); From 97a37d99af4769aeb37cb1e996b6e8b6fbc1a172 Mon Sep 17 00:00:00 2001 From: Scott Weiss Date: Mon, 9 Nov 2026 11:43:25 -0500 Subject: [PATCH 20/49] docs(controller): describe tcp_keepalive as the idle time before probes --- crates/controller/src/fleet_limits.rs | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/crates/controller/src/fleet_limits.rs b/crates/controller/src/fleet_limits.rs index 83f97867..9f6d973b 100644 --- a/crates/controller/src/fleet_limits.rs +++ b/crates/controller/src/fleet_limits.rs @@ -53,7 +53,8 @@ pub struct FleetLimits { /// PING waits before the connection is closed. pub keepalive_interval: Duration, pub keepalive_timeout: Duration, - /// TCP keepalive probe interval, for peers that vanish without closing. + /// How long a connection may be silent before TCP keepalive probes start, + /// for peers that vanish without closing. pub tcp_keepalive: Duration, /// Open connections allowed in total and from one client address (an IPv4 /// address or an IPv6 /64). A connection over either cap is closed as soon From b59e7462ae2d3bac3f766c40264733c19fc12a1c Mon Sep 17 00:00:00 2001 From: Scott Weiss Date: Mon, 9 Nov 2026 14:42:41 -0500 Subject: [PATCH 21/49] fix(controller): keep a connection busy only while Connect is open A response held the connection busy until hyper dropped its body. An HTTP/2 client that grants no flow-control window never lets the server send a body, so any call that returns one (the gateway JWKS route needs no credentials) kept the connection busy forever: never idle-closed, never evicted, while the client answers keepalive PINGs. A few addresses could fill the connection cap that way. Only Connect, whose response is the device's session, still keeps the connection busy until hyper drops its response body, normally when the stream ends. Any other request counts as in progress while its handler runs, so a stalled response body leaves the connection idle and subject to the idle close and eviction. --- crates/controller/src/fleet_limits.rs | 8 ++++ .../src/fleet_limits/connections.rs | 28 ++++++++---- crates/controller/src/fleet_limits/tests.rs | 44 +++++++++++++++++++ 3 files changed, 72 insertions(+), 8 deletions(-) diff --git a/crates/controller/src/fleet_limits.rs b/crates/controller/src/fleet_limits.rs index 9f6d973b..70e61f22 100644 --- a/crates/controller/src/fleet_limits.rs +++ b/crates/controller/src/fleet_limits.rs @@ -232,6 +232,14 @@ fn is_enrollment(path: &str) -> bool { .is_some_and(|method| matches!(method, "/BeginEnrollment" | "/CompleteEnrollment")) } +/// Whether `path` is the `Connect` stream, whose response lasts as long as the +/// device's session. +fn is_connect(path: &str) -> bool { + path.strip_prefix('/') + .and_then(|path| path.strip_prefix(SERVICE_NAME)) + .is_some_and(|method| method == "/Connect") +} + /// Whether the TLS client presented a certificate signed by the device CA /// (the handshake rejects any other). fn has_device_certificate(request: &http::Request) -> bool { diff --git a/crates/controller/src/fleet_limits/connections.rs b/crates/controller/src/fleet_limits/connections.rs index 73ff10d0..dc39a68f 100644 --- a/crates/controller/src/fleet_limits/connections.rs +++ b/crates/controller/src/fleet_limits/connections.rs @@ -12,6 +12,11 @@ //! an open `Connect` stream, is never idle. When the total cap is reached, the //! longest-idle connection is closed to make room instead of refusing the new //! one, so idle connections cannot crowd out devices that are using theirs. +//! +//! A request other than `Connect` is in progress only while its handler runs, +//! not while its response is sent. Otherwise a client could keep a connection +//! busy forever by granting no HTTP/2 flow-control window, so that the server +//! can never send it the response body. use std::{ collections::{HashMap, VecDeque}, @@ -42,7 +47,7 @@ use tonic::{ }; use tower_layer::Layer; -use super::{FleetLimits, ThrottledWarning, address_key}; +use super::{FleetLimits, ThrottledWarning, address_key, is_connect}; static TOTAL_CAP: ThrottledWarning = ThrottledWarning::new(); static PER_ADDRESS_CAP: ThrottledWarning = ThrottledWarning::new(); @@ -444,8 +449,8 @@ impl AsyncWrite for LimitedConnection { } } -/// Marks the connection busy until the request's response has been fully sent -/// (for `Connect`, until the stream ends). +/// Marks the connection busy while a request's handler runs, and for `Connect` +/// until its stream ends. struct RequestGuard { connections: Connections, peer: SocketAddr, @@ -523,18 +528,25 @@ where fn call(&mut self, request: http::Request) -> Self::Future { let guard = remote_addr(&request).and_then(|peer| self.connections.start_request(peer)); + let session = is_connect(request.uri().path()); let response = self.inner.call(request); Box::pin(async move { let response = response.await?; - Ok(match guard { - Some(guard) => response.map(|body| { + match guard { + // The stream is the device's session: keep the connection busy + // until it ends. + Some(guard) if session => Ok(response.map(|body| { Body::new(TrackedBody { inner: body, _guard: guard, }) - }), - None => response, - }) + })), + // Any other response is done once its handler returns. + guard => { + drop(guard); + Ok(response) + } + } }) } } diff --git a/crates/controller/src/fleet_limits/tests.rs b/crates/controller/src/fleet_limits/tests.rs index abbbb779..ee4fec28 100644 --- a/crates/controller/src/fleet_limits/tests.rs +++ b/crates/controller/src/fleet_limits/tests.rs @@ -594,6 +594,50 @@ async fn a_connection_idle_after_its_last_request_is_closed() { .await; } +#[tokio::test] +async fn a_response_the_client_never_accepts_does_not_keep_its_connection_busy() { + let mut limits = limits(); + limits.first_request = Duration::from_secs(5); + limits.idle_connection = Duration::from_millis(200); + let (address, connections) = serve_tracked( + limits, + None, + Slow { + delay: Duration::ZERO, + }, + ) + .await; + // A client that grants no HTTP/2 flow-control window, so the server can + // never send it a response body. It still answers keepalive PINGs. + let channel = tonic::transport::Endpoint::from_shared(format!("http://{address}")) + .unwrap() + .initial_stream_window_size(0) + .connect() + .await + .unwrap(); + let mut stalled = FleetAgentClient::new(channel); + let call = tokio::spawn(async move { + stalled + .begin_enrollment(BeginEnrollmentRequest::default()) + .await + }); + bounded(async { + while connections.open_count() == 0 { + time::sleep(Duration::from_millis(10)).await; + } + }) + .await; + + bounded(async { + while connections.open_count() > 0 { + time::sleep(Duration::from_millis(50)).await; + } + }) + .await; + let response = bounded(call).await.unwrap(); + assert!(response.is_err(), "the stalled response was delivered"); +} + #[tokio::test] async fn an_open_stream_keeps_its_connection_from_going_idle() { let mut limits = limits(); From fd5bd377d69e70ef893372910d284d4cb56b4a52 Mon Sep 17 00:00:00 2001 From: Scott Weiss Date: Mon, 9 Nov 2026 14:22:54 -0500 Subject: [PATCH 22/49] fix(controller): check the per-address cap before evicting at the total cap At the total connection cap, a new connection evicted the longest-idle connection before the per-address cap was checked. An address already at its cap could keep opening TCP connections (admission runs before TLS), each one closing some other client's idle connection and then being refused itself. The per-address cap is now checked first. Registration is split out of `admit` so the ordering can be tested without sockets. --- .../src/fleet_limits/connections.rs | 61 ++++++++++++------- crates/controller/src/fleet_limits/tests.rs | 23 +++++++ 2 files changed, 61 insertions(+), 23 deletions(-) diff --git a/crates/controller/src/fleet_limits/connections.rs b/crates/controller/src/fleet_limits/connections.rs index dc39a68f..cf3cf4b5 100644 --- a/crates/controller/src/fleet_limits/connections.rs +++ b/crates/controller/src/fleet_limits/connections.rs @@ -80,7 +80,7 @@ struct State { } /// One connection's request activity. -struct Activity { +pub(super) struct Activity { in_flight: AtomicUsize, /// Advanced each time the connection becomes idle. idle_epoch: AtomicU64, @@ -202,51 +202,66 @@ impl Connections { fn admit(&self, stream: TcpStream) -> Option { let peer = stream.peer_addr().ok()?; + let (activity, evicted) = self.register(peer)?; + Some(LimitedConnection { + stream, + peer, + activity, + connections: self.clone(), + evicted: Some(evicted), + timer: Box::pin(tokio::time::sleep_until( + Instant::now() + self.shared.first_request, + )), + closed: false, + }) + } + + /// Registers a connection from `peer` as idle, or returns `None` when the + /// caps refuse it. The receiver resolves if the connection is evicted. + pub(super) fn register( + &self, + peer: SocketAddr, + ) -> Option<(Arc, oneshot::Receiver<()>)> { let mut state = self.state(); - if state.open.len() >= self.shared.max_connections && !state.evict_idle() { - if let Some(events) = TOTAL_CAP.event(Instant::now()) { + // The per-address cap comes first, so an address over it cannot close + // another client's idle connection only to be refused itself. + let address = address_key(peer.ip()); + if state + .per_address + .get(&address) + .is_some_and(|count| *count >= self.shared.connections_per_address) + { + if let Some(events) = PER_ADDRESS_CAP.event(Instant::now()) { tracing::warn!( + %address, events, - "closed new fleet connections at the connection limit with no connection idle" + "closed new fleet connections over the per-address connection limit" ); } return None; } - let address = address_key(peer.ip()); - let count = state.per_address.entry(address).or_default(); - if *count >= self.shared.connections_per_address { - if let Some(events) = PER_ADDRESS_CAP.event(Instant::now()) { + if state.open.len() >= self.shared.max_connections && !state.evict_idle() { + if let Some(events) = TOTAL_CAP.event(Instant::now()) { tracing::warn!( - %address, events, - "closed new fleet connections over the per-address connection limit" + "closed new fleet connections at the connection limit with no connection idle" ); } return None; } - *count += 1; - let now = Instant::now(); + *state.per_address.entry(address).or_default() += 1; let (evict, evicted) = oneshot::channel(); let activity = Arc::new(Activity { in_flight: AtomicUsize::new(0), idle_epoch: AtomicU64::new(0), used: AtomicBool::new(false), - last_active: Mutex::new(now), + last_active: Mutex::new(Instant::now()), evict: Mutex::new(Some(evict)), waker: Mutex::new(None), }); state.open.insert(peer, activity.clone()); state.mark_idle(peer, 0); - drop(state); - Some(LimitedConnection { - stream, - peer, - activity, - connections: self.clone(), - evicted: Some(evicted), - timer: Box::pin(tokio::time::sleep_until(now + self.shared.first_request)), - closed: false, - }) + Some((activity, evicted)) } fn release(&self, peer: SocketAddr, activity: &Arc) { diff --git a/crates/controller/src/fleet_limits/tests.rs b/crates/controller/src/fleet_limits/tests.rs index ee4fec28..377a9308 100644 --- a/crates/controller/src/fleet_limits/tests.rs +++ b/crates/controller/src/fleet_limits/tests.rs @@ -473,6 +473,29 @@ async fn connections_over_the_per_address_cap_are_closed() { connections_over_cap(limits).await; } +#[test] +fn an_address_over_its_cap_does_not_evict_other_connections() { + use tokio::sync::oneshot::error::TryRecvError; + + let mut limits = limits(); + limits.max_connections = 2; + limits.connections_per_address = 1; + let connections = limits.connections(); + let first_address = |port| SocketAddr::from(([192, 0, 2, 1], port)); + let (_first, mut first_evicted) = connections + .register(first_address(1)) + .expect("first connection"); + let (_other, mut other_evicted) = connections + .register(SocketAddr::from(([192, 0, 2, 2], 1))) + .expect("a connection from another address"); + + // At the total cap, and over the first address's cap. + assert!(connections.register(first_address(2)).is_none()); + + assert_eq!(first_evicted.try_recv(), Err(TryRecvError::Empty)); + assert_eq!(other_evicted.try_recv(), Err(TryRecvError::Empty)); +} + #[tokio::test] async fn at_the_total_cap_the_longest_idle_connection_makes_room() { let mut limits = limits(); From 4c4ab7db051e8ec5d529fc17d2379bf8b4ab01f9 Mon Sep 17 00:00:00 2001 From: Scott Weiss Date: Mon, 9 Nov 2026 14:23:30 -0500 Subject: [PATCH 23/49] fix(controller): count a connection's requests under the registry lock `start_request` looked the connection up under the registry lock but incremented its in-flight count after releasing it, so `evict_idle` could see the connection idle and evict it just as a request started, ending that request. On the way out, the count dropped to zero before the connection was queued as idle with its new epoch, leaving a window in which an older queue entry still looked current. Both updates now happen under the lock `evict_idle` holds. The window is a few instructions wide, so there is no deterministic test for it. --- .../src/fleet_limits/connections.rs | 32 +++++++++++++------ 1 file changed, 23 insertions(+), 9 deletions(-) diff --git a/crates/controller/src/fleet_limits/connections.rs b/crates/controller/src/fleet_limits/connections.rs index cf3cf4b5..9bc3be79 100644 --- a/crates/controller/src/fleet_limits/connections.rs +++ b/crates/controller/src/fleet_limits/connections.rs @@ -283,8 +283,14 @@ impl Connections { } fn start_request(&self, peer: SocketAddr) -> Option { - let activity = self.state().open.get(&peer)?.clone(); - activity.in_flight.fetch_add(1, Ordering::AcqRel); + // Counted under the registry lock, which `evict_idle` holds, so a + // connection is never evicted once its request has started. + let activity = { + let state = self.state(); + let activity = state.open.get(&peer)?.clone(); + activity.in_flight.fetch_add(1, Ordering::AcqRel); + activity + }; activity.used.store(true, Ordering::Release); activity.touch(); Some(RequestGuard { @@ -475,18 +481,26 @@ struct RequestGuard { impl Drop for RequestGuard { fn drop(&mut self) { self.activity.touch(); - if self.activity.in_flight.fetch_sub(1, Ordering::AcqRel) == 1 { - let epoch = self.activity.idle_epoch.fetch_add(1, Ordering::AcqRel) + 1; - self.connections.state().mark_idle(self.peer, epoch); - if let Some(waker) = self + // Under the registry lock, like `start_request`, so `evict_idle` never + // sees the connection idle before it is queued with its new epoch. + let idle = { + let mut state = self.connections.state(); + let idle = self.activity.in_flight.fetch_sub(1, Ordering::AcqRel) == 1; + if idle { + let epoch = self.activity.idle_epoch.fetch_add(1, Ordering::AcqRel) + 1; + state.mark_idle(self.peer, epoch); + } + idle + }; + if idle + && let Some(waker) = self .activity .waker .lock() .expect("connection waker lock poisoned") .take() - { - waker.wake(); - } + { + waker.wake(); } } } From 7210a6bb67afc11bd51adb1b70b4e71b70440fa4 Mon Sep 17 00:00:00 2001 From: Scott Weiss Date: Mon, 9 Nov 2026 14:24:26 -0500 Subject: [PATCH 24/49] fix(controller): time out identity provider requests The OIDC HTTP client had no timeout, so an IdP that accepted a connection and never answered stalled whatever was waiting on it. A Connect stream waits on UserInfo at its periodic access-token check, and since that stream's task now also forwards config pushes and reacts to a newer stream from the same device, a hung IdP stalled all of those too. An enrollment waited until the fleet request timeout cancelled it, which can consume the pending enrollment and make the user sign in again. Every IdP request now gives up after 10 seconds (the JWKS fetch keeps its own limit of the same length). A timeout has no HTTP status, so the stream check treats it as the IdP being unavailable and keeps the stream. --- crates/controller/src/oidc.rs | 53 ++++++++++++++++++++++++++++++++++- 1 file changed, 52 insertions(+), 1 deletion(-) diff --git a/crates/controller/src/oidc.rs b/crates/controller/src/oidc.rs index eeeac882..55d403b7 100644 --- a/crates/controller/src/oidc.rs +++ b/crates/controller/src/oidc.rs @@ -22,6 +22,11 @@ use url::Url; const ENROLLMENT_LIFETIME: Duration = Duration::from_secs(10 * 60); const MAX_PENDING_ENROLLMENTS: usize = 1024; const JWKS_FETCH_TIMEOUT: Duration = Duration::from_secs(10); +/// Longest the controller waits for an identity provider request. Without a +/// limit, an IdP that accepts the connection and never answers stalls whatever +/// waits on it: an enrollment until the fleet request timeout cancels it, or a +/// device's `Connect` stream at its periodic access-token check. +const IDP_REQUEST_TIMEOUT: Duration = Duration::from_secs(10); const JWKS_REFRESH_COOLDOWN: Duration = Duration::from_secs(30); const ENROLLMENT_SCOPE: &str = "openid profile email offline_access"; // RFC 8628 section 3.2: clients poll every five seconds unless told otherwise. @@ -153,7 +158,7 @@ impl OidcProvider { redirect_uri: String, ) -> anyhow::Result { let _ = jsonwebtoken::crypto::aws_lc::DEFAULT_PROVIDER.install_default(); - let http = reqwest::Client::new(); + let http = idp_client(IDP_REQUEST_TIMEOUT)?; let discovery_url = format!( "{}/.well-known/openid-configuration", issuer.trim_end_matches('/') @@ -492,6 +497,15 @@ impl JwksCache { } } +/// The HTTP client for identity provider requests, which gives up on any one +/// request after `timeout`. +fn idp_client(timeout: Duration) -> anyhow::Result { + reqwest::Client::builder() + .timeout(timeout) + .build() + .context("build identity provider HTTP client") +} + async fn refresh_jwks(http: &reqwest::Client, uri: &Url) -> anyhow::Result { // Bound how long an unknown-key request can hold the single-flight refresh gate. fetch_jwks(refresh_jwks_request(http, uri)).await @@ -890,6 +904,43 @@ mod tests { drop(listener); } + #[tokio::test] + async fn an_identity_provider_that_never_answers_times_out() { + // An IdP that accepts the connection and never answers. + let listener = tokio::net::TcpListener::bind("127.0.0.1:0") + .await + .expect("bind silent IdP listener"); + let mut inner = Arc::into_inner(provider(None).inner).expect("sole owner"); + inner.userinfo_endpoint = format!( + "http://{}/userinfo", + listener.local_addr().expect("silent IdP local addr") + ) + .parse() + .expect("parse silent IdP URL"); + inner.http = super::idp_client(Duration::from_millis(200)).expect("build IdP client"); + let provider = OidcProvider { + inner: Arc::new(inner), + }; + + let error = tokio::time::timeout( + Duration::from_secs(5), + provider.authenticate_access_token("token"), + ) + .await + .expect("the UserInfo request was not bounded") + .expect_err("the silent IdP answered"); + + // A request error without an HTTP status, which `Connect` treats as the + // IdP being unavailable and keeps the stream. + let request_error = error + .chain() + .find_map(|cause| cause.downcast_ref::()) + .expect("a request error"); + assert!(request_error.is_timeout(), "{error:#}"); + assert!(request_error.status().is_none(), "{error:#}"); + drop(listener); + } + #[tokio::test] async fn unknown_key_refreshes_and_replaces_jwks() { let cache = cache(jwks(&["old"])); From c8955fb4c269b839d55a121f76b2ee1e6c6f85b7 Mon Sep 17 00:00:00 2001 From: Scott Weiss Date: Mon, 9 Nov 2026 18:37:11 -0500 Subject: [PATCH 25/49] test(controller): cover fleet limits on the production connection path The fleet limit tests served HTTP/2 only, but the controller also accepts HTTP/1, so its connections first go through hyper-util's protocol detection: the read that held connections open before the HTTP/2 preface. The test server now accepts HTTP/1 the same way. New tests: - a TLS connection that finishes the handshake and sends nothing is closed (the original hold, end to end with a raw rustls client); - a TLS connection is closed once idle after its last request; - enrollment gets the small message limit even from a client with a device certificate, so the path rule is tested apart from the certificate rule. --- crates/controller/src/fleet_limits/tests.rs | 112 +++++++++++++++++++- 1 file changed, 107 insertions(+), 5 deletions(-) diff --git a/crates/controller/src/fleet_limits/tests.rs b/crates/controller/src/fleet_limits/tests.rs index 377a9308..9d229c98 100644 --- a/crates/controller/src/fleet_limits/tests.rs +++ b/crates/controller/src/fleet_limits/tests.rs @@ -122,7 +122,9 @@ async fn serve_tracked( let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); let address = listener.local_addr().unwrap(); let connections = limits.connections(); - let mut server = limits.server(&connections); + // HTTP/1 is accepted as the controller does (`main.rs`), so connections go + // through hyper-util's protocol detection, as they do in production. + let mut server = limits.server(&connections).accept_http1(true); if let Some(tls) = tls { server = server.tls_config(limits.tls(tls)).unwrap(); } @@ -351,7 +353,7 @@ fn pki() -> Pki { } /// Serves with the production TLS shape: device CA, client certificate optional. -async fn serve_mtls(limits: FleetLimits, pki: &Pki) -> SocketAddr { +async fn serve_mtls(limits: FleetLimits, pki: &Pki) -> (SocketAddr, Connections) { let tls = ServerTlsConfig::new() .identity(pki.server.clone()) .client_ca_root(pki.ca.clone()) @@ -364,7 +366,6 @@ async fn serve_mtls(limits: FleetLimits, pki: &Pki) -> SocketAddr { }, ) .await - .0 } async fn tls_client( @@ -397,7 +398,7 @@ fn renewal(bytes: usize) -> RenewDeviceCertificateRequest { #[tokio::test] async fn device_calls_keep_the_larger_message_limit() { let pki = pki(); - let address = serve_mtls(small_messages(), &pki).await; + let (address, _) = serve_mtls(small_messages(), &pki).await; let mut client = tls_client(address, &pki, true).await; // Over the unauthenticated limit but within the device limit: reaches the handler. @@ -415,7 +416,7 @@ async fn device_calls_keep_the_larger_message_limit() { #[tokio::test] async fn callers_without_a_device_certificate_get_the_small_limit_on_every_call() { let pki = pki(); - let address = serve_mtls(small_messages(), &pki).await; + let (address, _) = serve_mtls(small_messages(), &pki).await; let mut client = tls_client(address, &pki, false).await; let error = bounded(client.renew_device_certificate(renewal(512))) @@ -425,6 +426,107 @@ async fn callers_without_a_device_certificate_get_the_small_limit_on_every_call( assert_eq!(error.code(), Code::OutOfRange, "{error}"); } +#[tokio::test] +async fn enrollment_gets_the_small_limit_even_with_a_device_certificate() { + let pki = pki(); + let (address, _) = serve_mtls(small_messages(), &pki).await; + let mut client = tls_client(address, &pki, true).await; + + // Over the unauthenticated limit but within the device limit. + let error = bounded(client.begin_enrollment(BeginEnrollmentRequest { + hostname: "h".repeat(512), + ..Default::default() + })) + .await + .expect_err("a large enrollment request reached the handler"); + + assert_eq!(error.code(), Code::OutOfRange, "{error}"); +} + +#[tokio::test] +async fn a_tls_connection_that_never_makes_a_request_is_closed() { + use std::io::Read; + + use rustls::pki_types::{CertificateDer, pem::PemObject}; + + let pki = pki(); + let mut limits = limits(); + limits.first_request = Duration::from_millis(300); + let (address, _) = serve_mtls(limits, &pki).await; + let ca = pki.ca.clone().into_inner(); + + // Finish the TLS handshake, then send nothing: not even the HTTP/2 + // preface, so the server is still detecting the protocol. + let closed = tokio::task::spawn_blocking(move || { + let mut roots = rustls::RootCertStore::empty(); + roots + .add(CertificateDer::from_pem_slice(&ca).expect("parse device CA")) + .expect("trust device CA"); + let config = rustls::ClientConfig::builder_with_provider(std::sync::Arc::new( + rustls::crypto::aws_lc_rs::default_provider(), + )) + .with_safe_default_protocol_versions() + .expect("TLS versions") + .with_root_certificates(roots) + .with_no_client_auth(); + let mut tls = rustls::ClientConnection::new( + std::sync::Arc::new(config), + "localhost".try_into().expect("server name"), + ) + .expect("TLS client"); + let mut socket = std::net::TcpStream::connect(address).expect("connect"); + socket + .set_read_timeout(Some(Duration::from_secs(5))) + .expect("read timeout"); + while tls.is_handshaking() { + tls.complete_io(&mut socket).expect("TLS handshake"); + } + + let mut stream = rustls::Stream::new(&mut tls, &mut socket); + let mut buffer = [0_u8; 1024]; + loop { + match stream.read(&mut buffer) { + Ok(0) => return true, + Ok(_) => {} + Err(error) + if matches!( + error.kind(), + std::io::ErrorKind::WouldBlock | std::io::ErrorKind::TimedOut + ) => + { + return false; + } + Err(_) => return true, + } + } + }) + .await + .unwrap(); + + assert!(closed, "the server kept a silent TLS connection open"); +} + +#[tokio::test] +async fn a_tls_connection_idle_after_its_last_request_is_closed() { + let pki = pki(); + let mut limits = limits(); + limits.first_request = Duration::from_secs(5); + limits.idle_connection = Duration::from_millis(200); + let (address, connections) = serve_mtls(limits, &pki).await; + let mut client = tls_client(address, &pki, true).await; + bounded(client.begin_enrollment(BeginEnrollmentRequest::default())) + .await + .unwrap(); + assert_eq!(connections.open_count(), 1); + + bounded(async { + while connections.open_count() > 0 { + time::sleep(Duration::from_millis(50)).await; + } + }) + .await; +} + /// Whether the server keeps `stream` open for a while. An accepted connection /// receives the server's HTTP/2 SETTINGS and then waits for the client; a /// rejected one is closed (end of stream) right away. From a84bb31de0f1b166bafa0169d597f0c6d9a2d7a8 Mon Sep 17 00:00:00 2001 From: Scott Weiss Date: Mon, 9 Nov 2026 14:27:09 -0500 Subject: [PATCH 26/49] docs(deploy): note that devices behind one address share an enrollment budget With the default external traffic policy, devices behind one node address share that address's enrollment budget as well as its connection limit, so a large rollout can be held to about one enrollment per second per node. Say so, say that Local does not help behind a load balancer that proxies connections, and name the enrollment settings to raise. --- deploy/README.md | 10 ++++++---- deploy/helm/agentdesktop-controller/values.yaml | 13 ++++++++----- 2 files changed, 14 insertions(+), 9 deletions(-) diff --git a/deploy/README.md b/deploy/README.md index 26143cb8..70e8aa21 100644 --- a/deploy/README.md +++ b/deploy/README.md @@ -43,10 +43,12 @@ See `values.yaml` for all chart settings. The Service exposes only the TLS-protected fleet API. The controller limits connections and enrollment per client address. With the Kubernetes default external traffic policy, it may see node addresses instead of devices' own, and -many devices then share one address's limits. Set -`service.externalTrafficPolicy` to `Local` to keep devices' addresses when your -load balancer sends traffic only to nodes that run the controller, or raise -`fleetLimits.connectionsPerAddress`. +many devices then share one address's connection limit and enrollment budget. +Set `service.externalTrafficPolicy` to `Local` to keep devices' addresses when +your load balancer sends traffic only to nodes that run the controller. With the +default policy, or behind a load balancer that proxies connections, raise +`fleetLimits.connectionsPerAddress`, and for large rollouts raise +`fleetLimits.enrollmentBurst` or shorten `fleetLimits.enrollmentRefill`. To inspect the loopback-only management UI: diff --git a/deploy/helm/agentdesktop-controller/values.yaml b/deploy/helm/agentdesktop-controller/values.yaml index 6fb4f8dd..a03cf56b 100644 --- a/deploy/helm/agentdesktop-controller/values.yaml +++ b/deploy/helm/agentdesktop-controller/values.yaml @@ -45,11 +45,14 @@ service: port: 443 # The controller limits connections and enrollment per client address. With # the Kubernetes default (Cluster), it may see node addresses instead of - # devices' own, so many devices share one address's limits. Set Local to keep - # devices' addresses, if your load balancer only sends traffic to nodes that - # run the controller (NodePort clients must connect to such a node). - # Otherwise raise fleetLimits.connectionsPerAddress. Applies only to - # LoadBalancer and NodePort services. + # devices' own, so many devices share one address's connection limit and + # enrollment budget. Set Local to keep devices' addresses, if your load + # balancer only sends traffic to nodes that run the controller (NodePort + # clients must connect to such a node). Local does not help behind a load + # balancer that proxies connections. Otherwise raise + # fleetLimits.connectionsPerAddress, and for large rollouts raise + # fleetLimits.enrollmentBurst or shorten fleetLimits.enrollmentRefill. + # Applies only to LoadBalancer and NodePort services. externalTrafficPolicy: "" resources: {} From 4b8a4f23897eac6450466ce0a486ab7047e4ff1d Mon Sep 17 00:00:00 2001 From: Scott Weiss Date: Mon, 9 Nov 2026 14:27:09 -0500 Subject: [PATCH 27/49] docs(config): describe maxMessageBytes as a two-way limit; schema rejects zero maxMessageBytes also caps the controller's responses, and devices accept at most 4 MiB however high it is set, so the description now says what raising it does. The schema gave the integer fleetLimits fields a minimum of 0 while the parser rejects 0; it now says 1. --- crates/core/src/config.rs | 11 ++++++++++- schema/controller-config.json | 14 +++++++------- schema/controller-config.md | 2 +- 3 files changed, 18 insertions(+), 9 deletions(-) diff --git a/crates/core/src/config.rs b/crates/core/src/config.rs index e204921a..0daf5dcf 100644 --- a/crates/core/src/config.rs +++ b/crates/core/src/config.rs @@ -585,20 +585,28 @@ pub struct ControllerFleetLimitsConfig { /// Open fleet connections allowed in total. At the limit, the connection /// idle the longest is closed to admit a new one; when none is idle, the /// new one is closed. Defaults to `50000`. + #[cfg_attr(feature = "schema", schemars(range(min = 1)))] pub max_connections: usize, /// Open fleet connections allowed from one client address: an IPv4 /// address, or an IPv6 /64. Raise it when more devices share an address, /// such as behind a NAT or a load balancer that does not preserve client /// addresses. Defaults to `4096`. + #[cfg_attr(feature = "schema", schemars(range(min = 1)))] pub connections_per_address: usize, - /// Largest request a device may send, in bytes. Defaults to `4194304` + /// Largest message, in bytes, that a device and the controller may send + /// each other. Requests from clients without a device certificate, such as + /// enrollment, stay limited to 64 KiB. Devices accept at most 4 MiB, so + /// raising this only raises what devices may send. Defaults to `4194304` /// (4 MiB). + #[cfg_attr(feature = "schema", schemars(range(min = 1)))] pub max_message_bytes: usize, /// Enrollment requests in progress at once across all clients. Enrollment /// calls the identity provider. Defaults to `64`. + #[cfg_attr(feature = "schema", schemars(range(min = 1)))] pub enrollment_in_flight: usize, /// Enrollments one client address may start at once; it then gets one more /// every `enrollmentRefill`. Defaults to `60`. + #[cfg_attr(feature = "schema", schemars(range(min = 1)))] pub enrollment_burst: u32, /// Time for a client address to get back one enrollment. Defaults to `1s`. #[serde(with = "humantime_serde")] @@ -607,6 +615,7 @@ pub struct ControllerFleetLimitsConfig { /// Messages one device may send at once on its `Connect` stream; it then /// gets one more every `deviceMessageRefill`. A device over its budget is /// disconnected. Defaults to `600`. + #[cfg_attr(feature = "schema", schemars(range(min = 1)))] pub device_message_burst: u32, /// Time for a device to get back one message. Defaults to `50ms`. #[serde(with = "humantime_serde")] diff --git a/schema/controller-config.json b/schema/controller-config.json index d12a8667..b1c5a15f 100644 --- a/schema/controller-config.json +++ b/schema/controller-config.json @@ -214,14 +214,14 @@ "type": "integer", "format": "uint", "default": 4096, - "minimum": 0 + "minimum": 1 }, "deviceMessageBurst": { "description": "Messages one device may send at once on its `Connect` stream; it then\ngets one more every `deviceMessageRefill`. A device over its budget is\ndisconnected. Defaults to `600`.", "type": "integer", "format": "uint32", "default": 600, - "minimum": 0 + "minimum": 1 }, "deviceMessageRefill": { "description": "Time for a device to get back one message. Defaults to `50ms`.", @@ -233,14 +233,14 @@ "type": "integer", "format": "uint32", "default": 60, - "minimum": 0 + "minimum": 1 }, "enrollmentInFlight": { "description": "Enrollment requests in progress at once across all clients. Enrollment\ncalls the identity provider. Defaults to `64`.", "type": "integer", "format": "uint", "default": 64, - "minimum": 0 + "minimum": 1 }, "enrollmentRefill": { "description": "Time for a client address to get back one enrollment. Defaults to `1s`.", @@ -252,14 +252,14 @@ "type": "integer", "format": "uint", "default": 50000, - "minimum": 0 + "minimum": 1 }, "maxMessageBytes": { - "description": "Largest request a device may send, in bytes. Defaults to `4194304`\n(4 MiB).", + "description": "Largest message, in bytes, that a device and the controller may send\neach other. Requests from clients without a device certificate, such as\nenrollment, stay limited to 64 KiB. Devices accept at most 4 MiB, so\nraising this only raises what devices may send. Defaults to `4194304`\n(4 MiB).", "type": "integer", "format": "uint", "default": 4194304, - "minimum": 0 + "minimum": 1 } }, "additionalProperties": false diff --git a/schema/controller-config.md b/schema/controller-config.md index 2c363606..140e3272 100644 --- a/schema/controller-config.md +++ b/schema/controller-config.md @@ -26,7 +26,7 @@ |`fleetLimits.enrollmentInFlight`|integer|Enrollment requests in progress at once across all clients. Enrollment
calls the identity provider. Defaults to `64`.| |`fleetLimits.enrollmentRefill`|string|Time for a client address to get back one enrollment. Defaults to `1s`.| |`fleetLimits.maxConnections`|integer|Open fleet connections allowed in total. At the limit, the connection
idle the longest is closed to admit a new one; when none is idle, the
new one is closed. Defaults to `50000`.| -|`fleetLimits.maxMessageBytes`|integer|Largest request a device may send, in bytes. Defaults to `4194304`
(4 MiB).| +|`fleetLimits.maxMessageBytes`|integer|Largest message, in bytes, that a device and the controller may send
each other. Requests from clients without a device certificate, such as
enrollment, stay limited to 64 KiB. Devices accept at most 4 MiB, so
raising this only raises what devices may send. Defaults to `4194304`
(4 MiB).| |`fleetListen`|string|Address on which the device-facing gRPC fleet API listens.| |`gatewayJwt`|object|LLM gateway JWT signing settings.| |`gatewayJwt.issuer`|string|Issuer claim placed in generated JWTs.| From 7fd2ceb867483bb8a3e3f3ecfdd878b9db22670f Mon Sep 17 00:00:00 2001 From: Scott Weiss Date: Mon, 9 Nov 2026 14:43:29 -0500 Subject: [PATCH 28/49] test(controller): keep the first-request deadline out of the idle-close tests Three idle-close tests set the first-request deadline to 5 seconds, the same as the bound each test waits within. With request tracking broken, those connections were still closed by the first-request deadline just before the bound, so a slow run could let the regression pass. The deadline is now 60 seconds, so only the idle close can end the connection in time. --- crates/controller/src/fleet_limits/tests.rs | 9 ++++++--- 1 file changed, 6 insertions(+), 3 deletions(-) diff --git a/crates/controller/src/fleet_limits/tests.rs b/crates/controller/src/fleet_limits/tests.rs index 9d229c98..2cf8c195 100644 --- a/crates/controller/src/fleet_limits/tests.rs +++ b/crates/controller/src/fleet_limits/tests.rs @@ -510,7 +510,8 @@ async fn a_tls_connection_that_never_makes_a_request_is_closed() { async fn a_tls_connection_idle_after_its_last_request_is_closed() { let pki = pki(); let mut limits = limits(); - limits.first_request = Duration::from_secs(5); + // Far past the test's bound, so only the idle close can end the connection. + limits.first_request = Duration::from_secs(60); limits.idle_connection = Duration::from_millis(200); let (address, connections) = serve_mtls(limits, &pki).await; let mut client = tls_client(address, &pki, true).await; @@ -695,7 +696,8 @@ async fn a_connection_that_never_makes_a_request_is_closed() { #[tokio::test] async fn a_connection_idle_after_its_last_request_is_closed() { let mut limits = limits(); - limits.first_request = Duration::from_secs(5); + // Far past the test's bound, so only the idle close can end the connection. + limits.first_request = Duration::from_secs(60); limits.idle_connection = Duration::from_millis(200); let (address, connections) = serve_tracked( limits, @@ -722,7 +724,8 @@ async fn a_connection_idle_after_its_last_request_is_closed() { #[tokio::test] async fn a_response_the_client_never_accepts_does_not_keep_its_connection_busy() { let mut limits = limits(); - limits.first_request = Duration::from_secs(5); + // Far past the test's bound, so only the idle close can end the connection. + limits.first_request = Duration::from_secs(60); limits.idle_connection = Duration::from_millis(200); let (address, connections) = serve_tracked( limits, From 95a3688d6d0d09859bbf24384610457ee2c5765e Mon Sep 17 00:00:00 2001 From: Scott Weiss Date: Mon, 9 Nov 2026 18:10:27 -0500 Subject: [PATCH 29/49] fix(controller): free a Connect stream's connection when its session ends A Connect stream kept its connection busy until hyper dropped the response body. A device that grants no HTTP/2 flow-control window never lets the server send that body, so after the session ended (the device half-closed, or opened a newer stream) hyper kept the body and the connection stayed busy: never idle-closed, never evicted, while the device answered PINGs. Each such connection cost one stream-open token, about 20 per device certificate and then one a minute. The request layer now puts a SessionBusy handle in a Connect request's extensions. Connect claims a release from it and holds that in its session task, so the connection goes idle when the session ends or when hyper drops the body, whichever comes first. A handler that claims no release keeps the previous behavior. --- crates/controller/src/fleet_limits.rs | 4 +- .../src/fleet_limits/connections.rs | 100 ++++++++++++++---- crates/controller/src/fleet_limits/tests.rs | 43 +++++++- crates/controller/src/service.rs | 6 +- 4 files changed, 128 insertions(+), 25 deletions(-) diff --git a/crates/controller/src/fleet_limits.rs b/crates/controller/src/fleet_limits.rs index 70e61f22..1b8457ca 100644 --- a/crates/controller/src/fleet_limits.rs +++ b/crates/controller/src/fleet_limits.rs @@ -36,7 +36,9 @@ use tower_layer::{Identity, Stack}; mod connections; -pub use connections::{Connections, LimitedConnection, LimitedIncoming, TrackRequests}; +pub use connections::{ + Connections, LimitedConnection, LimitedIncoming, ReleaseOnDrop, SessionBusy, TrackRequests, +}; /// Limits applied to the fleet gRPC listener. #[derive(Clone, Copy, Debug, PartialEq)] diff --git a/crates/controller/src/fleet_limits/connections.rs b/crates/controller/src/fleet_limits/connections.rs index 9bc3be79..7ee133c6 100644 --- a/crates/controller/src/fleet_limits/connections.rs +++ b/crates/controller/src/fleet_limits/connections.rs @@ -14,9 +14,10 @@ //! one, so idle connections cannot crowd out devices that are using theirs. //! //! A request other than `Connect` is in progress only while its handler runs, -//! not while its response is sent. Otherwise a client could keep a connection -//! busy forever by granting no HTTP/2 flow-control window, so that the server -//! can never send it the response body. +//! not while its response is sent, and a `Connect` stream only while its +//! session lasts. Otherwise a client could keep a connection busy forever by +//! granting no HTTP/2 flow-control window, so that the server can never send +//! the response body and hyper never drops it. use std::{ collections::{HashMap, VecDeque}, @@ -471,7 +472,7 @@ impl AsyncWrite for LimitedConnection { } /// Marks the connection busy while a request's handler runs, and for `Connect` -/// until its stream ends. +/// while its session lasts (see `SessionBusy`). struct RequestGuard { connections: Connections, peer: SocketAddr, @@ -555,35 +556,92 @@ where self.inner.poll_ready(context) } - fn call(&mut self, request: http::Request) -> Self::Future { + fn call(&mut self, mut request: http::Request) -> Self::Future { let guard = remote_addr(&request).and_then(|peer| self.connections.start_request(peer)); - let session = is_connect(request.uri().path()); + // `Connect` is the device's session, so its connection stays busy past + // the handler: see `SessionBusy`. Any other request is done once its + // handler returns. + let (guard, session) = match guard { + Some(guard) if is_connect(request.uri().path()) => { + let session = SessionBusy::new(guard); + request.extensions_mut().insert(session.clone()); + (None, Some(session)) + } + guard => (guard, None), + }; let response = self.inner.call(request); Box::pin(async move { - let response = response.await?; - match guard { - // The stream is the device's session: keep the connection busy - // until it ends. - Some(guard) if session => Ok(response.map(|body| { + let response = response.await; + drop(guard); + let response = response?; + Ok(match session { + Some(session) => response.map(|body| { Body::new(TrackedBody { inner: body, - _guard: guard, + _release: session.release_on_drop(), }) - })), - // Any other response is done once its handler returns. - guard => { - drop(guard); - Ok(response) - } - } + }), + None => response, + }) }) } } -/// A response body that keeps its connection busy until it is dropped. +/// Keeps a `Connect` stream's connection busy. The connection goes idle when +/// the first release is dropped: the response body holds one, and the handler +/// can claim another from the request extensions (`SessionBusy::of`) and hold +/// it in the session's task. Without the handler's release, a client that +/// grants no HTTP/2 flow-control window could end the session and still keep +/// the connection busy, because hyper keeps a response body it cannot send. +/// If every handle is dropped without a release, the connection goes idle too. +#[derive(Clone)] +pub struct SessionBusy { + guard: Arc>>, +} + +impl SessionBusy { + fn new(guard: RequestGuard) -> Self { + Self { + guard: Arc::new(Mutex::new(Some(guard))), + } + } + + /// The handle for `request`'s connection, when it is a tracked `Connect`. + pub fn of(request: &tonic::Request) -> Option { + request.extensions().get::().cloned() + } + + /// Marks the connection idle when the returned value is dropped, unless + /// another release already has. + pub fn release_on_drop(&self) -> ReleaseOnDrop { + ReleaseOnDrop { + guard: self.guard.clone(), + } + } +} + +/// Releases a `SessionBusy` when dropped. +pub struct ReleaseOnDrop { + guard: Arc>>, +} + +impl Drop for ReleaseOnDrop { + fn drop(&mut self) { + // Taken out first, so the guard's own drop (which takes the registry + // lock) runs without this lock held. + let guard = self + .guard + .lock() + .expect("connection session lock poisoned") + .take(); + drop(guard); + } +} + +/// A `Connect` response body; dropping it releases its connection's session. struct TrackedBody { inner: Body, - _guard: RequestGuard, + _release: ReleaseOnDrop, } impl HttpBody for TrackedBody { diff --git a/crates/controller/src/fleet_limits/tests.rs b/crates/controller/src/fleet_limits/tests.rs index 2cf8c195..42e09088 100644 --- a/crates/controller/src/fleet_limits/tests.rs +++ b/crates/controller/src/fleet_limits/tests.rs @@ -41,7 +41,8 @@ async fn bounded(call: impl std::future::Future) -> T { } /// A fleet service whose calls take `delay`; `Connect` sends a message after each -/// `delay`, twice. +/// `delay`, twice, then ends its session, holding the connection's session +/// release as the controller's `Connect` does. #[derive(Clone)] struct Slow { delay: Duration, @@ -91,11 +92,13 @@ impl FleetAgent for Slow { async fn connect( &self, - _: Request>, + request: Request>, ) -> Result, Status> { + let session = super::SessionBusy::of(&request).map(|busy| busy.release_on_drop()); let (sender, receiver) = mpsc::channel(2); let delay = self.delay; tokio::spawn(async move { + let _session = session; for _ in 0..2 { time::sleep(delay).await; if sender.send(Ok(ControllerMessage::default())).await.is_err() { @@ -766,6 +769,42 @@ async fn a_response_the_client_never_accepts_does_not_keep_its_connection_busy() assert!(response.is_err(), "the stalled response was delivered"); } +#[tokio::test] +async fn a_stalled_connect_response_frees_its_connection_when_the_session_ends() { + let mut limits = limits(); + // Far past the test's bound, so only the idle close can end the connection. + limits.first_request = Duration::from_secs(60); + limits.idle_connection = Duration::from_millis(200); + let (address, connections) = serve_tracked( + limits, + None, + Slow { + delay: Duration::ZERO, + }, + ) + .await; + // A client that grants no flow-control window: the session's messages can + // never be sent, so hyper keeps the response body after the session ends. + let channel = tonic::transport::Endpoint::from_shared(format!("http://{address}")) + .unwrap() + .initial_stream_window_size(0) + .connect() + .await + .unwrap(); + let (_outbound, receiver) = mpsc::channel::(1); + let _stream = bounded(FleetAgentClient::new(channel).connect(ReceiverStream::new(receiver))) + .await + .expect("open the stream"); + assert_eq!(connections.open_count(), 1); + + bounded(async { + while connections.open_count() > 0 { + time::sleep(Duration::from_millis(50)).await; + } + }) + .await; +} + #[tokio::test] async fn an_open_stream_keeps_its_connection_from_going_idle() { let mut limits = limits(); diff --git a/crates/controller/src/service.rs b/crates/controller/src/service.rs index d638fcc2..d9adf1ea 100644 --- a/crates/controller/src/service.rs +++ b/crates/controller/src/service.rs @@ -22,7 +22,7 @@ use crate::{ daemon_config::DaemonConfigStore, database::{Database, DevicePrincipal}, device_ca::DeviceCertificateIssuer, - fleet_limits::{DeviceStreams, EnrollmentLimiter, FLEET_LIMITS, FleetLimits}, + fleet_limits::{DeviceStreams, EnrollmentLimiter, FLEET_LIMITS, FleetLimits, SessionBusy}, gateway_jwt::GatewayJwtIssuer, oidc::{OidcPrincipal, OidcProvider}, }; @@ -274,6 +274,9 @@ impl FleetAgent for FleetAgentService { bearer_credential(request.metadata())?.to_owned(), ); + // The connection counts as busy while this session runs, and no longer: + // hyper may keep a response body that a device will not accept. + let session_busy = SessionBusy::of(&request).map(|busy| busy.release_on_drop()); let mut inbound = request.into_inner(); let (sender, receiver) = mpsc::channel(8); let mut daemon_updates = self.daemon_config.subscribe(); @@ -295,6 +298,7 @@ impl FleetAgent for FleetAgentService { let auth_sender = sender; let mut device_stream = self.device_streams.open(&authenticated_device_id); tokio::spawn(async move { + let _session_busy = session_busy; let enrolled_principal = match database.device_principal(&authenticated_device_id).await { Ok(principal) => principal, From 0d8fade781164a5a4299bace27a5b73d78640029 Mon Sep 17 00:00:00 2001 From: Scott Weiss Date: Mon, 9 Nov 2026 18:10:58 -0500 Subject: [PATCH 30/49] fix(controller): never let a Connect session wait on a device that stopped reading The Connect session task sent config pushes and its closing statuses with `send().await` on an 8-message channel. A device that stopped reading could fill it, and the task then waited forever: it no longer saw a newer stream from the same device, the client going away, or the access-token check, and it held the session (and so the connection) open. Sends in the session loop now use `try_send`. A config push that finds the channel full closes the stream, and the device reconnects to get the current configuration. A closing status that does not fit is dropped, since the stream ends anyway. This path needs a device certificate and OIDC to exercise, so it has no test. --- crates/controller/src/service.rs | 37 ++++++++++++++++---------------- 1 file changed, 19 insertions(+), 18 deletions(-) diff --git a/crates/controller/src/service.rs b/crates/controller/src/service.rs index d9adf1ea..fbbd5d54 100644 --- a/crates/controller/src/service.rs +++ b/crates/controller/src/service.rs @@ -317,15 +317,20 @@ impl FleetAgent for FleetAgentService { break; } let daemon = daemon_updates.borrow_and_update().clone(); - if let Some(daemon) = daemon - && auth_sender - .send(Ok(ControllerMessage { - message: Some(controller_message::Message::DaemonConfig(daemon)), - })) - .await - .is_err() - { - break; + if let Some(daemon) = daemon { + // Sends in this loop never wait for room: a device that + // stopped reading would hold the task, and the session, + // forever. + match auth_sender.try_send(Ok(ControllerMessage { + message: Some(controller_message::Message::DaemonConfig(daemon)), + })) { + Ok(()) => {} + Err(mpsc::error::TrySendError::Full(_)) => { + warn!(device_id = %authenticated_device_id, "closing stream after the device stopped reading it"); + break; + } + Err(mpsc::error::TrySendError::Closed(_)) => break, + } } } _ = auth_interval.tick() => { @@ -335,8 +340,7 @@ impl FleetAgent for FleetAgentService { Ok(_) => { warn!(device_id = %authenticated_device_id, "closing stream after OIDC principal changed"); let _ = auth_sender - .send(Err(Status::permission_denied("OIDC principal changed"))) - .await; + .try_send(Err(Status::permission_denied("OIDC principal changed"))); break; } Err(error) => { @@ -346,7 +350,7 @@ impl FleetAgent for FleetAgentService { continue; } warn!(device_id = %authenticated_device_id, "closing stream after OIDC access token revalidation failed"); - let _ = auth_sender.send(Err(status)).await; + let _ = auth_sender.try_send(Err(status)); break; } } @@ -354,8 +358,7 @@ impl FleetAgent for FleetAgentService { _ = &mut device_stream.superseded => { info!(device_id = %authenticated_device_id, "closing stream replaced by a newer stream from the same device"); let _ = auth_sender - .send(Err(Status::aborted("replaced by a newer stream from this device"))) - .await; + .try_send(Err(Status::aborted("replaced by a newer stream from this device"))); break; } message = inbound.message() => match message { @@ -363,8 +366,7 @@ impl FleetAgent for FleetAgentService { if !device_stream.take_message(time::Instant::now()) { warn!(device_id = %authenticated_device_id, "closing stream after the device exceeded its message rate"); let _ = auth_sender - .send(Err(Status::resource_exhausted("too many messages on this stream"))) - .await; + .try_send(Err(Status::resource_exhausted("too many messages on this stream"))); break; } if let Some(agent_message::Message::Hello(hello)) = &message.message @@ -376,8 +378,7 @@ impl FleetAgent for FleetAgentService { "device certificate and hello identity do not match" ); let _ = auth_sender - .send(Err(Status::permission_denied("hello does not match the device certificate"))) - .await; + .try_send(Err(Status::permission_denied("hello does not match the device certificate"))); break; } if let Err(err) = From 54bf011cd26915a010c1abbe603cb79381a024b8 Mon Sep 17 00:00:00 2001 From: Scott Weiss Date: Mon, 9 Nov 2026 18:11:26 -0500 Subject: [PATCH 31/49] refactor(controller): require the fleet limits when building the fleet service FleetAgentService::new filled its enrollment and Connect limits from the built-in defaults, and startup replaced them with the configured ones through `with_fleet_limits`. A code path that built the service without that call would silently ignore `fleetLimits` in controller.yaml. `new` now takes the limits, so every caller has to choose them, and `with_fleet_limits` is gone. --- crates/controller/src/application.rs | 4 +-- crates/controller/src/service.rs | 47 ++++++++++++++++++---------- 2 files changed, 32 insertions(+), 19 deletions(-) diff --git a/crates/controller/src/application.rs b/crates/controller/src/application.rs index 045ee59e..09971e72 100644 --- a/crates/controller/src/application.rs +++ b/crates/controller/src/application.rs @@ -110,9 +110,9 @@ pub async fn run( daemon_config, gateway_jwt_issuer, Some(device_certificate_issuer), + &limits, ) - .with_usage_collection(usage_ingest.is_some()) - .with_fleet_limits(&limits); + .with_usage_collection(usage_ingest.is_some()); let certificate = std::fs::read(&tls.certificate) .with_context(|| format!("read TLS certificate from {}", tls.certificate.display()))?; diff --git a/crates/controller/src/service.rs b/crates/controller/src/service.rs index fbbd5d54..9a1acef1 100644 --- a/crates/controller/src/service.rs +++ b/crates/controller/src/service.rs @@ -22,7 +22,7 @@ use crate::{ daemon_config::DaemonConfigStore, database::{Database, DevicePrincipal}, device_ca::DeviceCertificateIssuer, - fleet_limits::{DeviceStreams, EnrollmentLimiter, FLEET_LIMITS, FleetLimits, SessionBusy}, + fleet_limits::{DeviceStreams, EnrollmentLimiter, FleetLimits, SessionBusy}, gateway_jwt::GatewayJwtIssuer, oidc::{OidcPrincipal, OidcProvider}, }; @@ -40,12 +40,15 @@ pub struct FleetAgentService { } impl FleetAgentService { + /// `limits` sets the enrollment and `Connect` limits; pass the configured + /// ones (`FLEET_LIMITS.configured(...)`), the same the fleet listener uses. pub fn new( oidc: Option, database: Database, daemon_config: DaemonConfigStore, gateway_jwt_issuer: Option, device_certificate_issuer: Option, + limits: &FleetLimits, ) -> Self { Self { oidc, @@ -54,18 +57,11 @@ impl FleetAgentService { gateway_jwt_issuer, device_certificate_issuer, usage_enabled: false, - enrollment: FLEET_LIMITS.enrollment(), - device_streams: FLEET_LIMITS.device_streams(), + enrollment: limits.enrollment(), + device_streams: limits.device_streams(), } } - /// Replaces the default enrollment and `Connect` limits. - pub fn with_fleet_limits(mut self, limits: &FleetLimits) -> Self { - self.enrollment = limits.enrollment(); - self.device_streams = limits.device_streams(); - self - } - /// Tells devices that the controller collects gateway usage. pub fn with_usage_collection(mut self, enabled: bool) -> Self { self.usage_enabled = enabled; @@ -731,9 +727,14 @@ mod tests { enrollment_completion_wait: std::time::Duration::from_millis(10), ..FLEET_LIMITS }; - let service = - FleetAgentService::new(None, database, DaemonConfigStore::new(None), None, None) - .with_fleet_limits(&full); + let service = FleetAgentService::new( + None, + database, + DaemonConfigStore::new(None), + None, + None, + &full, + ); // Without the limiter these fail later, with "OIDC enrollment is disabled". let begin = service @@ -764,8 +765,14 @@ mod tests { .enroll_device(device_id, "host", "issuer", "subject", None) .await .expect("enroll device"); - let service = - FleetAgentService::new(None, database, DaemonConfigStore::new(None), None, None); + let service = FleetAgentService::new( + None, + database, + DaemonConfigStore::new(None), + None, + None, + &crate::fleet_limits::FLEET_LIMITS, + ); let principal = |issuer: &str, subject: &str| OidcPrincipal { issuer: issuer.to_owned(), subject: subject.to_owned(), @@ -818,8 +825,14 @@ mod tests { .enroll_device(device_id, "host", "", "subject", None) .await .expect("enroll legacy device"); - let service = - FleetAgentService::new(None, database, DaemonConfigStore::new(None), None, None); + let service = FleetAgentService::new( + None, + database, + DaemonConfigStore::new(None), + None, + None, + &crate::fleet_limits::FLEET_LIMITS, + ); let status = service .authenticate_device_certificate_identity( From 7e77782624cca498471e2345617781def6fc8477 Mon Sep 17 00:00:00 2001 From: Scott Weiss Date: Tue, 10 Nov 2026 11:47:26 -0500 Subject: [PATCH 32/49] test(controller): check that every fleetLimits setting reaches the fleet limits The test set two of the eight settings, so dropping a copy from `configured` left that setting silently ignored with every test passing. It now names every setting with a non-default value and compares the whole result. --- crates/controller/src/fleet_limits/tests.rs | 35 ++++++++++++++++----- 1 file changed, 27 insertions(+), 8 deletions(-) diff --git a/crates/controller/src/fleet_limits/tests.rs b/crates/controller/src/fleet_limits/tests.rs index 42e09088..c014b389 100644 --- a/crates/controller/src/fleet_limits/tests.rs +++ b/crates/controller/src/fleet_limits/tests.rs @@ -844,17 +844,36 @@ fn the_configuration_defaults_match_the_default_limits() { } #[test] -fn configured_limits_replace_the_defaults() { +fn configured_limits_replace_every_deployment_limit() { use agentdesktop_core::config::ControllerFleetLimitsConfig; - let limits = FLEET_LIMITS.configured(&ControllerFleetLimitsConfig { - connections_per_address: 20_000, - device_message_refill: Duration::from_millis(10), - ..ControllerFleetLimitsConfig::default() - }); + // The literal names every setting, each with a value that is not its + // default, so a setting `configured` does not copy fails here. + let config = ControllerFleetLimitsConfig { + max_connections: 8, + connections_per_address: 7, + max_message_bytes: 6, + enrollment_in_flight: 5, + enrollment_burst: 4, + enrollment_refill: Duration::from_millis(3), + device_message_burst: 2, + device_message_refill: Duration::from_millis(1), + }; - assert_eq!(limits.connections_per_address, 20_000); - assert_eq!(limits.device_message_refill, Duration::from_millis(10)); + assert_eq!( + FLEET_LIMITS.configured(&config), + FleetLimits { + max_connections: 8, + connections_per_address: 7, + max_message_bytes: 6, + enrollment_in_flight: 5, + enrollment_burst: 4, + enrollment_refill: Duration::from_millis(3), + device_message_burst: 2, + device_message_refill: Duration::from_millis(1), + ..FLEET_LIMITS + } + ); } #[test] From e6615416404a9cd7a41918c71aa9040ee0f2b7f1 Mon Sep 17 00:00:00 2001 From: Scott Weiss Date: Tue, 10 Nov 2026 11:48:34 -0500 Subject: [PATCH 33/49] fix(controller): say what a throttled warning's count covers A throttled warning counts every refusal since the last warning, from any client or device, and reports it with the first refusal after the interval. It named only the client or device that triggered it, as if the count were that one's. The fields are now `latest_address` or `latest_device_id` and `refused_since_last_warning` (or `closed_since_last_warning`). --- crates/controller/src/fleet_limits.rs | 21 +++++++++++-------- .../src/fleet_limits/connections.rs | 14 ++++++------- 2 files changed, 19 insertions(+), 16 deletions(-) diff --git a/crates/controller/src/fleet_limits.rs b/crates/controller/src/fleet_limits.rs index 1b8457ca..7ce17eae 100644 --- a/crates/controller/src/fleet_limits.rs +++ b/crates/controller/src/fleet_limits.rs @@ -358,10 +358,10 @@ impl EnrollmentLimiter { if let Some(peer) = peer && !self.take(peer.ip(), Instant::now()) { - if let Some(events) = ENROLLMENT_RATE.event(Instant::now()) { + if let Some(refused) = ENROLLMENT_RATE.event(Instant::now()) { tracing::warn!( - peer = %address_key(peer.ip()), - events, + latest_address = %address_key(peer.ip()), + refused_since_last_warning = refused, "refused enrollment requests over the per-address rate limit" ); } @@ -402,9 +402,9 @@ impl EnrollmentLimiter { } fn in_flight_refusal() -> Status { - if let Some(events) = ENROLLMENT_IN_FLIGHT.event(Instant::now()) { + if let Some(refused) = ENROLLMENT_IN_FLIGHT.event(Instant::now()) { tracing::warn!( - events, + refused_since_last_warning = refused, "refused enrollment requests with every in-flight slot taken" ); } @@ -507,10 +507,10 @@ impl DeviceStreams { if admitted { Ok(()) } else { - if let Some(events) = STREAM_OPENS.event(Instant::now()) { + if let Some(refused) = STREAM_OPENS.event(Instant::now()) { tracing::warn!( - device_id, - events, + latest_device_id = device_id, + refused_since_last_warning = refused, "refused Connect streams from devices that open them too often" ); } @@ -579,7 +579,10 @@ fn refilled(bucket: Bucket, now: Instant, refill: Duration, burst: f64) -> f64 { /// A warning that fires once per refused request or connection. A flood would /// log a line per event, so this logs at most one per `WARNING_INTERVAL` and -/// says how many events it stands for. +/// says how many events it stands for: every event since the last warning, from +/// any client or device, reported with the first event after the interval +/// (however much later that is). A warning that names a client or device names +/// only the latest one. pub(crate) struct ThrottledWarning { state: Mutex, } diff --git a/crates/controller/src/fleet_limits/connections.rs b/crates/controller/src/fleet_limits/connections.rs index 7ee133c6..961c8856 100644 --- a/crates/controller/src/fleet_limits/connections.rs +++ b/crates/controller/src/fleet_limits/connections.rs @@ -232,19 +232,19 @@ impl Connections { .get(&address) .is_some_and(|count| *count >= self.shared.connections_per_address) { - if let Some(events) = PER_ADDRESS_CAP.event(Instant::now()) { + if let Some(closed) = PER_ADDRESS_CAP.event(Instant::now()) { tracing::warn!( - %address, - events, + latest_address = %address, + closed_since_last_warning = closed, "closed new fleet connections over the per-address connection limit" ); } return None; } if state.open.len() >= self.shared.max_connections && !state.evict_idle() { - if let Some(events) = TOTAL_CAP.event(Instant::now()) { + if let Some(closed) = TOTAL_CAP.event(Instant::now()) { tracing::warn!( - events, + closed_since_last_warning = closed, "closed new fleet connections at the connection limit with no connection idle" ); } @@ -363,9 +363,9 @@ impl LimitedConnection { { self.evicted = None; self.closed = true; - if let Some(events) = EVICTIONS.event(Instant::now()) { + if let Some(closed) = EVICTIONS.event(Instant::now()) { tracing::warn!( - events, + closed_since_last_warning = closed, "closed idle fleet connections to make room at the connection limit" ); } From a5342e3dc9dee2f2f360513820b8e1564cff5960 Mon Sep 17 00:00:00 2001 From: Scott Weiss Date: Tue, 10 Nov 2026 11:49:31 -0500 Subject: [PATCH 34/49] docs(controller): correct two comments and explain the enrollment check order The connection caps comment said a connection over either cap is closed on accept; at the total cap an idle connection is evicted instead. The device-grant comment described releasing the reservation by hand, which its Drop impl now does. BeginEnrollment spends the address's budget before it looks for a slot on purpose: requests refused for want of a slot still count, so the addresses flooding the slots cannot keep their budget to take the slots back as they free up. --- crates/controller/src/fleet_limits.rs | 17 +++++++++++------ crates/controller/src/oidc.rs | 6 +++--- 2 files changed, 14 insertions(+), 9 deletions(-) diff --git a/crates/controller/src/fleet_limits.rs b/crates/controller/src/fleet_limits.rs index 7ce17eae..34a88801 100644 --- a/crates/controller/src/fleet_limits.rs +++ b/crates/controller/src/fleet_limits.rs @@ -59,17 +59,18 @@ pub struct FleetLimits { /// for peers that vanish without closing. pub tcp_keepalive: Duration, /// Open connections allowed in total and from one client address (an IPv4 - /// address or an IPv6 /64). A connection over either cap is closed as soon - /// as it is accepted. The - /// per-address cap is high because many devices can share an address behind - /// NAT, or appear to when a load balancer rewrites client addresses. + /// address or an IPv6 /64). A new connection over the per-address cap is + /// closed as soon as it is accepted. At the total cap, the longest-idle + /// connection is closed to admit a new one, and the new one is closed only + /// when no connection is idle. The per-address cap is high because many + /// devices can share an address behind NAT, or appear to when a load + /// balancer rewrites client addresses. pub max_connections: usize, pub connections_per_address: usize, /// Time a new connection has to make its first request, and time a /// connection may stay without any request in progress after its last one /// finished. A connection with a request in progress (including an open - /// `Connect` stream) is never idle. At the total cap, the longest-idle - /// connection is closed to admit a new one. + /// `Connect` stream) is never idle. pub first_request: Duration, pub idle_connection: Duration, /// Concurrent requests one connection may have in progress, and concurrent @@ -355,6 +356,10 @@ impl EnrollmentLimiter { /// Admits one `BeginEnrollment` from `peer`, or rejects it with /// `RESOURCE_EXHAUSTED`. Hold the returned permit until the request finishes. pub fn admit(&self, peer: Option) -> Result { + // The address's budget is spent before a slot is sought, so a request + // refused for want of a slot still counts against it. Otherwise, while + // a flood holds every slot, the flooding addresses would keep their + // whole budget and take the slots back the moment they free up. if let Some(peer) = peer && !self.take(peer.ip(), Instant::now()) { diff --git a/crates/controller/src/oidc.rs b/crates/controller/src/oidc.rs index 55d403b7..def11089 100644 --- a/crates/controller/src/oidc.rs +++ b/crates/controller/src/oidc.rs @@ -280,9 +280,9 @@ impl OidcProvider { Ok(response) } - /// The IdP round-trip and response validation for [`begin_device`], split out so - /// its error paths all flow through one `?`-propagating function and the caller - /// can release the reservation on any of them uniformly. + /// The IdP round-trip and response validation for [`begin_device`]. On any + /// error the caller's [`Reservation`] is dropped, which releases the pending + /// enrollment. async fn request_device_grant( &self, enrollment_id: &str, From f5f2757a11f3e6ef9cd02f30170c73c5e141411d Mon Sep 17 00:00:00 2001 From: Scott Weiss Date: Tue, 10 Nov 2026 11:49:43 -0500 Subject: [PATCH 35/49] docs(config): describe the burst settings as bursts, not concurrency `enrollmentBurst` and `deviceMessageBurst` said "at once", which reads like a limit on requests in progress; that is what `enrollmentInFlight` is. --- crates/controller/src/fleet_limits.rs | 6 +++--- crates/core/src/config.rs | 10 +++++----- schema/controller-config.json | 4 ++-- schema/controller-config.md | 4 ++-- 4 files changed, 12 insertions(+), 12 deletions(-) diff --git a/crates/controller/src/fleet_limits.rs b/crates/controller/src/fleet_limits.rs index 34a88801..cd314e80 100644 --- a/crates/controller/src/fleet_limits.rs +++ b/crates/controller/src/fleet_limits.rs @@ -90,7 +90,7 @@ pub struct FleetLimits { /// also bounds the load a flood can put on it. pub enrollment_in_flight: usize, /// Per-client-address `BeginEnrollment` budget: up to `enrollment_burst` - /// requests at once, refilled at one request per `enrollment_refill`. + /// requests in a burst, refilled at one request per `enrollment_refill`. /// Generous, because many devices can share one address behind NAT. pub enrollment_burst: u32, pub enrollment_refill: Duration, @@ -99,7 +99,7 @@ pub struct FleetLimits { /// already signed in, and a refusal makes them start over. pub enrollment_completion_wait: Duration, /// Messages a device may send on its `Connect` streams: up to - /// `device_message_burst` at once, refilled at one per + /// `device_message_burst` in a burst, refilled at one per /// `device_message_refill`. The budget belongs to the device, so /// reconnecting does not reset it. Devices send a heartbeat every 30 /// seconds, inventory periodically and one telemetry event per tool event, @@ -107,7 +107,7 @@ pub struct FleetLimits { pub device_message_burst: u32, pub device_message_refill: Duration, /// How often a device may open a `Connect` stream: up to - /// `stream_open_burst` at once, refilled at one per `stream_open_refill`. + /// `stream_open_burst` in a burst, refilled at one per `stream_open_refill`. /// Each open checks the access token with the identity provider. pub stream_open_burst: u32, pub stream_open_refill: Duration, diff --git a/crates/core/src/config.rs b/crates/core/src/config.rs index 0daf5dcf..ec3063cb 100644 --- a/crates/core/src/config.rs +++ b/crates/core/src/config.rs @@ -604,17 +604,17 @@ pub struct ControllerFleetLimitsConfig { /// calls the identity provider. Defaults to `64`. #[cfg_attr(feature = "schema", schemars(range(min = 1)))] pub enrollment_in_flight: usize, - /// Enrollments one client address may start at once; it then gets one more - /// every `enrollmentRefill`. Defaults to `60`. + /// Enrollments one client address may start in a burst, after which it may + /// start one more every `enrollmentRefill`. Defaults to `60`. #[cfg_attr(feature = "schema", schemars(range(min = 1)))] pub enrollment_burst: u32, /// Time for a client address to get back one enrollment. Defaults to `1s`. #[serde(with = "humantime_serde")] #[cfg_attr(feature = "schema", schemars(with = "String"))] pub enrollment_refill: Duration, - /// Messages one device may send at once on its `Connect` stream; it then - /// gets one more every `deviceMessageRefill`. A device over its budget is - /// disconnected. Defaults to `600`. + /// Messages one device may send in a burst on its `Connect` stream, after + /// which it may send one more every `deviceMessageRefill`. A device over + /// its budget is disconnected. Defaults to `600`. #[cfg_attr(feature = "schema", schemars(range(min = 1)))] pub device_message_burst: u32, /// Time for a device to get back one message. Defaults to `50ms`. diff --git a/schema/controller-config.json b/schema/controller-config.json index b1c5a15f..646722dd 100644 --- a/schema/controller-config.json +++ b/schema/controller-config.json @@ -217,7 +217,7 @@ "minimum": 1 }, "deviceMessageBurst": { - "description": "Messages one device may send at once on its `Connect` stream; it then\ngets one more every `deviceMessageRefill`. A device over its budget is\ndisconnected. Defaults to `600`.", + "description": "Messages one device may send in a burst on its `Connect` stream, after\nwhich it may send one more every `deviceMessageRefill`. A device over\nits budget is disconnected. Defaults to `600`.", "type": "integer", "format": "uint32", "default": 600, @@ -229,7 +229,7 @@ "default": "50ms" }, "enrollmentBurst": { - "description": "Enrollments one client address may start at once; it then gets one more\nevery `enrollmentRefill`. Defaults to `60`.", + "description": "Enrollments one client address may start in a burst, after which it may\nstart one more every `enrollmentRefill`. Defaults to `60`.", "type": "integer", "format": "uint32", "default": 60, diff --git a/schema/controller-config.md b/schema/controller-config.md index 140e3272..8c835e00 100644 --- a/schema/controller-config.md +++ b/schema/controller-config.md @@ -20,9 +20,9 @@ |`databaseUrl`|string|SQLite or PostgreSQL URL used for controller state.| |`fleetLimits`|object|Limits on the device-facing fleet API that depend on the deployment.| |`fleetLimits.connectionsPerAddress`|integer|Open fleet connections allowed from one client address: an IPv4
address, or an IPv6 /64. Raise it when more devices share an address,
such as behind a NAT or a load balancer that does not preserve client
addresses. Defaults to `4096`.| -|`fleetLimits.deviceMessageBurst`|integer|Messages one device may send at once on its `Connect` stream; it then
gets one more every `deviceMessageRefill`. A device over its budget is
disconnected. Defaults to `600`.| +|`fleetLimits.deviceMessageBurst`|integer|Messages one device may send in a burst on its `Connect` stream, after
which it may send one more every `deviceMessageRefill`. A device over
its budget is disconnected. Defaults to `600`.| |`fleetLimits.deviceMessageRefill`|string|Time for a device to get back one message. Defaults to `50ms`.| -|`fleetLimits.enrollmentBurst`|integer|Enrollments one client address may start at once; it then gets one more
every `enrollmentRefill`. Defaults to `60`.| +|`fleetLimits.enrollmentBurst`|integer|Enrollments one client address may start in a burst, after which it may
start one more every `enrollmentRefill`. Defaults to `60`.| |`fleetLimits.enrollmentInFlight`|integer|Enrollment requests in progress at once across all clients. Enrollment
calls the identity provider. Defaults to `64`.| |`fleetLimits.enrollmentRefill`|string|Time for a client address to get back one enrollment. Defaults to `1s`.| |`fleetLimits.maxConnections`|integer|Open fleet connections allowed in total. At the limit, the connection
idle the longest is closed to admit a new one; when none is idle, the
new one is closed. Defaults to `50000`.| From 56a31be23b54be62b408f8e2a49204dc5fc55670 Mon Sep 17 00:00:00 2001 From: Scott Weiss Date: Tue, 10 Nov 2026 11:50:23 -0500 Subject: [PATCH 36/49] fix(controller): stop re-arming a busy connection's idle timer on every read A busy connection has no idle deadline, but each read armed its timer for an idle period from now, so a connection with an open Connect stream reset its timer on every read. The guard of the connection's last request already wakes it when it goes idle, so a busy connection now arms nothing. --- .../src/fleet_limits/connections.rs | 30 +++++++++---------- 1 file changed, 15 insertions(+), 15 deletions(-) diff --git a/crates/controller/src/fleet_limits/connections.rs b/crates/controller/src/fleet_limits/connections.rs index 961c8856..e06d244b 100644 --- a/crates/controller/src/fleet_limits/connections.rs +++ b/crates/controller/src/fleet_limits/connections.rs @@ -388,22 +388,22 @@ impl LimitedConnection { } } loop { - let now = Instant::now(); - let target = match self.activity.idle_deadline(&self.connections.shared) { - Some(deadline) if deadline <= now => { - self.closed = true; - tracing::debug!(peer = %self.peer, "closing idle fleet connection"); - return Err(io::Error::new( - io::ErrorKind::TimedOut, - "idle fleet connection closed", - )); - } - Some(deadline) => deadline, - // Busy: check again after an idle period. - None => now + self.connections.shared.idle, + // Busy: no deadline. The guard of the last request to finish wakes + // this connection (through the waker registered above) when it + // goes idle. + let Some(deadline) = self.activity.idle_deadline(&self.connections.shared) else { + return Ok(()); }; - if self.timer.deadline() != target { - self.timer.as_mut().reset(target); + if deadline <= Instant::now() { + self.closed = true; + tracing::debug!(peer = %self.peer, "closing idle fleet connection"); + return Err(io::Error::new( + io::ErrorKind::TimedOut, + "idle fleet connection closed", + )); + } + if self.timer.deadline() != deadline { + self.timer.as_mut().reset(deadline); } if self.timer.as_mut().poll(context).is_pending() { return Ok(()); From 2c23e265cb84fedf6ecfb5e16889b033cba2e2bc Mon Sep 17 00:00:00 2001 From: Scott Weiss Date: Tue, 10 Nov 2026 11:51:43 -0500 Subject: [PATCH 37/49] test(controller): lower only the limits each fleet test exercises The shared helper said untouched limits stayed at their production values, but it shortened the handshake, request and keepalive limits for every test. Tests meant to show a connection staying open then depended on a 100 ms PING reply, and keepalive woke every idle connection often enough that no test noticed when a finished request failed to wake its connection. Each test now starts from the defaults and lowers what it exercises; the stalled Connect test now fails if that wake is missing. --- crates/controller/src/fleet_limits/tests.rs | 94 ++++++++++----------- 1 file changed, 47 insertions(+), 47 deletions(-) diff --git a/crates/controller/src/fleet_limits/tests.rs b/crates/controller/src/fleet_limits/tests.rs index c014b389..7f8f298f 100644 --- a/crates/controller/src/fleet_limits/tests.rs +++ b/crates/controller/src/fleet_limits/tests.rs @@ -22,17 +22,6 @@ use tonic::{ use super::{Connections, FLEET_LIMITS, FleetLimits}; -/// Limits small enough for tests to observe; anything a test does not exercise -/// stays at the production value. -fn limits() -> FleetLimits { - let mut limits = FLEET_LIMITS; - limits.tls_handshake = Duration::from_millis(200); - limits.request = Duration::from_millis(200); - limits.keepalive_interval = Duration::from_millis(100); - limits.keepalive_timeout = Duration::from_millis(100); - limits -} - /// Fails the test, instead of hanging it, when the server never ends the wait. async fn bounded(call: impl std::future::Future) -> T { time::timeout(Duration::from_secs(5), call) @@ -156,8 +145,12 @@ async fn until_closed(mut stream: TcpStream) { async fn a_client_that_never_finishes_the_tls_handshake_is_disconnected() { let certified = rcgen::generate_simple_self_signed(vec!["localhost".to_owned()]).unwrap(); let identity = Identity::from_pem(certified.cert.pem(), certified.signing_key.serialize_pem()); + let limits = FleetLimits { + tls_handshake: Duration::from_millis(200), + ..FLEET_LIMITS + }; let address = serve( - limits(), + limits, Some(identity), Slow { delay: Duration::ZERO, @@ -173,8 +166,12 @@ async fn a_client_that_never_finishes_the_tls_handshake_is_disconnected() { #[tokio::test] async fn a_request_that_runs_too_long_is_ended() { + let limits = FleetLimits { + request: Duration::from_millis(200), + ..FLEET_LIMITS + }; let address = serve( - limits(), + limits, None, Slow { delay: Duration::from_secs(30), @@ -195,9 +192,13 @@ async fn a_request_that_runs_too_long_is_ended() { #[tokio::test] async fn the_request_timeout_does_not_end_an_open_stream() { + let limits = FleetLimits { + request: Duration::from_millis(200), + ..FLEET_LIMITS + }; // Each message arrives after more than the request timeout. - let delay = limits().request * 2; - let address = serve(limits(), None, Slow { delay }).await; + let delay = limits.request * 2; + let address = serve(limits, None, Slow { delay }).await; let (_outbound, receiver) = mpsc::channel::(1); let mut stream = bounded(client(address).await.connect(ReceiverStream::new(receiver))) @@ -215,8 +216,13 @@ async fn the_request_timeout_does_not_end_an_open_stream() { #[tokio::test] async fn a_connection_that_stops_answering_pings_is_closed() { + let limits = FleetLimits { + keepalive_interval: Duration::from_millis(100), + keepalive_timeout: Duration::from_millis(100), + ..FLEET_LIMITS + }; let address = serve( - limits(), + limits, None, Slow { delay: Duration::ZERO, @@ -257,16 +263,12 @@ async fn two_calls_on_one_connection(limits: FleetLimits) -> Duration { started.elapsed() } -fn one_at_a_time() -> FleetLimits { - let mut limits = limits(); - limits.request = Duration::from_secs(5); - limits -} - #[tokio::test] async fn requests_beyond_the_per_connection_limit_wait() { - let mut limits = one_at_a_time(); - limits.requests_per_connection = 1; + let limits = FleetLimits { + requests_per_connection: 1, + ..FLEET_LIMITS + }; let elapsed = two_calls_on_one_connection(limits).await; @@ -278,8 +280,10 @@ async fn requests_beyond_the_per_connection_limit_wait() { #[tokio::test] async fn streams_beyond_the_per_connection_limit_wait() { - let mut limits = one_at_a_time(); - limits.streams_per_connection = 1; + let limits = FleetLimits { + streams_per_connection: 1, + ..FLEET_LIMITS + }; let elapsed = two_calls_on_one_connection(limits).await; @@ -290,10 +294,11 @@ async fn streams_beyond_the_per_connection_limit_wait() { } fn small_messages() -> FleetLimits { - let mut limits = limits(); - limits.max_message_bytes = 1024; - limits.unauthenticated_message_bytes = 128; - limits + FleetLimits { + max_message_bytes: 1024, + unauthenticated_message_bytes: 128, + ..FLEET_LIMITS + } } #[tokio::test] @@ -453,7 +458,7 @@ async fn a_tls_connection_that_never_makes_a_request_is_closed() { use rustls::pki_types::{CertificateDer, pem::PemObject}; let pki = pki(); - let mut limits = limits(); + let mut limits = FLEET_LIMITS; limits.first_request = Duration::from_millis(300); let (address, _) = serve_mtls(limits, &pki).await; let ca = pki.ca.clone().into_inner(); @@ -512,7 +517,7 @@ async fn a_tls_connection_that_never_makes_a_request_is_closed() { #[tokio::test] async fn a_tls_connection_idle_after_its_last_request_is_closed() { let pki = pki(); - let mut limits = limits(); + let mut limits = FLEET_LIMITS; // Far past the test's bound, so only the idle close can end the connection. limits.first_request = Duration::from_secs(60); limits.idle_connection = Duration::from_millis(200); @@ -573,7 +578,7 @@ async fn connections_over_cap(limits: FleetLimits) { #[tokio::test] async fn connections_over_the_per_address_cap_are_closed() { - let mut limits = limits(); + let mut limits = FLEET_LIMITS; limits.connections_per_address = 2; connections_over_cap(limits).await; @@ -583,7 +588,7 @@ async fn connections_over_the_per_address_cap_are_closed() { fn an_address_over_its_cap_does_not_evict_other_connections() { use tokio::sync::oneshot::error::TryRecvError; - let mut limits = limits(); + let mut limits = FLEET_LIMITS; limits.max_connections = 2; limits.connections_per_address = 1; let connections = limits.connections(); @@ -604,7 +609,7 @@ fn an_address_over_its_cap_does_not_evict_other_connections() { #[tokio::test] async fn at_the_total_cap_the_longest_idle_connection_makes_room() { - let mut limits = limits(); + let mut limits = FLEET_LIMITS; limits.max_connections = 2; let address = serve( limits, @@ -628,9 +633,8 @@ async fn at_the_total_cap_the_longest_idle_connection_makes_room() { #[tokio::test] async fn at_the_total_cap_a_new_connection_is_closed_when_every_connection_is_busy() { - let mut limits = limits(); + let mut limits = FLEET_LIMITS; limits.max_connections = 1; - limits.request = Duration::from_secs(5); let address = serve( limits, None, @@ -654,7 +658,7 @@ async fn at_the_total_cap_a_new_connection_is_closed_when_every_connection_is_bu #[tokio::test] async fn a_connection_that_never_sends_anything_is_closed() { - let mut limits = limits(); + let mut limits = FLEET_LIMITS; limits.first_request = Duration::from_millis(200); let address = serve( limits, @@ -672,11 +676,8 @@ async fn a_connection_that_never_sends_anything_is_closed() { #[tokio::test] async fn a_connection_that_never_makes_a_request_is_closed() { - let mut limits = limits(); + let mut limits = FLEET_LIMITS; limits.first_request = Duration::from_millis(200); - // Keepalive alone would also close this connection; keep it out of the way. - limits.keepalive_interval = Duration::from_secs(60); - limits.keepalive_timeout = Duration::from_secs(60); let address = serve( limits, None, @@ -698,7 +699,7 @@ async fn a_connection_that_never_makes_a_request_is_closed() { #[tokio::test] async fn a_connection_idle_after_its_last_request_is_closed() { - let mut limits = limits(); + let mut limits = FLEET_LIMITS; // Far past the test's bound, so only the idle close can end the connection. limits.first_request = Duration::from_secs(60); limits.idle_connection = Duration::from_millis(200); @@ -726,7 +727,7 @@ async fn a_connection_idle_after_its_last_request_is_closed() { #[tokio::test] async fn a_response_the_client_never_accepts_does_not_keep_its_connection_busy() { - let mut limits = limits(); + let mut limits = FLEET_LIMITS; // Far past the test's bound, so only the idle close can end the connection. limits.first_request = Duration::from_secs(60); limits.idle_connection = Duration::from_millis(200); @@ -771,7 +772,7 @@ async fn a_response_the_client_never_accepts_does_not_keep_its_connection_busy() #[tokio::test] async fn a_stalled_connect_response_frees_its_connection_when_the_session_ends() { - let mut limits = limits(); + let mut limits = FLEET_LIMITS; // Far past the test's bound, so only the idle close can end the connection. limits.first_request = Duration::from_secs(60); limits.idle_connection = Duration::from_millis(200); @@ -807,8 +808,7 @@ async fn a_stalled_connect_response_frees_its_connection_when_the_session_ends() #[tokio::test] async fn an_open_stream_keeps_its_connection_from_going_idle() { - let mut limits = limits(); - limits.request = Duration::from_secs(5); + let mut limits = FLEET_LIMITS; limits.idle_connection = Duration::from_millis(200); // Messages arrive further apart than the idle limit. let address = serve( From 8a170a9d8cbd0751384712ea9a345c53f80034fb Mon Sep 17 00:00:00 2001 From: Scott Weiss Date: Tue, 10 Nov 2026 11:54:55 -0500 Subject: [PATCH 38/49] refactor(controller): build the fleet listener with every limit in one call Starting the fleet listener took six calls on FleetLimits, two of which had to share one connection registry, and any of them could be left out or mismatched without a compile error: separate registries closed every connection, open Connect streams included, 20 s after accept; serving with `serve(addr)` dropped the connection caps; skipping `tls` dropped the handshake timeout. The tests rebuilt this wiring themselves, so they would not have noticed. `FleetLimits::listener` now binds the listener with the TLS handshake timeout, the request and keepalive limits, the request layer and the connection registry, and `FleetListener::serve` adds the fleet service with its message limits. The tests serve through the same two calls. The defaults are no longer public: limits come from the controller configuration through `FleetLimits::from_config`, so code outside the module cannot pass the unconfigured defaults where the configured limits belong. A bind failure now names the fleet listener and its address. --- crates/controller/src/application.rs | 30 ++--- crates/controller/src/fleet_limits.rs | 117 ++++++++++++------ .../src/fleet_limits/connections.rs | 7 ++ crates/controller/src/fleet_limits/tests.rs | 25 ++-- crates/controller/src/service.rs | 13 +- 5 files changed, 116 insertions(+), 76 deletions(-) diff --git a/crates/controller/src/application.rs b/crates/controller/src/application.rs index 09971e72..ada2bcc7 100644 --- a/crates/controller/src/application.rs +++ b/crates/controller/src/application.rs @@ -2,7 +2,7 @@ use crate::{ daemon_config::{self, DaemonConfigStore}, database::Database, device_ca::DeviceCertificateIssuer, - fleet_limits::FLEET_LIMITS, + fleet_limits::FleetLimits, gateway_jwt::{self, GatewayJwtIssuer}, management::{self, Settings, State}, oidc::OidcProvider, @@ -103,7 +103,7 @@ pub async fn run( let device_certificate_issuer = DeviceCertificateIssuer::from_pem(ca_certificate, &ca_key) .context("initialize device certificate issuer")?; tracing::info!("device certificate issuance enabled"); - let limits = FLEET_LIMITS.configured(&config.fleet_limits); + let limits = FleetLimits::from_config(&config.fleet_limits); let service = FleetAgentService::new( Some(oidc), database, @@ -124,30 +124,22 @@ pub async fn run( tls.client_ca_certificate.display() ) })?; - let tls_config = limits.tls( - ServerTlsConfig::new() - .identity(Identity::from_pem(certificate, key)) - .client_ca_root(Certificate::from_pem(client_ca)) - .client_auth_optional(true), - ); - let fleet_connections = limits.connections(); - let mut server = limits - .server(&fleet_connections) - .accept_http1(true) - .tls_config(tls_config)?; - + let tls_config = ServerTlsConfig::new() + .identity(Identity::from_pem(certificate, key)) + .client_ca_root(Certificate::from_pem(client_ca)) + .client_auth_optional(true); let fleet_listen = config.fleet_listen; + let fleet_listener = limits.listener(fleet_listen, Some(tls_config))?; + let admin_listen = config.admin_listen; tracing::info!(listen = %fleet_listen, "fleet controller listening"); let fleet = async move { - let mut routes = Routes::from(gateway_jwt::routes(gateway_jwks)) - .add_service(limits.fleet_service(service)); + let mut routes = Routes::from(gateway_jwt::routes(gateway_jwks)); if let Some(usage_ingest) = usage_ingest { routes = routes.add_service(usage_ingest); } - server - .add_routes(routes) - .serve_with_incoming(limits.incoming(fleet_listen, &fleet_connections)?) + fleet_listener + .serve(service, routes) .await .context("serve fleet gRPC API") }; diff --git a/crates/controller/src/fleet_limits.rs b/crates/controller/src/fleet_limits.rs index cd314e80..d900edc1 100644 --- a/crates/controller/src/fleet_limits.rs +++ b/crates/controller/src/fleet_limits.rs @@ -7,7 +7,6 @@ use std::{ collections::HashMap, convert::Infallible, - io, net::{IpAddr, Ipv6Addr, SocketAddr}, sync::{ Arc, Mutex, @@ -19,6 +18,7 @@ use std::{ use agentdesktop_core::config::ControllerFleetLimitsConfig; use agentdesktop_proto::fleet::fleet_agent_server::{FleetAgent, FleetAgentServer, SERVICE_NAME}; +use anyhow::Context as _; use tokio::{ sync::{OwnedSemaphorePermit, Semaphore, oneshot}, time::Instant, @@ -27,6 +27,7 @@ use tonic::{ Status, codegen::{Service, http}, server::NamedService, + service::Routes, transport::{ Server, ServerTlsConfig, server::{TcpConnectInfo, TcpIncoming, TlsConnectInfo}, @@ -36,9 +37,8 @@ use tower_layer::{Identity, Stack}; mod connections; -pub use connections::{ - Connections, LimitedConnection, LimitedIncoming, ReleaseOnDrop, SessionBusy, TrackRequests, -}; +use connections::{Connections, LimitedIncoming, TrackRequests}; +pub use connections::{ReleaseOnDrop, SessionBusy}; /// Limits applied to the fleet gRPC listener. #[derive(Clone, Copy, Debug, PartialEq)] @@ -113,9 +113,9 @@ pub struct FleetLimits { pub stream_open_refill: Duration, } -/// Default limits. The controller configuration can change the -/// deployment-dependent ones (see `configured`). -pub const FLEET_LIMITS: FleetLimits = FleetLimits { +/// Default limits. The controller configuration replaces the +/// deployment-dependent ones (see `FleetLimits::from_config`). +const FLEET_LIMITS: FleetLimits = FleetLimits { tls_handshake: Duration::from_secs(10), request: Duration::from_secs(30), keepalive_interval: Duration::from_secs(30), @@ -140,9 +140,13 @@ pub const FLEET_LIMITS: FleetLimits = FleetLimits { }; impl FleetLimits { - /// These limits, with the deployment-dependent ones from the controller - /// configuration. - pub fn configured(mut self, config: &ControllerFleetLimitsConfig) -> Self { + /// The limits for a controller: the deployment-dependent ones from its + /// configuration, and the defaults for the rest. + pub fn from_config(config: &ControllerFleetLimitsConfig) -> Self { + FLEET_LIMITS.configured(config) + } + + fn configured(mut self, config: &ControllerFleetLimitsConfig) -> Self { self.max_connections = config.max_connections; self.connections_per_address = config.connections_per_address; self.max_message_bytes = config.max_message_bytes; @@ -154,43 +158,48 @@ impl FleetLimits { self } - /// Adds the handshake timeout to the fleet listener's TLS settings. - pub fn tls(&self, config: ServerTlsConfig) -> ServerTlsConfig { - config.timeout(self.tls_handshake) - } - - /// The registry that admits fleet connections and tracks their activity. - /// Pass the same one to `server` and `incoming`. - pub fn connections(&self) -> Connections { - Connections::new(self) - } - - /// A server builder with the fleet listener's limits applied. - pub fn server(&self, connections: &Connections) -> Server> { - Server::builder() + /// Binds the fleet listener on `address`, over TLS when `tls` is given, + /// with every connection and request limit applied. The listener serves + /// the fleet service with its message limits (`FleetListener::serve`). + pub fn listener( + &self, + address: SocketAddr, + tls: Option, + ) -> anyhow::Result { + // The listener and the request layer share one registry: the listener + // admits and closes connections, and the layer marks them busy. + let connections = Connections::new(self); + let mut server = Server::builder() .timeout(self.request) .http2_keepalive_interval(Some(self.keepalive_interval)) .http2_keepalive_timeout(Some(self.keepalive_timeout)) .concurrency_limit_per_connection(self.requests_per_connection) .max_concurrent_streams(self.streams_per_connection) - .layer(connections.track_requests()) - } - - /// Binds the fleet listener with TCP keepalive, admitting connections - /// through `connections`. Serve with `Server::serve_with_incoming`. - pub fn incoming( - &self, - address: SocketAddr, - connections: &Connections, - ) -> io::Result> { - let incoming = TcpIncoming::bind(address)? + // agentgateway fetches the gateway JWKS from this listener over HTTP/1. + .accept_http1(true) + .layer(connections.track_requests()); + if let Some(tls) = tls { + server = server + .tls_config(tls.timeout(self.tls_handshake)) + .context("configure fleet listener TLS")?; + } + let incoming = TcpIncoming::bind(address) + .with_context(|| format!("bind fleet listener on {address}"))? .with_nodelay(Some(true)) .with_keepalive(Some(self.tcp_keepalive)); - Ok(connections.incoming(incoming)) + let local_address = incoming + .local_addr() + .context("read fleet listener address")?; + Ok(FleetListener { + limits: *self, + server, + incoming: connections.incoming(incoming), + local_address, + }) } /// The fleet gRPC service with message size limits applied. - pub fn fleet_service(&self, service: T) -> FleetService { + fn fleet_service(&self, service: T) -> FleetService { let service = Arc::new(service); FleetService { unauthenticated: FleetAgentServer::from_arc(service.clone()) @@ -212,6 +221,40 @@ impl FleetLimits { } } +/// The bound fleet listener, from `FleetLimits::listener`. +pub struct FleetListener { + limits: FleetLimits, + server: Server>, + incoming: LimitedIncoming, + local_address: SocketAddr, +} + +impl FleetListener { + /// The address the listener is bound to. + pub fn local_address(&self) -> SocketAddr { + self.local_address + } + + /// Serves `service`, with the fleet message limits, and the other `routes` + /// until the listener fails. + pub async fn serve( + self, + service: T, + routes: Routes, + ) -> Result<(), tonic::transport::Error> { + let mut server = self.server; + server + .add_routes(routes.add_service(self.limits.fleet_service(service))) + .serve_with_incoming(self.incoming) + .await + } + + #[cfg(test)] + fn connections(&self) -> &Connections { + self.incoming.connections() + } +} + /// The fleet gRPC service, with a smaller request limit for enrollment and for /// callers without a device certificate. Tonic sets message limits per service, /// so those requests go to their own instance of the generated server. diff --git a/crates/controller/src/fleet_limits/connections.rs b/crates/controller/src/fleet_limits/connections.rs index e06d244b..dc3f26c6 100644 --- a/crates/controller/src/fleet_limits/connections.rs +++ b/crates/controller/src/fleet_limits/connections.rs @@ -313,6 +313,13 @@ pub struct LimitedIncoming { connections: Connections, } +#[cfg(test)] +impl LimitedIncoming { + pub(super) fn connections(&self) -> &Connections { + &self.connections + } +} + impl Stream for LimitedIncoming where S: Stream> + Unpin, diff --git a/crates/controller/src/fleet_limits/tests.rs b/crates/controller/src/fleet_limits/tests.rs index 7f8f298f..4e7dc6ff 100644 --- a/crates/controller/src/fleet_limits/tests.rs +++ b/crates/controller/src/fleet_limits/tests.rs @@ -10,14 +10,15 @@ use agentdesktop_proto::fleet::{ use futures_core::Stream; use tokio::{ io::{AsyncReadExt, AsyncWriteExt}, - net::{TcpListener, TcpStream}, + net::TcpStream, sync::mpsc, time, }; use tokio_stream::wrappers::ReceiverStream; use tonic::{ Code, Request, Response, Status, Streaming, - transport::{Identity, ServerTlsConfig, server::TcpIncoming}, + service::Routes, + transport::{Identity, ServerTlsConfig}, }; use super::{Connections, FLEET_LIMITS, FleetLimits}; @@ -111,17 +112,13 @@ async fn serve_tracked( tls: Option, service: Slow, ) -> (SocketAddr, Connections) { - let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); - let address = listener.local_addr().unwrap(); - let connections = limits.connections(); - // HTTP/1 is accepted as the controller does (`main.rs`), so connections go - // through hyper-util's protocol detection, as they do in production. - let mut server = limits.server(&connections).accept_http1(true); - if let Some(tls) = tls { - server = server.tls_config(limits.tls(tls)).unwrap(); - } - let router = server.add_service(limits.fleet_service(service)); - tokio::spawn(router.serve_with_incoming(connections.incoming(TcpIncoming::from(listener)))); + // The controller's own listener, so these tests exercise what it runs. + let listener = limits + .listener(SocketAddr::from(([127, 0, 0, 1], 0)), tls) + .unwrap(); + let address = listener.local_address(); + let connections = listener.connections().clone(); + tokio::spawn(listener.serve(service, Routes::default())); (address, connections) } @@ -591,7 +588,7 @@ fn an_address_over_its_cap_does_not_evict_other_connections() { let mut limits = FLEET_LIMITS; limits.max_connections = 2; limits.connections_per_address = 1; - let connections = limits.connections(); + let connections = Connections::new(&limits); let first_address = |port| SocketAddr::from(([192, 0, 2, 1], port)); let (_first, mut first_evicted) = connections .register(first_address(1)) diff --git a/crates/controller/src/service.rs b/crates/controller/src/service.rs index 9a1acef1..42910f5a 100644 --- a/crates/controller/src/service.rs +++ b/crates/controller/src/service.rs @@ -40,8 +40,9 @@ pub struct FleetAgentService { } impl FleetAgentService { - /// `limits` sets the enrollment and `Connect` limits; pass the configured - /// ones (`FLEET_LIMITS.configured(...)`), the same the fleet listener uses. + /// `limits` sets the enrollment and `Connect` limits; pass the ones from + /// the controller configuration (`FleetLimits::from_config`), the same the + /// fleet listener uses. pub fn new( oidc: Option, database: Database, @@ -712,7 +713,7 @@ mod tests { }; use tonic::{Code, Request}; - use crate::fleet_limits::{FLEET_LIMITS, FleetLimits}; + use crate::fleet_limits::FleetLimits; let path = std::env::temp_dir().join(format!( "agentdesktop-enrollment-limit-{}-{}.db", @@ -725,7 +726,7 @@ mod tests { let full = FleetLimits { enrollment_in_flight: 0, enrollment_completion_wait: std::time::Duration::from_millis(10), - ..FLEET_LIMITS + ..FleetLimits::from_config(&Default::default()) }; let service = FleetAgentService::new( None, @@ -771,7 +772,7 @@ mod tests { DaemonConfigStore::new(None), None, None, - &crate::fleet_limits::FLEET_LIMITS, + &crate::fleet_limits::FleetLimits::from_config(&Default::default()), ); let principal = |issuer: &str, subject: &str| OidcPrincipal { issuer: issuer.to_owned(), @@ -831,7 +832,7 @@ mod tests { DaemonConfigStore::new(None), None, None, - &crate::fleet_limits::FLEET_LIMITS, + &crate::fleet_limits::FleetLimits::from_config(&Default::default()), ); let status = service From 72f3efe8964950f68898a4000d1b8d6b197dc98c Mon Sep 17 00:00:00 2001 From: Scott Weiss Date: Tue, 10 Nov 2026 12:29:51 -0500 Subject: [PATCH 39/49] test(controller): fail when a fleet RPC has no limit class Enrollment and Connect were recognized by method name in two separate matches, so a new or renamed pre-certificate RPC would silently get neither the enrollment treatment nor a review of its limits. RPC_CLASSES now classifies every FleetAgent method as enrollment, device or Connect, and is_enrollment and is_connect read it. A test compares the table with the RPCs declared in fleet.proto and fails, naming the RPC, when one is missing, renamed or listed twice. --- crates/controller/src/fleet_limits.rs | 46 +++++++++++++--- crates/controller/src/fleet_limits/tests.rs | 59 ++++++++++++++++++++- 2 files changed, 98 insertions(+), 7 deletions(-) diff --git a/crates/controller/src/fleet_limits.rs b/crates/controller/src/fleet_limits.rs index d900edc1..dd8989d8 100644 --- a/crates/controller/src/fleet_limits.rs +++ b/crates/controller/src/fleet_limits.rs @@ -272,18 +272,52 @@ impl Clone for FleetService { } } +/// How the fleet limits treat a `FleetAgent` method. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +enum RpcClass { + /// Called before the device has a certificate. Its requests get the small + /// message limit whether or not a certificate is presented, and its handler + /// must take a permit from `EnrollmentLimiter` before any other work. + Enrollment, + /// Called with a device certificate: the device message limit. Without a + /// certificate it gets the small limit, and its handler rejects it. + Device, + /// The device's `Connect` stream. Its connection stays busy for as long as + /// the session lasts, not only while the handler runs (see `SessionBusy`). + Connect, +} + +/// Every `FleetAgent` method in `fleet.proto`, with its class. A test fails +/// when an RPC is missing here, so a new or renamed one cannot skip its limits. +const RPC_CLASSES: &[(&str, RpcClass)] = &[ + ("BeginEnrollment", RpcClass::Enrollment), + ("CompleteEnrollment", RpcClass::Enrollment), + ("RenewDeviceCertificate", RpcClass::Device), + ("GetLlmGatewayCredential", RpcClass::Device), + ("GetDeviceUsage", RpcClass::Device), + ("Connect", RpcClass::Connect), +]; + +/// The class of the `FleetAgent` method at `path`, or `None` for another path. +fn rpc_class(path: &str) -> Option { + let method = path + .strip_prefix('/')? + .strip_prefix(SERVICE_NAME)? + .strip_prefix('/')?; + RPC_CLASSES + .iter() + .find(|(name, _)| *name == method) + .map(|(_, class)| *class) +} + fn is_enrollment(path: &str) -> bool { - path.strip_prefix('/') - .and_then(|path| path.strip_prefix(SERVICE_NAME)) - .is_some_and(|method| matches!(method, "/BeginEnrollment" | "/CompleteEnrollment")) + rpc_class(path) == Some(RpcClass::Enrollment) } /// Whether `path` is the `Connect` stream, whose response lasts as long as the /// device's session. fn is_connect(path: &str) -> bool { - path.strip_prefix('/') - .and_then(|path| path.strip_prefix(SERVICE_NAME)) - .is_some_and(|method| method == "/Connect") + rpc_class(path) == Some(RpcClass::Connect) } /// Whether the TLS client presented a certificate signed by the device CA diff --git a/crates/controller/src/fleet_limits/tests.rs b/crates/controller/src/fleet_limits/tests.rs index 4e7dc6ff..e87419b8 100644 --- a/crates/controller/src/fleet_limits/tests.rs +++ b/crates/controller/src/fleet_limits/tests.rs @@ -21,7 +21,7 @@ use tonic::{ transport::{Identity, ServerTlsConfig}, }; -use super::{Connections, FLEET_LIMITS, FleetLimits}; +use super::{Connections, FLEET_LIMITS, FleetLimits, RPC_CLASSES, rpc_class}; /// Fails the test, instead of hanging it, when the server never ends the wait. async fn bounded(call: impl std::future::Future) -> T { @@ -873,6 +873,63 @@ fn configured_limits_replace_every_deployment_limit() { ); } +/// The RPC names in `fleet.proto`'s `FleetAgent` service. +fn fleet_rpcs(proto: &str) -> std::collections::BTreeSet<&str> { + let (_, service) = proto + .split_once("service FleetAgent {") + .expect("fleet.proto declares service FleetAgent"); + let (service, _) = service + .split_once('}') + .expect("service FleetAgent is closed"); + service + .lines() + .filter_map(|line| line.trim().strip_prefix("rpc ")) + .map(|rpc| rpc.split_once('(').map_or(rpc, |(name, _)| name).trim()) + .collect() +} + +#[test] +fn every_fleet_rpc_has_a_limit_class() { + let rpcs = fleet_rpcs(include_str!("../../../proto/proto/fleet.proto")); + let classified: std::collections::BTreeSet<&str> = + RPC_CLASSES.iter().map(|(name, _)| *name).collect(); + assert_eq!( + classified.len(), + RPC_CLASSES.len(), + "RPC_CLASSES lists an RPC twice" + ); + + let unclassified: Vec<_> = rpcs.difference(&classified).collect(); + let not_in_proto: Vec<_> = classified.difference(&rpcs).collect(); + assert!( + unclassified.is_empty() && not_in_proto.is_empty(), + "every FleetAgent RPC needs a class in RPC_CLASSES (fleet_limits.rs), which sets its \ + message limit and how long it keeps its connection busy; an Enrollment RPC's handler \ + must also take an EnrollmentLimiter permit. Not classified: {unclassified:?}. \ + Classified but not in fleet.proto: {not_in_proto:?}" + ); +} + +#[test] +fn fleet_paths_are_classified_by_the_table() { + use agentdesktop_proto::fleet::fleet_agent_server::SERVICE_NAME; + + for (name, class) in RPC_CLASSES { + assert_eq!( + rpc_class(&format!("/{SERVICE_NAME}/{name}")), + Some(*class), + "{name}" + ); + } + assert_eq!(rpc_class(&format!("/{SERVICE_NAME}/Unknown")), None); + assert_eq!(rpc_class("/other.Service/BeginEnrollment"), None); + assert!(super::is_enrollment(&format!( + "/{SERVICE_NAME}/CompleteEnrollment" + ))); + assert!(super::is_connect(&format!("/{SERVICE_NAME}/Connect"))); + assert!(!super::is_enrollment(&format!("/{SERVICE_NAME}/Connect"))); +} + #[test] fn a_repeated_warning_is_logged_once_per_interval_with_its_count() { let warning = super::ThrottledWarning::new(); From ba24b92e1f2a002dfb1248f134b78b8b5014a579 Mon Sep 17 00:00:00 2001 From: Scott Weiss Date: Tue, 10 Nov 2026 15:12:46 -0500 Subject: [PATCH 40/49] fix(controller): select the rustls provider before the fleet TLS listener starts tonic builds the fleet listener's TLS configuration from rustls's process-wide default provider and has no way to pass one. rustls picks that default by itself only when exactly one of its aws-lc-rs and ring backends is compiled in. A workspace-wide test build enables both (ring comes from the agent's test dependencies), so building the listener's TLS configuration panicked there, and the fleet TLS tests failed in CI. install_rustls_provider makes aws-lc-rs the default unless one is already set. The controller calls it at the start of application::run, and the fleet test server calls it too, so each TLS test works on its own whichever backends the build enables. --- crates/controller/src/application.rs | 3 +++ crates/controller/src/fleet_limits/tests.rs | 3 +++ crates/controller/src/lib.rs | 14 ++++++++++++++ 3 files changed, 20 insertions(+) diff --git a/crates/controller/src/application.rs b/crates/controller/src/application.rs index ada2bcc7..ee1a63be 100644 --- a/crates/controller/src/application.rs +++ b/crates/controller/src/application.rs @@ -25,6 +25,9 @@ pub async fn run( database: Database, extensions: management::Extensions, ) -> anyhow::Result<()> { + // Before any TLS configuration is built: the fleet listener's uses the + // process-wide rustls provider. + crate::install_rustls_provider(); // Fail before any network setup when the access mode has no provider. management::auth(&config.admin_access, &extensions)?; let tls = config.tls.files(); diff --git a/crates/controller/src/fleet_limits/tests.rs b/crates/controller/src/fleet_limits/tests.rs index e87419b8..f4135312 100644 --- a/crates/controller/src/fleet_limits/tests.rs +++ b/crates/controller/src/fleet_limits/tests.rs @@ -112,6 +112,9 @@ async fn serve_tracked( tls: Option, service: Slow, ) -> (SocketAddr, Connections) { + // As the controller does at startup, so TLS works whichever rustls + // backends the test build enables. + crate::install_rustls_provider(); // The controller's own listener, so these tests exercise what it runs. let listener = limits .listener(SocketAddr::from(([127, 0, 0, 1], 0)), tls) diff --git a/crates/controller/src/lib.rs b/crates/controller/src/lib.rs index 709fd868..82969132 100644 --- a/crates/controller/src/lib.rs +++ b/crates/controller/src/lib.rs @@ -8,3 +8,17 @@ pub mod management; pub mod oidc; pub mod service; pub mod usage; + +/// Makes aws-lc-rs rustls's process-wide crypto provider, unless one is +/// already installed. Call it before building any TLS configuration that uses +/// the process default. +/// +/// tonic builds the fleet listener's TLS configuration from the process default +/// and has no way to pass a provider. rustls picks the default on its own only +/// when exactly one of its `aws-lc-rs` and `ring` backends is compiled in, so a +/// build that also enables `ring` (a workspace-wide test build does) would +/// otherwise panic when the listener starts. +pub fn install_rustls_provider() { + // An error only means a provider is already installed. + let _ = rustls::crypto::aws_lc_rs::default_provider().install_default(); +} From d915119d740806aa5ae7be6dc297a9a00ffb04e9 Mon Sep 17 00:00:00 2001 From: Scott Weiss Date: Wed, 11 Nov 2026 12:22:26 -0500 Subject: [PATCH 41/49] fix(controller): limit the streams a client can make the fleet listener reset Tonic passes `None` for `http2_max_local_error_reset_streams` unless it is set, which overwrites hyper's and h2's default of 1024 and leaves no limit on the streams a client can make the server reset by breaking the protocol ("MadeYouReset"). Set it explicitly to that default. --- crates/controller/src/fleet_limits.rs | 9 ++++ crates/controller/src/fleet_limits/tests.rs | 49 +++++++++++++++++++++ 2 files changed, 58 insertions(+) diff --git a/crates/controller/src/fleet_limits.rs b/crates/controller/src/fleet_limits.rs index dd8989d8..00ee71bd 100644 --- a/crates/controller/src/fleet_limits.rs +++ b/crates/controller/src/fleet_limits.rs @@ -78,6 +78,13 @@ pub struct FleetLimits { /// cannot occupy an unbounded share of the server. pub requests_per_connection: usize, pub streams_per_connection: u32, + /// Streams the server may reset over a connection's lifetime because the + /// client broke the HTTP/2 protocol on them. Past this, the connection is + /// closed with `ENHANCE_YOUR_CALM`. A client can otherwise provoke resets + /// without end, which bypasses the protection against clients resetting + /// streams themselves ("MadeYouReset"). h2 and hyper default to 1024, but + /// tonic passes `None` when this is not set, which removes the limit. + pub error_resets_per_connection: usize, /// Largest encoded message an enrolled device's calls may send or receive. pub max_message_bytes: usize, /// Largest request from a caller without a device certificate: enrollment, @@ -127,6 +134,7 @@ const FLEET_LIMITS: FleetLimits = FleetLimits { idle_connection: Duration::from_secs(120), requests_per_connection: 32, streams_per_connection: 64, + error_resets_per_connection: 1024, max_message_bytes: 4 * 1024 * 1024, unauthenticated_message_bytes: 64 * 1024, enrollment_in_flight: 64, @@ -175,6 +183,7 @@ impl FleetLimits { .http2_keepalive_timeout(Some(self.keepalive_timeout)) .concurrency_limit_per_connection(self.requests_per_connection) .max_concurrent_streams(self.streams_per_connection) + .http2_max_local_error_reset_streams(Some(self.error_resets_per_connection)) // agentgateway fetches the gateway JWKS from this listener over HTTP/1. .accept_http1(true) .layer(connections.track_requests()); diff --git a/crates/controller/src/fleet_limits/tests.rs b/crates/controller/src/fleet_limits/tests.rs index f4135312..3a5e6dbe 100644 --- a/crates/controller/src/fleet_limits/tests.rs +++ b/crates/controller/src/fleet_limits/tests.rs @@ -241,6 +241,55 @@ async fn a_connection_that_stops_answering_pings_is_closed() { bounded(until_closed(stream)).await; } +/// The error code of the GOAWAY the server sends, or `None` if it closes the +/// connection without one. +async fn goaway_code(mut stream: TcpStream) -> Option { + loop { + let mut header = [0_u8; 9]; + stream.read_exact(&mut header).await.ok()?; + let length = u32::from_be_bytes([0, header[0], header[1], header[2]]) as usize; + let mut payload = vec![0_u8; length]; + stream.read_exact(&mut payload).await.ok()?; + // GOAWAY: last stream ID, then the error code. + if header[3] == 0x7 { + return Some(u32::from_be_bytes(payload[4..8].try_into().unwrap())); + } + } +} + +#[tokio::test] +async fn a_client_that_provokes_too_many_stream_resets_is_disconnected() { + let limits = FleetLimits { + error_resets_per_connection: 1, + ..FLEET_LIMITS + }; + let address = serve( + limits, + None, + Slow { + delay: Duration::ZERO, + }, + ) + .await; + let mut stream = TcpStream::connect(address).await.unwrap(); + + // The preface and an empty SETTINGS frame, then two PRIORITY frames that + // make streams 1 and 3 depend on themselves. Each is a stream error, so the + // server resets the stream. + let mut frames = + b"PRI * HTTP/2.0\r\n\r\nSM\r\n\r\n\x00\x00\x00\x04\x00\x00\x00\x00\x00".to_vec(); + for stream_id in [1_u32, 3] { + frames.extend_from_slice(&[0, 0, 5, 0x2, 0]); + frames.extend_from_slice(&stream_id.to_be_bytes()); + frames.extend_from_slice(&stream_id.to_be_bytes()); + frames.push(0); + } + stream.write_all(&frames).await.unwrap(); + + // ENHANCE_YOUR_CALM + assert_eq!(bounded(goaway_code(stream)).await, Some(0xb)); +} + /// How long two concurrent half-second calls take on one connection. async fn two_calls_on_one_connection(limits: FleetLimits) -> Duration { let delay = Duration::from_millis(500); From 03bd372e7cf21a019fcf55c687534dcc33e9aeac Mon Sep 17 00:00:00 2001 From: Scott Weiss Date: Wed, 11 Nov 2026 12:23:00 -0500 Subject: [PATCH 42/49] refactor(controller): keep one per-connection stream limit The per-connection request limit (32) only bounded handlers running at once, below the stream limit (64) that already bounds every open stream, including ones whose handler has returned. Drop it, rename the stream limit to say it counts concurrent streams, and document why it is set. --- crates/controller/src/fleet_limits.rs | 19 ++++++++++--------- crates/controller/src/fleet_limits/tests.rs | 17 +---------------- 2 files changed, 11 insertions(+), 25 deletions(-) diff --git a/crates/controller/src/fleet_limits.rs b/crates/controller/src/fleet_limits.rs index 00ee71bd..8864402b 100644 --- a/crates/controller/src/fleet_limits.rs +++ b/crates/controller/src/fleet_limits.rs @@ -73,11 +73,14 @@ pub struct FleetLimits { /// `Connect` stream) is never idle. pub first_request: Duration, pub idle_connection: Duration, - /// Concurrent requests one connection may have in progress, and concurrent - /// HTTP/2 streams it may open. Further requests wait, so one connection - /// cannot occupy an unbounded share of the server. - pub requests_per_connection: usize, - pub streams_per_connection: u32, + /// HTTP/2 streams one connection may have open at once + /// (`SETTINGS_MAX_CONCURRENT_STREAMS`). The client waits for one to close + /// before opening another, so one connection cannot occupy an unbounded + /// share of the server. Streams count until they close, not only while + /// their handler runs, so this also bounds open `Connect` streams and + /// responses the client does not read. hyper defaults to 200, but tonic + /// passes `None` when this is not set, which removes the limit. + pub concurrent_streams_per_connection: u32, /// Streams the server may reset over a connection's lifetime because the /// client broke the HTTP/2 protocol on them. Past this, the connection is /// closed with `ENHANCE_YOUR_CALM`. A client can otherwise provoke resets @@ -132,8 +135,7 @@ const FLEET_LIMITS: FleetLimits = FleetLimits { connections_per_address: 4_096, first_request: Duration::from_secs(20), idle_connection: Duration::from_secs(120), - requests_per_connection: 32, - streams_per_connection: 64, + concurrent_streams_per_connection: 64, error_resets_per_connection: 1024, max_message_bytes: 4 * 1024 * 1024, unauthenticated_message_bytes: 64 * 1024, @@ -181,8 +183,7 @@ impl FleetLimits { .timeout(self.request) .http2_keepalive_interval(Some(self.keepalive_interval)) .http2_keepalive_timeout(Some(self.keepalive_timeout)) - .concurrency_limit_per_connection(self.requests_per_connection) - .max_concurrent_streams(self.streams_per_connection) + .max_concurrent_streams(self.concurrent_streams_per_connection) .http2_max_local_error_reset_streams(Some(self.error_resets_per_connection)) // agentgateway fetches the gateway JWKS from this listener over HTTP/1. .accept_http1(true) diff --git a/crates/controller/src/fleet_limits/tests.rs b/crates/controller/src/fleet_limits/tests.rs index 3a5e6dbe..3454e0a6 100644 --- a/crates/controller/src/fleet_limits/tests.rs +++ b/crates/controller/src/fleet_limits/tests.rs @@ -312,25 +312,10 @@ async fn two_calls_on_one_connection(limits: FleetLimits) -> Duration { started.elapsed() } -#[tokio::test] -async fn requests_beyond_the_per_connection_limit_wait() { - let limits = FleetLimits { - requests_per_connection: 1, - ..FLEET_LIMITS - }; - - let elapsed = two_calls_on_one_connection(limits).await; - - assert!( - elapsed >= Duration::from_millis(950), - "ran concurrently: {elapsed:?}" - ); -} - #[tokio::test] async fn streams_beyond_the_per_connection_limit_wait() { let limits = FleetLimits { - streams_per_connection: 1, + concurrent_streams_per_connection: 1, ..FLEET_LIMITS }; From 381f8c961f29ac32fced0cb787b5a5a48113d777 Mon Sep 17 00:00:00 2001 From: Scott Weiss Date: Wed, 11 Nov 2026 12:23:25 -0500 Subject: [PATCH 43/49] refactor(controller): make the fleet service wrapper private Only FleetListener::serve builds it, so its generic stays out of the public API. --- crates/controller/src/fleet_limits.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/crates/controller/src/fleet_limits.rs b/crates/controller/src/fleet_limits.rs index 8864402b..b08a9d44 100644 --- a/crates/controller/src/fleet_limits.rs +++ b/crates/controller/src/fleet_limits.rs @@ -268,7 +268,7 @@ impl FleetListener { /// The fleet gRPC service, with a smaller request limit for enrollment and for /// callers without a device certificate. Tonic sets message limits per service, /// so those requests go to their own instance of the generated server. -pub struct FleetService { +struct FleetService { unauthenticated: FleetAgentServer, devices: FleetAgentServer, } From 3f4a88dad9a40cd7ba0908db335d7da1b501f691 Mon Sep 17 00:00:00 2001 From: Scott Weiss Date: Wed, 11 Nov 2026 12:23:50 -0500 Subject: [PATCH 44/49] refactor(controller): make token-bucket refill and take methods of Bucket --- crates/controller/src/fleet_limits.rs | 69 +++++++++++++-------------- 1 file changed, 33 insertions(+), 36 deletions(-) diff --git a/crates/controller/src/fleet_limits.rs b/crates/controller/src/fleet_limits.rs index b08a9d44..5a76d442 100644 --- a/crates/controller/src/fleet_limits.rs +++ b/crates/controller/src/fleet_limits.rs @@ -423,12 +423,38 @@ impl PeerBuckets { } } +/// A token bucket: up to `burst` tokens, refilled at one per `refill`. #[derive(Clone, Copy)] struct Bucket { tokens: f64, updated: Instant, } +impl Bucket { + fn full(burst: f64, now: Instant) -> Self { + Self { + tokens: burst, + updated: now, + } + } + + /// Refills the bucket for the time since its last use, then takes one token. + fn take(&mut self, now: Instant, refill: Duration, burst: f64) -> bool { + self.tokens = self.refilled(now, refill, burst); + self.updated = now; + if self.tokens < 1.0 { + return false; + } + self.tokens -= 1.0; + true + } + + fn refilled(self, now: Instant, refill: Duration, burst: f64) -> f64 { + let elapsed = now.saturating_duration_since(self.updated); + (self.tokens + elapsed.as_secs_f64() / refill.as_secs_f64()).min(burst) + } +} + impl EnrollmentLimiter { fn new(limits: &FleetLimits) -> Self { Self { @@ -482,14 +508,9 @@ impl EnrollmentLimiter { fn take(&self, peer: IpAddr, now: Instant) -> bool { let mut peers = self.peers.lock().expect("enrollment limiter lock poisoned"); - let bucket = peers.bucket( - address_key(peer), - Bucket { - tokens: self.burst, - updated: now, - }, - ); - take_token(bucket, now, self.refill, self.burst) + peers + .bucket(address_key(peer), Bucket::full(self.burst, now)) + .take(now, self.refill, self.burst) } } @@ -503,17 +524,6 @@ fn in_flight_refusal() -> Status { Status::resource_exhausted("too many enrollment requests in progress; retry later") } -/// Refills `bucket` for the time since its last use, then takes one token. -fn take_token(bucket: &mut Bucket, now: Instant, refill: Duration, burst: f64) -> bool { - bucket.tokens = refilled(*bucket, now, refill, burst); - bucket.updated = now; - if bucket.tokens < 1.0 { - return false; - } - bucket.tokens -= 1.0; - true -} - /// `Connect` state by device. A device has at most one open stream: opening a /// new stream tells the device's previous one to close, so a device that /// reconnects (after sleep, for example) is never blocked by a stream the @@ -571,14 +581,8 @@ impl DeviceStreams { .entry(device_id.to_owned()) .or_insert_with(|| DeviceState { current: None, - messages: Bucket { - tokens: f64::from(self.limits.device_message_burst), - updated: now, - }, - opens: Bucket { - tokens: f64::from(self.limits.stream_open_burst), - updated: now, - }, + messages: Bucket::full(f64::from(self.limits.device_message_burst), now), + opens: Bucket::full(f64::from(self.limits.stream_open_burst), now), }); f(device) } @@ -589,8 +593,7 @@ impl DeviceStreams { pub fn admit_open(&self, device_id: &str) -> Result<(), Status> { let limits = self.limits; let admitted = self.with_device(device_id, |device| { - take_token( - &mut device.opens, + device.opens.take( Instant::now(), limits.stream_open_refill, f64::from(limits.stream_open_burst), @@ -635,8 +638,7 @@ impl DeviceStream { pub fn take_message(&self, now: Instant) -> bool { let limits = self.streams.limits; self.streams.with_device(&self.device_id, |device| { - take_token( - &mut device.messages, + device.messages.take( now, limits.device_message_refill, f64::from(limits.device_message_burst), @@ -664,11 +666,6 @@ impl Drop for DeviceStream { } } -fn refilled(bucket: Bucket, now: Instant, refill: Duration, burst: f64) -> f64 { - let elapsed = now.saturating_duration_since(bucket.updated); - (bucket.tokens + elapsed.as_secs_f64() / refill.as_secs_f64()).min(burst) -} - /// A warning that fires once per refused request or connection. A flood would /// log a line per event, so this logs at most one per `WARNING_INTERVAL` and /// says how many events it stands for: every event since the last warning, from From bbed8a925828d88b0199fa0636c9351e33e3638f Mon Sep 17 00:00:00 2001 From: Scott Weiss Date: Wed, 11 Nov 2026 12:25:20 -0500 Subject: [PATCH 45/49] refactor(controller): embed the fleet limits configuration in FleetLimits FleetLimits copied each configurable limit from ControllerFleetLimitsConfig and repeated its default, and two tests kept the copies in sync. Hold the configuration as one field instead, so its defaults live only in the configuration (now a const, so FLEET_LIMITS stays a const), and drop the copy and both tests. --- crates/controller/src/fleet_limits.rs | 82 +++++-------------- .../src/fleet_limits/connections.rs | 4 +- crates/controller/src/fleet_limits/tests.rs | 79 ++++-------------- crates/controller/src/service.rs | 4 +- crates/core/src/config.rs | 27 +++--- 5 files changed, 59 insertions(+), 137 deletions(-) diff --git a/crates/controller/src/fleet_limits.rs b/crates/controller/src/fleet_limits.rs index 5a76d442..2f718cab 100644 --- a/crates/controller/src/fleet_limits.rs +++ b/crates/controller/src/fleet_limits.rs @@ -43,6 +43,9 @@ pub use connections::{ReleaseOnDrop, SessionBusy}; /// Limits applied to the fleet gRPC listener. #[derive(Clone, Copy, Debug, PartialEq)] pub struct FleetLimits { + /// The limits that depend on the deployment, set by the controller + /// configuration (`fleetLimits`), where each is documented. + pub configured: ControllerFleetLimitsConfig, /// Time a client has to finish the TLS handshake after connecting. Tonic sets /// no limit by default, so a client could hold a connection open forever by /// never sending its ClientHello. @@ -58,15 +61,6 @@ pub struct FleetLimits { /// How long a connection may be silent before TCP keepalive probes start, /// for peers that vanish without closing. pub tcp_keepalive: Duration, - /// Open connections allowed in total and from one client address (an IPv4 - /// address or an IPv6 /64). A new connection over the per-address cap is - /// closed as soon as it is accepted. At the total cap, the longest-idle - /// connection is closed to admit a new one, and the new one is closed only - /// when no connection is idle. The per-address cap is high because many - /// devices can share an address behind NAT, or appear to when a load - /// balancer rewrites client addresses. - pub max_connections: usize, - pub connections_per_address: usize, /// Time a new connection has to make its first request, and time a /// connection may stay without any request in progress after its last one /// finished. A connection with a request in progress (including an open @@ -88,34 +82,15 @@ pub struct FleetLimits { /// streams themselves ("MadeYouReset"). h2 and hyper default to 1024, but /// tonic passes `None` when this is not set, which removes the limit. pub error_resets_per_connection: usize, - /// Largest encoded message an enrolled device's calls may send or receive. - pub max_message_bytes: usize, /// Largest request from a caller without a device certificate: enrollment, /// which runs before a device has one, and any other call made without one /// (rejected by its handler, but only after the request has been decoded). /// Enrollment requests are small (a hostname and a CSR). pub unauthenticated_message_bytes: usize, - /// Enrollment requests (`BeginEnrollment`, `CompleteEnrollment`) running at - /// once across all clients. Enrollment calls the identity provider, so this - /// also bounds the load a flood can put on it. - pub enrollment_in_flight: usize, - /// Per-client-address `BeginEnrollment` budget: up to `enrollment_burst` - /// requests in a burst, refilled at one request per `enrollment_refill`. - /// Generous, because many devices can share one address behind NAT. - pub enrollment_burst: u32, - pub enrollment_refill: Duration, /// How long `CompleteEnrollment` waits for an in-flight slot before it is /// refused. It waits rather than failing at once because the user has /// already signed in, and a refusal makes them start over. pub enrollment_completion_wait: Duration, - /// Messages a device may send on its `Connect` streams: up to - /// `device_message_burst` in a burst, refilled at one per - /// `device_message_refill`. The budget belongs to the device, so - /// reconnecting does not reset it. Devices send a heartbeat every 30 - /// seconds, inventory periodically and one telemetry event per tool event, - /// so this is far above normal use; a stream over budget is closed. - pub device_message_burst: u32, - pub device_message_refill: Duration, /// How often a device may open a `Connect` stream: up to /// `stream_open_burst` in a burst, refilled at one per `stream_open_refill`. /// Each open checks the access token with the identity provider. @@ -123,28 +98,21 @@ pub struct FleetLimits { pub stream_open_refill: Duration, } -/// Default limits. The controller configuration replaces the -/// deployment-dependent ones (see `FleetLimits::from_config`). +/// Default limits. The controller configuration replaces `configured` (see +/// `FleetLimits::from_config`). const FLEET_LIMITS: FleetLimits = FleetLimits { + configured: ControllerFleetLimitsConfig::DEFAULT, tls_handshake: Duration::from_secs(10), request: Duration::from_secs(30), keepalive_interval: Duration::from_secs(30), keepalive_timeout: Duration::from_secs(20), tcp_keepalive: Duration::from_secs(60), - max_connections: 50_000, - connections_per_address: 4_096, first_request: Duration::from_secs(20), idle_connection: Duration::from_secs(120), concurrent_streams_per_connection: 64, error_resets_per_connection: 1024, - max_message_bytes: 4 * 1024 * 1024, unauthenticated_message_bytes: 64 * 1024, - enrollment_in_flight: 64, - enrollment_burst: 60, - enrollment_refill: Duration::from_secs(1), enrollment_completion_wait: Duration::from_secs(10), - device_message_burst: 600, - device_message_refill: Duration::from_millis(50), stream_open_burst: 20, stream_open_refill: Duration::from_secs(60), }; @@ -153,19 +121,10 @@ impl FleetLimits { /// The limits for a controller: the deployment-dependent ones from its /// configuration, and the defaults for the rest. pub fn from_config(config: &ControllerFleetLimitsConfig) -> Self { - FLEET_LIMITS.configured(config) - } - - fn configured(mut self, config: &ControllerFleetLimitsConfig) -> Self { - self.max_connections = config.max_connections; - self.connections_per_address = config.connections_per_address; - self.max_message_bytes = config.max_message_bytes; - self.enrollment_in_flight = config.enrollment_in_flight; - self.enrollment_burst = config.enrollment_burst; - self.enrollment_refill = config.enrollment_refill; - self.device_message_burst = config.device_message_burst; - self.device_message_refill = config.device_message_refill; - self + Self { + configured: *config, + ..FLEET_LIMITS + } } /// Binds the fleet listener on `address`, over TLS when `tls` is given, @@ -214,10 +173,10 @@ impl FleetLimits { FleetService { unauthenticated: FleetAgentServer::from_arc(service.clone()) .max_decoding_message_size(self.unauthenticated_message_bytes) - .max_encoding_message_size(self.max_message_bytes), + .max_encoding_message_size(self.configured.max_message_bytes), devices: FleetAgentServer::from_arc(service) - .max_decoding_message_size(self.max_message_bytes) - .max_encoding_message_size(self.max_message_bytes), + .max_decoding_message_size(self.configured.max_message_bytes) + .max_encoding_message_size(self.configured.max_message_bytes), } } @@ -458,10 +417,10 @@ impl Bucket { impl EnrollmentLimiter { fn new(limits: &FleetLimits) -> Self { Self { - in_flight: Arc::new(Semaphore::new(limits.enrollment_in_flight)), + in_flight: Arc::new(Semaphore::new(limits.configured.enrollment_in_flight)), peers: Arc::default(), - burst: f64::from(limits.enrollment_burst), - refill: limits.enrollment_refill, + burst: f64::from(limits.configured.enrollment_burst), + refill: limits.configured.enrollment_refill, completion_wait: limits.enrollment_completion_wait, } } @@ -572,7 +531,8 @@ impl DeviceStreams { .devices .lock() .expect("device stream registry lock poisoned"); - let max_devices = MAX_TRACKED_DEVICES.max(self.limits.max_connections.saturating_mul(2)); + let max_devices = + MAX_TRACKED_DEVICES.max(self.limits.configured.max_connections.saturating_mul(2)); if devices.len() >= max_devices && !devices.contains_key(device_id) { devices.retain(|_, device| device.current.is_some()); } @@ -581,7 +541,7 @@ impl DeviceStreams { .entry(device_id.to_owned()) .or_insert_with(|| DeviceState { current: None, - messages: Bucket::full(f64::from(self.limits.device_message_burst), now), + messages: Bucket::full(f64::from(self.limits.configured.device_message_burst), now), opens: Bucket::full(f64::from(self.limits.stream_open_burst), now), }); f(device) @@ -640,8 +600,8 @@ impl DeviceStream { self.streams.with_device(&self.device_id, |device| { device.messages.take( now, - limits.device_message_refill, - f64::from(limits.device_message_burst), + limits.configured.device_message_refill, + f64::from(limits.configured.device_message_burst), ) }) } diff --git a/crates/controller/src/fleet_limits/connections.rs b/crates/controller/src/fleet_limits/connections.rs index dc3f26c6..cfd06886 100644 --- a/crates/controller/src/fleet_limits/connections.rs +++ b/crates/controller/src/fleet_limits/connections.rs @@ -170,8 +170,8 @@ impl Connections { pub(super) fn new(limits: &FleetLimits) -> Self { Self { shared: Arc::new(Shared { - max_connections: limits.max_connections, - connections_per_address: limits.connections_per_address, + max_connections: limits.configured.max_connections, + connections_per_address: limits.configured.connections_per_address, first_request: limits.first_request, idle: limits.idle_connection, state: Mutex::default(), diff --git a/crates/controller/src/fleet_limits/tests.rs b/crates/controller/src/fleet_limits/tests.rs index 3454e0a6..4771deb7 100644 --- a/crates/controller/src/fleet_limits/tests.rs +++ b/crates/controller/src/fleet_limits/tests.rs @@ -328,11 +328,12 @@ async fn streams_beyond_the_per_connection_limit_wait() { } fn small_messages() -> FleetLimits { - FleetLimits { - max_message_bytes: 1024, + let mut limits = FleetLimits { unauthenticated_message_bytes: 128, ..FLEET_LIMITS - } + }; + limits.configured.max_message_bytes = 1024; + limits } #[tokio::test] @@ -613,7 +614,7 @@ async fn connections_over_cap(limits: FleetLimits) { #[tokio::test] async fn connections_over_the_per_address_cap_are_closed() { let mut limits = FLEET_LIMITS; - limits.connections_per_address = 2; + limits.configured.connections_per_address = 2; connections_over_cap(limits).await; } @@ -623,8 +624,8 @@ fn an_address_over_its_cap_does_not_evict_other_connections() { use tokio::sync::oneshot::error::TryRecvError; let mut limits = FLEET_LIMITS; - limits.max_connections = 2; - limits.connections_per_address = 1; + limits.configured.max_connections = 2; + limits.configured.connections_per_address = 1; let connections = Connections::new(&limits); let first_address = |port| SocketAddr::from(([192, 0, 2, 1], port)); let (_first, mut first_evicted) = connections @@ -644,7 +645,7 @@ fn an_address_over_its_cap_does_not_evict_other_connections() { #[tokio::test] async fn at_the_total_cap_the_longest_idle_connection_makes_room() { let mut limits = FLEET_LIMITS; - limits.max_connections = 2; + limits.configured.max_connections = 2; let address = serve( limits, None, @@ -668,7 +669,7 @@ async fn at_the_total_cap_the_longest_idle_connection_makes_room() { #[tokio::test] async fn at_the_total_cap_a_new_connection_is_closed_when_every_connection_is_busy() { let mut limits = FLEET_LIMITS; - limits.max_connections = 1; + limits.configured.max_connections = 1; let address = serve( limits, None, @@ -867,49 +868,6 @@ async fn an_open_stream_keeps_its_connection_from_going_idle() { } } -#[test] -fn the_configuration_defaults_match_the_default_limits() { - use agentdesktop_core::config::ControllerFleetLimitsConfig; - - assert_eq!( - FLEET_LIMITS.configured(&ControllerFleetLimitsConfig::default()), - FLEET_LIMITS - ); -} - -#[test] -fn configured_limits_replace_every_deployment_limit() { - use agentdesktop_core::config::ControllerFleetLimitsConfig; - - // The literal names every setting, each with a value that is not its - // default, so a setting `configured` does not copy fails here. - let config = ControllerFleetLimitsConfig { - max_connections: 8, - connections_per_address: 7, - max_message_bytes: 6, - enrollment_in_flight: 5, - enrollment_burst: 4, - enrollment_refill: Duration::from_millis(3), - device_message_burst: 2, - device_message_refill: Duration::from_millis(1), - }; - - assert_eq!( - FLEET_LIMITS.configured(&config), - FleetLimits { - max_connections: 8, - connections_per_address: 7, - max_message_bytes: 6, - enrollment_in_flight: 5, - enrollment_burst: 4, - enrollment_refill: Duration::from_millis(3), - device_message_burst: 2, - device_message_refill: Duration::from_millis(1), - ..FLEET_LIMITS - } - ); -} - /// The RPC names in `fleet.proto`'s `FleetAgent` service. fn fleet_rpcs(proto: &str) -> std::collections::BTreeSet<&str> { let (_, service) = proto @@ -988,12 +946,11 @@ mod enrollment { use super::Duration; fn limits(in_flight: usize, burst: u32) -> FleetLimits { - FleetLimits { - enrollment_in_flight: in_flight, - enrollment_burst: burst, - enrollment_refill: Duration::from_secs(10), - ..FLEET_LIMITS - } + let mut limits = FLEET_LIMITS; + limits.configured.enrollment_in_flight = in_flight; + limits.configured.enrollment_burst = burst; + limits.configured.enrollment_refill = Duration::from_secs(10); + limits } fn peer(last: u8) -> IpAddr { @@ -1227,8 +1184,8 @@ mod device_streams { #[test] fn a_device_gets_its_message_burst_then_the_refill_rate() { let mut limits = FLEET_LIMITS; - limits.device_message_burst = 3; - limits.device_message_refill = Duration::from_secs(1); + limits.configured.device_message_burst = 3; + limits.configured.device_message_refill = Duration::from_secs(1); let stream = limits.device_streams().open("device-a"); let start = Instant::now(); @@ -1243,8 +1200,8 @@ mod device_streams { #[test] fn reconnecting_does_not_reset_the_message_budget() { let mut limits = FLEET_LIMITS; - limits.device_message_burst = 2; - limits.device_message_refill = Duration::from_secs(3600); + limits.configured.device_message_burst = 2; + limits.configured.device_message_refill = Duration::from_secs(3600); let streams = limits.device_streams(); let now = Instant::now(); let first = streams.open("device-a"); diff --git a/crates/controller/src/service.rs b/crates/controller/src/service.rs index 42910f5a..433985f2 100644 --- a/crates/controller/src/service.rs +++ b/crates/controller/src/service.rs @@ -723,11 +723,11 @@ mod tests { let database = Database::connect(&format!("sqlite://{}?mode=rwc", path.display())) .await .expect("connect database"); - let full = FleetLimits { - enrollment_in_flight: 0, + let mut full = FleetLimits { enrollment_completion_wait: std::time::Duration::from_millis(10), ..FleetLimits::from_config(&Default::default()) }; + full.configured.enrollment_in_flight = 0; let service = FleetAgentService::new( None, database, diff --git a/crates/core/src/config.rs b/crates/core/src/config.rs index ec3063cb..1bb19938 100644 --- a/crates/core/src/config.rs +++ b/crates/core/src/config.rs @@ -578,7 +578,7 @@ pub enum ControllerAdminTransport { /// Limits on the device-facing fleet API that depend on the deployment: how /// many devices connect, how many share one client address, and how much they /// send. Unset fields keep their defaults. -#[derive(Clone, Debug, Deserialize, PartialEq, Serialize)] +#[derive(Clone, Copy, Debug, Deserialize, PartialEq, Serialize)] #[cfg_attr(feature = "schema", derive(schemars::JsonSchema))] #[serde(rename_all = "camelCase", default, deny_unknown_fields)] pub struct ControllerFleetLimitsConfig { @@ -623,18 +623,23 @@ pub struct ControllerFleetLimitsConfig { pub device_message_refill: Duration, } +impl ControllerFleetLimitsConfig { + /// The defaults, usable in constants. + pub const DEFAULT: Self = Self { + max_connections: 50_000, + connections_per_address: 4_096, + max_message_bytes: 4 * 1024 * 1024, + enrollment_in_flight: 64, + enrollment_burst: 60, + enrollment_refill: Duration::from_secs(1), + device_message_burst: 600, + device_message_refill: Duration::from_millis(50), + }; +} + impl Default for ControllerFleetLimitsConfig { fn default() -> Self { - Self { - max_connections: 50_000, - connections_per_address: 4_096, - max_message_bytes: 4 * 1024 * 1024, - enrollment_in_flight: 64, - enrollment_burst: 60, - enrollment_refill: Duration::from_secs(1), - device_message_burst: 600, - device_message_refill: Duration::from_millis(50), - } + Self::DEFAULT } } From df701e1e380109a947115728080f2804c2b08943 Mon Sep 17 00:00:00 2001 From: Scott Weiss Date: Wed, 11 Nov 2026 12:27:45 -0500 Subject: [PATCH 46/49] fix(controller): slow a device over its message budget instead of closing its stream A device that spent its Connect message budget had its stream closed. The agent sends one telemetry message per tool call, so a busy device could pass the budget, reconnect, spend it again from its queued telemetry, use up its stream-open budget, and report enrollment as failed meanwhile. Instead, stop reading the device's messages until the budget refills. HTTP/2 flow control then makes the device wait to send, so it stays connected and is still held to the same rate; the agent's own queue takes the backlog, and refuses new events once full. Config pushes, access-token revalidation and stream replacement keep running in the controller while reads wait. Throttled reads are logged with the throttled warnings. --- crates/controller/src/fleet_limits.rs | 45 +++++++++++++-- crates/controller/src/fleet_limits/tests.rs | 61 ++++++++++++++++++--- crates/controller/src/service.rs | 17 +++--- crates/core/src/config.rs | 5 +- schema/controller-config.json | 2 +- schema/controller-config.md | 2 +- 6 files changed, 106 insertions(+), 26 deletions(-) diff --git a/crates/controller/src/fleet_limits.rs b/crates/controller/src/fleet_limits.rs index 2f718cab..6fb8c760 100644 --- a/crates/controller/src/fleet_limits.rs +++ b/crates/controller/src/fleet_limits.rs @@ -382,6 +382,9 @@ impl PeerBuckets { } } +/// The longest `Bucket::spend` asks a caller to wait. +const LONGEST_WAIT: Duration = Duration::from_secs(24 * 60 * 60); + /// A token bucket: up to `burst` tokens, refilled at one per `refill`. #[derive(Clone, Copy)] struct Bucket { @@ -408,6 +411,21 @@ impl Bucket { true } + /// Refills the bucket, then spends one token even when none is left (the + /// bucket goes into debt). Returns when the next token is available: + /// `None` when one is left now. + fn spend(&mut self, now: Instant, refill: Duration, burst: f64) -> Option { + self.tokens = self.refilled(now, refill, burst) - 1.0; + self.updated = now; + if self.tokens >= 1.0 { + return None; + } + // Capped, so an absurdly long refill setting cannot overflow the instant. + let wait = Duration::try_from_secs_f64(refill.as_secs_f64() * (1.0 - self.tokens)) + .map_or(LONGEST_WAIT, |wait| wait.min(LONGEST_WAIT)); + Some(now + wait) + } + fn refilled(self, now: Instant, refill: Duration, burst: f64) -> f64 { let elapsed = now.saturating_duration_since(self.updated); (self.tokens + elapsed.as_secs_f64() / refill.as_secs_f64()).min(burst) @@ -594,16 +612,29 @@ impl DeviceStreams { } impl DeviceStream { - /// Takes one message from the device's budget, or returns false when it is spent. - pub fn take_message(&self, now: Instant) -> bool { + /// Spends one message from the device's budget. Returns when the budget + /// allows the next one: `None` when it allows it now, or the instant it + /// has refilled enough. Do not read another message from the device + /// before then. + pub fn spend_message(&self, now: Instant) -> Option { let limits = self.streams.limits; - self.streams.with_device(&self.device_id, |device| { - device.messages.take( + let next = self.streams.with_device(&self.device_id, |device| { + device.messages.spend( now, limits.configured.device_message_refill, f64::from(limits.configured.device_message_burst), ) - }) + }); + if next.is_some() + && let Some(throttled) = DEVICE_MESSAGES.event(now) + { + tracing::warn!( + latest_device_id = %self.device_id, + throttled_since_last_warning = throttled, + "slowed reading from devices that send messages faster than their budget" + ); + } + next } } @@ -626,7 +657,8 @@ impl Drop for DeviceStream { } } -/// A warning that fires once per refused request or connection. A flood would +/// A warning that fires once per refused or throttled request, connection or +/// message. A flood would /// log a line per event, so this logs at most one per `WARNING_INTERVAL` and /// says how many events it stands for: every event since the last warning, from /// any client or device, reported with the first event after the interval @@ -672,6 +704,7 @@ impl ThrottledWarning { static ENROLLMENT_RATE: ThrottledWarning = ThrottledWarning::new(); static ENROLLMENT_IN_FLIGHT: ThrottledWarning = ThrottledWarning::new(); static STREAM_OPENS: ThrottledWarning = ThrottledWarning::new(); +static DEVICE_MESSAGES: ThrottledWarning = ThrottledWarning::new(); #[cfg(test)] mod tests; diff --git a/crates/controller/src/fleet_limits/tests.rs b/crates/controller/src/fleet_limits/tests.rs index 4771deb7..364a763f 100644 --- a/crates/controller/src/fleet_limits/tests.rs +++ b/crates/controller/src/fleet_limits/tests.rs @@ -1188,13 +1188,53 @@ mod device_streams { limits.configured.device_message_refill = Duration::from_secs(1); let stream = limits.device_streams().open("device-a"); let start = Instant::now(); + let second = Duration::from_secs(1); - for _ in 0..3 { - assert!(stream.take_message(start)); - } - assert!(!stream.take_message(start), "burst exceeded"); - assert!(stream.take_message(start + Duration::from_secs(1))); - assert!(!stream.take_message(start + Duration::from_secs(1))); + assert_eq!(stream.spend_message(start), None); + assert_eq!(stream.spend_message(start), None); + assert_eq!( + stream.spend_message(start), + Some(start + second), + "the burst is spent" + ); + assert_eq!( + stream.spend_message(start + second), + Some(start + 2 * second), + "one message per refill" + ); + } + + #[test] + fn a_message_read_before_the_budget_refills_delays_the_next_one() { + let mut limits = FLEET_LIMITS; + limits.configured.device_message_burst = 1; + limits.configured.device_message_refill = Duration::from_secs(1); + let stream = limits.device_streams().open("device-a"); + let start = Instant::now(); + + assert_eq!( + stream.spend_message(start), + Some(start + Duration::from_secs(1)) + ); + + assert_eq!( + stream.spend_message(start), + Some(start + Duration::from_secs(2)) + ); + } + + #[test] + fn an_absurdly_long_refill_caps_the_wait() { + let mut limits = FLEET_LIMITS; + limits.configured.device_message_burst = 1; + limits.configured.device_message_refill = Duration::MAX; + let stream = limits.device_streams().open("device-a"); + let start = Instant::now(); + + assert_eq!( + stream.spend_message(start), + Some(start + super::super::LONGEST_WAIT) + ); } #[test] @@ -1205,13 +1245,16 @@ mod device_streams { let streams = limits.device_streams(); let now = Instant::now(); let first = streams.open("device-a"); - assert!(first.take_message(now)); - assert!(first.take_message(now)); + assert_eq!(first.spend_message(now), None); + assert!(first.spend_message(now).is_some(), "the burst is spent"); drop(first); let second = streams.open("device-a"); - assert!(!second.take_message(now), "a new stream got a fresh budget"); + assert!( + second.spend_message(now).is_some(), + "a new stream got a fresh budget" + ); } #[test] diff --git a/crates/controller/src/service.rs b/crates/controller/src/service.rs index 433985f2..d19ba1e9 100644 --- a/crates/controller/src/service.rs +++ b/crates/controller/src/service.rs @@ -306,6 +306,11 @@ impl FleetAgent for FleetAgentService { }; let mut auth_interval = time::interval(std::time::Duration::from_secs(5 * 60)); auth_interval.tick().await; + // Once the device has spent its message budget, its next message + // is not read until the budget refills. HTTP/2 flow control then + // makes the device wait to send, so a busy device is slowed down + // rather than disconnected, while the other branches keep running. + let mut next_read: Option = None; loop { tokio::select! { _ = auth_sender.closed() => break, @@ -358,14 +363,12 @@ impl FleetAgent for FleetAgentService { .try_send(Err(Status::aborted("replaced by a newer stream from this device"))); break; } - message = inbound.message() => match message { + _ = time::sleep_until(next_read.unwrap_or_else(time::Instant::now)), if next_read.is_some() => { + next_read = None; + } + message = inbound.message(), if next_read.is_none() => match message { Ok(Some(message)) => { - if !device_stream.take_message(time::Instant::now()) { - warn!(device_id = %authenticated_device_id, "closing stream after the device exceeded its message rate"); - let _ = auth_sender - .try_send(Err(Status::resource_exhausted("too many messages on this stream"))); - break; - } + next_read = device_stream.spend_message(time::Instant::now()); if let Some(agent_message::Message::Hello(hello)) = &message.message && hello.device_id != authenticated_device_id { diff --git a/crates/core/src/config.rs b/crates/core/src/config.rs index 1bb19938..3af35579 100644 --- a/crates/core/src/config.rs +++ b/crates/core/src/config.rs @@ -613,8 +613,9 @@ pub struct ControllerFleetLimitsConfig { #[cfg_attr(feature = "schema", schemars(with = "String"))] pub enrollment_refill: Duration, /// Messages one device may send in a burst on its `Connect` stream, after - /// which it may send one more every `deviceMessageRefill`. A device over - /// its budget is disconnected. Defaults to `600`. + /// which it may send one more every `deviceMessageRefill`. Once a device + /// has sent its burst, the controller reads its messages no faster than + /// that, and the device waits to send more. Defaults to `600`. #[cfg_attr(feature = "schema", schemars(range(min = 1)))] pub device_message_burst: u32, /// Time for a device to get back one message. Defaults to `50ms`. diff --git a/schema/controller-config.json b/schema/controller-config.json index 646722dd..967281a4 100644 --- a/schema/controller-config.json +++ b/schema/controller-config.json @@ -217,7 +217,7 @@ "minimum": 1 }, "deviceMessageBurst": { - "description": "Messages one device may send in a burst on its `Connect` stream, after\nwhich it may send one more every `deviceMessageRefill`. A device over\nits budget is disconnected. Defaults to `600`.", + "description": "Messages one device may send in a burst on its `Connect` stream, after\nwhich it may send one more every `deviceMessageRefill`. Once a device\nhas sent its burst, the controller reads its messages no faster than\nthat, and the device waits to send more. Defaults to `600`.", "type": "integer", "format": "uint32", "default": 600, diff --git a/schema/controller-config.md b/schema/controller-config.md index 8c835e00..f9ee69ee 100644 --- a/schema/controller-config.md +++ b/schema/controller-config.md @@ -20,7 +20,7 @@ |`databaseUrl`|string|SQLite or PostgreSQL URL used for controller state.| |`fleetLimits`|object|Limits on the device-facing fleet API that depend on the deployment.| |`fleetLimits.connectionsPerAddress`|integer|Open fleet connections allowed from one client address: an IPv4
address, or an IPv6 /64. Raise it when more devices share an address,
such as behind a NAT or a load balancer that does not preserve client
addresses. Defaults to `4096`.| -|`fleetLimits.deviceMessageBurst`|integer|Messages one device may send in a burst on its `Connect` stream, after
which it may send one more every `deviceMessageRefill`. A device over
its budget is disconnected. Defaults to `600`.| +|`fleetLimits.deviceMessageBurst`|integer|Messages one device may send in a burst on its `Connect` stream, after
which it may send one more every `deviceMessageRefill`. Once a device
has sent its burst, the controller reads its messages no faster than
that, and the device waits to send more. Defaults to `600`.| |`fleetLimits.deviceMessageRefill`|string|Time for a device to get back one message. Defaults to `50ms`.| |`fleetLimits.enrollmentBurst`|integer|Enrollments one client address may start in a burst, after which it may
start one more every `enrollmentRefill`. Defaults to `60`.| |`fleetLimits.enrollmentInFlight`|integer|Enrollment requests in progress at once across all clients. Enrollment
calls the identity provider. Defaults to `64`.| From 0bbe4d94e062087f90bd8807b5c5bbb687196ff3 Mon Sep 17 00:00:00 2001 From: Scott Weiss Date: Thu, 12 Nov 2026 12:58:25 -0500 Subject: [PATCH 47/49] fix(controller): check fleet connection deadlines on writes too A client that sends HTTP/2 PINGs without reading the replies fills the server's send buffer. h2 flushes each PONG before it reads the next frame, so once the socket's send side is full the connection stops reading. The first-request deadline, the idle deadline and eviction were checked only on reads, so over plain TCP none of them fired: the connection held its slot, and an eviction aimed at it pushed the open count past maxConnections. Check them on writes and flushes as well, so closing a connection no longer depends on the layers above reading from it. In testing, a stalled TLS connection still closed at its first-request deadline without this change. Shutdown stays unchecked so a connection being closed can still shut down cleanly. The new tests flood PINGs without reading, against the first-request deadline and at the total connection cap. Both fail without this change. --- .../src/fleet_limits/connections.rs | 14 ++- crates/controller/src/fleet_limits/tests.rs | 102 ++++++++++++++++-- 2 files changed, 105 insertions(+), 11 deletions(-) diff --git a/crates/controller/src/fleet_limits/connections.rs b/crates/controller/src/fleet_limits/connections.rs index cfd06886..5fae48ed 100644 --- a/crates/controller/src/fleet_limits/connections.rs +++ b/crates/controller/src/fleet_limits/connections.rs @@ -344,7 +344,9 @@ where } /// An accepted fleet connection. It holds its slots until it closes, and fails -/// its next read once it is evicted or has been idle too long. +/// its next read, write or flush once it is evicted or has been idle too long. +/// Writes are checked too because a client that never reads fills the server's +/// send buffer, and HTTP/2 then stops reading until it can write. pub struct LimitedConnection { stream: TcpStream, peer: SocketAddr, @@ -454,6 +456,9 @@ impl AsyncWrite for LimitedConnection { context: &mut Context<'_>, buffer: &[u8], ) -> Poll> { + if let Err(error) = self.check_open(context) { + return Poll::Ready(Err(error)); + } Pin::new(&mut self.stream).poll_write(context, buffer) } @@ -462,6 +467,9 @@ impl AsyncWrite for LimitedConnection { context: &mut Context<'_>, buffers: &[io::IoSlice<'_>], ) -> Poll> { + if let Err(error) = self.check_open(context) { + return Poll::Ready(Err(error)); + } Pin::new(&mut self.stream).poll_write_vectored(context, buffers) } @@ -470,9 +478,13 @@ impl AsyncWrite for LimitedConnection { } fn poll_flush(mut self: Pin<&mut Self>, context: &mut Context<'_>) -> Poll> { + if let Err(error) = self.check_open(context) { + return Poll::Ready(Err(error)); + } Pin::new(&mut self.stream).poll_flush(context) } + // Unchecked, so a connection being closed can still shut down cleanly. fn poll_shutdown(mut self: Pin<&mut Self>, context: &mut Context<'_>) -> Poll> { Pin::new(&mut self.stream).poll_shutdown(context) } diff --git a/crates/controller/src/fleet_limits/tests.rs b/crates/controller/src/fleet_limits/tests.rs index 364a763f..32ffb475 100644 --- a/crates/controller/src/fleet_limits/tests.rs +++ b/crates/controller/src/fleet_limits/tests.rs @@ -233,10 +233,7 @@ async fn a_connection_that_stops_answering_pings_is_closed() { // Open an HTTP/2 connection (preface and an empty SETTINGS frame), then go // silent: never acknowledge the server's SETTINGS or PINGs. - stream - .write_all(b"PRI * HTTP/2.0\r\n\r\nSM\r\n\r\n\x00\x00\x00\x04\x00\x00\x00\x00\x00") - .await - .unwrap(); + stream.write_all(PREFACE).await.unwrap(); bounded(until_closed(stream)).await; } @@ -276,8 +273,7 @@ async fn a_client_that_provokes_too_many_stream_resets_is_disconnected() { // The preface and an empty SETTINGS frame, then two PRIORITY frames that // make streams 1 and 3 depend on themselves. Each is a stream error, so the // server resets the stream. - let mut frames = - b"PRI * HTTP/2.0\r\n\r\nSM\r\n\r\n\x00\x00\x00\x04\x00\x00\x00\x00\x00".to_vec(); + let mut frames = PREFACE.to_vec(); for stream_id in [1_u32, 3] { frames.extend_from_slice(&[0, 0, 5, 0x2, 0]); frames.extend_from_slice(&stream_id.to_be_bytes()); @@ -724,10 +720,7 @@ async fn a_connection_that_never_makes_a_request_is_closed() { let mut stream = TcpStream::connect(address).await.unwrap(); // A valid HTTP/2 preface and SETTINGS frame, then no request. - stream - .write_all(b"PRI * HTTP/2.0\r\n\r\nSM\r\n\r\n\x00\x00\x00\x04\x00\x00\x00\x00\x00") - .await - .unwrap(); + stream.write_all(PREFACE).await.unwrap(); bounded(until_closed(stream)).await; } @@ -868,6 +861,95 @@ async fn an_open_stream_keeps_its_connection_from_going_idle() { } } +/// The HTTP/2 client preface and an empty SETTINGS frame. +const PREFACE: &[u8] = b"PRI * HTTP/2.0\r\n\r\nSM\r\n\r\n\x00\x00\x00\x04\x00\x00\x00\x00\x00"; + +/// An HTTP/2 PING frame: length 8, type 6, no flags, stream 0, then its payload. +const PING: &[u8] = &[0, 0, 8, 6, 0, 0, 0, 0, 0, 1, 2, 3, 4, 5, 6, 7, 8]; + +/// Batches of PINGs sent before a flood gives up on the server stalling. +const FLOOD_BATCHES: usize = 1000; + +/// Sends the HTTP/2 preface, then PINGs without reading the server's replies, +/// until a write stalls or fails. The unread replies fill the server's send +/// buffer, and HTTP/2 stops reading until it can write them. +async fn flood_pings(stream: &mut TcpStream) { + stream.write_all(PREFACE).await.unwrap(); + let pings = PING.repeat(4096); + for _ in 0..FLOOD_BATCHES { + match time::timeout(Duration::from_millis(500), stream.write_all(&pings)).await { + Ok(Ok(())) => {} + Ok(Err(_)) | Err(_) => return, + } + } + panic!("the server read every PING without stalling or closing the connection"); +} + +/// Connects with a small receive buffer, so the server's unread replies fill +/// its send buffer quickly. +async fn connect_small_window(address: SocketAddr) -> TcpStream { + let socket = tokio::net::TcpSocket::new_v4().unwrap(); + socket.set_recv_buffer_size(4096).unwrap(); + socket.connect(address).await.unwrap() +} + +/// Waits until at most `count` connections are open. +async fn until_open_at_most(connections: &Connections, count: usize) { + while connections.open_count() > count { + time::sleep(Duration::from_millis(50)).await; + } +} + +#[tokio::test] +async fn a_connection_that_floods_pings_without_reading_is_closed() { + let mut limits = FLEET_LIMITS; + // Long enough that the server has stopped reading when it passes. + limits.first_request = Duration::from_secs(2); + let (address, connections) = serve_tracked( + limits, + None, + Slow { + delay: Duration::ZERO, + }, + ) + .await; + let mut stream = connect_small_window(address).await; + + flood_pings(&mut stream).await; + + bounded(until_open_at_most(&connections, 0)).await; +} + +#[tokio::test] +async fn an_evicted_connection_that_floods_pings_is_closed() { + let mut limits = FLEET_LIMITS; + limits.configured.max_connections = 1; + // Far past the test's bound, so only eviction can end the flooding connection. + limits.first_request = Duration::from_secs(60); + let (address, connections) = serve_tracked( + limits, + None, + Slow { + delay: Duration::ZERO, + }, + ) + .await; + let mut flooder = connect_small_window(address).await; + flood_pings(&mut flooder).await; + + // At the cap, the new connection evicts the flooder, which has made no + // request. The server answers its preface once it is admitted. + let mut newcomer = TcpStream::connect(address).await.unwrap(); + newcomer.write_all(PREFACE).await.unwrap(); + let mut settings = [0_u8; 9]; + bounded(newcomer.read_exact(&mut settings)) + .await + .expect("the new connection was refused"); + + bounded(until_open_at_most(&connections, 1)).await; + assert!(stays_open(newcomer).await, "the new connection was closed"); +} + /// The RPC names in `fleet.proto`'s `FleetAgent` service. fn fleet_rpcs(proto: &str) -> std::collections::BTreeSet<&str> { let (_, service) = proto From 025ac7a324144671e47ee4be9f9fcbea130b0c1f Mon Sep 17 00:00:00 2001 From: Scott Weiss Date: Thu, 12 Nov 2026 16:45:25 -0500 Subject: [PATCH 48/49] fix(controller): only device requests keep a fleet connection busy A client without a certificate could send a request's HEADERS and never its body. tonic waits for the body before the handler runs, so the request kept its connection busy, safe from the idle deadline and from eviction, until the 30 s request timeout, and the client could open the next one before then. At the defaults, 13 addresses could hold every connection. Count a request as busy only when the caller has a device certificate. Other requests still count as activity, and their connections stay under the caps, the deadlines and the request timeout, but they can be evicted. Connect requires a device certificate, so device sessions stay protected. At the per-address cap, close that address's longest-idle connection instead of refusing the new one, as the total cap already does. Otherwise an anonymous client that shares an address with devices, such as a node address under the Cluster traffic policy, could fill the address's slots and lock those devices out. Tests that used a plain-TCP Connect as the busy connection now use a device certificate. --- .../src/fleet_limits/connections.rs | 223 ++++++++----- crates/controller/src/fleet_limits/tests.rs | 301 ++++++++++++++---- crates/core/src/config.rs | 12 +- schema/controller-config.json | 4 +- schema/controller-config.md | 4 +- 5 files changed, 391 insertions(+), 153 deletions(-) diff --git a/crates/controller/src/fleet_limits/connections.rs b/crates/controller/src/fleet_limits/connections.rs index 5fae48ed..f14eea48 100644 --- a/crates/controller/src/fleet_limits/connections.rs +++ b/crates/controller/src/fleet_limits/connections.rs @@ -8,16 +8,25 @@ //! //! So each connection's request activity is tracked. A connection is closed if //! it makes no request soon after it is accepted, or stays idle too long after -//! its last request finishes. A connection with a request in progress, including -//! an open `Connect` stream, is never idle. When the total cap is reached, the -//! longest-idle connection is closed to make room instead of refusing the new -//! one, so idle connections cannot crowd out devices that are using theirs. +//! its last request finishes. A connection with a device's request in progress, +//! including an open `Connect` stream, is busy: never idle. When a cap is +//! reached, the longest-idle connection is closed to make room instead of +//! refusing the new one: at the total cap any connection, and at the +//! per-address cap one from the same address. So idle connections cannot crowd +//! out devices that are using theirs. //! -//! A request other than `Connect` is in progress only while its handler runs, -//! not while its response is sent, and a `Connect` stream only while its -//! session lasts. Otherwise a client could keep a connection busy forever by -//! granting no HTTP/2 flow-control window, so that the server can never send -//! the response body and hyper never drops it. +//! Only a request from a caller with a device certificate makes its connection +//! busy. Anyone can open a request and never send its body, and tonic waits for +//! the body before the handler runs, so an anonymous request would otherwise +//! hold its connection busy until the request timeout, and the next one after +//! it. An anonymous request still counts as activity, and its connection stays +//! subject to the caps, deadlines and request timeout, but it can be evicted. +//! +//! A device's request other than `Connect` is in progress only while its +//! handler runs, not while its response is sent, and a `Connect` stream only +//! while its session lasts. Otherwise a client could keep a connection busy +//! forever by granting no HTTP/2 flow-control window, so that the server can +//! never send the response body and hyper never drops it. use std::{ collections::{HashMap, VecDeque}, @@ -48,7 +57,7 @@ use tonic::{ }; use tower_layer::Layer; -use super::{FleetLimits, ThrottledWarning, address_key, is_connect}; +use super::{FleetLimits, ThrottledWarning, address_key, has_device_certificate, is_connect}; static TOTAL_CAP: ThrottledWarning = ThrottledWarning::new(); static PER_ADDRESS_CAP: ThrottledWarning = ThrottledWarning::new(); @@ -74,10 +83,58 @@ struct Shared { struct State { open: HashMap>, /// Open connections by `address_key`. - per_address: HashMap, - /// Connections in the order they became idle. An entry is stale once its - /// connection makes a request or closes; `Activity::idle_epoch` tells. - idle: VecDeque<(SocketAddr, u64)>, + per_address: HashMap, + /// Connections in the order they became idle. + idle: IdleQueue, +} + +/// One address key's open connections. +#[derive(Default)] +struct Address { + count: usize, + /// This address's connections in the order they became idle. + idle: IdleQueue, +} + +/// Connections in the order they became idle. An entry is stale once its +/// connection starts a device request or closes; `Activity::idle_epoch` tells. +#[derive(Default)] +struct IdleQueue(VecDeque<(SocketAddr, u64)>); + +impl IdleQueue { + /// Queues `peer`, dropping stale entries once they dominate, so the queue + /// stays bounded by its `open` connections. Each address has a queue, so + /// the slack is small. + fn push( + &mut self, + peer: SocketAddr, + epoch: u64, + open: usize, + connections: &HashMap>, + ) { + self.0.push_back((peer, epoch)); + if self.0.len() > 2 * open + 8 { + self.0.retain(|(peer, epoch)| { + connections + .get(peer) + .is_some_and(|activity| activity.is_idle_at(*epoch)) + }); + } + } + + /// Closes the longest-idle connection in the queue. Returns false when + /// none is idle. + fn evict(&mut self, connections: &HashMap>) -> bool { + while let Some((peer, epoch)) = self.0.pop_front() { + if connections + .get(&peer) + .is_some_and(|activity| activity.is_idle_at(epoch) && activity.evict()) + { + return true; + } + } + false + } } /// One connection's request activity. @@ -122,46 +179,46 @@ impl Activity { + allowed, ) } + + /// Whether the connection is idle and has been since it was queued at `epoch`. + fn is_idle_at(&self, epoch: u64) -> bool { + self.in_flight.load(Ordering::Acquire) == 0 + && self.idle_epoch.load(Ordering::Acquire) == epoch + } + + /// Tells the connection to close. Returns false if it already was told. + fn evict(&self) -> bool { + self.evict + .lock() + .expect("connection eviction lock poisoned") + .take() + .is_some_and(|evict| { + let _ = evict.send(()); + true + }) + } } impl State { /// Closes the longest-idle connection. Returns false when every open - /// connection has a request in progress. + /// connection has a device request in progress. fn evict_idle(&mut self) -> bool { - while let Some((peer, epoch)) = self.idle.pop_front() { - let Some(activity) = self.open.get(&peer) else { - continue; - }; - if activity.in_flight.load(Ordering::Acquire) > 0 - || activity.idle_epoch.load(Ordering::Acquire) != epoch - { - continue; - } - if let Some(evict) = activity - .evict - .lock() - .expect("connection eviction lock poisoned") - .take() - { - let _ = evict.send(()); - return true; - } - } - false + self.idle.evict(&self.open) + } + + /// Closes `address`'s longest-idle connection. Returns false when each of + /// its connections has a device request in progress. + fn evict_idle_from(&mut self, address: IpAddr) -> bool { + self.per_address + .get_mut(&address) + .is_some_and(|connections| connections.idle.evict(&self.open)) } fn mark_idle(&mut self, peer: SocketAddr, epoch: u64) { - self.idle.push_back((peer, epoch)); - // Drop stale entries once they dominate, so the queue stays bounded by - // the number of open connections. - if self.idle.len() > 2 * self.open.len() + 64 { - let open = &self.open; - self.idle.retain(|(peer, epoch)| { - open.get(peer).is_some_and(|activity| { - activity.in_flight.load(Ordering::Acquire) == 0 - && activity.idle_epoch.load(Ordering::Acquire) == *epoch - }) - }); + let open = self.open.len(); + self.idle.push(peer, epoch, open, &self.open); + if let Some(address) = self.per_address.get_mut(&address_key(peer.ip())) { + address.idle.push(peer, epoch, address.count, &self.open); } } } @@ -194,7 +251,8 @@ impl Connections { } } - /// A layer that marks a connection busy while a request (or stream) is open. + /// A layer that records each request on its connection, and marks the + /// connection busy while a device's request (or stream) is open. pub fn track_requests(&self) -> TrackRequests { TrackRequests { connections: self.clone(), @@ -224,24 +282,26 @@ impl Connections { peer: SocketAddr, ) -> Option<(Arc, oneshot::Receiver<()>)> { let mut state = self.state(); - // The per-address cap comes first, so an address over it cannot close - // another client's idle connection only to be refused itself. let address = address_key(peer.ip()); - if state + let at_address_cap = state .per_address .get(&address) - .is_some_and(|count| *count >= self.shared.connections_per_address) - { - if let Some(closed) = PER_ADDRESS_CAP.event(Instant::now()) { - tracing::warn!( - latest_address = %address, - closed_since_last_warning = closed, - "closed new fleet connections over the per-address connection limit" - ); + .is_some_and(|connections| connections.count >= self.shared.connections_per_address); + if at_address_cap { + // Room comes from the address's own connections only, so a client + // at its cap never closes another client's connection. The one it + // replaces frees a slot under the total cap too. + if !state.evict_idle_from(address) { + if let Some(closed) = PER_ADDRESS_CAP.event(Instant::now()) { + tracing::warn!( + latest_address = %address, + closed_since_last_warning = closed, + "closed new fleet connections at the per-address connection limit with no connection from the address idle" + ); + } + return None; } - return None; - } - if state.open.len() >= self.shared.max_connections && !state.evict_idle() { + } else if state.open.len() >= self.shared.max_connections && !state.evict_idle() { if let Some(closed) = TOTAL_CAP.event(Instant::now()) { tracing::warn!( closed_since_last_warning = closed, @@ -250,7 +310,7 @@ impl Connections { } return None; } - *state.per_address.entry(address).or_default() += 1; + state.per_address.entry(address).or_default().count += 1; let (evict, evicted) = oneshot::channel(); let activity = Arc::new(Activity { in_flight: AtomicUsize::new(0), @@ -275,21 +335,25 @@ impl Connections { state.open.remove(&peer); } let address = address_key(peer.ip()); - if let Some(count) = state.per_address.get_mut(&address) { - *count -= 1; - if *count == 0 { + if let Some(connections) = state.per_address.get_mut(&address) { + connections.count -= 1; + if connections.count == 0 { state.per_address.remove(&address); } } } - fn start_request(&self, peer: SocketAddr) -> Option { - // Counted under the registry lock, which `evict_idle` holds, so a - // connection is never evicted once its request has started. + /// Records a request on `peer`'s connection; a `busy` one keeps the + /// connection from going idle until the returned guard is dropped. + pub(super) fn start_request(&self, peer: SocketAddr, busy: bool) -> Option { + // Counted under the registry lock, which eviction holds, so a + // connection is never evicted once its busy request has started. let activity = { let state = self.state(); let activity = state.open.get(&peer)?.clone(); - activity.in_flight.fetch_add(1, Ordering::AcqRel); + if busy { + activity.in_flight.fetch_add(1, Ordering::AcqRel); + } activity }; activity.used.store(true, Ordering::Release); @@ -298,6 +362,7 @@ impl Connections { connections: self.clone(), peer, activity, + busy, }) } @@ -490,18 +555,23 @@ impl AsyncWrite for LimitedConnection { } } -/// Marks the connection busy while a request's handler runs, and for `Connect` -/// while its session lasts (see `SessionBusy`). -struct RequestGuard { +/// A request in progress. A device's request marks its connection busy while +/// its handler runs, and for `Connect` while its session lasts (see +/// `SessionBusy`); any request counts as activity when it starts and ends. +pub(super) struct RequestGuard { connections: Connections, peer: SocketAddr, activity: Arc, + busy: bool, } impl Drop for RequestGuard { fn drop(&mut self) { self.activity.touch(); - // Under the registry lock, like `start_request`, so `evict_idle` never + if !self.busy { + return; + } + // Under the registry lock, like `start_request`, so eviction never // sees the connection idle before it is queued with its new epoch. let idle = { let mut state = self.connections.state(); @@ -576,12 +646,15 @@ where } fn call(&mut self, mut request: http::Request) -> Self::Future { - let guard = remote_addr(&request).and_then(|peer| self.connections.start_request(peer)); + // Only a device's request makes its connection busy; see the module docs. + let busy = has_device_certificate(&request); + let guard = + remote_addr(&request).and_then(|peer| self.connections.start_request(peer, busy)); // `Connect` is the device's session, so its connection stays busy past // the handler: see `SessionBusy`. Any other request is done once its // handler returns. let (guard, session) = match guard { - Some(guard) if is_connect(request.uri().path()) => { + Some(guard) if guard.busy && is_connect(request.uri().path()) => { let session = SessionBusy::new(guard); request.extensions_mut().insert(session.clone()); (None, Some(session)) diff --git a/crates/controller/src/fleet_limits/tests.rs b/crates/controller/src/fleet_limits/tests.rs index 32ffb475..880ff8a7 100644 --- a/crates/controller/src/fleet_limits/tests.rs +++ b/crates/controller/src/fleet_limits/tests.rs @@ -393,13 +393,9 @@ fn pki() -> Pki { /// Serves with the production TLS shape: device CA, client certificate optional. async fn serve_mtls(limits: FleetLimits, pki: &Pki) -> (SocketAddr, Connections) { - let tls = ServerTlsConfig::new() - .identity(pki.server.clone()) - .client_ca_root(pki.ca.clone()) - .client_auth_optional(true); - serve_tracked( + serve_mtls_with( limits, - Some(tls), + pki, Slow { delay: Duration::ZERO, }, @@ -407,25 +403,39 @@ async fn serve_mtls(limits: FleetLimits, pki: &Pki) -> (SocketAddr, Connections) .await } -async fn tls_client( - address: SocketAddr, +/// Like `serve_mtls`, serving `service`. +async fn serve_mtls_with( + limits: FleetLimits, pki: &Pki, - device: bool, -) -> FleetAgentClient { + service: Slow, +) -> (SocketAddr, Connections) { + let tls = ServerTlsConfig::new() + .identity(pki.server.clone()) + .client_ca_root(pki.ca.clone()) + .client_auth_optional(true); + serve_tracked(limits, Some(tls), service).await +} + +/// A TLS endpoint for `address`, presenting the device certificate when `device`. +fn tls_endpoint(address: SocketAddr, pki: &Pki, device: bool) -> tonic::transport::Endpoint { let mut tls = tonic::transport::ClientTlsConfig::new() .ca_certificate(pki.ca.clone()) .domain_name("localhost"); if device { tls = tls.identity(pki.device.clone()); } - let channel = tonic::transport::Endpoint::from_shared(format!("https://{address}")) + tonic::transport::Endpoint::from_shared(format!("https://{address}")) .unwrap() .tls_config(tls) .unwrap() - .connect() - .await - .unwrap(); - FleetAgentClient::new(channel) +} + +async fn tls_client( + address: SocketAddr, + pki: &Pki, + device: bool, +) -> FleetAgentClient { + FleetAgentClient::new(tls_endpoint(address, pki, device).connect().await.unwrap()) } fn renewal(bytes: usize) -> RenewDeviceCertificateRequest { @@ -584,7 +594,10 @@ async fn stays_open(mut stream: TcpStream) -> bool { closed.is_err() } -async fn connections_over_cap(limits: FleetLimits) { +#[tokio::test] +async fn at_the_per_address_cap_the_addresses_longest_idle_connection_makes_room() { + let mut limits = FLEET_LIMITS; + limits.configured.connections_per_address = 2; let address = serve( limits, None, @@ -593,30 +606,50 @@ async fn connections_over_cap(limits: FleetLimits) { }, ) .await; - let first = TcpStream::connect(address).await.unwrap(); - let second = TcpStream::connect(address).await.unwrap(); - - let third = TcpStream::connect(address).await.unwrap(); - bounded(until_closed(third)).await; - - // Closing a connection frees its slot. - drop(first); - time::sleep(Duration::from_millis(100)).await; - let replacement = TcpStream::connect(address).await.unwrap(); - assert!(stays_open(replacement).await, "a freed slot was not reused"); - drop(second); + let oldest = TcpStream::connect(address).await.unwrap(); + time::sleep(Duration::from_millis(50)).await; + let newer = TcpStream::connect(address).await.unwrap(); + time::sleep(Duration::from_millis(50)).await; + + let newest = TcpStream::connect(address).await.unwrap(); + + bounded(until_closed(oldest)).await; + assert!(stays_open(newest).await, "the new connection was refused"); + drop(newer); } #[tokio::test] -async fn connections_over_the_per_address_cap_are_closed() { +async fn at_the_per_address_cap_a_new_connection_is_closed_when_every_connection_is_busy() { + let pki = pki(); let mut limits = FLEET_LIMITS; - limits.configured.connections_per_address = 2; + limits.configured.connections_per_address = 1; + let (address, _) = serve_mtls_with( + limits, + &pki, + Slow { + delay: Duration::from_millis(500), + }, + ) + .await; + let (_outbound, receiver) = mpsc::channel::(1); + let mut stream = bounded( + tls_client(address, &pki, true) + .await + .connect(ReceiverStream::new(receiver)), + ) + .await + .expect("open the stream") + .into_inner(); + + let newcomer = TcpStream::connect(address).await.unwrap(); + bounded(until_closed(newcomer)).await; - connections_over_cap(limits).await; + // The busy connection was kept. + assert!(bounded(stream.message()).await.unwrap().is_some()); } #[test] -fn an_address_over_its_cap_does_not_evict_other_connections() { +fn at_its_cap_an_address_evicts_only_its_own_idle_connections() { use tokio::sync::oneshot::error::TryRecvError; let mut limits = FLEET_LIMITS; @@ -627,17 +660,70 @@ fn an_address_over_its_cap_does_not_evict_other_connections() { let (_first, mut first_evicted) = connections .register(first_address(1)) .expect("first connection"); + let other = SocketAddr::from(([192, 0, 2, 2], 1)); let (_other, mut other_evicted) = connections - .register(SocketAddr::from(([192, 0, 2, 2], 1))) + .register(other) .expect("a connection from another address"); - // At the total cap, and over the first address's cap. - assert!(connections.register(first_address(2)).is_none()); + // At the total cap and at the first address's cap: the first address's + // own connection makes room, though the other one has been idle as long. + assert!(connections.register(first_address(2)).is_some()); + assert_eq!(first_evicted.try_recv(), Ok(())); + assert_eq!(other_evicted.try_recv(), Err(TryRecvError::Empty)); - assert_eq!(first_evicted.try_recv(), Err(TryRecvError::Empty)); + // A busy connection is never evicted, so its address is refused. + let _busy = connections + .start_request(other, true) + .expect("busy request"); + assert!( + connections + .register(SocketAddr::from(([192, 0, 2, 2], 2))) + .is_none() + ); assert_eq!(other_evicted.try_recv(), Err(TryRecvError::Empty)); } +#[test] +fn a_connection_made_idle_by_its_last_request_is_evictable_at_its_address_cap() { + use tokio::sync::oneshot::error::TryRecvError; + + let mut limits = FLEET_LIMITS; + limits.configured.connections_per_address = 1; + let connections = Connections::new(&limits); + let peer = |port| SocketAddr::from(([192, 0, 2, 1], port)); + let (_first, mut first_evicted) = connections.register(peer(1)).expect("first connection"); + let busy = connections + .start_request(peer(1), true) + .expect("busy request"); + assert!(connections.register(peer(2)).is_none()); + + // The request's end queues the connection as idle again. + drop(busy); + assert_eq!(first_evicted.try_recv(), Err(TryRecvError::Empty)); + assert!(connections.register(peer(3)).is_some()); + assert_eq!(first_evicted.try_recv(), Ok(())); +} + +#[test] +fn a_request_without_a_device_certificate_does_not_keep_its_connection_from_eviction() { + let mut limits = FLEET_LIMITS; + limits.configured.max_connections = 1; + let connections = Connections::new(&limits); + let (_anonymous, mut evicted) = connections + .register(SocketAddr::from(([192, 0, 2, 1], 1))) + .expect("first connection"); + let _request = connections + .start_request(SocketAddr::from(([192, 0, 2, 1], 1)), false) + .expect("anonymous request"); + + assert!( + connections + .register(SocketAddr::from(([192, 0, 2, 2], 1))) + .is_some() + ); + assert_eq!(evicted.try_recv(), Ok(())); +} + #[tokio::test] async fn at_the_total_cap_the_longest_idle_connection_makes_room() { let mut limits = FLEET_LIMITS; @@ -664,21 +750,26 @@ async fn at_the_total_cap_the_longest_idle_connection_makes_room() { #[tokio::test] async fn at_the_total_cap_a_new_connection_is_closed_when_every_connection_is_busy() { + let pki = pki(); let mut limits = FLEET_LIMITS; limits.configured.max_connections = 1; - let address = serve( + let (address, _) = serve_mtls_with( limits, - None, + &pki, Slow { delay: Duration::from_millis(500), }, ) .await; let (_outbound, receiver) = mpsc::channel::(1); - let mut stream = bounded(client(address).await.connect(ReceiverStream::new(receiver))) - .await - .expect("open the stream") - .into_inner(); + let mut stream = bounded( + tls_client(address, &pki, true) + .await + .connect(ReceiverStream::new(receiver)), + ) + .await + .expect("open the stream") + .into_inner(); let newcomer = TcpStream::connect(address).await.unwrap(); bounded(until_closed(newcomer)).await; @@ -755,22 +846,15 @@ async fn a_connection_idle_after_its_last_request_is_closed() { #[tokio::test] async fn a_response_the_client_never_accepts_does_not_keep_its_connection_busy() { + let pki = pki(); let mut limits = FLEET_LIMITS; // Far past the test's bound, so only the idle close can end the connection. limits.first_request = Duration::from_secs(60); limits.idle_connection = Duration::from_millis(200); - let (address, connections) = serve_tracked( - limits, - None, - Slow { - delay: Duration::ZERO, - }, - ) - .await; - // A client that grants no HTTP/2 flow-control window, so the server can + let (address, connections) = serve_mtls(limits, &pki).await; + // A device that grants no HTTP/2 flow-control window, so the server can // never send it a response body. It still answers keepalive PINGs. - let channel = tonic::transport::Endpoint::from_shared(format!("http://{address}")) - .unwrap() + let channel = tls_endpoint(address, &pki, true) .initial_stream_window_size(0) .connect() .await @@ -800,22 +884,15 @@ async fn a_response_the_client_never_accepts_does_not_keep_its_connection_busy() #[tokio::test] async fn a_stalled_connect_response_frees_its_connection_when_the_session_ends() { + let pki = pki(); let mut limits = FLEET_LIMITS; // Far past the test's bound, so only the idle close can end the connection. limits.first_request = Duration::from_secs(60); limits.idle_connection = Duration::from_millis(200); - let (address, connections) = serve_tracked( - limits, - None, - Slow { - delay: Duration::ZERO, - }, - ) - .await; - // A client that grants no flow-control window: the session's messages can + let (address, connections) = serve_mtls(limits, &pki).await; + // A device that grants no flow-control window: the session's messages can // never be sent, so hyper keeps the response body after the session ends. - let channel = tonic::transport::Endpoint::from_shared(format!("http://{address}")) - .unwrap() + let channel = tls_endpoint(address, &pki, true) .initial_stream_window_size(0) .connect() .await @@ -836,22 +913,27 @@ async fn a_stalled_connect_response_frees_its_connection_when_the_session_ends() #[tokio::test] async fn an_open_stream_keeps_its_connection_from_going_idle() { + let pki = pki(); let mut limits = FLEET_LIMITS; limits.idle_connection = Duration::from_millis(200); // Messages arrive further apart than the idle limit. - let address = serve( + let (address, _) = serve_mtls_with( limits, - None, + &pki, Slow { delay: Duration::from_millis(500), }, ) .await; let (_outbound, receiver) = mpsc::channel::(1); - let mut stream = bounded(client(address).await.connect(ReceiverStream::new(receiver))) - .await - .expect("open the stream") - .into_inner(); + let mut stream = bounded( + tls_client(address, &pki, true) + .await + .connect(ReceiverStream::new(receiver)), + ) + .await + .expect("open the stream") + .into_inner(); for message in 0..2 { let received = bounded(stream.message()) @@ -861,6 +943,85 @@ async fn an_open_stream_keeps_its_connection_from_going_idle() { } } +/// A HEADERS frame opening a gRPC request for the fleet `method` on +/// `stream_id`. With `end_stream` the request ends with an empty body; without +/// it, the server waits for a body that this client never sends. +fn grpc_request_headers(stream_id: u32, method: &str, end_stream: bool) -> Vec { + use agentdesktop_proto::fleet::fleet_agent_server::SERVICE_NAME; + + let path = format!("/{SERVICE_NAME}/{method}"); + assert!(path.len() < 127, "the path needs a one-byte HPACK length"); + // HPACK: `:method: POST` and `:scheme: http` from the static table, then + // `:path`, `content-type` and `te` as literals without indexing. + let mut block = vec![0x83, 0x86, 0x04, path.len() as u8]; + block.extend_from_slice(path.as_bytes()); + block.extend_from_slice(&[0x0f, 0x10, 16]); + block.extend_from_slice(b"application/grpc"); + block.extend_from_slice(&[0x00, 2]); + block.extend_from_slice(b"te"); + block.push(8); + block.extend_from_slice(b"trailers"); + + let mut frame = (block.len() as u32).to_be_bytes()[1..].to_vec(); + // Type HEADERS; flags END_HEADERS, plus END_STREAM. + frame.extend_from_slice(&[0x1, if end_stream { 0x5 } else { 0x4 }]); + frame.extend_from_slice(&stream_id.to_be_bytes()); + frame.extend_from_slice(&block); + frame +} + +/// Reads frames from `stream` until a HEADERS frame arrives on `stream_id`. +async fn until_response_headers(stream: &mut TcpStream, stream_id: u32) { + loop { + let mut header = [0_u8; 9]; + stream.read_exact(&mut header).await.unwrap(); + let length = u32::from_be_bytes([0, header[0], header[1], header[2]]); + let mut payload = vec![0_u8; length as usize]; + stream.read_exact(&mut payload).await.unwrap(); + let frame_stream = + u32::from_be_bytes([header[5], header[6], header[7], header[8]]) & 0x7fff_ffff; + if header[3] == 0x1 && frame_stream == stream_id { + return; + } + } +} + +#[tokio::test] +async fn a_request_without_a_body_does_not_keep_its_connection_from_eviction() { + let mut limits = FLEET_LIMITS; + limits.configured.max_connections = 1; + // Far past the test's bound, so only eviction can close the connection. + limits.first_request = Duration::from_secs(60); + let address = serve( + limits, + None, + Slow { + delay: Duration::ZERO, + }, + ) + .await; + let mut anonymous = TcpStream::connect(address).await.unwrap(); + anonymous.write_all(PREFACE).await.unwrap(); + // A request whose body never comes: tonic waits for it until the request + // timeout, before any handler runs. + anonymous + .write_all(&grpc_request_headers(1, "BeginEnrollment", false)) + .await + .unwrap(); + // A request that ends at once. Its response shows the server has started + // both. + anonymous + .write_all(&grpc_request_headers(3, "BeginEnrollment", true)) + .await + .unwrap(); + bounded(until_response_headers(&mut anonymous, 3)).await; + + let newcomer = TcpStream::connect(address).await.unwrap(); + + bounded(until_closed(anonymous)).await; + assert!(stays_open(newcomer).await, "the new connection was refused"); +} + /// The HTTP/2 client preface and an empty SETTINGS frame. const PREFACE: &[u8] = b"PRI * HTTP/2.0\r\n\r\nSM\r\n\r\n\x00\x00\x00\x04\x00\x00\x00\x00\x00"; diff --git a/crates/core/src/config.rs b/crates/core/src/config.rs index 3af35579..d4fd835e 100644 --- a/crates/core/src/config.rs +++ b/crates/core/src/config.rs @@ -584,13 +584,17 @@ pub enum ControllerAdminTransport { pub struct ControllerFleetLimitsConfig { /// Open fleet connections allowed in total. At the limit, the connection /// idle the longest is closed to admit a new one; when none is idle, the - /// new one is closed. Defaults to `50000`. + /// new one is closed. A connection is idle unless a client with a device + /// certificate has a request or `Connect` stream open on it. Defaults to + /// `50000`. #[cfg_attr(feature = "schema", schemars(range(min = 1)))] pub max_connections: usize, /// Open fleet connections allowed from one client address: an IPv4 - /// address, or an IPv6 /64. Raise it when more devices share an address, - /// such as behind a NAT or a load balancer that does not preserve client - /// addresses. Defaults to `4096`. + /// address, or an IPv6 /64. At the limit, that address's connection idle + /// the longest is closed to admit a new one; when none is idle, the new one + /// is closed. Raise it when more devices share an address, such as behind a + /// NAT or a load balancer that does not preserve client addresses. Defaults + /// to `4096`. #[cfg_attr(feature = "schema", schemars(range(min = 1)))] pub connections_per_address: usize, /// Largest message, in bytes, that a device and the controller may send diff --git a/schema/controller-config.json b/schema/controller-config.json index 967281a4..c889ba9c 100644 --- a/schema/controller-config.json +++ b/schema/controller-config.json @@ -210,7 +210,7 @@ "type": "object", "properties": { "connectionsPerAddress": { - "description": "Open fleet connections allowed from one client address: an IPv4\naddress, or an IPv6 /64. Raise it when more devices share an address,\nsuch as behind a NAT or a load balancer that does not preserve client\naddresses. Defaults to `4096`.", + "description": "Open fleet connections allowed from one client address: an IPv4\naddress, or an IPv6 /64. At the limit, that address's connection idle\nthe longest is closed to admit a new one; when none is idle, the new one\nis closed. Raise it when more devices share an address, such as behind a\nNAT or a load balancer that does not preserve client addresses. Defaults\nto `4096`.", "type": "integer", "format": "uint", "default": 4096, @@ -248,7 +248,7 @@ "default": "1s" }, "maxConnections": { - "description": "Open fleet connections allowed in total. At the limit, the connection\nidle the longest is closed to admit a new one; when none is idle, the\nnew one is closed. Defaults to `50000`.", + "description": "Open fleet connections allowed in total. At the limit, the connection\nidle the longest is closed to admit a new one; when none is idle, the\nnew one is closed. A connection is idle unless a client with a device\ncertificate has a request or `Connect` stream open on it. Defaults to\n`50000`.", "type": "integer", "format": "uint", "default": 50000, diff --git a/schema/controller-config.md b/schema/controller-config.md index f9ee69ee..a133e104 100644 --- a/schema/controller-config.md +++ b/schema/controller-config.md @@ -19,13 +19,13 @@ |`daemonConfig.revision`|integer|Monotonically increasing revision assigned to the daemon configuration.| |`databaseUrl`|string|SQLite or PostgreSQL URL used for controller state.| |`fleetLimits`|object|Limits on the device-facing fleet API that depend on the deployment.| -|`fleetLimits.connectionsPerAddress`|integer|Open fleet connections allowed from one client address: an IPv4
address, or an IPv6 /64. Raise it when more devices share an address,
such as behind a NAT or a load balancer that does not preserve client
addresses. Defaults to `4096`.| +|`fleetLimits.connectionsPerAddress`|integer|Open fleet connections allowed from one client address: an IPv4
address, or an IPv6 /64. At the limit, that address's connection idle
the longest is closed to admit a new one; when none is idle, the new one
is closed. Raise it when more devices share an address, such as behind a
NAT or a load balancer that does not preserve client addresses. Defaults
to `4096`.| |`fleetLimits.deviceMessageBurst`|integer|Messages one device may send in a burst on its `Connect` stream, after
which it may send one more every `deviceMessageRefill`. Once a device
has sent its burst, the controller reads its messages no faster than
that, and the device waits to send more. Defaults to `600`.| |`fleetLimits.deviceMessageRefill`|string|Time for a device to get back one message. Defaults to `50ms`.| |`fleetLimits.enrollmentBurst`|integer|Enrollments one client address may start in a burst, after which it may
start one more every `enrollmentRefill`. Defaults to `60`.| |`fleetLimits.enrollmentInFlight`|integer|Enrollment requests in progress at once across all clients. Enrollment
calls the identity provider. Defaults to `64`.| |`fleetLimits.enrollmentRefill`|string|Time for a client address to get back one enrollment. Defaults to `1s`.| -|`fleetLimits.maxConnections`|integer|Open fleet connections allowed in total. At the limit, the connection
idle the longest is closed to admit a new one; when none is idle, the
new one is closed. Defaults to `50000`.| +|`fleetLimits.maxConnections`|integer|Open fleet connections allowed in total. At the limit, the connection
idle the longest is closed to admit a new one; when none is idle, the
new one is closed. A connection is idle unless a client with a device
certificate has a request or `Connect` stream open on it. Defaults to
`50000`.| |`fleetLimits.maxMessageBytes`|integer|Largest message, in bytes, that a device and the controller may send
each other. Requests from clients without a device certificate, such as
enrollment, stay limited to 64 KiB. Devices accept at most 4 MiB, so
raising this only raises what devices may send. Defaults to `4194304`
(4 MiB).| |`fleetListen`|string|Address on which the device-facing gRPC fleet API listens.| |`gatewayJwt`|object|LLM gateway JWT signing settings.| From 523222e25c532c1c134f4d10ff2f6dae54c261e1 Mon Sep 17 00:00:00 2001 From: Scott Weiss Date: Thu, 12 Nov 2026 17:21:44 -0500 Subject: [PATCH 49/49] test(controller): run the fleet limit tests against the real fleet handlers No test ran the controller's own Connect or enrollment handlers behind the fleet listener. A handler could stop holding what keeps its limit in force and every test would still pass: - If Connect dropped its session release early, the connection would go idle while the session was still open. The idle deadline would then close every device stream that stays quiet longer than the idle limit. - If BeginEnrollment or CompleteEnrollment dropped its in-flight permit once admitted, enrollmentInFlight would no longer limit the requests waiting on the identity provider. Add tests that run FleetAgentService against an identity provider on a local port. A device session keeps its connection busy through quiet spells longer than the idle limit and frees it when the session ends. Each enrollment handler holds its slot while it waits on the identity provider and frees it when it ends. The test device certificate now carries a device identity, which the Connect handler requires, and the serve helpers accept any fleet service. --- crates/controller/src/fleet_limits/tests.rs | 349 +++++++++++++++++++- 1 file changed, 345 insertions(+), 4 deletions(-) diff --git a/crates/controller/src/fleet_limits/tests.rs b/crates/controller/src/fleet_limits/tests.rs index 880ff8a7..27aa07dc 100644 --- a/crates/controller/src/fleet_limits/tests.rs +++ b/crates/controller/src/fleet_limits/tests.rs @@ -110,7 +110,7 @@ async fn serve(limits: FleetLimits, identity: Option, service: Slow) - async fn serve_tracked( limits: FleetLimits, tls: Option, - service: Slow, + service: impl FleetAgent, ) -> (SocketAddr, Connections) { // As the controller does at startup, so TLS works whichever rustls // backends the test build enables. @@ -357,7 +357,11 @@ async fn an_oversized_enrollment_request_is_rejected() { assert_eq!(error.code(), Code::OutOfRange, "{error}"); } -/// A device CA with a server certificate for `localhost` and a device certificate. +/// The device that `Pki`'s device certificate identifies. +const DEVICE_ID: &str = "7ca03414-bb20-4c80-98ef-7b0538b988ba"; + +/// A device CA with a server certificate for `localhost` and a certificate for +/// the device `DEVICE_ID`, which the controller issues. struct Pki { ca: tonic::transport::Certificate, server: Identity, @@ -387,7 +391,7 @@ fn pki() -> Pki { Pki { ca: tonic::transport::Certificate::from_pem(ca.pem()), server: leaf("localhost"), - device: leaf("device"), + device: leaf(&format!("{DEVICE_ID}.device.agentdesktop.invalid")), } } @@ -407,7 +411,7 @@ async fn serve_mtls(limits: FleetLimits, pki: &Pki) -> (SocketAddr, Connections) async fn serve_mtls_with( limits: FleetLimits, pki: &Pki, - service: Slow, + service: impl FleetAgent, ) -> (SocketAddr, Connections) { let tls = ServerTlsConfig::new() .identity(pki.server.clone()) @@ -1518,3 +1522,340 @@ mod device_streams { .expect("another device is unaffected"); } } + +/// The controller's own fleet service behind the fleet limits, with its +/// identity provider on a local port. +mod fleet_service { + use std::{ + net::SocketAddr, + path::PathBuf, + sync::{ + Arc, + atomic::{AtomicBool, Ordering}, + }, + time::Duration, + }; + + use agentdesktop_proto::fleet::{ + AgentMessage, BeginEnrollmentRequest, CompleteEnrollmentRequest, DaemonConfig, + fleet_agent_server::FleetAgent, + }; + use axum::{ + http::{StatusCode, Uri}, + response::{IntoResponse, Response}, + }; + use base64::{Engine, engine::general_purpose::URL_SAFE_NO_PAD}; + use serde_json::json; + use tokio::{sync::mpsc, time}; + use tokio_stream::wrappers::ReceiverStream; + use tonic::{Code, Request, Status}; + + use super::{ + super::{Connections, FLEET_LIMITS, FleetLimits}, + DEVICE_ID, Pki, bounded, pki, serve_mtls_with, tls_client, until_open_at_most, + }; + use crate::{ + daemon_config::DaemonConfigStore, database::Database, oidc::OidcProvider, + service::FleetAgentService, + }; + + /// The user the identity provider signs in, who enrolled `DEVICE_ID`. + const SUBJECT: &str = "user"; + + /// The `authorization` header a device sends. The identity provider + /// accepts any access token. + const AUTHORIZATION: &str = "Bearer access-token"; + + /// One enrollment slot, which `CompleteEnrollment` waits for only briefly. + const ONE_ENROLLMENT_SLOT: FleetLimits = { + let mut limits = FLEET_LIMITS; + limits.configured.enrollment_in_flight = 1; + limits.enrollment_completion_wait = Duration::from_millis(50); + limits + }; + + /// An identity provider on a local port. It answers discovery, its key set + /// and UserInfo (for `SUBJECT`) until `stall` is called; from then on it + /// holds every request without answering. + struct IdentityProvider { + issuer: String, + stalled: Arc, + /// The path of each request held since `stall`. + held: mpsc::UnboundedReceiver, + } + + impl IdentityProvider { + async fn start() -> Self { + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let issuer = format!("http://{}", listener.local_addr().unwrap()); + let stalled = Arc::new(AtomicBool::new(false)); + let (hold, held) = mpsc::unbounded_channel(); + let endpoints = Endpoints { + issuer: issuer.clone(), + stalled: stalled.clone(), + hold, + }; + let app = axum::Router::new().fallback(move |uri: Uri| endpoints.clone().answer(uri)); + tokio::spawn(async move { axum::serve(listener, app).await }); + Self { + issuer, + stalled, + held, + } + } + + /// The controller's client of this provider, built as at startup. + async fn client(&self) -> OidcProvider { + OidcProvider::discover( + self.issuer.clone(), + "fleet".to_owned(), + "http://127.0.0.1/callback".to_owned(), + ) + .await + .unwrap() + } + + /// From now on, holds every request without answering. + fn stall(&self) { + self.stalled.store(true, Ordering::SeqCst); + } + + /// Waits until a request is held, and returns its path. + async fn held_request(&mut self) -> String { + bounded(self.held.recv()) + .await + .expect("the identity provider stopped") + } + } + + /// The identity provider's server. + #[derive(Clone)] + struct Endpoints { + issuer: String, + stalled: Arc, + hold: mpsc::UnboundedSender, + } + + impl Endpoints { + async fn answer(self, uri: Uri) -> Response { + if self.stalled.load(Ordering::SeqCst) { + let _ = self.hold.send(uri.path().to_owned()); + return std::future::pending().await; + } + let issuer = &self.issuer; + let body = match uri.path() { + "/.well-known/openid-configuration" => json!({ + "issuer": issuer, + "authorization_endpoint": format!("{issuer}/authorize"), + "token_endpoint": format!("{issuer}/token"), + "userinfo_endpoint": format!("{issuer}/userinfo"), + "jwks_uri": format!("{issuer}/keys"), + "device_authorization_endpoint": format!("{issuer}/device"), + }), + "/keys" => json!({ "keys": [{ "kty": "oct", "kid": "current", "k": "a2V5" }] }), + "/userinfo" => json!({ "sub": SUBJECT }), + _ => return StatusCode::NOT_FOUND.into_response(), + }; + axum::Json(body).into_response() + } + } + + /// The controller's fleet service, signing users in with an + /// `IdentityProvider`, with `DEVICE_ID` enrolled and a daemon configuration + /// to send it. + struct Controller { + limits: FleetLimits, + service: FleetAgentService, + idp: IdentityProvider, + database: PathBuf, + } + + impl Controller { + async fn start(limits: FleetLimits) -> Self { + let idp = IdentityProvider::start().await; + let database = std::env::temp_dir().join(format!( + "agentdesktop-fleet-service-{}-{}.db", + std::process::id(), + rand::random::() + )); + let devices = Database::connect(&format!("sqlite://{}?mode=rwc", database.display())) + .await + .unwrap(); + devices + .enroll_device(DEVICE_ID, "host", &idp.issuer, SUBJECT, None) + .await + .unwrap(); + let config = DaemonConfig { + revision: 1, + yaml: b"{}".to_vec(), + sha256: Vec::new(), + }; + let service = FleetAgentService::new( + Some(idp.client().await), + devices, + DaemonConfigStore::new(Some(config)), + None, + None, + &limits, + ); + Self { + limits, + service, + idp, + database, + } + } + + /// Serves the fleet service on the fleet listener, over TLS with + /// `pki`'s device CA. + async fn serve(&self, pki: &Pki) -> (SocketAddr, Connections) { + serve_mtls_with(self.limits, pki, self.service.clone()).await + } + + fn enrollment(&self) -> Enrollment { + Enrollment(self.service.clone()) + } + } + + impl Drop for Controller { + fn drop(&mut self) { + let _ = std::fs::remove_file(&self.database); + } + } + + /// Calls a fleet service's enrollment handlers. + #[derive(Clone)] + struct Enrollment(FleetAgentService); + + impl Enrollment { + /// Begins a browser enrollment, which needs nothing from the identity + /// provider, and returns its ID. + async fn begin(self) -> Result { + let request = BeginEnrollmentRequest { + hostname: "host".to_owned(), + code_challenge: "c".repeat(43), + device_authorization: false, + }; + let response = self.0.begin_enrollment(Request::new(request)).await?; + Ok(response.into_inner().enrollment_id) + } + + /// Begins an enrollment with the device grant, which asks the identity + /// provider for a device code. + async fn begin_with_device_grant(self) -> Result { + let request = BeginEnrollmentRequest { + hostname: "host".to_owned(), + code_challenge: String::new(), + device_authorization: true, + }; + let response = self.0.begin_enrollment(Request::new(request)).await?; + Ok(response.into_inner().enrollment_id) + } + + /// Completes `enrollment_id` with an ID token signed by a key the + /// controller has not seen, so it asks the identity provider for its + /// keys again before anything else. + async fn complete(self, enrollment_id: String) -> Result<(), Status> { + let header = json!({ "alg": "RS256", "kid": "rotated" }).to_string(); + let mut request = Request::new(CompleteEnrollmentRequest { + enrollment_id, + id_token: format!("{}.e30.c2ln", URL_SAFE_NO_PAD.encode(header)), + ..Default::default() + }); + request + .metadata_mut() + .insert("authorization", AUTHORIZATION.parse().unwrap()); + self.0.complete_enrollment(request).await.map(drop) + } + + /// Asserts that every slot is taken: neither handler is admitted. + async fn assert_full(self) { + let begun = self.clone().begin().await.expect_err("an enrollment began"); + assert_eq!(begun.code(), Code::ResourceExhausted, "{begun}"); + let completed = self + .complete(String::new()) + .await + .expect_err("an enrollment completed"); + assert_eq!(completed.code(), Code::ResourceExhausted, "{completed}"); + } + } + + /// A device's `Connect` session keeps its connection busy for as long as + /// it lasts, through quiet spells longer than the idle limit, and frees it + /// when it ends. + #[tokio::test] + async fn a_device_session_keeps_its_connection_busy_until_it_ends() { + let mut limits = FLEET_LIMITS; + limits.idle_connection = Duration::from_millis(200); + let controller = Controller::start(limits).await; + let pki = pki(); + let (address, connections) = controller.serve(&pki).await; + + let (outbound, receiver) = mpsc::channel::(1); + let mut request = Request::new(ReceiverStream::new(receiver)); + request + .metadata_mut() + .insert("authorization", AUTHORIZATION.parse().unwrap()); + let mut session = bounded(tls_client(address, &pki, true).await.connect(request)) + .await + .expect("open the session") + .into_inner(); + let config = bounded(session.message()) + .await + .expect("receive the configuration"); + assert!( + config.is_some(), + "the session ended before its configuration" + ); + + // Nothing is sent either way for several idle limits. + let quiet = time::timeout(5 * limits.idle_connection, session.message()).await; + assert!(quiet.is_err(), "the quiet session ended: {quiet:?}"); + assert_eq!(connections.open_count(), 1); + + // The device ends the session, so its connection goes idle and closes. + drop(outbound); + bounded(until_open_at_most(&connections, 0)).await; + } + + /// `BeginEnrollment` holds its in-flight slot while it waits on the + /// identity provider, and frees it when it ends. + #[tokio::test] + async fn a_beginning_enrollment_holds_its_slot_while_it_runs() { + let mut controller = Controller::start(ONE_ENROLLMENT_SLOT).await; + controller.idp.stall(); + let running = tokio::spawn(controller.enrollment().begin_with_device_grant()); + assert_eq!(controller.idp.held_request().await, "/device"); + + controller.enrollment().assert_full().await; + + running.abort(); + let _ = running.await; + controller + .enrollment() + .begin() + .await + .expect("the slot was not freed"); + } + + /// `CompleteEnrollment` holds its in-flight slot while it waits on the + /// identity provider, and frees it when it ends. + #[tokio::test] + async fn a_completing_enrollment_holds_its_slot_while_it_runs() { + let mut controller = Controller::start(ONE_ENROLLMENT_SLOT).await; + let enrollment_id = controller.enrollment().begin().await.unwrap(); + controller.idp.stall(); + let running = tokio::spawn(controller.enrollment().complete(enrollment_id)); + assert_eq!(controller.idp.held_request().await, "/keys"); + + controller.enrollment().assert_full().await; + + running.abort(); + let _ = running.await; + controller + .enrollment() + .begin() + .await + .expect("the slot was not freed"); + } +}