From c1b0b7a3c5a8752edce496dc32096eef3870af54 Mon Sep 17 00:00:00 2001
From: Evan Rittenhouse <erittenhouse@cloudflare.com>
Date: Thu, 29 Aug 2024 09:51:00 -0500
Subject: [PATCH] Add stream_bytes action to prompt and qlog

---
 h3i/src/lib.rs               | 17 -----------------
 h3i/src/prompts/h3/mod.rs    | 16 ++++++++++++++++
 h3i/src/recordreplay/qlog.rs | 32 ++++++++++++++++++++------------
 3 files changed, 36 insertions(+), 29 deletions(-)

diff --git a/h3i/src/lib.rs b/h3i/src/lib.rs
index 64c2acb52d..0f97aede7e 100644
--- a/h3i/src/lib.rs
+++ b/h3i/src/lib.rs
@@ -51,7 +51,6 @@ use qlog::events::EventData;
 pub use quiche;
 use quiche::h3::NameValue;
 
-use smallvec::smallvec;
 use smallvec::SmallVec;
 
 #[derive(Default)]
@@ -106,22 +105,6 @@ fn fake_packet_header() -> PacketHeader {
     }
 }
 
-fn fake_packet_with_stream_fin(stream_id: u64, fin: bool) -> Option<EventData> {
-    if !fin {
-        return None;
-    }
-
-    let frames = Some(smallvec![QuicFrame::Stream {
-        stream_id,
-        offset: 0,
-        length: 0,
-        fin: Some(fin),
-        raw: None
-    }]);
-
-    Some(fake_packet_sent(frames))
-}
-
 fn fake_packet_sent(frames: Option<SmallVec<[QuicFrame; 1]>>) -> EventData {
     EventData::PacketSent(PacketSent {
         header: fake_packet_header(),
diff --git a/h3i/src/prompts/h3/mod.rs b/h3i/src/prompts/h3/mod.rs
index 806f8fd027..f5389cc607 100644
--- a/h3i/src/prompts/h3/mod.rs
+++ b/h3i/src/prompts/h3/mod.rs
@@ -77,6 +77,7 @@ pub type PromptedFrame = (u64, quiche::h3::frame::Frame);
 
 thread_local! {static CONNECTION_IDLE_TIMEOUT: RefCell<u64> = const { RefCell::new(0) }}
 
+// TODO(erittenhouse): exploring generating prompts at compile-time
 const HEADERS: &str = "headers";
 const HEADERS_NO_PSEUDO: &str = "headers_no_pseudo";
 const DATA: &str = "data";
@@ -92,6 +93,7 @@ const OPEN_UNI_STREAM: &str = "open_uni_stream";
 const RESET_STREAM: &str = "reset_stream";
 const STOP_SENDING: &str = "stop_sending";
 const CONNECTION_CLOSE: &str = "connection_close";
+const STREAM_BYTES: &str = "stream_bytes";
 
 const COMMIT: &str = "commit";
 const FLUSH_PACKETS: &str = "flush_packets";
@@ -158,6 +160,7 @@ impl Prompter {
             PUSH_PROMISE => prompt_push_promise(),
             PRIORITY_UPDATE => priority::prompt_priority(),
             CONNECTION_CLOSE => prompt_connection_close(),
+            STREAM_BYTES => prompt_stream_bytes(),
             FLUSH_PACKETS => return PromptOutcome::Action(Action::FlushPackets),
             COMMIT => return PromptOutcome::Commit,
             WAIT => prompt_wait(),
@@ -254,6 +257,7 @@ fn action_suggester(val: &str) -> SuggestionResult<Vec<String>> {
         RESET_STREAM,
         STOP_SENDING,
         CONNECTION_CLOSE,
+        STREAM_BYTES,
         FLUSH_PACKETS,
         COMMIT,
         WAIT,
@@ -446,6 +450,18 @@ pub fn prompt_connection_close() -> InquireResult<Action> {
     })
 }
 
+pub fn prompt_stream_bytes() -> InquireResult<Action> {
+    let stream_id = h3::prompt_stream_id()?;
+    let bytes = Text::new("bytes:").prompt()?;
+    let fin_stream = prompt_fin_stream()?;
+
+    Ok(Action::StreamBytes {
+        stream_id,
+        fin_stream,
+        bytes: bytes.as_bytes().to_vec(),
+    })
+}
+
 fn validate_wait_period(period: &str) -> SuggestionResult<Validation> {
     let x = period.parse::<u64>();
 
diff --git a/h3i/src/recordreplay/qlog.rs b/h3i/src/recordreplay/qlog.rs
index 4ff7f7ac8b..ad3a60844b 100644
--- a/h3i/src/recordreplay/qlog.rs
+++ b/h3i/src/recordreplay/qlog.rs
@@ -37,6 +37,7 @@ use qlog::events::Event;
 use qlog::events::EventData;
 use qlog::events::ExData;
 use qlog::events::JsonEvent;
+use qlog::events::RawInfo;
 use quiche;
 use quiche::h3::frame::Frame;
 use quiche::h3::NameValue;
@@ -49,7 +50,6 @@ use crate::actions::h3::Action;
 use crate::actions::h3::WaitType;
 use crate::encode_header_block;
 use crate::fake_packet_sent;
-use crate::fake_packet_with_stream_fin;
 
 const HTTP3_CONTROL_STREAM_TYPE_ID: u64 = 0x0;
 const HTTP3_PUSH_STREAM_TYPE_ID: u64 = 0x1;
@@ -188,18 +188,26 @@ impl From<&Action> for QlogEvents {
             Action::StreamBytes {
                 stream_id,
                 fin_stream,
-                bytes: _,
+                bytes,
             } => {
-                if let Some(dummy) =
-                    fake_packet_with_stream_fin(*stream_id, *fin_stream)
-                {
-                    vec![QlogEvent::Event {
-                        data: Box::new(dummy),
-                        ex_data: BTreeMap::new(),
-                    }]
-                } else {
-                    vec![]
-                }
+                let len = bytes.len() as u64;
+                let ev = fake_packet_sent(Some(smallvec![QuicFrame::Stream {
+                    stream_id: *stream_id,
+                    fin: Some(*fin_stream),
+                    // ignore offset
+                    offset: 0,
+                    length: len,
+                    raw: Some(RawInfo {
+                        length: Some(len),
+                        payload_length: Some(len),
+                        data: String::from_utf8(bytes.clone()).ok()
+                    })
+                }]));
+
+                vec![QlogEvent::Event {
+                    data: Box::new(ev),
+                    ex_data: BTreeMap::new(),
+                }]
             },
 
             Action::ResetStream {