wiremux 0.11.7

//! Incremental SSE frame reader. Feed TCP chunks; emit complete frames.

use super::RawSse;

/// Hostile or buggy stream with no newline must not grow unbounded.
/// Joined `data:` payload for one frame uses the same cap.
pub const MAX_SSE_PENDING: usize = 16 * 1024 * 1024;

/// Incremental parser for `event:` / `data:` frames.
///
/// Bytes stay raw until a newline so a multi-byte UTF-8 character split
/// across chunks is not turned into U+FFFD.
#[derive(Debug, Default)]
pub struct SseFrameReader {
    buffer: Vec<u8>,
    event: Option<String>,
    data: Vec<String>,
    data_bytes: usize,
    has_fields: bool,
}

impl SseFrameReader {
    /// Empty reader.
    #[must_use]
    pub fn new() -> Self {
        Self::default()
    }

    /// Append `bytes` and return every complete frame found so far.
    pub fn feed(&mut self, bytes: &[u8]) -> Result<Vec<RawSse>, String> {
        if self.buffer.len().saturating_add(bytes.len()) > MAX_SSE_PENDING
            && !bytes.contains(&b'\n')
            && !self.buffer.contains(&b'\n')
        {
            self.reset();
            return Err(format!("SSE line exceeds {MAX_SSE_PENDING} bytes"));
        }
        self.buffer.extend_from_slice(bytes);
        let mut out = Vec::new();
        let mut cursor = 0;
        while let Some(rel) = self.buffer[cursor..].iter().position(|&b| b == b'\n') {
            if rel > MAX_SSE_PENDING {
                self.reset();
                return Err(format!("SSE line exceeds {MAX_SSE_PENDING} bytes"));
            }
            let end = cursor + rel;
            let mut line_bytes = self.buffer[cursor..end].to_vec();
            if line_bytes.last() == Some(&b'\r') {
                line_bytes.pop();
            }
            let line = String::from_utf8_lossy(&line_bytes);
            cursor = end + 1;
            if let Some(frame) = self.push_line(&line)? {
                out.push(frame);
            }
        }
        // One drain for the consumed prefix. A per-line drain copies the tail.
        self.buffer.drain(..cursor);
        if self.buffer.len() > MAX_SSE_PENDING {
            self.reset();
            return Err(format!("SSE line exceeds {MAX_SSE_PENDING} bytes"));
        }
        Ok(out)
    }

    /// A frame has started and has not closed with a blank line.
    #[cfg(feature = "proxy")]
    pub fn has_open_frame(&self) -> bool {
        self.has_fields || !self.buffer.is_empty()
    }

    /// Bytes of the frame still open, so a later chunk can be joined
    /// without forwarding the partial payload early.
    #[cfg(feature = "proxy")]
    pub fn pending_wire_bytes(&self) -> Vec<u8> {
        let mut out = Vec::new();
        if let Some(event) = &self.event {
            out.extend(format!("event: {event}\n").into_bytes());
        }
        for line in &self.data {
            out.extend(format!("data: {line}\n").into_bytes());
        }
        out.extend_from_slice(&self.buffer);
        out
    }

    /// Emit a last frame if fields are pending (EOF without a blank line).
    pub fn drain(&mut self) -> Option<RawSse> {
        self.finish().ok().flatten()
    }

    /// EOF. A pending line that breaks the data cap is an error.
    pub fn finish(&mut self) -> Result<Option<RawSse>, String> {
        if !self.buffer.is_empty() {
            let line = String::from_utf8_lossy(&self.buffer).into_owned();
            self.buffer.clear();
            match self.push_line(line.trim_end_matches('\r')) {
                Ok(Some(frame)) => return Ok(Some(frame)),
                Ok(None) => {}
                Err(err) => return Err(err),
            }
        }
        Ok(self.take_frame())
    }

    fn push_line(&mut self, line: &str) -> Result<Option<RawSse>, String> {
        if line.is_empty() {
            return Ok(self.take_frame());
        }
        if line.starts_with(':') {
            return Ok(None);
        }
        let (name, value) = match line.split_once(':') {
            Some((name, value)) => (name, value.strip_prefix(' ').unwrap_or(value)),
            None => (line, ""),
        };
        match name {
            "event" => {
                self.event = Some(value.to_string());
                self.has_fields = true;
            }
            "data" => {
                let extra = if self.data.is_empty() {
                    value.len()
                } else {
                    value.len().saturating_add(1)
                };
                if self.data_bytes.saturating_add(extra) > MAX_SSE_PENDING {
                    self.reset();
                    return Err(format!("SSE data exceeds {MAX_SSE_PENDING} bytes"));
                }
                self.data_bytes = self.data_bytes.saturating_add(extra);
                self.data.push(value.to_string());
                self.has_fields = true;
            }
            "id" | "retry" => self.has_fields = true,
            _ => {}
        }
        Ok(None)
    }

