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];
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum Bom {
Undecided,
Done,
}
#[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 {
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;
}
}
_ => {}
}
}
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(),
});
}
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)),
}
}
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)
}
#[doc(hidden)]
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
pub struct Stats {
pub bytes_scanned: u64,
pub retry_updates: u64,
pub buffered_bytes: usize,
}
#[derive(Debug, Clone)]
pub struct SseParser {
buf: Vec<u8>,
consumed: usize,
scan_pos: usize,
held_cr: bool,
bom: Bom,
finished: bool,
core: Core,
max_buffer: Option<usize>,
discarding: bool,
in_dropped_line: bool,
overflows: u64,
bytes_scanned: u64,
}
impl Default for SseParser {
fn default() -> Self {
Self::new()
}
}
impl SseParser {
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,
}
}
pub fn with_last_event_id(id: impl Into<String>) -> Self {
let mut parser = Self::new();
parser.set_last_event_id(id);
parser
}
pub fn max_buffer(mut self, limit: usize) -> Self {
self.max_buffer = Some(limit);
self
}
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;
}
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();
}
pub fn next_event(&mut self) -> Option<Event> {
self.core.events.pop_front()
}
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();
}
pub fn last_event_id(&self) -> &str {
&self.core.last_event_id
}
pub fn retry(&self) -> Option<Duration> {
self.core.retry
}
#[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(),
}
}
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
}
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;
}
}
}
}
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);
if self.exceeds_cap(0) {
self.overflow();
}
}
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();
}
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();
}
}
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);
}
}