1use alloc::vec::Vec;
2use core::fmt::Debug;
3
4use crate::error::InvalidMessage;
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 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: &'a [u8]) -> Self {
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<Self, 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) -> &'a [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<&'a [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<'a>: 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<'a>) -> 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 ///
111 /// Returns `Err(InvalidMessage::ExcessData(_))` if
112 /// `Self::read` does not read the entirety of `bytes`.
113 fn read_bytes(bytes: &'a [u8]) -> Result<Self, InvalidMessage> {
114 let mut reader = Reader::init(bytes);
115 Self::read(&mut reader).and_then(|r| {
116 reader.expect_empty("read_bytes")?;
117 Ok(r)
118 })
119 }
120}
121
122impl Codec<'_> for u8 {
123 fn encode(&self, bytes: &mut Vec<u8>) {
124 bytes.push(*self);
125 }
126
127 fn read(r: &mut Reader<'_>) -> Result<Self, InvalidMessage> {
128 match r.take(length:1) {
129 Some(&[byte: u8]) => Ok(byte),
130 _ => Err(InvalidMessage::MissingData("u8")),
131 }
132 }
133}
134
135pub(crate) fn put_u16(v: u16, out: &mut [u8]) {
136 let out: &mut [u8; 2] = (&mut out[..2]).try_into().unwrap();
137 *out = u16::to_be_bytes(self:v);
138}
139
140impl Codec<'_> for u16 {
141 fn encode(&self, bytes: &mut Vec<u8>) {
142 let mut b16: [u8; 2] = [0u8; 2];
143 put_u16(*self, &mut b16);
144 bytes.extend_from_slice(&b16);
145 }
146
147 fn read(r: &mut Reader<'_>) -> Result<Self, InvalidMessage> {
148 match r.take(length:2) {
149 Some(&[b1: u8, b2: u8]) => Ok(Self::from_be_bytes([b1, b2])),
150 _ => Err(InvalidMessage::MissingData("u16")),
151 }
152 }
153}
154
155// Make a distinct type for u24, even though it's a u32 underneath
156#[allow(non_camel_case_types)]
157#[derive(Debug, Copy, Clone)]
158pub struct u24(pub u32);
159
160#[cfg(any(target_pointer_width = "32", target_pointer_width = "64"))]
161impl From<u24> for usize {
162 #[inline]
163 fn from(v: u24) -> Self {
164 v.0 as Self
165 }
166}
167
168impl Codec<'_> for u24 {
169 fn encode(&self, bytes: &mut Vec<u8>) {
170 let be_bytes: [u8; _] = u32::to_be_bytes(self.0);
171 bytes.extend_from_slice(&be_bytes[1..]);
172 }
173
174 fn read(r: &mut Reader<'_>) -> Result<Self, InvalidMessage> {
175 match r.take(length:3) {
176 Some(&[a: u8, b: u8, c: u8]) => Ok(Self(u32::from_be_bytes([0, a, b, c]))),
177 _ => Err(InvalidMessage::MissingData("u24")),
178 }
179 }
180}
181
182impl Codec<'_> for u32 {
183 fn encode(&self, bytes: &mut Vec<u8>) {
184 bytes.extend(Self::to_be_bytes(*self));
185 }
186
187 fn read(r: &mut Reader<'_>) -> Result<Self, InvalidMessage> {
188 match r.take(length:4) {
189 Some(&[a: u8, b: u8, c: u8, d: u8]) => Ok(Self::from_be_bytes([a, b, c, d])),
190 _ => Err(InvalidMessage::MissingData("u32")),
191 }
192 }
193}
194
195pub(crate) fn put_u64(v: u64, bytes: &mut [u8]) {
196 let bytes: &mut [u8; 8] = (&mut bytes[..8]).try_into().unwrap();
197 *bytes = u64::to_be_bytes(self:v);
198}
199
200impl Codec<'_> for u64 {
201 fn encode(&self, bytes: &mut Vec<u8>) {
202 let mut b64: [u8; 8] = [0u8; 8];
203 put_u64(*self, &mut b64);
204 bytes.extend_from_slice(&b64);
205 }
206
207 fn read(r: &mut Reader<'_>) -> Result<Self, InvalidMessage> {
208 match r.take(length:8) {
209 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])),
210 _ => Err(InvalidMessage::MissingData("u64")),
211 }
212 }
213}
214
215/// Implement `Codec` for lists of elements that implement `TlsListElement`.
216///
217/// `TlsListElement` provides the size of the length prefix for the list.
218impl<'a, T: Codec<'a> + TlsListElement + Debug> Codec<'a> for Vec<T> {
219 fn encode(&self, bytes: &mut Vec<u8>) {
220 let nest = LengthPrefixedBuffer::new(T::SIZE_LEN, bytes);
221
222 for i in self {
223 i.encode(nest.buf);
224 }
225 }
226
227 fn read(r: &mut Reader<'a>) -> Result<Self, InvalidMessage> {
228 let len = match T::SIZE_LEN {
229 ListLength::U8 => usize::from(u8::read(r)?),
230 ListLength::U16 => usize::from(u16::read(r)?),
231 ListLength::U24 { max, error } => match usize::from(u24::read(r)?) {
232 len if len > max => return Err(error),
233 len => len,
234 },
235 };
236
237 let mut sub = r.sub(len)?;
238 let mut ret = Self::new();
239 while sub.any_left() {
240 ret.push(T::read(&mut sub)?);
241 }
242
243 Ok(ret)
244 }
245}
246
247/// A trait for types that can be encoded and decoded in a list.
248///
249/// This trait is used to implement `Codec` for `Vec<T>`. Lists in the TLS wire format are
250/// prefixed with a length, the size of which depends on the type of the list elements.
251/// As such, the `Codec` implementation for `Vec<T>` requires an implementation of this trait
252/// for its element type `T`.
253pub(crate) trait TlsListElement {
254 const SIZE_LEN: ListLength;
255}
256
257/// The length of the length prefix for a list.
258///
259/// The types that appear in lists are limited to three kinds of length prefixes:
260/// 1, 2, and 3 bytes. For the latter kind, we require a `TlsListElement` implementer
261/// to specify a maximum length and error if the actual length is larger.
262pub(crate) enum ListLength {
263 U8,
264 U16,
265 U24 { max: usize, error: InvalidMessage },
266}
267
268/// Tracks encoding a length-delimited structure in a single pass.
269pub(crate) struct LengthPrefixedBuffer<'a> {
270 pub(crate) buf: &'a mut Vec<u8>,
271 len_offset: usize,
272 size_len: ListLength,
273}
274
275impl<'a> LengthPrefixedBuffer<'a> {
276 /// Inserts a dummy length into `buf`, and remembers where it went.
277 ///
278 /// After this, the body of the length-delimited structure should be appended to `LengthPrefixedBuffer::buf`.
279 /// The length header is corrected in `LengthPrefixedBuffer::drop`.
280 pub(crate) fn new(size_len: ListLength, buf: &'a mut Vec<u8>) -> Self {
281 let len_offset: usize = buf.len();
282 buf.extend(iter:match size_len {
283 ListLength::U8 => &[0xff][..],
284 ListLength::U16 => &[0xff, 0xff],
285 ListLength::U24 { .. } => &[0xff, 0xff, 0xff],
286 });
287
288 Self {
289 buf,
290 len_offset,
291 size_len,
292 }
293 }
294}
295
296impl Drop for LengthPrefixedBuffer<'_> {
297 /// Goes back and corrects the length previously inserted at the start of the structure.
298 fn drop(&mut self) {
299 match self.size_len {
300 ListLength::U8 => {
301 let len = self.buf.len() - self.len_offset - 1;
302 debug_assert!(len <= 0xff);
303 self.buf[self.len_offset] = len as u8;
304 }
305 ListLength::U16 => {
306 let len = self.buf.len() - self.len_offset - 2;
307 debug_assert!(len <= 0xffff);
308 let out: &mut [u8; 2] = (&mut self.buf[self.len_offset..self.len_offset + 2])
309 .try_into()
310 .unwrap();
311 *out = u16::to_be_bytes(len as u16);
312 }
313 ListLength::U24 { .. } => {
314 let len = self.buf.len() - self.len_offset - 3;
315 debug_assert!(len <= 0xff_ffff);
316 let len_bytes = u32::to_be_bytes(len as u32);
317 let out: &mut [u8; 3] = (&mut self.buf[self.len_offset..self.len_offset + 3])
318 .try_into()
319 .unwrap();
320 out.copy_from_slice(&len_bytes[1..]);
321 }
322 }
323 }
324}
325
326#[cfg(test)]
327mod tests {
328 use std::prelude::v1::*;
329 use std::vec;
330
331 use super::*;
332
333 #[test]
334 fn interrupted_length_prefixed_buffer_leaves_maximum_length() {
335 let mut buf = Vec::new();
336 let nested = LengthPrefixedBuffer::new(ListLength::U16, &mut buf);
337 nested.buf.push(0xaa);
338 assert_eq!(nested.buf, &vec![0xff, 0xff, 0xaa]);
339 // <- if the buffer is accidentally read here, there is no possibility
340 // that the contents of the length-prefixed buffer are interpreted
341 // as a subsequent encoding (perhaps allowing injection of a different
342 // extension)
343 drop(nested);
344 assert_eq!(buf, vec![0x00, 0x01, 0xaa]);
345 }
346}
347