fix(peer): bound aggregate control response buffering

Read the four-byte response length explicitly, acquire bytes from one 32 MiB connector-wide budget before buffering, and require exact body plus EOF. One-way requests now check EOF without allocating an 8 MiB decoder.

Test Plan:
- just test
- focused complete, oversized, truncated, trailing, one-way, and shared-budget tests
- git diff --check
This commit is contained in:
2026-09-12 12:41:39 +02:00
parent 145610c0a1
commit dfc858d0e2
2 changed files with 290 additions and 37 deletions
+216 -34
View File
@@ -6,7 +6,8 @@ use std::{
time::Duration,
};
use futures::{SinkExt, StreamExt};
use bytes::Bytes;
use futures::SinkExt;
use if_addrs::{IfAddr, Interface, get_if_addrs};
use lanspread_proto::{
ChangeHint,
@@ -18,8 +19,9 @@ use lanspread_proto::{
Request,
Response,
};
use tokio::io::{AsyncRead, AsyncReadExt as _};
use tokio_util::{
codec::{FramedRead, FramedWrite, LengthDelimitedCodec},
codec::{FramedWrite, LengthDelimitedCodec},
sync::CancellationToken,
};
@@ -72,25 +74,13 @@ async fn send_oneway_request(
let mut conn = connector.connect(endpoint).await?;
let stream = conn.open_bidirectional_stream().await?;
let (rx, tx) = stream.split();
let mut framed_rx = FramedRead::new(rx, control_codec());
let (mut rx, tx) = stream.split();
let mut framed_tx = FramedWrite::new(tx, control_codec());
framed_tx.send(request.encode()?).await?;
framed_tx.close().await?;
match framed_rx.next().await {
None => {}
Some(Ok(frame)) => {
let response = Response::decode(frame.freeze())?;
eyre::bail!(
"peer {} unexpectedly responded to one-way request: {response:?}",
endpoint.addr
);
}
Some(Err(error)) => return Err(error.into()),
}
Ok(())
require_oneway_response_eof(&mut rx, endpoint.addr).await
})
.await
}
@@ -124,27 +114,13 @@ async fn exchange_request(
let mut conn = connector.connect(endpoint).await?;
let stream = conn.open_bidirectional_stream().await?;
let (rx, tx) = stream.split();
let mut framed_rx = FramedRead::new(rx, control_codec());
let (mut rx, tx) = stream.split();
let mut framed_tx = FramedWrite::new(tx, control_codec());
framed_tx.send(request.encode()?).await?;
framed_tx.close().await?;
let frame = framed_rx
.next()
.await
.ok_or_else(|| eyre::eyre!("peer {} returned no control response", endpoint.addr))??;
let response = Response::decode(frame.freeze())?;
match framed_rx.next().await {
None => Ok(response),
Some(Ok(_)) => eyre::bail!(
"peer {} returned more than one control response frame",
endpoint.addr
),
Some(Err(error)) => Err(error.into()),
}
read_control_response(&mut rx, connector, endpoint.addr).await
})
.await
}
@@ -185,6 +161,106 @@ fn control_codec() -> LengthDelimitedCodec {
.new_codec()
}
const CONTROL_FRAME_LENGTH_BYTES: usize = size_of::<u32>();
const CONTROL_RESPONSE_READ_CHUNK_BYTES: usize = 64 * 1024;
/// Reads one response without allowing its untrusted length prefix to reserve
/// the full frame independently of every other outbound request.
async fn read_control_response<R>(
reader: &mut R,
connector: &QuicConnector,
peer_addr: SocketAddr,
) -> eyre::Result<Response>
where
R: AsyncRead + Unpin,
{
let mut length_bytes = [0; CONTROL_FRAME_LENGTH_BYTES];
let mut length_bytes_read = 0;
while length_bytes_read < length_bytes.len() {
let read = reader.read(&mut length_bytes[length_bytes_read..]).await?;
if read == 0 {
if length_bytes_read == 0 {
eyre::bail!("peer {peer_addr} returned no control response");
}
eyre::bail!(
"peer {peer_addr} returned a truncated control response length prefix ({length_bytes_read} of {} bytes)",
length_bytes.len()
);
}
length_bytes_read += read;
}
let frame_len_u32 = u32::from_be_bytes(length_bytes);
let frame_len = usize::try_from(frame_len_u32)
.map_err(|_| eyre::eyre!("control response length does not fit usize"))?;
if frame_len > MAX_CONTROL_FRAME_BYTES {
eyre::bail!(
"peer {peer_addr} declared a {frame_len}-byte control response; maximum is {MAX_CONTROL_FRAME_BYTES} bytes"
);
}
let _response_bytes = connector
.reserve_control_response_bytes(frame_len_u32)
.await?;
let mut frame = Vec::new();
while frame.len() < frame_len {
if frame.len() == frame.capacity() {
let requested_capacity = frame
.capacity()
.saturating_mul(2)
.max(CONTROL_RESPONSE_READ_CHUNK_BYTES)
.min(frame_len);
frame
.try_reserve_exact(requested_capacity - frame.len())
.map_err(|error| {
eyre::eyre!("failed to reserve control response bytes: {error}")
})?;
}
let chunk_start = frame.len();
let chunk_end = frame.capacity().min(frame_len);
frame.resize(chunk_end, 0);
let mut body_bytes_read = chunk_start;
while body_bytes_read < chunk_end {
let read = reader.read(&mut frame[body_bytes_read..chunk_end]).await?;
if read == 0 {
frame.truncate(body_bytes_read);
eyre::bail!(
"peer {peer_addr} returned a truncated control response body ({} of {frame_len} bytes)",
frame.len()
);
}
body_bytes_read += read;
}
}
require_control_response_eof(reader, peer_addr).await?;
let response = Response::decode(Bytes::from(frame))?;
Ok(response)
}
async fn require_control_response_eof<R>(reader: &mut R, peer_addr: SocketAddr) -> eyre::Result<()>
where
R: AsyncRead + Unpin,
{
let mut trailing = [0];
if reader.read(&mut trailing).await? != 0 {
eyre::bail!("peer {peer_addr} returned trailing bytes after one control response frame");
}
Ok(())
}
async fn require_oneway_response_eof<R>(reader: &mut R, peer_addr: SocketAddr) -> eyre::Result<()>
where
R: AsyncRead + Unpin,
{
let mut unexpected = [0];
if reader.read(&mut unexpected).await? != 0 {
eyre::bail!("peer {peer_addr} unexpectedly responded to a one-way request");
}
Ok(())
}
async fn run_short_request<T>(
peer_addr: SocketAddr,
cancellation: &CancellationToken,
@@ -311,10 +387,20 @@ mod tests {
time::Duration,
};
use tokio::sync::oneshot;
use lanspread_proto::{
ControlErrorCode,
ControlMessage as _,
MAX_CONTROL_FRAME_BYTES,
Response,
};
use tokio::{
io::{AsyncWriteExt as _, DuplexStream, duplex},
sync::oneshot,
};
use tokio_util::sync::CancellationToken;
use super::run_network_operation;
use super::{read_control_response, require_oneway_response_eof, run_network_operation};
use crate::quic_runtime::QuicConnector;
struct DropProbe(Arc<AtomicBool>);
@@ -328,6 +414,102 @@ mod tests {
SocketAddr::from(([127, 0, 0, 1], 12345))
}
async fn closed_reader(bytes: &[u8]) -> DuplexStream {
let (mut writer, reader) = duplex(bytes.len().saturating_add(1));
writer
.write_all(bytes)
.await
.expect("test response should fit its duplex buffer");
writer
.shutdown()
.await
.expect("test response writer should close");
reader
}
fn framed_response(body: &[u8]) -> Vec<u8> {
let mut frame = u32::try_from(body.len())
.expect("test response length should fit u32")
.to_be_bytes()
.to_vec();
frame.extend_from_slice(body);
frame
}
#[tokio::test]
async fn bounded_response_reader_accepts_one_complete_frame() {
let encoded = Response::Error(ControlErrorCode::InvalidRequest)
.encode()
.expect("test response should encode");
let mut reader = closed_reader(&framed_response(&encoded)).await;
let connector = QuicConnector::unavailable();
assert!(matches!(
read_control_response(&mut reader, &connector, peer_addr())
.await
.expect("one complete response should decode"),
Response::Error(ControlErrorCode::InvalidRequest)
));
}
#[tokio::test]
async fn bounded_response_reader_rejects_oversized_prefix_before_body() {
let declared =
u32::try_from(MAX_CONTROL_FRAME_BYTES + 1).expect("control frame bound should fit u32");
let mut reader = closed_reader(&declared.to_be_bytes()).await;
let connector = QuicConnector::unavailable();
let error = read_control_response(&mut reader, &connector, peer_addr())
.await
.expect_err("oversized response must be rejected");
assert!(format!("{error:#}").contains("maximum is"));
}
#[tokio::test]
async fn bounded_response_reader_rejects_truncated_body() {
let mut response = 10_u32.to_be_bytes().to_vec();
response.extend_from_slice(b"short");
let mut reader = closed_reader(&response).await;
let connector = QuicConnector::unavailable();
let error = read_control_response(&mut reader, &connector, peer_addr())
.await
.expect_err("truncated response must be rejected");
assert!(format!("{error:#}").contains("truncated control response body (5 of 10 bytes)"));
}
#[tokio::test]
async fn bounded_response_reader_rejects_trailing_frame_bytes() {
let encoded = Response::Error(ControlErrorCode::InvalidRequest)
.encode()
.expect("test response should encode");
let mut response = framed_response(&encoded);
response.extend_from_slice(&1_u32.to_be_bytes());
response.push(b'x');
let mut reader = closed_reader(&response).await;
let connector = QuicConnector::unavailable();
let error = read_control_response(&mut reader, &connector, peer_addr())
.await
.expect_err("trailing response frame must be rejected");
assert!(format!("{error:#}").contains("trailing bytes"));
}
#[tokio::test]
async fn one_way_response_path_requires_immediate_eof() {
let mut empty = closed_reader(&[]).await;
require_oneway_response_eof(&mut empty, peer_addr())
.await
.expect("empty response side should be accepted");
let mut response = closed_reader(&[0]).await;
assert!(
require_oneway_response_eof(&mut response, peer_addr())
.await
.is_err()
);
}
#[tokio::test]
async fn cancellation_drops_operation_scope_before_returning() {
let cancellation = CancellationToken::new();
+74 -3
View File
@@ -33,7 +33,10 @@ use s2n_quic_core::{
path::mtu,
time::{Clock, Timestamp},
};
use tokio::task::{JoinError, JoinHandle};
use tokio::{
sync::{OwnedSemaphorePermit, Semaphore},
task::{JoinError, JoinHandle},
};
use tokio_util::sync::{CancellationToken, WaitForCancellationFutureOwned};
use crate::{
@@ -62,6 +65,12 @@ pub(crate) const SERVER_CONTROL_RECEIVE_WINDOW: u64 = MAX_REQUEST_FRAME_BYTES as
/// open request stream may hold its full per-stream allowance.
pub(crate) const SERVER_CONNECTION_RECEIVE_WINDOW: u64 =
SERVER_CONTROL_RECEIVE_WINDOW * MAX_OPEN_BIDIRECTIONAL_STREAMS;
/// Aggregate application buffer budget for concurrently received control responses.
///
/// Response frames may legitimately approach 8 MiB, but discovery and refresh
/// tasks must not each reserve that allowance independently after reading only a
/// length prefix. All connector clones in one network generation share this pool.
pub(crate) const MAX_INFLIGHT_CONTROL_RESPONSE_BYTES: usize = 32 * 1024 * 1024;
pub(crate) fn quic_client_limits() -> eyre::Result<Limits> {
Ok(Limits::default()
@@ -349,6 +358,7 @@ fn join_result(result: Result<(), JoinError>, expected_abort: bool) -> eyre::Res
#[derive(Clone, Debug)]
pub(crate) struct QuicConnector {
client: Option<QuicClient>,
control_response_bytes: Arc<Semaphore>,
}
impl QuicConnector {
@@ -362,9 +372,22 @@ impl QuicConnector {
Ok(PeerConnection::new(client.connect(connect).await?))
}
pub(crate) async fn reserve_control_response_bytes(
&self,
byte_count: u32,
) -> eyre::Result<OwnedSemaphorePermit> {
Arc::clone(&self.control_response_bytes)
.acquire_many_owned(byte_count)
.await
.map_err(|_| eyre::eyre!("control-response byte budget was closed"))
}
#[cfg(test)]
pub(crate) fn unavailable() -> Self {
Self { client: None }
Self {
client: None,
control_response_bytes: Arc::new(Semaphore::new(MAX_INFLIGHT_CONTROL_RESPONSE_BYTES)),
}
}
}
@@ -427,6 +450,7 @@ pub(crate) fn start_quic_client() -> eyre::Result<(QuicClientRuntime, QuicConnec
let endpoint = endpoint_control.take_started()?;
let connector = QuicConnector {
client: Some(client.clone()),
control_response_bytes: Arc::new(Semaphore::new(MAX_INFLIGHT_CONTROL_RESPONSE_BYTES)),
};
Ok((QuicClientRuntime { client, endpoint }, connector))
@@ -468,10 +492,17 @@ impl Drop for PeerConnection {
mod tests {
use std::{sync::Arc, time::Duration};
use lanspread_proto::MAX_CONTROL_FRAME_BYTES;
use tokio::sync::Notify;
use tokio_util::sync::CancellationToken;
use super::{EndpointTask, EndpointTaskControl, start_quic_client};
use super::{
EndpointTask,
EndpointTaskControl,
MAX_INFLIGHT_CONTROL_RESPONSE_BYTES,
QuicConnector,
start_quic_client,
};
#[test]
fn endpoint_control_rejects_take_before_provider_start() {
@@ -490,6 +521,46 @@ mod tests {
.expect("client endpoint should join cleanly");
}
#[tokio::test]
async fn connector_clones_share_and_release_the_control_response_budget() {
assert_eq!(
MAX_INFLIGHT_CONTROL_RESPONSE_BYTES % MAX_CONTROL_FRAME_BYTES,
0
);
let max_frame =
u32::try_from(MAX_CONTROL_FRAME_BYTES).expect("control frame bound should fit u32");
let connector = QuicConnector::unavailable();
let connector_clone = connector.clone();
let mut permits = Vec::new();
for _ in 0..MAX_INFLIGHT_CONTROL_RESPONSE_BYTES / MAX_CONTROL_FRAME_BYTES {
permits.push(
connector
.reserve_control_response_bytes(max_frame)
.await
.expect("configured aggregate budget should admit its exact capacity"),
);
}
let blocked = connector_clone.reserve_control_response_bytes(1);
tokio::pin!(blocked);
tokio::select! {
biased;
result = blocked.as_mut() => {
panic!("response budget exceeded its configured capacity: {result:?}");
}
() = tokio::task::yield_now() => {}
}
let released = permits
.pop()
.expect("the exact-capacity fixture should hold at least one permit");
drop(released);
let reacquired = blocked
.await
.expect("dropping a permit should release shared response capacity");
drop(reacquired);
}
#[tokio::test]
async fn endpoint_owner_signals_and_joins_its_task() {
let stop = CancellationToken::new();