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]
199 #[must_use]
200 pub fn with_limit(max_payload_size: NonZeroUsize) -> Self {
201 Self {
202 mode: Mode::Bom { bytes_read: 0 },
203 last_event_id: None,
204 staged_last_event_id: None,
205 last_event_id_buf: vec![],
206 event_buf: vec![],
207 data_buf: vec![],
208 retry_buf: None,
209 max_payload_size,
210 corrupted: false,
211 }
212 }
213
214 #[inline]
216 #[must_use]
217 pub fn last_event_id(&self) -> Option<&Arc<str>> {
218 self.last_event_id.as_ref()
219 }
220
221 #[inline]
232 pub fn reconnect_with_id(&mut self, id: Option<Arc<str>>) {
233 self.last_event_id = id;
234 self.reconnect();
235 }
236
237 #[inline]
246 pub fn clear(&mut self) {
247 self.reconnect_with_id(None);
248 }
249
250 #[inline]
259 pub fn reconnect(&mut self) {
260 self.mode = Mode::Bom { bytes_read: 0 };
261 self.clear_bufs();
262 self.corrupted = false;
263 }
264
265 fn mark_corrupted(&mut self) {
266 self.clear_bufs();
267 self.corrupted = true
268 }
269
270 fn push_data_newline(&mut self) -> Result<(), PayloadTooLargeError> {
278 if self.max_payload_size.get() < self.data_buf.len() {
279 self.mark_corrupted();
280 return Err(PayloadTooLargeError);
281 }
282 self.data_buf.push(b'\n');
283 Ok(())
284 }
285
286 #[inline]
287 fn clear_bufs(&mut self) {
288 self.data_buf.clear();
289 self.event_buf.clear();
290 self.last_event_id_buf.clear();
291 self.staged_last_event_id = self.last_event_id.clone();
292 }
293
294 fn dispatch(&mut self, cr: bool) -> Option<SseEvent> {
295 self.mode = match cr {
296 true => Mode::PostCr,
297 false => Mode::Field(None),
298 };
299
300 if self.corrupted {
301 self.corrupted = false;
302 return None;
304 }
305
306 self.last_event_id = self.staged_last_event_id.clone();
307
308 match self.data_buf.last() {
309 Some(b'\n') => {
310 self.data_buf.pop();
311 }
312 Some(_) => {}
313 None => {
314 self.event_buf.clear();
315 return None;
316 }
317 }
318
319 let data = String::from_utf8_lossy(&self.data_buf).into_owned();
320 self.data_buf.clear();
321
322 let event = match &*self.event_buf {
323 b"" => Cow::Borrowed("message"),
324 event_buf => Cow::Owned(String::from_utf8_lossy(event_buf).into_owned()),
325 };
326 self.event_buf.clear();
327
328 Some(SseEvent::Message(MessageEvent {
329 data,
330 event,
331 last_event_id: self.last_event_id.clone(),
332 }))
333 }
334
335 pub fn next(&mut self, buf: &mut impl Buf) -> Option<Result<SseEvent, PayloadTooLargeError>> {
379 loop {
399 let chunk = buf.chunk();
400 if chunk.is_empty() {
401 return None;
402 }
403
404 match &mut self.mode {
405 Mode::Bom { bytes_read } => {
406 let b0 = chunk[0];
407
408 const BOM: &[u8; 3] = b"\xef\xbb\xbf";
409
410 if b0 != BOM[*bytes_read as usize] {
411 self.mode = match *bytes_read {
412 0 => Mode::Field(None),
413 _ => Mode::Ignore,
414 };
415 continue;
416 }
417
418 buf.advance(1);
419 *bytes_read += 1;
420
421 if BOM.len() <= *bytes_read as usize {
422 self.mode = Mode::Field(None);
423 }
424 }
425 Mode::Field(None) => {
426 let b0 = chunk[0];
427 buf.advance(1);
428 let mode = match b0 {
429 b'd' if !self.corrupted => ValueMode::Data,
430 b'e' if !self.corrupted => ValueMode::Event,
431 b'i' if !self.corrupted => ValueMode::Id,
432 b'r' => ValueMode::Retry,
433
434 b'\n' | b'\r' => match self.dispatch(b0 == b'\r') {
435 Some(ev) => return Some(Ok(ev)),
436 None => continue,
437 },
438
439 _ => {
440 self.mode = Mode::Ignore;
441 continue;
442 }
443 };
444 self.mode = Mode::Field(Some((mode, NonZeroU8::new(1).unwrap())));
445 }
446 &mut Mode::Field(Some((mode, ref mut len))) => {
447 let cmp = &mode.field_name().as_bytes()[len.get() as usize..];
448 if iter::zip(chunk, cmp).any(|(ch0, ch1)| ch0 != ch1) {
449 self.mode = Mode::Ignore;
450 continue;
451 }
452 let Some(&b_post) = chunk.get(cmp.len()) else {
453 *len = NonZeroU8::new(len.get() + chunk.len() as u8).unwrap();
454 buf.advance(chunk.len());
455 continue;
456 };
457 buf.advance(cmp.len() + 1);
458
459 match b_post {
460 b'\n' => self.mode = Mode::Field(None),
461 b'\r' => self.mode = Mode::PostCr,
462 b':' => {
463 match mode {
464 ValueMode::Data => {}
465 ValueMode::Event => self.event_buf.clear(),
466 ValueMode::Id => self.last_event_id_buf.clear(),
467 ValueMode::Retry => self.retry_buf = None,
468 }
469
470 self.mode = Mode::PostColon(mode);
471 continue;
472 }
473 _ => {
474 self.mode = Mode::Ignore;
475 continue;
476 }
477 }
478
479 match mode {
484 ValueMode::Data => {
485 if let Err(err) = self.push_data_newline() {
486 return Some(Err(err));
487 }
488 }
489 ValueMode::Event => self.event_buf.clear(),
490 ValueMode::Id => {
491 self.last_event_id_buf.clear();
492 self.staged_last_event_id = None;
493 }
494 ValueMode::Retry => {}
495 }
496 }
497 Mode::Value(ValueMode::Retry) => {
498 let mut advanced = 0;
499 let mut return_event = false;
500
501 for &b in chunk {
502 advanced += 1;
503 match b {
504 b'0'..=b'9' => {
505 let digit = (b & 0xf) as _;
506
507 let retry_buf = self.retry_buf.unwrap_or(0);
508 let Some(retry_buf) = retry_buf.checked_mul(10) else {
509 self.mode = Mode::Ignore;
510 break;
511 };
512 let Some(retry_buf) = retry_buf.checked_add(digit) else {
513 self.mode = Mode::Ignore;
514 break;
515 };
516 self.retry_buf = Some(retry_buf);
517 }
518 b'\r' => {
519 self.mode = Mode::PostCr;
520 return_event = true;
521 break;
522 }
523 b'\n' => {
524 self.mode = Mode::Field(None);
525 return_event = true;
526 break;
527 }
528 _ => {
529 self.mode = Mode::Ignore;
530 break;
531 }
532 }
533 }
534
535 buf.advance(advanced);
536
537 if let (true, Some(retry_buf)) = (return_event, self.retry_buf) {
538 return Some(Ok(SseEvent::Retry(retry_buf)));
539 }
540 }
541 Mode::Value(ValueMode::Data) => {
542 match consume_until_newline(
543 &mut self.mode,
544 Some(&mut self.data_buf),
545 self.max_payload_size,
546 buf,
547 ) {
548 Ok(true) => {
549 if let Err(err) = self.push_data_newline() {
550 return Some(Err(err));
551 }
552 }
553 Ok(false) => {}
554 Err(err) => {
555 self.mark_corrupted();
556 return Some(Err(err));
557 }
558 }
559 }
560 Mode::Value(ValueMode::Event) => {
561 if let Err(err) = consume_until_newline(
562 &mut self.mode,
563 Some(&mut self.event_buf),
564 self.max_payload_size,
565 buf,
566 ) {
567 self.mark_corrupted();
568 return Some(Err(err));
569 }
570 }
571 Mode::Value(ValueMode::Id) => {
572 match consume_until_newline(
573 &mut self.mode,
574 Some(&mut self.last_event_id_buf),
575 self.max_payload_size,
576 buf,
577 ) {
578 Ok(true) => {
579 if memchr(0, &self.last_event_id_buf).is_none() {
580 self.staged_last_event_id = match &*self.last_event_id_buf {
581 [] => None,
582 buf => Some(String::from_utf8_lossy(buf).into()),
583 };
584 }
585 self.last_event_id_buf.clear();
586 }
587 Ok(false) => {}
588 Err(err) => {
589 self.mark_corrupted();
590 return Some(Err(err));
591 }
592 }
593 }
594 Mode::Ignore => {
595 consume_until_newline(&mut self.mode, None, self.max_payload_size, buf)
596 .expect("there should be no payload to grow too large");
597 }
598 Mode::PostCr => {
599 if chunk[0] == b'\n' {
600 buf.advance(1);
601 }
602 self.mode = Mode::Field(None);
603 }
604 Mode::PostColon(value) => {
605 if chunk[0] == b' ' {
606 buf.advance(1);
607 }
608 self.mode = Mode::Value(*value);
609 }
610 }
611 }
612 }
613}
614
615impl Default for SseDecoder {
616 fn default() -> Self {
617 Self::new()
618 }
619}
620
621impl fmt::Debug for SseDecoder {
622 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
623 f.debug_struct("SseDecoder")
624 .field("mode", &self.mode)
625 .field(
626 "last_event_id",
627 &self.last_event_id.as_deref().map(ShowBigStr),
628 )
629 .field(
630 "staged_last_event_id",
631 &self.staged_last_event_id.as_deref().map(ShowBigStr),
632 )
633 .field("last_event_id_buf", &ShowBigBuf(&self.last_event_id_buf))
634 .field("event_buf", &ShowBigBuf(&self.event_buf))
635 .field("data_buf", &ShowBigBuf(&self.data_buf))
636 .field("retry_buf", &self.retry_buf)
637 .field("max_payload_size", &self.max_payload_size)
638 .finish()
639 }
640}
641
642fn consume_until_newline(
643 mode: &mut Mode,
644 mut out: Option<&mut Vec<u8>>,
645 max_size: NonZeroUsize,
646 buf: &mut impl Buf,
647) -> Result<bool, PayloadTooLargeError> {
648 loop {
649 let chunk = buf.chunk();
650 if chunk.is_empty() {
651 return Ok(false);
652 };
653
654 let Some(i) = memchr2(b'\r', b'\n', chunk) else {
655 if let Some(out) = out.as_deref_mut() {
656 if max_size.get() < out.len() + chunk.len() {
657 out.clear();
658 *mode = Mode::Ignore;
659 return Err(PayloadTooLargeError);
660 }
661 out.extend_from_slice(chunk);
662 }
663 buf.advance(chunk.len());
664 continue;
665 };
666
667 if let Some(out) = out.as_deref_mut() {
668 if max_size.get() < out.len() + i {
669 out.clear();
670 *mode = Mode::Ignore;
671 return Err(PayloadTooLargeError);
672 }
673 out.extend_from_slice(&chunk[..i]);
674 }
675
676 *mode = match chunk[i] {
677 b'\r' => Mode::PostCr,
678 b'\n' => Mode::Field(None),
679 _ => unreachable!(),
680 };
681
682 buf.advance(i + 1);
683
684 return Ok(true);
685 }
686}
687
688#[test]
689fn hard_parse() -> Result<(), PayloadTooLargeError> {
690 use core::slice;
691
692 let bytes = "\u{FEFF}data: x
694
695:
696
697event: my-event\r
698data:line1
699data: line2
700data
701:
702id: my-id
703:should be ignored too\rretry:42
704retry:
705
706data:second
707
708data:ignored
709";
710
711 let mut decoder = SseDecoder::new();
712
713 let events = bytes
714 .bytes()
715 .filter_map(|b| decoder.next(&mut slice::from_ref(&b)))
716 .collect::<Result<Vec<_>, PayloadTooLargeError>>()?;
717
718 let id = Some("my-id".into());
719
720 assert_eq!(
721 events,
722 &[
723 SseEvent::Message(MessageEvent {
724 event: "message".into(),
725 data: "x".into(),
726 last_event_id: None
727 }),
728 SseEvent::Retry(42),
729 SseEvent::Message(MessageEvent {
730 event: "my-event".into(),
731 data: "line1\nline2\n".into(),
732 last_event_id: id.clone()
733 }),
734 SseEvent::Message(MessageEvent {
735 event: "message".into(),
736 data: "second".into(),
737 last_event_id: id.clone()
738 })
739 ]
740 );
741 Ok(())
742}
743
744#[test]
749fn test_valueless_fields() {
750 fn messages(bytes: &str) -> Vec<(String, String, Option<String>)> {
751 let mut decoder = SseDecoder::new();
752 let mut buf = bytes.as_bytes();
753 let mut out = vec![];
754 while let Some(SseEvent::Message(msg)) = decoder.next(&mut buf).transpose().unwrap() {
755 out.push((
756 msg.event.into_owned(),
757 msg.data,
758 msg.last_event_id.map(|id| id.to_string()),
759 ));
760 }
761 out
762 }
763
764 assert_eq!(
766 messages("event: custom\nevent\ndata: x\n\n"),
767 [("message".into(), "x".into(), None)],
768 );
769
770 for bytes in [
772 "id: abc\ndata: 1\n\nid\ndata: 2\n\n",
773 "id: abc\ndata: 1\n\nid:\ndata: 2\n\n",
774 ] {
775 assert_eq!(
776 messages(bytes),
777 [
778 ("message".into(), "1".into(), Some("abc".into())),
779 ("message".into(), "2".into(), None),
780 ],
781 "for {bytes:?}",
782 );
783 }
784
785 assert_eq!(
787 messages("data: a\ndata\ndata: b\n\n"),
788 [("message".into(), "a\n\nb".into(), None)],
789 );
790}
791
792#[test]
797fn test_valueless_data_respects_the_limit() {
798 let limit = NonZeroUsize::new(10).unwrap();
799
800 let input = "data\n".repeat(100) + "\n";
802 let mut buf = input.as_bytes();
803 let mut decoder = SseDecoder::with_limit(limit);
804 assert_eq!(decoder.next(&mut buf), Some(Err(PayloadTooLargeError)));
805
806 assert_eq!(decoder.next(&mut buf), None);
808 assert!(buf.is_empty());
809
810 let mut buf: &[u8] = b"data: after\n\n";
811 assert_eq!(
812 decoder.next(&mut buf),
813 Some(Ok(SseEvent::Message(MessageEvent {
814 event: "message".into(),
815 data: "after".into(),
816 last_event_id: None,
817 }))),
818 );
819
820 let input = "data\n".repeat(limit.get() + 1) + "\n";
822 let mut buf = input.as_bytes();
823 let mut decoder = SseDecoder::with_limit(limit);
824 assert_eq!(
825 decoder.next(&mut buf),
826 Some(Ok(SseEvent::Message(MessageEvent {
827 event: "message".into(),
828 data: "\n".repeat(limit.get()),
829 last_event_id: None,
830 }))),
831 );
832}
833
834#[test]
835fn test_reconnect() {
836 let mut stream1: &[u8] = b"
837id: my-id
838
839event: my-event
840data:line1
841:
842data: line2
843id: ignored1
844";
845
846 let mut stream2: &[u8] = b"
847
848data: data
849
850id: final
851
852id: ignored2
853";
854
855 let my_id = Some("my-id".into());
856
857 let mut decoder = SseDecoder::new();
858
859 assert_eq!(decoder.next(&mut stream1), None);
860 assert_eq!(decoder.last_event_id(), my_id.as_ref());
861 assert!(stream1.is_empty());
862
863 decoder.reconnect();
864
865 assert_eq!(
867 decoder.next(&mut stream2),
868 Some(Ok(SseEvent::Message(MessageEvent {
869 event: "message".into(),
870 data: "data".into(),
871 last_event_id: my_id,
872 }))),
873 );
874 assert_eq!(decoder.next(&mut stream2), None);
875 assert_eq!(decoder.last_event_id().map(|id| &**id), Some("final"));
876 assert!(stream2.is_empty());
877}
878
879#[test]
880fn test_limits() {
881 let mut stream: &[u8] = b"
882data: 0123456789
883id: my-id
884
885id: 01234567890
886data: thing
887event: ev
888
889data: mid
890
891event: jojo
892id: ignored
893data: 01234
894data: 56789
895retry: 10
896
897event: final
898data
899
900";
901
902 let my_id = Some("my-id".into());
903
904 let mut decoder = SseDecoder::with_limit(NonZeroUsize::new(10).unwrap());
905
906 assert_eq!(
907 decoder.next(&mut stream),
908 Some(Ok(SseEvent::Message(MessageEvent {
909 event: "message".into(),
910 data: "0123456789".into(),
911 last_event_id: my_id.clone(),
912 })))
913 );
914 assert_eq!(decoder.next(&mut stream), Some(Err(PayloadTooLargeError)));
915 assert_eq!(
916 decoder.next(&mut stream),
917 Some(Ok(SseEvent::Message(MessageEvent {
918 event: "message".into(),
919 data: "mid".into(),
920 last_event_id: my_id.clone()
921 })))
922 );
923 assert_eq!(decoder.next(&mut stream), Some(Err(PayloadTooLargeError)));
924 assert_eq!(decoder.next(&mut stream), Some(Ok(SseEvent::Retry(10))));
925 assert_eq!(
926 decoder.next(&mut stream),
927 Some(Ok(SseEvent::Message(MessageEvent {
928 event: "final".into(),
929 data: "".into(),
930 last_event_id: my_id.clone()
931 })))
932 );
933 assert!(stream.is_empty());
934}