Skip to main content

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}
31
32impl Default for SseDecoder {
33    fn default() -> Self {
34        Self::new(1024 * 1024)
35    }
36}
37
38impl SseDecoder {
39    pub fn new(max_line_bytes: usize) -> Self {
40        Self {
41            line: Vec::new(),
42            prev_cr: false,
43            first: true,
44            event: String::new(),
45            data: Vec::new(),
46            has_data: false,
47            id: None,
48            retry_ms: None,
49            max_line: max_line_bytes,
50        }
51    }
52
53    pub fn push(&mut self, bytes: &[u8]) -> Result<Vec<SseEvent>, HttpError> {
54        let mut out = Vec::new();
55        for &b in bytes {
56            if self.prev_cr {
57                self.prev_cr = false;
58                if b == b'\n' {
59                    continue; // second half of CRLF
60                }
61            }
62            match b {
63                b'\n' | b'\r' => {
64                    self.prev_cr = b == b'\r';
65                    let line = std::mem::take(&mut self.line);
66                    if let Some(ev) = self.handle_line(&line)? {
67                        out.push(ev);
68                    }
69                }
70                _ => {
71                    self.line.push(b);
72                    if self.line.len() > self.max_line {
73                        return Err(HttpError::Stream("SSE line exceeds limit".into()));
74                    }
75                }
76            }
77        }
78        Ok(out)
79    }
80
81    fn handle_line(&mut self, raw: &[u8]) -> Result<Option<SseEvent>, HttpError> {
82        let mut line = raw;
83        if self.first {
84            self.first = false;
85            if line.starts_with(&[0xEF, 0xBB, 0xBF]) {
86                line = &line[3..];
87            }
88        }
89        if line.is_empty() {
90            return Ok(self.dispatch());
91        }
92        let text = std::str::from_utf8(line)
93            .map_err(|_| HttpError::Stream("SSE line is not valid UTF-8".into()))?;
94        if text.starts_with(':') {
95            return Ok(None);
96        }
97        let (field, value) = match text.split_once(':') {
98            Some((f, v)) => (f, v.strip_prefix(' ').unwrap_or(v)),
99            None => (text, ""),
100        };
101        match field {
102            "event" => self.event = value.to_string(),
103            "data" => {
104                self.data.push(value.to_string());
105                self.has_data = true;
106            }
107            "id" if !value.contains('\0') => self.id = Some(value.to_string()),
108            "retry" => {
109                if let Ok(ms) = value.parse::<u64>() {
110                    self.retry_ms = Some(ms);
111                }
112            }
113            _ => {}
114        }
115        Ok(None)
116    }
117
118    fn dispatch(&mut self) -> Option<SseEvent> {
119        let had = self.has_data;
120        let event = std::mem::take(&mut self.event);
121        let data = std::mem::take(&mut self.data);
122        self.has_data = false;
123        let retry_ms = self.retry_ms.take();
124        if !had {
125            return None;
126        }
127        Some(SseEvent {
128            event: if event.is_empty() {
129                "message".into()
130            } else {
131                event
132            },
133            data: data.join("\n"),
134            id: self.id.clone(),
135            retry_ms,
136        })
137    }
138}
139
140/// Blocking iterator of events over any reader. An event cut off by EOF is
141/// dropped, per the EventSource spec.
142pub struct SseReader<R: Read> {
143    reader: R,
144    decoder: SseDecoder,
145    queue: VecDeque<SseEvent>,
146    done: bool,
147}
148
149impl<R: Read> SseReader<R> {
150    pub fn new(reader: R) -> Self {
151        Self {
152            reader,
153            decoder: SseDecoder::default(),
154            queue: VecDeque::new(),
155            done: false,
156        }
157    }
158}
159
160impl<R: Read> Iterator for SseReader<R> {
161    type Item = Result<SseEvent, HttpError>;
162    fn next(&mut self) -> Option<Self::Item> {
163        let mut buf = [0u8; 8192];
164        loop {
165            if let Some(ev) = self.queue.pop_front() {
166                return Some(Ok(ev));
167            }
168            if self.done {
169                return None;
170            }
171            match self.reader.read(&mut buf) {
172                Ok(0) => {
173                    self.done = true;
174                }
175                Ok(n) => match self.decoder.push(&buf[..n]) {
176                    Ok(events) => self.queue.extend(events),
177                    Err(e) => {
178                        self.done = true;
179                        return Some(Err(e));
180                    }
181                },
182                Err(e) if e.kind() == std::io::ErrorKind::Interrupted => {}
183                Err(e) => {
184                    self.done = true;
185                    return Some(Err(match e.kind() {
186                        std::io::ErrorKind::TimedOut => HttpError::Timeout(e.to_string()),
187                        _ => HttpError::Transport {
188                            message: e.to_string(),
189                            connect_phase: false,
190                        },
191                    }));
192                }
193            }
194        }
195    }
196}