timvw_rust_runtime/
sse.rs

1//! Server-sent events (`text/event-stream`): a parser, the events of a response, and their `data:` decoded as JSON.
2//!
3//! [`Response::events`] gives the events of a response as they arrive; [`Events::json`] decodes each `data:` payload
4//! to a type of your own. A `data: [DONE]` event or the end of the body ends the stream. [`Decoder`] is the parser
5//! underneath, for a body read some other way.
6//!
7//! ```no_run
8//! # async fn demo(client: timvw_rust_runtime::Client) -> Result<(), timvw_rust_runtime::Error> {
9//! use futures_util::StreamExt;
10//!
11//! #[derive(serde::Deserialize)]
12//! struct Chunk {
13//!     id: String,
14//! }
15//!
16//! let response = client
17//!     .request(timvw_rust_runtime::Method::POST, "https://api.example.com/v1/stream")
18//!     .json(&serde_json::json!({"stream": true}))
19//!     .send()
20//!     .await?;
21//! let mut chunks = response.events().json::<Chunk>();
22//! while let Some(chunk) = chunks.next().await {
23//!     println!("{}", chunk?.id);
24//! }
25//! # Ok(())
26//! # }
27//! ```
28
29use std::collections::VecDeque;
30use std::fmt;
31use std::marker::PhantomData;
32use std::pin::Pin;
33use std::task::{Context, Poll};
34
35use futures_core::Stream;
36use serde::de::DeserializeOwned;
37
38use crate::{Error, Response};
39
40/// The `data:` value that ends a stream.
41pub const DONE: &str = "[DONE]";
42
43/// One server-sent event, before its data is decoded.
44#[derive(Clone, Debug, Default, PartialEq, Eq)]
45pub struct Event {
46    /// The `event:` field, when the event has one.
47    pub event: Option<String>,
48    /// The `data:` lines of the event, joined with `\n`.
49    pub data: String,
50    /// The last `id:` field seen on the stream, when there was one.
51    pub id: Option<String>,
52}
53
54/// An incremental `text/event-stream` parser, following the HTML standard: lines end with CR, LF or CRLF, a line
55/// starting with `:` is a comment, an empty line dispatches the event, and an event without `data:` is not
56/// dispatched. Feed it the body as it arrives.
57#[derive(Clone, Debug, Default)]
58pub struct Decoder {
59    buffer: Vec<u8>,
60    after_cr: bool,
61    started: bool,
62    event: Option<String>,
63    data: String,
64    has_data: bool,
65    id: Option<String>,
66}
67
68impl Decoder {
69    /// A parser at the start of a stream.
70    pub fn new() -> Self {
71        Self::default()
72    }
73
74    /// Takes the next chunk of the body and appends the events it completes to `events`.
75    pub fn feed(&mut self, chunk: &[u8], events: &mut VecDeque<Event>) {
76        let mut chunk = chunk;
77        if !self.started {
78            // A UTF-8 byte order mark at the start of the stream is skipped.
79            let pending = self.buffer.len() + chunk.len();
80            let mut head = self.buffer.clone();
81            head.extend_from_slice(&chunk[..chunk.len().min(3)]);
82            if pending < 3 && b"\xEF\xBB\xBF".starts_with(&head) {
83                self.buffer.extend_from_slice(chunk);
84                return;
85            }
86            self.started = true;
87            let mut all = std::mem::take(&mut self.buffer);
88            all.extend_from_slice(chunk);
89            let rest = all.strip_prefix(b"\xEF\xBB\xBF").unwrap_or(&all).to_vec();
90            self.feed_lines(&rest, events);
91            return;
92        }
93        if self.after_cr {
94            if let Some(rest) = chunk.strip_prefix(b"\n") {
95                chunk = rest;
96            }
97            self.after_cr = false;
98        }
99        self.feed_lines(chunk, events);
100    }
101
102    fn feed_lines(&mut self, chunk: &[u8], events: &mut VecDeque<Event>) {
103        let mut start = 0;
104        let mut i = 0;
105        while i < chunk.len() {
106            match chunk[i] {
107                b'\n' | b'\r' => {
108                    let mut line = std::mem::take(&mut self.buffer);
109                    line.extend_from_slice(&chunk[start..i]);
110                    self.line(&line, events);
111                    if chunk[i] == b'\r' {
112                        if i + 1 < chunk.len() {
113                            if chunk[i + 1] == b'\n' {
114                                i += 1;
115                            }
116                        } else {
117                            self.after_cr = true;
118                        }
119                    }
120                    start = i + 1;
121                }
122                _ => {}
123            }
124            i += 1;
125        }
126        self.buffer.extend_from_slice(&chunk[start..]);
127    }
128
129    fn line(&mut self, line: &[u8], events: &mut VecDeque<Event>) {
130        if line.is_empty() {
131            if self.has_data {
132                events.push_back(Event {
133                    event: self.event.take(),
134                    data: std::mem::take(&mut self.data),
135                    id: self.id.clone(),
136                });
137            }
138            self.event = None;
139            self.data.clear();
140            self.has_data = false;
141            return;
142        }
143        if line[0] == b':' {
144            return;
145        }
146        let line = String::from_utf8_lossy(line);
147        let (field, value) = match line.find(':') {
148            Some(colon) => {
149                let value = &line[colon + 1..];
150                (&line[..colon], value.strip_prefix(' ').unwrap_or(value))
151            }
152            None => (&line[..], ""),
153        };
154        match field {
155            "data" => {
156                if self.has_data {
157                    self.data.push('\n');
158                }
159                self.data.push_str(value);
160                self.has_data = true;
161            }
162            "event" => self.event = Some(value.to_owned()),
163            "id" if !value.contains('\0') => self.id = Some(value.to_owned()),
164            _ => {}
165        }
166    }
167
168    /// The body ended: an event not completed by an empty line is discarded, as the standard requires.
169    pub fn finish(&mut self) {
170        self.buffer.clear();
171        self.event = None;
172        self.data.clear();
173        self.has_data = false;
174    }
175}
176
177/// The events of a `text/event-stream` body, as they come. A `data: [DONE]` event ends it, as does the end of the body
178/// or a failure reading it (yielded as `Err` first); the body is dropped then, which frees the connection.
179pub struct Events {
180    bytes: Option<Body>,
181    decoder: Decoder,
182    ready: VecDeque<Event>,
183}
184
185type Chunk = Box<dyn AsRef<[u8]> + Send>;
186type Body = Pin<Box<dyn Stream<Item = Result<Chunk, Error>> + Send>>;
187
188impl Events {
189    /// The events of a response body. The status and the content type are not checked.
190    pub fn from_response(response: Response) -> Self {
191        Self::from_bytes(response.bytes_stream())
192    }
193
194    /// The events of a body read as a stream of chunks.
195    pub fn from_bytes<S, B>(bytes: S) -> Self
196    where
197        S: Stream<Item = Result<B, Error>> + Send + 'static,
198        B: AsRef<[u8]> + Send + 'static,
199    {
200        Self {
201            bytes: Some(Box::pin(Chunks(Box::pin(bytes)))),
202            decoder: Decoder::new(),
203            ready: VecDeque::new(),
204        }
205    }
206}
207
208impl Stream for Events {
209    type Item = Result<Event, Error>;
210
211    fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
212        let this = &mut *self;
213        loop {
214            if let Some(event) = this.ready.pop_front() {
215                if event.data == DONE {
216                    this.ready.clear();
217                    this.bytes = None;
218                    return Poll::Ready(None);
219                }
220                return Poll::Ready(Some(Ok(event)));
221            }
222            let Some(bytes) = this.bytes.as_mut() else {
223                return Poll::Ready(None);
224            };
225            match bytes.as_mut().poll_next(cx) {
226                Poll::Pending => return Poll::Pending,
227                Poll::Ready(Some(Ok(chunk))) => {
228                    this.decoder.feed((*chunk).as_ref(), &mut this.ready)
229                }
230                Poll::Ready(Some(Err(error))) => {
231                    this.bytes = None;
232                    this.decoder.finish();
233                    return Poll::Ready(Some(Err(error)));
234                }
235                Poll::Ready(None) => {
236                    this.bytes = None;
237                    this.decoder.finish();
238                }
239            }
240        }
241    }
242}
243
244impl fmt::Debug for Events {
245    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
246        f.debug_struct("Events")
247            .field("open", &self.bytes.is_some())
248            .field("ready", &self.ready)
249            .finish_non_exhaustive()
250    }
251}
252
253/// A byte stream whose chunks are boxed, so that [`Events`] holds any body.
254struct Chunks<S>(Pin<Box<S>>);
255
256impl<S, B> Stream for Chunks<S>
257where
258    S: Stream<Item = Result<B, Error>>,
259    B: AsRef<[u8]> + Send + 'static,
260{
261    type Item = Result<Chunk, Error>;
262
263    fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
264        self.0
265            .as_mut()
266            .poll_next(cx)
267            .map(|item| item.map(|chunk| chunk.map(|bytes| Box::new(bytes) as Chunk)))
268    }
269}
270
271impl Events {
272    /// Decodes the `data:` payload of every event as JSON, to `T`.
273    ///
274    /// A payload that does not decode yields an error for which [`Error::is_decode`] is true, and the stream goes on
275    /// with the next event; an event with an empty payload yields nothing. A failure reading the body yields an error
276    /// and ends the stream.
277    pub fn json<T: DeserializeOwned>(self) -> JsonEvents<T> {
278        JsonEvents {
279            events: self,
280            types: PhantomData,
281        }
282    }
283}
284
285/// The events of a stream with their `data:` decoded as JSON: [`Events::json`].
286pub struct JsonEvents<T> {
287    events: Events,
288    types: PhantomData<fn() -> T>,
289}
290
291impl<T> JsonEvents<T> {
292    /// The events not yet read, undecoded.
293    pub fn into_events(self) -> Events {
294        self.events
295    }
296}
297
298impl<T: DeserializeOwned> Stream for JsonEvents<T> {
299    type Item = Result<T, Error>;
300
301    fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
302        loop {
303            match Pin::new(&mut self.events).poll_next(cx) {
304                Poll::Ready(Some(Ok(event))) if event.data.trim().is_empty() => {}
305                Poll::Ready(Some(Ok(event))) => {
306                    return Poll::Ready(Some(
307                        serde_json::from_str(&event.data).map_err(Error::decode),
308                    ));
309                }
310                Poll::Ready(Some(Err(error))) => return Poll::Ready(Some(Err(error))),
311                Poll::Ready(None) => return Poll::Ready(None),
312                Poll::Pending => return Poll::Pending,
313            }
314        }
315    }
316}
317
318impl<T> fmt::Debug for JsonEvents<T> {
319    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
320        f.debug_struct("JsonEvents")
321            .field("events", &self.events)
322            .finish()
323    }
324}
325
326#[cfg(test)]
327mod tests {
328    use super::*;
329    use futures_util::StreamExt;
330
331    fn events(chunks: &[&'static str]) -> Events {
332        let chunks: Vec<Result<&'static [u8], Error>> =
333            chunks.iter().map(|chunk| Ok(chunk.as_bytes())).collect();
334        Events::from_bytes(futures_util::stream::iter(chunks))
335    }
336
337    #[test]
338    fn parses_across_chunks_and_line_endings() {
339        let mut decoder = Decoder::new();
340        let mut out = VecDeque::new();
341        decoder.feed(b"\xEF\xBB\xBFevent: a\r", &mut out);
342        decoder.feed(b"\ndata: 1\r\ndata:", &mut out);
343        decoder.feed(b"2\n: comment\nid: 7\n\ndata: x\n", &mut out);
344        decoder.finish();
345        assert_eq!(
346            out.into_iter().collect::<Vec<_>>(),
347            vec![Event {
348                event: Some("a".into()),
349                data: "1\n2".into(),
350                id: Some("7".into()),
351            }]
352        );
353    }
354
355    #[tokio::test]
356    async fn done_ends_the_stream() {
357        let data: Vec<_> = events(&["data: 1\n\n", "data: [DONE]\n\ndata: 2\n\n"])
358            .map(|event| event.unwrap().data)
359            .collect()
360            .await;
361        assert_eq!(data, vec!["1"]);
362    }
363
364    #[tokio::test]
365    async fn decodes_json_and_goes_on_after_a_bad_payload() {
366        #[derive(serde::Deserialize, Debug, PartialEq)]
367        struct Chunk {
368            n: u32,
369        }
370        let items: Vec<_> =
371            events(&["data: {\"n\":1}\n\ndata: oops\n\ndata:\n\ndata: {\"n\":2}\n\n"])
372                .json::<Chunk>()
373                .collect()
374                .await;
375        assert_eq!(items.len(), 3);
376        assert_eq!(items[0].as_ref().unwrap(), &Chunk { n: 1 });
377        assert!(items[1].as_ref().unwrap_err().is_decode());
378        assert_eq!(items[2].as_ref().unwrap(), &Chunk { n: 2 });
379    }
380}
timvw_rust_runtime/sse.rs

timvw_rust_runtime/
sse.rs
