h3i: implement expected frames

Expected frames allow the user to specify a list of frames that the h3i
client expects to receive over a given connection.

If h3i sees all of the exected frames over the course of the connection,
it will pre-emptively close the connection with a CONNECTION_CLOSE
frame. If h3i does _not_ see all of the expected frames, the resulting
ConnectionSummary will contain a list of the missing target frames for
future inspection.

This gives users a way to close tests out without waiting for the idle
timeout, or adding Wait/ConnectionClose actions to the end of each test.
This should vastly speed up test suites that have a large number of h3i
tests.

h3i/examples/content_length_mismatch.rs

@@ -67,7 +67,7 @@ fn main() {
     ];
 
     let summary =
-        sync_client::connect(config, &actions).expect("connection failed");
+        sync_client::connect(config, &actions, None).expect("connection failed");
 
     println!(
         "=== received connection summary! ===\n\n{}",

h3i/src/client/connection_summary.rs

@@ -39,6 +39,7 @@ use std::collections::HashMap;
 use std::iter::FromIterator;
 
 use crate::frame::EnrichedHeaders;
+use crate::frame::ExpectedFrame;
 use crate::frame::H3iFrame;
 
 /// Maximum length of any serialized element's unstructured data such as reason
@@ -74,22 +75,32 @@ impl Serialize for ConnectionSummary {
             self.path_stats.iter().map(SerializablePathStats).collect();
         state.serialize_field("path_stats", &p)?;
         state.serialize_field("error", &self.conn_close_details)?;
+        state.serialize_field(
+            "missed_expected_frames",
+            &self.stream_map.missing_frames(),
+        )?;
         state.end()
     }
 }
 
 /// A read-only aggregation of frames received over a connection, mapped to the
 /// stream ID over which they were received.
 #[derive(Clone, Debug, Default, Serialize)]
-pub struct StreamMap(HashMap<u64, Vec<H3iFrame>>);
+pub struct StreamMap {
+    map: HashMap<u64, Vec<H3iFrame>>,
+    expected_frames: Option<ExpectedFrames>,
+}
 
 impl<T> From<T> for StreamMap
 where
     T: IntoIterator<Item = (u64, Vec<H3iFrame>)>,
 {
     fn from(value: T) -> Self {
         let map = HashMap::from_iter(value);
-        Self(map)
+        Self {
+            map,
+            expected_frames: None,
+        }
     }
 }
 
