1use crate::error::InvalidMessage;
2
3use alloc::vec::Vec;
4use core::fmt::Debug;
5
6/// Wrapper over a slice of bytes that allows reading chunks from
7/// with the current position state held using a cursor.
8///
9/// A new reader for a sub section of the the buffer can be created
10/// using the `sub` function or a section of a certain length can
11/// be obtained using the `take` function
12pub struct Reader<'a> {
13 /// The underlying buffer storing the readers content
14 buffer: &'a [u8],
15 /// Stores the current reading position for the buffer
16 cursor: usize,
17}
18
19impl<'a> Reader<'a> {
20 /// Creates a new Reader of the provided `bytes` slice with
21 /// the initial cursor position of zero.
22 pub fn init(bytes: &[u8]) -> Reader {
23 Reader {
24 buffer: bytes,
25 cursor: 0,
26 }
27 }
28
29 /// Attempts to create a new Reader on a sub section of this
30 /// readers bytes by taking a slice of the provided `length`
31 /// will return None if there is not enough bytes
32 pub fn sub(&mut self, length: usize) -> Result<Reader, InvalidMessage> {
33 match self.take(length) {
34 Some(bytes) => Ok(Reader::init(bytes)),
35 None => Err(InvalidMessage::MessageTooShort),
36 }
37 }
38
39 /// Borrows a slice of all the remaining bytes
40 /// that appear after the cursor position.
41 ///
42 /// Moves the cursor to the end of the buffer length.
43 pub fn rest(&mut self) -> &[u8] {
44 let rest = &self.buffer[self.cursor..];
45 self.cursor = self.buffer.len();
46 rest
47 }
48
49 /// Attempts to borrow a slice of bytes from the current
50 /// cursor position of `length` if there is not enough
51 /// bytes remaining after the cursor to take the length
52 /// then None is returned instead.
53 pub fn take(&mut self, length: usize) -> Option<&[u8]> {
54 if self.left() < length {
55 return None;
56 }
57 let current = self.cursor;
58 self.cursor += length;
59 Some(&self.buffer[current..current + length])
60 }
61
62 /// Used to check whether the reader has any content left
63 /// after the cursor (cursor has not reached end of buffer)
64 pub fn any_left(&self) -> bool {
65 self.cursor < self.buffer.len()
66 }
67
68 pub fn expect_empty(&self, name: &'static str) -> Result<(), InvalidMessage> {
69 match self.any_left() {
70 true => Err(InvalidMessage::TrailingData(name)),
71 false => Ok(()),
72 }
73 }
74
75 /// Returns the cursor position which is also the number
76 /// of bytes that have been read from the buffer.
77 pub fn used(&self) -> usize {
78 self.cursor
79 }
80
81 /// Returns the number of bytes that are still able to be
82 /// read (The number of remaining takes)
83 pub fn left(&self) -> usize {
84 self.buffer.len() - self.cursor
85 }
86}
87
88/// Trait for implementing encoding and decoding functionality
89/// on something.
90pub trait Codec: Debug + Sized {
91 /// Function for encoding itself by appending itself to
92 /// the provided vec of bytes.
93 fn encode(&self, bytes: &mut Vec<u8>);
94
95 /// Function for decoding itself from the provided reader
96 /// will return Some if the decoding was successful or
97 /// None if it was not.
98 fn read(_: &mut Reader) -> Result<Self, InvalidMessage>;
99
100 /// Convenience function for encoding the implementation
101 /// into a vec and returning it
102 fn get_encoding(&self) -> Vec<u8> {
103 let mut bytes = Vec::new();
104 self.encode(&mut bytes);
105 bytes
106 }
107
108 /// Function for wrapping a call to the read function in
109 /// a Reader for the slice of bytes provided
110 fn read_bytes(bytes: &[u8]) -> Result<Self, InvalidMessage> {
111 let mut reader = Reader::init(bytes);
112 Self::read(&mut reader)
113 }
114}
115
116impl Codec for u8 {
117 fn encode(&self, bytes: &mut Vec<u8>) {
118 bytes.push(*self);
119 }
120
121 fn read(r: &mut Reader) -> Result<Self, InvalidMessage> {
122 match r.take(length:1) {
123 Some(&[byte: u8]) => Ok(byte),
124 _ => Err(InvalidMessage::MissingData("u8")),
125 }
126 }
127}
128
129pub(crate) fn put_u16(v: u16, out: &mut [u8]) {
130 let out: &mut [u8; 2] = (&mut out[..2]).try_into().unwrap();
131 *out = u16::to_be_bytes(self:v);
132}
133
134impl Codec for u16 {
135 fn encode(&self, bytes: &mut Vec<u8>) {
136 let mut b16: [u8; 2] = [0u8; 2];
137 put_u16(*self, &mut b16);
138 bytes.extend_from_slice(&b16);
139 }
140
141 fn read(r: &mut Reader) -> Result<Self, InvalidMessage> {
142 match r.take(length:2) {
143 Some(&[b1: u8, b2: u8]) => Ok(Self::from_be_bytes([b1, b2])),
144 _ => Err(InvalidMessage::MissingData("u8")),
145 }
146 }
147}
148
149// Make a distinct type for u24, even though it's a u32 underneath
150#[allow(non_camel_case_types)]
151#[derive(Debug, Copy, Clone)]
152pub struct u24(pub u32);
153
154#[cfg(any(target_pointer_width = "32", target_pointer_width = "64"))]
155impl From<u24> for usize {
156 #[inline]
157 fn from(v: u24) -> Self {
158 v.0 as Self
159 }
160}
161
162impl Codec for u24 {
163 fn encode(&self, bytes: &mut Vec<u8>) {
164 let be_bytes: [u8; 4] = u32::to_be_bytes(self.0);
165 bytes.extend_from_slice(&be_bytes[1..]);
166 }
167
168 fn read(r: &mut Reader) -> Result<Self, InvalidMessage> {
169 match r.take(length:3) {
170 Some(&[a: u8, b: u8, c: u8]) => Ok(Self(u32::from_be_bytes([0, a, b, c]))),
171 _ => Err(InvalidMessage::MissingData("u24")),
172 }
173 }
174}
175
176impl Codec for u32 {
177 fn encode(&self, bytes: &mut Vec<u8>) {
178 bytes.extend(Self::to_be_bytes(*self));
179 }
180
181 fn read(r: &mut Reader) -> Result<Self, InvalidMessage> {
182 match r.take(length:4) {
183 Some(&[a: u8, b: u8, c: u8, d: u8]) => Ok(Self::from_be_bytes([a, b, c, d])),
184 _ => Err(InvalidMessage::MissingData("u32")),
185 }
186 }
187}
188
189pub(crate) fn put_u64(v: u64, bytes: &mut [u8]) {
190 let bytes: &mut [u8; 8] = (&mut bytes[..8]).try_into().unwrap();
191 *bytes = u64::to_be_bytes(self:v);
192}
193
194impl Codec for u64 {
195 fn encode(&self, bytes: &mut Vec<u8>) {
196 let mut b64: [u8; 8] = [0u8; 8];
197 put_u64(*self, &mut b64);
198 bytes.extend_from_slice(&b64);
199 }
200
201 fn read(r: &mut Reader) -> Result<Self, InvalidMessage> {
202 match r.take(length:8) {
203 Some(&[a: u8, b: u8, c: u8, d: u8, e: u8, f: u8, g: u8, h: u8]) => Ok(Self::from_be_bytes([a, b, c, d, e, f, g, h])),
204 _ => Err(InvalidMessage::MissingData("u64")),
205 }
206 }
207}
208
209/// Implement `Codec` for lists of elements that implement `TlsListElement`.
210///
211/// `TlsListElement` provides the size of the length prefix for the list.
212impl<T: Codec + TlsListElement + Debug> Codec for Vec<T> {
213 fn encode(&self, bytes: &mut Vec<u8>) {
214 let nest = LengthPrefixedBuffer::new(T::SIZE_LEN, bytes);
215
216 for i in self {
217 i.encode(nest.buf);
218 }
219 }
220
221 fn read(r: &mut Reader) -> Result<Self, InvalidMessage> {
222 let len = match T::SIZE_LEN {
223 ListLength::U8 => usize::from(u8::read(r)?),
224 ListLength::U16 => usize::from(u16::read(r)?),
225 ListLength::U24 { max } => Ord::min(usize::from(u24::read(r)?), max),
226 };
227
228 let mut sub = r.sub(len)?;
229 let mut ret = Self::new();
230 while sub.any_left() {
231 ret.push(T::read(&mut sub)?);
232 }
233
234 Ok(ret)
235 }
236}
237
238/// A trait for types that can be encoded and decoded in a list.
239///
240/// This trait is used to implement `Codec` for `Vec<T>`. Lists in the TLS wire format are
241/// prefixed with a length, the size of which depends on the type of the list elements.
242/// As such, the `Codec` implementation for `Vec<T>` requires an implementation of this trait
243/// for its element type `T`.
244pub(crate) trait TlsListElement {
245 const SIZE_LEN: ListLength;
246}
247
248/// The length of the length prefix for a list.
249///
250/// The types that appear in lists are limited to three kinds of length prefixes:
251/// 1, 2, and 3 bytes. For the latter kind, we require a `TlsListElement` implementer
252/// to specify a maximum length.
253pub(crate) enum ListLength {
254 U8,
255 U16,
256 U24 { max: usize },
257}
258
259/// Tracks encoding a length-delimited structure in a single pass.
260pub(crate) struct LengthPrefixedBuffer<'a> {
261 pub(crate) buf: &'a mut Vec<u8>,
262 len_offset: usize,
263 size_len: ListLength,
264}
265
266impl<'a> LengthPrefixedBuffer<'a> {
267 /// Inserts a dummy length into `buf`, and remembers where it went.
268 ///
269 /// After this, the body of the length-delimited structure should be appended to `LengthPrefixedBuffer::buf`.
270 /// The length header is corrected in `LengthPrefixedBuffer::drop`.
271 pub(crate) fn new(size_len: ListLength, buf: &'a mut Vec<u8>) -> LengthPrefixedBuffer<'a> {
272 let len_offset: usize = buf.len();
273 buf.extend(iter:match size_len {
274 ListLength::U8 => &[0xff][..],
275 ListLength::U16 => &[0xff, 0xff],
276 ListLength::U24 { .. } => &[0xff, 0xff, 0xff],
277 });
278
279 Self {
280 buf,
281 len_offset,
282 size_len,
283 }
284 }
285}
286
287impl<'a> Drop for LengthPrefixedBuffer<'a> {
288 /// Goes back and corrects the length previously inserted at the start of the structure.
289 fn drop(&mut self) {
290 match self.size_len {
291 ListLength::U8 => {
292 let len = self.buf.len() - self.len_offset - 1;
293 debug_assert!(len <= 0xff);
294 self.buf[self.len_offset] = len as u8;
295 }
296 ListLength::U16 => {
297 let len = self.buf.len() - self.len_offset - 2;
298 debug_assert!(len <= 0xffff);
299 let out: &mut [u8; 2] = (&mut self.buf[self.len_offset..self.len_offset + 2])
300 .try_into()
301 .unwrap();
302 *out = u16::to_be_bytes(len as u16);
303 }
304 ListLength::U24 { .. } => {
305 let len = self.buf.len() - self.len_offset - 3;
306 debug_assert!(len <= 0xff_ffff);
307 let len_bytes = u32::to_be_bytes(len as u32);
308 let out: &mut [u8; 3] = (&mut self.buf[self.len_offset..self.len_offset + 3])
309 .try_into()
310 .unwrap();
311 out.copy_from_slice(&len_bytes[1..]);
312 }
313 }
314 }
315}
316
317#[cfg(test)]
318mod tests {
319 use super::*;
320
321 #[test]
322 fn interrupted_length_prefixed_buffer_leaves_maximum_length() {
323 let mut buf = Vec::new();
324 let nested = LengthPrefixedBuffer::new(ListLength::U16, &mut buf);
325 nested.buf.push(0xaa);
326 assert_eq!(nested.buf, &vec![0xff, 0xff, 0xaa]);
327 // <- if the buffer is accidentally read here, there is no possiblity
328 // that the contents of the length-prefixed buffer are interpretted
329 // as a subsequent encoding (perhaps allowing injection of a different
330 // extension)
331 drop(nested);
332 assert_eq!(buf, vec![0x00, 0x01, 0xaa]);
333 }
334}
335