1use alloc::{borrow::Cow, string::String, sync::Arc, vec, vec::Vec};
2use core::{
3 fmt, iter,
4 num::{NonZeroU8, NonZeroUsize},
5 str,
6};
7use thiserror::Error;
8
9use bytes::Buf;
10use memchr::{memchr, memchr2};
11
12#[derive(Debug, Clone, PartialEq, Eq)]
14#[cfg_attr(feature = "serde", derive(serde::Deserialize, serde::Serialize))]
15pub struct MessageEvent {
16 pub event: Cow<'static, str>,
18 pub data: String,
20 pub last_event_id: Option<Arc<str>>,
22}
23
24#[derive(Debug, Clone, PartialEq, Eq)]
26#[cfg_attr(feature = "serde", derive(serde::Deserialize, serde::Serialize))]
27pub enum SseEvent {
28 Message(MessageEvent),
30 Retry(u32),
32}
33
34#[derive(Debug, Clone, Copy, PartialEq, Eq, Default, Hash, Error)]
36#[cfg_attr(feature = "serde", derive(serde::Deserialize, serde::Serialize))]
37#[error("payload exceeded the allotted buffer size limit")]
38pub struct PayloadTooLargeError;
39
40const MAX_DEBUG_SIZE: usize = 200;
41
42struct ShowBigStr<'a>(&'a str);
43
44impl fmt::Debug for ShowBigStr<'_> {
45 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
46 let mut end = self.0.len().min(MAX_DEBUG_SIZE);
47 while !self.0.is_char_boundary(end) {
48 end -= 1;
49 }
50 let s = &self.0[..end];
51
52 fmt::Debug::fmt(s, f)?;
53 if end < self.0.len() {
54 write!(f, "... ({} bytes total)", self.0.len())?;
55 }
56
57 Ok(())
58 }
59}
60
61struct ShowBigBuf<'a>(&'a [u8]);
62
63impl fmt::Debug for ShowBigBuf<'_> {
64 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
65 let (buf, truncated) = match self.0.len() {
66 ..=MAX_DEBUG_SIZE => (self.0, false),
67 _ => (&self.0[..MAX_DEBUG_SIZE], true),
68 };
69
70 let mut chunks = buf.utf8_chunks().peekable();
71
72 f.write_str("\"")?;
73 while let Some(chunk) = chunks.next() {
74 fmt::Display::fmt(&chunk.valid().escape_debug(), f)?;
75
76 let invalid = chunk.invalid();
77 if invalid.is_empty() {
78 continue;
79 }
80
81 if truncated && chunks.peek().is_none() {
84 break;
85 }
86
87 for &byte in invalid {
88 write!(f, "\\x{byte:02X}")?;
89 }
90 }
91
92 if truncated {
93 write!(f, "\"... ({} bytes total)", self.0.len())?;
94 } else {
95 f.write_str("\"")?;
96 }
97
98 Ok(())
99 }
100}
101
102#[derive(Debug, Clone, Copy)]
103enum ValueMode {
104 Data,
105 Event,
106 Retry,
107 Id,
108}
109
110impl ValueMode {
111 pub const fn field_name(self) -> &'static str {
112 match self {
113 Self::Data => "data",
114 Self::Event => "event",
115 Self::Retry => "retry",
116 Self::Id => "id",
117 }
118 }
119}
120
121#[derive(Debug, Clone, Copy)]
122enum Mode {
123 Bom { bytes_read: u8 },
124 Field(Option<(ValueMode, NonZeroU8)>),
125 Value(ValueMode),
126 Ignore,
127 PostCr,
128 PostColon(ValueMode),
129}
130
131#[derive(Clone)]
136pub struct SseDecoder {
137 mode: Mode,
138 last_event_id: Option<Arc<str>>,
139 staged_last_event_id: Option<Arc<str>>,
140 last_event_id_buf: Vec<u8>,
141 event_buf: Vec<u8>,
142 data_buf: Vec<u8>,
143 retry_buf: Option<u32>,
144 max_payload_size: NonZeroUsize,
145 corrupted: bool,
146}
147
148impl SseDecoder {
149 #[inline]
165 #[must_use]
166 pub fn new() -> Self {
167 Self::with_limit(NonZeroUsize::new(512 * 1024).unwrap())
168 }
169
170 #[inline]
194 #[must_use]
195 pub fn with_limit(max_payload_size: NonZeroUsize) -> Self {
196 Self {
197 mode: Mode::Bom { bytes_read: 0 },
198 last_event_id: None,
199 staged_last_event_id: None,
200 last_event_id_buf: vec![],
201 event_buf: vec![],
202 data_buf: vec![],
203 retry_buf: None,
204 max_payload_size,
205 corrupted: false,
206 }
207 }
208
209 #[inline]
211 #[must_use]
212 pub fn last_event_id(&self) -> Option<&Arc<str>> {
213 self.last_event_id.as_ref()
214 }
215
216 #[inline]
227 pub fn reconnect_with_id(&mut self, id: Option<Arc<str>>) {
228 self.last_event_id = id;
229 self.reconnect();
230 }
231
232 #[inline]
241 pub fn clear(&mut self) {
242 self.reconnect_with_id(None);
243 }
244
245 #[inline]
254 pub fn reconnect(&mut self) {
255 self.mode = Mode::Bom { bytes_read: 0 };
256 self.clear_bufs();
257 self.corrupted = false;
258 }
259
260 fn mark_corrupted(&mut self) {
261 self.clear_bufs();
262 self.corrupted = true
263 }
264
265 #[inline]
266 fn clear_bufs(&mut self) {
267 self.data_buf.clear();
268 self.event_buf.clear();
269 self.last_event_id_buf.clear();
270 self.staged_last_event_id = self.last_event_id.clone();
271 }
272
273 fn dispatch(&mut self, cr: bool) -> Option<SseEvent> {
274 self.mode = match cr {
275 true => Mode::PostCr,
276 false => Mode::Field(None),
277 };
278
279 if self.corrupted {
280 self.corrupted = false;
281 return None;
283 }
284
285 self.last_event_id = self.staged_last_event_id.clone();
286
287 match self.data_buf.last() {
288 Some(b'\n') => {
289 self.data_buf.pop();
290 }
291 Some(_) => {}
292 None => {
293 self.event_buf.clear();
294 return None;
295 }
296 }
297
298 let data = String::from_utf8_lossy(&self.data_buf).into_owned();
299 self.data_buf.clear();
300
301 let event = match &*self.event_buf {
302 b"" => Cow::Borrowed("message"),
303 event_buf => Cow::Owned(String::from_utf8_lossy(event_buf).into_owned()),
304 };
305 self.event_buf.clear();
306
307 Some(SseEvent::Message(MessageEvent {
308 data,
309 event,
310 last_event_id: self.last_event_id.clone(),
311 }))
312 }
313
314 pub fn next(&mut self, buf: &mut impl Buf) -> Option<Result<SseEvent, PayloadTooLargeError>> {
346 loop {
366 let chunk = buf.chunk();
367 if chunk.is_empty() {
368 return None;
369 }
370
371 match &mut self.mode {
372 Mode::Bom { bytes_read } => {
373 let b0 = chunk[0];
374
375 const BOM: &[u8; 3] = b"\xef\xbb\xbf";
376
377 if b0 != BOM[*bytes_read as usize] {
378 self.mode = match *bytes_read {
379 0 => Mode::Field(None),
380 _ => Mode::Ignore,
381 };
382 continue;
383 }
384
385 buf.advance(1);
386 *bytes_read += 1;
387
388 if BOM.len() <= *bytes_read as usize {
389 self.mode = Mode::Field(None);
390 }
391 }
392 Mode::Field(None) => {
393 let b0 = chunk[0];
394 buf.advance(1);
395 let mode = match b0 {
396 b'd' if !self.corrupted => ValueMode::Data,
397 b'e' if !self.corrupted => ValueMode::Event,
398 b'i' if !self.corrupted => ValueMode::Id,
399 b'r' => ValueMode::Retry,
400
401 b'\n' | b'\r' => match self.dispatch(b0 == b'\r') {
402 Some(ev) => return Some(Ok(ev)),
403 None => continue,
404 },
405
406 _ => {
407 self.mode = Mode::Ignore;
408 continue;
409 }
410 };
411 self.mode = Mode::Field(Some((mode, NonZeroU8::new(1).unwrap())));
412 }
413 &mut Mode::Field(Some((mode, ref mut len))) => {
414 let cmp = &mode.field_name().as_bytes()[len.get() as usize..];
415 if iter::zip(chunk, cmp).any(|(ch0, ch1)| ch0 != ch1) {
416 self.mode = Mode::Ignore;
417 continue;
418 }
419 let Some(&b_post) = chunk.get(cmp.len()) else {
420 *len = NonZeroU8::new(len.get() + chunk.len() as u8).unwrap();
421 buf.advance(chunk.len());
422 continue;
423 };
424 buf.advance(cmp.len() + 1);
425
426 match b_post {
427 b'\n' => self.mode = Mode::Field(None),
428 b'\r' => self.mode = Mode::PostCr,
429 b':' => {
430 match mode {
431 ValueMode::Data => {}
432 ValueMode::Event => self.event_buf.clear(),
433 ValueMode::Id => self.last_event_id_buf.clear(),
434 ValueMode::Retry => self.retry_buf = None,
435 }
436
437 self.mode = Mode::PostColon(mode);
438 continue;
439 }
440 _ => {
441 self.mode = Mode::Ignore;
442 continue;
443 }
444 }
445
446 match mode {
447 ValueMode::Data => self.data_buf.push(b'\n'),
448 ValueMode::Id => self.last_event_id_buf.clear(),
449 ValueMode::Event | ValueMode::Retry => {}
450 }
451 }
452 Mode::Value(ValueMode::Retry) => {
453 let mut advanced = 0;
454 let mut return_event = false;
455
456 for &b in chunk {
457 advanced += 1;
458 match b {
459 b'0'..=b'9' => {
460 let digit = (b & 0xf) as _;
461
462 let retry_buf = self.retry_buf.unwrap_or(0);
463 let Some(retry_buf) = retry_buf.checked_mul(10) else {
464 self.mode = Mode::Ignore;
465 break;
466 };
467 let Some(retry_buf) = retry_buf.checked_add(digit) else {
468 self.mode = Mode::Ignore;
469 break;
470 };
471 self.retry_buf = Some(retry_buf);
472 }
473 b'\r' => {
474 self.mode = Mode::PostCr;
475 return_event = true;
476 break;
477 }
478 b'\n' => {
479 self.mode = Mode::Field(None);
480 return_event = true;
481 break;
482 }
483 _ => {
484 self.mode = Mode::Ignore;
485 break;
486 }
487 }
488 }
489
490 buf.advance(advanced);
491
492 if let (true, Some(retry_buf)) = (return_event, self.retry_buf) {
493 return Some(Ok(SseEvent::Retry(retry_buf)));
494 }
495 }
496 Mode::Value(ValueMode::Data) => {
497 match consume_until_newline(
498 &mut self.mode,
499 Some(&mut self.data_buf),
500 self.max_payload_size,
501 buf,
502 ) {
503 Ok(true) => self.data_buf.push(b'\n'),
504 Ok(false) => {}
505 Err(err) => {
506 self.mark_corrupted();
507 return Some(Err(err));
508 }
509 }
510 }
511 Mode::Value(ValueMode::Event) => {
512 if let Err(err) = consume_until_newline(
513 &mut self.mode,
514 Some(&mut self.event_buf),
515 self.max_payload_size,
516 buf,
517 ) {
518 self.mark_corrupted();
519 return Some(Err(err));
520 }
521 }
522 Mode::Value(ValueMode::Id) => {
523 match consume_until_newline(
524 &mut self.mode,
525 Some(&mut self.last_event_id_buf),
526 self.max_payload_size,
527 buf,
528 ) {
529 Ok(true) => {
530 if memchr(0, &self.last_event_id_buf).is_none() {
531 self.staged_last_event_id = match &*self.last_event_id_buf {
532 [] => None,
533 buf => Some(String::from_utf8_lossy(buf).into()),
534 };
535 }
536 self.last_event_id_buf.clear();
537 }
538 Ok(false) => {}
539 Err(err) => {
540 self.mark_corrupted();
541 return Some(Err(err));
542 }
543 }
544 }
545 Mode::Ignore => {
546 consume_until_newline(&mut self.mode, None, self.max_payload_size, buf)
547 .expect("there should be no payload to grow too large");
548 }
549 Mode::PostCr => {
550 if chunk[0] == b'\n' {
551 buf.advance(1);
552 }
553 self.mode = Mode::Field(None);
554 }
555 Mode::PostColon(value) => {
556 if chunk[0] == b' ' {
557 buf.advance(1);
558 }
559 self.mode = Mode::Value(*value);
560 }
561 }
562 }
563 }
564}
565
566impl Default for SseDecoder {
567 fn default() -> Self {
568 Self::new()
569 }
570}
571
572impl fmt::Debug for SseDecoder {
573 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
574 f.debug_struct("SseDecoder")
575 .field("mode", &self.mode)
576 .field(
577 "last_event_id",
578 &self.last_event_id.as_deref().map(ShowBigStr),
579 )
580 .field(
581 "staged_last_event_id",
582 &self.staged_last_event_id.as_deref().map(ShowBigStr),
583 )
584 .field("last_event_id_buf", &ShowBigBuf(&self.last_event_id_buf))
585 .field("event_buf", &ShowBigBuf(&self.event_buf))
586 .field("data_buf", &ShowBigBuf(&self.data_buf))
587 .field("retry_buf", &self.retry_buf)
588 .field("max_payload_size", &self.max_payload_size)
589 .finish()
590 }
591}
592
593fn consume_until_newline(
594 mode: &mut Mode,
595 mut out: Option<&mut Vec<u8>>,
596 max_size: NonZeroUsize,
597 buf: &mut impl Buf,
598) -> Result<bool, PayloadTooLargeError> {
599 loop {
600 let chunk = buf.chunk();
601 if chunk.is_empty() {
602 return Ok(false);
603 };
604
605 let Some(i) = memchr2(b'\r', b'\n', chunk) else {
606 if let Some(out) = out.as_deref_mut() {
607 if max_size.get() < out.len() + chunk.len() {
608 out.clear();
609 *mode = Mode::Ignore;
610 return Err(PayloadTooLargeError);
611 }
612 out.extend_from_slice(chunk);
613 }
614 buf.advance(chunk.len());
615 continue;
616 };
617
618 if let Some(out) = out.as_deref_mut() {
619 if max_size.get() < out.len() + i {
620 out.clear();
621 *mode = Mode::Ignore;
622 return Err(PayloadTooLargeError);
623 }
624 out.extend_from_slice(&chunk[..i]);
625 }
626
627 *mode = match chunk[i] {
628 b'\r' => Mode::PostCr,
629 b'\n' => Mode::Field(None),
630 _ => unreachable!(),
631 };
632
633 buf.advance(i + 1);
634
635 return Ok(true);
636 }
637}
638
639#[test]
640fn hard_parse() -> Result<(), PayloadTooLargeError> {
641 use core::slice;
642
643 let bytes = "\u{FEFF}data: x
645
646:
647
648event: my-event\r
649data:line1
650data: line2
651data
652:
653id: my-id
654:should be ignored too\rretry:42
655retry:
656
657data:second
658
659data:ignored
660";
661
662 let mut decoder = SseDecoder::new();
663
664 let events = bytes
665 .bytes()
666 .filter_map(|b| decoder.next(&mut slice::from_ref(&b)))
667 .collect::<Result<Vec<_>, PayloadTooLargeError>>()?;
668
669 let id = Some("my-id".into());
670
671 assert_eq!(
672 events,
673 &[
674 SseEvent::Message(MessageEvent {
675 event: "message".into(),
676 data: "x".into(),
677 last_event_id: None
678 }),
679 SseEvent::Retry(42),
680 SseEvent::Message(MessageEvent {
681 event: "my-event".into(),
682 data: "line1\nline2\n".into(),
683 last_event_id: id.clone()
684 }),
685 SseEvent::Message(MessageEvent {
686 event: "message".into(),
687 data: "second".into(),
688 last_event_id: id.clone()
689 })
690 ]
691 );
692 Ok(())
693}
694
695#[test]
696fn test_reconnect() {
697 let mut stream1: &[u8] = b"
698id: my-id
699
700event: my-event
701data:line1
702:
703data: line2
704id: ignored1
705";
706
707 let mut stream2: &[u8] = b"
708
709data: data
710
711id: final
712
713id: ignored2
714";
715
716 let my_id = Some("my-id".into());
717
718 let mut decoder = SseDecoder::new();
719
720 assert_eq!(decoder.next(&mut stream1), None);
721 assert_eq!(decoder.last_event_id(), my_id.as_ref());
722 assert!(stream1.is_empty());
723
724 decoder.reconnect();
725
726 assert_eq!(
728 decoder.next(&mut stream2),
729 Some(Ok(SseEvent::Message(MessageEvent {
730 event: "message".into(),
731 data: "data".into(),
732 last_event_id: my_id,
733 }))),
734 );
735 assert_eq!(decoder.next(&mut stream2), None);
736 assert_eq!(decoder.last_event_id().map(|id| &**id), Some("final"));
737 assert!(stream2.is_empty());
738}
739
740#[test]
741fn test_limits() {
742 let mut stream: &[u8] = b"
743data: 0123456789
744id: my-id
745
746id: 01234567890
747data: thing
748event: ev
749
750data: mid
751
752event: jojo
753id: ignored
754data: 01234
755data: 56789
756retry: 10
757
758event: final
759data
760
761";
762
763 let my_id = Some("my-id".into());
764
765 let mut decoder = SseDecoder::with_limit(NonZeroUsize::new(10).unwrap());
766
767 assert_eq!(
768 decoder.next(&mut stream),
769 Some(Ok(SseEvent::Message(MessageEvent {
770 event: "message".into(),
771 data: "0123456789".into(),
772 last_event_id: my_id.clone(),
773 })))
774 );
775 assert_eq!(decoder.next(&mut stream), Some(Err(PayloadTooLargeError)));
776 assert_eq!(
777 decoder.next(&mut stream),
778 Some(Ok(SseEvent::Message(MessageEvent {
779 event: "message".into(),
780 data: "mid".into(),
781 last_event_id: my_id.clone()
782 })))
783 );
784 assert_eq!(decoder.next(&mut stream), Some(Err(PayloadTooLargeError)));
785 assert_eq!(decoder.next(&mut stream), Some(Ok(SseEvent::Retry(10))));
786 assert_eq!(
787 decoder.next(&mut stream),
788 Some(Ok(SseEvent::Message(MessageEvent {
789 event: "final".into(),
790 data: "".into(),
791 last_event_id: my_id.clone()
792 })))
793 );
794 assert!(stream.is_empty());
795}