krabka_protocol/
tagged_fields.rs

1//! KIP-482 flexible-version tagged fields.
2
3use bytes::{Buf, BufMut, Bytes, BytesMut};
4
5use crate::{
6    ProtocolError,
7    primitives::varint::{get_uvarint, put_uvarint, uvarint_len},
8};
9
10/// An unknown tagged field that was preserved verbatim during decode.
11#[derive(Debug, Clone, PartialEq, Eq)]
12pub struct UnknownTaggedField {
13    pub tag: u32,
14    pub bytes: Bytes,
15}
16
17/// A collection of tagged fields that the schema does not declare. Generated
18/// message types hold a `Vec<UnknownTaggedField>`, sorted by tag, so that
19/// values round-trip without information loss.
20#[derive(Debug, Clone, Default, PartialEq, Eq)]
21pub struct UnknownTaggedFields(pub Vec<UnknownTaggedField>);
22
23impl UnknownTaggedFields {
24    #[must_use]
25    pub fn is_empty(&self) -> bool {
26        self.0.is_empty()
27    }
28    #[must_use]
29    pub fn len(&self) -> usize {
30        self.0.len()
31    }
32}
33
34/// Reads the tagged-fields trailer at the current position of `buf`.
35///
36/// This function calls `known` for each entry whose tag is in the schema.
37/// `known` must consume the field's payload, a `size`-byte slice, and return
38/// `Ok(true)` if it recognized the tag. Every entry for which `known` returns
39/// `Ok(false)` goes into the unknown vec instead.
40/// # Errors
41/// Returns the underlying protocol error when input is truncated, contains an invalid length or tag, or cannot be encoded for the selected version.
42pub fn read_tagged_fields<B, F>(
43    buf: &mut B,
44    mut known: F,
45) -> Result<UnknownTaggedFields, ProtocolError>
46where
47    B: Buf,
48    F: FnMut(u32, &mut &[u8]) -> Result<bool, ProtocolError>,
49{
50    let count = get_uvarint(buf)? as usize;
51    let mut unknown = Vec::new();
52    let mut last_tag: Option<u32> = None;
53    for _ in 0..count {
54        let tag = get_uvarint(buf)?;
55        if let Some(prev) = last_tag
56            && tag <= prev
57        {
58            return Err(ProtocolError::InvalidValue(
59                "tagged fields not strictly ascending",
60            ));
61        }
62        last_tag = Some(tag);
63        let size = get_uvarint(buf)? as usize;
64        if buf.remaining() < size {
65            return Err(ProtocolError::UnexpectedEof {
66                needed: size - buf.remaining(),
67            });
68        }
69        // Copy the payload so we can hand a slice to the closure or store it.
70        let mut payload = vec![0u8; size];
71        buf.copy_to_slice(&mut payload);
72        let mut slice = &payload[..];
73        if !known(tag, &mut slice)? {
74            unknown.push(UnknownTaggedField {
75                tag,
76                bytes: Bytes::from(payload),
77            });
78        } else if !slice.is_empty() {
79            return Err(ProtocolError::InvalidValue(
80                "tagged field decoder did not consume all bytes",
81            ));
82        }
83    }
84    Ok(UnknownTaggedFields(unknown))
85}
86
87/// Helper for generated code that emits tagged fields. Call
88/// `WriteTaggedFields::new()`, then `add` for each known tag and payload
89/// encoder, then `write` to flush the fields and merge them with `unknown`.
90pub struct WriteTaggedFields {
91    entries: Vec<(u32, Bytes)>,
92}
93
94impl Default for WriteTaggedFields {
95    fn default() -> Self {
96        Self::new()
97    }
98}
99
100impl WriteTaggedFields {
101    #[must_use]
102    pub fn new() -> Self {
103        Self {
104            entries: Vec::new(),
105        }
106    }
107
108    pub fn add(&mut self, tag: u32, payload: Bytes) {
109        self.entries.push((tag, payload));
110    }
111
112    /// # Panics
113    /// Panics if a value previously validated by the protocol type no longer satisfies its encoded-length or field-range invariant.
114    pub fn write<B: BufMut>(mut self, buf: &mut B, unknown: &UnknownTaggedFields) {
115        for u in &unknown.0 {
116            self.entries.push((u.tag, u.bytes.clone()));
117        }
118        self.entries.sort_by_key(|(t, _)| *t);
119        put_uvarint(
120            buf,
121            u32::try_from(self.entries.len()).expect("too many tagged fields"),
122        );
123        for (tag, payload) in self.entries {
124            put_uvarint(buf, tag);
125            put_uvarint(
126                buf,
127                u32::try_from(payload.len()).expect("tagged field too large"),
128            );
129            buf.put_slice(&payload);
130        }
131    }
132}
133
134/// Predicted length of the tagged-fields trailer.
135#[must_use]
136/// # Panics
137/// Panics if a value previously validated by the protocol type no longer satisfies its encoded-length or field-range invariant.
138pub fn tagged_fields_len(known: &[(u32, usize)], unknown: &UnknownTaggedFields) -> usize {
139    let total = known.len() + unknown.0.len();
140    let mut n = uvarint_len(u32::try_from(total).unwrap());
141    for (tag, len) in known {
142        n += uvarint_len(*tag) + uvarint_len(u32::try_from(*len).unwrap()) + *len;
143    }
144    for u in &unknown.0 {
145        n +=
146            uvarint_len(u.tag) + uvarint_len(u32::try_from(u.bytes.len()).unwrap()) + u.bytes.len();
147    }
148    n
149}
150
151/// Encodes a value into a freshly allocated `Bytes`. This materializes a
152/// tagged-field payload before the outer trailer is sized.
153///
154/// The write closure can return an error, for example from a nested struct
155/// encode. The error propagates as a panic, because a tagged-field encoding
156/// failure shows a bug in the emitter's `encoded_len` prediction.
157/// # Panics
158/// Panics if a value previously validated by the protocol type no longer satisfies its encoded-length or field-range invariant.
159pub fn encode_to_bytes<F>(predicted_len: usize, write: F) -> Bytes
160where
161    F: FnOnce(&mut BytesMut) -> Result<(), crate::ProtocolError>,
162{
163    let mut buf = BytesMut::with_capacity(predicted_len);
164    write(&mut buf).expect("tagged-field encode failed: emitter bug");
165    debug_assert_eq!(buf.len(), predicted_len, "encoded_len lied");
166    buf.freeze()
167}
168
169#[cfg(test)]
170mod tests {
171
172    use super::*;
173
174    #[test]
175    fn empty_tagged_fields() {
176        let buf = [0x00u8];
177        let mut cur = &buf[..];
178        let unknown = read_tagged_fields(&mut cur, |_, _| Ok(false)).unwrap();
179        assert2::assert!(unknown.is_empty());
180        assert2::assert!(cur.is_empty());
181    }
182
183    #[test]
184    fn unknown_tagged_fields_preserved() {
185        // count=2, tag=5, size=3, payload=[10,20,30], tag=8, size=1, payload=[40]
186        let buf = [0x02, 0x05, 0x03, 10, 20, 30, 0x08, 0x01, 40];
187        let mut cur = &buf[..];
188        let unknown = read_tagged_fields(&mut cur, |_, _| Ok(false)).unwrap();
189        let expected = UnknownTaggedFields(vec![
190            UnknownTaggedField {
191                tag: 5,
192                bytes: Bytes::from_static(&[10, 20, 30]),
193            },
194            UnknownTaggedField {
195                tag: 8,
196                bytes: Bytes::from_static(&[40]),
197            },
198        ]);
199        assert2::assert!(unknown == expected);
200    }
201
202    #[test]
203    fn truncated_tagged_field_reports_missing_payload_bytes() {
204        // count=1, tag=5, size=4, payload has only one byte.
205        let buf = [0x01, 0x05, 0x04, 0xAA];
206        let mut cur = &buf[..];
207        let err = read_tagged_fields(&mut cur, |_, _| Ok(false)).unwrap_err();
208
209        assert2::assert!(matches!(err, ProtocolError::UnexpectedEof { needed: 3 }));
210    }
211
212    #[test]
213    fn tagged_fields_len_matches_serialized_len_for_known_and_unknown() {
214        let known = [(1, 2), (300, 3)];
215        let unknown = UnknownTaggedFields(vec![
216            UnknownTaggedField {
217                tag: 7,
218                bytes: Bytes::from_static(&[0xA0]),
219            },
220            UnknownTaggedField {
221                tag: 400,
222                bytes: Bytes::from_static(&[0xB0, 0xB1, 0xB2, 0xB3]),
223            },
224        ]);
225        let mut w = WriteTaggedFields::new();
226        w.add(1, Bytes::from_static(&[0x01, 0x02]));
227        w.add(300, Bytes::from_static(&[0x03, 0x04, 0x05]));
228        let mut out = BytesMut::new();
229        w.write(&mut out, &unknown);
230
231        assert2::assert!(tagged_fields_len(&known, &unknown) == out.len());
232    }
233
234    #[test]
235    fn tagged_fields_len_accounts_for_count_varint_width() {
236        let unknown = UnknownTaggedFields(
237            (0..128)
238                .map(|tag| UnknownTaggedField {
239                    tag,
240                    bytes: Bytes::new(),
241                })
242                .collect(),
243        );
244        let mut out = BytesMut::new();
245        WriteTaggedFields::new().write(&mut out, &unknown);
246
247        assert2::assert!(tagged_fields_len(&[], &unknown) == out.len());
248    }
249
250    #[test]
251    fn ascending_order_enforced() {
252        // count=2, tag=5..., tag=3...  — invalid (descending)
253        let buf = [0x02, 0x05, 0x01, 0x00, 0x03, 0x01, 0x00];
254        let mut cur = &buf[..];
255        assert2::assert!(read_tagged_fields(&mut cur, |_, _| Ok(false)).is_err());
256    }
257
258    #[test]
259    fn write_merges_known_and_unknown_sorted() {
260        let mut w = WriteTaggedFields::new();
261        w.add(10, Bytes::from_static(&[0xAA]));
262        let unknown = UnknownTaggedFields(vec![UnknownTaggedField {
263            tag: 5,
264            bytes: Bytes::from_static(&[0xBB]),
265        }]);
266        let mut out = BytesMut::new();
267        w.write(&mut out, &unknown);
268        // Expect: count=2, tag=5,len=1,0xBB, tag=10,len=1,0xAA
269        assert2::assert!(&out[..] == &[0x02, 0x05, 0x01, 0xBB, 0x0A, 0x01, 0xAA]);
270    }
271}
krabka_protocol/tagged_fields.rs

krabka_protocol/
tagged_fields.rs
