use crate::common;
use bytes::Bytes;
use camber::http::Router;
use camber::runtime;
use futures_util::FutureExt;
use hyper::body::{Body, Frame};
use hyper::server::conn::http1;
use hyper::service::service_fn;
use hyper_util::rt::TokioIo;
use std::convert::Infallible;
use std::future::Future;
use std::pin::Pin;
use std::sync::Arc;
use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
use std::task::{Context, Poll};
use std::time::Duration;
const DOWNSTREAM_WINDOW_BYTES: usize = 1024;
const UPSTREAM_FRAME_BYTES: usize = 64 * 1024;
const UPSTREAM_FRAME_LIMIT: usize = 256;
const FIRST_BOUNDARY: &[u8] = b"[first-upstream-frame]";
const FINAL_BOUNDARY: &[u8] = b"[final-upstream-frame]";
const FRAME_BOUNDARY: &[u8] = b"[ordered-frame:]";
const EVENT_TIMEOUT: Duration = Duration::from_secs(30);
#[derive(Debug, PartialEq, Eq)]
enum BodyPollCheckpoint {
Data(usize),
End(usize),
}
struct BodyControl {
end_at: AtomicUsize,
poll_count: AtomicUsize,
complete: AtomicBool,
}
impl BodyControl {
fn new() -> Self {
Self {
end_at: AtomicUsize::new(UPSTREAM_FRAME_LIMIT),
poll_count: AtomicUsize::new(0),
complete: AtomicBool::new(false),
}
}
}
struct CheckpointBody {
control: Arc<BodyControl>,
checkpoints: tokio::sync::mpsc::UnboundedSender<BodyPollCheckpoint>,
next_frame: usize,
}
impl Body for CheckpointBody {
type Data = Bytes;
type Error = Infallible;
fn poll_frame(
mut self: Pin<&mut Self>,
_context: &mut Context<'_>,
) -> Poll<Option<Result<Frame<Self::Data>, Self::Error>>> {
let end_at = self.control.end_at.load(Ordering::Acquire);
self.control.poll_count.fetch_add(1, Ordering::AcqRel);
match self.next_frame == end_at {
true => {
self.control.complete.store(true, Ordering::Release);
self.checkpoints
.send(BodyPollCheckpoint::End(end_at))
.expect("body checkpoint receiver remains live");
Poll::Ready(None)
}
false => {
let index = self.next_frame;
self.next_frame += 1;
self.checkpoints
.send(BodyPollCheckpoint::Data(index))
.expect("body checkpoint receiver remains live");
Poll::Ready(Some(Ok(Frame::data(ordered_frame(index, end_at)))))
}
}
}
}
#[derive(Debug)]
struct ConnectionPollReport {
body_polls_before: usize,
body_polls_after: usize,
connection_finished: bool,
}
struct GatedConnection<F> {
connection: Pin<Box<F>>,
permits: tokio::sync::mpsc::UnboundedReceiver<()>,
reports: tokio::sync::mpsc::UnboundedSender<ConnectionPollReport>,
body_poll_count: Arc<BodyControl>,
continuous: Arc<AtomicBool>,
}
impl<F> Future for GatedConnection<F>
where
F: Future,
{
type Output = F::Output;
fn poll(mut self: Pin<&mut Self>, context: &mut Context<'_>) -> Poll<Self::Output> {
if self.continuous.load(Ordering::Acquire) {
return self.connection.as_mut().poll(context);
}
match self.permits.poll_recv(context) {
Poll::Ready(Some(())) => {}
Poll::Ready(None) | Poll::Pending => return Poll::Pending,
}
let body_polls_before = self.body_poll_count.poll_count.load(Ordering::Acquire);
let result = self.connection.as_mut().poll(context);
let report = ConnectionPollReport {
body_polls_before,
body_polls_after: self.body_poll_count.poll_count.load(Ordering::Acquire),
connection_finished: result.is_ready(),
};
self.reports
.send(report)
.expect("connection poll report receiver remains live");
if result.is_pending() {
context.waker().wake_by_ref();
}
result
}
}
struct ConnectionController {
permits: tokio::sync::mpsc::UnboundedSender<()>,
reports: tokio::sync::mpsc::UnboundedReceiver<ConnectionPollReport>,
continuous: Arc<AtomicBool>,
}
impl ConnectionController {
async fn poll_once(&mut self) -> ConnectionPollReport {
self.permits
.send(())
.expect("gated upstream connection remains live");
self.reports
.recv()
.await
.expect("gated connection reports every permitted poll")
}
fn run_continuously(&self) {
self.continuous.store(true, Ordering::Release);
self.permits
.send(())
.expect("gated upstream connection remains live");
}
}
async fn event<F>(name: &str, future: F) -> F::Output
where
F: Future,
{
tokio::time::timeout(EVENT_TIMEOUT, future)
.await
.unwrap_or_else(|_| panic!("timed out waiting for {name}"))
}
fn ordered_frame(index: usize, end_at: usize) -> Bytes {
match (index, index + 1 == end_at) {
(0, _) => padded_frame(FIRST_BOUNDARY, index),
(_, true) => Bytes::from_static(FINAL_BOUNDARY),
_ => padded_frame(FRAME_BOUNDARY, index),
}
}
fn padded_frame(boundary: &[u8], index: usize) -> Bytes {
let mut bytes = Vec::with_capacity(UPSTREAM_FRAME_BYTES);
bytes.extend_from_slice(boundary);
bytes.extend_from_slice(&(index as u64).to_be_bytes());
bytes.resize(UPSTREAM_FRAME_BYTES, (index % 251) as u8);
bytes.into()
}
fn expected_body(frame_count: usize) -> Vec<u8> {
(0..frame_count)
.flat_map(|index| ordered_frame(index, frame_count))
.collect()
}
async fn spawn_gated_upstream(
stream: tokio::net::TcpStream,
body_control: Arc<BodyControl>,
checkpoints: tokio::sync::mpsc::UnboundedSender<BodyPollCheckpoint>,
) -> (
ConnectionController,
tokio::task::JoinHandle<Result<(), hyper::Error>>,
) {
let body = CheckpointBody {
control: Arc::clone(&body_control),
checkpoints,
next_frame: 0,
};
let body = Arc::new(std::sync::Mutex::new(Some(body)));
let service = service_fn(move |_request| {
let body = body
.lock()
.unwrap_or_else(|error| error.into_inner())
.take()
.expect("the upstream serves one request");
async move { Ok::<_, Infallible>(hyper::Response::new(body)) }
});
let connection = http1::Builder::new().serve_connection(TokioIo::new(stream), service);
let (permit_tx, permit_rx) = tokio::sync::mpsc::unbounded_channel();
let (report_tx, report_rx) = tokio::sync::mpsc::unbounded_channel();
let continuous = Arc::new(AtomicBool::new(false));
let gated = GatedConnection {
connection: Box::pin(connection),
permits: permit_rx,
reports: report_tx,
body_poll_count: body_control,
continuous: Arc::clone(&continuous),
};
let task = tokio::spawn(gated);
(
ConnectionController {
permits: permit_tx,
reports: report_rx,
continuous,
},
task,
)
}
async fn response_head_while_driving(
response: h2::client::ResponseFuture,
upstream: &mut ConnectionController,
) -> ::http::Response<h2::RecvStream> {
let mut response = Box::pin(response);
loop {
if let Some(result) = response.as_mut().now_or_never() {
return result.expect("proxy returns an HTTP/2 response head");
}
let report = upstream.poll_once().await;
assert!(
!report.connection_finished,
"upstream connection ended early"
);
}
}
async fn data_while_driving(
body: &mut h2::RecvStream,
upstream: &mut ConnectionController,
) -> Option<Result<Bytes, h2::Error>> {
let mut data = Box::pin(body.data());
loop {
if let Some(result) = data.as_mut().now_or_never() {
return result;
}
let report = upstream.poll_once().await;
assert!(
!report.connection_finished,
"upstream connection ended early"
);
}
}
fn drain_data_checkpoints(
checkpoints: &mut tokio::sync::mpsc::UnboundedReceiver<BodyPollCheckpoint>,
) -> Vec<usize> {
std::iter::from_fn(|| checkpoints.try_recv().ok())
.map(|checkpoint| match checkpoint {
BodyPollCheckpoint::Data(index) => index,
BodyPollCheckpoint::End(count) => {
panic!("upstream body ended at {count} frames before flow-control release")
}
})
.collect()
}
struct FlowControlScenario {
body_control: Arc<BodyControl>,
checkpoints: tokio::sync::mpsc::UnboundedReceiver<BodyPollCheckpoint>,
client: h2::client::SendRequest<Bytes>,
downstream_connection: tokio::task::JoinHandle<Result<(), h2::Error>>,
upstream: ConnectionController,
upstream_task: tokio::task::JoinHandle<Result<(), hyper::Error>>,
body: h2::RecvStream,
actual: Vec<u8>,
}
async fn downstream_request(
proxy_addr: std::net::SocketAddr,
) -> (
h2::client::SendRequest<Bytes>,
tokio::task::JoinHandle<Result<(), h2::Error>>,
h2::client::ResponseFuture,
) {
let tcp = event(
"the local proxy connection",
tokio::net::TcpStream::connect(proxy_addr),
)
.await
.expect("connect local HTTP/2 client");
let mut builder = h2::client::Builder::new();
builder
.initial_window_size(DOWNSTREAM_WINDOW_BYTES as u32)
.initial_connection_window_size(DOWNSTREAM_WINDOW_BYTES as u32);
let (mut client, connection) = event(
"the local HTTP/2 handshake",
builder.handshake::<_, Bytes>(tcp),
)
.await
.expect("complete local HTTP/2 handshake");
let downstream_connection = tokio::spawn(connection);
let request = ::http::Request::get(format!("http://{proxy_addr}/api/flow"))
.version(::http::Version::HTTP_2)
.body(())
.expect("build HTTP/2 flow-control request");
client = event("HTTP/2 request readiness", client.ready())
.await
.expect("HTTP/2 request sender becomes ready");
let (response, _) = client
.send_request(request, true)
.expect("send HTTP/2 flow-control request");
(client, downstream_connection, response)
}
async fn start_flow_control_scenario() -> FlowControlScenario {
let listener = tokio::net::TcpListener::bind("127.0.0.1:0")
.await
.expect("bind test-owned upstream listener");
let upstream_addr = listener
.local_addr()
.expect("read test-owned upstream address");
let body_control = Arc::new(BodyControl::new());
let (checkpoint_tx, checkpoints) = tokio::sync::mpsc::unbounded_channel();
let mut proxy = Router::new();
proxy.proxy_stream("/api", &format!("http://{upstream_addr}"));
let proxy_addr = common::spawn_server(proxy);
let (client, downstream_connection, response) = downstream_request(proxy_addr).await;
let (upstream_stream, _) = event("the proxy's upstream connection", listener.accept())
.await
.expect("accept the proxy's upstream connection");
let (mut upstream, upstream_task) =
spawn_gated_upstream(upstream_stream, Arc::clone(&body_control), checkpoint_tx).await;
let response = event(
"the proxied response head",
response_head_while_driving(response, &mut upstream),
)
.await;
assert_eq!(response.status(), 200);
let mut body = response.into_body();
let first = event(
"the first decoded downstream DATA frame",
data_while_driving(&mut body, &mut upstream),
)
.await
.expect("stream contains a first DATA frame")
.expect("first DATA frame decodes successfully");
assert!(
first.starts_with(FIRST_BOUNDARY),
"first decoded DATA omitted the first-frame boundary"
);
assert!(
!body_control.complete.load(Ordering::Acquire),
"upstream completed before the first downstream DATA frame was observed"
);
FlowControlScenario {
body_control,
checkpoints,
client,
downstream_connection,
upstream,
upstream_task,
body,
actual: first.to_vec(),
}
}
async fn fill_downstream_window(scenario: &mut FlowControlScenario) {
while scenario.actual.len() < DOWNSTREAM_WINDOW_BYTES {
let chunk = event(
"DATA that consumes the configured downstream window",
data_while_driving(&mut scenario.body, &mut scenario.upstream),
)
.await
.expect("stream remains open while filling its receive window")
.expect("window-filling DATA decodes successfully");
scenario.actual.extend_from_slice(&chunk);
}
assert_eq!(
scenario.actual.len(),
DOWNSTREAM_WINDOW_BYTES,
"decoded DATA must consume the deliberately small receive window exactly"
);
}
async fn await_upstream_stall(scenario: &mut FlowControlScenario) -> usize {
let blocked_report = event("upstream transport saturation", async {
loop {
let report = scenario.upstream.poll_once().await;
match report.body_polls_after == report.body_polls_before {
true => return report,
false => {}
}
}
})
.await;
assert!(!blocked_report.connection_finished);
let emitted_before_release = blocked_report.body_polls_after;
assert!(
emitted_before_release > 1,
"the direct upstream body must enter transport flow before it stalls"
);
assert!(
!scenario.body_control.complete.load(Ordering::Acquire),
"upstream exhausted its safety frame limit instead of being flow-controlled"
);
emitted_before_release
}
async fn assert_capacity_withheld(
scenario: &mut FlowControlScenario,
emitted_before_release: usize,
) {
let checkpoints = drain_data_checkpoints(&mut scenario.checkpoints);
assert_eq!(
checkpoints,
(0..emitted_before_release).collect::<Vec<_>>(),
"upstream body poll checkpoints were missing or reordered"
);
let withheld = event(
"the explicit withheld-capacity body-poll probe",
scenario.upstream.poll_once(),
)
.await;
assert_eq!(
withheld.body_polls_after, withheld.body_polls_before,
"proxy polled another upstream body frame while downstream capacity was withheld"
);
assert!(
matches!(
scenario.checkpoints.try_recv(),
Err(tokio::sync::mpsc::error::TryRecvError::Empty)
),
"a body poll checkpoint appeared while downstream capacity was withheld"
);
}
async fn release_downstream_capacity(
scenario: &mut FlowControlScenario,
emitted_before_release: usize,
) -> usize {
let final_frame_count = emitted_before_release + 2;
scenario
.body_control
.end_at
.store(final_frame_count, Ordering::Release);
scenario
.body
.flow_control()
.release_capacity(scenario.actual.len())
.expect("release retained downstream HTTP/2 capacity");
let resumed = event("upstream body polling after flow-control release", async {
loop {
let report = scenario.upstream.poll_once().await;
match report.body_polls_after > report.body_polls_before {
true => return report,
false => {
let chunk = scenario
.body
.data()
.await
.expect("buffered proxy DATA precedes resumed upstream polling")
.expect("buffered proxy DATA decodes during flow-control release");
scenario.actual.extend_from_slice(&chunk);
scenario
.body
.flow_control()
.release_capacity(chunk.len())
.expect("return capacity while draining the proxy transport");
}
}
}
})
.await;
assert!(!resumed.connection_finished);
assert_eq!(
event(
"the resumed body poll checkpoint",
scenario.checkpoints.recv()
)
.await,
Some(BodyPollCheckpoint::Data(emitted_before_release)),
"downstream capacity release did not resume at the exact upstream boundary"
);
final_frame_count
}
async fn finish_flow_control_scenario(mut scenario: FlowControlScenario, final_frame_count: usize) {
scenario.upstream.run_continuously();
while let Some(chunk) = event("ordered downstream completion", scenario.body.data()).await {
let chunk = chunk.expect("downstream DATA decodes after capacity release");
scenario.actual.extend_from_slice(&chunk);
scenario
.body
.flow_control()
.release_capacity(chunk.len())
.expect("return consumed HTTP/2 capacity");
}
assert_eq!(
event(
"the terminal upstream body checkpoint",
scenario.checkpoints.recv()
)
.await,
Some(BodyPollCheckpoint::Data(final_frame_count - 1))
);
assert_eq!(
event("upstream body completion", scenario.checkpoints.recv()).await,
Some(BodyPollCheckpoint::End(final_frame_count))
);
assert_eq!(
scenario.actual,
expected_body(final_frame_count),
"proxied bytes or frame boundaries changed"
);
assert!(scenario.actual.starts_with(FIRST_BOUNDARY));
assert!(scenario.actual.ends_with(FINAL_BOUNDARY));
drop(scenario.client);
scenario.downstream_connection.abort();
scenario.upstream_task.abort();
runtime::request_shutdown();
}
#[camber::test]
async fn adversarial_proxy_stream_obeys_downstream_http2_flow_control() {
let mut scenario = start_flow_control_scenario().await;
fill_downstream_window(&mut scenario).await;
let emitted_before_release = await_upstream_stall(&mut scenario).await;
assert_capacity_withheld(&mut scenario, emitted_before_release).await;
let final_frame_count =
release_downstream_capacity(&mut scenario, emitted_before_release).await;
finish_flow_control_scenario(scenario, final_frame_count).await;
}