diff --git a/rs/web-transport-quinn/src/error.rs b/rs/web-transport-quinn/src/error.rs index 10b34797..11880256 100644 --- a/rs/web-transport-quinn/src/error.rs +++ b/rs/web-transport-quinn/src/error.rs @@ -1,4 +1,4 @@ -use std::sync::Arc; +use std::sync::{Arc, OnceLock}; use thiserror::Error; @@ -67,6 +67,72 @@ impl From for SessionError { } } +/// The session's first close reason and its application-code space. +#[derive(Debug)] +pub(crate) struct CloseReason { + // The raw QUIC connection, whose application codes bypass the HTTP/3 mapping. + raw: Option, + reason: OnceLock, +} + +impl CloseReason { + pub(crate) fn http3() -> Self { + Self { + raw: None, + reason: OnceLock::new(), + } + } + + pub(crate) fn raw(conn: quinn::Connection) -> Self { + Self { + raw: Some(conn), + reason: OnceLock::new(), + } + } + + pub(crate) fn get(&self) -> Option<&SessionError> { + self.reason.get() + } + + pub(crate) fn set(&self, err: SessionError) -> Result<(), SessionError> { + self.reason.set(err) + } + + pub(crate) fn map(&self, err: SessionError) -> SessionError { + if !matches!( + &err, + SessionError::ConnectionError(_) + | SessionError::WebTransportError(WebTransportError::Closed(..)) + | SessionError::SendDatagramError(quinn::SendDatagramError::ConnectionLost(_)) + ) { + return err; + } + if let Some(reason) = self.get() { + return reason.clone(); + } + + // `err` was already decoded as HTTP/3, so a raw session decodes the connection's own reason. + // Raw sessions have no capsule task, so the peer close is latched here. + let Some(quinn::ConnectionError::ApplicationClosed(close)) = + self.raw.as_ref().and_then(|conn| conn.close_reason()) + else { + return err; + }; + self.reason + .get_or_init(|| match u32::try_from(close.error_code.into_inner()) { + Ok(code) => WebTransportError::Closed( + code, + String::from_utf8_lossy(&close.reason).into_owned(), + ) + .into(), + Err(_) => { + SessionError::ConnectionError(quinn::ConnectionError::ApplicationClosed(close)) + } + }) + .clone() + } +} + /// The QUIC code to reset or stop a stream with. Raw QUIC sends the code as is, while /// HTTP/3 maps it into its own code space. pub(crate) fn encode_stream_code(code: u32, raw: bool) -> quinn::VarInt { @@ -319,3 +385,28 @@ impl web_transport_trait::Error for ReadError { } } } + +#[cfg(test)] +mod tests { + use super::*; + use web_transport_trait::Error as _; + + fn peer_close(code: u64) -> SessionError { + quinn::ConnectionError::ApplicationClosed(quinn::ApplicationClose { + error_code: quinn::VarInt::from_u64(code).unwrap(), + reason: b"peer closed".as_slice().into(), + }) + .into() + } + + #[test] + fn http3_close_keeps_its_code_space() { + let h3 = CloseReason::http3(); + assert_eq!(h3.map(peer_close(4075)).session_error(), None); + assert_eq!( + h3.map(peer_close(web_transport_proto::error_to_http3(4075))) + .session_error(), + Some((4075, "peer closed".into())) + ); + } +} diff --git a/rs/web-transport-quinn/src/recv.rs b/rs/web-transport-quinn/src/recv.rs index 43f3978c..c7e40c4e 100644 --- a/rs/web-transport-quinn/src/recv.rs +++ b/rs/web-transport-quinn/src/recv.rs @@ -1,8 +1,9 @@ +use crate::error::CloseReason; use std::{ future::Future, io, pin::Pin, - sync::{Arc, OnceLock}, + sync::Arc, task::{ready, Context, Poll}, }; @@ -17,17 +18,13 @@ use crate::{ #[derive(Debug)] pub struct RecvStream { inner: quinn::RecvStream, - error: Arc>, + error: Arc, // Raw QUIC carries stream codes as is; HTTP/3 maps them into its own code space. raw: bool, } impl RecvStream { - pub(crate) fn new( - stream: quinn::RecvStream, - error: Arc>, - raw: bool, - ) -> Self { + pub(crate) fn new(stream: quinn::RecvStream, error: Arc, raw: bool) -> Self { Self { inner: stream, error, @@ -50,7 +47,10 @@ impl RecvStream { return ReadError::SessionError(err.clone()); } } - e + match e { + ReadError::SessionError(err) => ReadError::SessionError(self.error.map(err)), + err => err, + } } /// Tell the other end to stop sending data with the given error code. See [`quinn::RecvStream::stop`]. @@ -115,7 +115,7 @@ impl RecvStream { Ok(None) => Ok(None), Ok(Some(code)) => Ok(decode_stream_code(code, self.raw)), Err(quinn::ResetError::ConnectionLost(conn_err)) => { - Err(self.error.get().cloned().unwrap_or_else(|| conn_err.into())) + Err(self.error.map(conn_err.into())) } Err(quinn::ResetError::ZeroRttRejected) => unreachable!("0-RTT not supported"), } diff --git a/rs/web-transport-quinn/src/send.rs b/rs/web-transport-quinn/src/send.rs index e4e9901a..60173e25 100644 --- a/rs/web-transport-quinn/src/send.rs +++ b/rs/web-transport-quinn/src/send.rs @@ -1,7 +1,8 @@ +use crate::error::CloseReason; use std::{ io, pin::Pin, - sync::{Arc, OnceLock}, + sync::Arc, task::{ready, Context, Poll}, }; @@ -19,7 +20,7 @@ use crate::{ #[derive(Debug)] pub struct SendStream { stream: quinn::SendStream, - error: Arc>, + error: Arc, // Raw QUIC carries stream codes as is; HTTP/3 maps them into its own code space. raw: bool, @@ -29,11 +30,7 @@ pub struct SendStream { } impl SendStream { - pub(crate) fn new( - stream: quinn::SendStream, - error: Arc>, - raw: bool, - ) -> Self { + pub(crate) fn new(stream: quinn::SendStream, error: Arc, raw: bool) -> Self { Self { stream, error, @@ -60,7 +57,10 @@ impl SendStream { return WriteError::SessionError(err.clone()); } } - e + match e { + WriteError::SessionError(err) => WriteError::SessionError(self.error.map(err)), + err => err, + } } /// Abruptly reset the stream with the provided error code. See [`quinn::SendStream::reset`]. @@ -80,7 +80,7 @@ impl SendStream { Ok(Some(code)) => Ok(decode_stream_code(code, self.raw)), Ok(None) => Ok(None), Err(quinn::StoppedError::ConnectionLost(conn_err)) => { - Err(self.error.get().cloned().unwrap_or_else(|| conn_err.into())) + Err(self.error.map(conn_err.into())) } Err(quinn::StoppedError::ZeroRttRejected) => unreachable!("0-RTT not supported"), } @@ -260,9 +260,7 @@ impl web_transport_trait::poll::SendStream for SendStream { }, Ok(None) => Ok(()), Err(quinn::StoppedError::ConnectionLost(conn_err)) => { - Err(WriteError::SessionError( - error.get().cloned().unwrap_or_else(|| conn_err.into()), - )) + Err(WriteError::SessionError(error.map(conn_err.into()))) } Err(quinn::StoppedError::ZeroRttRejected) => unreachable!("0-RTT not supported"), } diff --git a/rs/web-transport-quinn/src/session.rs b/rs/web-transport-quinn/src/session.rs index e68dfd79..ba70161b 100644 --- a/rs/web-transport-quinn/src/session.rs +++ b/rs/web-transport-quinn/src/session.rs @@ -1,3 +1,4 @@ +use crate::error::CloseReason; use std::{ fmt, future::Future, @@ -82,7 +83,7 @@ pub struct Session { // 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>, + error: Arc, // The request sent by the client, or None for a raw QUIC session. request: Option, @@ -141,7 +142,7 @@ impl Session { let mut header_datagram = Vec::new(); session_id.encode(&mut header_datagram); - let error: Arc> = Arc::new(OnceLock::new()); + let error = Arc::new(CloseReason::http3()); // Accept logic is stateful, so use an Arc to share it. let accept = SessionAccept::new(conn.clone(), session_id, error.clone()); @@ -179,11 +180,7 @@ impl Session { // 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>, - ) { + 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); @@ -401,7 +398,10 @@ impl Session { 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. + // this is a no-op. A raw peer close that already arrived is latched first. + if let Some(err) = self.conn.close_reason() { + self.map_error(err); + } let err = SessionError::ConnectionError(quinn::ConnectionError::LocallyClosed); if self.error.set(err).is_err() { return; @@ -515,27 +515,14 @@ impl Session { } // The owned-argument form, for futures that outlive the borrow that created them. - fn map_error_owned(error: &OnceLock, e: impl Into) -> SessionError { + fn map_error_owned(error: &CloseReason, e: impl Into) -> SessionError { map_error_with(error, e) } } /// Replace connection-level errors with the stored session error if available. -fn map_error_with(stored: &OnceLock, e: impl Into) -> SessionError { - { - let e = e.into(); - if let Some(err) = stored.get() { - if matches!( - &e, - SessionError::ConnectionError(_) - | SessionError::WebTransportError(WebTransportError::Closed(..)) - | SessionError::SendDatagramError(quinn::SendDatagramError::ConnectionLost(_)) - ) { - return err.clone(); - } - } - e - } +fn map_error_with(stored: &CloseReason, e: impl Into) -> SessionError { + stored.map(e.into()) } impl Session { @@ -558,6 +545,7 @@ impl Session { /// (from `into_0rtt`, say) is fine, it just has no ALPN to report until it is done. pub fn raw(conn: quinn::Connection) -> Self { Self { + error: Arc::new(CloseReason::raw(conn.clone())), conn, session_id: None, header_uni: Default::default(), @@ -575,7 +563,6 @@ impl Session { parked_accept_bi: Default::default(), settings: None, connect_send: Arc::new(Mutex::new(None)), - error: Arc::new(OnceLock::new()), request: None, response: None, alpn: Default::default(), @@ -721,7 +708,7 @@ pub struct SessionAccept { session_id: VarInt, // Shared session error for propagation to accepted streams. - error: Arc>, + error: Arc, // We also need to keep a reference to the qpack streams if the endpoint (incorrectly) creates them. // Again, this is just so they don't get closed until we drop the session. @@ -754,7 +741,7 @@ impl SessionAccept { pub(crate) fn new( conn: quinn::Connection, session_id: VarInt, - error: Arc>, + 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 { @@ -1086,7 +1073,7 @@ impl Session { async fn open_uni_owned( conn: quinn::Connection, header: Bytes, - error: Arc>, + error: Arc, raw: bool, ) -> Result { let mut send = conn @@ -1108,7 +1095,7 @@ impl Session { async fn open_bi_owned( conn: quinn::Connection, header: Bytes, - error: Arc>, + error: Arc, raw: bool, ) -> Result<(SendStream, RecvStream), SessionError> { let (mut send, recv) = conn @@ -1171,7 +1158,7 @@ impl Session { async fn read_datagram_owned( conn: quinn::Connection, session_id: Option, - error: Arc>, + error: Arc, ) -> Result { let mut datagram = conn .read_datagram() diff --git a/rs/web-transport-quinn/tests/raw.rs b/rs/web-transport-quinn/tests/raw.rs index 58a8da4d..bdc66023 100644 --- a/rs/web-transport-quinn/tests/raw.rs +++ b/rs/web-transport-quinn/tests/raw.rs @@ -236,6 +236,83 @@ async fn raw_session_streams_omit_webtransport_header() -> Result<()> { Ok(()) } +#[tokio::test] +async fn raw_peer_close_code_is_reported_by_session_and_streams() -> Result<()> { + use web_transport_trait::Error as _; + let (client, peer) = connect_raw().await?; + let session = Session::raw(client); + let mut send = session.open_uni().await?; + send.write_all(b"hello").await?; + let mut peer_recv = peer.accept_uni().await?; + let mut hello = [0; 5]; + peer_recv.read_exact(&mut hello).await?; + let mut peer_send = peer.open_uni().await?; + peer_send.write_all(b"hello").await?; + let mut recv = session.accept_uni().await?; + recv.read_exact(&mut hello).await?; + peer.close(4075u32.into(), b"application close"); + let closed = session.closed().await; + assert_eq!( + closed.session_error(), + Some((4075, "application close".into())) + ); + assert_eq!( + session.close_reason().unwrap().session_error(), + closed.session_error() + ); + assert_eq!( + send.write_all(b"after close") + .await + .unwrap_err() + .session_error(), + closed.session_error() + ); + assert_eq!( + recv.read(&mut hello).await.unwrap_err().session_error(), + closed.session_error() + ); + let stopped = futures::future::poll_fn(|cx| { + web_transport_trait::poll::SendStream::poll_closed(&mut send, cx) + }) + .await + .unwrap_err(); + assert_eq!(stopped.session_error(), closed.session_error()); + Ok(()) +} + +#[tokio::test] +async fn raw_peer_close_survives_local_close() -> Result<()> { + use web_transport_trait::Error as _; + let (client, peer) = connect_raw().await?; + let session = Session::raw(client.clone()); + peer.close(4075u32.into(), b"application close"); + // Wait on the quinn connection so the session never observes the close first. + client.closed().await; + session.close(1, b"cleanup"); + assert_eq!( + session.close_reason().unwrap().session_error(), + Some((4075, "application close".into())) + ); + Ok(()) +} + +#[tokio::test] +async fn raw_peer_close_above_u32_is_not_a_session_error() -> Result<()> { + use web_transport_trait::Error as _; + // Includes a code that would decode as WebTransport code 7 under the HTTP/3 mapping. + for code in [u32::MAX as u64 + 1, web_transport_proto::error_to_http3(7)] { + let (client, peer) = connect_raw().await?; + let session = Session::raw(client); + peer.close(quinn::VarInt::from_u64(code)?, b"too big"); + assert_eq!( + session.closed().await.session_error(), + None, + "code {code:#x}" + ); + } + Ok(()) +} + /// A raw session sends RESET_STREAM and STOP_SENDING codes as is, so a plain QUIC /// peer agrees on them, and still reads the HTTP/3-mapped codes an older raw peer sends. #[tokio::test]