wechat_protocol/
protobuf.rs

1//! Bounded numbered protobuf primitives. Unknown fields and full u64 values survive.
2use crate::MAX_MESSAGE_BYTES;
3use anyhow::{bail, ensure, Result};
4
5#[derive(Clone, PartialEq, Eq)]
6pub enum Value {
7    Varint(u64),
8    Fixed64([u8; 8]),
9    Bytes(Vec<u8>),
10    Fixed32([u8; 4]),
11}
12#[derive(Clone, PartialEq, Eq)]
13pub struct Field {
14    pub number: u32,
15    pub value: Value,
16}
17
18pub fn varint(mut value: u64) -> Vec<u8> {
19    let mut out = Vec::new();
20    loop {
21        let b = (value & 127) as u8;
22        value >>= 7;
23        out.push(b | if value > 0 { 128 } else { 0 });
24        if value == 0 {
25            break;
26        }
27    }
28    out
29}
30pub fn read_varint(input: &[u8], offset: &mut usize) -> Result<u64> {
31    let mut out = 0u64;
32    for i in 0..10 {
33        let b = *input
34            .get(*offset)
35            .ok_or_else(|| anyhow::anyhow!("Truncated varint"))?;
36        *offset += 1;
37        ensure!(i != 9 || b <= 1, "Varint overflow");
38        out |= ((b & 127) as u64) << (i * 7);
39        if b & 128 == 0 {
40            return Ok(out);
41        }
42    }
43    bail!("Unterminated varint")
44}
45fn tag(number: u32, wire: u64) -> Vec<u8> {
46    assert!(number > 0 && number < 1 << 29);
47    varint((number as u64) * 8 + wire)
48}
49pub fn uint(number: u32, value: u64) -> Vec<u8> {
50    [tag(number, 0), varint(value)].concat()
51}
52pub fn int32(number: u32, value: i32) -> Vec<u8> {
53    uint(number, value as i64 as u64)
54}
55pub fn bytes(number: u32, value: &[u8]) -> Vec<u8> {
56    [tag(number, 2), varint(value.len() as u64), value.to_vec()].concat()
57}
58pub fn message(parts: &[Vec<u8>]) -> Result<Vec<u8>> {
59    let size = parts.iter().try_fold(0usize, |n, p| {
60        n.checked_add(p.len())
61            .ok_or_else(|| anyhow::anyhow!("Message length overflow"))
62    })?;
63    ensure!(size <= MAX_MESSAGE_BYTES, "Message too large");
64    Ok(parts.concat())
65}
66pub fn buffer(value: &[u8]) -> Vec<u8> {
67    [uint(1, value.len() as u64), bytes(2, value)].concat()
68}
69pub fn decode(input: &[u8]) -> Result<Vec<Field>> {
70    ensure!(input.len() <= MAX_MESSAGE_BYTES, "Message too large");
71    let mut offset = 0;
72    let mut out = Vec::new();
73    while offset < input.len() {
74        ensure!(out.len() < 65536, "Too many fields");
75        let key = read_varint(input, &mut offset)?;
76        let number = key >> 3;
77        ensure!(number > 0 && number < 1 << 29, "Invalid field tag");
78        let value = match key & 7 {
79            0 => Value::Varint(read_varint(input, &mut offset)?),
80            wire @ (1 | 2 | 5) => {
81                let n = match wire {
82                    1 => 8,
83                    5 => 4,
84                    _ => usize::try_from(read_varint(input, &mut offset)?)?,
85                };
86                ensure!(
87                    n <= MAX_MESSAGE_BYTES && n <= input.len() - offset,
88                    "Truncated field"
89                );
90                let value = &input[offset..offset + n];
91                offset += n;
92                match wire {
93                    1 => Value::Fixed64(value.try_into()?),
94                    5 => Value::Fixed32(value.try_into()?),
95                    _ => Value::Bytes(value.to_vec()),
96                }
97            }
98            _ => bail!("Unsupported wire type"),
99        };
100        out.push(Field {
101            number: number as u32,
102            value,
103        });
104    }
105    Ok(out)
106}
107pub fn optional(fields: &[Field], number: u32) -> Result<Option<&Value>> {
108    let mut found = fields.iter().filter(|f| f.number == number);
109    let first = found.next();
110    ensure!(found.next().is_none(), "Ambiguous field {number}");
111    Ok(first.map(|f| &f.value))
112}
113pub fn scalar(fields: &[Field], number: u32) -> Result<u64> {
114    match optional(fields, number)? {
115        Some(Value::Varint(v)) => Ok(*v),
116        _ => bail!("Missing scalar field {number}"),
117    }
118}
119pub fn data(fields: &[Field], number: u32) -> Result<&[u8]> {
120    match optional(fields, number)? {
121        Some(Value::Bytes(v)) => Ok(v),
122        _ => bail!("Missing bytes field {number}"),
123    }
124}
125pub fn optional_data(fields: &[Field], number: u32) -> Result<Option<&[u8]>> {
126    match optional(fields, number)? {
127        Some(Value::Bytes(v)) => Ok(Some(v)),
128        None => Ok(None),
129        _ => bail!("Invalid bytes field {number}"),
130    }
131}
132pub fn unbuffer(input: &[u8]) -> Result<Vec<u8>> {
133    let f = decode(input)?;
134    let raw = optional_data(&f, 2)?.unwrap_or_default();
135    ensure!(scalar(&f, 1)? == raw.len() as u64, "Buffer length mismatch");
136    Ok(raw.to_vec())
137}
138pub fn base_code(fields: &[Field]) -> Result<i32> {
139    let base = decode(data(fields, 1)?)?;
140    Ok(scalar(&base, 1)? as u32 as i32)
141}
142
143#[cfg(test)]
144mod tests {
145    use super::*;
146    #[test]
147    fn u64_and_unknowns() {
148        let input = message(&[uint(1, u64::MAX), bytes(900, b"unknown"), uint(1, 7)]).unwrap();
149        let f = decode(&input).unwrap();
150        assert_eq!(f.len(), 3);
151        assert!(scalar(&f, 1).is_err());
152        assert!(matches!(f[0].value, Value::Varint(u64::MAX)));
153    }
154    #[test]
155    fn malformed_rejected() {
156        for x in [
157            vec![0],
158            vec![8, 255],
159            vec![10, 3, 0],
160            vec![8, 255, 255, 255, 255, 255, 255, 255, 255, 255, 2],
161            vec![11],
162        ] {
163            assert!(decode(&x).is_err());
164        }
165    }
166}
wechat_protocol/protobuf.rs

wechat_protocol/
protobuf.rs
