buffa/
unknown_fields.rs

1//! Unknown field preservation for round-trip fidelity.
2
3use alloc::vec::Vec;
4use bytes::BufMut;
5
6/// A collection of unknown fields encountered during decoding.
7///
8/// When a message is decoded with a schema that doesn't include all fields
9/// present on the wire, the unknown fields are stored here so they can be
10/// re-encoded without data loss.
11#[derive(Clone, Debug, Default, PartialEq, Hash)]
12pub struct UnknownFields {
13    fields: Vec<UnknownField>,
14}
15
16impl UnknownFields {
17    /// Creates an empty set of unknown fields.
18    pub fn new() -> Self {
19        Self::default()
20    }
21
22    /// Returns `true` if no unknown fields have been recorded.
23    pub fn is_empty(&self) -> bool {
24        self.fields.is_empty()
25    }
26
27    /// Returns the number of unknown fields recorded.
28    pub fn len(&self) -> usize {
29        self.fields.len()
30    }
31
32    /// Appends an unknown field.
33    pub fn push(&mut self, field: UnknownField) {
34        self.fields.push(field);
35    }
36
37    /// Returns an iterator over the recorded unknown fields.
38    pub fn iter(&self) -> core::slice::Iter<'_, UnknownField> {
39        self.fields.iter()
40    }
41
42    /// Removes all recorded unknown fields.
43    pub fn clear(&mut self) {
44        self.fields.clear();
45    }
46
47    /// Retain only the fields for which the predicate returns `true`.
48    ///
49    /// Used by [`ExtensionSet::set_extension`](crate::ExtensionSet::set_extension)
50    /// and [`clear_extension`](crate::ExtensionSet::clear_extension) to remove
51    /// prior occurrences at a given field number.
52    pub fn retain(&mut self, f: impl FnMut(&UnknownField) -> bool) {
53        self.fields.retain(f);
54    }
55
56    /// Compute the encoded size of all unknown fields.
57    pub fn encoded_len(&self) -> usize {
58        self.fields.iter().map(|f| f.encoded_len()).sum()
59    }
60
61    /// Re-encode all unknown fields to `buf` in their original wire format.
62    pub fn write_to(&self, buf: &mut impl BufMut) {
63        for field in &self.fields {
64            field.write_to(buf);
65        }
66    }
67
68    /// Decode a concatenation of wire-format fields into [`UnknownFields`].
69    ///
70    /// Reads tag/data pairs until `data` is exhausted. Each field is decoded
71    /// via [`decode_unknown_field`](crate::encoding::decode_unknown_field) with
72    /// the full [`RECURSION_LIMIT`](crate::message::RECURSION_LIMIT) budget.
73    ///
74    /// Used by [`GroupCodec`](crate::extension::codecs::GroupCodec) to turn a
75    /// message's encoded bytes back into the inner-field representation that
76    /// [`UnknownFieldData::Group`] wraps.
77    ///
78    /// # Errors
79    ///
80    /// Returns [`DecodeError`](crate::DecodeError) if `data` contains a
81    /// malformed tag, truncated field, or exceeds the recursion limit.
82    pub fn decode_from_slice(mut data: &[u8]) -> Result<Self, crate::DecodeError> {
83        use crate::encoding::{decode_unknown_field, Tag};
84        use crate::message::RECURSION_LIMIT;
85        let mut out = Self::new();
86        while !data.is_empty() {
87            let tag = Tag::decode(&mut data)?;
88            out.push(decode_unknown_field(tag, &mut data, RECURSION_LIMIT)?);
89        }
90        Ok(out)
91    }
92}
93
94impl<'a> IntoIterator for &'a UnknownFields {
95    type Item = &'a UnknownField;
96    type IntoIter = core::slice::Iter<'a, UnknownField>;
97
98    fn into_iter(self) -> Self::IntoIter {
99        self.fields.iter()
100    }
101}
102
103impl IntoIterator for UnknownFields {
104    type Item = UnknownField;
105    type IntoIter = alloc::vec::IntoIter<UnknownField>;
106
107    fn into_iter(self) -> Self::IntoIter {
108        self.fields.into_iter()
109    }
110}
111
112/// A single unknown field (field number + wire data).
113#[derive(Clone, Debug, PartialEq, Hash)]
114pub struct UnknownField {
115    pub number: u32,
116    pub data: UnknownFieldData,
117}
118
119impl UnknownField {
120    /// Compute the encoded size of this field (tag + data).
121    pub fn encoded_len(&self) -> usize {
122        let tag_len =
123            crate::encoding::varint_len(((self.number as u64) << 3) | self.data.wire_type_value());
124        tag_len + self.data.encoded_len(self.number)
125    }
126
127    /// Re-encode this field (tag + data) to `buf` in its original wire format.
128    pub fn write_to(&self, buf: &mut impl BufMut) {
129        use crate::encoding::encode_varint;
130        let tag_value = ((self.number as u64) << 3) | self.data.wire_type_value();
131        encode_varint(tag_value, buf);
132        match &self.data {
133            UnknownFieldData::Varint(v) => encode_varint(*v, buf),
134            UnknownFieldData::Fixed64(v) => buf.put_u64_le(*v),
135            UnknownFieldData::Fixed32(v) => buf.put_u32_le(*v),
136            UnknownFieldData::LengthDelimited(data) => {
137                encode_varint(data.len() as u64, buf);
138                buf.put_slice(data);
139            }
140            UnknownFieldData::Group(fields) => {
141                fields.write_to(buf);
142                // End-group tag (wire type 4, same field number).
143                encode_varint(((self.number as u64) << 3) | 4, buf);
144            }
145        }
146    }
147}
148
149/// The wire data for an unknown field.
150#[derive(Clone, Debug, PartialEq, Hash)]
151pub enum UnknownFieldData {
152    Varint(u64),
153    Fixed64(u64),
154    Fixed32(u32),
155    LengthDelimited(Vec<u8>),
156    Group(UnknownFields),
157}
158
159impl UnknownFieldData {
160    fn wire_type_value(&self) -> u64 {
161        match self {
162            Self::Varint(_) => 0,
163            Self::Fixed64(_) => 1,
164            Self::LengthDelimited(_) => 2,
165            Self::Group(_) => 3,
166            Self::Fixed32(_) => 5,
167        }
168    }
169
170    fn encoded_len(&self, field_number: u32) -> usize {
171        match self {
172            Self::Varint(v) => crate::encoding::varint_len(*v),
173            Self::Fixed64(_) => 8,
174            Self::Fixed32(_) => 4,
175            Self::LengthDelimited(data) => {
176                crate::encoding::varint_len(data.len() as u64) + data.len()
177            }
178            Self::Group(fields) => {
179                // Group content + end-group tag (wire type 4, same field number).
180                let end_tag_len = crate::encoding::varint_len((field_number as u64) << 3 | 4);
181                fields.encoded_len() + end_tag_len
182            }
183        }
184    }
185}
186
187#[cfg(feature = "arbitrary")]
188impl<'a> arbitrary::Arbitrary<'a> for UnknownFields {
189    fn arbitrary(_u: &mut arbitrary::Unstructured<'a>) -> arbitrary::Result<Self> {
190        Ok(UnknownFields::new())
191    }
192}
193
194#[cfg(test)]
195mod tests {
196    use super::*;
197    use crate::encoding::MAX_FIELD_NUMBER;
198
199    fn varint_field(number: u32, value: u64) -> UnknownField {
200        UnknownField {
201            number,
202            data: UnknownFieldData::Varint(value),
203        }
204    }
205
206    fn fixed32_field(number: u32, value: u32) -> UnknownField {
207        UnknownField {
208            number,
209            data: UnknownFieldData::Fixed32(value),
210        }
211    }
212
213    fn fixed64_field(number: u32, value: u64) -> UnknownField {
214        UnknownField {
215            number,
216            data: UnknownFieldData::Fixed64(value),
217        }
218    }
219
220    fn ld_field(number: u32, data: Vec<u8>) -> UnknownField {
221        UnknownField {
222            number,
223            data: UnknownFieldData::LengthDelimited(data),
224        }
225    }
226
227    fn group_field(number: u32, fields: UnknownFields) -> UnknownField {
228        UnknownField {
229            number,
230            data: UnknownFieldData::Group(fields),
231        }
232    }
233
234    #[test]
235    fn test_empty_encoded_len() {
236        assert_eq!(UnknownFields::new().encoded_len(), 0);
237    }
238
239    #[test]
240    fn test_varint_encoded_len() {
241        // field 1, value 0: tag=(1<<3)|0=8 (1 byte) + varint(0) (1 byte) = 2
242        assert_eq!(varint_field(1, 0).encoded_len(), 2);
243        // field 1, value 128: tag (1 byte) + varint(128) (2 bytes) = 3
244        assert_eq!(varint_field(1, 128).encoded_len(), 3);
245    }
246
247    #[test]
248    fn test_fixed32_encoded_len() {
249        // field 1: tag=(1<<3)|5=13 (1 byte) + 4 bytes = 5
250        assert_eq!(fixed32_field(1, 0xDEAD_BEEF).encoded_len(), 5);
251    }
252
253    #[test]
254    fn test_fixed64_encoded_len() {
255        // field 1: tag=(1<<3)|1=9 (1 byte) + 8 bytes = 9
256        assert_eq!(fixed64_field(1, 0xDEAD_BEEF_CAFE_BABE).encoded_len(), 9);
257    }
258
259    #[test]
260    fn test_length_delimited_encoded_len() {
261        // field 1, 3-byte data: tag (1) + len_varint(3)=1 + 3 = 5
262        assert_eq!(ld_field(1, vec![0xAA, 0xBB, 0xCC]).encoded_len(), 5);
263        // field 1, empty: tag (1) + len_varint(0)=1 + 0 = 2
264        assert_eq!(ld_field(1, vec![]).encoded_len(), 2);
265        // field 1, 128-byte data: tag (1) + len_varint(128)=2 + 128 = 131
266        assert_eq!(ld_field(1, vec![0u8; 128]).encoded_len(), 131);
267    }
268
269    #[test]
270    fn test_group_encoded_len() {
271        // field 1, empty group: start_tag (1) + end_tag (1) = 2
272        assert_eq!(group_field(1, UnknownFields::new()).encoded_len(), 2);
273        // field 1, group containing one varint(0) field (2 bytes):
274        //   start_tag (1) + inner (2) + end_tag (1) = 4
275        let mut inner = UnknownFields::new();
276        inner.push(varint_field(1, 0));
277        assert_eq!(group_field(1, inner).encoded_len(), 4);
278    }
279
280    #[test]
281    fn test_large_field_number_two_byte_tag() {
282        // field 16: tag=(16<<3)|0=128, varint_len(128)=2; varint(0)=1 → total 3
283        assert_eq!(varint_field(16, 0).encoded_len(), 3);
284    }
285
286    #[test]
287    fn test_multiple_fields_sum() {
288        let mut fields = UnknownFields::new();
289        fields.push(varint_field(1, 0)); // tag(1) + varint(0)(1) = 2
290        fields.push(fixed32_field(2, 0)); // tag=(2<<3)|5=21(1) + 4 = 5
291        assert_eq!(fields.encoded_len(), 7);
292    }
293
294    #[test]
295    fn test_into_iter_by_ref() {
296        let mut fields = UnknownFields::new();
297        fields.push(varint_field(1, 0));
298        fields.push(varint_field(2, 0));
299        let numbers: Vec<u32> = (&fields).into_iter().map(|f| f.number).collect();
300        assert_eq!(numbers, vec![1, 2]);
301        // `for f in &fields` syntax works:
302        let mut count = 0;
303        for _ in &fields {
304            count += 1;
305        }
306        assert_eq!(count, 2);
307    }
308
309    // ── encoded_len vs write_to consistency ──────────────────────────────
310    //
311    // Size/write divergence corrupts downstream encoding (compute_size
312    // feeds the length prefix of length-delimited messages). These tests
313    // build fields spanning the full variant matrix + boundary field
314    // numbers, and assert write_to produces exactly encoded_len bytes.
315
316    fn assert_len_matches_write(f: UnknownField) {
317        let mut fields = UnknownFields::new();
318        let claimed = f.encoded_len();
319        fields.push(f);
320        let mut buf = alloc::vec::Vec::new();
321        fields.write_to(&mut buf);
322        assert_eq!(
323            buf.len(),
324            claimed,
325            "encoded_len={claimed} but wrote {} bytes for {fields:?}",
326            buf.len()
327        );
328    }
329
330    #[test]
331    fn test_encoded_len_matches_write_varint() {
332        for &num in &[1, 15, 16, 2047, 2048, MAX_FIELD_NUMBER] {
333            for &val in &[0, 1, 127, 128, u32::MAX as u64, u64::MAX] {
334                assert_len_matches_write(varint_field(num, val));
335            }
336        }
337    }
338
339    #[test]
340    fn test_encoded_len_matches_write_fixed() {
341        for &num in &[1, 16, MAX_FIELD_NUMBER] {
342            assert_len_matches_write(fixed32_field(num, 0));
343            assert_len_matches_write(fixed32_field(num, u32::MAX));
344            assert_len_matches_write(fixed64_field(num, 0));
345            assert_len_matches_write(fixed64_field(num, u64::MAX));
346        }
347    }
348
349    #[test]
350    fn test_encoded_len_matches_write_length_delimited() {
351        for &num in &[1, 16, MAX_FIELD_NUMBER] {
352            assert_len_matches_write(ld_field(num, alloc::vec![]));
353            assert_len_matches_write(ld_field(num, alloc::vec![0xAB]));
354            assert_len_matches_write(ld_field(num, alloc::vec![0; 127]));
355            assert_len_matches_write(ld_field(num, alloc::vec![0; 128]));
356        }
357    }
358
359    #[test]
360    fn test_encoded_len_matches_write_group() {
361        for &num in &[1, 16, MAX_FIELD_NUMBER] {
362            // Empty group.
363            assert_len_matches_write(group_field(num, UnknownFields::new()));
364            // Group with mixed children.
365            let mut inner = UnknownFields::new();
366            inner.push(varint_field(2, 42));
367            inner.push(ld_field(3, alloc::vec![1, 2, 3]));
368            assert_len_matches_write(group_field(num, inner.clone()));
369            // Nested group.
370            let mut inner2 = UnknownFields::new();
371            inner2.push(group_field(4, inner));
372            assert_len_matches_write(group_field(num, inner2));
373        }
374    }
375
376    #[test]
377    fn test_encoded_len_matches_write_multi_field() {
378        // All variants in one UnknownFields.
379        let mut fields = UnknownFields::new();
380        fields.push(varint_field(1, 42));
381        fields.push(fixed32_field(2, 0xDEADBEEF));
382        fields.push(fixed64_field(3, u64::MAX));
383        fields.push(ld_field(4, alloc::vec![1, 2, 3, 4, 5]));
384        let mut nested = UnknownFields::new();
385        nested.push(varint_field(10, 7));
386        fields.push(group_field(5, nested));
387        let claimed = fields.encoded_len();
388        let mut buf = alloc::vec::Vec::new();
389        fields.write_to(&mut buf);
390        assert_eq!(buf.len(), claimed);
391    }
392
393    // ── decode_from_slice ────────────────────────────────────────────────
394
395    #[test]
396    fn test_decode_from_slice_empty() {
397        let out = UnknownFields::decode_from_slice(&[]).unwrap();
398        assert!(out.is_empty());
399    }
400
401    #[test]
402    fn test_decode_from_slice_roundtrip() {
403        // Build a multi-variant field set, write, decode back, compare.
404        let mut orig = UnknownFields::new();
405        orig.push(varint_field(1, 42));
406        orig.push(fixed32_field(2, 0xDEAD_BEEF));
407        orig.push(fixed64_field(3, u64::MAX));
408        orig.push(ld_field(4, vec![1, 2, 3]));
409        let mut nested = UnknownFields::new();
410        nested.push(varint_field(10, 7));
411        orig.push(group_field(5, nested));
412
413        let mut buf = Vec::new();
414        orig.write_to(&mut buf);
415        let decoded = UnknownFields::decode_from_slice(&buf).unwrap();
416        assert_eq!(decoded, orig);
417    }
418
419    #[test]
420    fn test_decode_from_slice_truncated() {
421        // Tag 0x08 (field 1, varint) with no value — truncated.
422        assert!(UnknownFields::decode_from_slice(&[0x08]).is_err());
423    }
424}
buffa/unknown_fields.rs

buffa/
unknown_fields.rs
