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 6620999a..ee1a63be 100644 --- a/crates/controller/src/application.rs +++ b/crates/controller/src/application.rs @@ -2,18 +2,18 @@ use crate::{ daemon_config::{self, DaemonConfigStore}, database::Database, device_ca::DeviceCertificateIssuer, + fleet_limits::FleetLimits, gateway_jwt::{self, GatewayJwtIssuer}, management::{self, Settings, State}, oidc::OidcProvider, service::FleetAgentService, }; use agentdesktop_core::config::ControllerConfig; -use agentdesktop_proto::fleet::fleet_agent_server::FleetAgentServer; 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. @@ -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(); @@ -103,12 +106,14 @@ 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 = FleetLimits::from_config(&config.fleet_limits); let service = FleetAgentService::new( Some(oidc), database, daemon_config, gateway_jwt_issuer, Some(device_certificate_issuer), + &limits, ) .with_usage_collection(usage_ingest.is_some()); @@ -126,22 +131,18 @@ pub async fn run( .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)?; - 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(FleetAgentServer::new(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(fleet_listen) + 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 new file mode 100644 index 00000000..6fb8c760 --- /dev/null +++ b/crates/controller/src/fleet_limits.rs @@ -0,0 +1,710 @@ +//! 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::{ + collections::HashMap, + convert::Infallible, + net::{IpAddr, Ipv6Addr, SocketAddr}, + sync::{ + Arc, Mutex, + atomic::{AtomicU64, Ordering}, + }, + task::{Context, Poll}, + time::Duration, +}; + +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, +}; +use tonic::{ + Status, + codegen::{Service, http}, + server::NamedService, + service::Routes, + transport::{ + Server, ServerTlsConfig, + server::{TcpConnectInfo, TcpIncoming, TlsConnectInfo}, + }, +}; +use tower_layer::{Identity, Stack}; + +mod connections; + +use connections::{Connections, LimitedIncoming, TrackRequests}; +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. + 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, + /// 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, + /// How long a connection may be silent before TCP keepalive probes start, + /// for peers that vanish without closing. + pub tcp_keepalive: Duration, + /// 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. + pub first_request: Duration, + pub idle_connection: Duration, + /// 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 + /// 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 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, + /// 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, + /// 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. + pub stream_open_burst: u32, + pub stream_open_refill: Duration, +} + +/// 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), + first_request: Duration::from_secs(20), + idle_connection: Duration::from_secs(120), + concurrent_streams_per_connection: 64, + error_resets_per_connection: 1024, + unauthenticated_message_bytes: 64 * 1024, + enrollment_completion_wait: Duration::from_secs(10), + stream_open_burst: 20, + stream_open_refill: Duration::from_secs(60), +}; + +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 { + Self { + configured: *config, + ..FLEET_LIMITS + } + } + + /// 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)) + .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) + .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)); + 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. + fn fleet_service(&self, service: T) -> FleetService { + let service = Arc::new(service); + FleetService { + unauthenticated: FleetAgentServer::from_arc(service.clone()) + .max_decoding_message_size(self.unauthenticated_message_bytes) + .max_encoding_message_size(self.configured.max_message_bytes), + devices: FleetAgentServer::from_arc(service) + .max_decoding_message_size(self.configured.max_message_bytes) + .max_encoding_message_size(self.configured.max_message_bytes), + } + } + + pub fn enrollment(&self) -> EnrollmentLimiter { + EnrollmentLimiter::new(self) + } + + /// Per-device `Connect` state: one stream per device, and its budgets. + pub fn device_streams(&self) -> DeviceStreams { + DeviceStreams::new(self) + } +} + +/// 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. +struct FleetService { + unauthenticated: FleetAgentServer, + devices: FleetAgentServer, +} + +impl Clone for FleetService { + fn clone(&self) -> Self { + Self { + unauthenticated: self.unauthenticated.clone(), + devices: self.devices.clone(), + } + } +} + +/// 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 { + 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 { + rpc_class(path) == Some(RpcClass::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 { + request + .extensions() + .get::>() + .and_then(TlsConnectInfo::peer_certs) + .is_some_and(|certificates| !certificates.is_empty()) +} + +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.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()) || !has_device_certificate(&request) { + self.unauthenticated.call(request) + } else { + self.devices.call(request) + } + } +} + +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>, + burst: f64, + refill: Duration, + completion_wait: Duration, +} + +/// 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() + } +} + +/// 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 { + 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 + } + + /// 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) + } +} + +impl EnrollmentLimiter { + fn new(limits: &FleetLimits) -> Self { + Self { + in_flight: Arc::new(Semaphore::new(limits.configured.enrollment_in_flight)), + peers: Arc::default(), + burst: f64::from(limits.configured.enrollment_burst), + refill: limits.configured.enrollment_refill, + completion_wait: limits.enrollment_completion_wait, + } + } + + /// 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()) + { + if let Some(refused) = ENROLLMENT_RATE.event(Instant::now()) { + tracing::warn!( + latest_address = %address_key(peer.ip()), + refused_since_last_warning = refused, + "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(|_| in_flight_refusal()) + } + + /// 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), + _ => Err(in_flight_refusal()), + } + } + + fn take(&self, peer: IpAddr, now: Instant) -> bool { + let mut peers = self.peers.lock().expect("enrollment limiter lock poisoned"); + peers + .bucket(address_key(peer), Bucket::full(self.burst, now)) + .take(now, self.refill, self.burst) + } +} + +fn in_flight_refusal() -> Status { + if let Some(refused) = ENROLLMENT_IN_FLIGHT.event(Instant::now()) { + tracing::warn!( + refused_since_last_warning = refused, + "refused enrollment requests with every in-flight slot taken" + ); + } + Status::resource_exhausted("too many enrollment requests in progress; retry later") +} + +/// `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 { + 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, +} + +/// 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 +/// 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 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"); + 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()); + } + let now = Instant::now(); + let device = devices + .entry(device_id.to_owned()) + .or_insert_with(|| DeviceState { + current: None, + 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) + } + + /// 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| { + device.opens.take( + Instant::now(), + limits.stream_open_refill, + f64::from(limits.stream_open_burst), + ) + }); + if admitted { + Ok(()) + } else { + if let Some(refused) = STREAM_OPENS.event(Instant::now()) { + tracing::warn!( + latest_device_id = device_id, + refused_since_last_warning = refused, + "refused Connect streams from devices that open them 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.with_device(device_id, |device| { + device.current.replace((generation, replace)) + }); + if let Some((_, previous)) = previous { + let _ = previous.send(()); + } + DeviceStream { + streams: self.clone(), + device_id: device_id.to_owned(), + generation, + superseded, + } + } +} + +impl DeviceStream { + /// 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; + 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 + } +} + +impl Drop for DeviceStream { + fn drop(&mut self) { + let mut devices = self + .streams + .devices + .lock() + .expect("device stream registry lock poisoned"); + // A newer stream for the same device keeps its registration. + if let Some(device) = devices.get_mut(&self.device_id) + && device + .current + .as_ref() + .is_some_and(|(generation, _)| *generation == self.generation) + { + device.current = None; + } + } +} + +/// 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 +/// (however much later that is). A warning that names a client or device names +/// only the latest one. +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(); +static DEVICE_MESSAGES: 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 new file mode 100644 index 00000000..f14eea48 --- /dev/null +++ b/crates/controller/src/fleet_limits/connections.rs @@ -0,0 +1,757 @@ +//! 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 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. +//! +//! 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}, + 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, ThrottledWarning, address_key, has_device_certificate, is_connect}; + +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 +/// 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>, + /// Open connections by `address_key`. + 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. +pub(super) 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, + ) + } + + /// 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 device request in progress. + fn evict_idle(&mut self) -> bool { + 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) { + 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); + } + } +} + +impl Connections { + pub(super) fn new(limits: &FleetLimits) -> Self { + Self { + shared: Arc::new(Shared { + 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(), + }), + } + } + + 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 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(), + } + } + + 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(); + let address = address_key(peer.ip()); + let at_address_cap = state + .per_address + .get(&address) + .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; + } + } 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, + "closed new fleet connections at the connection limit with no connection idle" + ); + } + return None; + } + state.per_address.entry(address).or_default().count += 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(Instant::now()), + evict: Mutex::new(Some(evict)), + waker: Mutex::new(None), + }); + state.open.insert(peer, activity.clone()); + state.mark_idle(peer, 0); + Some((activity, evicted)) + } + + 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); + } + let address = address_key(peer.ip()); + if let Some(connections) = state.per_address.get_mut(&address) { + connections.count -= 1; + if connections.count == 0 { + state.per_address.remove(&address); + } + } + } + + /// 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(); + if busy { + activity.in_flight.fetch_add(1, Ordering::AcqRel); + } + activity + }; + activity.used.store(true, Ordering::Release); + activity.touch(); + Some(RequestGuard { + connections: self.clone(), + peer, + activity, + busy, + }) + } + + #[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, +} + +#[cfg(test)] +impl LimitedIncoming { + pub(super) fn connections(&self) -> &Connections { + &self.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, 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, + 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; + if let Some(closed) = EVICTIONS.event(Instant::now()) { + tracing::warn!( + closed_since_last_warning = closed, + "closed idle fleet connections to make room at the connection limit" + ); + } + 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 { + // 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 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(()); + } + } + } +} + +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> { + if let Err(error) = self.check_open(context) { + return Poll::Ready(Err(error)); + } + 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> { + if let Err(error) = self.check_open(context) { + return Poll::Ready(Err(error)); + } + 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> { + 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) + } +} + +/// 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(); + 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(); + 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(); + } + } +} + +/// 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, mut request: http::Request) -> Self::Future { + // 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 guard.busy && 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; + drop(guard); + let response = response?; + Ok(match session { + Some(session) => response.map(|body| { + Body::new(TrackedBody { + inner: body, + _release: session.release_on_drop(), + }) + }), + None => response, + }) + }) + } +} + +/// 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, + _release: ReleaseOnDrop, +} + +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 new file mode 100644 index 00000000..27aa07dc --- /dev/null +++ b/crates/controller/src/fleet_limits/tests.rs @@ -0,0 +1,1861 @@ +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_client::FleetAgentClient, + fleet_agent_server::FleetAgent, +}; +use futures_core::Stream; +use tokio::{ + io::{AsyncReadExt, AsyncWriteExt}, + net::TcpStream, + sync::mpsc, + time, +}; +use tokio_stream::wrappers::ReceiverStream; +use tonic::{ + Code, Request, Response, Status, Streaming, + service::Routes, + transport::{Identity, ServerTlsConfig}, +}; + +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 { + 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, then ends its session, holding the connection's session +/// release as the controller's `Connect` does. +#[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: 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() { + 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 tls = identity.map(|identity| ServerTlsConfig::new().identity(identity)); + serve_tracked(limits, tls, service).await.0 +} + +/// Like `serve`, with any TLS configuration, also returning the connection registry. +async fn serve_tracked( + limits: FleetLimits, + tls: Option, + service: impl FleetAgent, +) -> (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) + .unwrap(); + let address = listener.local_address(); + let connections = listener.connections().clone(); + tokio::spawn(listener.serve(service, Routes::default())); + (address, connections) +} + +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]; + 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 limits = FleetLimits { + tls_handshake: Duration::from_millis(200), + ..FLEET_LIMITS + }; + 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; +} + +#[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, + 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() { + 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 (_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}"); + } +} + +#[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, + 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(PREFACE).await.unwrap(); + + 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 = 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()); + 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); + 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() +} + +#[tokio::test] +async fn streams_beyond_the_per_connection_limit_wait() { + let limits = FleetLimits { + concurrent_streams_per_connection: 1, + ..FLEET_LIMITS + }; + + let elapsed = two_calls_on_one_connection(limits).await; + + assert!( + elapsed >= Duration::from_millis(950), + "ran concurrently: {elapsed:?}" + ); +} + +fn small_messages() -> FleetLimits { + let mut limits = FleetLimits { + unauthenticated_message_bytes: 128, + ..FLEET_LIMITS + }; + limits.configured.max_message_bytes = 1024; + 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}"); +} + +/// 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, + 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(&format!("{DEVICE_ID}.device.agentdesktop.invalid")), + } +} + +/// Serves with the production TLS shape: device CA, client certificate optional. +async fn serve_mtls(limits: FleetLimits, pki: &Pki) -> (SocketAddr, Connections) { + serve_mtls_with( + limits, + pki, + Slow { + delay: Duration::ZERO, + }, + ) + .await +} + +/// Like `serve_mtls`, serving `service`. +async fn serve_mtls_with( + limits: FleetLimits, + pki: &Pki, + service: impl FleetAgent, +) -> (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()); + } + tonic::transport::Endpoint::from_shared(format!("https://{address}")) + .unwrap() + .tls_config(tls) + .unwrap() +} + +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 { + 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(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}"); +} + +#[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 = FLEET_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 = 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_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. +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() +} + +#[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, + 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; + + 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_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 = 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; + + // The busy connection was kept. + assert!(bounded(stream.message()).await.unwrap().is_some()); +} + +#[test] +fn at_its_cap_an_address_evicts_only_its_own_idle_connections() { + use tokio::sync::oneshot::error::TryRecvError; + + let mut limits = FLEET_LIMITS; + 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 + .register(first_address(1)) + .expect("first connection"); + let other = SocketAddr::from(([192, 0, 2, 2], 1)); + let (_other, mut other_evicted) = connections + .register(other) + .expect("a connection from another address"); + + // 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)); + + // 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; + limits.configured.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; + + 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 pki = pki(); + let mut limits = FLEET_LIMITS; + limits.configured.max_connections = 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; + + // 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 = FLEET_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 = FLEET_LIMITS; + limits.first_request = Duration::from_millis(200); + 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(PREFACE).await.unwrap(); + + bounded(until_closed(stream)).await; +} + +#[tokio::test] +async fn a_connection_idle_after_its_last_request_is_closed() { + 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; + 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 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_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 = tls_endpoint(address, &pki, true) + .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 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_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 = tls_endpoint(address, &pki, true) + .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 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_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(); + + 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}"); + } +} + +/// 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"; + +/// 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 + .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(); + 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}; + + use tokio::time::Instant; + + use super::super::{FLEET_LIMITS, FleetLimits, PEERS_PER_GENERATION, address_key}; + use super::Duration; + + fn limits(in_flight: usize, burst: u32) -> FleetLimits { + 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 { + 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 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(); + + let _first = limiter.admit(None).expect("first request"); + 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 the_peer_table_stays_bounded_while_addresses_keep_changing() { + let limiter = limits(64, 1).enrollment(); + let now = Instant::now(); + + for index in 0..5 * PEERS_PER_GENERATION { + assert!(limiter.take(numbered(index), now)); + } + + 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::")); + } +} + +mod device_streams { + use tokio::{sync::oneshot::error::TryRecvError, time::Instant}; + + 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 = streams(); + 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 = streams(); + 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 = streams(); + drop(streams.open("device-a")); + + assert!( + streams.devices.lock().unwrap()["device-a"] + .current + .is_none(), + "the closed stream is still registered" + ); + } + + #[test] + fn a_device_gets_its_message_burst_then_the_refill_rate() { + let mut limits = FLEET_LIMITS; + 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(); + let second = 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] + fn reconnecting_does_not_reset_the_message_budget() { + let mut limits = FLEET_LIMITS; + 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"); + 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.spend_message(now).is_some(), + "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"); + } +} + +/// 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"); + } +} diff --git a/crates/controller/src/lib.rs b/crates/controller/src/lib.rs index 6e88b029..82969132 100644 --- a/crates/controller/src/lib.rs +++ b/crates/controller/src/lib.rs @@ -2,8 +2,23 @@ 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; 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(); +} diff --git a/crates/controller/src/oidc.rs b/crates/controller/src/oidc.rs index f6c8672d..def11089 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. @@ -32,6 +37,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, @@ -120,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('/') @@ -226,26 +264,25 @@ 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?; - - 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 reservation = Reservation { + provider: Some(self.inner.clone()), + enrollment_id: enrollment_id.clone(), + }; + + 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 - /// 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, @@ -300,13 +337,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. @@ -467,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 @@ -837,6 +876,71 @@ 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 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"])); diff --git a/crates/controller/src/service.rs b/crates/controller/src/service.rs index 6fe3570b..d19ba1e9 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::{DeviceStreams, EnrollmentLimiter, FleetLimits, SessionBusy}, gateway_jwt::GatewayJwtIssuer, oidc::{OidcPrincipal, OidcProvider}, }; @@ -34,15 +35,21 @@ pub struct FleetAgentService { gateway_jwt_issuer: Option, device_certificate_issuer: Option, usage_enabled: bool, + enrollment: EnrollmentLimiter, + device_streams: DeviceStreams, } impl FleetAgentService { + /// `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, daemon_config: DaemonConfigStore, gateway_jwt_issuer: Option, device_certificate_issuer: Option, + limits: &FleetLimits, ) -> Self { Self { oidc, @@ -51,6 +58,8 @@ impl FleetAgentService { gateway_jwt_issuer, device_certificate_issuer, usage_enabled: false, + enrollment: limits.enrollment(), + device_streams: limits.device_streams(), } } @@ -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_completion().await?; let oidc = self .oidc .as_ref() @@ -244,7 +255,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") @@ -252,6 +271,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(); @@ -265,28 +287,15 @@ 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); tokio::spawn(async move { + let _session_busy = session_busy; let enrolled_principal = match database.device_principal(&authenticated_device_id).await { Ok(principal) => principal, @@ -297,8 +306,35 @@ 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, + changed = daemon_updates.changed() => { + if changed.is_err() { + break; + } + let daemon = daemon_updates.borrow_and_update().clone(); + 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() => { let (oidc, access_token) = &oauth_revalidation; match oidc.authenticate_access_token(access_token).await { @@ -306,8 +342,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) => { @@ -317,13 +352,23 @@ 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; } } } - message = inbound.message() => match message { + _ = &mut device_stream.superseded => { + info!(device_id = %authenticated_device_id, "closing stream replaced by a newer stream from the same device"); + let _ = auth_sender + .try_send(Err(Status::aborted("replaced by a newer stream from this device"))); + break; + } + _ = 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)) => { + 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 { @@ -332,6 +377,8 @@ impl FleetAgent for FleetAgentService { claimed_device_id = %hello.device_id, "device certificate and hello identity do not match" ); + let _ = auth_sender + .try_send(Err(Status::permission_denied("hello does not match the device certificate"))); break; } if let Err(err) = @@ -662,6 +709,51 @@ 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::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 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, + DaemonConfigStore::new(None), + None, + None, + &full, + ); + + // 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"; @@ -677,8 +769,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::FleetLimits::from_config(&Default::default()), + ); let principal = |issuer: &str, subject: &str| OidcPrincipal { issuer: issuer.to_owned(), subject: subject.to_owned(), @@ -731,8 +829,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::FleetLimits::from_config(&Default::default()), + ); let status = service .authenticate_device_certificate_identity( diff --git a/crates/core/src/config.rs b/crates/core/src/config.rs index 84ad790b..d4fd835e 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,79 @@ 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, Copy, 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. 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. 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 + /// 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 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 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`. + #[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")] + #[cfg_attr(feature = "schema", schemars(with = "String"))] + 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::DEFAULT + } +} + /// OpenID Connect settings used for interactive device enrollment. #[derive(Clone, Debug, Deserialize, Serialize)] #[cfg_attr(feature = "schema", derive(schemars::JsonSchema))] @@ -1415,6 +1491,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 +2658,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/README.md b/deploy/README.md index 6da1b48b..70e8aa21 100644 --- a/deploy/README.md +++ b/deploy/README.md @@ -40,8 +40,17 @@ 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. 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 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: ```console kubectl port-forward deployment/agentdesktop 8080:8080 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/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..a03cf56b 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: {} @@ -38,6 +43,17 @@ tlsSecretName: agentdesktop-controller-tls service: type: LoadBalancer 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 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: {} diff --git a/schema/controller-config.json b/schema/controller-config.json index 774c2998..c889ba9c 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. 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, + "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`. 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, + "minimum": 1 + }, + "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 in a burst, after which it may\nstart one more every `enrollmentRefill`. Defaults to `60`.", + "type": "integer", + "format": "uint32", + "default": 60, + "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": 1 + }, + "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. 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, + "minimum": 1 + }, + "maxMessageBytes": { + "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": 1 + } + }, + "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..a133e104 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. 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. 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.| |`gatewayJwt.issuer`|string|Issuer claim placed in generated JWTs.|