sse_stream/
stream.rs

1use std::{
2    collections::VecDeque,
3    num::ParseIntError,
4    str::Utf8Error,
5    task::{ready, Context, Poll},
6};
7
8use crate::Sse;
9use bytes::Buf;
10use futures_util::{stream::MapOk, Stream, TryStreamExt};
11use http_body::{Body, Frame};
12use http_body_util::{BodyDataStream, StreamBody};
13
14#[derive(Debug)]
15enum BomHeaderState {
16    NotFoundYet,
17    Parsing,
18    Consumed,
19}
20
21const BOM_HEADER: &[u8] = b"\xEF\xBB\xBF";
22
23pin_project_lite::pin_project! {
24    pub struct SseStream<B: Body> {
25        #[pin]
26        body: BodyDataStream<B>,
27        parsed: VecDeque<Sse>,
28        current: Option<Sse>,
29        unfinished_line: Vec<u8>,
30        mark_last_chunk_ending_with_cr: bool,
31        bom_header_state: BomHeaderState,
32    }
33}
34
35pub type ByteStreamBody<S, D> = StreamBody<MapOk<S, fn(D) -> Frame<D>>>;
36impl<E, S, D> SseStream<ByteStreamBody<S, D>>
37where
38    S: Stream<Item = Result<D, E>>,
39    E: std::error::Error,
40    D: Buf,
41    StreamBody<ByteStreamBody<S, D>>: Body,
42{
43    /// Alias of [`from_bytes_stream`](Self::from_bytes_stream).
44    #[deprecated(
45        since = "0.2.4",
46        note = "It's a typo, use `from_bytes_stream` instead. This method will be removed in 0.3.0"
47    )]
48    pub fn from_byte_stream(stream: S) -> Self {
49        Self::from_bytes_stream(stream)
50    }
51
52    /// Create a new [`SseStream`] from a stream of [`Bytes`](bytes::Bytes).
53    ///
54    /// This is useful when you interact with clients don't provide response body directly like reqwest.
55    pub fn from_bytes_stream(stream: S) -> Self {
56        let stream = stream.map_ok(http_body::Frame::data as fn(D) -> Frame<D>);
57        let body = StreamBody::new(stream);
58        Self {
59            body: BodyDataStream::new(body),
60            parsed: VecDeque::new(),
61            current: None,
62            unfinished_line: Vec::new(),
63            mark_last_chunk_ending_with_cr: false,
64            bom_header_state: BomHeaderState::NotFoundYet,
65        }
66    }
67}
68
69impl<B: Body> SseStream<B> {
70    /// Create a new [`SseStream`] from a [`Body`].
71    pub fn new(body: B) -> Self {
72        Self {
73            body: BodyDataStream::new(body),
74            parsed: VecDeque::new(),
75            current: None,
76            unfinished_line: Vec::new(),
77            mark_last_chunk_ending_with_cr: false,
78            bom_header_state: BomHeaderState::NotFoundYet,
79        }
80    }
81}
82
83pub enum Error {
84    Body(Box<dyn std::error::Error + Send + Sync>),
85    InvalidLine,
86    DuplicatedEventLine,
87    DuplicatedIdLine,
88    DuplicatedRetry,
89    Utf8Parse(Utf8Error),
90    IntParse(ParseIntError),
91}
92
93impl std::fmt::Display for Error {
94    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
95        match self {
96            Error::Body(e) => write!(f, "body error: {}", e),
97            Error::InvalidLine => write!(f, "invalid line"),
98            Error::DuplicatedEventLine => write!(f, "duplicated event line"),
99            Error::DuplicatedIdLine => write!(f, "duplicated id line"),
100            Error::DuplicatedRetry => write!(f, "duplicated retry line"),
101            Error::Utf8Parse(e) => write!(f, "utf8 parse error: {}", e),
102            Error::IntParse(e) => write!(f, "int parse error: {}", e),
103        }
104    }
105}
106
107impl std::fmt::Debug for Error {
108    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
109        match self {
110            Error::Body(e) => write!(f, "Body({:?})", e),
111            Error::InvalidLine => write!(f, "InvalidLine"),
112            Error::DuplicatedEventLine => write!(f, "DuplicatedEventLine"),
113            Error::DuplicatedIdLine => write!(f, "DuplicatedIdLine"),
114            Error::DuplicatedRetry => write!(f, "DuplicatedRetry"),
115            Error::Utf8Parse(e) => write!(f, "Utf8Parse({:?})", e),
116            Error::IntParse(e) => write!(f, "IntParse({:?})", e),
117        }
118    }
119}
120
121impl std::error::Error for Error {
122    fn description(&self) -> &str {
123        match self {
124            Error::Body(_) => "body error",
125            Error::InvalidLine => "invalid line",
126            Error::DuplicatedEventLine => "duplicated event line",
127            Error::DuplicatedIdLine => "duplicated id line",
128            Error::DuplicatedRetry => "duplicated retry line",
129            Error::Utf8Parse(_) => "utf8 parse error",
130            Error::IntParse(_) => "int parse error",
131        }
132    }
133
134    fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
135        match self {
136            Error::Body(e) => Some(e.as_ref()),
137            Error::Utf8Parse(e) => Some(e),
138            Error::IntParse(e) => Some(e),
139            _ => None,
140        }
141    }
142}
143
144impl<B: Body> Stream for SseStream<B>
145where
146    B::Error: std::error::Error + Send + Sync + 'static,
147{
148    type Item = Result<Sse, Error>;
149
150    fn poll_next(
151        mut self: std::pin::Pin<&mut Self>,
152        cx: &mut Context<'_>,
153    ) -> Poll<Option<Self::Item>> {
154        let this = self.as_mut().project();
155        if let Some(sse) = this.parsed.pop_front() {
156            return Poll::Ready(Some(Ok(sse)));
157        }
158        let next_data = ready!(this.body.poll_next(cx));
159        match next_data {
160            Some(Ok(mut data)) => {
161                loop {
162                    let mut bytes = data.chunk();
163                    let chunk_size = bytes.len();
164
165                    if *this.mark_last_chunk_ending_with_cr {
166                        if !bytes.is_empty() && bytes[0] == b'\n' {
167                            bytes = &bytes[1..];
168                        }
169                        *this.mark_last_chunk_ending_with_cr = false;
170                    }
171
172                    if bytes.is_empty() {
173                        return self.poll_next(cx);
174                    }
175                    if let BomHeaderState::NotFoundYet = this.bom_header_state {
176                        if bytes[0] == BOM_HEADER[0] {
177                            *this.bom_header_state = BomHeaderState::Parsing;
178                        }
179                    }
180                    // handling situation when the last line is end with `'\r'`. The next chunk may start with `'\n'`, but we should treat them as one line.
181                    if bytes.last().is_some_and(|b| *b == b'\r') {
182                        *this.mark_last_chunk_ending_with_cr = true;
183                    }
184                    let mut lines = bytes.chunk_by(|line_end, line_start| {
185                        !(
186                            // for line ending with `\n`, it can be either `\n` or `\r\n`
187                            *line_end == b'\n' ||
188                            // for line ending with `\r`
189                            (*line_end == b'\r' && *line_start != b'\n')
190                        )
191                    });
192                    let first_line = lines.next().expect("frame is empty");
193
194                    let mut new_unfinished_line = Vec::new();
195                    let mut first_line = if !this.unfinished_line.is_empty() {
196                        this.unfinished_line.extend(first_line);
197                        std::mem::swap(&mut new_unfinished_line, this.unfinished_line);
198                        new_unfinished_line.as_ref()
199                    } else {
200                        first_line
201                    };
202
203                    if let BomHeaderState::Parsing = this.bom_header_state {
204                        if first_line.len() > BOM_HEADER.len() {
205                            if let Some(stripped) = first_line.strip_prefix(BOM_HEADER) {
206                                first_line = stripped
207                            }
208                            // we only check the BOM header only ONCE in the whole stream, it happens instantly when we receive the first line with enough length.
209                            *this.bom_header_state = BomHeaderState::Consumed;
210                        } else {
211                            this.unfinished_line.extend(first_line);
212                            return self.poll_next(cx);
213                        }
214                    }
215
216                    let mut lines = std::iter::once(first_line).chain(lines);
217                    *this.unfinished_line = loop {
218                        let Some(line) = lines.next() else {
219                            break Vec::new();
220                        };
221                        let line = if line.ends_with(b"\r\n") {
222                            &line[..line.len() - 2]
223                        } else if line.ends_with(b"\n") || line.ends_with(b"\r") {
224                            &line[..line.len() - 1]
225                        } else {
226                            break line.to_vec();
227                        };
228
229                        if line.is_empty() {
230                            if let Some(sse) = this.current.take() {
231                                this.parsed.push_back(sse);
232                            }
233                            continue;
234                        }
235                        // find comma
236                        let Some(comma_index) = line.iter().position(|b| *b == b':') else {
237                            #[cfg(feature = "tracing")]
238                            tracing::warn!(?line, "invalid line, missing `:`");
239                            return Poll::Ready(Some(Err(Error::InvalidLine)));
240                        };
241                        let field_name = &line[..comma_index];
242                        let field_value = if line.len() > comma_index + 1 {
243                            let field_value = &line[comma_index + 1..];
244                            if field_value.starts_with(b" ") {
245                                &field_value[1..]
246                            } else {
247                                field_value
248                            }
249                        } else {
250                            b""
251                        };
252                        match field_name {
253                            b"data" => {
254                                let data_line =
255                                    std::str::from_utf8(field_value).map_err(Error::Utf8Parse)?;
256                                // merge data lines
257                                if let Some(Sse { data, .. }) = this.current.as_mut() {
258                                    if data.is_none() {
259                                        data.replace(data_line.to_owned());
260                                    } else {
261                                        let data = data.as_mut().unwrap();
262                                        data.push('\n');
263                                        data.push_str(data_line);
264                                    }
265                                } else {
266                                    this.current.replace(Sse {
267                                        event: None,
268                                        data: Some(data_line.to_owned()),
269                                        id: None,
270                                        retry: None,
271                                    });
272                                }
273                            }
274                            b"event" => {
275                                let event_value =
276                                    std::str::from_utf8(field_value).map_err(Error::Utf8Parse)?;
277                                if let Some(Sse { event, .. }) = this.current.as_mut() {
278                                    if event.is_some() {
279                                        return Poll::Ready(Some(Err(Error::DuplicatedEventLine)));
280                                    } else {
281                                        event.replace(event_value.to_owned());
282                                    }
283                                } else {
284                                    this.current.replace(Sse {
285                                        event: Some(event_value.to_owned()),
286                                        ..Default::default()
287                                    });
288                                }
289                            }
290                            b"id" => {
291                                // Per spec: if the id field value contains U+0000 NULL,
292                                // the entire field MUST be ignored.
293                                if field_value.contains(&0u8) {
294                                    #[cfg(feature = "tracing")]
295                                    tracing::warn!(
296                                        ?line,
297                                        "id field contains NULL byte, ignoring per spec"
298                                    );
299                                    continue;
300                                }
301                                let id_value =
302                                    std::str::from_utf8(field_value).map_err(Error::Utf8Parse)?;
303                                if let Some(Sse { id, .. }) = this.current.as_mut() {
304                                    if id.is_some() {
305                                        return Poll::Ready(Some(Err(Error::DuplicatedIdLine)));
306                                    } else {
307                                        id.replace(id_value.to_owned());
308                                    }
309                                } else {
310                                    this.current.replace(Sse {
311                                        id: Some(id_value.to_owned()),
312                                        ..Default::default()
313                                    });
314                                }
315                            }
316                            b"retry" => {
317                                let retry_value = std::str::from_utf8(field_value)
318                                    .map_err(Error::Utf8Parse)?
319                                    .trim_ascii();
320                                let retry_value =
321                                    retry_value.parse::<u64>().map_err(Error::IntParse)?;
322                                if let Some(Sse { retry, .. }) = this.current.as_mut() {
323                                    if retry.is_some() {
324                                        return Poll::Ready(Some(Err(Error::DuplicatedRetry)));
325                                    } else {
326                                        retry.replace(retry_value);
327                                    }
328                                } else {
329                                    this.current.replace(Sse {
330                                        retry: Some(retry_value),
331                                        ..Default::default()
332                                    });
333                                }
334                            }
335                            b"" => {
336                                #[cfg(feature = "tracing")]
337                                if tracing::enabled!(tracing::Level::DEBUG) {
338                                    // a comment
339                                    let comment = std::str::from_utf8(field_value)
340                                        .map_err(Error::Utf8Parse)?;
341                                    tracing::debug!(?comment, "sse comment line");
342                                }
343                            }
344                            _line => {
345                                #[cfg(feature = "tracing")]
346                                if tracing::enabled!(tracing::Level::WARN) {
347                                    tracing::warn!(line = ?_line, "invalid line: unknown field");
348                                }
349                                return Poll::Ready(Some(Err(Error::InvalidLine)));
350                            }
351                        }
352                    };
353                    data.advance(chunk_size);
354                    if !data.has_remaining() {
355                        break;
356                    }
357                }
358                self.poll_next(cx)
359            }
360            Some(Err(e)) => Poll::Ready(Some(Err(Error::Body(Box::new(e))))),
361            None => {
362                // When data stream terminated without empty line, we should discard last incomplate message.
363                Poll::Ready(None)
364            }
365        }
366    }
367}
sse_stream/stream.rs

sse_stream/
stream.rs
