From e0cae64e42b346450ed48e32e86b5ce8a2353bbb Mon Sep 17 00:00:00 2001 From: ai-and-i <138349030+ai-and-i@users.noreply.github.com> Date: Wed, 25 Feb 2026 23:37:54 -0800 Subject: [PATCH 1/8] Implement W3C-compliant CloseWebTransportSession capsule for browser compatibility Write a CloseWebTransportSession capsule inside an HTTP/3 DATA frame on the CONNECT stream (RFC 9297), enabling browsers to receive close codes and reasons via the W3C WebTransport API's `WebTransport.closed` promise. Adds Http3CapsuleReader to handle capsules split across multiple DATA frames, and reworks Session::close() to deliver the capsule before closing the QUIC connection. Co-Authored-By: Claude Opus 4.6 --- web-transport-proto/src/capsule.rs | 185 ++++++++++++++++++++++++++++- web-transport-quinn/Cargo.toml | 1 + web-transport-quinn/src/session.rs | 165 +++++++++++++++++++++---- 3 files changed, 325 insertions(+), 26 deletions(-) diff --git a/web-transport-proto/src/capsule.rs b/web-transport-proto/src/capsule.rs index a018b3cd..ee7d76b0 100644 --- a/web-transport-proto/src/capsule.rs +++ b/web-transport-proto/src/capsule.rs @@ -3,7 +3,7 @@ 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 @@ -232,6 +232,84 @@ 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(&mut self.stream).await { + Ok(v) => Frame(v), + Err(_) => return Ok(false), + }; + 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 { + if len > 0 { + let mut skip = vec![0u8; len]; + self.stream.read_exact(&mut skip).await?; + } + 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 +571,109 @@ 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-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/session.rs b/web-transport-quinn/src/session.rs index 5ab7b952..1fc93951 100644 --- a/web-transport-quinn/src/session.rs +++ b/web-transport-quinn/src/session.rs @@ -44,6 +44,14 @@ 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>>, + + // Close info received from the peer via a CloseWebTransportSession capsule. + // Set by the background task when a capsule is received. + close_capsule: Arc>>, + // The request sent by the client. request: ConnectRequest, @@ -71,6 +79,8 @@ impl Session { // Accept logic is stateful, so use an Arc to share it. let accept = SessionAccept::new(conn.clone(), session_id); + let close_capsule = Arc::new(Mutex::new(None)); + let this = Self { conn, accept: Some(Arc::new(Mutex::new(accept))), @@ -79,40 +89,60 @@ impl Session { header_bi, header_datagram, settings: Some(Arc::new(settings)), + connect_send: Arc::new(Mutex::new(Some(connect.send))), + close_capsule: close_capsule.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(); + // Run a background task to read capsules from the CONNECT recv stream. + let conn2 = this.conn.clone(); + let connect_send2 = this.connect_send.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()); + let close_info = Self::run_closed(connect.recv).await; + + // Send capsule close info (if any) for Session::closed() to consume. + // This MUST happen before conn.close() so the info is available + // when conn.closed() resolves on other tasks. + if let Some(ref info) = close_info { + *close_capsule.lock().unwrap() = Some(info.clone()); + } + + // Take the send stream so that any concurrent/future close() call + // sees None and doesn't try to write a capsule that can't be delivered + // (the peer already closed). + let _send = connect_send2.lock().unwrap().take(); + + // Close the QUIC connection. + 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(); + conn2.close(http3_code, &[]); }); this } - // Keep reading from the control stream until it's closed. - async fn run_closed(&mut self, mut connect: Connected) -> (u32, String) { + // 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 run_closed(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; } } } @@ -256,26 +286,109 @@ 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()`] if they need to ensure the capsule + /// has been delivered before proceeding. 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() + if self.session_id.is_some() { + // Take the send stream. Only the first caller gets it. + let send = self.connect_send.lock().unwrap().take(); + + if let Some(send) = send { + let conn = self.conn.clone(); + let capsule = web_transport_proto::Capsule::CloseWebTransportSession { + code, + reason: String::from_utf8_lossy(reason).into_owned(), + }; + let timeout = self.rtt() * 3; + + tokio::spawn(async move { + Self::close_with_capsule(conn, send, capsule, code, timeout).await; + }); + } else { + // Send stream already taken — close() was already called. + // The first call is handling the capsule + connection close. + tracing::debug!("close() called more than once; ignoring duplicate"); + } } else { - code.into() - }; + // Raw QUIC mode: no capsule needed. + self.conn.close(code.into(), reason); + } + } - self.conn.close(code, 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); + VarInt::from_u32(capsule_bytes.len() as u32).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; + } + + // 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. pub async fn closed(&self) -> SessionError { - self.conn.closed().await.into() + let err = self.conn.closed().await; + + // Check if the background task received a CloseWebTransportSession capsule. + // The background task always sends the capsule info before calling conn.close(), + // so the value is guaranteed to be visible here when conn.closed() resolves. + if let Some((code, reason)) = self.close_capsule.lock().unwrap().clone() { + return WebTransportError::Closed(code, reason).into(); + } + + err.into() } /// Return why the session was closed, or None if it's not closed. See [`quinn::Connection::close_reason`]. pub fn close_reason(&self) -> Option { + if let Some((code, reason)) = self.close_capsule.lock().unwrap().clone() { + return Some(WebTransportError::Closed(code, reason).into()); + } self.conn.close_reason().map(Into::into) } @@ -304,6 +417,8 @@ impl Session { header_datagram: Default::default(), accept: None, settings: None, + connect_send: Arc::new(Mutex::new(None)), + close_capsule: Arc::new(Mutex::new(None)), request: request.into(), response: response.into(), } From 16069d1850866713848fdc65c8295f02a4271179 Mon Sep 17 00:00:00 2001 From: ai-and-i <138349030+ai-and-i@users.noreply.github.com> Date: Thu, 26 Feb 2026 10:14:06 -0800 Subject: [PATCH 2/8] Propagate close capsule info to stream, accept, and datagram errors MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit When a peer sends a CloseWebTransportSession capsule, the background task stores the (code, reason) and calls conn.close(), causing all active Quinn operations to fail with LocallyClosed. Previously only Session::closed() checked the capsule—now every error path maps LocallyClosed to WebTransportError::Closed(code, reason). - Add CloseCapsule type alias and map_conn_error/map_session_error/ map_send_datagram_error helpers on Session and SessionAccept - Thread close_capsule into RecvStream and SendStream; add map_read_error/map_write_error helpers - Store local close info in close() so locally-initiated closes also surface WebTransportError::Closed consistently Co-Authored-By: Claude Opus 4.6 --- web-transport-quinn/src/recv.rs | 69 +++++++++++-- web-transport-quinn/src/send.rs | 57 ++++++++-- web-transport-quinn/src/session.rs | 160 +++++++++++++++++++++-------- 3 files changed, 228 insertions(+), 58 deletions(-) diff --git a/web-transport-quinn/src/recv.rs b/web-transport-quinn/src/recv.rs index 6734c202..796983b5 100644 --- a/web-transport-quinn/src/recv.rs +++ b/web-transport-quinn/src/recv.rs @@ -6,17 +6,49 @@ use std::{ use bytes::Bytes; -use crate::{ReadError, ReadExactError, ReadToEndError, SessionError}; +use crate::{ + session::CloseCapsule, ReadError, ReadExactError, ReadToEndError, SessionError, + WebTransportError, +}; /// A stream that can be used to recieve bytes. See [`quinn::RecvStream`]. #[derive(Debug)] pub struct RecvStream { inner: quinn::RecvStream, + close_capsule: CloseCapsule, } impl RecvStream { - pub(crate) fn new(stream: quinn::RecvStream) -> Self { - Self { inner: stream } + pub(crate) fn new(stream: quinn::RecvStream, close_capsule: CloseCapsule) -> Self { + Self { + inner: stream, + close_capsule, + } + } + + fn map_read_error(&self, e: quinn::ReadError) -> ReadError { + if let quinn::ReadError::ConnectionLost(ref conn_err) = e { + if matches!(conn_err, quinn::ConnectionError::LocallyClosed) { + if let Some((code, reason)) = self.close_capsule.lock().unwrap().clone() { + return ReadError::SessionError(WebTransportError::Closed(code, reason).into()); + } + } + } + e.into() + } + + fn map_read_exact_error(&self, e: quinn::ReadExactError) -> ReadExactError { + if let quinn::ReadExactError::ReadError(read_err) = e { + return self.map_read_error(read_err).into(); + } + e.into() + } + + fn map_read_to_end_error(&self, e: quinn::ReadToEndError) -> ReadToEndError { + if let quinn::ReadToEndError::Read(read_err) = e { + return self.map_read_error(read_err).into(); + } + e.into() } /// Tell the other end to stop sending data with the given error code. See [`quinn::RecvStream::stop`]. @@ -31,12 +63,18 @@ impl RecvStream { /// Read some data into the buffer and return the amount read. See [`quinn::RecvStream::read`]. pub async fn read(&mut self, buf: &mut [u8]) -> Result, ReadError> { - self.inner.read(buf).await.map_err(Into::into) + self.inner + .read(buf) + .await + .map_err(|e| self.map_read_error(e)) } /// Fill the entire buffer with data. See [`quinn::RecvStream::read_exact`]. pub async fn read_exact(&mut self, buf: &mut [u8]) -> Result<(), ReadExactError> { - self.inner.read_exact(buf).await.map_err(Into::into) + self.inner + .read_exact(buf) + .await + .map_err(|e| self.map_read_exact_error(e)) } /// Read a chunk of data from the stream. See [`quinn::RecvStream::read_chunk`]. @@ -48,17 +86,23 @@ impl RecvStream { self.inner .read_chunk(max_length, ordered) .await - .map_err(Into::into) + .map_err(|e| self.map_read_error(e)) } /// Read chunks of data from the stream. See [`quinn::RecvStream::read_chunks`]. pub async fn read_chunks(&mut self, bufs: &mut [Bytes]) -> Result, ReadError> { - self.inner.read_chunks(bufs).await.map_err(Into::into) + self.inner + .read_chunks(bufs) + .await + .map_err(|e| self.map_read_error(e)) } /// Read until the end of the stream or the limit is hit. See [`quinn::RecvStream::read_to_end`]. pub async fn read_to_end(&mut self, size_limit: usize) -> Result, ReadToEndError> { - self.inner.read_to_end(size_limit).await.map_err(Into::into) + self.inner + .read_to_end(size_limit) + .await + .map_err(|e| self.map_read_to_end_error(e)) } /// Block until the stream has been reset and return the error code. See [`quinn::RecvStream::received_reset`]. @@ -70,7 +114,14 @@ 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)) => { + if matches!(conn_err, quinn::ConnectionError::LocallyClosed) { + if let Some((code, reason)) = self.close_capsule.lock().unwrap().clone() { + return Err(WebTransportError::Closed(code, reason).into()); + } + } + 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..92754cd6 100644 --- a/web-transport-quinn/src/send.rs +++ b/web-transport-quinn/src/send.rs @@ -6,7 +6,7 @@ use std::{ use bytes::{Buf, Bytes}; -use crate::{ClosedStream, SessionError, WriteError}; +use crate::{session::CloseCapsule, ClosedStream, SessionError, WebTransportError, WriteError}; /// A stream that can be used to send bytes. See [`quinn::SendStream`]. /// @@ -15,11 +15,28 @@ use crate::{ClosedStream, SessionError, WriteError}; #[derive(Debug)] pub struct SendStream { stream: quinn::SendStream, + close_capsule: CloseCapsule, } impl SendStream { - pub(crate) fn new(stream: quinn::SendStream) -> Self { - Self { stream } + pub(crate) fn new(stream: quinn::SendStream, close_capsule: CloseCapsule) -> Self { + Self { + stream, + close_capsule, + } + } + + fn map_write_error(&self, e: quinn::WriteError) -> WriteError { + if let quinn::WriteError::ConnectionLost(ref conn_err) = e { + if matches!(conn_err, quinn::ConnectionError::LocallyClosed) { + if let Some((code, reason)) = self.close_capsule.lock().unwrap().clone() { + return WriteError::SessionError( + WebTransportError::Closed(code, reason).into(), + ); + } + } + } + e.into() } /// Abruptly reset the stream with the provided error code. See [`quinn::SendStream::reset`]. @@ -38,7 +55,14 @@ 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)) => { + if matches!(conn_err, quinn::ConnectionError::LocallyClosed) { + if let Some((code, reason)) = self.close_capsule.lock().unwrap().clone() { + return Err(WebTransportError::Closed(code, reason).into()); + } + } + Err(conn_err.into()) + } Err(quinn::StoppedError::ZeroRttRejected) => unreachable!("0-RTT not supported"), } } @@ -47,27 +71,42 @@ impl SendStream { /// Write some data to the stream, returning the size written. See [`quinn::SendStream::write`]. pub async fn write(&mut self, buf: &[u8]) -> Result { - self.stream.write(buf).await.map_err(Into::into) + self.stream + .write(buf) + .await + .map_err(|e| self.map_write_error(e)) } /// Write all of the data to the stream. See [`quinn::SendStream::write_all`]. pub async fn write_all(&mut self, buf: &[u8]) -> Result<(), WriteError> { - self.stream.write_all(buf).await.map_err(Into::into) + self.stream + .write_all(buf) + .await + .map_err(|e| self.map_write_error(e)) } /// Write chunks of data to the stream. See [`quinn::SendStream::write_chunks`]. pub async fn write_chunks(&mut self, bufs: &mut [Bytes]) -> Result { - self.stream.write_chunks(bufs).await.map_err(Into::into) + self.stream + .write_chunks(bufs) + .await + .map_err(|e| self.map_write_error(e)) } /// Write a chunk of data to the stream. See [`quinn::SendStream::write_chunk`]. pub async fn write_chunk(&mut self, buf: Bytes) -> Result<(), WriteError> { - self.stream.write_chunk(buf).await.map_err(Into::into) + self.stream + .write_chunk(buf) + .await + .map_err(|e| self.map_write_error(e)) } /// Write all of the chunks of data to the stream. See [`quinn::SendStream::write_all_chunks`]. pub async fn write_all_chunks(&mut self, bufs: &mut [Bytes]) -> Result<(), WriteError> { - self.stream.write_all_chunks(bufs).await.map_err(Into::into) + self.stream + .write_all_chunks(bufs) + .await + .map_err(|e| self.map_write_error(e)) } /// Mark the stream as finished, such that no more data can be written. See [`quinn::SendStream::finish`]. diff --git a/web-transport-quinn/src/session.rs b/web-transport-quinn/src/session.rs index 1fc93951..a5d4d6ef 100644 --- a/web-transport-quinn/src/session.rs +++ b/web-transport-quinn/src/session.rs @@ -16,6 +16,10 @@ use crate::{ ClientError, Connected, RecvStream, SendStream, SessionError, Settings, WebTransportError, }; +/// Shared close capsule info set by the background task when a +/// `CloseWebTransportSession` capsule is received from the peer. +pub(crate) type CloseCapsule = Arc>>; + /// An established WebTransport session, acting like a full QUIC connection. See [`quinn::Connection`]. /// /// It is important to remember that WebTransport is layered on top of QUIC: @@ -76,11 +80,11 @@ impl Session { let mut header_datagram = Vec::new(); session_id.encode(&mut header_datagram); - // Accept logic is stateful, so use an Arc to share it. - let accept = SessionAccept::new(conn.clone(), session_id); - let close_capsule = Arc::new(Mutex::new(None)); + // Accept logic is stateful, so use an Arc to share it. + let accept = SessionAccept::new(conn.clone(), session_id, close_capsule.clone()); + let this = Self { conn, accept: Some(Arc::new(Mutex::new(accept))), @@ -172,57 +176,81 @@ 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 + poll_fn(|cx| accept.lock().unwrap().poll_accept_uni(cx)) + .await + .map_err(|e| self.map_session_error(e)) } else { self.conn .accept_uni() .await - .map(RecvStream::new) - .map_err(Into::into) + .map(|recv| RecvStream::new(recv, self.close_capsule.clone())) + .map_err(|e| self.map_conn_error(e)) } } /// 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 + poll_fn(|cx| accept.lock().unwrap().poll_accept_bi(cx)) + .await + .map_err(|e| self.map_session_error(e)) } else { self.conn .accept_bi() .await - .map(|(send, recv)| (SendStream::new(send), RecvStream::new(recv))) - .map_err(Into::into) + .map(|(send, recv)| { + ( + SendStream::new(send, self.close_capsule.clone()), + RecvStream::new(recv, self.close_capsule.clone()), + ) + }) + .map_err(|e| self.map_conn_error(e)) } } /// Open a new unidirectional stream. See [`quinn::Connection::open_uni`]. pub async fn open_uni(&self) -> Result { - let mut send = self.conn.open_uni().await?; + let mut send = self + .conn + .open_uni() + .await + .map_err(|e| self.map_conn_error(e))?; // Set the stream priority to max and then write the stream header. // Otherwise the application could write data with lower priority than the header, resulting in queuing. // Also the header is very important for determining the session ID without reliable reset. send.set_priority(i32::MAX).ok(); - Self::write_full(&mut send, &self.header_uni).await?; + Self::write_full(&mut send, &self.header_uni) + .await + .map_err(|e| self.map_session_error(e))?; // Reset the stream priority back to the default of 0. send.set_priority(0).ok(); - Ok(SendStream::new(send)) + Ok(SendStream::new(send, self.close_capsule.clone())) } /// Open a new bidirectional stream. See [`quinn::Connection::open_bi`]. pub async fn open_bi(&self) -> Result<(SendStream, RecvStream), SessionError> { - let (mut send, recv) = self.conn.open_bi().await?; + let (mut send, recv) = self + .conn + .open_bi() + .await + .map_err(|e| self.map_conn_error(e))?; // Set the stream priority to max and then write the stream header. // Otherwise the application could write data with lower priority than the header, resulting in queuing. // Also the header is very important for determining the session ID without reliable reset. send.set_priority(i32::MAX).ok(); - Self::write_full(&mut send, &self.header_bi).await?; + Self::write_full(&mut send, &self.header_bi) + .await + .map_err(|e| self.map_session_error(e))?; // Reset the stream priority back to the default of 0. send.set_priority(0).ok(); - Ok((SendStream::new(send), RecvStream::new(recv))) + Ok(( + SendStream::new(send, self.close_capsule.clone()), + RecvStream::new(recv, self.close_capsule.clone()), + )) } /// Asynchronously receives an application datagram from the remote peer. @@ -235,7 +263,7 @@ impl Session { .conn .read_datagram() .await - .map_err(SessionError::from)?; + .map_err(|e| self.map_session_error(SessionError::from(e)))?; let mut cursor = Cursor::new(&datagram); @@ -259,7 +287,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()); @@ -268,12 +296,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) + }; - Ok(()) + result.map_err(|e| self.map_send_datagram_error(e)) } /// Computes the maximum size of datagrams that may be passed to @@ -301,10 +329,20 @@ impl Session { let send = self.connect_send.lock().unwrap().take(); if let Some(send) = send { + let reason = String::from_utf8_lossy(reason).into_owned(); + + // Store the local close info so closed()/map_conn_error() can + // return WebTransportError::Closed even for locally-initiated closes. + // Use get_or_insert to avoid overwriting a peer capsule that arrived first. + self.close_capsule + .lock() + .unwrap() + .get_or_insert_with(|| (code, reason.clone())); + let conn = self.conn.clone(); let capsule = web_transport_proto::Capsule::CloseWebTransportSession { code, - reason: String::from_utf8_lossy(reason).into_owned(), + reason, }; let timeout = self.rtt() * 3; @@ -372,24 +410,48 @@ impl Session { /// If the peer sent a `CloseWebTransportSession` capsule, the returned error will be /// [`WebTransportError::Closed`] with the code and reason from the capsule. pub async fn closed(&self) -> SessionError { - let err = self.conn.closed().await; + self.map_conn_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(|e| self.map_conn_error(e)) + } - // Check if the background task received a CloseWebTransportSession capsule. - // The background task always sends the capsule info before calling conn.close(), - // so the value is guaranteed to be visible here when conn.closed() resolves. - if let Some((code, reason)) = self.close_capsule.lock().unwrap().clone() { - return WebTransportError::Closed(code, reason).into(); + /// Convert a `quinn::ConnectionError` to `SessionError`, checking the + /// close_capsule so that `LocallyClosed` becomes `WebTransportError::Closed` + /// when a close capsule was received from the peer. + fn map_conn_error(&self, e: quinn::ConnectionError) -> SessionError { + if matches!(e, quinn::ConnectionError::LocallyClosed) { + if let Some((code, reason)) = self.close_capsule.lock().unwrap().clone() { + return WebTransportError::Closed(code, reason).into(); + } } + e.into() + } - err.into() + /// Convert a `SessionError` that wraps `LocallyClosed` into + /// `WebTransportError::Closed` when a close capsule was received. + fn map_session_error(&self, e: SessionError) -> SessionError { + if let SessionError::ConnectionError(ref conn_err) = e { + if matches!(conn_err, quinn::ConnectionError::LocallyClosed) { + if let Some((code, reason)) = self.close_capsule.lock().unwrap().clone() { + return WebTransportError::Closed(code, reason).into(); + } + } + } + e } - /// Return why the session was closed, or None if it's not closed. See [`quinn::Connection::close_reason`]. - pub fn close_reason(&self) -> Option { - if let Some((code, reason)) = self.close_capsule.lock().unwrap().clone() { - return Some(WebTransportError::Closed(code, reason).into()); + fn map_send_datagram_error(&self, e: quinn::SendDatagramError) -> SessionError { + if let quinn::SendDatagramError::ConnectionLost(ref conn_err) = e { + if matches!(conn_err, quinn::ConnectionError::LocallyClosed) { + if let Some((code, reason)) = self.close_capsule.lock().unwrap().clone() { + return WebTransportError::Closed(code, reason).into(); + } + } } - self.conn.close_reason().map(Into::into) + e.into() } async fn write_full(send: &mut quinn::SendStream, buf: &[u8]) -> Result<(), SessionError> { @@ -486,10 +548,17 @@ pub struct SessionAccept { // Keep track of work being done to read/write the WebTransport stream header. pending_uni: FuturesUnordered>>, pending_bi: FuturesUnordered>>, + + // Close capsule info from the peer, shared with Session. + close_capsule: CloseCapsule, } impl SessionAccept { - pub(crate) fn new(conn: quinn::Connection, session_id: VarInt) -> Self { + pub(crate) fn new( + conn: quinn::Connection, + session_id: VarInt, + close_capsule: CloseCapsule, + ) -> Self { // Create a stream that just outputs new streams, so it's easy to call from poll. let accept_uni = Box::pin(futures::stream::unfold(conn.clone(), |conn| async { Some((conn.accept_uni().await, conn)) @@ -510,7 +579,18 @@ impl SessionAccept { pending_uni: FuturesUnordered::new(), pending_bi: FuturesUnordered::new(), + + close_capsule, + } + } + + fn map_conn_error(&self, e: quinn::ConnectionError) -> SessionError { + if matches!(e, quinn::ConnectionError::LocallyClosed) { + if let Some((code, reason)) = self.close_capsule.lock().unwrap().clone() { + return WebTransportError::Closed(code, reason).into(); + } } + e.into() } // This is poll-based because we accept and decode streams in parallel. @@ -524,7 +604,7 @@ 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 = res.map_err(|e| self.map_conn_error(e))?; let pending = Self::decode_uni(recv, self.session_id); self.pending_uni.push(Box::pin(pending)); @@ -545,7 +625,7 @@ impl SessionAccept { // Decide if we keep looping based on the type. match typ { StreamUni::WEBTRANSPORT => { - let recv = RecvStream::new(recv); + let recv = RecvStream::new(recv, self.close_capsule.clone()); return Poll::Ready(Ok(recv)); } StreamUni::QPACK_DECODER => { @@ -595,7 +675,7 @@ 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) = res.map_err(|e| self.map_conn_error(e))?; let pending = Self::decode_bi(send, recv, self.session_id); self.pending_bi.push(Box::pin(pending)); @@ -615,8 +695,8 @@ impl SessionAccept { 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); + let send = SendStream::new(send, self.close_capsule.clone()); + let recv = RecvStream::new(recv, self.close_capsule.clone()); return Poll::Ready(Ok((send, recv))); } From 7ed42c2f4a7532bfc95c36c403d9ae11987927ca Mon Sep 17 00:00:00 2001 From: ai-and-i <138349030+ai-and-i@users.noreply.github.com> Date: Thu, 26 Feb 2026 21:41:20 -0800 Subject: [PATCH 3/8] Replace CloseCapsule mutex with Arc> MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Use OnceLock for set-once, first-writer-wins semantics between close() and the background capsule reader task. This eliminates the Mutex> in favor of lock-free reads on every operation's error path. - close() uses error.set() as the sole guard — if it wins, it owns the capsule write and connection teardown; if the background task already set a remote close error, close() returns early - Background task symmetrically bails out if close() won the race - Consolidated three near-identical error mappers (map_conn_error, map_session_error, map_send_datagram_error) into a single map_error(impl Into) - closed() always awaits conn.closed() to ensure the capsule is delivered before returning, unlike quinn's early-return pattern - Removed connect_send reference from background task since the OnceLock guard makes it unnecessary Co-Authored-By: Claude Opus 4.6 --- web-transport-quinn/src/recv.rs | 50 +++----- web-transport-quinn/src/send.rs | 28 ++--- web-transport-quinn/src/session.rs | 193 ++++++++++++----------------- 3 files changed, 102 insertions(+), 169 deletions(-) diff --git a/web-transport-quinn/src/recv.rs b/web-transport-quinn/src/recv.rs index 796983b5..96208b07 100644 --- a/web-transport-quinn/src/recv.rs +++ b/web-transport-quinn/src/recv.rs @@ -1,52 +1,32 @@ use std::{ io, pin::Pin, + sync::{Arc, OnceLock}, task::{Context, Poll}, }; use bytes::Bytes; -use crate::{ - session::CloseCapsule, ReadError, ReadExactError, ReadToEndError, SessionError, - WebTransportError, -}; +use crate::{ReadError, ReadExactError, ReadToEndError, SessionError}; /// A stream that can be used to recieve bytes. See [`quinn::RecvStream`]. #[derive(Debug)] pub struct RecvStream { inner: quinn::RecvStream, - close_capsule: CloseCapsule, + error: Arc>, } impl RecvStream { - pub(crate) fn new(stream: quinn::RecvStream, close_capsule: CloseCapsule) -> Self { + pub(crate) fn new(stream: quinn::RecvStream, error: Arc>) -> Self { Self { inner: stream, - close_capsule, + error, } } fn map_read_error(&self, e: quinn::ReadError) -> ReadError { - if let quinn::ReadError::ConnectionLost(ref conn_err) = e { - if matches!(conn_err, quinn::ConnectionError::LocallyClosed) { - if let Some((code, reason)) = self.close_capsule.lock().unwrap().clone() { - return ReadError::SessionError(WebTransportError::Closed(code, reason).into()); - } - } - } - e.into() - } - - fn map_read_exact_error(&self, e: quinn::ReadExactError) -> ReadExactError { - if let quinn::ReadExactError::ReadError(read_err) = e { - return self.map_read_error(read_err).into(); - } - e.into() - } - - fn map_read_to_end_error(&self, e: quinn::ReadToEndError) -> ReadToEndError { - if let quinn::ReadToEndError::Read(read_err) = e { - return self.map_read_error(read_err).into(); + if let Some(err) = self.error.get() { + return ReadError::SessionError(err.clone()); } e.into() } @@ -74,7 +54,10 @@ impl RecvStream { self.inner .read_exact(buf) .await - .map_err(|e| self.map_read_exact_error(e)) + .map_err(|e| match e { + quinn::ReadExactError::ReadError(e) => self.map_read_error(e).into(), + e => e.into(), + }) } /// Read a chunk of data from the stream. See [`quinn::RecvStream::read_chunk`]. @@ -102,7 +85,10 @@ impl RecvStream { self.inner .read_to_end(size_limit) .await - .map_err(|e| self.map_read_to_end_error(e)) + .map_err(|e| match e { + quinn::ReadToEndError::Read(e) => self.map_read_error(e).into(), + e => e.into(), + }) } /// Block until the stream has been reset and return the error code. See [`quinn::RecvStream::received_reset`]. @@ -115,10 +101,8 @@ impl RecvStream { web_transport_proto::error_from_http3(code.into_inner()).unwrap(), )), Err(quinn::ResetError::ConnectionLost(conn_err)) => { - if matches!(conn_err, quinn::ConnectionError::LocallyClosed) { - if let Some((code, reason)) = self.close_capsule.lock().unwrap().clone() { - return Err(WebTransportError::Closed(code, reason).into()); - } + if let Some(err) = self.error.get() { + return Err(err.clone()); } Err(conn_err.into()) } diff --git a/web-transport-quinn/src/send.rs b/web-transport-quinn/src/send.rs index 92754cd6..504ec8ac 100644 --- a/web-transport-quinn/src/send.rs +++ b/web-transport-quinn/src/send.rs @@ -1,12 +1,13 @@ use std::{ io, pin::Pin, + sync::{Arc, OnceLock}, task::{Context, Poll}, }; use bytes::{Buf, Bytes}; -use crate::{session::CloseCapsule, ClosedStream, SessionError, WebTransportError, WriteError}; +use crate::{ClosedStream, SessionError, WriteError}; /// A stream that can be used to send bytes. See [`quinn::SendStream`]. /// @@ -15,26 +16,17 @@ use crate::{session::CloseCapsule, ClosedStream, SessionError, WebTransportError #[derive(Debug)] pub struct SendStream { stream: quinn::SendStream, - close_capsule: CloseCapsule, + error: Arc>, } impl SendStream { - pub(crate) fn new(stream: quinn::SendStream, close_capsule: CloseCapsule) -> Self { - Self { - stream, - close_capsule, - } + pub(crate) fn new(stream: quinn::SendStream, error: Arc>) -> Self { + Self { stream, error } } fn map_write_error(&self, e: quinn::WriteError) -> WriteError { - if let quinn::WriteError::ConnectionLost(ref conn_err) = e { - if matches!(conn_err, quinn::ConnectionError::LocallyClosed) { - if let Some((code, reason)) = self.close_capsule.lock().unwrap().clone() { - return WriteError::SessionError( - WebTransportError::Closed(code, reason).into(), - ); - } - } + if let Some(err) = self.error.get() { + return WriteError::SessionError(err.clone()); } e.into() } @@ -56,10 +48,8 @@ impl SendStream { Ok(Some(code)) => Ok(web_transport_proto::error_from_http3(code.into_inner())), Ok(None) => Ok(None), Err(quinn::StoppedError::ConnectionLost(conn_err)) => { - if matches!(conn_err, quinn::ConnectionError::LocallyClosed) { - if let Some((code, reason)) = self.close_capsule.lock().unwrap().clone() { - return Err(WebTransportError::Closed(code, reason).into()); - } + if let Some(err) = self.error.get() { + return Err(err.clone()); } Err(conn_err.into()) } diff --git a/web-transport-quinn/src/session.rs b/web-transport-quinn/src/session.rs index a5d4d6ef..84df7b93 100644 --- a/web-transport-quinn/src/session.rs +++ b/web-transport-quinn/src/session.rs @@ -4,7 +4,7 @@ use std::{ io::Cursor, ops::Deref, pin::Pin, - sync::{Arc, Mutex}, + sync::{Arc, Mutex, OnceLock}, task::{ready, Context, Poll}, }; @@ -16,10 +16,6 @@ use crate::{ ClientError, Connected, RecvStream, SendStream, SessionError, Settings, WebTransportError, }; -/// Shared close capsule info set by the background task when a -/// `CloseWebTransportSession` capsule is received from the peer. -pub(crate) type CloseCapsule = Arc>>; - /// An established WebTransport session, acting like a full QUIC connection. See [`quinn::Connection`]. /// /// It is important to remember that WebTransport is layered on top of QUIC: @@ -52,9 +48,10 @@ pub struct Session { // Wrapped in Arc>> so close() can take it exactly once. connect_send: Arc>>, - // Close info received from the peer via a CloseWebTransportSession capsule. - // Set by the background task when a capsule is received. - close_capsule: 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, @@ -80,10 +77,10 @@ impl Session { let mut header_datagram = Vec::new(); session_id.encode(&mut header_datagram); - let close_capsule = Arc::new(Mutex::new(None)); + 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, close_capsule.clone()); + let accept = SessionAccept::new(conn.clone(), session_id, error.clone()); let this = Self { conn, @@ -94,31 +91,27 @@ impl Session { header_datagram, settings: Some(Arc::new(settings)), connect_send: Arc::new(Mutex::new(Some(connect.send))), - close_capsule: close_capsule.clone(), + error: error.clone(), request: connect.request.clone(), response: connect.response.clone(), }; // Run a background task to read capsules from the CONNECT recv stream. let conn2 = this.conn.clone(); - let connect_send2 = this.connect_send.clone(); tokio::spawn(async move { let close_info = Self::run_closed(connect.recv).await; + let code = close_info.as_ref().map_or(0, |(c, _)| *c); - // Send capsule close info (if any) for Session::closed() to consume. - // This MUST happen before conn.close() so the info is available - // when conn.closed() resolves on other tasks. - if let Some(ref info) = close_info { - *close_capsule.lock().unwrap() = Some(info.clone()); + // Try to record the remote close error. If close() already set + // the error, it owns the connection teardown, so we bail out. + let err: SessionError = match close_info { + Some((code, reason)) => WebTransportError::Closed(code, reason).into(), + None => quinn::ConnectionError::LocallyClosed.into(), + }; + if error.set(err).is_err() { + return; } - // Take the send stream so that any concurrent/future close() call - // sees None and doesn't try to write a capsule that can't be delivered - // (the peer already closed). - let _send = connect_send2.lock().unwrap().take(); - - // Close the QUIC connection. - 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(); @@ -178,13 +171,13 @@ impl Session { if let Some(accept) = &self.accept { poll_fn(|cx| accept.lock().unwrap().poll_accept_uni(cx)) .await - .map_err(|e| self.map_session_error(e)) + .map_err(|e| self.map_error(e)) } else { self.conn .accept_uni() .await - .map(|recv| RecvStream::new(recv, self.close_capsule.clone())) - .map_err(|e| self.map_conn_error(e)) + .map(|recv| RecvStream::new(recv, self.error.clone())) + .map_err(|e| self.map_error(e)) } } @@ -193,28 +186,24 @@ impl Session { if let Some(accept) = &self.accept { poll_fn(|cx| accept.lock().unwrap().poll_accept_bi(cx)) .await - .map_err(|e| self.map_session_error(e)) + .map_err(|e| self.map_error(e)) } else { self.conn .accept_bi() .await .map(|(send, recv)| { ( - SendStream::new(send, self.close_capsule.clone()), - RecvStream::new(recv, self.close_capsule.clone()), + SendStream::new(send, self.error.clone()), + RecvStream::new(recv, self.error.clone()), ) }) - .map_err(|e| self.map_conn_error(e)) + .map_err(|e| self.map_error(e)) } } /// Open a new unidirectional stream. See [`quinn::Connection::open_uni`]. pub async fn open_uni(&self) -> Result { - let mut send = self - .conn - .open_uni() - .await - .map_err(|e| self.map_conn_error(e))?; + let mut send = self.conn.open_uni().await.map_err(|e| self.map_error(e))?; // Set the stream priority to max and then write the stream header. // Otherwise the application could write data with lower priority than the header, resulting in queuing. @@ -222,20 +211,16 @@ impl Session { send.set_priority(i32::MAX).ok(); Self::write_full(&mut send, &self.header_uni) .await - .map_err(|e| self.map_session_error(e))?; + .map_err(|e| self.map_error(e))?; // Reset the stream priority back to the default of 0. send.set_priority(0).ok(); - Ok(SendStream::new(send, self.close_capsule.clone())) + Ok(SendStream::new(send, self.error.clone())) } /// Open a new bidirectional stream. See [`quinn::Connection::open_bi`]. pub async fn open_bi(&self) -> Result<(SendStream, RecvStream), SessionError> { - let (mut send, recv) = self - .conn - .open_bi() - .await - .map_err(|e| self.map_conn_error(e))?; + let (mut send, recv) = self.conn.open_bi().await.map_err(|e| self.map_error(e))?; // Set the stream priority to max and then write the stream header. // Otherwise the application could write data with lower priority than the header, resulting in queuing. @@ -243,13 +228,13 @@ impl Session { send.set_priority(i32::MAX).ok(); Self::write_full(&mut send, &self.header_bi) .await - .map_err(|e| self.map_session_error(e))?; + .map_err(|e| self.map_error(e))?; // Reset the stream priority back to the default of 0. send.set_priority(0).ok(); Ok(( - SendStream::new(send, self.close_capsule.clone()), - RecvStream::new(recv, self.close_capsule.clone()), + SendStream::new(send, self.error.clone()), + RecvStream::new(recv, self.error.clone()), )) } @@ -263,7 +248,7 @@ impl Session { .conn .read_datagram() .await - .map_err(|e| self.map_session_error(SessionError::from(e)))?; + .map_err(|e| self.map_error(e))?; let mut cursor = Cursor::new(&datagram); @@ -301,7 +286,13 @@ impl Session { self.conn.send_datagram(data) }; - result.map_err(|e| self.map_send_datagram_error(e)) + result.map_err(|e| match self.error.get() { + Some(SessionError::ConnectionError(conn_err)) => { + quinn::SendDatagramError::ConnectionLost(conn_err.clone()).into() + } + Some(err) => err.clone(), + None => e.into(), + }) } /// Computes the maximum size of datagrams that may be passed to @@ -321,38 +312,31 @@ impl Session { /// 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()`] if they need to ensure the capsule - /// has been delivered before proceeding. + /// 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]) { + // 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. Only the first caller gets it. + // 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(); - - // Store the local close info so closed()/map_conn_error() can - // return WebTransportError::Closed even for locally-initiated closes. - // Use get_or_insert to avoid overwriting a peer capsule that arrived first. - self.close_capsule - .lock() - .unwrap() - .get_or_insert_with(|| (code, reason.clone())); - let conn = self.conn.clone(); - let capsule = web_transport_proto::Capsule::CloseWebTransportSession { - code, - reason, - }; + let capsule = + web_transport_proto::Capsule::CloseWebTransportSession { code, reason }; let timeout = self.rtt() * 3; tokio::spawn(async move { Self::close_with_capsule(conn, send, capsule, code, timeout).await; }); - } else { - // Send stream already taken — close() was already called. - // The first call is handling the capsule + connection close. - tracing::debug!("close() called more than once; ignoring duplicate"); } } else { // Raw QUIC mode: no capsule needed. @@ -409,47 +393,24 @@ impl Session { /// /// 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.map_conn_error(self.conn.closed().await) + 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(|e| self.map_conn_error(e)) + self.conn.close_reason().map(|e| self.map_error(e)) } - /// Convert a `quinn::ConnectionError` to `SessionError`, checking the - /// close_capsule so that `LocallyClosed` becomes `WebTransportError::Closed` - /// when a close capsule was received from the peer. - fn map_conn_error(&self, e: quinn::ConnectionError) -> SessionError { - if matches!(e, quinn::ConnectionError::LocallyClosed) { - if let Some((code, reason)) = self.close_capsule.lock().unwrap().clone() { - return WebTransportError::Closed(code, reason).into(); - } - } - e.into() - } - - /// Convert a `SessionError` that wraps `LocallyClosed` into - /// `WebTransportError::Closed` when a close capsule was received. - fn map_session_error(&self, e: SessionError) -> SessionError { - if let SessionError::ConnectionError(ref conn_err) = e { - if matches!(conn_err, quinn::ConnectionError::LocallyClosed) { - if let Some((code, reason)) = self.close_capsule.lock().unwrap().clone() { - return WebTransportError::Closed(code, reason).into(); - } - } - } - e - } - - fn map_send_datagram_error(&self, e: quinn::SendDatagramError) -> SessionError { - if let quinn::SendDatagramError::ConnectionLost(ref conn_err) = e { - if matches!(conn_err, quinn::ConnectionError::LocallyClosed) { - if let Some((code, reason)) = self.close_capsule.lock().unwrap().clone() { - return WebTransportError::Closed(code, reason).into(); - } - } + /// 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() } @@ -480,7 +441,7 @@ impl Session { accept: None, settings: None, connect_send: Arc::new(Mutex::new(None)), - close_capsule: Arc::new(Mutex::new(None)), + error: Arc::new(OnceLock::new()), request: request.into(), response: response.into(), } @@ -549,15 +510,15 @@ pub struct SessionAccept { pending_uni: FuturesUnordered>>, pending_bi: FuturesUnordered>>, - // Close capsule info from the peer, shared with Session. - close_capsule: CloseCapsule, + // Session error, shared with Session. + error: Arc>, } impl SessionAccept { pub(crate) fn new( conn: quinn::Connection, session_id: VarInt, - close_capsule: CloseCapsule, + error: Arc>, ) -> Self { // Create a stream that just outputs new streams, so it's easy to call from poll. let accept_uni = Box::pin(futures::stream::unfold(conn.clone(), |conn| async { @@ -580,15 +541,13 @@ impl SessionAccept { pending_uni: FuturesUnordered::new(), pending_bi: FuturesUnordered::new(), - close_capsule, + error, } } - fn map_conn_error(&self, e: quinn::ConnectionError) -> SessionError { - if matches!(e, quinn::ConnectionError::LocallyClosed) { - if let Some((code, reason)) = self.close_capsule.lock().unwrap().clone() { - return WebTransportError::Closed(code, reason).into(); - } + fn map_error(&self, e: impl Into) -> SessionError { + if let Some(err) = self.error.get() { + return err.clone(); } e.into() } @@ -604,7 +563,7 @@ 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.map_err(|e| self.map_conn_error(e))?; + let recv = res.map_err(|e| self.map_error(e))?; let pending = Self::decode_uni(recv, self.session_id); self.pending_uni.push(Box::pin(pending)); @@ -625,7 +584,7 @@ impl SessionAccept { // Decide if we keep looping based on the type. match typ { StreamUni::WEBTRANSPORT => { - let recv = RecvStream::new(recv, self.close_capsule.clone()); + let recv = RecvStream::new(recv, self.error.clone()); return Poll::Ready(Ok(recv)); } StreamUni::QPACK_DECODER => { @@ -675,7 +634,7 @@ 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.map_err(|e| self.map_conn_error(e))?; + let (send, recv) = res.map_err(|e| self.map_error(e))?; let pending = Self::decode_bi(send, recv, self.session_id); self.pending_bi.push(Box::pin(pending)); @@ -695,8 +654,8 @@ impl SessionAccept { if let Some((send, recv)) = res { // Wrap the streams in our own types for correct error codes. - let send = SendStream::new(send, self.close_capsule.clone()); - let recv = RecvStream::new(recv, self.close_capsule.clone()); + let send = SendStream::new(send, self.error.clone()); + let recv = RecvStream::new(recv, self.error.clone()); return Poll::Ready(Ok((send, recv))); } From 815b61c719d91850678912a146f7d3a5d32bd148 Mon Sep 17 00:00:00 2001 From: ai-and-i <138349030+ai-and-i@users.noreply.github.com> Date: Sat, 28 Feb 2026 10:53:18 -0800 Subject: [PATCH 4/8] Fix lost-waker bug in concurrent accept_bi/accept_uni calls SessionAccept used a single unfold stream for accepting connections, which only stores one waker. When multiple tasks call accept_bi() or accept_uni() concurrently, only the last caller's waker is retained and earlier callers hang forever. Add per-method waker vecs (bi_wakers, uni_wakers) to SessionAccept. On Pending, push the caller's waker; on Ready, drain and wake all stored wakers so other callers retry. Replace the ready!() macro with explicit match to ensure the waker push isn't bypassed by its hidden early return on Poll::Pending. Co-Authored-By: Claude Opus 4.6 --- web-transport-quinn/src/session.rs | 59 ++++++++++++++++++++++++------ 1 file changed, 48 insertions(+), 11 deletions(-) diff --git a/web-transport-quinn/src/session.rs b/web-transport-quinn/src/session.rs index 84df7b93..6ee76371 100644 --- a/web-transport-quinn/src/session.rs +++ b/web-transport-quinn/src/session.rs @@ -5,7 +5,7 @@ use std::{ ops::Deref, pin::Pin, sync::{Arc, Mutex, OnceLock}, - task::{ready, Context, Poll}, + task::{Context, Poll, Waker}, }; use bytes::{Bytes, BytesMut}; @@ -512,6 +512,12 @@ pub struct SessionAccept { // Session error, shared with Session. error: Arc>, + + // 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 { @@ -542,6 +548,9 @@ impl SessionAccept { pending_bi: FuturesUnordered::new(), error, + + bi_wakers: Vec::new(), + uni_wakers: Vec::new(), } } @@ -563,7 +572,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.map_err(|e| self.map_error(e))?; + let recv = match res { + Ok(recv) => recv, + Err(e) => { + for waker in self.uni_wakers.drain(..) { + waker.wake(); + } + return Poll::Ready(Err(self.map_error(e))); + } + }; let pending = Self::decode_uni(recv, self.session_id); self.pending_uni.push(Box::pin(pending)); @@ -571,20 +588,26 @@ 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 => { + 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, self.error.clone()); + for waker in self.uni_wakers.drain(..) { + waker.wake(); + } return Poll::Ready(Ok(recv)); } StreamUni::QPACK_DECODER => { @@ -634,7 +657,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.map_err(|e| self.map_error(e))?; + let (send, recv) = match res { + Ok(pair) => pair, + Err(e) => { + for waker in self.bi_wakers.drain(..) { + waker.wake(); + } + return Poll::Ready(Err(self.map_error(e))); + } + }; let pending = Self::decode_bi(send, recv, self.session_id); self.pending_bi.push(Box::pin(pending)); @@ -642,20 +673,26 @@ 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 => { + 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, self.error.clone()); let recv = RecvStream::new(recv, self.error.clone()); + for waker in self.bi_wakers.drain(..) { + waker.wake(); + } return Poll::Ready(Ok((send, recv))); } From 9546358cfaa197c8d73f216e7c3bf904b65f4907 Mon Sep 17 00:00:00 2001 From: Luke Curley Date: Sat, 28 Feb 2026 20:18:35 -0800 Subject: [PATCH 5/8] Simplify the PR a bit. --- web-transport-proto/src/capsule.rs | 8 +- web-transport-quinn/src/recv.rs | 48 ++---------- web-transport-quinn/src/send.rs | 42 ++-------- web-transport-quinn/src/session.rs | 120 ++++++++++------------------- 4 files changed, 59 insertions(+), 159 deletions(-) diff --git a/web-transport-proto/src/capsule.rs b/web-transport-proto/src/capsule.rs index ee7d76b0..03ea4f67 100644 --- a/web-transport-proto/src/capsule.rs +++ b/web-transport-proto/src/capsule.rs @@ -5,10 +5,7 @@ use tokio::io::{AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt}; 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)] @@ -216,9 +213,6 @@ pub enum CapsuleError { #[error("message too long")] MessageTooLong, - #[error("unknown capsule type: {0:?}")] - UnknownType(VarInt), - #[error("varint decode error: {0:?}")] VarInt(#[from] VarIntUnexpectedEnd), diff --git a/web-transport-quinn/src/recv.rs b/web-transport-quinn/src/recv.rs index 96208b07..2cda7293 100644 --- a/web-transport-quinn/src/recv.rs +++ b/web-transport-quinn/src/recv.rs @@ -1,7 +1,6 @@ use std::{ io, pin::Pin, - sync::{Arc, OnceLock}, task::{Context, Poll}, }; @@ -13,22 +12,11 @@ use crate::{ReadError, ReadExactError, ReadToEndError, SessionError}; #[derive(Debug)] pub struct RecvStream { inner: quinn::RecvStream, - error: Arc>, } impl RecvStream { - pub(crate) fn new(stream: quinn::RecvStream, error: Arc>) -> Self { - Self { - inner: stream, - error, - } - } - - fn map_read_error(&self, e: quinn::ReadError) -> ReadError { - if let Some(err) = self.error.get() { - return ReadError::SessionError(err.clone()); - } - e.into() + pub(crate) fn new(stream: quinn::RecvStream) -> Self { + Self { inner: stream } } /// Tell the other end to stop sending data with the given error code. See [`quinn::RecvStream::stop`]. @@ -43,21 +31,12 @@ impl RecvStream { /// Read some data into the buffer and return the amount read. See [`quinn::RecvStream::read`]. pub async fn read(&mut self, buf: &mut [u8]) -> Result, ReadError> { - self.inner - .read(buf) - .await - .map_err(|e| self.map_read_error(e)) + self.inner.read(buf).await.map_err(Into::into) } /// Fill the entire buffer with data. See [`quinn::RecvStream::read_exact`]. pub async fn read_exact(&mut self, buf: &mut [u8]) -> Result<(), ReadExactError> { - self.inner - .read_exact(buf) - .await - .map_err(|e| match e { - quinn::ReadExactError::ReadError(e) => self.map_read_error(e).into(), - e => e.into(), - }) + self.inner.read_exact(buf).await.map_err(Into::into) } /// Read a chunk of data from the stream. See [`quinn::RecvStream::read_chunk`]. @@ -69,15 +48,12 @@ impl RecvStream { self.inner .read_chunk(max_length, ordered) .await - .map_err(|e| self.map_read_error(e)) + .map_err(Into::into) } /// Read chunks of data from the stream. See [`quinn::RecvStream::read_chunks`]. pub async fn read_chunks(&mut self, bufs: &mut [Bytes]) -> Result, ReadError> { - self.inner - .read_chunks(bufs) - .await - .map_err(|e| self.map_read_error(e)) + self.inner.read_chunks(bufs).await.map_err(Into::into) } /// Read until the end of the stream or the limit is hit. See [`quinn::RecvStream::read_to_end`]. @@ -85,10 +61,7 @@ impl RecvStream { self.inner .read_to_end(size_limit) .await - .map_err(|e| match e { - quinn::ReadToEndError::Read(e) => self.map_read_error(e).into(), - e => e.into(), - }) + .map_err(Into::into) } /// Block until the stream has been reset and return the error code. See [`quinn::RecvStream::received_reset`]. @@ -100,12 +73,7 @@ impl RecvStream { Ok(Some(code)) => Ok(Some( web_transport_proto::error_from_http3(code.into_inner()).unwrap(), )), - Err(quinn::ResetError::ConnectionLost(conn_err)) => { - if let Some(err) = self.error.get() { - return Err(err.clone()); - } - Err(conn_err.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 504ec8ac..0335aaa2 100644 --- a/web-transport-quinn/src/send.rs +++ b/web-transport-quinn/src/send.rs @@ -1,7 +1,6 @@ use std::{ io, pin::Pin, - sync::{Arc, OnceLock}, task::{Context, Poll}, }; @@ -16,19 +15,11 @@ use crate::{ClosedStream, SessionError, WriteError}; #[derive(Debug)] pub struct SendStream { stream: quinn::SendStream, - error: Arc>, } impl SendStream { - pub(crate) fn new(stream: quinn::SendStream, error: Arc>) -> Self { - Self { stream, error } - } - - fn map_write_error(&self, e: quinn::WriteError) -> WriteError { - if let Some(err) = self.error.get() { - return WriteError::SessionError(err.clone()); - } - e.into() + pub(crate) fn new(stream: quinn::SendStream) -> Self { + Self { stream } } /// Abruptly reset the stream with the provided error code. See [`quinn::SendStream::reset`]. @@ -47,12 +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(conn_err)) => { - if let Some(err) = self.error.get() { - return Err(err.clone()); - } - Err(conn_err.into()) - } + Err(quinn::StoppedError::ConnectionLost(conn_err)) => Err(conn_err.into()), Err(quinn::StoppedError::ZeroRttRejected) => unreachable!("0-RTT not supported"), } } @@ -61,34 +47,22 @@ impl SendStream { /// Write some data to the stream, returning the size written. See [`quinn::SendStream::write`]. pub async fn write(&mut self, buf: &[u8]) -> Result { - self.stream - .write(buf) - .await - .map_err(|e| self.map_write_error(e)) + self.stream.write(buf).await.map_err(Into::into) } /// Write all of the data to the stream. See [`quinn::SendStream::write_all`]. pub async fn write_all(&mut self, buf: &[u8]) -> Result<(), WriteError> { - self.stream - .write_all(buf) - .await - .map_err(|e| self.map_write_error(e)) + self.stream.write_all(buf).await.map_err(Into::into) } /// Write chunks of data to the stream. See [`quinn::SendStream::write_chunks`]. pub async fn write_chunks(&mut self, bufs: &mut [Bytes]) -> Result { - self.stream - .write_chunks(bufs) - .await - .map_err(|e| self.map_write_error(e)) + self.stream.write_chunks(bufs).await.map_err(Into::into) } /// Write a chunk of data to the stream. See [`quinn::SendStream::write_chunk`]. pub async fn write_chunk(&mut self, buf: Bytes) -> Result<(), WriteError> { - self.stream - .write_chunk(buf) - .await - .map_err(|e| self.map_write_error(e)) + self.stream.write_chunk(buf).await.map_err(Into::into) } /// Write all of the chunks of data to the stream. See [`quinn::SendStream::write_all_chunks`]. @@ -96,7 +70,7 @@ impl SendStream { self.stream .write_all_chunks(bufs) .await - .map_err(|e| self.map_write_error(e)) + .map_err(Into::into) } /// Mark the stream as finished, such that no more data can be written. See [`quinn::SendStream::finish`]. diff --git a/web-transport-quinn/src/session.rs b/web-transport-quinn/src/session.rs index 6ee76371..06464f66 100644 --- a/web-transport-quinn/src/session.rs +++ b/web-transport-quinn/src/session.rs @@ -6,6 +6,7 @@ use std::{ pin::Pin, sync::{Arc, Mutex, OnceLock}, task::{Context, Poll, Waker}, + time::Duration, }; use bytes::{Bytes, BytesMut}; @@ -80,7 +81,7 @@ impl Session { 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, error.clone()); + let accept = SessionAccept::new(conn.clone(), session_id); let this = Self { conn, @@ -102,20 +103,27 @@ impl Session { let close_info = Self::run_closed(connect.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. let err: SessionError = match close_info { - Some((code, reason)) => WebTransportError::Closed(code, reason).into(), + Some((code, reason)) => { + let err = WebTransportError::Closed(code, reason.clone()); + if error.set(err.into()).is_err() { + return; + } + conn2.close(http3_code, reason.as_bytes()); + return; + } None => quinn::ConnectionError::LocallyClosed.into(), }; if error.set(err).is_err() { return; } - - let http3_code: quinn::VarInt = web_transport_proto::error_to_http3(code) - .try_into() - .unwrap(); - conn2.close(http3_code, &[]); + conn2.close(http3_code, b""); }); this @@ -169,73 +177,50 @@ 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 - .map_err(|e| self.map_error(e)) + Ok(poll_fn(|cx| accept.lock().unwrap().poll_accept_uni(cx)).await?) } else { - self.conn - .accept_uni() - .await - .map(|recv| RecvStream::new(recv, self.error.clone())) - .map_err(|e| self.map_error(e)) + 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 - .map_err(|e| self.map_error(e)) + Ok(poll_fn(|cx| accept.lock().unwrap().poll_accept_bi(cx)).await?) } else { - self.conn - .accept_bi() - .await - .map(|(send, recv)| { - ( - SendStream::new(send, self.error.clone()), - RecvStream::new(recv, self.error.clone()), - ) - }) - .map_err(|e| self.map_error(e)) + let (send, recv) = self.conn.accept_bi().await?; + Ok((SendStream::new(send), RecvStream::new(recv))) } } /// Open a new unidirectional stream. See [`quinn::Connection::open_uni`]. pub async fn open_uni(&self) -> Result { - let mut send = self.conn.open_uni().await.map_err(|e| self.map_error(e))?; + let mut send = self.conn.open_uni().await?; // Set the stream priority to max and then write the stream header. // Otherwise the application could write data with lower priority than the header, resulting in queuing. // Also the header is very important for determining the session ID without reliable reset. send.set_priority(i32::MAX).ok(); - Self::write_full(&mut send, &self.header_uni) - .await - .map_err(|e| self.map_error(e))?; + Self::write_full(&mut send, &self.header_uni).await?; // Reset the stream priority back to the default of 0. send.set_priority(0).ok(); - Ok(SendStream::new(send, self.error.clone())) + Ok(SendStream::new(send)) } /// Open a new bidirectional stream. See [`quinn::Connection::open_bi`]. pub async fn open_bi(&self) -> Result<(SendStream, RecvStream), SessionError> { - let (mut send, recv) = self.conn.open_bi().await.map_err(|e| self.map_error(e))?; + let (mut send, recv) = self.conn.open_bi().await?; // Set the stream priority to max and then write the stream header. // Otherwise the application could write data with lower priority than the header, resulting in queuing. // Also the header is very important for determining the session ID without reliable reset. send.set_priority(i32::MAX).ok(); - Self::write_full(&mut send, &self.header_bi) - .await - .map_err(|e| self.map_error(e))?; + Self::write_full(&mut send, &self.header_bi).await?; // Reset the stream priority back to the default of 0. send.set_priority(0).ok(); - Ok(( - SendStream::new(send, self.error.clone()), - RecvStream::new(recv, self.error.clone()), - )) + Ok((SendStream::new(send), RecvStream::new(recv))) } /// Asynchronously receives an application datagram from the remote peer. @@ -244,11 +229,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(|e| self.map_error(e))?; + let mut datagram = self.conn.read_datagram().await?; let mut cursor = Cursor::new(&datagram); @@ -286,13 +267,8 @@ impl Session { self.conn.send_datagram(data) }; - result.map_err(|e| match self.error.get() { - Some(SessionError::ConnectionError(conn_err)) => { - quinn::SendDatagramError::ConnectionLost(conn_err.clone()).into() - } - Some(err) => err.clone(), - None => e.into(), - }) + result?; + Ok(()) } /// Computes the maximum size of datagrams that may be passed to @@ -332,7 +308,7 @@ impl Session { let conn = self.conn.clone(); let capsule = web_transport_proto::Capsule::CloseWebTransportSession { code, reason }; - let timeout = self.rtt() * 3; + let timeout = (self.rtt() * 3).max(Duration::from_millis(100)); tokio::spawn(async move { Self::close_with_capsule(conn, send, capsule, code, timeout).await; @@ -510,9 +486,6 @@ pub struct SessionAccept { pending_uni: FuturesUnordered>>, pending_bi: FuturesUnordered>>, - // Session error, shared with Session. - error: Arc>, - // 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. @@ -521,11 +494,7 @@ pub struct SessionAccept { } impl SessionAccept { - pub(crate) fn new( - conn: quinn::Connection, - session_id: VarInt, - error: Arc>, - ) -> Self { + pub(crate) fn new(conn: quinn::Connection, session_id: VarInt) -> Self { // Create a stream that just outputs new streams, so it's easy to call from poll. let accept_uni = Box::pin(futures::stream::unfold(conn.clone(), |conn| async { Some((conn.accept_uni().await, conn)) @@ -547,20 +516,11 @@ impl SessionAccept { pending_uni: FuturesUnordered::new(), pending_bi: FuturesUnordered::new(), - error, - bi_wakers: Vec::new(), uni_wakers: Vec::new(), } } - fn map_error(&self, e: impl Into) -> SessionError { - if let Some(err) = self.error.get() { - return err.clone(); - } - e.into() - } - // This is poll-based because we accept and decode streams in parallel. // In async land I would use tokio::JoinSet, but that requires a runtime. // It's better to use FuturesUnordered instead because it's agnostic. @@ -578,7 +538,7 @@ impl SessionAccept { for waker in self.uni_wakers.drain(..) { waker.wake(); } - return Poll::Ready(Err(self.map_error(e))); + return Poll::Ready(Err(e.into())); } }; let pending = Self::decode_uni(recv, self.session_id); @@ -596,7 +556,9 @@ impl SessionAccept { continue; } Poll::Ready(None) | Poll::Pending => { - self.uni_wakers.push(cx.waker().clone()); + if !self.uni_wakers.iter().any(|w| w.will_wake(cx.waker())) { + self.uni_wakers.push(cx.waker().clone()); + } return Poll::Pending; } }; @@ -604,7 +566,7 @@ impl SessionAccept { // Decide if we keep looping based on the type. match typ { StreamUni::WEBTRANSPORT => { - let recv = RecvStream::new(recv, self.error.clone()); + let recv = RecvStream::new(recv); for waker in self.uni_wakers.drain(..) { waker.wake(); } @@ -663,7 +625,7 @@ impl SessionAccept { for waker in self.bi_wakers.drain(..) { waker.wake(); } - return Poll::Ready(Err(self.map_error(e))); + return Poll::Ready(Err(e.into())); } }; let pending = Self::decode_bi(send, recv, self.session_id); @@ -681,15 +643,17 @@ impl SessionAccept { continue; } Poll::Ready(None) | Poll::Pending => { - self.bi_wakers.push(cx.waker().clone()); + 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, self.error.clone()); - let recv = RecvStream::new(recv, self.error.clone()); + let send = SendStream::new(send); + let recv = RecvStream::new(recv); for waker in self.bi_wakers.drain(..) { waker.wake(); } From 08dcde7e4a3b9485b988efd02dfe940e09c411a0 Mon Sep 17 00:00:00 2001 From: Luke Curley Date: Sat, 28 Feb 2026 20:19:12 -0800 Subject: [PATCH 6/8] just fix --- web-transport-proto/src/capsule.rs | 35 ++++++++++++++++++++++++------ web-transport-quinn/src/recv.rs | 5 +---- web-transport-quinn/src/send.rs | 5 +---- 3 files changed, 30 insertions(+), 15 deletions(-) diff --git a/web-transport-proto/src/capsule.rs b/web-transport-proto/src/capsule.rs index 03ea4f67..c95d3d68 100644 --- a/web-transport-proto/src/capsule.rs +++ b/web-transport-proto/src/capsule.rs @@ -598,7 +598,10 @@ mod tests { #[tokio::test] async fn test_http3_reader_single_capsule_single_frame() { - let capsule = Capsule::CloseWebTransportSession { code: 42, reason: "bye".into() }; + 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()); @@ -606,7 +609,10 @@ mod tests { #[tokio::test] async fn test_http3_reader_capsule_split_across_frames() { - let capsule = Capsule::CloseWebTransportSession { code: 100, reason: "split".into() }; + 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); @@ -614,7 +620,10 @@ mod tests { #[tokio::test] async fn test_http3_reader_capsule_split_across_three_frames() { - let capsule = Capsule::CloseWebTransportSession { code: 200, reason: "three".into() }; + 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); @@ -623,7 +632,10 @@ mod tests { #[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 c2 = Capsule::CloseWebTransportSession { + code: 7, + reason: "done".into(), + }; let mut bytes = encode_capsule(&c1); bytes.extend_from_slice(&encode_capsule(&c2)); @@ -634,7 +646,10 @@ mod tests { #[tokio::test] async fn test_http3_reader_skips_non_data_frames() { - let capsule = Capsule::CloseWebTransportSession { code: 0, reason: String::new() }; + 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); @@ -658,12 +673,18 @@ mod tests { // Missing: length + payload let mut reader = reader_from(wrap_in_data_frame(&partial)); - assert!(matches!(reader.read().await, Err(CapsuleError::UnexpectedEnd))); + 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 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))); diff --git a/web-transport-quinn/src/recv.rs b/web-transport-quinn/src/recv.rs index 2cda7293..429281c6 100644 --- a/web-transport-quinn/src/recv.rs +++ b/web-transport-quinn/src/recv.rs @@ -58,10 +58,7 @@ impl RecvStream { /// Read until the end of the stream or the limit is hit. See [`quinn::RecvStream::read_to_end`]. pub async fn read_to_end(&mut self, size_limit: usize) -> Result, ReadToEndError> { - self.inner - .read_to_end(size_limit) - .await - .map_err(Into::into) + self.inner.read_to_end(size_limit).await.map_err(Into::into) } /// Block until the stream has been reset and return the error code. See [`quinn::RecvStream::received_reset`]. diff --git a/web-transport-quinn/src/send.rs b/web-transport-quinn/src/send.rs index 0335aaa2..65bfdaeb 100644 --- a/web-transport-quinn/src/send.rs +++ b/web-transport-quinn/src/send.rs @@ -67,10 +67,7 @@ impl SendStream { /// Write all of the chunks of data to the stream. See [`quinn::SendStream::write_all_chunks`]. pub async fn write_all_chunks(&mut self, bufs: &mut [Bytes]) -> Result<(), WriteError> { - self.stream - .write_all_chunks(bufs) - .await - .map_err(Into::into) + self.stream.write_all_chunks(bufs).await.map_err(Into::into) } /// Mark the stream as finished, such that no more data can be written. See [`quinn::SendStream::finish`]. From 4c8e27a982743bcf5cebd5c058e6012305b3a1c8 Mon Sep 17 00:00:00 2001 From: Luke Curley Date: Sat, 28 Feb 2026 20:35:35 -0800 Subject: [PATCH 7/8] Extract spawn body into run_recv method Co-Authored-By: Claude Opus 4.6 --- web-transport-quinn/src/session.rs | 61 +++++++++++++++++------------- 1 file changed, 35 insertions(+), 26 deletions(-) diff --git a/web-transport-quinn/src/session.rs b/web-transport-quinn/src/session.rs index 06464f66..c3bcef12 100644 --- a/web-transport-quinn/src/session.rs +++ b/web-transport-quinn/src/session.rs @@ -99,40 +99,49 @@ impl Session { // Run a background task to read capsules from the CONNECT recv stream. let conn2 = this.conn.clone(); - tokio::spawn(async move { - let close_info = Self::run_closed(connect.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. - let err: SessionError = match close_info { - Some((code, reason)) => { - let err = WebTransportError::Closed(code, reason.clone()); - if error.set(err.into()).is_err() { - return; - } - conn2.close(http3_code, reason.as_bytes()); + tokio::spawn(Self::run_recv(conn2, connect.recv, error)); + + this + } + + // 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; } - None => quinn::ConnectionError::LocallyClosed.into(), - }; - if error.set(err).is_err() { - return; + conn.close(http3_code, reason.as_bytes()); } - conn2.close(http3_code, b""); - }); - - this + 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 run_closed(recv: quinn::RecvStream) -> Option<(u32, String)> { + async fn read_capsules(recv: quinn::RecvStream) -> Option<(u32, String)> { let mut reader = web_transport_proto::Http3CapsuleReader::new(recv); loop { match reader.read().await { From 975034f2451f0765685bd7036bf039630d0fc398 Mon Sep 17 00:00:00 2001 From: Luke Curley Date: Sun, 1 Mar 2026 13:27:55 -0800 Subject: [PATCH 8/8] PR review. --- web-transport-proto/src/capsule.rs | 21 ++++++++++++--------- web-transport-proto/src/varint.rs | 27 +++++++++++++++++++++------ web-transport-quinn/src/session.rs | 7 ++++++- 3 files changed, 39 insertions(+), 16 deletions(-) diff --git a/web-transport-proto/src/capsule.rs b/web-transport-proto/src/capsule.rs index c95d3d68..ca62781a 100644 --- a/web-transport-proto/src/capsule.rs +++ b/web-transport-proto/src/capsule.rs @@ -78,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 @@ -273,9 +274,10 @@ impl Http3CapsuleReader { /// Returns `false` on EOF. async fn read_data_frame(&mut self) -> Result { loop { - let frame_type = match VarInt::read(&mut self.stream).await { - Ok(v) => Frame(v), - Err(_) => return Ok(false), + 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 @@ -287,9 +289,10 @@ impl Http3CapsuleReader { } if frame_type != Frame::DATA { - if len > 0 { - let mut skip = vec![0u8; len]; - self.stream.read_exact(&mut skip).await?; + 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; } 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/src/session.rs b/web-transport-quinn/src/session.rs index c3bcef12..b45f534a 100644 --- a/web-transport-quinn/src/session.rs +++ b/web-transport-quinn/src/session.rs @@ -350,7 +350,12 @@ impl Session { let mut frame = Vec::new(); Frame::DATA.encode(&mut frame); - VarInt::from_u32(capsule_bytes.len() as u32).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.