http_streams_core/protobuf_len_codec.rs
1//! Decoding LEB128 length-prefixed protobuf messages.
2
3use crate::error::{StreamError, StreamErrorKind};
4use bytes::{Buf, BytesMut};
5use std::marker::PhantomData;
6
7/// A [`Decoder`](tokio_util::codec::Decoder) that yields one message per length-prefixed frame.
8#[derive(Clone, Debug)]
9pub struct ProtobufLenPrefixCodec<T> {
10 max_length: usize,
11 cursor: ProtobufCursor,
12 /// `fn() -> T` rather than `T`: a bare `PhantomData<T>` would make this codec `!Send`
13 /// whenever `T` is, and the crates built on this one promise `Send` streams for item
14 /// types that carry no such bound. The item type is produced, never held, so this is also
15 /// the honest variance.
16 _ph: PhantomData<fn() -> T>,
17}
18
19#[derive(Clone, Debug)]
20struct ProtobufCursor {
21 /// The length of the message currently being accumulated.
22 ///
23 /// `Option`, not a bare `usize` with zero meaning "none": a protobuf message all of whose
24 /// fields hold their default values encodes to *zero bytes*, so `0` is a perfectly ordinary
25 /// frame length and conflating it with "no length read yet" would make the codec try to
26 /// parse the following frame's length prefix as that message's body.
27 expected_len: Option<usize>,
28}
29
30impl<T> ProtobufLenPrefixCodec<T> {
31 /// A codec that rejects any single message longer than `max_length` bytes.
32 pub fn new_with_max_length(max_length: usize) -> Self {
33 let initial_cursor = ProtobufCursor { expected_len: None };
34
35 ProtobufLenPrefixCodec {
36 max_length,
37 cursor: initial_cursor,
38 _ph: PhantomData,
39 }
40 }
41}
42
43impl<T> tokio_util::codec::Decoder for ProtobufLenPrefixCodec<T>
44where
45 T: prost::Message + Default,
46{
47 /// Always `Ok(_)`: a message that fails to decode is reported as terminal, matching the
48 /// behaviour this codec has always had.
49 type Item = Result<T, StreamError>;
50 type Error = StreamError;
51
52 fn decode(&mut self, buf: &mut BytesMut) -> Result<Option<Self::Item>, StreamError> {
53 // Loops rather than returning `Ok(None)` after reading a length prefix.
54 //
55 // `Ok(None)` means "I need more bytes", and `FramedRead` acts on it: during `decode_eof`
56 // it ends the stream outright. A codec that consumed a length prefix and then said
57 // "need more bytes" while a complete message sat in the buffer would silently drop that
58 // message and every one after it. That happens whenever two or more whole frames are
59 // buffered when the body ends. So `Ok(None)` is returned only when the buffer genuinely
60 // cannot yield anything further.
61 loop {
62 // The length is checked before the buffer-empty guard, deliberately: a zero-length
63 // frame is complete with an empty buffer, and returning "need more bytes" for one
64 // would drop it.
65 let Some(expected_len) = self.cursor.expected_len else {
66 if buf.is_empty() {
67 return Ok(None);
68 }
69 match read_varint(buf)? {
70 // Made progress: go round again, the body may already be buffered.
71 Some(len) => {
72 self.cursor.expected_len = Some(len as usize);
73 continue;
74 }
75 // A partial varint. This one really does need more bytes.
76 None => return Ok(None),
77 }
78 };
79
80 if expected_len > self.max_length {
81 return Err(StreamError::new(
82 StreamErrorKind::MaxLenReachedError,
83 None,
84 Some("Max object length reached".into()),
85 ));
86 }
87
88 if buf.len() < expected_len {
89 return Ok(None);
90 }
91
92 let obj_bytes = buf.copy_to_bytes(expected_len);
93 self.cursor.expected_len = None;
94 return prost::Message::decode(obj_bytes)
95 .map(|item| Some(Ok(item)))
96 .map_err(|err| {
97 StreamError::new(StreamErrorKind::CodecError, Some(Box::new(err)), None)
98 });
99 }
100 }
101
102 fn decode_eof(&mut self, buf: &mut BytesMut) -> Result<Option<Self::Item>, StreamError> {
103 self.decode(buf)
104 }
105}
106
107/// Reads a LEB128 varint from the front of `buf`, consuming it.
108///
109/// Returns `Ok(None)` when the buffer holds only part of one, leaving `buf` untouched so the
110/// next call can retry with more bytes.
111fn read_varint(buf: &mut BytesMut) -> Result<Option<u64>, StreamError> {
112 let bytes = buf.chunk();
113 if bytes.is_empty() {
114 return Ok(None);
115 }
116
117 // The single-byte case is both the overwhelmingly common one and the only one that needs
118 // no lookahead at all.
119 if bytes[0] < 0x80 {
120 let value = u64::from(bytes[0]);
121 buf.advance(1);
122 return Ok(Some(value));
123 }
124
125 // `decode_varint_slice` requires either a full 10 bytes or a visible terminator; without
126 // one of those the varint is genuinely incomplete.
127 if bytes.len() > 10 || bytes[bytes.len() - 1] < 0x80 {
128 let (value, advance) = decode_varint_slice(bytes)?;
129 buf.advance(advance);
130 return Ok(Some(value));
131 }
132
133 Ok(None)
134}
135
136/// This function is copied from Prost, since it is not available as public API yet optimized for performance.
137///
138/// Decodes a LEB128-encoded variable length integer from the slice, returning the value and the
139/// number of bytes read.
140///
141/// Based loosely on [`ReadVarint64FromArray`][1] with a varint overflow check from
142/// [`ConsumeVarint`][2].
143///
144/// ## Safety
145///
146/// The caller must ensure that `bytes` is non-empty and either `bytes.len() >= 10` or the last
147/// element in bytes is < `0x80`.
148///
149/// [1]: https://github.com/google/protobuf/blob/3.3.x/src/google/protobuf/io/coded_stream.cc#L365-L406
150/// [2]: https://github.com/protocolbuffers/protobuf-go/blob/v1.27.1/encoding/protowire/wire.go#L358
151#[inline]
152fn decode_varint_slice(bytes: &[u8]) -> Result<(u64, usize), StreamError> {
153 // Fully unrolled varint decoding loop. Splitting into 32-bit pieces gives better performance.
154
155 // Use assertions to ensure memory safety, but it should always be optimized after inline.
156 assert!(!bytes.is_empty());
157 assert!(bytes.len() > 10 || bytes[bytes.len() - 1] < 0x80);
158
159 let mut b: u8 = bytes[0];
160 let mut part0: u32 = u32::from(b);
161 if b < 0x80 {
162 return Ok((u64::from(part0), 1));
163 };
164 part0 -= 0x80;
165 b = bytes[1];
166 part0 += u32::from(b) << 7;
167 if b < 0x80 {
168 return Ok((u64::from(part0), 2));
169 };
170 part0 -= 0x80 << 7;
171 b = bytes[2];
172 part0 += u32::from(b) << 14;
173 if b < 0x80 {
174 return Ok((u64::from(part0), 3));
175 };
176 part0 -= 0x80 << 14;
177 b = bytes[3];
178 part0 += u32::from(b) << 21;
179 if b < 0x80 {
180 return Ok((u64::from(part0), 4));
181 };
182 part0 -= 0x80 << 21;
183 let value = u64::from(part0);
184
185 b = bytes[4];
186 let mut part1: u32 = u32::from(b);
187 if b < 0x80 {
188 return Ok((value + (u64::from(part1) << 28), 5));
189 };
190 part1 -= 0x80;
191 b = bytes[5];
192 part1 += u32::from(b) << 7;
193 if b < 0x80 {
194 return Ok((value + (u64::from(part1) << 28), 6));
195 };
196 part1 -= 0x80 << 7;
197 b = bytes[6];
198 part1 += u32::from(b) << 14;
199 if b < 0x80 {
200 return Ok((value + (u64::from(part1) << 28), 7));
201 };
202 part1 -= 0x80 << 14;
203 b = bytes[7];
204 part1 += u32::from(b) << 21;
205 if b < 0x80 {
206 return Ok((value + (u64::from(part1) << 28), 8));
207 };
208 part1 -= 0x80 << 21;
209 let value = value + ((u64::from(part1)) << 28);
210
211 b = bytes[8];
212 let mut part2: u32 = u32::from(b);
213 if b < 0x80 {
214 return Ok((value + (u64::from(part2) << 56), 9));
215 };
216 part2 -= 0x80;
217 b = bytes[9];
218 part2 += u32::from(b) << 7;
219 // Check for u64::MAX overflow. See [`ConsumeVarint`][1] for details.
220 // [1]: https://github.com/protocolbuffers/protobuf-go/blob/v1.27.1/encoding/protowire/wire.go#L358
221 if b < 0x02 {
222 return Ok((value + (u64::from(part2) << 56), 10));
223 };
224
225 // We have overrun the maximum size of a varint (10 bytes) or the final byte caused an overflow.
226 // Assume the data is corrupt.
227 Err(StreamError::new(
228 StreamErrorKind::CodecError,
229 None,
230 Some("invalid varint".into()),
231 ))
232}