sse_parse/

parser.rs

use alloc::collections::VecDeque;
use alloc::string::String;
use alloc::vec::Vec;
use core::mem;
use core::time::Duration;

use crate::Event;

const BOM: [u8; 3] = [0xEF, 0xBB, 0xBF];

/// Whether the one-time leading BOM check at stream offset 0 has completed.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum Bom {
    Undecided,
    Done,
}

/// The WHATWG per-stream state: the three line buffers plus the two pieces
/// of cross-event state, and the queue of completed events.
#[derive(Debug, Clone, Default)]
struct Core {
    data: String,
    event_type: String,
    id_buf: String,
    last_event_id: String,
    retry: Option<Duration>,
    pending_retry: Option<Duration>,
    retry_updates: u64,
    events: VecDeque<Event>,
}

impl Core {
    /// Classify one completed line and apply its field, per the field
    /// processing rules of the HTML Living Standard's "Interpreting an
    /// event stream" section.
    fn process_line(&mut self, line: &[u8]) {
        if line.is_empty() {
            self.dispatch();
            return;
        }
        if line[0] == b':' {
            return;
        }
        let (name, value) = match memchr::memchr(b':', line) {
            Some(i) => {
                let value = &line[i + 1..];
                (&line[..i], value.strip_prefix(b" ").unwrap_or(value))
            }
            None => (line, &line[line.len()..]),
        };
        match name {
            b"event" => {
                self.event_type.clear();
                push_lossy(&mut self.event_type, value);
            }
            b"data" => {
                push_lossy(&mut self.data, value);
                self.data.push('\n');
            }
            b"id" => {
                if !value.contains(&0) {
                    self.id_buf.clear();
                    push_lossy(&mut self.id_buf, value);
                }
            }
            b"retry" => {
                if let Some(ms) = parse_retry(value) {
                    let retry = Duration::from_millis(ms);
                    self.retry = Some(retry);
                    self.pending_retry = Some(retry);
                    self.retry_updates += 1;
                }
            }
            _ => {}
        }
    }

    /// The blank-line "dispatch the event" steps of the HTML Living
    /// Standard's "Interpreting an event stream" section, in order.
    fn dispatch(&mut self) {
        if self.last_event_id != self.id_buf {
            self.last_event_id.clone_from(&self.id_buf);
        }
        if self.data.is_empty() {
            self.event_type.clear();
            return;
        }
        if self.data.ends_with('\n') {
            self.data.pop();
        }
        let event = if self.event_type.is_empty() {
            String::from("message")
        } else {
            mem::take(&mut self.event_type)
        };
        self.events.push_back(Event {
            event,
            data: mem::take(&mut self.data),
            id: self.last_event_id.clone(),
            retry: self.pending_retry.take(),
        });
    }

    /// Drop the block under construction. The last event ID buffer is stream
    /// state and survives.
    fn discard_block(&mut self) {
        self.data.clear();
        self.event_type.clear();
    }

    fn block_len(&self) -> usize {
        self.data.len() + self.event_type.len()
    }
}

fn push_lossy(dst: &mut String, bytes: &[u8]) {
    match core::str::from_utf8(bytes) {
        Ok(s) => dst.push_str(s),
        Err(_) => dst.push_str(&String::from_utf8_lossy(bytes)),
    }
}

/// The spec's `retry` rule plus a pinned edge decision: only ASCII
/// digits, at least one, and the value must fit in a `u64`.
fn parse_retry(value: &[u8]) -> Option<u64> {
    if value.is_empty() {
        return None;
    }
    let mut ms: u64 = 0;
    for &b in value {
        if !b.is_ascii_digit() {
            return None;
        }
        ms = ms.checked_mul(10)?.checked_add(u64::from(b - b'0'))?;
    }
    Some(ms)
}

/// Test and diagnostics hook. Not part of the stable API.
#[doc(hidden)]
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
pub struct Stats {
    /// Bytes examined by terminator scanning plus bytes handed to line
    /// processing. Bounded by twice the bytes fed.
    pub bytes_scanned: u64,
    /// Number of times a valid `retry` field updated the reconnection time.
    pub retry_updates: u64,
    /// Bytes currently held: the rolling input buffer plus the block under
    /// construction.
    pub buffered_bytes: usize,
}

