objstore_r/
sse.rs

1// Author: spicyzhug@gmail.com
2// Copyright: lituus-io, all rights reserved
3//
4//! Incremental server-sent-events decoder, for streaming responses read
5//! through [`crate::transport::BodyStream`].
6//!
7//! Chunks arrive split at arbitrary byte positions (even inside a UTF-8
8//! sequence or a `\r\n`). [`SseDecoder`] buffers only the unconsumed tail,
9//! compacts lazily, and hands out each event's `data` payload by reference to
10//! its internal buffer — no allocation per event once warm. The buffered
11//! event size is capped so a hostile stream cannot grow memory unboundedly.
12
13use crate::{Error, Result};
14
15/// Largest single event accepted, bytes.
16pub const MAX_EVENT_BYTES: usize = 8 << 20;
17
18/// Incremental SSE decoder. See the module docs.
19#[derive(Debug, Default)]
20pub struct SseDecoder {
21    buf: Vec<u8>,
22    start: usize,
23    scanned: usize,
24    data: String,
25}
26
27impl SseDecoder {
28    /// A fresh decoder.
29    #[must_use]
30    pub fn new() -> Self {
31        Self::default()
32    }
33
34    /// Append a received chunk.
35    pub fn push(&mut self, chunk: &[u8]) -> Result<()> {
36        if self.start > 0 && self.start * 2 >= self.buf.len() {
37            self.buf.drain(..self.start);
38            self.scanned -= self.start;
39            self.start = 0;
40        }
41        if self.buf.len() - self.start + chunk.len() > MAX_EVENT_BYTES {
42            return Err(Error::TooLarge { limit: MAX_EVENT_BYTES as u64 });
43        }
44        self.buf.extend_from_slice(chunk);
45        Ok(())
46    }
47
48    /// The next complete event's `data` payload (multi-line data joined with
49    /// `\n`), or `None` until more input arrives. Events without data
50    /// (comments, keep-alives, `event:`/`id:`-only) are skipped.
51    pub fn next_data(&mut self) -> Result<Option<&str>> {
52        loop {
53            let Some((end, next)) = self.find_boundary() else { return Ok(None) };
54            let event = &self.buf[self.start..end];
55            self.start = next;
56            self.scanned = next;
57            self.data.clear();
58            let mut has_data = false;
59            for line in event.split(|&b| b == b'\n') {
60                let line = line.strip_suffix(b"\r").unwrap_or(line);
61                let Some(rest) = line.strip_prefix(b"data:") else { continue };
62                let rest = rest.strip_prefix(b" ").unwrap_or(rest);
63                if has_data {
64                    self.data.push('\n');
65                }
66                has_data = true;
67                let text = std::str::from_utf8(rest)
68                    .map_err(|_| Error::Protocol { message: "sse data is not utf-8".into() })?;
69                self.data.push_str(text);
70            }
71            if has_data {
72                return Ok(Some(&self.data));
73            }
74        }
75    }
76
77    /// Unconsumed bytes remain (a truncated final event).
78    #[must_use]
79    pub fn has_partial(&self) -> bool {
80        self.buf[self.start..].iter().any(|b| !b.is_ascii_whitespace())
81    }
82
83    /// Find the end of the first complete event: a blank line (`\n\n`,
84    /// `\r\n\r\n`, or mixed). Returns (event end, next event start).
85    fn find_boundary(&mut self) -> Option<(usize, usize)> {
86        let b = &self.buf;
87        let mut i = self.scanned.max(self.start);
88        while i < b.len() {
89            if b[i] == b'\n' {
90                // Look for the next line being empty: "\n\n" or "\n\r\n".
91                if b.get(i + 1) == Some(&b'\n') {
92                    return Some((i, i + 2));
93                }
94                if b.get(i + 1) == Some(&b'\r') && b.get(i + 2) == Some(&b'\n') {
95                    return Some((i, i + 3));
96                }
97                if i + 2 >= b.len() {
98                    // Might complete with the next chunk: rescan from here.
99                    self.scanned = i;
100                    return None;
101                }
102            }
103            i += 1;
104        }
105        self.scanned = b.len();
106        None
107    }
108}
109
110#[cfg(test)]
111mod tests {
112    use super::*;
113
114    fn collect(chunks: &[&[u8]]) -> Vec<String> {
115        let mut d = SseDecoder::new();
116        let mut out = Vec::new();
117        for c in chunks {
118            d.push(c).unwrap();
119            while let Some(data) = d.next_data().unwrap() {
120                out.push(data.to_owned());
121            }
122        }
123        out
124    }
125
126    #[test]
127    fn whole_events() {
128        assert_eq!(collect(&[b"data: a\n\ndata: b\n\n"]), ["a", "b"]);
129    }
130
131    #[test]
132    fn events_split_at_every_byte() {
133        let src =
134            b"data: {\"x\":\"\xc3\xa9\"}\r\n\r\n: keep-alive\n\nevent: e\ndata: l1\ndata: l2\n\n";
135        let chunks: Vec<&[u8]> = src.chunks(1).collect();
136        assert_eq!(collect(&chunks), ["{\"x\":\"\u{e9}\"}", "l1\nl2"]);
137    }
138
139    #[test]
140    fn no_data_events_are_skipped() {
141        assert_eq!(collect(&[b": comment\n\nid: 1\n\n"]), Vec::<String>::new());
142    }
143
144    #[test]
145    fn partial_tail_is_reported() {
146        let mut d = SseDecoder::new();
147        d.push(b"data: a\n\ndata: trunc").unwrap();
148        assert_eq!(d.next_data().unwrap(), Some("a"));
149        assert_eq!(d.next_data().unwrap(), None);
150        assert!(d.has_partial());
151    }
152
153    #[test]
154    fn oversize_events_are_refused() {
155        let mut d = SseDecoder::new();
156        let big = vec![b'x'; MAX_EVENT_BYTES];
157        d.push(&big).unwrap();
158        assert!(d.push(b"y").is_err());
159    }
160
161    #[test]
162    fn buffer_compacts_under_steady_streaming() {
163        let mut d = SseDecoder::new();
164        for _ in 0..10_000 {
165            d.push(b"data: 0123456789\n\n").unwrap();
166            assert_eq!(d.next_data().unwrap(), Some("0123456789"));
167        }
168        assert!(d.buf.capacity() < 4096, "consumed events do not accumulate");
169    }
170}
objstore_r/sse.rs

objstore_r/
sse.rs
