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

@@ -66,8 +66,12 @@ fn main() {
         },
     ];
 
-    let summary =
-        sync_client::connect(config, &actions).expect("connection failed");
+    // This example doesn't use expected frames, since we manually close the
+    // connection upon receiving a HEADERS frame on stream 0.
+    let expected_frames = None;
+
+    let summary = sync_client::connect(config, &actions, expected_frames)
+        .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,36 @@ 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_expected_frames(),
+        )?;
         state.end()
     }
 }
 
 /// A read-only aggregation of frames received over a connection, mapped to the
 /// stream ID over which they were received.
+///
+/// [`StreamMap`] also contains the [`ExpectedFrames`] for the connection so
+/// that its state can be updated as new frames are received.
 #[derive(Clone, Debug, Default, Serialize)]
-pub struct StreamMap(HashMap<u64, Vec<H3iFrame>>);
+pub struct StreamMap {
+    stream_frame_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)
+        let stream_frame_map = HashMap::from_iter(value);
+
+        Self {
+            stream_frame_map,
+            expected_frames: None,
+        }
     }
 }
 
@@ -113,7 +128,7 @@ impl StreamMap {
     /// assert_eq!(stream_map.all_frames(), vec![headers]);
     /// ```
     pub fn all_frames(&self) -> Vec<H3iFrame> {
-        self.0
+        self.stream_frame_map
             .values()
             .flatten()
             .map(Clone::clone)
@@ -140,7 +155,10 @@ 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.stream_frame_map
+            .get(&stream_id)
+            .cloned()
+            .unwrap_or_default()
     }
 
     /// Check if a provided [`H3iFrame`] was received, regardless of what stream
@@ -155,8 +173,6 @@ impl StreamMap {
     /// use quiche::h3::Header;
     /// use std::iter::FromIterator;
     ///
-    /// let mut stream_map = StreamMap::default();
-    ///
     /// let h = Header::new(b"hello", b"world");
     /// let headers = H3iFrame::Headers(EnrichedHeaders::from(vec![h]));
     ///
@@ -178,8 +194,6 @@ impl StreamMap {
     /// use quiche::h3::Header;
     /// use std::iter::FromIterator;
     ///
-    /// let mut stream_map = StreamMap::default();
-    ///
     /// let h = Header::new(b"hello", b"world");
     /// let headers = H3iFrame::Headers(EnrichedHeaders::from(vec![h]));
     ///
@@ -189,7 +203,10 @@ 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.stream_frame_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 +230,7 @@ impl StreamMap {
     /// assert!(!stream_map.is_empty());
     /// ```
     pub fn is_empty(&self) -> bool {
-        self.0.is_empty()
+        self.stream_frame_map.is_empty()
     }
 
     /// See all HEADERS received on a given stream.
@@ -227,8 +244,6 @@ impl StreamMap {
     /// use quiche::h3::Header;
     /// use std::iter::FromIterator;
     ///
-    /// let mut stream_map = StreamMap::default();
-    ///
     /// let h = Header::new(b"hello", b"world");
     /// let enriched = EnrichedHeaders::from(vec![h]);
     /// let headers = H3iFrame::Headers(enriched.clone());
@@ -246,8 +261,119 @@ impl StreamMap {
             .collect()
     }
 
+    /// If all [`ExpectedFrame`]s were seen. If no frames were expected, this
+    /// will return `false`.
+    pub fn all_expected_frames_seen(&self) -> bool {
+        if let Some(expected) = self.expected_frames.as_ref() {
+            expected.saw_all_frames()
+        } else {
+            false
+        }
+    }
+
+    /// The set of all [`ExpectedFrame`]s that were _not_ seen on the
+    /// connection. Returns `None` if
+    pub fn missing_expected_frames(&self) -> Option<Vec<ExpectedFrame>> {
+        self.expected_frames.as_ref().map(|e| e.missing_frames())
+    }
+
+    ///  Not `pub` as users aren't expected to build their own [`StreamMap`]s.
+    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.stream_frame_map
+            .entry(stream_id)
+            .or_default()
+            .push(frame);
+    }
+
+    /// 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,
+    ) {
+        if let Some(ConnectionError {
+            is_app,
+            error_code,
+            reason,
+        }) = self.expected_frames.as_ref().map(|ef| &ef.close_with)
+        {
+            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 {
+    /// Create a new [`ExpectedFrames`]. If all expected frames are received,
+    /// h3i will close the connection with an application-level
+    /// CONNECTION_CLOSE frame with error code 0x100.
+    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(),
+        })
+    }
+
+    /// Create a new [`ExpectedFrames`] with a custom close frame. When all
+    /// expected frames are received, h3i will close the connection with the
+    /// level, error code, and reason from `close_with`.
+    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 +530,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 +549,77 @@ impl Serialize for SerializableConnectionError<'_> {
         state.end()
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::frame::EnrichedHeaders;
+    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")]),
+        ];
+        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.all_expected_frames_seen());
+    }
+
+    #[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.all_expected_frames_seen());
+        assert_eq!(stream_map.missing_expected_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
@@ -72,9 +82,13 @@ impl Client for SyncClient {
 /// Constructs a socket and [quiche::Connection] based on the provided `args`,
 /// then iterates over `actions`.
 ///
+/// If `expected_frames` are specified, h3i will close the connection
+/// immediately upon receiving all of the supplied frames rather than waiting
+/// for the idle timeout. See [`ExpectedFrames`] for details.
+///
 /// 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 +156,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 {
@@ -277,6 +290,10 @@ pub fn connect(
                 wait_cleared = true;
             }
 
+            if client.streams.all_expected_frames_seen() {
+                client.streams.close_due_to_expected_frames(&mut conn);
+            }
+
             if wait_cleared {
                 check_duration_and_do_actions(
                     &mut wait_duration,