hidane_core/export/
wire.rs

1//! The protobuf wire format, by hand: the export's messages are proto1 (App Engine's
2//! `EntityProto`), whose groups and field order prost does not reproduce.
3
4/// Appends fields to a buffer, in the order they are written.
5#[derive(Default)]
6pub struct Writer {
7    buf: Vec<u8>,
8}
9
10const VARINT: u32 = 0;
11const FIXED64: u32 = 1;
12const LEN: u32 = 2;
13const START_GROUP: u32 = 3;
14const END_GROUP: u32 = 4;
15const FIXED32: u32 = 5;
16
17impl Writer {
18    pub fn into_bytes(self) -> Vec<u8> {
19        self.buf
20    }
21
22    fn varint(&mut self, mut value: u64) {
23        while value >= 0x80 {
24            #[allow(clippy::cast_possible_truncation)]
25            self.buf.push((value as u8) | 0x80);
26            value >>= 7;
27        }
28        #[allow(clippy::cast_possible_truncation)]
29        self.buf.push(value as u8);
30    }
31
32    fn tag(&mut self, field: u32, wire_type: u32) {
33        self.varint(u64::from(field << 3 | wire_type));
34    }
35
36    pub fn uint(&mut self, field: u32, value: u64) {
37        self.tag(field, VARINT);
38        self.varint(value);
39    }
40
41    /// An `int64`: negative values take ten bytes, as in protobuf.
42    pub fn int(&mut self, field: u32, value: i64) {
43        #[allow(clippy::cast_sign_loss)]
44        self.uint(field, value as u64);
45    }
46
47    pub fn double(&mut self, field: u32, value: f64) {
48        self.tag(field, FIXED64);
49        self.buf.extend_from_slice(&value.to_bits().to_le_bytes());
50    }
51
52    pub fn bytes(&mut self, field: u32, value: &[u8]) {
53        self.tag(field, LEN);
54        self.varint(value.len() as u64);
55        self.buf.extend_from_slice(value);
56    }
57
58    /// A length-delimited message.
59    pub fn message(&mut self, field: u32, write: impl FnOnce(&mut Self)) {
60        let mut inner = Self::default();
61        write(&mut inner);
62        self.bytes(field, &inner.buf);
63    }
64
65    /// A proto1 group: the fields between a start and an end tag.
66    pub fn group(&mut self, field: u32, write: impl FnOnce(&mut Self)) {
67        self.tag(field, START_GROUP);
68        write(self);
69        self.tag(field, END_GROUP);
70    }
71}
72
73/// A field's value as read.
74#[derive(Debug, Clone, Copy)]
75pub enum Field<'a> {
76    Varint(u64),
77    Fixed64(u64),
78    /// Skipped: no field of the export is a `fixed32`.
79    Fixed32,
80    /// A length-delimited field: bytes, a string or a message.
81    Bytes(&'a [u8]),
82    /// The fields of a group.
83    Group(&'a [u8]),
84}
85
86impl<'a> Field<'a> {
87    pub fn as_varint(self) -> Option<u64> {
88        match self {
89            Self::Varint(v) => Some(v),
90            _ => None,
91        }
92    }
93
94    pub fn as_fixed64(self) -> Option<u64> {
95        match self {
96            Self::Fixed64(v) => Some(v),
97            _ => None,
98        }
99    }
100
101    pub fn as_bytes(self) -> Option<&'a [u8]> {
102        match self {
103            Self::Bytes(b) => Some(b),
104            _ => None,
105        }
106    }
107
108    /// The fields of a message or of a group.
109    pub fn fields(self) -> Option<Reader<'a>> {
110        match self {
111            Self::Bytes(b) | Self::Group(b) => Some(Reader::new(b)),
112            _ => None,
113        }
114    }
115}
116
117/// Malformed protobuf.
118#[derive(Debug, Clone, Copy, PartialEq, Eq)]
119pub struct Malformed;
120
121/// Reads fields one by one, in the order they were written.
122#[derive(Debug, Clone)]
123pub struct Reader<'a> {
124    data: &'a [u8],
125    pos: usize,
126}
127
128impl<'a> Reader<'a> {
129    pub fn new(data: &'a [u8]) -> Self {
130        Self { data, pos: 0 }
131    }
132
133    fn varint(&mut self) -> Result<u64, Malformed> {
134        let mut value = 0u64;
135        for shift in (0..64).step_by(7) {
136            let byte = *self.data.get(self.pos).ok_or(Malformed)?;
137            self.pos += 1;
138            value |= u64::from(byte & 0x7f) << shift;
139            if byte < 0x80 {
140                return Ok(value);
141            }
142        }
143        Err(Malformed)
144    }
145
146    fn take(&mut self, len: usize) -> Result<&'a [u8], Malformed> {
147        let end = self.pos.checked_add(len).ok_or(Malformed)?;
148        let slice = self.data.get(self.pos..end).ok_or(Malformed)?;
149        self.pos = end;
150        Ok(slice)
151    }
152
153    /// The next field number and wire type, or `None` at the end.
154    fn tag(&mut self) -> Result<Option<(u32, u32)>, Malformed> {
155        if self.pos == self.data.len() {
156            return Ok(None);
157        }
158        let tag = u32::try_from(self.varint()?).map_err(|_| Malformed)?;
159        if tag >> 3 == 0 {
160            return Err(Malformed);
161        }
162        Ok(Some((tag >> 3, tag & 7)))
163    }
164
165    fn value(&mut self, field: u32, wire_type: u32) -> Result<Field<'a>, Malformed> {
166        Ok(match wire_type {
167            VARINT => Field::Varint(self.varint()?),
168            FIXED64 => Field::Fixed64(u64::from_le_bytes(
169                self.take(8)?.try_into().map_err(|_| Malformed)?,
170            )),
171            FIXED32 => {
172                self.take(4)?;
173                Field::Fixed32
174            }
175            LEN => {
176                let len = usize::try_from(self.varint()?).map_err(|_| Malformed)?;
177                Field::Bytes(self.take(len)?)
178            }
179            START_GROUP => {
180                let start = self.pos;
181                loop {
182                    let end = self.pos;
183                    let (inner, inner_type) = self.tag()?.ok_or(Malformed)?;
184                    if inner_type == END_GROUP {
185                        if inner != field {
186                            return Err(Malformed);
187                        }
188                        return Ok(Field::Group(&self.data[start..end]));
189                    }
190                    self.value(inner, inner_type)?;
191                }
192            }
193            _ => return Err(Malformed),
194        })
195    }
196}
197
198impl<'a> Iterator for Reader<'a> {
199    type Item = Result<(u32, Field<'a>), Malformed>;
200
201    fn next(&mut self) -> Option<Self::Item> {
202        let result = (|| {
203            let Some((field, wire_type)) = self.tag()? else {
204                return Ok(None);
205            };
206            Ok(Some((field, self.value(field, wire_type)?)))
207        })();
208        match result {
209            Ok(Some(item)) => Some(Ok(item)),
210            Ok(None) => None,
211            Err(err) => {
212                // Stop after an error.
213                self.pos = self.data.len();
214                Some(Err(err))
215            }
216        }
217    }
218}
219
220#[cfg(test)]
221mod tests {
222    use super::*;
223
224    #[test]
225    fn groups_nest_and_round_trip() {
226        let mut w = Writer::default();
227        w.uint(1, 300);
228        w.group(12, |g| {
229            g.bytes(13, b"app");
230            g.group(14, |e| {
231                e.bytes(15, b"kind");
232                e.int(16, -3);
233            });
234        });
235        w.double(4, 1.5);
236        let bytes = w.into_bytes();
237
238        let fields: Vec<_> = Reader::new(&bytes).collect::<Result<_, _>>().unwrap();
239        assert_eq!(fields.len(), 3);
240        assert_eq!(fields[0].1.as_varint(), Some(300));
241        let reference: Vec<_> = fields[1]
242            .1
243            .fields()
244            .unwrap()
245            .collect::<Result<_, _>>()
246            .unwrap();
247        assert_eq!(reference[0].1.as_bytes(), Some(&b"app"[..]));
248        let element: Vec<_> = reference[1]
249            .1
250            .fields()
251            .unwrap()
252            .collect::<Result<_, _>>()
253            .unwrap();
254        #[allow(clippy::cast_possible_wrap)]
255        let id = element[1].1.as_varint().unwrap() as i64;
256        assert_eq!(id, -3);
257        assert_eq!(fields[2].1.as_fixed64(), Some(1.5f64.to_bits()));
258    }
259
260    #[test]
261    fn truncated_input_is_malformed() {
262        let mut w = Writer::default();
263        w.bytes(3, b"hello");
264        let bytes = w.into_bytes();
265        assert!(Reader::new(&bytes[..4]).any(|f| f.is_err()));
266        // A group without its end tag.
267        assert!(Reader::new(&[0x63, 0x08, 0x01]).any(|f| f.is_err()));
268    }
269}
hidane_core/export/wire.rs

hidane_core/export/
wire.rs
