gproxy_protocol/codec/
sse.rs

1use super::{CodecError, CodecErrorKind, CodecErrorStage, CodecLimits};
2use bytes::Bytes;
3
4/// Parsed SSE data event. `id` and `retry` are effective connection settings,
5/// including values supplied in earlier control-only blocks.
6#[derive(Debug, Clone, PartialEq, Eq)]
7pub struct SseEvent {
8    pub event: Option<String>,
9    pub id: Option<String>,
10    pub data: String,
11    pub retry: Option<u64>,
12}
13
14#[derive(Debug, Clone, PartialEq, Eq)]
15pub enum SseFrame {
16    Event(SseEvent),
17    Done,
18}
19
20fn invalid(stage: CodecErrorStage, msg: &'static str) -> CodecError {
21    CodecError::new(CodecErrorKind::Invalid, stage, msg)
22}
23
24fn limit(n: u64, max: u64) -> Result<(), CodecError> {
25    if n > max {
26        Err(CodecError::new(
27            CodecErrorKind::Limit,
28            CodecErrorStage::Stream,
29            "SSE size limit exceeded",
30        ))
31    } else {
32        Ok(())
33    }
34}
35
36fn text(bytes: &[u8]) -> Result<&str, CodecError> {
37    std::str::from_utf8(bytes).map_err(|e| {
38        CodecError::with_source(
39            CodecErrorKind::Utf8,
40            CodecErrorStage::Stream,
41            "invalid SSE UTF-8",
42            e,
43        )
44    })
45}
46
47/// Incremental LF/CRLF/CR framing. Invalid UTF-8 is rejected rather than
48/// replaced. `finish` rejects a pending data event instead of dispatching an
49/// event without its required blank-line delimiter. A final partial comment or
50/// field line is also rejected. No vendor lifecycle is inferred from EOF.
51pub struct SseDecoder {
52    pending: Vec<u8>,
53    data: Vec<u8>,
54    event: Option<String>,
55    id: Option<String>,
56    retry: Option<u64>,
57    has_data: bool,
58    first_line: bool,
59    skip_lf: bool,
60    seen: u64,
61    failed: bool,
62    finished: bool,
63    limits: CodecLimits,
64}
65
66impl SseDecoder {
67    pub fn new(limits: CodecLimits) -> Self {
68        Self {
69            pending: Vec::new(),
70            data: Vec::new(),
71            event: None,
72            id: None,
73            retry: None,
74            has_data: false,
75            first_line: true,
76            skip_lf: false,
77            seen: 0,
78            failed: false,
79            finished: false,
80            limits,
81        }
82    }
83    pub fn last_event_id(&self) -> Option<&str> {
84        self.id.as_deref()
85    }
86    pub fn retry(&self) -> Option<u64> {
87        self.retry
88    }
89    pub fn push(&mut self, chunk: &[u8]) -> Result<Vec<SseFrame>, CodecError> {
90        if self.failed || self.finished {
91            return Err(invalid(CodecErrorStage::Stream, "SSE decoder terminated"));
92        }
93        let result = self.push_inner(chunk);
94        if result.is_err() {
95            self.failed = true;
96            self.pending.clear();
97            self.data.clear();
98        }
99        result
100    }
101    fn push_inner(&mut self, chunk: &[u8]) -> Result<Vec<SseFrame>, CodecError> {
102        self.seen = self.seen.saturating_add(chunk.len() as u64);
103        limit(self.seen, self.limits.max_body_bytes)?;
104        let mut frames = Vec::new();
105        for &b in chunk {
106            if self.skip_lf {
107                self.skip_lf = false;
108                if b == b'\n' {
109                    continue;
110                }
111            }
112            if b == b'\r' || b == b'\n' {
113                let line = std::mem::take(&mut self.pending);
114                self.line(&line, &mut frames)?;
115                self.skip_lf = b == b'\r';
116            } else {
117                limit(
118                    (self.pending.len() as u64).saturating_add(1),
119                    self.limits.max_line_bytes.saturating_add({
120                        let bom = [0xef, 0xbb, 0xbf];
121                        if !self.first_line {
122                            0
123                        } else if self.pending.starts_with(&bom) {
124                            3
125                        } else if self.pending.len() < 3
126                            && bom.starts_with(&self.pending)
127                            && b == bom[self.pending.len()]
128                        {
129                            (self.pending.len() + 1) as u64
130                        } else {
131                            0
132                        }
133                    }),
134                )?;
135                limit(
136                    (self.pending.len() as u64)
137                        .saturating_add(self.data.len() as u64)
138                        .saturating_add(1),
139                    self.limits.max_buffer_bytes,
140                )?;
141                self.pending.push(b);
142            }
143        }
144        Ok(frames)
145    }
146    fn line(&mut self, line: &[u8], frames: &mut Vec<SseFrame>) -> Result<(), CodecError> {
147        let line = if self.first_line {
148            self.first_line = false;
149            line.strip_prefix(&[0xef, 0xbb, 0xbf]).unwrap_or(line)
150        } else {
151            line
152        };
153        limit(line.len() as u64, self.limits.max_line_bytes)?;
154        text(line)?;
155        if line.is_empty() {
156            return self.dispatch(frames);
157        }
158        if line.first() == Some(&b':') {
159            return Ok(());
160        }
161        let (name, value) = match line.iter().position(|b| *b == b':') {
162            Some(i) => (&line[..i], &line[i + 1..]),
163            None => (line, &[][..]),
164        };
165        let value = value.strip_prefix(b" ").unwrap_or(value);
166        match text(name)? {
167            "data" => {
168                let n = (self.data.len() as u64)
169                    .saturating_add(value.len() as u64)
170                    .saturating_add(u64::from(self.has_data));
171                limit(
172                    n,
173                    self.limits
174                        .max_value_bytes
175                        .min(self.limits.max_buffer_bytes),
176                )?;
177                if self.has_data {
178                    self.data.push(b'\n');
179                }
180                self.data.extend_from_slice(value);
181                self.has_data = true;
182            }
183            "event" => self.event = Some(text(value)?.to_owned()),
184            "id" => {
185                let id = text(value)?;
186                if !id.contains('\0') {
187                    self.id = Some(id.to_owned());
188                }
189            }
190            "retry" => {
191                if !value.is_empty()
192                    && value.iter().all(u8::is_ascii_digit)
193                    && let Ok(retry) = text(value)?.parse()
194                {
195                    self.retry = Some(retry);
196                }
197            }
198            _ => {}
199        }
200        Ok(())
201    }
202    fn dispatch(&mut self, frames: &mut Vec<SseFrame>) -> Result<(), CodecError> {
203        let event = self.event.take();
204        if !self.has_data {
205            return Ok(());
206        }
207        let data = text(&self.data)?.to_owned();
208        if data == "[DONE]" {
209            frames.push(SseFrame::Done);
210        } else {
211            frames.push(SseFrame::Event(SseEvent {
212                event,
213                id: self.id.clone(),
214                retry: self.retry,
215                data,
216            }));
217        }
218        self.data.clear();
219        self.has_data = false;
220        Ok(())
221    }
222    pub fn finish(&mut self) -> Result<Vec<SseFrame>, CodecError> {
223        if self.failed || self.finished {
224            return Err(invalid(CodecErrorStage::Finish, "SSE decoder terminated"));
225        }
226        self.finished = true;
227        if !self.pending.is_empty() || self.has_data {
228            self.failed = true;
229            return Err(CodecError::new(
230                CodecErrorKind::UnexpectedEof,
231                CodecErrorStage::Finish,
232                "unterminated SSE data event or line",
233            ));
234        }
235        self.event = None;
236        Ok(Vec::new())
237    }
238}
239
240fn line(
241    out: &mut Vec<u8>,
242    field: &str,
243    value: &str,
244    limits: CodecLimits,
245) -> Result<(), CodecError> {
246    if value.contains(['\r', '\n']) || (field == "id" && value.contains('\0')) {
247        return Err(invalid(CodecErrorStage::Start, "invalid SSE field value"));
248    }
249    let n = (field.len() as u64)
250        .saturating_add(2)
251        .saturating_add(value.len() as u64);
252    limit(n, limits.max_line_bytes)?;
253    limit(
254        (out.len() as u64).saturating_add(n).saturating_add(1),
255        limits.max_body_bytes,
256    )?;
257    out.extend_from_slice(field.as_bytes());
258    out.extend_from_slice(b": ");
259    out.extend_from_slice(value.as_bytes());
260    out.push(b'\n');
261    Ok(())
262}
263
264/// Encodes one event. Use SseEncoder for cumulative limits over a whole body.
265pub fn encode_sse_event(event: &SseEvent, limits: CodecLimits) -> Result<Bytes, CodecError> {
266    limit(event.data.len() as u64, limits.max_value_bytes)?;
267    let mut output = Vec::new();
268    if let Some(v) = &event.event {
269        line(&mut output, "event", v, limits)?;
270    }
271    if let Some(v) = &event.id {
272        line(&mut output, "id", v, limits)?;
273    }
274    if let Some(v) = event.retry {
275        line(&mut output, "retry", &v.to_string(), limits)?;
276    }
277    for data in event.data.split('\n') {
278        line(&mut output, "data", data, limits)?;
279    }
280    limit(
281        (output.len() as u64).saturating_add(1),
282        limits.max_body_bytes,
283    )?;
284    output.push(b'\n');
285    Ok(Bytes::from(output))
286}
287
288/// Literal marker bytes, not a complete streaming encoder; SseEncoder::done
289/// accounts this marker against the enclosing body limit.
290pub fn encode_sse_done() -> Bytes {
291    Bytes::from_static(b"data: [DONE]\n\n")
292}
293
294pub struct SseEncoder {
295    limits: CodecLimits,
296    seen: u64,
297    terminated: bool,
298}
299
300impl SseEncoder {
301    pub fn new(limits: CodecLimits) -> Self {
302        Self {
303            limits,
304            seen: 0,
305            terminated: false,
306        }
307    }
308    pub fn event(&mut self, event: &SseEvent) -> Result<Bytes, CodecError> {
309        if self.terminated {
310            return Err(invalid(CodecErrorStage::Stream, "SSE encoder terminated"));
311        }
312        let result = (|| {
313            let bytes = encode_sse_event(event, self.limits)?;
314            self.account(bytes)
315        })();
316        if result.is_err() {
317            self.terminated = true;
318        }
319        result
320    }
321    fn account(&mut self, bytes: Bytes) -> Result<Bytes, CodecError> {
322        limit(
323            self.seen.saturating_add(bytes.len() as u64),
324            self.limits.max_body_bytes,
325        )?;
326        self.seen += bytes.len() as u64;
327        Ok(bytes)
328    }
329    pub fn done(&mut self) -> Result<Bytes, CodecError> {
330        if self.terminated {
331            return Err(invalid(CodecErrorStage::Stream, "SSE encoder terminated"));
332        }
333        self.terminated = true;
334        let bytes = encode_sse_event(
335            &SseEvent {
336                event: None,
337                id: None,
338                retry: None,
339                data: "[DONE]".into(),
340            },
341            self.limits,
342        )?;
343        self.account(bytes)
344    }
345}
gproxy_protocol/codec/sse.rs

gproxy_protocol/codec/
sse.rs
