Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
133 changes: 34 additions & 99 deletions src/proto/h2/upgrade.rs
Original file line number Diff line number Diff line change
@@ -1,11 +1,8 @@
use std::future::Future;
use std::io::Cursor;
use std::pin::Pin;
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::Arc;
use std::task::{Context, Poll};

use atomic_waker::AtomicWaker;
use bytes::{Buf, Bytes};
use futures_channel::{mpsc, oneshot};
use futures_core::{ready, Stream};
Expand All @@ -23,23 +20,18 @@ pub(super) fn pair<B>(
) -> (H2Upgraded, UpgradedSendStreamTask<B>) {
let (tx, rx) = mpsc::channel(1);
let (error_tx, error_rx) = oneshot::channel();
let close_notify = Arc::new(UpgradedCloseNotify::new());

(
H2Upgraded {
send_stream: UpgradedSendStreamBridge {
tx,
error_rx,
close_notify: close_notify.clone(),
},
send_stream: UpgradedSendStreamBridge { tx, error_rx },
recv_stream,
ping,
buf: Bytes::new(),
},
UpgradedSendStreamTask {
h2_tx: send_stream,
rx,
close_notify,
buffered: None,
error_tx: Some(error_tx),
},
)
Expand All @@ -55,46 +47,6 @@ pub(super) struct H2Upgraded {
struct UpgradedSendStreamBridge {
tx: mpsc::Sender<Cursor<Box<[u8]>>>,
error_rx: oneshot::Receiver<crate::Error>,
close_notify: Arc<UpgradedCloseNotify>,
}

impl Drop for UpgradedSendStreamBridge {
fn drop(&mut self) {
self.close_notify.close();
}
}

struct UpgradedCloseNotify {
closed: AtomicBool,
task: AtomicWaker,
}

impl UpgradedCloseNotify {
fn new() -> Self {
Self {
closed: AtomicBool::new(false),
task: AtomicWaker::new(),
}
}

fn close(&self) {
self.closed.store(true, Ordering::Release);
self.task.wake();
}

fn poll_closed(&self, cx: &mut Context<'_>) -> Poll<()> {
if self.closed.load(Ordering::Acquire) {
return Poll::Ready(());
}

self.task.register(cx.waker());

if self.closed.load(Ordering::Acquire) {
Poll::Ready(())
} else {
Poll::Pending
}
}
}

pin_project! {
Expand All @@ -104,7 +56,7 @@ pin_project! {
h2_tx: SendStream<SendBuf<B>>,
#[pin]
rx: mpsc::Receiver<Cursor<Box<[u8]>>>,
close_notify: Arc<UpgradedCloseNotify>,
buffered: Option<Cursor<Box<[u8]>>>,
error_tx: Option<oneshot::Sender<crate::Error>>,
}
}
Expand All @@ -123,35 +75,6 @@ where
// one of the sides hanging up, so the task doesn't live around
// longer than it's meant to.
loop {
// we don't have the next chunk of data yet, so just reserve 1 byte to make
// sure there's some capacity available. h2 will handle the capacity management
// for the actual body chunk.
me.h2_tx.reserve_capacity(1);

let h2_has_capacity = if me.h2_tx.capacity() == 0 {
// poll_capacity oddly needs a loop
loop {
match me.h2_tx.poll_capacity(cx) {
Poll::Ready(Some(Ok(0))) => {}
Poll::Ready(Some(Ok(_))) => break true,
Poll::Ready(Some(Err(e))) => {
return Poll::Ready(Err(crate::Error::new_body_write(e)))
}
Poll::Ready(None) => {
// None means the stream is no longer in a
// streaming state, we either finished it
// somehow, or the remote reset us.
return Poll::Ready(Err(crate::Error::new_body_write(
"send stream capacity unexpectedly closed",
)));
}
Poll::Pending => break false,
}
}
} else {
true
};

match me.h2_tx.poll_reset(cx) {
Poll::Ready(Ok(reason)) => {
trace!("stream received RST_STREAM: {:?}", reason);
Expand All @@ -165,30 +88,43 @@ where
Poll::Pending => (),
}

// If h2 has no capacity, don't pull another item from the mpsc
// receiver. That would free a channel slot and let the writer
// enqueue more data without h2 backpressure.
//
// Still allow the task to finish once the upgraded write side is
// gone and the mpsc queue is empty.
if !h2_has_capacity {
// `size_hint` reads the queued message count without popping,
// so an accepted write stays queued until h2 capacity returns.
if me.rx.size_hint().0 == 0 && me.close_notify.poll_closed(cx).is_ready() {
me.h2_tx
.send_data(SendBuf::None, true)
.map_err(crate::Error::new_body_write)?;
return Poll::Ready(Ok(()));
// A write taken from the mpsc receiver waits here for h2 capacity,
// and the next one isn't pulled until it has been handed to h2, so
// the writer still sees h2 backpressure.
if me.buffered.is_some() {
// poll_capacity oddly needs a loop
while me.h2_tx.capacity() == 0 {
match ready!(me.h2_tx.poll_capacity(cx)) {
Some(Ok(0)) => {}
Some(Ok(_)) => break,
Some(Err(e)) => return Poll::Ready(Err(crate::Error::new_body_write(e))),
None => {
// None means the stream is no longer in a
// streaming state, we either finished it
// somehow, or the remote reset us.
return Poll::Ready(Err(crate::Error::new_body_write(
"send stream capacity unexpectedly closed",
)));
}
}
}

return Poll::Pending;
let cursor = me.buffered.take().expect("checked is_some above");
me.h2_tx
.send_data(SendBuf::Cursor(cursor), false)
.map_err(crate::Error::new_body_write)?;
continue;
}

match me.rx.as_mut().poll_next(cx) {
Poll::Ready(Some(cursor)) => {
me.h2_tx
.send_data(SendBuf::Cursor(cursor), false)
.map_err(crate::Error::new_body_write)?;
// Only reserve capacity once there is something to send.
// Reserving while idle, even a single byte, pins that
// capacity on the connection-level window (#4003). As in
// `PipeToSendStream`, h2 raises the request to the
// buffered length inside `send_data`.
me.h2_tx.reserve_capacity(1);
*me.buffered = Some(cursor);
}
Poll::Ready(None) => {
me.h2_tx
Expand Down Expand Up @@ -328,7 +264,6 @@ impl Write for H2Upgraded {
cx: &mut Context<'_>,
) -> Poll<Result<(), std::io::Error>> {
self.send_stream.tx.close_channel();
self.send_stream.close_notify.close();
match Pin::new(&mut self.send_stream.error_rx).poll(cx) {
Poll::Ready(Ok(reason)) => Poll::Ready(Err(io_error(reason))),
Poll::Ready(Err(_task_dropped)) => Poll::Ready(Ok(())),
Expand Down
112 changes: 112 additions & 0 deletions tests/client.rs
Original file line number Diff line number Diff line change
Expand Up @@ -3315,6 +3315,118 @@ mod conn {
drop(tx_a);
let _ = tokio::time::timeout(Duration::from_secs(5), a_handle).await;
}

// https://github.com/hyperium/hyper/issues/4003, for HTTP/2 CONNECT
//
// Like `h2_idle_stream_does_not_pin_connection_window`, but the idle
// stream is the send side of an `Upgraded` tunnel. It must not reserve
// connection-level flow control capacity while it has nothing to write.
#[tokio::test]
async fn h2_idle_upgraded_does_not_pin_connection_window() {
// One byte short of the initial connection-level window.
const STREAM_A_LEN: usize = 65534;

let (client_io, server_io, _) = setup_duplex_test_server();
let (stream_a_full_tx, stream_a_full_rx) = oneshot::channel::<()>();
let (stream_a_done_tx, stream_a_done_rx) = oneshot::channel::<()>();
let (stream_b_got_tx, stream_b_got_rx) = oneshot::channel::<usize>();

// Raw h2 server that never calls `release_capacity`, so it never
// sends a connection-level WINDOW_UPDATE.
tokio::spawn(async move {
let mut h2 = h2::server::handshake(server_io).await.unwrap();
let mut stream_a_full_tx = Some(stream_a_full_tx);
let mut stream_a_done_rx = Some(stream_a_done_rx);
let mut stream_b_got_tx = Some(stream_b_got_tx);
while let Some(result) = h2.accept().await {
let (req, mut respond) = result.unwrap();
if req.method() == Method::CONNECT {
let full_tx = stream_a_full_tx.take().unwrap();
let done_rx = stream_a_done_rx.take().unwrap();
tokio::spawn(async move {
let _send = respond.send_response(Response::new(()), false).unwrap();
let mut body = req.into_body();
let mut received = 0usize;
while received < STREAM_A_LEN {
match body.data().await {
Some(Ok(frame)) => received += frame.len(),
_ => return,
}
}
let _ = full_tx.send(());
// Hold on to the recv stream, dropping it would release
// its capacity and send a WINDOW_UPDATE.
let _ = done_rx.await;
drop(body);
});
} else {
let got_tx = stream_b_got_tx.take().unwrap();
tokio::spawn(async move {
let mut body = req.into_body();
let mut received = 0usize;
if let Some(Ok(frame)) = body.data().await {
received += frame.len();
}
let _ = got_tx.send(received);
let mut send = respond.send_response(Response::new(()), false).unwrap();
let _ = send.send_data(Bytes::from_static(b"ok"), true);
});
}
}
});

let io = TokioIo::new(client_io);
let (mut client, conn) = conn::http2::Builder::new(TokioExecutor)
.handshake::<_, Full<Bytes>>(io)
.await
.expect("http handshake");
tokio::spawn(async move {
let _ = conn.await;
});

// Stream A: a tunnel that writes STREAM_A_LEN bytes and then stays
// open, leaving one byte of connection window.
let req_a = Request::connect("localhost")
.body(Full::new(Bytes::new()))
.unwrap();
let res_a = client.send_request(req_a).await.expect("send_request A");
assert_eq!(res_a.status(), StatusCode::OK);
let mut upgraded = TokioIo::new(hyper::upgrade::on(res_a).await.unwrap());
upgraded
.write_all(&vec![b'A'; STREAM_A_LEN])
.await
.expect("write to tunnel");

tokio::time::timeout(Duration::from_secs(5), stream_a_full_rx)
.await
.expect("server should receive all of stream A in time")
.expect("stream_a_full_rx");

// Let the tunnel's send task park waiting for more writes.
for _ in 0..16 {
tokio::task::yield_now().await;
}

// Stream B: one byte of body, which needs the last byte of window.
let req_b = Request::post("http://localhost/b")
.body(Full::new(Bytes::from_static(b"b")))
.unwrap();
let b_fut = client.send_request(req_b);

let received_b = tokio::time::timeout(Duration::from_secs(5), stream_b_got_rx)
.await
.expect("stream B must reach the server even while the tunnel is idle")
.expect("stream_b_got_rx");
assert_eq!(
received_b, 1,
"stream B should deliver its single body byte"
);

let _ = tokio::time::timeout(Duration::from_secs(5), b_fut).await;

let _ = stream_a_done_tx.send(());
drop(upgraded);
}
}

trait FutureHyperExt: TryFuture {
Expand Down