Fix close capsules

web-transport-proto/src/capsule.rs

@@ -3,12 +3,9 @@ use std::sync::Arc;
 use bytes::{Buf, BufMut, Bytes, BytesMut};
 use tokio::io::{AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt};
 
-use crate::{VarInt, VarIntUnexpectedEnd, MAX_FRAME_SIZE};
+use crate::{Frame, VarInt, VarIntUnexpectedEnd, MAX_FRAME_SIZE};
 
-// The spec (draft-ietf-webtrans-http3-06) says the type is 0x2843, which would
-// varint-encode to 0x68 0x43. However, actual wire data shows 0x43 0x28 which
-// decodes to 808. There may be a discrepancy in implementations or specs.
-// Using 0x2843 as specified in the standard.
+// CloseWebTransportSession capsule type (draft-ietf-webtrans-http3-06).
 const CLOSE_WEBTRANSPORT_SESSION_TYPE: u64 = 0x2843;
 
 #[derive(Debug, Clone, PartialEq, Eq)]
@@ -81,9 +78,10 @@ impl Capsule {
     ///
     /// Returns `Ok(None)` if the stream is cleanly closed (EOF before any bytes).
     pub async fn read<S: AsyncRead + Unpin>(stream: &mut S) -> Result<Option<Self>, CapsuleError> {
-        let typ = match VarInt::read(stream).await {
-            Ok(v) => v,
-            Err(_) => return Ok(None), // Clean EOF
+        let typ = match VarInt::read_optional(stream).await {
+            Ok(Some(v)) => v,
+            Ok(None) => return Ok(None), // Clean EOF
+            Err(_) => return Err(CapsuleError::UnexpectedEnd),
         };
         let length = VarInt::read(stream)
             .await
@@ -216,9 +214,6 @@ pub enum CapsuleError {
     #[error("message too long")]
     MessageTooLong,
 
-    #[error("unknown capsule type: {0:?}")]
-    UnknownType(VarInt),
-
     #[error("varint decode error: {0:?}")]
     VarInt(#[from] VarIntUnexpectedEnd),
 
@@ -232,6 +227,86 @@ impl From<std::io::Error> 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<S> {
+    stream: S,
+    buf: BytesMut,
+}
+
+impl<S: AsyncRead + Unpin> Http3CapsuleReader<S> {
+    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<Option<Capsule>, 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<bool, CapsuleError> {
+        loop {
+            let frame_type = match VarInt::read_optional(&mut self.stream).await {
+                Ok(Some(v)) => Frame(v),
+                Ok(None) => return Ok(false),
+                Err(_) => return Err(CapsuleError::UnexpectedEnd),
+            };
+            let len = VarInt::read(&mut self.stream)
+                .await
+                .map_err(|_| CapsuleError::UnexpectedEnd)?
+                .into_inner() as usize;
+
+            if len > MAX_FRAME_SIZE as usize {
+                return Err(CapsuleError::MessageTooLong);
+            }
+
+            if frame_type != Frame::DATA {
+                let mut limited = (&mut self.stream).take(len as u64);
+                let n = tokio::io::copy(&mut limited, &mut tokio::io::sink()).await?;
+                if (n as usize) < len {
+                    return Err(CapsuleError::UnexpectedEnd);
+                }
+                continue;
+            }
+
+            if len > 0 {
+                let start = self.buf.len();
+                self.buf.resize(start + len, 0);
+                self.stream.read_exact(&mut self.buf[start..]).await?;
+            }
+            return Ok(true);
+        }
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;
@@ -493,4 +568,130 @@ mod tests {
         let err = Capsule::read(&mut cursor).await.unwrap_err();
         assert!(matches!(err, CapsuleError::UnexpectedEnd));
     }
+
+    fn encode_capsule(c: &Capsule) -> Vec<u8> {
+        let mut buf = Vec::new();
+        c.encode(&mut buf);
+        buf
+    }
+
+    fn wrap_in_data_frame(payload: &[u8]) -> Vec<u8> {
+        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<u8> {
+        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<u8>) -> Http3CapsuleReader<std::io::Cursor<Vec<u8>>> {
+        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);
+    }
 }

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<S: AsyncRead + Unpin>(stream: &mut S) -> Result<Self, VarIntUnexpectedEnd> {
+        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<S: AsyncRead + Unpin>(
+        stream: &mut S,
+    ) -> Result<Option<Self>, 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<B: BufMut>(&self, w: &mut 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"

web-transport-quinn/src/recv.rs

@@ -70,7 +70,7 @@ impl RecvStream {
             Ok(Some(code)) => Ok(Some(
                 web_transport_proto::error_from_http3(code.into_inner()).unwrap(),
             )),
-            Err(quinn::ResetError::ConnectionLost(e)) => Err(e.into()),
+            Err(quinn::ResetError::ConnectionLost(conn_err)) => Err(conn_err.into()),
             Err(quinn::ResetError::ZeroRttRejected) => unreachable!("0-RTT not supported"),
         }
     }

web-transport-quinn/src/send.rs

@@ -38,7 +38,7 @@ impl SendStream {
         match self.stream.stopped().await {
             Ok(Some(code)) => Ok(web_transport_proto::error_from_http3(code.into_inner())),
             Ok(None) => Ok(None),
-            Err(quinn::StoppedError::ConnectionLost(e)) => Err(e.into()),
+            Err(quinn::StoppedError::ConnectionLost(conn_err)) => Err(conn_err.into()),
             Err(quinn::StoppedError::ZeroRttRejected) => unreachable!("0-RTT not supported"),
         }
     }