Skip to main content

eventsource_stream2/
event_stream.rs

1#[cfg(not(feature = "std"))]
2use alloc::string::{FromUtf8Error, String, ToString};
3
4#[cfg(feature = "std")]
5use std::string::FromUtf8Error;
6
7use crate::event::Event;
8use crate::parser::{is_bom, is_lf, line, RawEventLine};
9use crate::utf8_stream::{Utf8Stream, Utf8StreamError};
10use core::fmt;
11use core::pin::Pin;
12use core::time::Duration;
13use futures_core::stream::Stream;
14use futures_core::task::{Context, Poll};
15use pin_project_lite::pin_project;
16
17#[derive(Default, Debug)]
18struct EventBuilder {
19    event: Event,
20    is_complete: bool,
21}
22
23impl EventBuilder {
24    /// From the HTML spec
25    ///
26    /// -> If the field name is "event"
27    ///    Set the event type buffer to field value.
28    ///
29    /// -> If the field name is "data"
30    ///    Append the field value to the data buffer, then append a single U+000A LINE FEED (LF)
31    ///    character to the data buffer.
32    ///
33    /// -> If the field name is "id"
34    ///    If the field value does not contain U+0000 NULL, then set the last event ID buffer
35    ///    to the field value. Otherwise, ignore the field.
36    ///
37    /// -> If the field name is "retry"
38    ///    If the field value consists of only ASCII digits, then interpret the field value as
39    ///    an integer in base ten, and set the event stream's reconnection time to that integer.
40    ///    Otherwise, ignore the field.
41    ///
42    /// -> Otherwise
43    ///    The field is ignored.
44    fn add(&mut self, line: RawEventLine) {
45        match line {
46            RawEventLine::Field(field, val) => {
47                let val = val.unwrap_or("");
48                match field {
49                    "event" => {
50                        self.event.event = val.to_string();
51                    }
52                    "data" => {
53                        self.event.data.push_str(val);
54                        self.event.data.push('\u{000A}');
55                    }
56                    "id" => {
57                        if !val.contains('\u{0000}') {
58                            self.event.id = val.to_string()
59                        }
60                    }
61                    "retry" => {
62                        if let Ok(val) = val.parse::<u64>() {
63                            self.event.retry = Some(Duration::from_millis(val))
64                        }
65                    }
66                    _ => {}
67                }
68            }
69            RawEventLine::Comment => {}
70            RawEventLine::Empty => self.is_complete = true,
71        }
72    }
73
74    /// From the HTML spec
75    ///
76    /// 1. Set the last event ID string of the event source to the value of the last event ID buffer. The buffer does not get reset, so the last event ID string of the event source remains set to this value until the next time it is set by the server.
77    /// 2. If the data buffer is an empty string, set the data buffer and the event type buffer to the empty string and return.
78    /// 3. If the data buffer's last character is a U+000A LINE FEED (LF) character, then remove the last character from the data buffer.
79    /// 4. Let event be the result of creating an event using MessageEvent, in the relevant Realm of the EventSource object.
80    /// 5. Initialize event's type attribute to message, its data attribute to data, its origin attribute to the serialization of the origin of the event stream's final URL (i.e., the URL after redirects), and its lastEventId attribute to the last event ID string of the event source.
81    /// 6. If the event type buffer has a value other than the empty string, change the type of the newly created event to equal the value of the event type buffer.
82    /// 7. Set the data buffer and the event type buffer to the empty string.
83    /// 8. Queue a task which, if the readyState attribute is set to a value other than CLOSED, dispatches the newly created event at the EventSource object.
84    fn dispatch(&mut self) -> Option<Event> {
85        let builder = core::mem::take(self);
86        let mut event = builder.event;
87        self.event.id = event.id.clone();
88
89        if event.data.is_empty() {
90            return None;
91        }
92
93        if is_lf(event.data.chars().next_back().unwrap()) {
94            event.data.pop();
95        }
96
97        if event.event.is_empty() {
98            event.event = "message".to_string();
99        }
100
101        Some(event)
102    }
103}
104
105#[derive(Debug, Clone, Copy)]
106pub enum EventStreamState {
107    NotStarted,
108    Started,
109    Terminated,
110}
111
112impl EventStreamState {
113    fn is_terminated(self) -> bool {
114        matches!(self, Self::Terminated)
115    }
116    fn is_started(self) -> bool {
117        matches!(self, Self::Started)
118    }
119}
120
121pin_project! {
122/// A Stream of events
123pub struct EventStream<S> {
124    #[pin]
125    stream: Utf8Stream<S>,
126    buffer: String,
127    builder: EventBuilder,
128    state: EventStreamState,
129    last_event_id: String,
130}
131}
132
133impl<S> EventStream<S> {
134    /// Initialize the EventStream with a Stream
135    pub fn new(stream: S) -> Self {
136        Self {
137            stream: Utf8Stream::new(stream),
138            buffer: String::new(),
139            builder: EventBuilder::default(),
140            state: EventStreamState::NotStarted,
141            last_event_id: String::new(),
142        }
143    }
144
145    /// Set the last event ID of the stream. Useful for initializing the stream with a previous
146    /// last event ID
147    pub fn set_last_event_id(&mut self, id: impl Into<String>) {
148        self.last_event_id = id.into();
149    }
150
151    /// Get the last event ID of the stream
152    pub fn last_event_id(&self) -> &str {
153        &self.last_event_id
154    }
155}
156
157/// Error thrown while parsing an event line
158#[derive(Debug, PartialEq)]
159pub enum EventStreamError<E> {
160    /// Source stream is not valid UTF8
161    Utf8(FromUtf8Error),
162    /// Underlying source stream error
163    Transport(E),
164}
165
166impl<E> From<Utf8StreamError<E>> for EventStreamError<E> {
167    fn from(err: Utf8StreamError<E>) -> Self {
168        match err {
169            Utf8StreamError::Utf8(err) => Self::Utf8(err),
170            Utf8StreamError::Transport(err) => Self::Transport(err),
171        }
172    }
173}
174
175impl<E> fmt::Display for EventStreamError<E>
176where
177    E: fmt::Display,
178{
179    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
180        match self {
181            Self::Utf8(err) => f.write_fmt(format_args!("UTF8 error: {}", err)),
182            Self::Transport(err) => f.write_fmt(format_args!("Transport error: {}", err)),
183        }
184    }
185}
186
187#[cfg(feature = "std")]
188impl<E> std::error::Error for EventStreamError<E> where E: fmt::Display + fmt::Debug + Send + Sync {}
189
190fn parse_event(buffer: &mut String, builder: &mut EventBuilder) -> Option<Event> {
191    if buffer.is_empty() {
192        return None;
193    }
194    loop {
195        let (rem, next_line) = line(buffer.as_ref())?;
196        builder.add(next_line);
197        let consumed = buffer.len() - rem.len();
198        let rem = buffer.split_off(consumed);
199        *buffer = rem;
200        if builder.is_complete {
201            if let Some(event) = builder.dispatch() {
202                return Some(event);
203            }
204        }
205    }
206}
207
208impl<S, B, E> Stream for EventStream<S>
209where
210    S: Stream<Item = Result<B, E>>,
211    B: AsRef<[u8]>,
212{
213    type Item = Result<Event, EventStreamError<E>>;
214
215    fn poll_next(self: Pin<&mut Self>, cx: &mut Context) -> Poll<Option<Self::Item>> {
216        let mut this = self.project();
217
218        if let Some(event) = parse_event(this.buffer, this.builder) {
219            *this.last_event_id = event.id.clone();
220            return Poll::Ready(Some(Ok(event)));
221        }
222
223        if this.state.is_terminated() {
224            return Poll::Ready(None);
225        }
226
227        loop {
228            match this.stream.as_mut().poll_next(cx) {
229                Poll::Ready(Some(Ok(string))) => {
230                    if string.is_empty() {
231                        continue;
232                    }
233
234                    let slice = if this.state.is_started() {
235                        &string
236                    } else {
237                        *this.state = EventStreamState::Started;
238                        if is_bom(string.chars().next().unwrap()) {
239                            &string[1..]
240                        } else {
241                            &string
242                        }
243                    };
244                    this.buffer.push_str(slice);
245
246                    if let Some(event) = parse_event(this.buffer, this.builder) {
247                        *this.last_event_id = event.id.clone();
248                        return Poll::Ready(Some(Ok(event)));
249                    }
250                }
251                Poll::Ready(Some(Err(err))) => return Poll::Ready(Some(Err(err.into()))),
252                Poll::Ready(None) => {
253                    *this.state = EventStreamState::Terminated;
254                    return Poll::Ready(None);
255                }
256                Poll::Pending => return Poll::Pending,
257            }
258        }
259    }
260}
261
262#[cfg(test)]
263mod tests {
264    use super::*;
265    use futures::prelude::*;
266
267    #[tokio::test]
268    async fn valid_data_fields() {
269        assert_eq!(
270            EventStream::new(futures::stream::iter(vec![Ok::<_, ()>(
271                "data: Hello, world!\n\n"
272            )]))
273            .try_collect::<Vec<_>>()
274            .await
275            .unwrap(),
276            vec![Event {
277                event: "message".to_string(),
278                data: "Hello, world!".to_string(),
279                ..Default::default()
280            }]
281        );
282        assert_eq!(
283            EventStream::new(futures::stream::iter(vec![
284                Ok::<_, ()>("data: Hello,"),
285                Ok::<_, ()>(" world!\n\n")
286            ]))
287            .try_collect::<Vec<_>>()
288            .await
289            .unwrap(),
290            vec![Event {
291                event: "message".to_string(),
292                data: "Hello, world!".to_string(),
293                ..Default::default()
294            }]
295        );
296        assert_eq!(
297            EventStream::new(futures::stream::iter(vec![
298                Ok::<_, ()>("data: Hello,"),
299                Ok::<_, ()>(""),
300                Ok::<_, ()>(" world!\n\n")
301            ]))
302            .try_collect::<Vec<_>>()
303            .await
304            .unwrap(),
305            vec![Event {
306                event: "message".to_string(),
307                data: "Hello, world!".to_string(),
308                ..Default::default()
309            }]
310        );
311        assert_eq!(
312            EventStream::new(futures::stream::iter(vec![Ok::<_, ()>(
313                "data: Hello, world!\n"
314            )]))
315            .try_collect::<Vec<_>>()
316            .await
317            .unwrap(),
318            vec![]
319        );
320        assert_eq!(
321            EventStream::new(futures::stream::iter(vec![Ok::<_, ()>(
322                "data: Hello,\ndata: world!\n\n"
323            )]))
324            .try_collect::<Vec<_>>()
325            .await
326            .unwrap(),
327            vec![Event {
328                event: "message".to_string(),
329                data: "Hello,\nworld!".to_string(),
330                ..Default::default()
331            }]
332        );
333        assert_eq!(
334            EventStream::new(futures::stream::iter(vec![Ok::<_, ()>(
335                "data: Hello,\n\ndata: world!\n\n"
336            )]))
337            .try_collect::<Vec<_>>()
338            .await
339            .unwrap(),
340            vec![
341                Event {
342                    event: "message".to_string(),
343                    data: "Hello,".to_string(),
344                    ..Default::default()
345                },
346                Event {
347                    event: "message".to_string(),
348                    data: "world!".to_string(),
349                    ..Default::default()
350                }
351            ]
352        );
353    }
354
355    #[tokio::test]
356    async fn spec_examples() {
357        assert_eq!(
358            EventStream::new(futures::stream::iter(vec![Ok::<_, ()>(
359                "data: This is the first message.
360
361data: This is the second message, it
362data: has two lines.
363
364data: This is the third message.
365
366"
367            )]))
368            .try_collect::<Vec<_>>()
369            .await
370            .unwrap(),
371            vec![
372                Event {
373                    event: "message".to_string(),
374                    data: "This is the first message.".to_string(),
375                    ..Default::default()
376                },
377                Event {
378                    event: "message".to_string(),
379                    data: "This is the second message, it\nhas two lines.".to_string(),
380                    ..Default::default()
381                },
382                Event {
383                    event: "message".to_string(),
384                    data: "This is the third message.".to_string(),
385                    ..Default::default()
386                }
387            ]
388        );
389        assert_eq!(
390            EventStream::new(futures::stream::iter(vec![Ok::<_, ()>(
391                "event: add
392data: 73857293
393
394event: remove
395data: 2153
396
397event: add
398data: 113411
399
400"
401            )]))
402            .try_collect::<Vec<_>>()
403            .await
404            .unwrap(),
405            vec![
406                Event {
407                    event: "add".to_string(),
408                    data: "73857293".to_string(),
409                    ..Default::default()
410                },
411                Event {
412                    event: "remove".to_string(),
413                    data: "2153".to_string(),
414                    ..Default::default()
415                },
416                Event {
417                    event: "add".to_string(),
418                    data: "113411".to_string(),
419                    ..Default::default()
420                }
421            ]
422        );
423        assert_eq!(
424            EventStream::new(futures::stream::iter(vec![Ok::<_, ()>(
425                "data: YHOO
426data: +2
427data: 10
428
429"
430            )]))
431            .try_collect::<Vec<_>>()
432            .await
433            .unwrap(),
434            vec![Event {
435                event: "message".to_string(),
436                data: "YHOO\n+2\n10".to_string(),
437                ..Default::default()
438            },]
439        );
440        assert_eq!(
441            EventStream::new(futures::stream::iter(vec![Ok::<_, ()>(
442                ": test stream
443
444data: first event
445id: 1
446
447data:second event
448id
449
450data:  third event
451
452"
453            )]))
454            .try_collect::<Vec<_>>()
455            .await
456            .unwrap(),
457            vec![
458                Event {
459                    event: "message".to_string(),
460                    id: "1".to_string(),
461                    data: "first event".to_string(),
462                    ..Default::default()
463                },
464                Event {
465                    event: "message".to_string(),
466                    data: "second event".to_string(),
467                    ..Default::default()
468                },
469                Event {
470                    event: "message".to_string(),
471                    data: " third event".to_string(),
472                    ..Default::default()
473                }
474            ]
475        );
476        assert_eq!(
477            EventStream::new(futures::stream::iter(vec![Ok::<_, ()>(
478                "data
479
480data
481data
482
483data:
484"
485            )]))
486            .try_collect::<Vec<_>>()
487            .await
488            .unwrap(),
489            vec![
490                Event {
491                    event: "message".to_string(),
492                    data: "".to_string(),
493                    ..Default::default()
494                },
495                Event {
496                    event: "message".to_string(),
497                    data: "\n".to_string(),
498                    ..Default::default()
499                },
500            ]
501        );
502        assert_eq!(
503            EventStream::new(futures::stream::iter(vec![Ok::<_, ()>(
504                "data:test
505
506data: test
507
508"
509            )]))
510            .try_collect::<Vec<_>>()
511            .await
512            .unwrap(),
513            vec![
514                Event {
515                    event: "message".to_string(),
516                    data: "test".to_string(),
517                    ..Default::default()
518                },
519                Event {
520                    event: "message".to_string(),
521                    data: "test".to_string(),
522                    ..Default::default()
523                },
524            ]
525        );
526    }
527}