fixup! Bound concurrent request buffering

Limit concurrent body reads per source IP across connections so one
peer cannot occupy the global pool with replacement stalled requests.
Keep the global limit and body deadline, and release peer records when
reads finish, fail, or are cancelled.

Test repeated saturation while another peer reads and authenticates a
request, plus global capacity, address normalization, and cleanup.

This change was made with assistance from Codex.

Author: benthecarman

ldk-server/Cargo.toml

@@ -39,4 +39,5 @@ default = []
 experimental-lsps2-support = []
 
 [dev-dependencies]
+tokio = { version = "1.38.0", default-features = false, features = ["test-util"] }
 futures-util = "0.3.31"

ldk-server/src/main.rs

@@ -717,7 +717,7 @@ fn main() {
 				},
 				res = grpc_listener.accept() => {
 					match res {
-						Ok((stream, _)) => {
+						Ok((stream, peer_addr)) => {
 							let handshake_permit =
 								match Arc::clone(&tls_handshake_semaphore).try_acquire_owned() {
 									Ok(permit) => permit,
@@ -733,6 +733,7 @@ fn main() {
 								metrics_auth_header.clone(),
 								event_sender.clone(),
 								shutdown_rx.clone(),
+								peer_addr.ip(),
 							);
 							let acceptor = tls_acceptor.clone();
 							runtime.spawn(async move {

ldk-server/src/service.rs

@@ -7,9 +7,11 @@
 // You may not use this file except in accordance with one or both of these
 // licenses.
 
+use std::collections::BTreeMap;
 use std::future::Future;
+use std::net::IpAddr;
 use std::pin::Pin;
-use std::sync::Arc;
+use std::sync::{Arc, Mutex};
 use std::time::Duration;
 
 use http_body_util::{BodyExt, Limited};
@@ -44,7 +46,7 @@ use ldk_server_grpc::grpc::{
 	GRPC_STATUS_UNAUTHENTICATED, GRPC_STATUS_UNAVAILABLE, GRPC_STATUS_UNIMPLEMENTED,
 };
 use prost::Message;
-use tokio::sync::{broadcast, mpsc, Semaphore};
+use tokio::sync::{broadcast, mpsc};
 
 use crate::api::bolt11_claim_for_id::handle_bolt11_claim_for_id_request;
 use crate::api::bolt11_fail_for_id::handle_bolt11_fail_for_id_request;
@@ -101,13 +103,57 @@ const GRPC_SERVICE_PREFIX: &str = "/api.LightningNode/";
 // Maximum request body size: 10 MB
 const MAX_BODY_SIZE: usize = 10 * 1024 * 1024;
 const MAX_CONCURRENT_BODY_READS: usize = 8;
+// Share this limit across all connections from the same source IP. Authentication
+// covers the body, so it cannot be checked before reserving buffer space.
+const MAX_BODY_READS_PER_PEER: usize = 2;
 // A client that stalls mid-body would otherwise hold one of the few body-read slots
 // indefinitely.
 const REQUEST_BODY_TIMEOUT: Duration = Duration::from_secs(30);
-static REQUEST_BODY_SEMAPHORE: Semaphore = Semaphore::const_new(MAX_CONCURRENT_BODY_READS);
+static REQUEST_BODY_LIMITER: RequestBodyLimiter = RequestBodyLimiter::new();
+
+// Track only active reads; the map has at most MAX_CONCURRENT_BODY_READS entries.
+struct RequestBodyLimiter {
+	active: Mutex<BTreeMap<IpAddr, usize>>,
+}
+
+impl RequestBodyLimiter {
+	const fn new() -> Self {
+		Self { active: Mutex::new(BTreeMap::new()) }
+	}
+
+	fn try_acquire(&self, peer_ip: IpAddr) -> Result<RequestBodyPermit<'_>, GrpcStatus> {
+		// Treat IPv4 and its IPv4-mapped IPv6 representation as the same peer.
+		let peer_ip = peer_ip.to_canonical();
+		let mut active = self.active.lock().unwrap();
+		if active.get(&peer_ip).copied().unwrap_or(0) >= MAX_BODY_READS_PER_PEER
+			|| active.values().sum::<usize>() >= MAX_CONCURRENT_BODY_READS
+		{
+			return Err(GrpcStatus::new(GRPC_STATUS_UNAVAILABLE, "Too many concurrent requests"));
+		}
+		*active.entry(peer_ip).or_default() += 1;
+		Ok(RequestBodyPermit { limiter: self, peer_ip })
+	}
+}
+
+struct RequestBodyPermit<'a> {
+	limiter: &'a RequestBodyLimiter,
+	peer_ip: IpAddr,
+}
+
+impl Drop for RequestBodyPermit<'_> {
+	fn drop(&mut self) {
+		let mut active = self.limiter.active.lock().unwrap();
+		let count = active.get_mut(&self.peer_ip).unwrap();
+		*count -= 1;
+		if *count == 0 {
+			active.remove(&self.peer_ip);
+		}
+	}
+}
 
 #[derive(Clone)]
 pub(crate) struct NodeService {
+	peer_ip: IpAddr,
 	context: Arc<Context>,
 	api_key: String,
 	metrics: Option<Arc<Metrics>>,
@@ -120,10 +166,10 @@ impl NodeService {
 	pub(crate) fn new(
 		node: Arc<Node>, api_key: String, metrics: Option<Arc<Metrics>>,
 		metrics_auth_header: Option<String>, event_sender: broadcast::Sender<EventEnvelope>,
-		shutdown_rx: tokio::sync::watch::Receiver<bool>,
+		shutdown_rx: tokio::sync::watch::Receiver<bool>, peer_ip: IpAddr,
 	) -> Self {
 		let context = Arc::new(Context { node });
-		Self { context, api_key, metrics, metrics_auth_header, event_sender, shutdown_rx }
+		Self { context, api_key, metrics, metrics_auth_header, event_sender, shutdown_rx, peer_ip }
 	}
 }
 
@@ -275,36 +321,23 @@ impl Service<Request<Incoming>> for NodeService {
 		let event_sender = self.event_sender.clone();
 		let shutdown_rx = self.shutdown_rx.clone();
 		let (request_parts, request_body) = req.into_parts();
+		let peer_ip = self.peer_ip;
 		let future: Self::Future = Box::pin(async move {
-			let body_permit = match REQUEST_BODY_SEMAPHORE.try_acquire() {
-				Ok(permit) => permit,
-				Err(_) => {
-					return Ok(grpc_error_response(GrpcStatus::new(
-						GRPC_STATUS_UNAVAILABLE,
-						"Too many concurrent requests",
-					)));
-				},
-			};
 			let content_length = match request_content_length(&request_parts.headers) {
 				Ok(content_length) => content_length,
 				Err(status) => return Ok(grpc_error_response(status)),
 			};
-			let body_bytes = match tokio::time::timeout(
-				REQUEST_BODY_TIMEOUT,
-				read_request_body(request_body, content_length),
+			let body_bytes = match read_request_body(
+				request_body,
+				content_length,
+				peer_ip,
+				&REQUEST_BODY_LIMITER,
 			)
 			.await
 			{
-				Ok(Ok(bytes)) => bytes,
-				Ok(Err(status)) => return Ok(grpc_error_response(status)),
-				Err(_) => {
-					return Ok(grpc_error_response(GrpcStatus::new(
-						GRPC_STATUS_UNAVAILABLE,
-						"Timed out reading request body",
-					)));
-				},
+				Ok(bytes) => bytes,
+				Err(status) => return Ok(grpc_error_response(status)),
 			};
-			drop(body_permit);
 
 			let auth_req = Request::from_parts(request_parts, ());
 			if let Err(e) = validate_auth(&auth_req, &api_key, &body_bytes) {
@@ -627,21 +660,32 @@ fn validate_request_body_len(
 	Ok(())
 }
 
-async fn read_request_body(
-	body: Incoming, content_length: Option<u64>,
-) -> Result<bytes::Bytes, GrpcStatus> {
-	let limited_body = Limited::new(body, MAX_BODY_SIZE);
-	let bytes = match limited_body.collect().await {
-		Ok(collected) => collected.to_bytes(),
-		Err(_) => {
-			return Err(GrpcStatus::new(
-				GRPC_STATUS_INVALID_ARGUMENT,
-				"Request body too large or failed to read",
-			));
-		},
-	};
-	validate_request_body_len(content_length, bytes.len())?;
-	Ok(bytes)
+async fn read_request_body<B>(
+	body: B, content_length: Option<u64>, peer_ip: IpAddr, limiter: &RequestBodyLimiter,
+) -> Result<bytes::Bytes, GrpcStatus>
+where
+	B: hyper::body::Body<Data = bytes::Bytes>,
+	B::Error: Into<Box<dyn std::error::Error + Send + Sync>>,
+{
+	let _permit = limiter.try_acquire(peer_ip)?;
+	tokio::time::timeout(REQUEST_BODY_TIMEOUT, async move {
+		let limited_body = Limited::new(body, MAX_BODY_SIZE);
+		let bytes = match limited_body.collect().await {
+			Ok(collected) => collected.to_bytes(),
+			Err(_) => {
+				return Err(GrpcStatus::new(
+					GRPC_STATUS_INVALID_ARGUMENT,
+					"Request body too large or failed to read",
+				));
+			},
+		};
+		validate_request_body_len(content_length, bytes.len())?;
+		Ok(bytes)
+	})
+	.await
+	.unwrap_or_else(|_| {
+		Err(GrpcStatus::new(GRPC_STATUS_UNAVAILABLE, "Timed out reading request body"))
+	})
 }
 
 /// Map an `LdkServerError` to a `GrpcStatus`.
@@ -659,6 +703,131 @@ pub(crate) fn ldk_error_to_grpc_status(e: LdkServerError) -> GrpcStatus {
 mod tests {
 	use super::*;
 
+	fn stalled_body(
+	) -> impl hyper::body::Body<Data = bytes::Bytes, Error = std::convert::Infallible> {
+		http_body_util::StreamBody::new(futures_util::stream::pending::<
+			Result<hyper::body::Frame<bytes::Bytes>, std::convert::Infallible>,
+		>())
+	}
+
+	#[test]
+	fn test_request_body_sustained_peer_saturation() {
+		tokio::runtime::Builder::new_current_thread().enable_all().build().unwrap().block_on(
+			async {
+				tokio::time::pause();
+				let limiter = RequestBodyLimiter::new();
+				let attacker = "192.0.2.1".parse().unwrap();
+				let client = "192.0.2.2".parse().unwrap();
+
+				// Refill the attacker's slots after each deadline, as replacement
+				// streams on either existing or new connections would do.
+				for _ in 0..10 {
+					let mut stalled = Vec::new();
+					for _ in 0..MAX_BODY_READS_PER_PEER {
+						let mut read =
+							Box::pin(read_request_body(stalled_body(), None, attacker, &limiter));
+						assert!(futures_util::poll!(&mut read).is_pending());
+						stalled.push(read);
+					}
+					for _ in 0..MAX_CONCURRENT_BODY_READS {
+						let err = read_request_body(stalled_body(), None, attacker, &limiter)
+							.await
+							.unwrap_err();
+						assert_eq!(err.code, GRPC_STATUS_UNAVAILABLE);
+						assert_eq!(err.message, "Too many concurrent requests");
+					}
+
+					// A correctly signed request can still collect its body and
+					// authenticate while the attacker's bodies remain unfinished.
+					let body = encode_grpc_frame(&[]);
+					let timestamp = std::time::SystemTime::now()
+						.duration_since(std::time::UNIX_EPOCH)
+						.unwrap()
+						.as_secs();
+					let hmac = compute_hmac("test_key", timestamp, &body);
+					let req = create_test_request(Some(format!("HMAC {timestamp}:{hmac}")));
+					let received = read_request_body(
+						http_body_util::Full::new(body.clone()),
+						Some(body.len() as u64),
+						client,
+						&limiter,
+					)
+					.await
+					.unwrap();
+					assert_eq!(received, body);
+					assert!(validate_auth(&req, "test_key", &received).is_ok());
+
+					tokio::time::advance(REQUEST_BODY_TIMEOUT).await;
+					for read in stalled {
+						let err = read.await.unwrap_err();
+						assert_eq!(err.code, GRPC_STATUS_UNAVAILABLE);
+						assert_eq!(err.message, "Timed out reading request body");
+					}
+					assert!(limiter.active.lock().unwrap().is_empty());
+				}
+			},
+		);
+	}
+
+	#[test]
+	fn test_request_body_global_limit_and_peer_cleanup() {
+		let limiter = RequestBodyLimiter::new();
+		let mut permits = Vec::new();
+		for i in 1..=MAX_CONCURRENT_BODY_READS {
+			permits.push(limiter.try_acquire(IpAddr::from([192, 0, 2, i as u8])).unwrap());
+		}
+		let next_peer = "192.0.2.100".parse().unwrap();
+		assert!(limiter.try_acquire(next_peer).is_err());
+		assert_eq!(limiter.active.lock().unwrap().len(), MAX_CONCURRENT_BODY_READS);
+		permits.pop();
+		permits.push(limiter.try_acquire(next_peer).unwrap());
+		drop(permits);
+		assert!(limiter.active.lock().unwrap().is_empty());
+
+		let peer = "192.0.2.1".parse().unwrap();
+		let mapped_peer = "::ffff:192.0.2.1".parse().unwrap();
+		let _first = limiter.try_acquire(peer).unwrap();
+		let _second = limiter.try_acquire(mapped_peer).unwrap();
+		assert!(limiter.try_acquire(peer).is_err());
+		assert!(limiter.try_acquire(mapped_peer).is_err());
+	}
+
+	#[test]
+	fn test_request_body_releases_slots_on_cancel_and_error() {
+		tokio::runtime::Builder::new_current_thread().enable_all().build().unwrap().block_on(
+			async {
+				let limiter = RequestBodyLimiter::new();
+				let peer = "192.0.2.1".parse().unwrap();
+				let mut read = Box::pin(read_request_body(stalled_body(), None, peer, &limiter));
+				assert!(futures_util::poll!(&mut read).is_pending());
+				drop(read);
+				assert!(limiter.active.lock().unwrap().is_empty());
+
+				let err = read_request_body(
+					http_body_util::Full::new(bytes::Bytes::from_static(b"short")),
+					Some(10),
+					peer,
+					&limiter,
+				)
+				.await
+				.unwrap_err();
+				assert_eq!(err.code, GRPC_STATUS_INVALID_ARGUMENT);
+				assert!(limiter.active.lock().unwrap().is_empty());
+
+				let err = read_request_body(
+					http_body_util::Full::new(bytes::Bytes::from(vec![0; MAX_BODY_SIZE + 1])),
+					None,
+					peer,
+					&limiter,
+				)
+				.await
+				.unwrap_err();
+				assert_eq!(err.code, GRPC_STATUS_INVALID_ARGUMENT);
+				assert!(limiter.active.lock().unwrap().is_empty());
+			},
+		);
+	}
+
 	fn compute_hmac(api_key: &str, timestamp: u64, body: &[u8]) -> String {
 		compute_auth_hmac(api_key, timestamp, body).to_string()
 	}