Fix error propagation when connection is closed (#181)

* Propagate session close error to all in-progress operations

Share Arc<OnceLock<SessionError>> with SendStream, RecvStream, and
SessionAccept so that when a CloseWebTransportSession capsule is
received (or close() is called), in-progress stream operations return
the meaningful session error instead of a generic LocallyClosed or
HTTP/3 teardown code.

Each stream type gets a map_error() method that replaces connection-level
errors and invalid HTTP/3 codes (InvalidReset/InvalidStopped) with the
stored session error. Session-level methods (accept, open, datagram) use
the existing Session::map_error(), which is tightened to only replace
connection-level errors.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

* Applied cargo fmt

* Address PR review comments

- Change map_error to accept impl Into<Error> to reduce boilerplate
- Only remap SessionError variants, not InvalidReset/InvalidStopped
- Remove unwrap() panic in received_reset for invalid reset codes

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

---------

Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com>
Co-authored-by: Luke Curley <kixelated@gmail.com>

Author: 3 people

web-transport-quinn/src/recv.rs

@@ -1,6 +1,7 @@
 use std::{
     io,
     pin::Pin,
+    sync::{Arc, OnceLock},
     task::{Context, Poll},
 };
 
@@ -12,11 +13,26 @@ use crate::{ReadError, ReadExactError, ReadToEndError, SessionError};
 #[derive(Debug)]
 pub struct RecvStream {
     inner: quinn::RecvStream,
+    error: Arc<OnceLock<SessionError>>,
 }
 
 impl RecvStream {
-    pub(crate) fn new(stream: quinn::RecvStream) -> Self {
-        Self { inner: stream }
+    pub(crate) fn new(stream: quinn::RecvStream, error: Arc<OnceLock<SessionError>>) -> Self {
+        Self {
+            inner: stream,
+            error,
+        }
+    }
+
+    /// Replace connection-level errors with the stored session error if available.
+    fn map_error(&self, e: impl Into<ReadError>) -> ReadError {
+        let e = e.into();
+        if let Some(err) = self.error.get() {
+            if matches!(&e, ReadError::SessionError(_)) {
+                return ReadError::SessionError(err.clone());
+            }
+        }
+        e
     }
 
     /// Tell the other end to stop sending data with the given error code. See [`quinn::RecvStream::stop`].
@@ -31,12 +47,15 @@ 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<Option<usize>, ReadError> {
-        self.inner.read(buf).await.map_err(Into::into)
+        self.inner.read(buf).await.map_err(|e| self.map_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| match e {
+            quinn::ReadExactError::ReadError(e) => self.map_error(e).into(),
+            e => e.into(),
+        })
     }
 
     /// Read a chunk of data from the stream. See [`quinn::RecvStream::read_chunk`].
@@ -48,17 +67,26 @@ impl RecvStream {
         self.inner
             .read_chunk(max_length, ordered)
             .await
-            .map_err(Into::into)
+            .map_err(|e| self.map_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<Option<usize>, ReadError> {
-        self.inner.read_chunks(bufs).await.map_err(Into::into)
+        self.inner
+            .read_chunks(bufs)
+            .await
+            .map_err(|e| self.map_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<Vec<u8>, ReadToEndError> {
-        self.inner.read_to_end(size_limit).await.map_err(Into::into)
+        self.inner
+            .read_to_end(size_limit)
+            .await
+            .map_err(|e| match e {
+                quinn::ReadToEndError::Read(e) => self.map_error(e).into(),
+                e => e.into(),
+            })
     }
 
     /// Block until the stream has been reset and return the error code. See [`quinn::RecvStream::received_reset`].
@@ -67,10 +95,10 @@ impl RecvStream {
     pub async fn received_reset(&mut self) -> Result<Option<u32>, SessionError> {
         match self.inner.received_reset().await {
             Ok(None) => Ok(None),
-            Ok(Some(code)) => Ok(Some(
-                web_transport_proto::error_from_http3(code.into_inner()).unwrap(),
-            )),
-            Err(quinn::ResetError::ConnectionLost(conn_err)) => Err(conn_err.into()),
+            Ok(Some(code)) => Ok(web_transport_proto::error_from_http3(code.into_inner())),
+            Err(quinn::ResetError::ConnectionLost(conn_err)) => {
+                Err(self.error.get().cloned().unwrap_or_else(|| conn_err.into()))
+            }
             Err(quinn::ResetError::ZeroRttRejected) => unreachable!("0-RTT not supported"),
         }
     }

web-transport-quinn/src/send.rs

@@ -1,6 +1,7 @@
 use std::{
     io,
     pin::Pin,
+    sync::{Arc, OnceLock},
     task::{Context, Poll},
 };
 
@@ -15,11 +16,23 @@ use crate::{ClosedStream, SessionError, WriteError};
 #[derive(Debug)]
 pub struct SendStream {
     stream: quinn::SendStream,
+    error: Arc<OnceLock<SessionError>>,
 }
 
 impl SendStream {
-    pub(crate) fn new(stream: quinn::SendStream) -> Self {
-        Self { stream }
+    pub(crate) fn new(stream: quinn::SendStream, error: Arc<OnceLock<SessionError>>) -> Self {
+        Self { stream, error }
+    }
+
+    /// Replace connection-level errors with the stored session error if available.
+    fn map_error(&self, e: impl Into<WriteError>) -> WriteError {
+        let e = e.into();
+        if let Some(err) = self.error.get() {
+            if matches!(&e, WriteError::SessionError(_)) {
+                return WriteError::SessionError(err.clone());
+            }
+        }
+        e
     }
 
     /// Abruptly reset the stream with the provided error code. See [`quinn::SendStream::reset`].
@@ -38,7 +51,9 @@ 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)) => Err(conn_err.into()),
+            Err(quinn::StoppedError::ConnectionLost(conn_err)) => {
+                Err(self.error.get().cloned().unwrap_or_else(|| conn_err.into()))
+            }
             Err(quinn::StoppedError::ZeroRttRejected) => unreachable!("0-RTT not supported"),
         }
     }
@@ -47,27 +62,39 @@ 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<usize, WriteError> {
-        self.stream.write(buf).await.map_err(Into::into)
+        self.stream.write(buf).await.map_err(|e| self.map_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_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<quinn::Written, WriteError> {
-        self.stream.write_chunks(bufs).await.map_err(Into::into)
+        self.stream
+            .write_chunks(bufs)
+            .await
+            .map_err(|e| self.map_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_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_error(e))
     }
 
     /// Mark the stream as finished, such that no more data can be written. See [`quinn::SendStream::finish`].

web-transport-quinn/src/session.rs

@@ -81,7 +81,7 @@ impl Session {
         let error: Arc<OnceLock<SessionError>> = Arc::new(OnceLock::new());
 
         // Accept logic is stateful, so use an Arc<Mutex> to share it.
-        let accept = SessionAccept::new(conn.clone(), session_id);
+        let accept = SessionAccept::new(conn.clone(), session_id, error.clone());
 
         let this = Self {
             conn,
@@ -186,50 +186,69 @@ impl Session {
     /// Accept a new unidirectional stream. See [`quinn::Connection::accept_uni`].
     pub async fn accept_uni(&self) -> Result<RecvStream, SessionError> {
         if let Some(accept) = &self.accept {
-            Ok(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_error(e))
         } else {
-            Ok(RecvStream::new(self.conn.accept_uni().await?))
+            let recv = self
+                .conn
+                .accept_uni()
+                .await
+                .map_err(|e| self.map_error(e))?;
+            Ok(RecvStream::new(recv, self.error.clone()))
         }
     }
 
     /// 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 {
-            Ok(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_error(e))
         } else {
-            let (send, recv) = self.conn.accept_bi().await?;
-            Ok((SendStream::new(send), RecvStream::new(recv)))
+            let (send, recv) = self.conn.accept_bi().await.map_err(|e| self.map_error(e))?;
+            Ok((
+                SendStream::new(send, self.error.clone()),
+                RecvStream::new(recv, self.error.clone()),
+            ))
         }
     }
 
     /// Open a new unidirectional stream. See [`quinn::Connection::open_uni`].
     pub async fn open_uni(&self) -> Result<SendStream, SessionError> {
-        let mut send = self.conn.open_uni().await?;
+        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.
         // 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_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.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?;
+        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.
         // 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_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.error.clone()),
+            RecvStream::new(recv, self.error.clone()),
+        ))
     }
 
     /// Asynchronously receives an application datagram from the remote peer.
@@ -238,7 +257,11 @@ 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<Bytes, SessionError> {
-        let mut datagram = self.conn.read_datagram().await?;
+        let mut datagram = self
+            .conn
+            .read_datagram()
+            .await
+            .map_err(|e| self.map_error(e))?;
 
         let mut cursor = Cursor::new(&datagram);
 
@@ -276,7 +299,7 @@ impl Session {
             self.conn.send_datagram(data)
         };
 
-        result?;
+        result.map_err(|e| self.map_error(e))?;
         Ok(())
     }
 
@@ -397,12 +420,19 @@ impl Session {
         self.conn.close_reason().map(|e| self.map_error(e))
     }
 
-    /// Convert an error to `SessionError`, using the stored session error if set.
+    /// Replace connection-level errors with the stored session error if available.
     fn map_error(&self, e: impl Into<SessionError>) -> SessionError {
+        let e = e.into();
         if let Some(err) = self.error.get() {
-            return err.clone();
+            if matches!(
+                &e,
+                SessionError::ConnectionError(_)
+                    | SessionError::SendDatagramError(quinn::SendDatagramError::ConnectionLost(_))
+            ) {
+                return err.clone();
+            }
         }
-        e.into()
+        e
     }
 
     async fn write_full(send: &mut quinn::SendStream, buf: &[u8]) -> Result<(), SessionError> {
@@ -488,6 +518,9 @@ type PendingBi = dyn Future<Output = Result<Option<(quinn::SendStream, quinn::Re
 pub struct SessionAccept {
     session_id: VarInt,
 
+    // Shared session error for propagation to accepted streams.
+    error: Arc<OnceLock<SessionError>>,
+
     // 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.
     qpack_encoder: Option<quinn::RecvStream>,
@@ -508,7 +541,11 @@ pub struct SessionAccept {
 }
 
 impl SessionAccept {
-    pub(crate) fn new(conn: quinn::Connection, session_id: VarInt) -> Self {
+    pub(crate) fn new(
+        conn: quinn::Connection,
+        session_id: VarInt,
+        error: Arc<OnceLock<SessionError>>,
+    ) -> 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))
@@ -520,6 +557,7 @@ impl SessionAccept {
 
         Self {
             session_id,
+            error,
 
             qpack_decoder: None,
             qpack_encoder: None,
@@ -580,7 +618,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.error.clone());
                     for waker in self.uni_wakers.drain(..) {
                         waker.wake();
                     }
@@ -666,8 +704,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.error.clone());
+                let recv = RecvStream::new(recv, self.error.clone());
                 for waker in self.bi_wakers.drain(..) {
                     waker.wake();
                 }