fix(adapters): incremental byte-level SSE parser (FWD-06)
Replace per-chunk parse_sse_frames with SseParser holding raw bytes across chunks: reassembles frames split at TCP boundaries (the review's silently-losing case), decodes UTF-8 per complete line so multi-byte chars split across chunks survive, caps the buffer at 1 MiB (SSE_EVENT_BUFFER_CAP) and dispatches a pending event at EOF. forward_stream threads the parser through its unfold state and emits a terminal error envelope on cap overflow. Existing single-chunk SSE test assertions preserved; added multi-chunk, split-UTF-8, EOF-dispatch, and cap tests. Verified: cargo test (219 pass), clippy -D warnings, fmt --check.
This commit is contained in:
+170
-58
@@ -263,6 +263,17 @@ pub(crate) async fn forward(
|
||||
}
|
||||
}
|
||||
|
||||
/// Converts a parsed SSE event into a response envelope, JSON-decoding
|
||||
/// the data payload when possible.
|
||||
fn sse_event_envelope(event: SseEvent, request_id: &str) -> ResponseEnvelope {
|
||||
let parsed = if event.data.trim().is_empty() {
|
||||
Value::Null
|
||||
} else {
|
||||
serde_json::from_str(&event.data).unwrap_or(Value::String(event.data.clone()))
|
||||
};
|
||||
ResponseEnvelope::ok(request_id, parsed)
|
||||
}
|
||||
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
pub(crate) fn forward_stream(
|
||||
http_client: &Arc<SharedHttpClient>,
|
||||
@@ -348,39 +359,53 @@ pub(crate) fn forward_stream(
|
||||
let request_id_inner = request_id.clone();
|
||||
Box::pin(
|
||||
stream::unfold(
|
||||
(response.bytes_stream(), String::new()),
|
||||
move |(mut bytes, mut buffer)| {
|
||||
(response.bytes_stream(), SseParser::new(), false),
|
||||
move |(mut bytes, mut parser, broken)| {
|
||||
let request_id = request_id_inner.clone();
|
||||
async move {
|
||||
if broken {
|
||||
return None;
|
||||
}
|
||||
match bytes.next().await {
|
||||
Some(Ok(chunk)) => {
|
||||
buffer.push_str(&String::from_utf8_lossy(&chunk));
|
||||
let (events, remaining) = parse_sse_frames(&buffer);
|
||||
let envelopes: Vec<ResponseEnvelope> = events
|
||||
.into_iter()
|
||||
.map(|e| {
|
||||
let parsed = if e.data.trim().is_empty() {
|
||||
Value::Null
|
||||
} else {
|
||||
serde_json::from_str(&e.data).unwrap_or(
|
||||
Value::String(e.data.clone()),
|
||||
)
|
||||
};
|
||||
ResponseEnvelope::ok(&request_id, parsed)
|
||||
})
|
||||
.collect();
|
||||
Some((envelopes, (bytes, remaining)))
|
||||
}
|
||||
Some(Ok(chunk)) => match parser.feed(&chunk, false) {
|
||||
Ok(events) => {
|
||||
let envelopes: Vec<ResponseEnvelope> = events
|
||||
.into_iter()
|
||||
.map(|e| sse_event_envelope(e, &request_id))
|
||||
.collect();
|
||||
Some((envelopes, (bytes, parser, false)))
|
||||
}
|
||||
Err(err) => {
|
||||
let error = CallError::internal(format!(
|
||||
"SSE parse error: {err}"
|
||||
));
|
||||
Some((
|
||||
vec![ResponseEnvelope::error(
|
||||
request_id, error,
|
||||
)],
|
||||
(bytes, parser, true),
|
||||
))
|
||||
}
|
||||
},
|
||||
Some(Err(err)) => {
|
||||
let error = CallError::internal(format!(
|
||||
"SSE stream error: {err}"
|
||||
));
|
||||
Some((
|
||||
vec![ResponseEnvelope::error(request_id, error)],
|
||||
(bytes, buffer),
|
||||
(bytes, parser, true),
|
||||
))
|
||||
}
|
||||
None => None,
|
||||
None => match parser.feed(&[], true) {
|
||||
Ok(events) if !events.is_empty() => {
|
||||
let envelopes: Vec<ResponseEnvelope> = events
|
||||
.into_iter()
|
||||
.map(|e| sse_event_envelope(e, &request_id))
|
||||
.collect();
|
||||
Some((envelopes, (bytes, parser, true)))
|
||||
}
|
||||
_ => None,
|
||||
},
|
||||
}
|
||||
}
|
||||
},
|
||||
@@ -395,48 +420,135 @@ pub(crate) fn forward_stream(
|
||||
Box::pin(sse)
|
||||
}
|
||||
|
||||
/// A parsed SSE event: the `data:` lines joined with `\n`.
|
||||
pub(crate) struct SseEvent {
|
||||
pub(crate) data: String,
|
||||
}
|
||||
|
||||
pub(crate) fn parse_sse_frames(buffer: &str) -> (Vec<SseEvent>, String) {
|
||||
let mut events = Vec::new();
|
||||
let text = if let Some(stripped) = buffer.strip_prefix('\u{feff}') {
|
||||
stripped
|
||||
} else {
|
||||
buffer
|
||||
};
|
||||
let lines: Vec<&str> = text.split('\n').collect();
|
||||
let mut data_buffer: Vec<String> = Vec::new();
|
||||
let mut remaining = String::new();
|
||||
#[derive(Debug, thiserror::Error)]
|
||||
pub(crate) enum SseParseError {
|
||||
#[error("SSE event buffer exceeded {SSE_EVENT_BUFFER_CAP} bytes without a complete event")]
|
||||
BufferOverflow,
|
||||
}
|
||||
|
||||
for (i, line) in lines.iter().enumerate() {
|
||||
if i == lines.len() - 1 {
|
||||
remaining = line.to_string();
|
||||
break;
|
||||
}
|
||||
let line = line.strip_suffix('\r').unwrap_or(line);
|
||||
if line.is_empty() {
|
||||
if !data_buffer.is_empty() {
|
||||
events.push(SseEvent {
|
||||
data: data_buffer.join("\n"),
|
||||
});
|
||||
}
|
||||
data_buffer.clear();
|
||||
continue;
|
||||
}
|
||||
if line.starts_with(':') {
|
||||
continue;
|
||||
}
|
||||
if let Some((field, value)) = line.split_once(':') {
|
||||
let value = value.strip_prefix(' ').unwrap_or(value);
|
||||
if field == "data" {
|
||||
data_buffer.push(value.to_string());
|
||||
}
|
||||
} else if line == "data" {
|
||||
data_buffer.push(String::new());
|
||||
/// Maximum size, in bytes, of the SSE parser's internal reassembly
|
||||
/// buffer. A single event (all `data:` lines plus framing) must fit
|
||||
/// within this budget; a stream emitting a longer partial line — or an
|
||||
/// unterminated event — trips `SseParseError::BufferOverflow` instead
|
||||
/// of buffering without bound.
|
||||
pub(crate) const SSE_EVENT_BUFFER_CAP: usize = 1024 * 1024;
|
||||
|
||||
/// Incremental byte-level SSE frame parser.
|
||||
///
|
||||
/// Holds raw bytes across `feed` calls so a frame split across TCP
|
||||
/// chunks is reassembled, and only decodes UTF-8 once a complete line
|
||||
/// (or EOF) bounds the decode window — a multi-byte character split at
|
||||
/// a chunk boundary is therefore not corrupted. Framing follows the
|
||||
/// WHATWG `text/event-stream` draft semantics for what the call
|
||||
/// protocol needs: lines split on `\n` with optional trailing `\r`;
|
||||
/// `data:` fields accumulate and join with `\n`; `event:`, `id:`, and
|
||||
/// `retry:` fields are accepted and ignored; a blank line dispatches
|
||||
/// the pending event; a pending event with data is dispatched at EOF.
|
||||
pub(crate) struct SseParser {
|
||||
buf: Vec<u8>,
|
||||
data_lines: Vec<String>,
|
||||
data_seen: bool,
|
||||
bom_stripped: bool,
|
||||
}
|
||||
|
||||
impl SseParser {
|
||||
pub(crate) fn new() -> Self {
|
||||
Self {
|
||||
buf: Vec::new(),
|
||||
data_lines: Vec::new(),
|
||||
data_seen: false,
|
||||
bom_stripped: false,
|
||||
}
|
||||
}
|
||||
|
||||
(events, remaining)
|
||||
/// Feeds one chunk and drains every complete event (a blank line
|
||||
/// dispatches; the last line stays buffered unless `eof`). With
|
||||
/// `eof`, also dispatches a pending event if it carries data
|
||||
/// lines, and flushes the buffer.
|
||||
pub(crate) fn feed(&mut self, chunk: &[u8], eof: bool) -> Result<Vec<SseEvent>, SseParseError> {
|
||||
self.buf.extend_from_slice(chunk);
|
||||
if self.buf.len() > SSE_EVENT_BUFFER_CAP {
|
||||
return Err(SseParseError::BufferOverflow);
|
||||
}
|
||||
let mut events = Vec::new();
|
||||
let mut start = 0usize;
|
||||
while let Some(nl) = self.buf[start..].iter().position(|&b| b == b'\n') {
|
||||
let end = start + nl;
|
||||
let line_end = if end > start && self.buf[end - 1] == b'\r' {
|
||||
end - 1
|
||||
} else {
|
||||
end
|
||||
};
|
||||
let line = self.buf[start..line_end].to_vec();
|
||||
if let Some(event) = self.parse_line(&line) {
|
||||
events.push(event);
|
||||
}
|
||||
start = end + 1;
|
||||
}
|
||||
if eof {
|
||||
if start < self.buf.len() {
|
||||
let line = self.buf[start..].to_vec();
|
||||
if let Some(event) = self.parse_line(&line) {
|
||||
events.push(event);
|
||||
}
|
||||
}
|
||||
if self.data_seen {
|
||||
if let Some(event) = self.complete_event() {
|
||||
events.push(event);
|
||||
}
|
||||
}
|
||||
self.buf.clear();
|
||||
} else {
|
||||
self.buf.drain(..start);
|
||||
}
|
||||
Ok(events)
|
||||
}
|
||||
|
||||
fn parse_line(&mut self, line: &[u8]) -> Option<SseEvent> {
|
||||
if !self.bom_stripped {
|
||||
self.bom_stripped = true;
|
||||
let bom = b"\xef\xbb\xbf";
|
||||
let line = if line.starts_with(bom) {
|
||||
&line[bom.len()..]
|
||||
} else {
|
||||
line
|
||||
};
|
||||
return self.parse_line(line);
|
||||
}
|
||||
let text = match std::str::from_utf8(line) {
|
||||
Ok(t) => t,
|
||||
Err(_) => return None,
|
||||
};
|
||||
if text.is_empty() {
|
||||
return if self.data_seen {
|
||||
self.complete_event()
|
||||
} else {
|
||||
None
|
||||
};
|
||||
}
|
||||
if text.starts_with(':') {
|
||||
return None;
|
||||
}
|
||||
let (field, value) = match text.split_once(':') {
|
||||
Some((f, v)) => (f, v.strip_prefix(' ').unwrap_or(v)),
|
||||
None => (text, ""),
|
||||
};
|
||||
if field == "data" {
|
||||
self.data_lines.push(value.to_string());
|
||||
self.data_seen = true;
|
||||
}
|
||||
None
|
||||
}
|
||||
|
||||
fn complete_event(&mut self) -> Option<SseEvent> {
|
||||
self.data_seen = false;
|
||||
Some(SseEvent {
|
||||
data: std::mem::take(&mut self.data_lines).join("\n"),
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user