Skip to main content

kacrab_protocol/
tagged.rs

1//! Tagged-fields section for flexible Kafka message versions.
2//!
3//! Wire format: `unsigned_varint(count)` then for each field
4//! `unsigned_varint(tag), unsigned_varint(size), <size bytes>`. Tags MUST be
5//! strictly ascending; readers reject duplicates and out-of-order tags.
6
7pub mod error;
8
9use bytes::{Buf, Bytes, BytesMut};
10
11pub use self::error::TaggedFieldError;
12use crate::primitives::{read_unsigned_varint, unsigned_varint_len, write_unsigned_varint};
13
14/// Result alias for tagged-field operations.
15pub type Result<T> = core::result::Result<T, TaggedFieldError>;
16
17/// A single raw tagged field. The body bytes are stored verbatim — interpretation
18/// is the caller's responsibility (the schema dictates the encoding per tag).
19#[derive(Debug, Clone, PartialEq, Eq)]
20pub struct RawTaggedField {
21    /// Tag number.
22    pub tag: u32,
23    /// Raw payload bytes.
24    pub data: Bytes,
25}
26
27/// Read the tagged-fields section.
28///
29/// Enforces ascending tag order and that each declared `size` is available
30/// in the buffer.
31pub fn read_tagged_fields(buf: &mut Bytes) -> Result<Vec<RawTaggedField>> {
32    let count = read_unsigned_varint(buf)?;
33    let count_usize = usize::try_from(count).map_err(|_| TaggedFieldError::CountOverflow {
34        count,
35        max: usize::MAX,
36    })?;
37    // Bound the pre-allocation by the bytes left in `buf`: every field needs at
38    // least a tag and a size varint, so `remaining` caps the possible count and
39    // a hostile `count` can't force a huge reservation. For valid input this is
40    // exactly `count_usize`.
41    let mut fields = Vec::with_capacity(count_usize.min(buf.remaining()));
42    let mut prev_tag: Option<u32> = None;
43
44    for _ in 0..count {
45        let tag = read_unsigned_varint(buf)?;
46        if let Some(prev) = prev_tag
47            && tag <= prev
48        {
49            return Err(TaggedFieldError::OutOfOrder {
50                tag,
51                prev_tag: prev,
52            });
53        }
54        prev_tag = Some(tag);
55
56        let raw_size = read_unsigned_varint(buf)?;
57        let size = usize::try_from(raw_size).map_err(|_| TaggedFieldError::SizeOverflow {
58            tag,
59            size: usize::MAX,
60            remaining: buf.remaining(),
61        })?;
62        let remaining = buf.remaining();
63        if size > remaining {
64            return Err(TaggedFieldError::SizeOverflow {
65                tag,
66                size,
67                remaining,
68            });
69        }
70        let bytes = buf.split_to(size);
71        fields.push(RawTaggedField { tag, data: bytes });
72    }
73
74    Ok(fields)
75}
76
77/// Write the tagged-fields section.
78///
79/// Caller must supply fields sorted by ascending tag; this is validated
80/// (returning [`TaggedFieldError::OutOfOrder`]), not re-sorted, to keep the
81/// wire encoding deterministic without surprise costs.
82pub fn write_tagged_fields(buf: &mut BytesMut, fields: &[RawTaggedField]) -> Result<()> {
83    let mut prev_tag: Option<u32> = None;
84    for field in fields {
85        if let Some(prev) = prev_tag
86            && field.tag <= prev
87        {
88            return Err(TaggedFieldError::OutOfOrder {
89                tag: field.tag,
90                prev_tag: prev,
91            });
92        }
93        prev_tag = Some(field.tag);
94    }
95
96    let field_count = u32::try_from(fields.len()).map_err(|_| TaggedFieldError::CountOverflow {
97        count: u32::MAX,
98        max: usize::try_from(u32::MAX).unwrap_or(usize::MAX),
99    })?;
100    write_unsigned_varint(buf, field_count);
101    for field in fields {
102        write_unsigned_varint(buf, field.tag);
103        let field_len =
104            u32::try_from(field.data.len()).map_err(|_| TaggedFieldError::FieldTooLarge {
105                tag: field.tag,
106                size: field.data.len(),
107                max: usize::try_from(u32::MAX).unwrap_or(usize::MAX),
108            })?;
109        write_unsigned_varint(buf, field_len);
110        buf.extend_from_slice(&field.data);
111    }
112    Ok(())
113}
114
115/// Encoded length of a tagged-fields section.
116pub fn tagged_fields_len(fields: &[RawTaggedField]) -> Result<usize> {
117    let mut prev_tag: Option<u32> = None;
118    for field in fields {
119        if let Some(prev) = prev_tag
120            && field.tag <= prev
121        {
122            return Err(TaggedFieldError::OutOfOrder {
123                tag: field.tag,
124                prev_tag: prev,
125            });
126        }
127        prev_tag = Some(field.tag);
128    }
129
130    let field_count = u32::try_from(fields.len()).map_err(|_| TaggedFieldError::CountOverflow {
131        count: u32::MAX,
132        max: usize::try_from(u32::MAX).unwrap_or(usize::MAX),
133    })?;
134    let mut len = unsigned_varint_len(field_count);
135    for field in fields {
136        let field_len =
137            u32::try_from(field.data.len()).map_err(|_| TaggedFieldError::FieldTooLarge {
138                tag: field.tag,
139                size: field.data.len(),
140                max: usize::try_from(u32::MAX).unwrap_or(usize::MAX),
141            })?;
142        len = len
143            .checked_add(unsigned_varint_len(field.tag))
144            .and_then(|len| len.checked_add(unsigned_varint_len(field_len)))
145            .and_then(|len| len.checked_add(field.data.len()))
146            .ok_or_else(|| TaggedFieldError::FieldTooLarge {
147                tag: field.tag,
148                size: field.data.len(),
149                max: usize::try_from(u32::MAX).unwrap_or(usize::MAX),
150            })?;
151    }
152    Ok(len)
153}