diff --git a/web-transport-proto/src/capsule.rs b/web-transport-proto/src/capsule.rs index a018b3cd..ca62781a 100644 --- a/web-transport-proto/src/capsule.rs +++ b/web-transport-proto/src/capsule.rs @@ -3,12 +3,9 @@ use std::sync::Arc; use bytes::{Buf, BufMut, Bytes, BytesMut}; use tokio::io::{AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt}; -use crate::{VarInt, VarIntUnexpectedEnd, MAX_FRAME_SIZE}; +use crate::{Frame, VarInt, VarIntUnexpectedEnd, MAX_FRAME_SIZE}; -// The spec (draft-ietf-webtrans-http3-06) says the type is 0x2843, which would -// varint-encode to 0x68 0x43. However, actual wire data shows 0x43 0x28 which -// decodes to 808. There may be a discrepancy in implementations or specs. -// Using 0x2843 as specified in the standard. +// CloseWebTransportSession capsule type (draft-ietf-webtrans-http3-06). const CLOSE_WEBTRANSPORT_SESSION_TYPE: u64 = 0x2843; #[derive(Debug, Clone, PartialEq, Eq)] @@ -81,9 +78,10 @@ impl Capsule { /// /// Returns `Ok(None)` if the stream is cleanly closed (EOF before any bytes). pub async fn read(stream: &mut S) -> Result, CapsuleError> { - let typ = match VarInt::read(stream).await { - Ok(v) => v, - Err(_) => return Ok(None), // Clean EOF + let typ = match VarInt::read_optional(stream).await { + Ok(Some(v)) => v, + Ok(None) => return Ok(None), // Clean EOF + Err(_) => return Err(CapsuleError::UnexpectedEnd), }; let length = VarInt::read(stream) .await @@ -216,9 +214,6 @@ pub enum CapsuleError { #[error("message too long")] MessageTooLong, - #[error("unknown capsule type: {0:?}")] - UnknownType(VarInt), - #[error("varint decode error: {0:?}")] VarInt(#[from] VarIntUnexpectedEnd), @@ -232,6 +227,86 @@ impl From for CapsuleError { } } +/// Reads capsules from an HTTP/3 CONNECT stream where capsule protocol +/// data is carried inside DATA frames (RFC 9297 Section 3.2). +/// +/// Handles capsules split across multiple DATA frames and multiple +/// capsules within a single DATA frame. +pub struct Http3CapsuleReader { + stream: S, + buf: BytesMut, +} + +impl Http3CapsuleReader { + pub fn new(stream: S) -> Self { + Self { + stream, + buf: BytesMut::new(), + } + } + + /// Read the next capsule. Returns `Ok(None)` on clean EOF. + pub async fn read(&mut self) -> Result, CapsuleError> { + loop { + if !self.buf.is_empty() { + let mut slice = &self.buf[..]; + match Capsule::decode(&mut slice) { + Ok(capsule) => { + self.buf.advance(self.buf.len() - slice.len()); + return Ok(Some(capsule)); + } + Err(CapsuleError::UnexpectedEnd | CapsuleError::VarInt(_)) => {} + Err(e) => return Err(e), + } + } + + if !self.read_data_frame().await? { + return if self.buf.is_empty() { + Ok(None) + } else { + Err(CapsuleError::UnexpectedEnd) + }; + } + } + } + + /// Read the next HTTP/3 DATA frame into the buffer, skipping non-DATA frames. + /// Returns `false` on EOF. + async fn read_data_frame(&mut self) -> Result { + loop { + let frame_type = match VarInt::read_optional(&mut self.stream).await { + Ok(Some(v)) => Frame(v), + Ok(None) => return Ok(false), + Err(_) => return Err(CapsuleError::UnexpectedEnd), + }; + let len = VarInt::read(&mut self.stream) + .await + .map_err(|_| CapsuleError::UnexpectedEnd)? + .into_inner() as usize; + + if len > MAX_FRAME_SIZE as usize { + return Err(CapsuleError::MessageTooLong); + } + + if frame_type != Frame::DATA { + let mut limited = (&mut self.stream).take(len as u64); + let n = tokio::io::copy(&mut limited, &mut tokio::io::sink()).await?; + if (n as usize) < len { + return Err(CapsuleError::UnexpectedEnd); + } + continue; + } + + if len > 0 { + let start = self.buf.len(); + self.buf.resize(start + len, 0); + self.stream.read_exact(&mut self.buf[start..]).await?; + } + return Ok(true); + } + } +} + #[cfg(test)] mod tests { use super::*; @@ -493,4 +568,130 @@ mod tests { let err = Capsule::read(&mut cursor).await.unwrap_err(); assert!(matches!(err, CapsuleError::UnexpectedEnd)); } + + fn encode_capsule(c: &Capsule) -> Vec { + let mut buf = Vec::new(); + c.encode(&mut buf); + buf + } + + fn wrap_in_data_frame(payload: &[u8]) -> Vec { + let mut frame = Vec::new(); + Frame::DATA.encode(&mut frame); + VarInt::from_u32(payload.len() as u32).encode(&mut frame); + frame.extend_from_slice(payload); + frame + } + + /// Concatenate multiple DATA-framed chunks into a single wire buffer. + fn split_into_data_frames(capsule_bytes: &[u8], splits: &[usize]) -> Vec { + let mut wire = Vec::new(); + let mut prev = 0; + for &s in splits { + wire.extend_from_slice(&wrap_in_data_frame(&capsule_bytes[prev..s])); + prev = s; + } + wire.extend_from_slice(&wrap_in_data_frame(&capsule_bytes[prev..])); + wire + } + + fn reader_from(wire: Vec) -> Http3CapsuleReader>> { + Http3CapsuleReader::new(std::io::Cursor::new(wire)) + } + + #[tokio::test] + async fn test_http3_reader_single_capsule_single_frame() { + let capsule = Capsule::CloseWebTransportSession { + code: 42, + reason: "bye".into(), + }; + let mut reader = reader_from(wrap_in_data_frame(&encode_capsule(&capsule))); + assert_eq!(reader.read().await.unwrap().unwrap(), capsule); + assert!(reader.read().await.unwrap().is_none()); + } + + #[tokio::test] + async fn test_http3_reader_capsule_split_across_frames() { + let capsule = Capsule::CloseWebTransportSession { + code: 100, + reason: "split".into(), + }; + let bytes = encode_capsule(&capsule); + let mut reader = reader_from(split_into_data_frames(&bytes, &[bytes.len() / 2])); + assert_eq!(reader.read().await.unwrap().unwrap(), capsule); + } + + #[tokio::test] + async fn test_http3_reader_capsule_split_across_three_frames() { + let capsule = Capsule::CloseWebTransportSession { + code: 200, + reason: "three".into(), + }; + let bytes = encode_capsule(&capsule); + let mut reader = reader_from(split_into_data_frames(&bytes, &[2, 5])); + assert_eq!(reader.read().await.unwrap().unwrap(), capsule); + } + + #[tokio::test] + async fn test_http3_reader_multiple_capsules_one_frame() { + let c1 = Capsule::Grease { num: 1 }; + let c2 = Capsule::CloseWebTransportSession { + code: 7, + reason: "done".into(), + }; + let mut bytes = encode_capsule(&c1); + bytes.extend_from_slice(&encode_capsule(&c2)); + + let mut reader = reader_from(wrap_in_data_frame(&bytes)); + assert_eq!(reader.read().await.unwrap().unwrap(), c1); + assert_eq!(reader.read().await.unwrap().unwrap(), c2); + } + + #[tokio::test] + async fn test_http3_reader_skips_non_data_frames() { + let capsule = Capsule::CloseWebTransportSession { + code: 0, + reason: String::new(), + }; + let mut wire = Vec::new(); + // HEADERS frame before the DATA frame. + Frame::HEADERS.encode(&mut wire); + VarInt::from_u32(5).encode(&mut wire); + wire.extend_from_slice(b"hello"); + wire.extend_from_slice(&wrap_in_data_frame(&encode_capsule(&capsule))); + + let mut reader = reader_from(wire); + assert_eq!(reader.read().await.unwrap().unwrap(), capsule); + } + + #[tokio::test] + async fn test_http3_reader_eof_returns_none() { + assert!(reader_from(vec![]).read().await.unwrap().is_none()); + } + + #[tokio::test] + async fn test_http3_reader_partial_capsule_at_eof() { + let mut partial = Vec::new(); + VarInt::from_u64(0x2843).unwrap().encode(&mut partial); + // Missing: length + payload + + let mut reader = reader_from(wrap_in_data_frame(&partial)); + assert!(matches!( + reader.read().await, + Err(CapsuleError::UnexpectedEnd) + )); + } + + #[tokio::test] + async fn test_http3_reader_empty_data_frame() { + let capsule = Capsule::CloseWebTransportSession { + code: 1, + reason: "ok".into(), + }; + let mut wire = wrap_in_data_frame(&[]); + wire.extend_from_slice(&wrap_in_data_frame(&encode_capsule(&capsule))); + + let mut reader = reader_from(wire); + assert_eq!(reader.read().await.unwrap().unwrap(), capsule); + } } diff --git a/web-transport-proto/src/varint.rs b/web-transport-proto/src/varint.rs index a0c2086a..dc44a3ae 100644 --- a/web-transport-proto/src/varint.rs +++ b/web-transport-proto/src/varint.rs @@ -163,16 +163,31 @@ impl VarInt { Ok(Self(x)) } - // Read a varint from the stream. + /// Read a varint from the stream. pub async fn read(stream: &mut S) -> Result { + Self::read_optional(stream) + .await? + .ok_or(VarIntUnexpectedEnd) + } + + /// Read a varint from the stream, returning `Ok(None)` on clean EOF. + /// + /// Returns `Ok(None)` if EOF is encountered before any byte is read, + /// `Err(VarIntUnexpectedEnd)` if EOF occurs mid-varint, and + /// `Ok(Some(varint))` on success. + pub async fn read_optional( + stream: &mut S, + ) -> Result, VarIntUnexpectedEnd> { // 8 bytes is the max size of a varint let mut buf = [0; 8]; // Read the first byte because it includes the length. - stream - .read_exact(&mut buf[0..1]) - .await - .map_err(|_| VarIntUnexpectedEnd)?; + // EOF here means clean end-of-stream. + match stream.read_exact(&mut buf[0..1]).await { + Ok(_) => {} + Err(e) if e.kind() == std::io::ErrorKind::UnexpectedEof => return Ok(None), + Err(_) => return Err(VarIntUnexpectedEnd), + } // 0b00 = 1, 0b01 = 2, 0b10 = 4, 0b11 = 8 let size = 1 << (buf[0] >> 6); @@ -185,7 +200,7 @@ impl VarInt { let mut cursor = Cursor::new(&buf[..size]); let v = VarInt::decode(&mut cursor).unwrap(); - Ok(v) + Ok(Some(v)) } pub fn encode(&self, w: &mut B) { diff --git a/web-transport-quinn/Cargo.toml b/web-transport-quinn/Cargo.toml index 8da01cc3..c3f0decd 100644 --- a/web-transport-quinn/Cargo.toml +++ b/web-transport-quinn/Cargo.toml @@ -40,6 +40,7 @@ thiserror = "2" tokio = { version = "1", default-features = false, features = [ "io-util", "macros", + "time", ] } tracing = "0.1" url = "2" diff --git a/web-transport-quinn/src/recv.rs b/web-transport-quinn/src/recv.rs index 6734c202..429281c6 100644 --- a/web-transport-quinn/src/recv.rs +++ b/web-transport-quinn/src/recv.rs @@ -70,7 +70,7 @@ impl RecvStream { Ok(Some(code)) => Ok(Some( web_transport_proto::error_from_http3(code.into_inner()).unwrap(), )), - Err(quinn::ResetError::ConnectionLost(e)) => Err(e.into()), + Err(quinn::ResetError::ConnectionLost(conn_err)) => Err(conn_err.into()), Err(quinn::ResetError::ZeroRttRejected) => unreachable!("0-RTT not supported"), } } diff --git a/web-transport-quinn/src/send.rs b/web-transport-quinn/src/send.rs index 836490f0..65bfdaeb 100644 --- a/web-transport-quinn/src/send.rs +++ b/web-transport-quinn/src/send.rs @@ -38,7 +38,7 @@ impl SendStream { match self.stream.stopped().await { Ok(Some(code)) => Ok(web_transport_proto::error_from_http3(code.into_inner())), Ok(None) => Ok(None), - Err(quinn::StoppedError::ConnectionLost(e)) => Err(e.into()), + Err(quinn::StoppedError::ConnectionLost(conn_err)) => Err(conn_err.into()), Err(quinn::StoppedError::ZeroRttRejected) => unreachable!("0-RTT not supported"), } } diff --git a/web-transport-quinn/src/session.rs b/web-transport-quinn/src/session.rs index 5ab7b952..b45f534a 100644 --- a/web-transport-quinn/src/session.rs +++ b/web-transport-quinn/src/session.rs @@ -4,8 +4,9 @@ use std::{ io::Cursor, ops::Deref, pin::Pin, - sync::{Arc, Mutex}, - task::{ready, Context, Poll}, + sync::{Arc, Mutex, OnceLock}, + task::{Context, Poll, Waker}, + time::Duration, }; use bytes::{Bytes, BytesMut}; @@ -44,6 +45,15 @@ pub struct Session { #[allow(dead_code)] settings: Option>, + // The send side of the CONNECT stream, used to write the CloseWebTransportSession capsule. + // Wrapped in Arc>> so close() can take it exactly once. + connect_send: Arc>>, + + // Session error, set once by either local close() or the background task + // when a remote CloseWebTransportSession capsule is received. + // Uses OnceLock for set-once, first-writer-wins semantics with lock-free reads. + error: Arc>, + // The request sent by the client. request: ConnectRequest, @@ -68,6 +78,8 @@ impl Session { let mut header_datagram = Vec::new(); session_id.encode(&mut header_datagram); + let error: Arc> = Arc::new(OnceLock::new()); + // Accept logic is stateful, so use an Arc to share it. let accept = SessionAccept::new(conn.clone(), session_id); @@ -79,40 +91,72 @@ impl Session { header_bi, header_datagram, settings: Some(Arc::new(settings)), + connect_send: Arc::new(Mutex::new(Some(connect.send))), + error: error.clone(), request: connect.request.clone(), response: connect.response.clone(), }; - // Run a background task to check if the connect stream is closed. - let mut this2 = this.clone(); - tokio::spawn(async move { - let (code, reason) = this2.run_closed(connect).await; - // TODO We shouldn't be closing the QUIC connection with the same error. - this2.close(code, reason.as_bytes()); - }); + // Run a background task to read capsules from the CONNECT recv stream. + let conn2 = this.conn.clone(); + tokio::spawn(Self::run_recv(conn2, connect.recv, error)); this } - // Keep reading from the control stream until it's closed. - async fn run_closed(&mut self, mut connect: Connected) -> (u32, String) { + // Read capsules from the CONNECT recv stream until it's closed, + // then record the close error and tear down the connection. + async fn run_recv( + conn: quinn::Connection, + recv: quinn::RecvStream, + error: Arc>, + ) { + let close_info = Self::read_capsules(recv).await; + let code = close_info.as_ref().map_or(0, |(c, _)| *c); + + let http3_code: quinn::VarInt = web_transport_proto::error_to_http3(code) + .try_into() + .unwrap(); + + // Try to record the remote close error. If close() already set + // the error, it owns the connection teardown, so we bail out. + match close_info { + Some((code, reason)) => { + let err = WebTransportError::Closed(code, reason.clone()); + if error.set(err.into()).is_err() { + return; + } + conn.close(http3_code, reason.as_bytes()); + } + None => { + let err = quinn::ConnectionError::LocallyClosed.into(); + if error.set(err).is_err() { + return; + } + conn.close(http3_code, b""); + } + }; + } + + // Keep reading capsules from the CONNECT recv stream until it's closed. + // Returns Some((code, reason)) if a CloseWebTransportSession capsule was received, + // or None if the stream closed without a capsule. + async fn read_capsules(recv: quinn::RecvStream) -> Option<(u32, String)> { + let mut reader = web_transport_proto::Http3CapsuleReader::new(recv); loop { - match web_transport_proto::Capsule::read(&mut connect.recv).await { + match reader.read().await { Ok(Some(web_transport_proto::Capsule::CloseWebTransportSession { code, reason, - })) => { - return (code, reason); - } + })) => return Some((code, reason)), Ok(Some(web_transport_proto::Capsule::Grease { .. })) => {} Ok(Some(web_transport_proto::Capsule::Unknown { typ, payload })) => { tracing::warn!(%typ, size = payload.len(), "unknown capsule"); } - Ok(None) => { - return (0, "stream closed".to_string()); - } - Err(_) => { - return (1, "capsule error".to_string()); + Ok(None) => return None, + Err(e) => { + tracing::warn!(?e, "failed to read capsule"); + return None; } } } @@ -142,26 +186,19 @@ impl Session { /// Accept a new unidirectional stream. See [`quinn::Connection::accept_uni`]. pub async fn accept_uni(&self) -> Result { if let Some(accept) = &self.accept { - poll_fn(|cx| accept.lock().unwrap().poll_accept_uni(cx)).await + Ok(poll_fn(|cx| accept.lock().unwrap().poll_accept_uni(cx)).await?) } else { - self.conn - .accept_uni() - .await - .map(RecvStream::new) - .map_err(Into::into) + Ok(RecvStream::new(self.conn.accept_uni().await?)) } } /// Accept a new bidirectional stream. See [`quinn::Connection::accept_bi`]. pub async fn accept_bi(&self) -> Result<(SendStream, RecvStream), SessionError> { if let Some(accept) = &self.accept { - poll_fn(|cx| accept.lock().unwrap().poll_accept_bi(cx)).await + Ok(poll_fn(|cx| accept.lock().unwrap().poll_accept_bi(cx)).await?) } else { - self.conn - .accept_bi() - .await - .map(|(send, recv)| (SendStream::new(send), RecvStream::new(recv))) - .map_err(Into::into) + let (send, recv) = self.conn.accept_bi().await?; + Ok((SendStream::new(send), RecvStream::new(recv))) } } @@ -201,11 +238,7 @@ impl Session { /// peer over the connection. /// It waits for a datagram to become available and returns the received bytes. pub async fn read_datagram(&self) -> Result { - let mut datagram = self - .conn - .read_datagram() - .await - .map_err(SessionError::from)?; + let mut datagram = self.conn.read_datagram().await?; let mut cursor = Cursor::new(&datagram); @@ -229,7 +262,7 @@ impl Session { /// Datagrams are unreliable and may be dropped or delivered out of order. /// The data must be smaller than [`max_datagram_size`](Self::max_datagram_size). pub fn send_datagram(&self, data: Bytes) -> Result<(), SessionError> { - if !self.header_datagram.is_empty() { + let result = if !self.header_datagram.is_empty() { // Unfortunately, we need to allocate/copy each datagram because of the Quinn API. // Pls go +1 if you care: https://github.com/quinn-rs/quinn/issues/1724 let mut buf = BytesMut::with_capacity(self.header_datagram.len() + data.len()); @@ -238,11 +271,12 @@ impl Session { buf.extend_from_slice(&self.header_datagram); buf.extend_from_slice(&data); - self.conn.send_datagram(buf.into())?; + self.conn.send_datagram(buf.into()) } else { - self.conn.send_datagram(data)?; - } + self.conn.send_datagram(data) + }; + result?; Ok(()) } @@ -256,27 +290,119 @@ impl Session { mtu.saturating_sub(self.header_datagram.len()) } - /// Immediately close the connection with an error code and reason. See [`quinn::Connection::close`]. + /// Close the session with an error code and reason. + /// + /// When there is a session ID (WebTransport over HTTP/3), a `CloseWebTransportSession` + /// capsule is written on the CONNECT stream before the QUIC connection is closed. + /// This allows browser clients to receive the close code and reason via `WebTransport.closed`. + /// + /// The capsule write and connection close happen asynchronously in a spawned task. + /// Callers should `await` [`Session::closed()`] to ensure the capsule has been + /// delivered. Session operations will fail once the QUIC connection is closed. pub fn close(&self, code: u32, reason: &[u8]) { - let code = if self.session_id.is_some() { - web_transport_proto::error_to_http3(code) - .try_into() - .unwrap() + // Record the local close error. First writer wins — if the background + // task already set a remote close error, or close() was already called, + // this is a no-op. + let err = SessionError::ConnectionError(quinn::ConnectionError::LocallyClosed); + if self.error.set(err).is_err() { + return; + } + + if self.session_id.is_some() { + // Take the send stream for the capsule write. + let send = self.connect_send.lock().unwrap().take(); + + if let Some(send) = send { + let reason = String::from_utf8_lossy(reason).into_owned(); + let conn = self.conn.clone(); + let capsule = + web_transport_proto::Capsule::CloseWebTransportSession { code, reason }; + let timeout = (self.rtt() * 3).max(Duration::from_millis(100)); + + tokio::spawn(async move { + Self::close_with_capsule(conn, send, capsule, code, timeout).await; + }); + } } else { - code.into() + // Raw QUIC mode: no capsule needed. + self.conn.close(code.into(), reason); + } + } + + /// Write the CloseWebTransportSession capsule, finish the stream, wait for + /// the peer to close the connection (or timeout), then force-close. + async fn close_with_capsule( + conn: quinn::Connection, + mut send: quinn::SendStream, + capsule: web_transport_proto::Capsule, + code: u32, + timeout: std::time::Duration, + ) { + let http3_code: quinn::VarInt = web_transport_proto::error_to_http3(code) + .try_into() + .unwrap(); + + // Encode the capsule, then wrap it in an HTTP/3 DATA frame. + // In HTTP/3, capsule data is carried inside DATA frames on the CONNECT + // stream (RFC 9297 Section 3.2). + let mut capsule_bytes = Vec::new(); + capsule.encode(&mut capsule_bytes); + + let mut frame = Vec::new(); + Frame::DATA.encode(&mut frame); + let Ok(len) = VarInt::try_from(capsule_bytes.len()) else { + tracing::warn!("capsule too large to encode as DATA frame"); + conn.close(http3_code, b""); + return; }; + len.encode(&mut frame); + frame.extend_from_slice(&capsule_bytes); + + // Write the DATA frame to the CONNECT send stream. + if let Err(e) = send.write_all(&frame).await { + tracing::warn!(?e, "failed to write CloseWebTransportSession capsule"); + conn.close(http3_code, b""); + return; + } + + // FIN the send stream so the peer knows no more capsules are coming. + if let Err(e) = send.finish() { + tracing::warn!(?e, "failed to finish CONNECT send stream"); + conn.close(http3_code, b""); + return; + } - self.conn.close(code, reason) + // Wait for the peer to close the CONNECT stream after receiving the capsule. + if tokio::time::timeout(timeout, conn.closed()).await.is_err() { + tracing::debug!("timeout waiting for peer to close; force-closing connection"); + conn.close(http3_code, b""); + } } /// Wait until the session is closed, returning the error. See [`quinn::Connection::closed`]. + /// + /// If the peer sent a `CloseWebTransportSession` capsule, the returned error will be + /// [`WebTransportError::Closed`] with the code and reason from the capsule. + /// + /// Unlike [`quinn::Connection::closed`], this does **not** return early when + /// [`close()`](Self::close) has been called. It waits for the underlying QUIC + /// connection to shut down, ensuring the `CloseWebTransportSession` capsule has + /// been delivered. Use [`close_reason()`](Self::close_reason) for a non-blocking check. pub async fn closed(&self) -> SessionError { - self.conn.closed().await.into() + self.map_error(self.conn.closed().await) } /// Return why the session was closed, or None if it's not closed. See [`quinn::Connection::close_reason`]. pub fn close_reason(&self) -> Option { - self.conn.close_reason().map(Into::into) + self.conn.close_reason().map(|e| self.map_error(e)) + } + + /// Convert an error to `SessionError`, using the stored session error if set. + fn map_error(&self, e: impl Into) -> SessionError { + if let Some(err) = self.error.get() { + return err.clone(); + } + e.into() } async fn write_full(send: &mut quinn::SendStream, buf: &[u8]) -> Result<(), SessionError> { @@ -304,6 +430,8 @@ impl Session { header_datagram: Default::default(), accept: None, settings: None, + connect_send: Arc::new(Mutex::new(None)), + error: Arc::new(OnceLock::new()), request: request.into(), response: response.into(), } @@ -371,6 +499,12 @@ pub struct SessionAccept { // Keep track of work being done to read/write the WebTransport stream header. pending_uni: FuturesUnordered>>, pending_bi: FuturesUnordered>>, + + // Wakers from concurrent callers of accept_bi / accept_uni. + // When one caller gets a stream, all others are woken so they can retry. + // This fixes the lost-waker bug where the unfold stream only stores one waker. + bi_wakers: Vec, + uni_wakers: Vec, } impl SessionAccept { @@ -395,6 +529,9 @@ impl SessionAccept { pending_uni: FuturesUnordered::new(), pending_bi: FuturesUnordered::new(), + + bi_wakers: Vec::new(), + uni_wakers: Vec::new(), } } @@ -409,7 +546,15 @@ impl SessionAccept { // Accept any new streams. if let Poll::Ready(Some(res)) = self.accept_uni.poll_next_unpin(cx) { // Start decoding the header and add the future to the list of pending streams. - let recv = res?; + let recv = match res { + Ok(recv) => recv, + Err(e) => { + for waker in self.uni_wakers.drain(..) { + waker.wake(); + } + return Poll::Ready(Err(e.into())); + } + }; let pending = Self::decode_uni(recv, self.session_id); self.pending_uni.push(Box::pin(pending)); @@ -417,20 +562,28 @@ impl SessionAccept { } // Poll the list of pending streams. - let (typ, recv) = match ready!(self.pending_uni.poll_next_unpin(cx)) { - Some(Ok(res)) => res, - Some(Err(err)) => { + let (typ, recv) = match self.pending_uni.poll_next_unpin(cx) { + Poll::Ready(Some(Ok(res))) => res, + Poll::Ready(Some(Err(err))) => { // Ignore the error, the stream was probably reset early. tracing::warn!(?err, "failed to decode unidirectional stream"); continue; } - None => return Poll::Pending, + Poll::Ready(None) | Poll::Pending => { + if !self.uni_wakers.iter().any(|w| w.will_wake(cx.waker())) { + self.uni_wakers.push(cx.waker().clone()); + } + return Poll::Pending; + } }; // Decide if we keep looping based on the type. match typ { StreamUni::WEBTRANSPORT => { let recv = RecvStream::new(recv); + for waker in self.uni_wakers.drain(..) { + waker.wake(); + } return Poll::Ready(Ok(recv)); } StreamUni::QPACK_DECODER => { @@ -480,7 +633,15 @@ impl SessionAccept { // Accept any new streams. if let Poll::Ready(Some(res)) = self.accept_bi.poll_next_unpin(cx) { // Start decoding the header and add the future to the list of pending streams. - let (send, recv) = res?; + let (send, recv) = match res { + Ok(pair) => pair, + Err(e) => { + for waker in self.bi_wakers.drain(..) { + waker.wake(); + } + return Poll::Ready(Err(e.into())); + } + }; let pending = Self::decode_bi(send, recv, self.session_id); self.pending_bi.push(Box::pin(pending)); @@ -488,20 +649,28 @@ impl SessionAccept { } // Poll the list of pending streams. - let res = match ready!(self.pending_bi.poll_next_unpin(cx)) { - Some(Ok(res)) => res, - Some(Err(err)) => { + let res = match self.pending_bi.poll_next_unpin(cx) { + Poll::Ready(Some(Ok(res))) => res, + Poll::Ready(Some(Err(err))) => { // Ignore the error, the stream was probably reset early. tracing::warn!(?err, "failed to decode bidirectional stream"); continue; } - None => return Poll::Pending, + Poll::Ready(None) | Poll::Pending => { + if !self.bi_wakers.iter().any(|w| w.will_wake(cx.waker())) { + self.bi_wakers.push(cx.waker().clone()); + } + return Poll::Pending; + } }; if let Some((send, recv)) = res { // Wrap the streams in our own types for correct error codes. let send = SendStream::new(send); let recv = RecvStream::new(recv); + for waker in self.bi_wakers.drain(..) { + waker.wake(); + } return Poll::Ready(Ok((send, recv))); }