/// Sans-IO incremental SSE parser.
///
/// Feed bytes with [`feed`](Self::feed), drain completed events with
/// [`next_event`](Self::next_event), and call [`finish`](Self::finish) at end
/// of stream. Chunk boundaries never affect the output: any chunking of a
/// byte stream yields the same events as feeding it whole.
///
/// Time is linear in the bytes fed. Each byte is scanned once for a line
/// terminator and decoded once as part of its completed line, whatever the
/// chunk size. Memory is proportional to the longest block in flight unless
/// [`max_buffer`](Self::max_buffer) caps it.
#[derive(Debug, Clone)]
pub struct SseParser {
    /// Rolling input buffer: `buf[..consumed]` is processed and awaits
    /// compaction, `buf[consumed..scan_pos]` is a scanned partial line, and
    /// `buf[scan_pos..]` is not yet scanned.
    buf: Vec<u8>,
    consumed: usize,
    scan_pos: usize,
    /// The last scanned byte was a CR that ended a line. The next byte decides
    /// whether it was the CR of a CRLF pair (then the LF is swallowed).
    held_cr: bool,
    bom: Bom,
    finished: bool,
    core: Core,
    max_buffer: Option<usize>,
    /// An oversized block is being thrown away until the next blank line.
    discarding: bool,
    /// Bytes of the current line were dropped, so its terminator must not be
    /// mistaken for a blank line.
    in_dropped_line: bool,
    overflows: u64,
    bytes_scanned: u64,
}

impl Default for SseParser {
    fn default() -> Self {
        Self::new()
    }
}

impl SseParser {
    /// A parser with an empty last event ID string and no buffer limit.
    pub fn new() -> Self {
        SseParser {
            buf: Vec::new(),
            consumed: 0,
            scan_pos: 0,
            held_cr: false,
            bom: Bom::Undecided,
            finished: false,
            core: Core::default(),
            max_buffer: None,
            discarding: false,
            in_dropped_line: false,
            overflows: 0,
            bytes_scanned: 0,
        }
    }

    /// A parser seeded with a last event ID string, for reconnection.
    ///
    /// Events dispatched before the stream's first `id` field report this ID.
    pub fn with_last_event_id(id: impl Into<String>) -> Self {
        let mut parser = Self::new();
        parser.set_last_event_id(id);
        parser
    }

    /// Cap the bytes held for one block.
    ///
    /// A block overflows when the raw bytes of one of its lines, before the
    /// terminator, plus the bytes already in the block's `data` and `event`
    /// buffers exceed `limit`, or when those buffers exceed `limit` after a
    /// line is decoded (invalid UTF-8 expands to the three-byte U+FFFD). The
    /// checks run on every completed line and on the unterminated line after
    /// every feed, so the outcome is the same for every chunking of the
    /// stream, and after each feed the parser holds at most `limit` bytes
    /// plus one chunk.
    ///
    /// On overflow the parser drops the block and ignores every following
    /// line up to the next blank line, where it resynchronizes. Fields
    /// processed before the overflow keep their effect: an `id` or `retry`
    /// seen earlier in the block stands, and the blank line still commits the
    /// last event ID. [`overflows`](Self::overflows) counts dropped blocks.
    /// The default is unlimited.
    pub fn max_buffer(mut self, limit: usize) -> Self {
        self.max_buffer = Some(limit);
        self
    }

    /// Number of blocks dropped because they exceeded the
    /// [`max_buffer`](Self::max_buffer) cap.
    pub fn overflows(&self) -> u64 {
        self.overflows
    }

    pub(crate) fn has_pending_events(&self) -> bool {
        !self.core.events.is_empty()
    }

    pub(crate) fn set_last_event_id(&mut self, id: impl Into<String>) {
        let id = id.into();
        self.core.id_buf.clone_from(&id);
        self.core.last_event_id = id;
    }

