kafka_protocol/protocol/
mod.rs

1//! Most types are used internally in encoding/decoding, and are not required by typical use cases
2//! for interacting with the protocol. However, types can be used for decoding partial messages,
3//! or rewriting parts of an encoded message.
4use std::cmp;
5use std::collections::BTreeMap;
6use std::ops::RangeBounds;
7use std::{borrow::Borrow, fmt::Display};
8
9use anyhow::{bail, Result};
10use buf::{ByteBuf, ByteBufMut};
11use bytes::Bytes;
12
13pub mod buf;
14pub mod types;
15
16mod str_bytes {
17    use bytes::Bytes;
18    use std::borrow::Borrow;
19    use std::convert::TryFrom;
20    use std::fmt::{Debug, Display, Formatter};
21    use std::ops::Deref;
22    use std::str::Utf8Error;
23
24    /// A string type backed by [Bytes].
25    #[derive(Clone, Hash, Ord, PartialOrd, PartialEq, Eq, Default)]
26    pub struct StrBytes(Bytes);
27
28    impl StrBytes {
29        /// Creates a new empty `StrBytes`.
30        ///
31        /// This will not allocate and the returned `StrBytes` handle will be empty.
32        pub const fn new() -> Self {
33            Self(Bytes::new())
34        }
35
36        /// Construct a [StrBytes] from the given [Bytes] instance,
37        /// checking that it contains valid UTF-8 data.
38        pub fn from_utf8(bytes: Bytes) -> Result<Self, Utf8Error> {
39            let _: &str = std::str::from_utf8(&bytes)?;
40            Ok(Self(bytes))
41        }
42
43        /// Construct a [StrBytes] from the provided static [str].
44        pub const fn from_static_str(s: &'static str) -> Self {
45            Self(Bytes::from_static(s.as_bytes()))
46        }
47
48        /// Construct a [StrBytes] from the provided [String] without additional allocations.
49        pub fn from_string(s: String) -> Self {
50            Self(Bytes::from(s.into_bytes()))
51        }
52
53        /// View the contents of this [StrBytes] as a [str] reference.
54        pub fn as_str(&self) -> &str {
55            // SAFETY: all methods of constructing `self` check that the backing data is valid utf8,
56            // and bytes::Bytes guarantees that its contents will not change unless we mutate it,
57            // and we never mutate it.
58            unsafe { std::str::from_utf8_unchecked(&self.0) }
59        }
60
61        /// Extract the underlying [Bytes].
62        pub fn into_bytes(self) -> Bytes {
63            self.0
64        }
65    }
66
67    impl TryFrom<Bytes> for StrBytes {
68        type Error = Utf8Error;
69
70        fn try_from(value: Bytes) -> Result<Self, Self::Error> {
71            StrBytes::from_utf8(value)
72        }
73    }
74
75    impl From<StrBytes> for Bytes {
76        fn from(value: StrBytes) -> Bytes {
77            value.0
78        }
79    }
80
81    impl From<String> for StrBytes {
82        fn from(value: String) -> Self {
83            Self::from_string(value)
84        }
85    }
86
87    impl From<&'static str> for StrBytes {
88        fn from(value: &'static str) -> Self {
89            Self::from_static_str(value)
90        }
91    }
92
93    impl Deref for StrBytes {
94        type Target = str;
95
96        fn deref(&self) -> &Self::Target {
97            self.as_str()
98        }
99    }
100
101    impl Debug for StrBytes {
102        fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
103            Debug::fmt(self.as_str(), f)
104        }
105    }
106
107    impl Display for StrBytes {
108        fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
109            std::fmt::Display::fmt(&**self, f)
110        }
111    }
112
113    impl PartialEq<str> for StrBytes {
114        fn eq(&self, other: &str) -> bool {
115            self.as_str().eq(other)
116        }
117    }
118
119    impl Borrow<[u8]> for StrBytes {
120        fn borrow(&self) -> &[u8] {
121            // Note that there is an equivalent Hash implementation between
122            // &[u8] and StrBytes, which makes this impl correct
123            // as described in the `std::borrow::Borrow` docs.
124            self.as_bytes()
125        }
126    }
127}
128
129pub use str_bytes::StrBytes;
130
131use crate::messages::{ApiKey, RequestHeader};
132
133pub(crate) trait NewType<Inner>: From<Inner> + Into<Inner> + Borrow<Inner> {}
134
135impl<T> NewType<T> for T {}
136
137pub(crate) trait Encoder<Value> {
138    fn encode<B: ByteBufMut>(&self, buf: &mut B, value: Value) -> Result<()>;
139    fn compute_size(&self, value: Value) -> Result<usize>;
140    fn fixed_size(&self) -> Option<usize> {
141        None
142    }
143}
144
145pub(crate) trait Decoder<Value> {
146    fn decode<B: ByteBuf>(&self, buf: &mut B) -> Result<Value>;
147}
148
149/// The range of versions (min, max) allowed for agiven message.
150#[derive(Debug, Copy, Clone, PartialEq)]
151pub struct VersionRange {
152    /// The minimum version in the range.
153    pub min: i16,
154    /// The maximum version in the range.
155    pub max: i16,
156}
157
158impl VersionRange {
159    /// Checks whether the version range contains no versions.
160    pub fn is_empty(&self) -> bool {
161        self.min > self.max
162    }
163
164    /// Finds the valid intersection with a provided other version range.
165    pub fn intersect(&self, other: &VersionRange) -> VersionRange {
166        VersionRange {
167            min: cmp::max(self.min, other.min),
168            max: cmp::min(self.max, other.max),
169        }
170    }
171}
172
173impl Display for VersionRange {
174    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
175        write!(f, "{}..{}", self.min, self.max)
176    }
177}
178
179/// An API request or response.
180///
181/// All API messages must provide a set of valid versions.
182pub trait Message: Sized {
183    /// The valid versions for this message.
184    const VERSIONS: VersionRange;
185    /// The deprecated versions for this message.
186    const DEPRECATED_VERSIONS: Option<VersionRange>;
187}
188
189/// An encodable message.
190pub trait Encodable: Sized {
191    /// Encode the message into the target buffer.
192    fn encode<B: ByteBufMut>(&self, buf: &mut B, version: i16) -> Result<()>;
193    /// Compute the total size of the message when encoded.
194    fn compute_size(&self, version: i16) -> Result<usize>;
195}
196
197/// A decodable message.
198pub trait Decodable: Sized {
199    /// Decode the message from the provided buffer and version.
200    fn decode<B: ByteBuf>(buf: &mut B, version: i16) -> Result<Self>;
201}
202
203/// Every message has a set of versions valid for a given header version.
204pub trait HeaderVersion {
205    /// Maps a header version to a given version for a particular API message.
206    fn header_version(version: i16) -> i16;
207}
208
209/// An API request.
210///
211/// Every abstract request must be able to provide the following items:
212/// - An API key mapped to this request.
213/// - A version based on a provided header version.
214pub trait Request: Message + Encodable + Decodable + HeaderVersion {
215    /// The API key of this request.
216    const KEY: i16;
217    /// The response associated with this request.
218    type Response: Message + Encodable + Decodable + HeaderVersion;
219}
220
221/// Decode the request header from the provided buffer.
222pub fn decode_request_header_from_buffer<B: ByteBuf>(buf: &mut B) -> Result<RequestHeader> {
223    let api_key = ApiKey::try_from(bytes::Buf::get_i16(&mut buf.peek_bytes(0..2)))
224        .map_err(|_| anyhow::Error::msg("Unknown API key"))?;
225    let api_version = bytes::Buf::get_i16(&mut buf.peek_bytes(2..4));
226    let header_version = api_key.request_header_version(api_version);
227    RequestHeader::decode(buf, header_version)
228}
229
230/// Encode the request header into the provided buffer.
231pub fn encode_request_header_into_buffer<B: ByteBufMut>(
232    buf: &mut B,
233    header: &RequestHeader,
234) -> Result<()> {
235    let api_key = ApiKey::try_from(header.request_api_key)
236        .map_err(|_| anyhow::Error::msg("Unknown API key"))?;
237    let version = api_key.request_header_version(header.request_api_version);
238    header.encode(buf, version)
239}
240
241pub(crate) fn write_unknown_tagged_fields<B: ByteBufMut, R: RangeBounds<i32>>(
242    buf: &mut B,
243    range: R,
244    unknown_tagged_fields: &BTreeMap<i32, Bytes>,
245) -> Result<()> {
246    for (&k, v) in unknown_tagged_fields.range(range) {
247        if v.len() > u32::MAX as usize {
248            bail!("Tagged field is too long to encode ({} bytes)", v.len());
249        }
250        types::UnsignedVarInt.encode(buf, k as u32)?;
251        types::UnsignedVarInt.encode(buf, v.len() as u32)?;
252        buf.put_slice(v);
253    }
254    Ok(())
255}
256
257pub(crate) fn compute_unknown_tagged_fields_size(
258    unknown_tagged_fields: &BTreeMap<i32, Bytes>,
259) -> Result<usize> {
260    let mut total_size = 0;
261    for (&k, v) in unknown_tagged_fields {
262        if v.len() > u32::MAX as usize {
263            bail!("Tagged field is too long to encode ({} bytes)", v.len());
264        }
265        total_size += types::UnsignedVarInt.compute_size(k as u32)?;
266        total_size += types::UnsignedVarInt.compute_size(v.len() as u32)?;
267        total_size += v.len();
268    }
269    Ok(total_size)
270}
kafka_protocol/protocol/mod.rs

kafka_protocol/protocol/
mod.rs