    fn take_frame(&mut self) -> Option<RawSse> {
        if !self.has_fields {
            return None;
        }
        let frame = RawSse {
            event: self.event.take(),
            data: self.data.join("\n"),
        };
        self.data.clear();
        self.data_bytes = 0;
        self.has_fields = false;
        Some(frame)
    }

    fn reset(&mut self) {
        self.buffer.clear();
        self.event = None;
        self.data.clear();
        self.data_bytes = 0;
        self.has_fields = false;
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn partial_line_stays_buffered() {
        let mut reader = SseFrameReader::new();
        assert!(reader.feed(b"data: hel").unwrap().is_empty());
        let frames = reader.feed(b"lo\n\n").unwrap();
        assert_eq!(frames.len(), 1);
        assert_eq!(frames[0].data, "hello");
    }

    #[test]
    fn keeps_event_field() {
        let mut reader = SseFrameReader::new();
        let frames = reader
            .feed(b"event: ping\ndata: {\"type\":\"ping\"}\n\n")
            .unwrap();
        assert_eq!(frames[0].event.as_deref(), Some("ping"));
        assert_eq!(frames[0].data, r#"{"type":"ping"}"#);
    }

    #[test]
    fn split_3byte_utf8_is_not_fffd() {
        let euro = "€";
        let bytes = euro.as_bytes();
        let mut reader = SseFrameReader::new();
        let mut first = b"data: ".to_vec();
        first.extend_from_slice(&bytes[..2]);
        assert!(reader.feed(&first).unwrap().is_empty());
        let mut second = bytes[2..].to_vec();
        second.extend_from_slice(b"\n\n");
        let frames = reader.feed(&second).unwrap();
        assert_eq!(frames[0].data, euro);
        assert!(!frames[0].data.contains('\u{FFFD}'));
    }

    #[test]
    fn drain_last_data_without_blank_line() {
        let mut reader = SseFrameReader::new();
        let items = reader.feed(b"data: first\n\ndata: last").unwrap();
        assert_eq!(items.len(), 1);
        assert_eq!(items[0].data, "first");
        let last = reader.drain().expect("drain last");
        assert_eq!(last.data, "last");
    }

    #[test]
    fn pending_over_cap_fails_closed() {
        let mut reader = SseFrameReader::new();
        let mut chunk = b"data: ".to_vec();
        chunk.resize(MAX_SSE_PENDING + 8, b'x');
        assert!(reader.feed(&chunk).is_err());
        assert!(reader.drain().is_none());
    }

    #[test]
    fn data_lines_over_cap_fails_closed() {
        let mut reader = SseFrameReader::new();
        let payload = "x".repeat(64 * 1024);
        let line = format!("data: {payload}\n");
        let mut saw_err = false;
        for _ in 0..(MAX_SSE_PENDING / payload.len() + 2) {
            match reader.feed(line.as_bytes()) {
                Ok(_) => {}
                Err(_) => {
                    saw_err = true;
                    break;
                }
            }
        }
        assert!(
            saw_err,
            "joined data: lines must fail closed at {MAX_SSE_PENDING} bytes"
        );
        assert!(reader.drain().is_none());
    }

    #[test]
    fn finish_pending_line_over_joined_cap_is_error() {
        let chunk = 64 * 1024;
        let payload = "x".repeat(chunk);
        let line = format!("data: {payload}\n");
        let n = MAX_SSE_PENDING / (chunk + 1);
        let mut reader = SseFrameReader::new();
        for _ in 0..n.saturating_sub(1) {
            reader.feed(line.as_bytes()).expect("under cap");
        }
        let pending = format!("data: {}", "y".repeat(chunk * 2));
        reader.feed(pending.as_bytes()).expect("buffer the tail");
        assert!(
            reader.finish().is_err(),
            "pending data line that crosses the joined cap must fail"
        );
    }

    #[test]
    fn many_short_frames_in_one_feed_stay_fast() {
        let n = 300_000;
        let body = "data: z\n\n".repeat(n);
        let mut reader = SseFrameReader::new();
        let started = std::time::Instant::now();
        let frames = reader.feed(body.as_bytes()).expect("feed");
        let elapsed = started.elapsed();
        assert_eq!(frames.len(), n);
        assert_eq!(frames.first().expect("first").data, "z");
        assert_eq!(frames.last().expect("last").data, "z");
        assert!(
            elapsed < std::time::Duration::from_secs(3),
            "one-chunk feed took {} ms",
            elapsed.as_millis()
        );
    }

    #[test]
    fn partial_tail_after_many_frames_completes_on_next_feed() {
        let mut reader = SseFrameReader::new();
        let first = reader.feed(b"data: a\n\ndata: hel").expect("first");
        assert_eq!(first.len(), 1);
        assert_eq!(first[0].data, "a");
        let second = reader.feed(b"lo\n\n").expect("second");
        assert_eq!(second.len(), 1);
        assert_eq!(second[0].data, "hello");
    }
}