rws_server/test_client/
sse.rs

1//! Server-Sent Events in tests: [`TestResponse::sse`](super::TestResponse::sse)
2//! returns a [`TestSse`], which reads the response's `text/event-stream`
3//! body — streamed (`Sse::stream`, `Response::stream_pipe`) as it arrives, or
4//! buffered (`Sse::into_response`) at once — and hands out [`SseFrame`]s.
5//!
6//! [`SseParser`] is the parser underneath, usable on its own: it follows the
7//! HTML Standard's event-stream rules (lines end in CRLF, LF or CR; `:`
8//! starts a comment; `data` lines join with `\n`; a block without `data` is
9//! not dispatched; the event type defaults to `message`). It loops over its
10//! input, never recursing, and refuses an event (or a line) over
11//! [`MAX_EVENT_BYTES`] instead of buffering without bound.
12
13use std::io::Read;
14use std::sync::mpsc::{self, Receiver, RecvTimeoutError};
15use std::time::{Duration, Instant};
16
17/// The most one event (all its lines) or one unfinished line may hold:
18/// 8 MiB. Past it [`SseParser::feed`] returns an error.
19pub const MAX_EVENT_BYTES: usize = 8 * 1024 * 1024;
20
21/// One dispatched Server-Sent Event.
22#[derive(Debug, Clone, PartialEq, Eq)]
23pub struct SseFrame {
24    /// The `event` field; `"message"` when the event had none, as a
25    /// browser's `EventSource` reports it.
26    pub event: String,
27    /// The `data` lines joined with `\n`.
28    pub data: String,
29    /// The `id` field of this event, if it had one (the stream's
30    /// last-event-ID is the last `Some` seen).
31    pub id: Option<String>,
32    /// The `retry` field (milliseconds), if this event had a valid one.
33    pub retry: Option<u64>,
34}
35
36/// An incremental `text/event-stream` parser: feed it bytes in any chunks,
37/// get back the events they completed.
38///
39/// ```rust
40/// # extern crate rws_server as rust_web_server;
41/// use rust_web_server::test_client::SseParser;
42///
43/// let mut parser = SseParser::new();
44/// assert!(parser.feed(b"event: tick\ndata: 1\n").unwrap().is_empty());
45/// let events = parser.feed(b"id: 7\n\n").unwrap();
46/// assert_eq!("tick", events[0].event);
47/// assert_eq!("1", events[0].data);
48/// assert_eq!(Some("7"), events[0].id.as_deref());
49/// ```
50#[derive(Debug, Default)]
51pub struct SseParser {
52    line: Vec<u8>,
53    /// The previous byte was a CR (so an LF now ends nothing).
54    after_cr: bool,
55    /// The first bytes are still to be checked for a BOM.
56    started: bool,
57    event: Option<String>,
58    data: String,
59    has_data: bool,
60    id: Option<String>,
61    retry: Option<u64>,
62    /// Bytes buffered in the current event's fields.
63    pending: usize,
64    comments: usize,
65}
66
67impl SseParser {
68    /// A parser at the start of a stream.
69    pub fn new() -> Self {
70        SseParser::default()
71    }
72
73    /// Comment lines (`: …`, keep-alives) seen so far.
74    pub fn comments(&self) -> usize {
75        self.comments
76    }
77
78    /// Parse `bytes`, returning the events whose blank line they contained.
79    /// An event or line over [`MAX_EVENT_BYTES`] is an error (the parser is
80    /// then reset to the start of a stream).
81    pub fn feed(&mut self, mut bytes: &[u8]) -> Result<Vec<SseFrame>, String> {
82        let mut events = Vec::new();
83        if !self.started && !bytes.is_empty() {
84            // A UTF-8 BOM may open the stream; it isn't part of a field.
85            // (A BOM split across chunks is rare enough to treat as data.)
86            if bytes.starts_with(&[0xEF, 0xBB, 0xBF]) {
87                bytes = &bytes[3..];
88            }
89            self.started = true;
90        }
91        for &byte in bytes {
92            match byte {
93                b'\n' if self.after_cr => self.after_cr = false,
94                b'\n' | b'\r' => {
95                    self.after_cr = byte == b'\r';
96                    let line = std::mem::take(&mut self.line);
97                    if let Some(event) = self.line_done(&line) {
98                        events.push(event);
99                    }
100                    self.line = line;
101                    self.line.clear();
102                }
103                _ => {
104                    self.after_cr = false;
105                    if self.line.len() + self.pending >= MAX_EVENT_BYTES {
106                        *self = SseParser::default();
107                        return Err(format!("an SSE event over {MAX_EVENT_BYTES} bytes"));
108                    }
109                    self.line.push(byte);
110                }
111            }
112        }
113        Ok(events)
114    }
115
116    /// Handle one complete line; a blank one dispatches the event.
117    fn line_done(&mut self, line: &[u8]) -> Option<SseFrame> {
118        if line.is_empty() {
119            let event = self.event.take();
120            let id = self.id.take();
121            let retry = self.retry.take();
122            let data = std::mem::take(&mut self.data);
123            let has_data = std::mem::replace(&mut self.has_data, false);
124            self.pending = 0;
125            if !has_data {
126                return None;
127            }
128            return Some(SseFrame { event: event.unwrap_or_else(|| "message".to_string()), data, id, retry });
129        }
130        if line[0] == b':' {
131            self.comments += 1;
132            return None;
133        }
134        let (name, value) = match line.iter().position(|&b| b == b':') {
135            Some(colon) => {
136                let value = &line[colon + 1..];
137                (&line[..colon], value.strip_prefix(b" ").unwrap_or(value))
138            }
139            None => (line, &b""[..]),
140        };
141        // Every line counts (an empty `data:` still adds a `\n`), so a
142        // flood of short lines is bounded as well as one long one.
143        self.pending += line.len() + 1;
144        let value = String::from_utf8_lossy(value).into_owned();
145        match name {
146            b"event" => self.event = Some(value),
147            b"data" => {
148                if self.has_data {
149                    self.data.push('\n');
150                }
151                self.data.push_str(&value);
152                self.has_data = true;
153            }
154            b"id" if !value.contains('\0') => self.id = Some(value),
155            b"retry" if !value.is_empty() && value.bytes().all(|b| b.is_ascii_digit()) => {
156                self.retry = value.parse().ok();
157            }
158            _ => {}
159        }
160        None
161    }
162}
163
164/// What the reading thread sends: a chunk, the end, or a read error.
165enum Chunk {
166    Bytes(Vec<u8>),
167    End,
168    Failed(String),
169}
170
171/// Server-Sent Events read from a test response; see the module docs.
172pub struct TestSse {
173    parser: SseParser,
174    ready: std::collections::VecDeque<SseFrame>,
175    chunks: Option<Receiver<Chunk>>,
176    ended: bool,
177    error: Option<String>,
178}
179
180impl TestSse {
181    /// Events from a buffered body.
182    pub fn from_bytes(body: &[u8]) -> TestSse {
183        let mut sse = TestSse::empty();
184        sse.take_in(body);
185        sse.ended = true;
186        sse
187    }
188
189    /// Events from a streamed body, read on a thread of its own so that
190    /// waits can time out. The thread ends with the stream, or at the
191    /// first read after this `TestSse` is dropped (for `Sse::stream`, the
192    /// next event or keep-alive comment).
193    pub fn from_reader(mut reader: Box<dyn Read + Send>) -> TestSse {
194        let (tx, rx) = mpsc::channel();
195        let spawned = std::thread::Builder::new().name("rws-test-sse".to_string()).spawn(move || {
196            let mut buf = vec![0u8; 8192];
197            loop {
198                let chunk = match reader.read(&mut buf) {
199                    Ok(0) => Chunk::End,
200                    Ok(n) => Chunk::Bytes(buf[..n].to_vec()),
201                    Err(e) if e.kind() == std::io::ErrorKind::Interrupted => continue,
202                    Err(e) => Chunk::Failed(e.to_string()),
203                };
204                let last = !matches!(chunk, Chunk::Bytes(_));
205                if tx.send(chunk).is_err() || last {
206                    return;
207                }
208            }
209        });
210        let mut sse = TestSse::empty();
211        match spawned {
212            Ok(_) => sse.chunks = Some(rx),
213            Err(e) => {
214                sse.ended = true;
215                sse.error = Some(format!("couldn't start the SSE reader thread: {e}"));
216            }
217        }
218        sse
219    }
220
221    fn empty() -> TestSse {
222        TestSse { parser: SseParser::new(), ready: Default::default(), chunks: None, ended: false, error: None }
223    }
224
225    fn take_in(&mut self, bytes: &[u8]) {
226        match self.parser.feed(bytes) {
227            Ok(events) => self.ready.extend(events),
228            Err(e) => {
229                self.error = Some(e);
230                self.ended = true;
231                self.chunks = None;
232            }
233        }
234    }
235
236    /// The next event, waiting at most `timeout`; `None` when the stream
237    /// ended (see [`is_ended`](Self::is_ended)) or the time ran out.
238    pub fn next_event(&mut self, timeout: Duration) -> Option<SseFrame> {
239        let deadline = Instant::now() + timeout;
240        loop {
241            if let Some(event) = self.ready.pop_front() {
242                return Some(event);
243            }
244            if self.ended {
245                return None;
246            }
247            let rx = self.chunks.as_ref()?;
248            let left = deadline.saturating_duration_since(Instant::now());
249            match rx.recv_timeout(left) {
250                Ok(Chunk::Bytes(bytes)) => self.take_in(&bytes),
251                Ok(Chunk::End) | Err(RecvTimeoutError::Disconnected) => self.finish(None),
252                Ok(Chunk::Failed(e)) => self.finish(Some(e)),
253                Err(RecvTimeoutError::Timeout) => return None,
254            }
255        }
256    }
257
258    fn finish(&mut self, error: Option<String>) {
259        self.ended = true;
260        self.chunks = None;
261        if error.is_some() {
262            self.error = error;
263        }
264    }
265
266    /// Up to `n` events, waiting at most `timeout` in all.
267    pub fn events(&mut self, n: usize, timeout: Duration) -> Vec<SseFrame> {
268        let deadline = Instant::now() + timeout;
269        let mut events = Vec::new();
270        while events.len() < n {
271            match self.next_event(deadline.saturating_duration_since(Instant::now())) {
272                Some(event) => events.push(event),
273                None => break,
274            }
275        }
276        events
277    }
278
279    /// Every event until the stream ends, waiting at most `timeout` in
280    /// all. For a stream that never ends (`Sse::stream` with a sender kept
281    /// alive) this waits the whole `timeout`: prefer [`events`](Self::events).
282    pub fn collect_until_end(&mut self, timeout: Duration) -> Vec<SseFrame> {
283        self.events(usize::MAX, timeout)
284    }
285
286    /// `true` once the stream has ended (and every event was read out of it).
287    pub fn is_ended(&self) -> bool {
288        self.ended && self.ready.is_empty()
289    }
290
291    /// Why the stream stopped, if it was a read or parse error.
292    pub fn error(&self) -> Option<&str> {
293        self.error.as_deref()
294    }
295
296    /// Comment lines (keep-alives) seen so far.
297    pub fn comments(&self) -> usize {
298        self.parser.comments()
299    }
300}
rws_server/test_client/sse.rs

rws_server/test_client/
sse.rs