    /// Feed one chunk of bytes.
    ///
    /// Every byte sequence is a valid event stream, so there is nothing to
    /// reject. `tests/no_panic.rs` and the fuzz target feed arbitrary bytes
    /// in arbitrary chunkings.
    ///
    /// Feeding after [`finish`](Self::finish) is a contract violation: it
    /// trips a debug assertion and is ignored in release builds.
    pub fn feed(&mut self, bytes: &[u8]) {
        debug_assert!(!self.finished, "SseParser::feed called after finish()");
        if self.finished || bytes.is_empty() {
            return;
        }
        self.buf.extend_from_slice(bytes);
        if self.bom == Bom::Undecided && !self.resolve_bom() {
            return;
        }
        self.scan();
        self.enforce_limit();
        self.compact();
    }

    /// Next completed event, in stream order, if any.
    pub fn next_event(&mut self) -> Option<Event> {
        self.core.events.pop_front()
    }

    /// Signal end of stream.
    ///
    /// A block with no closing blank line is discarded, per the WHATWG
    /// algorithm. Events completed before the end stay available through
    /// [`next_event`](Self::next_event). Calling this twice is harmless.
    pub fn finish(&mut self) {
        if self.finished {
            return;
        }
        self.finished = true;
        self.held_cr = false;
        self.buf = Vec::new();
        self.consumed = 0;
        self.scan_pos = 0;
        self.discarding = false;
        self.in_dropped_line = false;
        self.core.discard_block();
    }

    /// The last event ID string: the value a client sends back as
    /// `Last-Event-ID` on reconnection.
    ///
    /// Committed at every blank line, including blocks that dispatch no event.
    pub fn last_event_id(&self) -> &str {
        &self.core.last_event_id
    }

    /// Reconnection time from the most recent valid `retry` field, if any.
    ///
    /// Updated when the field is processed, whether or not its block
    /// dispatches an event.
    pub fn retry(&self) -> Option<Duration> {
        self.core.retry
    }

    /// Test and diagnostics hook. Not part of the stable API.
    #[doc(hidden)]
    pub fn stats(&self) -> Stats {
        Stats {
            bytes_scanned: self.bytes_scanned,
            retry_updates: self.core.retry_updates,
            buffered_bytes: self.buf.len() + self.core.block_len(),
        }
    }

    /// Strip one leading `EF BB BF` byte order mark at absolute offset 0,
    /// split-proof, per the event stream grammar of the HTML Living
    /// Standard's "Parsing an event stream" section.
    ///
    /// Returns false when more bytes are needed to decide. Only called while
    /// nothing has been consumed, so the buffer starts at stream offset 0.
    fn resolve_bom(&mut self) -> bool {
        let n = self.buf.len().min(BOM.len());
        if self.buf[..n] != BOM[..n] {
            self.bom = Bom::Done;
            return true;
        }
        if n < BOM.len() {
            return false;
        }
        self.bom = Bom::Done;
        self.consumed = BOM.len();
        self.scan_pos = BOM.len();
        true
    }

    /// Split completed lines at CR, LF, and CRLF, per the event stream
    /// grammar of the HTML Living Standard's "Parsing an event stream"
    /// section. Only `buf[scan_pos..]` is examined, so no byte is scanned
    /// twice.
    fn scan(&mut self) {
        if self.held_cr {
            if self.scan_pos >= self.buf.len() {
                return;
            }
            self.held_cr = false;
            if self.buf[self.scan_pos] == b'\n' {
                self.bytes_scanned += 1;
                self.scan_pos += 1;
                self.consumed = self.scan_pos;
            }
        }
        while self.scan_pos < self.buf.len() {
            let unscanned = &self.buf[self.scan_pos..];
            let i = match memchr::memchr2(b'\r', b'\n', unscanned) {
                Some(i) => i,
                None => {
                    self.bytes_scanned += unscanned.len() as u64;
                    self.scan_pos = self.buf.len();
                    break;
                }
            };
            let term = self.scan_pos + i;
            let is_cr = self.buf[term] == b'\r';
            self.bytes_scanned += i as u64 + 1;
            self.handle_line(term);
            self.consumed = term + 1;
            self.scan_pos = term + 1;
            if is_cr {
                if self.scan_pos < self.buf.len() {
                    if self.buf[self.scan_pos] == b'\n' {
                        self.bytes_scanned += 1;
                        self.scan_pos += 1;
                        self.consumed = self.scan_pos;
                    }
                } else {
                    self.held_cr = true;
                }
            }
        }
    }

