netget/server/fluent_forward/
codec.rs

1//! Native bounded MessagePack for Forward's typed JSON record surface.
2use anyhow::{bail, ensure, Context, Result};
3use serde::{Deserialize, Serialize};
4use serde_json::{Map, Number, Value};
5use std::io::{Read, Write};
6pub const MAX_FRAME_BYTES: usize = 256 * 1024;
7pub const MAX_RECORDS: usize = 256;
8pub const MAX_DEPTH: usize = 32;
9pub const MAX_VALUES: usize = 16384;
10pub const DEFAULT_LLM_FALLBACK: bool = false;
11#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
12#[serde(deny_unknown_fields)]
13pub struct Timestamp {
14    pub seconds: u64,
15    #[serde(default, skip_serializing_if = "Option::is_none")]
16    pub nanoseconds: Option<u32>,
17}
18#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
19#[serde(deny_unknown_fields)]
20pub struct Entry {
21    pub timestamp: Timestamp,
22    pub record: Map<String, Value>,
23}
24#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
25#[serde(deny_unknown_fields)]
26pub struct Batch {
27    pub tag: String,
28    pub entries: Vec<Entry>,
29    #[serde(default = "default_mode")]
30    pub mode: String,
31    #[serde(default)]
32    pub require_ack: bool,
33    #[serde(skip)]
34    pub chunk: Option<String>,
35}
36fn default_mode() -> String {
37    "forward".into()
38}
39#[derive(Clone, Debug, PartialEq)]
40pub enum Node {
41    Nil,
42    Bool(bool),
43    Int(i64),
44    Uint(u64),
45    Float(f64),
46    Str(Vec<u8>),
47    Bin(Vec<u8>),
48    Array(Vec<Node>),
49    Map(Vec<(Node, Node)>),
50    Ext(i8, Vec<u8>),
51}
52impl Node {
53    pub fn string(&self) -> Result<&str> {
54        match self {
55            Self::Str(b) => Ok(std::str::from_utf8(b)?),
56            _ => bail!("expected MessagePack string"),
57        }
58    }
59    fn json(self) -> Result<Value> {
60        Ok(match self {
61            Self::Nil => Value::Null,
62            Self::Bool(b) => Value::Bool(b),
63            Self::Int(i) => Value::Number(i.into()),
64            Self::Uint(i) => Value::Number(i.into()),
65            Self::Float(f) => Value::Number(Number::from_f64(f).context("nonfinite record value")?),
66            Self::Str(b) => Value::String(String::from_utf8(b)?),
67            Self::Array(a) => Value::Array(a.into_iter().map(Self::json).collect::<Result<_>>()?),
68            Self::Map(pairs) => {
69                let mut out = Map::new();
70                for (k, v) in pairs {
71                    let k = k.string()?.to_owned();
72                    ensure!(out.insert(k, v.json()?).is_none(), "duplicate record key");
73                }
74                Value::Object(out)
75            }
76            _ => bail!("binary/extension record values are outside the typed JSON surface"),
77        })
78    }
79    fn uint(&self) -> Result<u64> {
80        match self {
81            Self::Uint(n) => Ok(*n),
82            Self::Int(n) => Ok(u64::try_from(*n)?),
83            _ => bail!("expected nonnegative integer"),
84        }
85    }
86}
87#[derive(Debug)]
88enum ParseError {
89    Short,
90    Bad(&'static str),
91}
92fn take<'a>(b: &'a [u8], at: &mut usize, n: usize) -> std::result::Result<&'a [u8], ParseError> {
93    let end = at
94        .checked_add(n)
95        .ok_or(ParseError::Bad("length overflow"))?;
96    let out = b.get(*at..end).ok_or(ParseError::Short)?;
97    *at = end;
98    Ok(out)
99}
100fn uint(b: &[u8], at: &mut usize, n: usize) -> std::result::Result<u64, ParseError> {
101    let mut out = 0;
102    for v in take(b, at, n)? {
103        out = out * 256 + u64::from(*v);
104    }
105    Ok(out)
106}
107fn len(b: &[u8], at: &mut usize, n: usize) -> std::result::Result<usize, ParseError> {
108    let n = uint(b, at, n)? as usize;
109    if n > MAX_FRAME_BYTES {
110        return Err(ParseError::Bad("declared length exceeds frame limit"));
111    }
112    Ok(n)
113}
114fn parse(
115    b: &[u8],
116    at: &mut usize,
117    depth: usize,
118    budget: &mut usize,
119) -> std::result::Result<Node, ParseError> {
120    if depth > MAX_DEPTH {
121        return Err(ParseError::Bad("MessagePack nesting limit"));
122    }
123    if *budget == 0 {
124        return Err(ParseError::Bad("MessagePack value count limit"));
125    }
126    *budget -= 1;
127    let m = take(b, at, 1)?[0];
128    let (kind, count) = match m {
129        0x00..=0x7f => return Ok(Node::Uint(m as u64)),
130        0xe0..=0xff => return Ok(Node::Int(m as i8 as i64)),
131        0xc0 => return Ok(Node::Nil),
132        0xc2..=0xc3 => return Ok(Node::Bool(m == 0xc3)),
133        0xcc..=0xcf => return Ok(Node::Uint(uint(b, at, 1usize << (m - 0xcc))?)),
134        0xd0..=0xd3 => {
135            let size = 1usize << (m - 0xd0);
136            let n = uint(b, at, size)?;
137            return Ok(Node::Int(match size {
138                1 => n as i8 as i64,
139                2 => n as i16 as i64,
140                4 => n as i32 as i64,
141                _ => n as i64,
142            }));
143        }
144        0xca => return Ok(Node::Float(f32::from_bits(uint(b, at, 4)? as u32) as f64)),
145        0xcb => return Ok(Node::Float(f64::from_bits(uint(b, at, 8)?))),
146        0xa0..=0xbf => (0, (m & 31) as usize),
147        0xd9 => (0, len(b, at, 1)?),
148        0xda => (0, len(b, at, 2)?),
149        0xdb => (0, len(b, at, 4)?),
150        0xc4 => (1, len(b, at, 1)?),
151        0xc5 => (1, len(b, at, 2)?),
152        0xc6 => (1, len(b, at, 4)?),
153        0x90..=0x9f => (2, (m & 15) as usize),
154        0xdc => (2, len(b, at, 2)?),
155        0xdd => (2, len(b, at, 4)?),
156        0x80..=0x8f => (3, (m & 15) as usize),
157        0xde => (3, len(b, at, 2)?),
158        0xdf => (3, len(b, at, 4)?),
159        0xd4..=0xd8 => (4, 1usize << (m - 0xd4)),
160        0xc7 => (4, len(b, at, 1)?),
161        0xc8 => (4, len(b, at, 2)?),
162        0xc9 => (4, len(b, at, 4)?),
163        _ => return Err(ParseError::Bad("invalid MessagePack marker")),
164    };
165    Ok(match kind {
166        0 => Node::Str(take(b, at, count)?.to_vec()),
167        1 => Node::Bin(take(b, at, count)?.to_vec()),
168        2 => {
169            if count > *budget {
170                return Err(ParseError::Bad("declared array count limit"));
171            }
172            let mut out = Vec::new();
173            for _ in 0..count {
174                out.push(parse(b, at, depth + 1, budget)?);
175            }
176            Node::Array(out)
177        }
178        3 => {
179            if count > *budget / 2 {
180                return Err(ParseError::Bad("declared map count limit"));
181            }
182            let mut out = Vec::new();
183            for _ in 0..count {
184                let k = parse(b, at, depth + 1, budget)?;
185                let v = parse(b, at, depth + 1, budget)?;
186                out.push((k, v));
187            }
188            Node::Map(out)
189        }
190        _ => {
191            let t = take(b, at, 1)?[0] as i8;
192            Node::Ext(t, take(b, at, count)?.to_vec())
193        }
194    })
195}
196pub fn parse_one(b: &[u8]) -> Result<Option<(Node, usize)>> {
197    let mut at = 0;
198    match parse(b, &mut at, 0, &mut MAX_VALUES.clone()) {
199        Ok(n) => {
200            ensure!(at <= MAX_FRAME_BYTES, "MessagePack frame byte limit");
201            Ok(Some((n, at)))
202        }
203        Err(ParseError::Short) => {
204            ensure!(
205                b.len() <= MAX_FRAME_BYTES,
206                "incomplete MessagePack frame byte limit"
207            );
208            Ok(None)
209        }
210        Err(ParseError::Bad(message)) => bail!(message),
211    }
212}
213#[derive(Default)]
214pub struct Decoder {
215    bytes: Vec<u8>,
216}
217impl Decoder {
218    pub fn feed(&mut self, b: &[u8]) -> Result<()> {
219        ensure!(
220            self.bytes.len() + b.len() <= MAX_FRAME_BYTES + 8192,
221            "buffered byte limit"
222        );
223        self.bytes.extend_from_slice(b);
224        Ok(())
225    }
226    pub fn next_node(&mut self) -> Result<Option<Node>> {
227        if self.bytes.is_empty() {
228            return Ok(None);
229        }
230        if let Some((n, len)) = parse_one(&self.bytes)? {
231            self.bytes.drain(..len);
232            Ok(Some(n))
233        } else {
234            Ok(None)
235        }
236    }
237    pub fn finish(&self) -> Result<()> {
238        ensure!(self.bytes.is_empty(), "EOF in MessagePack frame");
239        Ok(())
240    }
241}
242fn strnode(s: &str) -> Node {
243    Node::Str(s.as_bytes().to_vec())
244}
245fn jsonnode(v: &Value, depth: usize, budget: &mut usize) -> Result<Node> {
246    ensure!(
247        depth <= MAX_DEPTH && *budget > 0,
248        "record nesting/value count limit"
249    );
250    *budget -= 1;
251    Ok(match v {
252        Value::Null => Node::Nil,
253        Value::Bool(b) => Node::Bool(*b),
254        Value::Number(n) => {
255            if let Some(n) = n.as_u64() {
256                Node::Uint(n)
257            } else if let Some(n) = n.as_i64() {
258                Node::Int(n)
259            } else {
260                let n = n.as_f64().context("invalid number")?;
261                ensure!(n.is_finite(), "nonfinite number");
262                Node::Float(n)
263            }
264        }
265        Value::String(s) => strnode(s),
266        Value::Array(a) => Node::Array(
267            a.iter()
268                .map(|v| jsonnode(v, depth + 1, budget))
269                .collect::<Result<_>>()?,
270        ),
271        Value::Object(map) => Node::Map(
272            map.iter()
273                .map(|(k, v)| {
274                    ensure!(*budget > 0, "record value count limit");
275                    *budget -= 1;
276                    Ok((strnode(k), jsonnode(v, depth + 1, budget)?))
277                })
278                .collect::<Result<_>>()?,
279        ),
280    })
281}
282fn header(
283    out: &mut Vec<u8>,
284    small: Option<(u8, usize)>,
285    m8: Option<u8>,
286    m16: u8,
287    m32: u8,
288    len: usize,
289) {
290    if let Some((m, max)) = small {
291        if len < max {
292            out.push(m | len as u8);
293            return;
294        }
295    }
296    if let Some(m) = m8 {
297        if len <= 255 {
298            out.extend_from_slice(&[m, len as u8]);
299            return;
300        }
301    }
302    if len <= 65535 {
303        out.push(m16);
304        out.extend_from_slice(&(len as u16).to_be_bytes());
305    } else {
306        out.push(m32);
307        out.extend_from_slice(&(len as u32).to_be_bytes());
308    }
309}
310fn emit(n: &Node, out: &mut Vec<u8>) -> Result<()> {
311    match n {
312        Node::Nil => out.push(0xc0),
313        Node::Bool(b) => out.push(if *b { 0xc3 } else { 0xc2 }),
314        Node::Uint(n) => {
315            if *n < 128 {
316                out.push(*n as u8)
317            } else {
318                out.push(0xcf);
319                out.extend_from_slice(&n.to_be_bytes())
320            }
321        }
322        Node::Int(n) => {
323            out.push(0xd3);
324            out.extend_from_slice(&n.to_be_bytes())
325        }
326        Node::Float(n) => {
327            ensure!(n.is_finite(), "nonfinite value");
328            out.push(0xcb);
329            out.extend_from_slice(&n.to_bits().to_be_bytes())
330        }
331        Node::Str(b) => {
332            header(out, Some((0xa0, 32)), Some(0xd9), 0xda, 0xdb, b.len());
333            out.extend_from_slice(b)
334        }
335        Node::Bin(b) => {
336            header(out, None, Some(0xc4), 0xc5, 0xc6, b.len());
337            out.extend_from_slice(b)
338        }
339        Node::Array(a) => {
340            header(out, Some((0x90, 16)), None, 0xdc, 0xdd, a.len());
341            for v in a {
342                emit(v, out)?;
343            }
344        }
345        Node::Map(a) => {
346            header(out, Some((0x80, 16)), None, 0xde, 0xdf, a.len());
347            for (k, v) in a {
348                emit(k, out)?;
349                emit(v, out)?;
350            }
351        }
352        Node::Ext(t, b) => {
353            ensure!(b.len() == 8, "only EventTime ext8 supported");
354            out.extend_from_slice(&[0xd7, *t as u8]);
355            out.extend_from_slice(b)
356        }
357    }
358    ensure!(out.len() <= MAX_FRAME_BYTES, "encoded frame byte limit");
359    Ok(())
360}
361pub fn encode_node(n: &Node) -> Result<Vec<u8>> {
362    let mut out = Vec::new();
363    emit(n, &mut out)?;
364    Ok(out)
365}
366fn timestamp(n: &Node) -> Result<Timestamp> {
367    match n {
368        Node::Ext(0, b) if b.len() == 8 => {
369            let seconds = u32::from_be_bytes(b[..4].try_into()?) as u64;
370            let nano = u32::from_be_bytes(b[4..].try_into()?);
371            ensure!(nano < 1_000_000_000, "EventTime nanoseconds out of range");
372            Ok(Timestamp {
373                seconds,
374                nanoseconds: Some(nano),
375            })
376        }
377        _ => Ok(Timestamp {
378            seconds: n.uint()?,
379            nanoseconds: None,
380        }),
381    }
382}
383fn timestampnode(ts: &Timestamp) -> Result<Node> {
384    if let Some(nano) = ts.nanoseconds {
385        ensure!(nano < 1_000_000_000, "nanoseconds must be <1e9");
386        let seconds = u32::try_from(ts.seconds).context("EventTime seconds exceed uint32")?;
387        let mut b = seconds.to_be_bytes().to_vec();
388        b.extend_from_slice(&nano.to_be_bytes());
389        Ok(Node::Ext(0, b))
390    } else {
391        Ok(Node::Uint(ts.seconds))
392    }
393}
394fn entry(n: Node) -> Result<Entry> {
395    let Node::Array(mut a) = n else {
396        bail!("entry must be array")
397    };
398    ensure!(a.len() == 2, "entry must contain timestamp and record");
399    let record = a
400        .pop()
401        .unwrap()
402        .json()?
403        .as_object()
404        .context("record must be map")?
405        .clone();
406    Ok(Entry {
407        timestamp: timestamp(&a[0])?,
408        record,
409    })
410}
411pub fn parse_batch(n: Node) -> Result<Option<Batch>> {
412    if n == Node::Nil {
413        return Ok(None);
414    }
415    let Node::Array(a) = n else {
416        bail!("Forward frame must be an array")
417    };
418    ensure!(a.len() >= 2 && a.len() <= 4, "invalid Forward frame arity");
419    let tag = a[0].string()?.to_owned();
420    ensure!(!tag.is_empty() && tag.len() <= 1024, "tag byte limit");
421    let single = matches!(&a[1], Node::Uint(_) | Node::Int(_) | Node::Ext(_, _));
422    let expected = if single { 3 } else { 2 };
423    ensure!(
424        a.len() == expected || a.len() == expected + 1,
425        "invalid mode arity"
426    );
427    let options = if a.len() == expected + 1 {
428        a.last()
429            .unwrap()
430            .clone()
431            .json()?
432            .as_object()
433            .context("options must be a map")?
434            .clone()
435    } else {
436        Map::new()
437    };
438    let chunk = options
439        .get("chunk")
440        .map(|v| {
441            let s = v.as_str().context("chunk must be string")?;
442            ensure!(
443                !s.is_empty() && s.len() <= 256 && s.bytes().all(|b| b.is_ascii_graphic()),
444                "chunk token byte limit"
445            );
446            Ok::<_, anyhow::Error>(s.to_owned())
447        })
448        .transpose()?;
449    let mut mode = "message";
450    let mut entries = if single {
451        vec![Entry {
452            timestamp: timestamp(&a[1])?,
453            record: a[2]
454                .clone()
455                .json()?
456                .as_object()
457                .context("record must be map")?
458                .clone(),
459        }]
460    } else {
461        match &a[1] {
462            Node::Array(records) => {
463                mode = "forward";
464                ensure!(records.len() <= MAX_RECORDS, "record count limit");
465                records
466                    .clone()
467                    .into_iter()
468                    .map(entry)
469                    .collect::<Result<Vec<_>>>()?
470            }
471            Node::Bin(bytes) | Node::Str(bytes) => {
472                mode = "packed";
473                let mut bytes = bytes.clone();
474                if let Some(compression) = options.get("compressed").filter(|v| *v != "text") {
475                    ensure!(compression == "gzip", "unsupported packed compression");
476                    ensure!(
477                        matches!(&a[1], Node::Bin(_)),
478                        "compressed packed entries must be bin"
479                    );
480                    mode = "compressed_packed";
481                    let mut reader = flate2::bufread::MultiGzDecoder::new(bytes.as_slice());
482                    let mut out = Vec::new();
483                    reader
484                        .by_ref()
485                        .take((MAX_FRAME_BYTES + 1) as u64)
486                        .read_to_end(&mut out)?;
487                    ensure!(
488                        out.len() <= MAX_FRAME_BYTES && reader.get_ref().is_empty(),
489                        "packed decompression byte limit/trailing data"
490                    );
491                    bytes = out;
492                }
493                let mut out = Vec::new();
494                let mut at = 0;
495                let mut budget = MAX_VALUES;
496                while at < bytes.len() {
497                    ensure!(out.len() < MAX_RECORDS, "record count limit");
498                    let node = parse(&bytes, &mut at, 0, &mut budget).map_err(|e| match e {
499                        ParseError::Short => anyhow::anyhow!("incomplete packed entry"),
500                        ParseError::Bad(message) => anyhow::anyhow!(message),
501                    })?;
502                    out.push(entry(node)?);
503                }
504                out
505            }
506            _ => bail!("invalid Forward carrier"),
507        }
508    };
509    ensure!(
510        !entries.is_empty() && entries.len() <= MAX_RECORDS,
511        "batch must have 1..256 entries"
512    );
513    if let Some(size) = options.get("size") {
514        ensure!(
515            size.as_u64() == Some(entries.len() as u64),
516            "size option disagrees with entries"
517        );
518    }
519    if single || mode == "forward" {
520        ensure!(
521            !options.contains_key("compressed"),
522            "compression requires PackedForward"
523        );
524    }
525    // Validate the record JSON bound across all entries and modes.
526    let mut budget = MAX_VALUES;
527    for e in &entries {
528        jsonnode(&Value::Object(e.record.clone()), 0, &mut budget)?;
529    }
530    Ok(Some(Batch {
531        tag,
532        entries: std::mem::take(&mut entries),
533        mode: mode.into(),
534        require_ack: chunk.is_some(),
535        chunk,
536    }))
537}
538pub fn encode_batch(b: &Batch, chunk: Option<&str>) -> Result<Vec<u8>> {
539    ensure!(!b.tag.is_empty() && b.tag.len() <= 1024, "tag byte limit");
540    ensure!(
541        !b.entries.is_empty() && b.entries.len() <= MAX_RECORDS,
542        "batch must contain 1..256 entries"
543    );
544    let mut budget = MAX_VALUES;
545    let entries = b
546        .entries
547        .iter()
548        .map(|e| {
549            Ok(Node::Array(vec![
550                timestampnode(&e.timestamp)?,
551                jsonnode(&Value::Object(e.record.clone()), 0, &mut budget)?,
552            ]))
553        })
554        .collect::<Result<Vec<_>>>()?;
555    let mut options = vec![(strnode("size"), Node::Uint(entries.len() as u64))];
556    if let Some(chunk) = chunk {
557        ensure!(
558            !chunk.is_empty() && chunk.len() <= 256 && chunk.bytes().all(|b| b.is_ascii_graphic()),
559            "chunk token byte limit"
560        );
561        options.push((strnode("chunk"), strnode(chunk)));
562    }
563    let mut root = vec![strnode(&b.tag)];
564    match b.mode.as_str() {
565        "message" => {
566            ensure!(entries.len() == 1, "message mode requires one entry");
567            let Node::Array(a) = &entries[0] else {
568                unreachable!()
569            };
570            root.extend(a.clone());
571        }
572        "forward" => root.push(Node::Array(entries)),
573        "packed" | "compressed_packed" => {
574            let mut bytes = Vec::new();
575            for e in entries {
576                emit(&e, &mut bytes)?;
577            }
578            if b.mode == "compressed_packed" {
579                let mut encoder =
580                    flate2::write::GzEncoder::new(Vec::new(), flate2::Compression::default());
581                encoder.write_all(&bytes)?;
582                bytes = encoder.finish()?;
583                options.push((strnode("compressed"), strnode("gzip")));
584            }
585            root.push(Node::Bin(bytes));
586        }
587        _ => bail!("mode must be message, forward, packed or compressed_packed"),
588    }
589    root.push(Node::Map(options));
590    let bytes = encode_node(&Node::Array(root))?;
591    let (node, _) = parse_one(&bytes)?.context("incomplete encoded frame")?;
592    // Packed records sit inside an opaque bin on the outer frame. Validate the
593    // inner stream too, so depth/value limits reject before any bytes are sent.
594    parse_batch(node)?.context("encoded batch unexpectedly empty")?;
595    Ok(bytes)
596}
597pub fn encode_ack(chunk: &str) -> Result<Vec<u8>> {
598    encode_node(&Node::Map(vec![(strnode("ack"), strnode(chunk))]))
599}
600pub fn parse_ack(n: Node) -> Result<String> {
601    let value = n.json()?;
602    let map = value.as_object().context("ACK must be map")?;
603    ensure!(map.len() == 1, "unexpected ACK fields");
604    Ok(map
605        .get("ack")
606        .and_then(Value::as_str)
607        .context("ACK string missing")?
608        .to_owned())
609}
netget/server/fluent_forward/codec.rs

netget/server/fluent_forward/
codec.rs
