gproxy_protocol/codec/
json.rs

1use super::{CodecError, CodecErrorKind, CodecErrorStage, CodecLimits};
2use bytes::Bytes;
3use serde::{Serialize, de::DeserializeOwned};
4use serde_json::Value;
5use std::io::{self, Write};
6
7fn error(kind: CodecErrorKind, stage: CodecErrorStage, message: &'static str) -> CodecError {
8    CodecError::new(kind, stage, message)
9}
10
11fn json_error(source: serde_json::Error, stage: CodecErrorStage) -> CodecError {
12    CodecError::with_source(
13        if source.is_eof() {
14            CodecErrorKind::UnexpectedEof
15        } else {
16            CodecErrorKind::Json
17        },
18        stage,
19        "invalid JSON",
20        source,
21    )
22}
23
24fn size(n: u64, max: u64) -> Result<(), CodecError> {
25    if n > max {
26        Err(error(
27            CodecErrorKind::Limit,
28            CodecErrorStage::Stream,
29            "JSON size limit exceeded",
30        ))
31    } else {
32        Ok(())
33    }
34}
35
36fn whitespace(b: u8) -> bool {
37    matches!(b, b' ' | b'\t' | b'\r' | b'\n')
38}
39
40pub struct JsonDecoder {
41    buffer: Vec<u8>,
42    limits: CodecLimits,
43    failed: bool,
44}
45
46impl JsonDecoder {
47    pub fn new(limits: CodecLimits) -> Self {
48        Self {
49            buffer: Vec::new(),
50            limits,
51            failed: false,
52        }
53    }
54    pub fn push(&mut self, chunk: &[u8]) -> Result<(), CodecError> {
55        if self.failed {
56            return Err(error(
57                CodecErrorKind::Invalid,
58                CodecErrorStage::Stream,
59                "JSON decoder terminated",
60            ));
61        }
62        let n = (self.buffer.len() as u64).saturating_add(chunk.len() as u64);
63        if let Err(e) = size(
64            n,
65            self.limits
66                .max_buffer_bytes
67                .min(self.limits.max_value_bytes)
68                .min(self.limits.max_body_bytes),
69        ) {
70            self.failed = true;
71            return Err(e);
72        }
73        self.buffer.extend_from_slice(chunk);
74        Ok(())
75    }
76    pub fn finish<T: DeserializeOwned>(self) -> Result<T, CodecError> {
77        if self.failed {
78            return Err(error(
79                CodecErrorKind::Invalid,
80                CodecErrorStage::Finish,
81                "JSON decoder terminated",
82            ));
83        }
84        serde_json::from_slice(&self.buffer).map_err(|e| json_error(e, CodecErrorStage::Finish))
85    }
86}
87
88pub fn decode_json<T: DeserializeOwned>(
89    bytes: &[u8],
90    limits: CodecLimits,
91) -> Result<T, CodecError> {
92    let mut d = JsonDecoder::new(limits);
93    d.push(bytes)?;
94    d.finish()
95}
96
97struct LimitedWriter {
98    bytes: Vec<u8>,
99    max: u64,
100    exceeded: bool,
101}
102
103impl Write for LimitedWriter {
104    fn write(&mut self, bytes: &[u8]) -> io::Result<usize> {
105        if (self.bytes.len() as u64).saturating_add(bytes.len() as u64) > self.max {
106            self.exceeded = true;
107            return Err(io::Error::other("JSON value exceeds limit"));
108        }
109        self.bytes.extend_from_slice(bytes);
110        Ok(bytes.len())
111    }
112    fn flush(&mut self) -> io::Result<()> {
113        Ok(())
114    }
115}
116
117pub fn encode_json<T: Serialize>(value: &T, limits: CodecLimits) -> Result<Bytes, CodecError> {
118    let mut writer = LimitedWriter {
119        bytes: Vec::new(),
120        max: limits.max_value_bytes.min(limits.max_body_bytes),
121        exceeded: false,
122    };
123    match serde_json::to_writer(&mut writer, value) {
124        Ok(()) => Ok(Bytes::from(writer.bytes)),
125        Err(_) if writer.exceeded => Err(error(
126            CodecErrorKind::Limit,
127            CodecErrorStage::Start,
128            "JSON value exceeds limit",
129        )),
130        Err(e) => Err(json_error(e, CodecErrorStage::Start)),
131    }
132}
133
134#[derive(Clone, Copy)]
135enum Phase {
136    Start,
137    Value { allow_end: bool },
138    InValue,
139    Separator,
140    Done,
141    Finished,
142    Failed,
143}
144
145/// Parses array elements only once a delimiter establishes the token boundary.
146/// The retained buffer is one value, independent of transport chunk size.
147pub struct JsonArrayDecoder {
148    value: Vec<u8>,
149    limits: CodecLimits,
150    phase: Phase,
151    depth: usize,
152    string: bool,
153    escape: bool,
154    seen: u64,
155}
156
157impl JsonArrayDecoder {
158    pub fn new(limits: CodecLimits) -> Self {
159        Self {
160            value: Vec::new(),
161            limits,
162            phase: Phase::Start,
163            depth: 0,
164            string: false,
165            escape: false,
166            seen: 0,
167        }
168    }
169    pub fn push(&mut self, chunk: &[u8]) -> Result<Vec<Value>, CodecError> {
170        if matches!(self.phase, Phase::Failed | Phase::Finished) {
171            return Err(error(
172                CodecErrorKind::Invalid,
173                CodecErrorStage::Stream,
174                "array decoder terminated",
175            ));
176        }
177        let result = self.push_inner(chunk);
178        if result.is_err() {
179            self.phase = Phase::Failed;
180            self.value.clear();
181        }
182        result
183    }
184    fn push_inner(&mut self, chunk: &[u8]) -> Result<Vec<Value>, CodecError> {
185        self.seen = self.seen.saturating_add(chunk.len() as u64);
186        size(self.seen, self.limits.max_body_bytes)?;
187        let mut values = Vec::new();
188        for &b in chunk {
189            self.byte(b, &mut values)?;
190        }
191        Ok(values)
192    }
193    fn byte(&mut self, b: u8, values: &mut Vec<Value>) -> Result<(), CodecError> {
194        match self.phase {
195            Phase::Start => {
196                if whitespace(b) {
197                    return Ok(());
198                }
199                if b != b'[' {
200                    return Err(error(
201                        CodecErrorKind::Invalid,
202                        CodecErrorStage::Stream,
203                        "array must start with '['",
204                    ));
205                }
206                self.phase = Phase::Value { allow_end: true };
207            }
208            Phase::Value { allow_end } => {
209                if whitespace(b) {
210                    return Ok(());
211                }
212                if b == b']' && allow_end {
213                    self.phase = Phase::Done;
214                    return Ok(());
215                }
216                if b == b']' || b == b',' {
217                    return Err(error(
218                        CodecErrorKind::Invalid,
219                        CodecErrorStage::Stream,
220                        "missing array value or trailing comma",
221                    ));
222                }
223                self.phase = Phase::InValue;
224                self.scan(b)?;
225            }
226            Phase::InValue => {
227                if !self.string && self.depth == 0 && (whitespace(b) || b == b',' || b == b']') {
228                    let value = serde_json::from_slice(&self.value)
229                        .map_err(|e| json_error(e, CodecErrorStage::Stream))?;
230                    self.value.clear();
231                    values.push(value);
232                    self.phase = Phase::Separator;
233                    self.byte(b, values)?;
234                } else {
235                    self.scan(b)?;
236                }
237            }
238            Phase::Separator => {
239                if whitespace(b) {
240                    return Ok(());
241                }
242                match b {
243                    b',' => self.phase = Phase::Value { allow_end: false },
244                    b']' => self.phase = Phase::Done,
245                    _ => {
246                        return Err(error(
247                            CodecErrorKind::Invalid,
248                            CodecErrorStage::Stream,
249                            "array needs comma or ']'",
250                        ));
251                    }
252                }
253            }
254            Phase::Done => {
255                if !whitespace(b) {
256                    return Err(error(
257                        CodecErrorKind::Invalid,
258                        CodecErrorStage::Stream,
259                        "data follows array",
260                    ));
261                }
262            }
263            Phase::Finished | Phase::Failed => {
264                return Err(error(
265                    CodecErrorKind::Invalid,
266                    CodecErrorStage::Stream,
267                    "array decoder terminated",
268                ));
269            }
270        }
271        Ok(())
272    }
273    fn scan(&mut self, b: u8) -> Result<(), CodecError> {
274        size(
275            (self.value.len() as u64).saturating_add(1),
276            self.limits
277                .max_value_bytes
278                .min(self.limits.max_buffer_bytes),
279        )?;
280        self.value.push(b);
281        if self.string {
282            if self.escape {
283                self.escape = false;
284            } else if b == b'\\' {
285                self.escape = true;
286            } else if b == b'"' {
287                self.string = false;
288            }
289        } else {
290            match b {
291                b'"' => self.string = true,
292                b'{' | b'[' => self.depth += 1,
293                b'}' | b']' => {
294                    self.depth = self.depth.checked_sub(1).ok_or_else(|| {
295                        error(
296                            CodecErrorKind::Invalid,
297                            CodecErrorStage::Stream,
298                            "unmatched JSON delimiter",
299                        )
300                    })?;
301                }
302                _ => {}
303            }
304        }
305        Ok(())
306    }
307    pub fn finish(&mut self) -> Result<(), CodecError> {
308        if matches!(self.phase, Phase::Done) {
309            self.phase = Phase::Finished;
310            Ok(())
311        } else {
312            let kind = if matches!(self.phase, Phase::Failed | Phase::Finished) {
313                CodecErrorKind::Invalid
314            } else {
315                CodecErrorKind::UnexpectedEof
316            };
317            self.phase = Phase::Failed;
318            self.value.clear();
319            Err(error(
320                kind,
321                CodecErrorStage::Finish,
322                "unfinished JSON array",
323            ))
324        }
325    }
326}
327
328pub struct JsonArrayEncoder {
329    limits: CodecLimits,
330    started: bool,
331    finished: bool,
332    failed: bool,
333    seen: u64,
334}
335
336impl JsonArrayEncoder {
337    pub fn new(limits: CodecLimits) -> Self {
338        Self {
339            limits,
340            started: false,
341            finished: false,
342            failed: false,
343            seen: 0,
344        }
345    }
346    pub fn push<T: Serialize>(&mut self, value: &T) -> Result<Bytes, CodecError> {
347        if self.finished || self.failed {
348            return Err(error(
349                CodecErrorKind::Invalid,
350                CodecErrorStage::Stream,
351                "array encoder terminated",
352            ));
353        }
354        let result = (|| {
355            let encoded = encode_json(value, self.limits)?;
356            let n = (encoded.len() as u64).saturating_add(1);
357            size(self.seen.saturating_add(n), self.limits.max_body_bytes)?;
358            let mut v = Vec::with_capacity(encoded.len() + 1);
359            v.push(if self.started { b',' } else { b'[' });
360            v.extend_from_slice(&encoded);
361            self.seen += n;
362            self.started = true;
363            Ok(Bytes::from(v))
364        })();
365        if result.is_err() {
366            self.failed = true;
367        }
368        result
369    }
370    pub fn finish(&mut self) -> Result<Bytes, CodecError> {
371        if self.finished || self.failed {
372            return Err(error(
373                CodecErrorKind::Invalid,
374                CodecErrorStage::Finish,
375                "array encoder terminated",
376            ));
377        }
378        self.finished = true;
379        let ending = if self.started {
380            b"]".as_slice()
381        } else {
382            b"[]".as_slice()
383        };
384        size(
385            self.seen.saturating_add(ending.len() as u64),
386            self.limits.max_body_bytes,
387        )?;
388        Ok(Bytes::copy_from_slice(ending))
389    }
390}
391
392/// Accepts LF/CRLF and an optional final complete JSON line without LF.
393pub struct NdjsonDecoder {
394    line: Vec<u8>,
395    limits: CodecLimits,
396    seen: u64,
397    finished: bool,
398    failed: bool,
399}
400
401impl NdjsonDecoder {
402    pub fn new(limits: CodecLimits) -> Self {
403        Self {
404            line: Vec::new(),
405            limits,
406            seen: 0,
407            finished: false,
408            failed: false,
409        }
410    }
411    pub fn push(&mut self, chunk: &[u8]) -> Result<Vec<Value>, CodecError> {
412        if self.finished || self.failed {
413            return Err(error(
414                CodecErrorKind::Invalid,
415                CodecErrorStage::Stream,
416                "NDJSON decoder terminated",
417            ));
418        }
419        let result = (|| {
420            self.seen = self.seen.saturating_add(chunk.len() as u64);
421            size(self.seen, self.limits.max_body_bytes)?;
422            let mut values = Vec::new();
423            for &b in chunk {
424                if b == b'\n' {
425                    self.flush(&mut values)?;
426                } else {
427                    size(
428                        (self.line.len() as u64).saturating_add(1),
429                        self.limits.max_line_bytes.min(self.limits.max_buffer_bytes),
430                    )?;
431                    self.line.push(b);
432                }
433            }
434            Ok(values)
435        })();
436        if result.is_err() {
437            self.failed = true;
438            self.line.clear();
439        }
440        result
441    }
442    fn flush(&mut self, values: &mut Vec<Value>) -> Result<(), CodecError> {
443        if self.line.last() == Some(&b'\r') {
444            self.line.pop();
445        }
446        if !self.line.iter().all(|b| whitespace(*b)) {
447            size(self.line.len() as u64, self.limits.max_value_bytes)?;
448            values.push(
449                serde_json::from_slice(&self.line)
450                    .map_err(|e| json_error(e, CodecErrorStage::Stream))?,
451            );
452        }
453        self.line.clear();
454        Ok(())
455    }
456    pub fn finish(&mut self) -> Result<Vec<Value>, CodecError> {
457        if self.finished || self.failed {
458            return Err(error(
459                CodecErrorKind::Invalid,
460                CodecErrorStage::Finish,
461                "NDJSON decoder terminated",
462            ));
463        }
464        self.finished = true;
465        let mut values = Vec::new();
466        self.flush(&mut values)?;
467        Ok(values)
468    }
469}
470
471pub struct NdjsonEncoder {
472    limits: CodecLimits,
473    seen: u64,
474    failed: bool,
475}
476
477impl NdjsonEncoder {
478    pub fn new(limits: CodecLimits) -> Self {
479        Self {
480            limits,
481            seen: 0,
482            failed: false,
483        }
484    }
485    pub fn encode<T: Serialize>(&mut self, value: &T) -> Result<Bytes, CodecError> {
486        if self.failed {
487            return Err(error(
488                CodecErrorKind::Invalid,
489                CodecErrorStage::Stream,
490                "NDJSON encoder terminated",
491            ));
492        }
493        let result = (|| {
494            let mut v = encode_json(value, self.limits)?.to_vec();
495            size(v.len() as u64, self.limits.max_line_bytes)?;
496            size(
497                self.seen.saturating_add(v.len() as u64).saturating_add(1),
498                self.limits.max_body_bytes,
499            )?;
500            v.push(b'\n');
501            self.seen += v.len() as u64;
502            Ok(Bytes::from(v))
503        })();
504        if result.is_err() {
505            self.failed = true;
506        }
507        result
508    }
509}
gproxy_protocol/codec/json.rs

gproxy_protocol/codec/
json.rs