@@ -113,7 +124,7 @@ impl StreamMap {
     /// assert_eq!(stream_map.all_frames(), vec![headers]);
     /// ```
     pub fn all_frames(&self) -> Vec<H3iFrame> {
-        self.0
+        self.map
             .values()
             .flatten()
             .map(Clone::clone)
@@ -140,7 +151,7 @@ impl StreamMap {
     /// assert_eq!(stream_map.stream(0), vec![headers]);
     /// ```
     pub fn stream(&self, stream_id: u64) -> Vec<H3iFrame> {
-        self.0.get(&stream_id).cloned().unwrap_or_default()
+        self.map.get(&stream_id).cloned().unwrap_or_default()
     }
 
     /// Check if a provided [`H3iFrame`] was received, regardless of what stream
@@ -189,7 +200,7 @@ impl StreamMap {
     pub fn received_frame_on_stream(
         &self, stream: u64, frame: &H3iFrame,
     ) -> bool {
-        self.0.get(&stream).map(|v| v.contains(frame)).is_some()
+        self.map.get(&stream).map(|v| v.contains(frame)).is_some()
     }
 
     /// Check if the stream map is empty, e.g., no frames were received.
@@ -213,7 +224,7 @@ impl StreamMap {
     /// assert!(!stream_map.is_empty());
     /// ```
     pub fn is_empty(&self) -> bool {
-        self.0.is_empty()
+        self.map.is_empty()
     }
 
     /// See all HEADERS received on a given stream.
@@ -246,8 +257,105 @@ impl StreamMap {
             .collect()
     }
 
+    pub(crate) fn new(expected_frames: Option<ExpectedFrames>) -> Self {
+        Self {
+            expected_frames,
+            ..Default::default()
+        }
+    }
+
     pub(crate) fn insert(&mut self, stream_id: u64, frame: H3iFrame) {
-        self.0.entry(stream_id).or_default().push(frame);
+        if let Some(expected) = self.expected_frames.as_mut() {
+            expected.receive_frame(stream_id, &frame);
+        }
+
+        self.map.entry(stream_id).or_default().push(frame);
+    }
+
+    pub(crate) fn check_expected_frames(&self) -> bool {
+        self.expected_frames
+            .as_ref()
+            .is_some_and(|e| e.saw_all_frames())
+    }
+
+    pub(crate) fn missing_frames(&self) -> Option<Vec<ExpectedFrame>> {
+        self.expected_frames.as_ref().map(|e| e.missing_frames())
+    }
+
+    /// Close a [`quiche::Connection`] with the CONNECTION_CLOSE frame specified by
+    /// [`ExpectedFrames`]. If no [`ExpectedFrames`] exist, this is a no-op.
+    pub(crate) fn close_due_to_expected_frames(
+        &self, qconn: &mut quiche::Connection,
+    ) {
+        let Some(ConnectionError {
+            is_app,
+            error_code,
+            reason,
+        }) = self.expected_frames.as_ref().map(|ef| &ef.close_with)
+        else {
+            return;
+        };
+
+        let _ = qconn.close(*is_app, *error_code, reason);
+    }
+}
+
+/// A container for frames that h3i expects to see over a given connection. If h3i receives all the
+/// frames it expects, it will send a CONNECTION_CLOSE frame to the server. This bypasses the idle
+/// timeout and vastly quickens test suites which depend heavily on h3i.
+///
+/// The specific CONNECTION_CLOSE frame can be customized by passing a [`ConnectionError`] to
+/// [`Self::new_with_close`]. h3i will send an application CONNECTION_CLOSE frame with error code
+/// 0x100 if this struct is constructed with the [`Self::new`] constructor.
+#[derive(Clone, Serialize, Debug)]
+pub struct ExpectedFrames {
+    missing: Vec<ExpectedFrame>,
+    #[serde(skip)]
+    close_with: ConnectionError,
+}
+
+impl ExpectedFrames {
+    pub fn new(frames: Vec<ExpectedFrame>) -> Self {
+        Self::new_with_close(
+            frames,
+            ConnectionError {
+                is_app: true,
+                error_code: quiche::h3::WireErrorCode::NoError as u64,
+                reason: b"saw all expected frames".to_vec(),
+            },
+        )
+    }
+
+    pub fn new_with_close(
+        frames: Vec<ExpectedFrame>, close_with: ConnectionError,
+    ) -> Self {
+        Self {
+            missing: frames,
+            close_with,
+        }
+    }
+
+    fn receive_frame(&mut self, stream_id: u64, frame: &H3iFrame) {
+        for (i, ef) in self.missing.iter_mut().enumerate() {
+            if ef.is_equivalent(frame) && ef.stream_id() == stream_id {
+                self.missing.remove(i);
+                break;
+            }
+        }
+    }
+
+    fn saw_all_frames(&self) -> bool {
+        self.missing.is_empty()
+    }
+
+    fn missing_frames(&self) -> Vec<ExpectedFrame> {
+        self.missing.clone()
+    }
+}
+
+impl From<Vec<ExpectedFrame>> for ExpectedFrames {
+    fn from(value: Vec<ExpectedFrame>) -> Self {
+        Self::new(value)
     }
 }
 
@@ -404,6 +512,7 @@ impl Serialize for SerializableStats<'_> {
 }
 
 /// A wrapper to help serialize a [quiche::ConnectionError]
+#[derive(Clone, Debug)]
 pub struct SerializableConnectionError<'a>(&'a quiche::ConnectionError);
 
 impl Serialize for SerializableConnectionError<'_> {
@@ -422,3 +531,77 @@ impl Serialize for SerializableConnectionError<'_> {
         state.end()
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use quiche::h3::Header;
+
+    fn h3i_frame() -> H3iFrame {
+        vec![Header::new(b"hello", b"world")].into()
+    }
+
+    #[test]
+    fn expected_frame() {
+        let frame = h3i_frame();
+        let mut expected =
+            ExpectedFrames::new(vec![ExpectedFrame::new(0, frame.clone())]);
+
+        expected.receive_frame(0, &frame);
+
+        assert!(expected.saw_all_frames());
+    }
+
+    #[test]
+    fn expected_frame_missing() {
+        let frame = h3i_frame();
+        let expected_frames = vec![
+            ExpectedFrame::new(0, frame.clone()),
+            ExpectedFrame::new(4, frame.clone()),
+            ExpectedFrame::new(8, vec![Header::new(b"go", b"jets")].into()),
+        ];
+        let mut expected = ExpectedFrames::new(expected_frames.clone());
+
+        expected.receive_frame(0, &frame);
+
+        assert!(!expected.saw_all_frames());
+        assert_eq!(expected.missing_frames(), expected_frames[1..].to_vec());
+    }
+
+    fn stream_map_data() -> Vec<H3iFrame> {
+        let headers =
+            H3iFrame::Headers(EnrichedHeaders::from(vec![Header::new(
+                b"hello", b"world",
+            )]));
+        let data = H3iFrame::QuicheH3(quiche::h3::frame::Frame::Data {
+            payload: b"hello world".to_vec(),
+        });
+
+        vec![headers, data]
+    }
+
+    #[test]
+    fn test_stream_map_expected_frames_with_none() {
+        let stream_map: StreamMap = vec![(0, stream_map_data())].into();
+        assert!(!stream_map.check_expected_frames());
+    }
+
+    #[test]
+    fn test_stream_map_expected_frames() {
+        let data = stream_map_data();
+        let mut stream_map = StreamMap::new(Some(
+            vec![
+                ExpectedFrame::new(0, data[0].clone()),
+                ExpectedFrame::new(0, data[1].clone()),
+            ]
+            .into(),
+        ));
+
+        stream_map.insert(0, data[0].clone());
+        assert!(!stream_map.check_expected_frames());
+        assert_eq!(
+            stream_map.missing_frames().unwrap(),
+            vec![ExpectedFrame::new(0, data[1].clone())]
+        );
+    }
+}

h3i/src/client/sync_client.rs

@@ -48,6 +48,7 @@ use crate::config::Config;
 
 use super::Client;
 use super::ConnectionSummary;
+use super::ExpectedFrames;
 use super::StreamMap;
 use super::StreamParserMap;
 
@@ -57,6 +58,15 @@ struct SyncClient {
     stream_parsers: StreamParserMap,
 }
 
+impl SyncClient {
+    fn new(expected_frames: Option<ExpectedFrames>) -> Self {
+        Self {
+            streams: StreamMap::new(expected_frames),
+            ..Default::default()
+        }
+    }
+}
+
 impl Client for SyncClient {
     fn stream_parsers_mut(&mut self) -> &mut StreamParserMap {
         &mut self.stream_parsers
@@ -74,7 +84,7 @@ impl Client for SyncClient {
 ///
 /// Returns a [ConnectionSummary] on success, [ClientError] on failure.
 pub fn connect(
-    args: Config, actions: &[Action],
+    args: Config, actions: &[Action], expected_frames: Option<ExpectedFrames>,
 ) -> std::result::Result<ConnectionSummary, ClientError> {
     let mut buf = [0; 65535];
     let mut out = [0; MAX_DATAGRAM_SIZE];
@@ -142,8 +152,7 @@ pub fn connect(
     let mut wait_duration = None;
     let mut wait_instant = None;
 
-    let mut client = SyncClient::default();
-
+    let mut client = SyncClient::new(expected_frames);
     let mut waiting_for = WaitingFor::default();
 
     loop {
@@ -248,8 +257,8 @@ pub fn connect(
 
         // Create a new application protocol session once the QUIC connection is
         // established.
-        if (conn.is_established() || conn.is_in_early_data()) &&
-            !app_proto_selected
+        if (conn.is_established() || conn.is_in_early_data())
+            && !app_proto_selected
         {
             app_proto_selected = true;
         }
@@ -277,6 +286,10 @@ pub fn connect(
                 wait_cleared = true;
             }
 
+            if client.streams.check_expected_frames() {
+                client.streams.close_due_to_expected_frames(&mut conn);
+            }
+
             if wait_cleared {
                 check_duration_and_do_actions(
                     &mut wait_duration,

h3i/src/config.rs

@@ -28,6 +28,7 @@
 use std::io;
 
 /// Server details and QUIC connection properties.
+#[derive(Clone)]
 pub struct Config {
     /// A string representing the host and port to connect to using the format
     /// `<host>:<port>`.