rightkit_http/
sse.rs

1//! Server-Sent Events decoder (WHATWG EventSource framing): LF, CR and CRLF line
2//! endings, multi-line `data`, comments, `id`, `retry`, BOM, bounded line size.
3
4use std::collections::VecDeque;
5use std::io::Read;
6
7use crate::error::HttpError;
8
9#[derive(Debug, Clone, Default, PartialEq, Eq)]
10pub struct SseEvent {
11    /// Event type; "message" when the stream gave none.
12    pub event: String,
13    pub data: String,
14    pub id: Option<String>,
15    pub retry_ms: Option<u64>,
16}
17
18/// Push-based decoder: feed bytes, receive completed events.
19#[derive(Debug)]
20pub struct SseDecoder {
21    line: Vec<u8>,
22    prev_cr: bool,
23    first: bool,
24    event: String,
25    data: Vec<String>,
26    has_data: bool,
27    id: Option<String>,
28    retry_ms: Option<u64>,
29    max_line: usize,
30    max_event: usize,
31    event_bytes: usize,
32    /// Deliver each `data:` line as its own event when its newline arrives (see [`SseDecoder::eager`]).
33    eager: bool,
34}
35
36impl Default for SseDecoder {
37    fn default() -> Self {
38        Self::new(1024 * 1024)
39    }
40}
41
42impl SseDecoder {
43    /// Decoder bounding each line; events are unbounded (use [`Self::with_limits`]).
44    pub fn new(max_line_bytes: usize) -> Self {
45        Self::with_limits(max_line_bytes, usize::MAX)
46    }
47
48    /// Decoder bounding each line and the accumulated bytes of one event
49    /// (event name plus all data lines, one separator byte per line).
50    pub fn with_limits(max_line_bytes: usize, max_event_bytes: usize) -> Self {
51        Self {
52            line: Vec::new(),
53            prev_cr: false,
54            first: true,
55            event: String::new(),
56            data: Vec::new(),
57            has_data: false,
58            id: None,
59            retry_ms: None,
60            max_line: max_line_bytes,
61            max_event: max_event_bytes,
62            event_bytes: 0,
63            eager: false,
64        }
65    }
66
67    /// Eager decoder with the default line bound. Every `data:` line becomes an event as soon
68    /// as its newline arrives, carrying that line as its `data`, instead of waiting for the
69    /// blank line that ends the frame. A provider whose daemon ends each frame with a newline
70    /// rather than a blank line needs this; the WHATWG framing is the default.
71    pub fn eager() -> Self {
72        Self {
73            eager: true,
74            ..Self::default()
75        }
76    }
77
78    /// Switch this decoder to eager mode (see [`SseDecoder::eager`]).
79    pub fn with_eager(mut self) -> Self {
80        self.eager = true;
81        self
82    }
83
84    /// True when bytes or a frame are buffered that no event has delivered yet: a partial line,
85    /// or `data` lines waiting for their blank line. A stream that ends while this is true was
86    /// cut off mid-event.
87    pub fn has_pending(&self) -> bool {
88        !self.line.is_empty() || self.has_data
89    }
90
91    /// Flush at end of stream. A final line without its newline is processed, and a complete
92    /// frame still waiting for its blank line is delivered, so a daemon that closes right after
93    /// its last event does not lose it. Returns at most one event. Errors (invalid UTF-8 in the
94    /// final line, for example) are reported, never dropped.
95    pub fn finish(&mut self) -> Result<Option<SseEvent>, HttpError> {
96        if !self.line.is_empty() {
97            let line = std::mem::take(&mut self.line);
98            if let Some(ev) = self.handle_line(&line)? {
99                return Ok(Some(ev));
100            }
101        }
102        Ok(self.dispatch())
103    }
104
105    pub fn push(&mut self, bytes: &[u8]) -> Result<Vec<SseEvent>, HttpError> {
106        let mut out = Vec::new();
107        for &b in bytes {
108            if self.prev_cr {
109                self.prev_cr = false;
110                if b == b'\n' {
111                    continue; // second half of CRLF
112                }
113            }
114            match b {
115                b'\n' | b'\r' => {
116                    self.prev_cr = b == b'\r';
117                    let line = std::mem::take(&mut self.line);
118                    if let Some(ev) = self.handle_line(&line)? {
119                        out.push(ev);
120                    }
121                }
122                _ => {
123                    self.line.push(b);
124                    if self.line.len() > self.max_line {
125                        return Err(HttpError::Stream("SSE line exceeds limit".into()));
126                    }
127                }
128            }
129        }
130        Ok(out)
131    }
132
133    fn handle_line(&mut self, raw: &[u8]) -> Result<Option<SseEvent>, HttpError> {
134        let mut line = raw;
135        if self.first {
136            self.first = false;
137            if line.starts_with(&[0xEF, 0xBB, 0xBF]) {
138                line = &line[3..];
139            }
140        }
141        if line.is_empty() {
142            return Ok(self.dispatch());
143        }
144        let text = std::str::from_utf8(line)
145            .map_err(|_| HttpError::Stream("SSE line is not valid UTF-8".into()))?;
146        if text.starts_with(':') {
147            return Ok(None);
148        }
149        let (field, value) = match text.split_once(':') {
150            Some((f, v)) => (f, v.strip_prefix(' ').unwrap_or(v)),
151            None => (text, ""),
152        };
153        if matches!(field, "event" | "data") {
154            self.event_bytes = self.event_bytes.saturating_add(value.len() + 1);
155            if self.event_bytes > self.max_event {
156                return Err(HttpError::Stream("SSE event exceeds limit".into()));
157            }
158        }
159        match field {
160            "event" => self.event = value.to_string(),
161            "data" if self.eager => {
162                self.event_bytes = 0;
163                let event = if self.event.is_empty() {
164                    "message".to_string()
165                } else {
166                    self.event.clone()
167                };
168                return Ok(Some(SseEvent {
169                    event,
170                    data: value.to_string(),
171                    id: self.id.clone(),
172                    retry_ms: self.retry_ms.take(),
173                }));
174            }
175            "data" => {
176                self.data.push(value.to_string());
177                self.has_data = true;
178            }
179            "id" if !value.contains('\0') => self.id = Some(value.to_string()),
180            "retry" => {
181                if let Ok(ms) = value.parse::<u64>() {
182                    self.retry_ms = Some(ms);
183                }
184            }
185            _ => {}
186        }
187        Ok(None)
188    }
189
190    fn dispatch(&mut self) -> Option<SseEvent> {
191        let had = self.has_data;
192        self.event_bytes = 0;
193        let event = std::mem::take(&mut self.event);
194        let data = std::mem::take(&mut self.data);
195        self.has_data = false;
196        let retry_ms = self.retry_ms.take();
197        if !had {
198            return None;
199        }
200        Some(SseEvent {
201            event: if event.is_empty() {
202                "message".into()
203            } else {
204                event
205            },
206            data: data.join("\n"),
207            id: self.id.clone(),
208            retry_ms,
209        })
210    }
211}
212
213/// What a reader does with an event cut off by end of stream.
214#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
215pub enum EofPolicy {
216    /// Drop it, per the EventSource spec (the default).
217    #[default]
218    Drop,
219    /// Deliver a complete final frame that lacks its blank line ([`SseDecoder::finish`]).
220    Deliver,
221    /// Report an error when the stream ends inside an event ([`SseDecoder::has_pending`]),
222    /// so a truncated stream fails closed.
223    Fail,
224}
225
226/// Blocking iterator of events over any reader. By default an event cut off by EOF is
227/// dropped, per the EventSource spec; see [`SseReader::eof_policy`].
228pub struct SseReader<R: Read> {
229    reader: R,
230    decoder: SseDecoder,
231    queue: VecDeque<SseEvent>,
232    done: bool,
233    eof: EofPolicy,
234}
235
236impl<R: Read> SseReader<R> {
237    pub fn new(reader: R) -> Self {
238        Self::with_decoder(reader, SseDecoder::default())
239    }
240
241    /// Read with a caller-built decoder (for example [`SseDecoder::eager`]).
242    pub fn with_decoder(reader: R, decoder: SseDecoder) -> Self {
243        Self {
244            reader,
245            decoder,
246            queue: VecDeque::new(),
247            done: false,
248            eof: EofPolicy::Drop,
249        }
250    }
251
252    /// Choose what happens to an event cut off by end of stream.
253    pub fn eof_policy(mut self, policy: EofPolicy) -> Self {
254        self.eof = policy;
255        self
256    }
257}
258
259impl<R: Read> Iterator for SseReader<R> {
260    type Item = Result<SseEvent, HttpError>;
261    fn next(&mut self) -> Option<Self::Item> {
262        let mut buf = [0u8; 8192];
263        loop {
264            if let Some(ev) = self.queue.pop_front() {
265                return Some(Ok(ev));
266            }
267            if self.done {
268                return None;
269            }
270            match self.reader.read(&mut buf) {
271                Ok(0) => {
272                    self.done = true;
273                    match self.eof {
274                        EofPolicy::Drop => {}
275                        EofPolicy::Deliver => match self.decoder.finish() {
276                            Ok(Some(ev)) => self.queue.push_back(ev),
277                            Ok(None) => {}
278                            Err(e) => return Some(Err(e)),
279                        },
280                        EofPolicy::Fail => {
281                            if self.decoder.has_pending() {
282                                return Some(Err(HttpError::Stream(
283                                    "SSE stream ended inside an event".into(),
284                                )));
285                            }
286                        }
287                    }
288                }
289                Ok(n) => match self.decoder.push(&buf[..n]) {
290                    Ok(events) => self.queue.extend(events),
291                    Err(e) => {
292                        self.done = true;
293                        return Some(Err(e));
294                    }
295                },
296                Err(e) if e.kind() == std::io::ErrorKind::Interrupted => {}
297                Err(e) => {
298                    self.done = true;
299                    return Some(Err(match e.kind() {
300                        std::io::ErrorKind::TimedOut => {
301                            HttpError::timeout(e.to_string(), crate::error::Phase::Body)
302                        }
303                        _ => HttpError::transport(e.to_string(), crate::error::Phase::Body),
304                    }));
305                }
306            }
307        }
308    }
309}
310
311#[cfg(test)]
312mod tests {
313    use super::*;
314
315    #[test]
316    fn eager_delivers_each_data_line_and_finish_flushes_the_tail() {
317        let mut d = SseDecoder::eager();
318        let evs = d.push(b"event: tick\ndata: one\ndata: two\n").unwrap();
319        assert_eq!(evs.len(), 2);
320        assert_eq!(evs[0].event, "tick");
321        assert_eq!(evs[0].data, "one");
322        assert_eq!(evs[1].data, "two");
323        assert!(!d.has_pending());
324        assert!(d.push(b"data: tail").unwrap().is_empty());
325        assert!(d.has_pending());
326        let last = d.finish().unwrap().unwrap();
327        assert_eq!(last.data, "tail");
328        assert!(!d.has_pending());
329    }
330
331    #[test]
332    fn default_mode_finish_delivers_a_frame_missing_its_blank_line() {
333        let mut d = SseDecoder::new(1024);
334        assert!(d.push(b"data: a\n").unwrap().is_empty());
335        assert!(d.has_pending());
336        let ev = d.finish().unwrap().unwrap();
337        assert_eq!(ev.data, "a");
338        assert_eq!(ev.event, "message");
339        assert!(!d.has_pending());
340    }
341
342    #[test]
343    fn reader_eof_policies_drop_deliver_or_fail() {
344        let input: &[u8] = b"data: a\n\ndata: b";
345        let dropped: Vec<String> = SseReader::new(input).map(|r| r.unwrap().data).collect();
346        assert_eq!(dropped, vec!["a"]);
347        let delivered: Vec<String> = SseReader::new(input)
348            .eof_policy(EofPolicy::Deliver)
349            .map(|r| r.unwrap().data)
350            .collect();
351        assert_eq!(delivered, vec!["a", "b"]);
352        let failed: Vec<Result<SseEvent, HttpError>> =
353            SseReader::new(input).eof_policy(EofPolicy::Fail).collect();
354        assert!(failed.last().unwrap().is_err());
355    }
356}
rightkit_http/sse.rs

rightkit_http/
sse.rs
