netget/server/thrift/
codec.rs

1//! Thrift's binary (strict, and the old non-strict form on read) and compact protocols over a
2//! byte buffer, to and from a typed value tree. Parsing reports `Incomplete` when the buffer ends
3//! early, so the unframed transport can read more.
4use anyhow::{anyhow, bail, ensure, Result};
5
6pub const MAX_MESSAGE: usize = 16 * 1024 * 1024;
7const MAX_DEPTH: usize = 64;
8
9pub const T_STOP: u8 = 0;
10pub const T_BOOL: u8 = 2;
11pub const T_BYTE: u8 = 3;
12pub const T_DOUBLE: u8 = 4;
13pub const T_I16: u8 = 6;
14pub const T_I32: u8 = 8;
15pub const T_I64: u8 = 10;
16pub const T_STRING: u8 = 11;
17pub const T_STRUCT: u8 = 12;
18pub const T_MAP: u8 = 13;
19pub const T_SET: u8 = 14;
20pub const T_LIST: u8 = 15;
21pub const T_UUID: u8 = 16;
22
23pub const CALL: u8 = 1;
24pub const REPLY: u8 = 2;
25pub const EXCEPTION: u8 = 3;
26pub const ONEWAY: u8 = 4;
27
28#[derive(Debug, Clone, PartialEq)]
29pub enum Tv {
30    Bool(bool),
31    Byte(i8),
32    I16(i16),
33    I32(i32),
34    I64(i64),
35    Double(f64),
36    Bin(Vec<u8>),
37    Uuid([u8; 16]),
38    Struct(Vec<(i16, Tv)>),
39    List(u8, Vec<Tv>),
40    Set(u8, Vec<Tv>),
41    Map(u8, u8, Vec<(Tv, Tv)>),
42}
43
44impl Tv {
45    pub fn ttype(&self) -> u8 {
46        match self {
47            Tv::Bool(_) => T_BOOL,
48            Tv::Byte(_) => T_BYTE,
49            Tv::I16(_) => T_I16,
50            Tv::I32(_) => T_I32,
51            Tv::I64(_) => T_I64,
52            Tv::Double(_) => T_DOUBLE,
53            Tv::Bin(_) => T_STRING,
54            Tv::Uuid(_) => T_UUID,
55            Tv::Struct(_) => T_STRUCT,
56            Tv::List(..) => T_LIST,
57            Tv::Set(..) => T_SET,
58            Tv::Map(..) => T_MAP,
59        }
60    }
61    pub fn field(&self, id: i16) -> Option<&Tv> {
62        match self {
63            Tv::Struct(f) => f.iter().find(|(i, _)| *i == id).map(|(_, v)| v),
64            _ => None,
65        }
66    }
67}
68
69#[derive(Debug, Clone, Copy, PartialEq)]
70pub enum Protocol {
71    Binary,
72    Compact,
73}
74
75#[derive(Debug, Clone, PartialEq)]
76pub struct Message {
77    pub name: String,
78    pub kind: u8,
79    pub seqid: i32,
80    pub body: Tv,
81}
82
83/// The buffer ended before the message did.
84#[derive(Debug)]
85pub struct Incomplete;
86impl std::fmt::Display for Incomplete {
87    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
88        f.write_str("incomplete message")
89    }
90}
91impl std::error::Error for Incomplete {}
92
93struct R<'a> {
94    b: &'a [u8],
95    pos: usize,
96    items: usize,
97}
98
99impl R<'_> {
100    fn take(&mut self, n: usize) -> Result<&[u8]> {
101        ensure!(n <= MAX_MESSAGE, "length {n} is over the message bound");
102        let end = self
103            .pos
104            .checked_add(n)
105            .ok_or_else(|| anyhow!("length overflow"))?;
106        let v = self
107            .b
108            .get(self.pos..end)
109            .ok_or_else(|| anyhow!(Incomplete))?;
110        self.pos = end;
111        Ok(v)
112    }
113    fn u8(&mut self) -> Result<u8> {
114        Ok(self.take(1)?[0])
115    }
116    fn i16(&mut self) -> Result<i16> {
117        Ok(i16::from_be_bytes(self.take(2)?.try_into()?))
118    }
119    fn i32(&mut self) -> Result<i32> {
120        Ok(i32::from_be_bytes(self.take(4)?.try_into()?))
121    }
122    fn i64(&mut self) -> Result<i64> {
123        Ok(i64::from_be_bytes(self.take(8)?.try_into()?))
124    }
125    fn count(&mut self, n: i64, min_each: usize) -> Result<usize> {
126        ensure!(n >= 0, "negative container size");
127        let n = n as usize;
128        self.items += n;
129        ensure!(self.items <= 4_000_000, "too many container elements");
130        // Every element takes at least `min_each` bytes, so a size beyond what remains cannot be honest.
131        if n.saturating_mul(min_each) > self.b.len() - self.pos {
132            if self.b.len() >= MAX_MESSAGE {
133                bail!("container size {n} cannot fit in the message");
134            }
135            return Err(anyhow!(Incomplete));
136        }
137        Ok(n)
138    }
139    fn varint(&mut self) -> Result<u64> {
140        let mut v = 0u64;
141        for shift in (0..70).step_by(7) {
142            let b = self.u8()?;
143            ensure!(shift < 64, "varint too long");
144            v |= ((b & 0x7f) as u64) << shift;
145            if b & 0x80 == 0 {
146                return Ok(v);
147            }
148        }
149        bail!("varint too long")
150    }
151    fn zigzag(&mut self) -> Result<i64> {
152        let v = self.varint()?;
153        Ok(((v >> 1) as i64) ^ -((v & 1) as i64))
154    }
155
156    // Binary protocol.
157    fn bin_value(&mut self, t: u8, depth: usize) -> Result<Tv> {
158        ensure!(depth <= MAX_DEPTH, "values nest too deeply");
159        Ok(match t {
160            T_BOOL => Tv::Bool(self.u8()? != 0),
161            T_BYTE => Tv::Byte(self.u8()? as i8),
162            T_I16 => Tv::I16(self.i16()?),
163            T_I32 => Tv::I32(self.i32()?),
164            T_I64 => Tv::I64(self.i64()?),
165            T_DOUBLE => Tv::Double(f64::from_bits(self.i64()? as u64)),
166            T_STRING => {
167                let n = self.i32()?;
168                ensure!(n >= 0, "negative string length");
169                Tv::Bin(self.take(n as usize)?.to_vec())
170            }
171            T_UUID => Tv::Uuid(self.take(16)?.try_into()?),
172            T_STRUCT => {
173                let mut fields = Vec::new();
174                loop {
175                    let ft = self.u8()?;
176                    if ft == T_STOP {
177                        break;
178                    }
179                    let id = self.i16()?;
180                    fields.push((id, self.bin_value(ft, depth + 1)?));
181                }
182                Tv::Struct(fields)
183            }
184            T_LIST | T_SET => {
185                let et = self.u8()?;
186                let n = self.i32()?;
187                let n = self.count(n as i64, 1)?;
188                let mut items = Vec::with_capacity(n.min(1024));
189                for _ in 0..n {
190                    items.push(self.bin_value(et, depth + 1)?);
191                }
192                if t == T_LIST {
193                    Tv::List(et, items)
194                } else {
195                    Tv::Set(et, items)
196                }
197            }
198            T_MAP => {
199                let (kt, vt) = (self.u8()?, self.u8()?);
200                let n = self.i32()?;
201                let n = self.count(n as i64, 2)?;
202                let mut items = Vec::with_capacity(n.min(1024));
203                for _ in 0..n {
204                    let k = self.bin_value(kt, depth + 1)?;
205                    let v = self.bin_value(vt, depth + 1)?;
206                    items.push((k, v));
207                }
208                Tv::Map(kt, vt, items)
209            }
210            other => bail!("unknown Thrift type {other}"),
211        })
212    }
213
214    // Compact protocol.
215    fn compact_type(t: u8) -> Result<u8> {
216        Ok(match t {
217            1 | 2 => T_BOOL,
218            3 => T_BYTE,
219            4 => T_I16,
220            5 => T_I32,
221            6 => T_I64,
222            7 => T_DOUBLE,
223            8 => T_STRING,
224            9 => T_LIST,
225            10 => T_SET,
226            11 => T_MAP,
227            12 => T_STRUCT,
228            13 => T_UUID,
229            other => bail!("unknown compact type {other}"),
230        })
231    }
232    fn compact_value(&mut self, t: u8, depth: usize) -> Result<Tv> {
233        ensure!(depth <= MAX_DEPTH, "values nest too deeply");
234        Ok(match t {
235            T_BOOL => Tv::Bool(self.u8()? == 1),
236            T_BYTE => Tv::Byte(self.u8()? as i8),
237            T_I16 => Tv::I16(i16::try_from(self.zigzag()?)?),
238            T_I32 => Tv::I32(i32::try_from(self.zigzag()?)?),
239            T_I64 => Tv::I64(self.zigzag()?),
240            T_DOUBLE => Tv::Double(f64::from_le_bytes(self.take(8)?.try_into()?)),
241            T_STRING => {
242                let n = self.varint()?;
243                Tv::Bin(self.take(usize::try_from(n)?)?.to_vec())
244            }
245            T_UUID => Tv::Uuid(self.take(16)?.try_into()?),
246            T_STRUCT => {
247                let mut fields = Vec::new();
248                let mut last: i16 = 0;
249                loop {
250                    let h = self.u8()?;
251                    if h == T_STOP {
252                        break;
253                    }
254                    let delta = (h >> 4) as i16;
255                    let id = if delta == 0 {
256                        i16::try_from(self.zigzag()?)?
257                    } else {
258                        last.checked_add(delta)
259                            .ok_or_else(|| anyhow!("field id overflow"))?
260                    };
261                    last = id;
262                    let ct = h & 0x0f;
263                    let v = if ct == 1 || ct == 2 {
264                        Tv::Bool(ct == 1)
265                    } else {
266                        self.compact_value(Self::compact_type(ct)?, depth + 1)?
267                    };
268                    fields.push((id, v));
269                }
270                Tv::Struct(fields)
271            }
272            T_LIST | T_SET => {
273                let h = self.u8()?;
274                let et = Self::compact_type(h & 0x0f)?;
275                let n = if h >> 4 == 0x0f {
276                    self.varint()? as i64
277                } else {
278                    (h >> 4) as i64
279                };
280                let n = self.count(n, 1)?;
281                let mut items = Vec::with_capacity(n.min(1024));
282                for _ in 0..n {
283                    items.push(self.compact_value(et, depth + 1)?);
284                }
285                if t == T_LIST {
286                    Tv::List(et, items)
287                } else {
288                    Tv::Set(et, items)
289                }
290            }
291            T_MAP => {
292                let n = self.varint()? as i64;
293                if n == 0 {
294                    return Ok(Tv::Map(T_STRING, T_STRING, vec![]));
295                }
296                let kv = self.u8()?;
297                let (kt, vt) = (Self::compact_type(kv >> 4)?, Self::compact_type(kv & 0x0f)?);
298                let n = self.count(n, 2)?;
299                let mut items = Vec::with_capacity(n.min(1024));
300                for _ in 0..n {
301                    let k = self.compact_value(kt, depth + 1)?;
302                    let v = self.compact_value(vt, depth + 1)?;
303                    items.push((k, v));
304                }
305                Tv::Map(kt, vt, items)
306            }
307            other => bail!("unknown Thrift type {other}"),
308        })
309    }
310}
311
312/// Parse one message from the front of `b`: the message, its protocol and the bytes it used.
313pub fn decode(b: &[u8]) -> Result<(Message, Protocol, usize)> {
314    let mut r = R {
315        b,
316        pos: 0,
317        items: 0,
318    };
319    let first = *b.first().ok_or_else(|| anyhow!(Incomplete))?;
320    let (name, kind, seqid, protocol) = if first == 0x82 {
321        r.u8()?;
322        let vt = r.u8()?;
323        ensure!(
324            vt & 0x1f == 1,
325            "unsupported compact protocol version {}",
326            vt & 0x1f
327        );
328        let seqid = r.varint()? as u32 as i32;
329        let n = r.varint()?;
330        let name = String::from_utf8(r.take(usize::try_from(n)?)?.to_vec())?;
331        (name, vt >> 5, seqid, Protocol::Compact)
332    } else if first & 0x80 != 0 {
333        let v = r.i32()? as u32;
334        ensure!(
335            v & 0xffff_0000 == 0x8001_0000,
336            "bad binary protocol version 0x{v:08x}"
337        );
338        let n = r.i32()?;
339        ensure!(n >= 0, "negative method name length");
340        let name = String::from_utf8(r.take(n as usize)?.to_vec())?;
341        let seqid = r.i32()?;
342        (name, (v & 0xff) as u8, seqid, Protocol::Binary)
343    } else {
344        // The old non-strict binary form: the name first.
345        let n = r.i32()?;
346        ensure!((0..=4096).contains(&n), "not a Thrift message");
347        let name = String::from_utf8(r.take(n as usize)?.to_vec())?;
348        let kind = r.u8()?;
349        let seqid = r.i32()?;
350        (name, kind, seqid, Protocol::Binary)
351    };
352    ensure!(
353        name.len() <= 256 && (CALL..=ONEWAY).contains(&kind),
354        "malformed message header"
355    );
356    let body = match protocol {
357        Protocol::Binary => r.bin_value(T_STRUCT, 0)?,
358        Protocol::Compact => r.compact_value(T_STRUCT, 0)?,
359    };
360    Ok((
361        Message {
362            name,
363            kind,
364            seqid,
365            body,
366        },
367        protocol,
368        r.pos,
369    ))
370}
371
372pub fn is_incomplete(e: &anyhow::Error) -> bool {
373    e.is::<Incomplete>()
374}
375
376struct W {
377    out: Vec<u8>,
378}
379
380impl W {
381    fn varint(&mut self, mut v: u64) {
382        loop {
383            if v < 0x80 {
384                self.out.push(v as u8);
385                return;
386            }
387            self.out.push((v as u8) | 0x80);
388            v >>= 7;
389        }
390    }
391    fn zigzag(&mut self, v: i64) {
392        self.varint(((v << 1) ^ (v >> 63)) as u64);
393    }
394    fn bin(&mut self, v: &Tv) {
395        match v {
396            Tv::Bool(b) => self.out.push(*b as u8),
397            Tv::Byte(n) => self.out.push(*n as u8),
398            Tv::I16(n) => self.out.extend(n.to_be_bytes()),
399            Tv::I32(n) => self.out.extend(n.to_be_bytes()),
400            Tv::I64(n) => self.out.extend(n.to_be_bytes()),
401            Tv::Double(f) => self.out.extend(f.to_bits().to_be_bytes()),
402            Tv::Bin(b) => {
403                self.out.extend((b.len() as i32).to_be_bytes());
404                self.out.extend(b);
405            }
406            Tv::Uuid(u) => self.out.extend(u),
407            Tv::Struct(fields) => {
408                for (id, f) in fields {
409                    self.out.push(f.ttype());
410                    self.out.extend(id.to_be_bytes());
411                    self.bin(f);
412                }
413                self.out.push(T_STOP);
414            }
415            Tv::List(et, items) | Tv::Set(et, items) => {
416                self.out.push(*et);
417                self.out.extend((items.len() as i32).to_be_bytes());
418                for i in items {
419                    self.bin(i);
420                }
421            }
422            Tv::Map(kt, vt, items) => {
423                self.out.extend([*kt, *vt]);
424                self.out.extend((items.len() as i32).to_be_bytes());
425                for (k, v) in items {
426                    self.bin(k);
427                    self.bin(v);
428                }
429            }
430        }
431    }
432    fn compact_code(t: u8) -> u8 {
433        match t {
434            T_BOOL => 1,
435            T_BYTE => 3,
436            T_I16 => 4,
437            T_I32 => 5,
438            T_I64 => 6,
439            T_DOUBLE => 7,
440            T_STRING => 8,
441            T_LIST => 9,
442            T_SET => 10,
443            T_MAP => 11,
444            T_STRUCT => 12,
445            T_UUID => 13,
446            _ => 0,
447        }
448    }
449    fn compact(&mut self, v: &Tv) {
450        match v {
451            Tv::Bool(b) => self.out.push(if *b { 1 } else { 2 }),
452            Tv::Byte(n) => self.out.push(*n as u8),
453            Tv::I16(n) => self.zigzag(*n as i64),
454            Tv::I32(n) => self.zigzag(*n as i64),
455            Tv::I64(n) => self.zigzag(*n),
456            Tv::Double(f) => self.out.extend(f.to_le_bytes()),
457            Tv::Bin(b) => {
458                self.varint(b.len() as u64);
459                self.out.extend(b);
460            }
461            Tv::Uuid(u) => self.out.extend(u),
462            Tv::Struct(fields) => {
463                let mut last: i16 = 0;
464                for (id, f) in fields {
465                    let code = match f {
466                        Tv::Bool(b) => {
467                            if *b {
468                                1
469                            } else {
470                                2
471                            }
472                        }
473                        other => Self::compact_code(other.ttype()),
474                    };
475                    let delta = id.wrapping_sub(last);
476                    if (1..=15).contains(&delta) {
477                        self.out.push(((delta as u8) << 4) | code);
478                    } else {
479                        self.out.push(code);
480                        self.zigzag(*id as i64);
481                    }
482                    last = *id;
483                    if !matches!(f, Tv::Bool(_)) {
484                        self.compact(f);
485                    }
486                }
487                self.out.push(T_STOP);
488            }
489            Tv::List(et, items) | Tv::Set(et, items) => {
490                let code = Self::compact_code(*et);
491                if items.len() < 15 {
492                    self.out.push(((items.len() as u8) << 4) | code);
493                } else {
494                    self.out.push(0xf0 | code);
495                    self.varint(items.len() as u64);
496                }
497                for i in items {
498                    // Booleans in containers are one byte: 1 true, 2 false, as the Java and
499                    // Python implementations write them (readers test for 1).
500                    if let Tv::Bool(b) = i {
501                        self.out.push(if *b { 1 } else { 2 });
502                    } else {
503                        self.compact(i);
504                    }
505                }
506            }
507            Tv::Map(kt, vt, items) => {
508                self.varint(items.len() as u64);
509                if !items.is_empty() {
510                    self.out
511                        .push((Self::compact_code(*kt) << 4) | Self::compact_code(*vt));
512                }
513                for (k, v) in items {
514                    for x in [k, v] {
515                        if let Tv::Bool(b) = x {
516                            self.out.push(if *b { 1 } else { 2 });
517                        } else {
518                            self.compact(x);
519                        }
520                    }
521                }
522            }
523        }
524    }
525}
526
527pub fn encode(m: &Message, protocol: Protocol) -> Vec<u8> {
528    let mut w = W { out: Vec::new() };
529    match protocol {
530        Protocol::Binary => {
531            w.out.extend((0x8001_0000u32 | m.kind as u32).to_be_bytes());
532            w.out.extend((m.name.len() as i32).to_be_bytes());
533            w.out.extend(m.name.as_bytes());
534            w.out.extend(m.seqid.to_be_bytes());
535            w.bin(&m.body);
536        }
537        Protocol::Compact => {
538            w.out.push(0x82);
539            w.out.push((m.kind << 5) | 1);
540            w.varint(m.seqid as u32 as u64);
541            w.varint(m.name.len() as u64);
542            w.out.extend(m.name.as_bytes());
543            w.compact(&m.body);
544        }
545    }
546    w.out
547}
548
549/// A TApplicationException body: 1: message, 2: type.
550pub fn application_exception(message: &str, kind: i32) -> Tv {
551    Tv::Struct(vec![
552        (1, Tv::Bin(message.as_bytes().to_vec())),
553        (2, Tv::I32(kind)),
554    ])
555}
556pub const UNKNOWN_METHOD: i32 = 1;
557pub const INTERNAL_ERROR: i32 = 6;
558pub const PROTOCOL_ERROR: i32 = 7;
netget/server/thrift/codec.rs

netget/server/thrift/
codec.rs