    /// The line `buf[consumed..end]` just got its terminator.
    fn handle_line(&mut self, end: usize) {
        let line = &self.buf[self.consumed..end];
        self.bytes_scanned += line.len() as u64;
        if self.discarding {
            if self.in_dropped_line {
                self.in_dropped_line = false;
            } else if line.is_empty() {
                self.discarding = false;
                self.core.process_line(line);
            }
            return;
        }
        if !line.is_empty() && self.exceeds_cap(line.len()) {
            self.overflow();
            return;
        }
        self.core.process_line(line);
        // Lossy decoding can grow a value: each invalid byte becomes the
        // three-byte U+FFFD. Re-check the decoded block so the held bytes
        // stay under the cap.
        if self.exceeds_cap(0) {
            self.overflow();
        }
    }

    /// Whether `line_bytes` of one line on top of the block in progress
    /// exceed the [`max_buffer`](Self::max_buffer) cap.
    fn exceeds_cap(&self, line_bytes: usize) -> bool {
        match self.max_buffer {
            Some(limit) => line_bytes.saturating_add(self.core.block_len()) > limit,
            None => false,
        }
    }

    fn overflow(&mut self) {
        self.overflows += 1;
        self.discarding = true;
        self.core.discard_block();
    }

    /// Apply the cap to the unterminated line left after a scan. While a
    /// block is being discarded, its pending bytes are dropped as they arrive
    /// and their terminator is remembered as the end of a dropped line.
    fn enforce_limit(&mut self) {
        if self.max_buffer.is_none() {
            return;
        }
        let pending = self.buf.len() - self.consumed;
        if pending == 0 {
            return;
        }
        if !self.discarding && self.exceeds_cap(pending) {
            self.overflow();
        }
        if self.discarding {
            self.in_dropped_line = true;
            self.consumed = self.buf.len();
        }
    }

    /// Reclaim consumed bytes. A move only happens once the consumed prefix
    /// is at least as long as the remainder, so moves cost amortized O(1)
    /// per byte fed.
    fn compact(&mut self) {
        if self.consumed == 0 {
            return;
        }
        if self.consumed == self.buf.len() {
            self.buf.clear();
        } else if self.consumed >= self.buf.len() / 2 {
            self.buf.drain(..self.consumed);
        } else {
            return;
        }
        self.scan_pos -= self.consumed;
        self.consumed = 0;
    }
}

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

    #[test]
    fn retry_parsing() {
        assert_eq!(parse_retry(b""), None);
        assert_eq!(parse_retry(b"0"), Some(0));
        assert_eq!(parse_retry(b"3000"), Some(3000));
        assert_eq!(parse_retry(b"30x0"), None);
        assert_eq!(parse_retry(b"12 "), None);
        assert_eq!(parse_retry(b"-1"), None);
        assert_eq!(parse_retry(b"18446744073709551615"), Some(u64::MAX));
        assert_eq!(parse_retry(b"18446744073709551616"), None);
        assert_eq!(parse_retry(b"99999999999999999999999"), None);
    }

    #[test]
    fn compaction_keeps_offsets_consistent() {
        let mut parser = SseParser::new();
        parser.feed(b"data: a\n\ndata: bbbbbbbbbbbbbbbbbbbb");
        assert_eq!(parser.next_event().unwrap().data, "a");
        parser.feed(b"\n\n");
        assert_eq!(parser.next_event().unwrap().data, "bbbbbbbbbbbbbbbbbbbb");
        assert_eq!(parser.buf.len(), 0);
        assert_eq!(parser.consumed, 0);
        assert_eq!(parser.scan_pos, 0);
    }
}