iggy_binary_protocol/
codec.rs

1// Licensed to the Apache Software Foundation (ASF) under one
2// or more contributor license agreements.  See the NOTICE file
3// distributed with this work for additional information
4// regarding copyright ownership.  The ASF licenses this file
5// to you under the Apache License, Version 2.0 (the
6// "License"); you may not use this file except in compliance
7// with the License.  You may obtain a copy of the License at
8//
9//   http://www.apache.org/licenses/LICENSE-2.0
10//
11// Unless required by applicable law or agreed to in writing,
12// software distributed under the License is distributed on an
13// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14// KIND, either express or implied.  See the License for the
15// specific language governing permissions and limitations
16// under the License.
17
18use crate::WireError;
19use bytes::{Bytes, BytesMut};
20
21/// Encode a wire type into a caller-owned buffer.
22///
23/// Buffer-first design: the caller controls allocation. The `to_bytes()`
24/// convenience method allocates when needed, but hot paths can reuse
25/// buffers via `encode()` directly.
26pub trait WireEncode {
27    /// Write the encoded representation into `buf`.
28    fn encode(&self, buf: &mut BytesMut);
29
30    /// Return the exact encoded size in bytes.
31    fn encoded_size(&self) -> usize;
32
33    /// Convenience: allocate a new [`Bytes`] and encode into it.
34    #[must_use]
35    fn to_bytes(&self) -> Bytes {
36        let mut buf = BytesMut::with_capacity(self.encoded_size());
37        self.encode(&mut buf);
38        buf.freeze()
39    }
40}
41
42/// Decode a wire type from a byte slice.
43///
44/// Takes `&[u8]` instead of [`Bytes`] to avoid requiring reference-counted
45/// ownership at the decode boundary.
46pub trait WireDecode: Sized {
47    /// Decode from `buf`, consuming exactly the bytes needed.
48    /// Returns the decoded value and the number of bytes consumed.
49    ///
50    /// # Errors
51    /// Returns `WireError` if the buffer is too short or contains invalid data.
52    fn decode(buf: &[u8]) -> Result<(Self, usize), WireError>;
53
54    /// Convenience: decode from the entire buffer, ignoring trailing bytes.
55    ///
56    /// # Errors
57    /// Returns `WireError` if decoding fails.
58    fn decode_from(buf: &[u8]) -> Result<Self, WireError> {
59        Self::decode(buf).map(|(val, _)| val)
60    }
61}
62
63/// Helper to read a `u8` from `buf` at `offset`.
64///
65/// # Errors
66/// Returns `WireError::UnexpectedEof` if `offset` is out of bounds.
67#[inline]
68pub fn read_u8(buf: &[u8], offset: usize) -> Result<u8, WireError> {
69    buf.get(offset)
70        .copied()
71        .ok_or_else(|| WireError::UnexpectedEof {
72            offset,
73            need: 1,
74            have: buf.len().saturating_sub(offset),
75        })
76}
77
78/// Helper to read a `u16` LE from `buf` at `offset`.
79///
80/// # Errors
81/// Returns `WireError::UnexpectedEof` if fewer than 2 bytes remain.
82#[allow(clippy::missing_panics_doc)]
83#[inline]
84pub fn read_u16_le(buf: &[u8], offset: usize) -> Result<u16, WireError> {
85    let end = offset
86        .checked_add(2)
87        .ok_or_else(|| WireError::UnexpectedEof {
88            offset,
89            need: 2,
90            have: buf.len().saturating_sub(offset),
91        })?;
92    let slice = buf
93        .get(offset..end)
94        .ok_or_else(|| WireError::UnexpectedEof {
95            offset,
96            need: 2,
97            have: buf.len().saturating_sub(offset),
98        })?;
99    Ok(u16::from_le_bytes(
100        slice.try_into().expect("slice is exactly 2 bytes"),
101    ))
102}
103
104/// Helper to read a `u32` LE from `buf` at `offset`.
105///
106/// # Errors
107/// Returns `WireError::UnexpectedEof` if fewer than 4 bytes remain.
108#[allow(clippy::missing_panics_doc)]
109#[inline]
110pub fn read_u32_le(buf: &[u8], offset: usize) -> Result<u32, WireError> {
111    let end = offset
112        .checked_add(4)
113        .ok_or_else(|| WireError::UnexpectedEof {
114            offset,
115            need: 4,
116            have: buf.len().saturating_sub(offset),
117        })?;
118    let slice = buf
119        .get(offset..end)
120        .ok_or_else(|| WireError::UnexpectedEof {
121            offset,
122            need: 4,
123            have: buf.len().saturating_sub(offset),
124        })?;
125    Ok(u32::from_le_bytes(
126        slice.try_into().expect("slice is exactly 4 bytes"),
127    ))
128}
129
130/// Helper to read a `u64` LE from `buf` at `offset`.
131///
132/// # Errors
133/// Returns `WireError::UnexpectedEof` if fewer than 8 bytes remain.
134#[allow(clippy::missing_panics_doc)]
135#[inline]
136pub fn read_u64_le(buf: &[u8], offset: usize) -> Result<u64, WireError> {
137    let end = offset
138        .checked_add(8)
139        .ok_or_else(|| WireError::UnexpectedEof {
140            offset,
141            need: 8,
142            have: buf.len().saturating_sub(offset),
143        })?;
144    let slice = buf
145        .get(offset..end)
146        .ok_or_else(|| WireError::UnexpectedEof {
147            offset,
148            need: 8,
149            have: buf.len().saturating_sub(offset),
150        })?;
151    Ok(u64::from_le_bytes(
152        slice.try_into().expect("slice is exactly 8 bytes"),
153    ))
154}
155
156/// Helper to read a UTF-8 string of `len` bytes from `buf` at `offset`.
157///
158/// # Errors
159/// Returns `WireError::UnexpectedEof` or `WireError::InvalidUtf8` on failure.
160#[inline]
161pub fn read_str(buf: &[u8], offset: usize, len: usize) -> Result<String, WireError> {
162    let end = offset
163        .checked_add(len)
164        .ok_or_else(|| WireError::UnexpectedEof {
165            offset,
166            need: len,
167            have: buf.len().saturating_sub(offset),
168        })?;
169    let slice = buf
170        .get(offset..end)
171        .ok_or_else(|| WireError::UnexpectedEof {
172            offset,
173            need: len,
174            have: buf.len().saturating_sub(offset),
175        })?;
176    std::str::from_utf8(slice)
177        .map(str::to_string)
178        .map_err(|_| WireError::InvalidUtf8 { offset })
179}
180
181/// Helper to read a `u128` LE from `buf` at `offset`.
182///
183/// # Errors
184/// Returns `WireError::UnexpectedEof` if fewer than 16 bytes remain.
185#[allow(clippy::missing_panics_doc)]
186#[inline]
187pub fn read_u128_le(buf: &[u8], offset: usize) -> Result<u128, WireError> {
188    let end = offset
189        .checked_add(16)
190        .ok_or_else(|| WireError::UnexpectedEof {
191            offset,
192            need: 16,
193            have: buf.len().saturating_sub(offset),
194        })?;
195    let slice = buf
196        .get(offset..end)
197        .ok_or_else(|| WireError::UnexpectedEof {
198            offset,
199            need: 16,
200            have: buf.len().saturating_sub(offset),
201        })?;
202    Ok(u128::from_le_bytes(
203        slice.try_into().expect("slice is exactly 16 bytes"),
204    ))
205}
206
207/// Helper to read an `f32` LE from `buf` at `offset`.
208///
209/// # Errors
210/// Returns `WireError::UnexpectedEof` if fewer than 4 bytes remain.
211#[allow(clippy::missing_panics_doc)]
212#[inline]
213pub fn read_f32_le(buf: &[u8], offset: usize) -> Result<f32, WireError> {
214    let end = offset
215        .checked_add(4)
216        .ok_or_else(|| WireError::UnexpectedEof {
217            offset,
218            need: 4,
219            have: buf.len().saturating_sub(offset),
220        })?;
221    let slice = buf
222        .get(offset..end)
223        .ok_or_else(|| WireError::UnexpectedEof {
224            offset,
225            need: 4,
226            have: buf.len().saturating_sub(offset),
227        })?;
228    Ok(f32::from_le_bytes(
229        slice.try_into().expect("slice is exactly 4 bytes"),
230    ))
231}
232
233/// Helper to read a byte slice of `len` bytes from `buf` at `offset`.
234///
235/// # Errors
236/// Returns `WireError::UnexpectedEof` if fewer than `len` bytes remain.
237#[inline]
238pub fn read_bytes(buf: &[u8], offset: usize, len: usize) -> Result<&[u8], WireError> {
239    let end = offset
240        .checked_add(len)
241        .ok_or_else(|| WireError::UnexpectedEof {
242            offset,
243            need: len,
244            have: buf.len().saturating_sub(offset),
245        })?;
246    buf.get(offset..end)
247        .ok_or_else(|| WireError::UnexpectedEof {
248            offset,
249            need: len,
250            have: buf.len().saturating_sub(offset),
251        })
252}
253
254/// Capacity to preallocate for `count` elements that still have to be decoded
255/// out of `remaining` bytes, each at least `min_element_size` bytes on the
256/// wire.
257///
258/// A count read off the wire is chosen by the peer, and `Vec::with_capacity`
259/// reserves for it before a single element has been bounds-checked. A hostile
260/// or corrupt frame claiming `u32::MAX` elements therefore aborts the process
261/// on allocation failure instead of returning a decode error. The bytes left
262/// to decode are the honest ceiling on how many elements can exist.
263///
264/// A zero `min_element_size` is treated as one byte rather than passing the
265/// wire count through: an element that occupies no bytes cannot exist, and
266/// trusting the count is the exact hole this function was written to close.
267#[must_use]
268#[inline]
269pub fn bounded_capacity(count: usize, remaining: usize, min_element_size: usize) -> usize {
270    count.min(remaining / min_element_size.max(1))
271}
272
273#[cfg(test)]
274mod tests {
275    use super::*;
276
277    #[test]
278    fn bounded_capacity_limits_allocation() {
279        assert_eq!(bounded_capacity(1_000_000, 100, 10), 10);
280        assert_eq!(bounded_capacity(5, 100, 10), 5);
281        assert_eq!(bounded_capacity(10, 100, 10), 10);
282        assert_eq!(bounded_capacity(11, 100, 10), 10);
283        assert_eq!(bounded_capacity(0, 100, 10), 0);
284        assert_eq!(bounded_capacity(100, 0, 10), 0);
285    }
286
287    #[test]
288    fn bounded_capacity_zero_element_size_still_bounds_by_the_buffer() {
289        assert_eq!(bounded_capacity(1_000_000, 100, 0), 100);
290    }
291}
iggy_binary_protocol/codec.rs

iggy_binary_protocol/
codec.rs
