gproxy_transform/transform/common/
sse.rs1#[derive(Debug, Clone, PartialEq, Eq)]
6pub struct SseFrame {
7 pub event: Option<String>,
8 pub data: String,
9}
10
11impl SseFrame {
12 pub fn data(data: impl Into<String>) -> Self {
13 Self {
14 event: None,
15 data: data.into(),
16 }
17 }
18
19 pub fn event(event: impl Into<String>, data: impl Into<String>) -> Self {
20 Self {
21 event: Some(event.into()),
22 data: data.into(),
23 }
24 }
25
26 pub fn encode(&self) -> String {
27 let mut encoded = String::new();
28 if let Some(event) = &self.event {
29 encoded.push_str("event: ");
30 encoded.push_str(event);
31 encoded.push('\n');
32 }
33 for line in self.data.lines() {
34 encoded.push_str("data: ");
35 encoded.push_str(line);
36 encoded.push('\n');
37 }
38 encoded.push('\n');
39 encoded
40 }
41}
42
43#[derive(Debug, Default)]
48pub struct SseDecoder {
49 buf: String,
50 utf8: super::utf8::Utf8StreamDecoder,
51}
52
53impl SseDecoder {
54 pub fn new() -> Self {
55 Self::default()
56 }
57
58 pub fn push(&mut self, chunk: &[u8]) -> Vec<SseFrame> {
62 self.utf8.decode_into(chunk, &mut self.buf);
63 if self.buf.contains('\r') {
64 self.buf = self.buf.replace("\r\n", "\n");
65 }
66 let mut frames = Vec::new();
67 while let Some(pos) = self.buf.find("\n\n") {
68 let raw: String = self.buf.drain(..pos + 2).collect();
69 if let Some(frame) = parse_frame(&raw) {
70 frames.push(frame);
71 }
72 }
73 frames
74 }
75
76 pub fn finish(&mut self) -> Option<SseFrame> {
79 self.utf8.flush(&mut self.buf);
80 let raw = std::mem::take(&mut self.buf);
81 parse_frame(&raw)
82 }
83}
84
85fn parse_frame(raw: &str) -> Option<SseFrame> {
86 let mut event = None;
87 let mut data_lines: Vec<&str> = Vec::new();
88 for line in raw.lines() {
89 if let Some(rest) = line.strip_prefix("event:") {
90 event = Some(rest.trim_start().to_owned());
91 } else if let Some(rest) = line.strip_prefix("data:") {
92 data_lines.push(rest.strip_prefix(' ').unwrap_or(rest));
93 }
94 }
95 if data_lines.is_empty() {
96 return None;
97 }
98 Some(SseFrame {
99 event,
100 data: data_lines.join("\n"),
101 })
102}
103
104#[cfg(test)]
105mod tests {
106 use super::*;
107
108 #[test]
109 fn frame_split_across_chunks() {
110 let mut d = SseDecoder::new();
111 assert!(d.push(b"event: ping\nda").is_empty());
112 let frames = d.push(b"ta: {\"a\":1}\n\n: comment\ndata: x");
113 assert_eq!(frames.len(), 1);
114 assert_eq!(frames[0].event.as_deref(), Some("ping"));
115 assert_eq!(frames[0].data, "{\"a\":1}");
116 assert_eq!(d.finish().unwrap().data, "x");
118 }
119
120 #[test]
121 fn crlf_and_multiline_data() {
122 let mut d = SseDecoder::new();
123 let frames = d.push(b"data: l1\r\ndata: l2\r\n\r\n");
124 assert_eq!(frames.len(), 1);
125 assert_eq!(frames[0].data, "l1\nl2");
126 }
127
128 #[test]
129 fn multibyte_char_split_across_chunks() {
130 let mut d = SseDecoder::new();
132 let bytes = "data: 汉字\n\n".as_bytes();
133 assert!(d.push(&bytes[..7]).is_empty()); let frames = d.push(&bytes[7..]);
135 assert_eq!(frames.len(), 1);
136 assert_eq!(frames[0].data, "汉字");
137 }
138}