rustls/msgs/
codec.rs

Help
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(1) {
129            Some(&[byte]) => 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(v);
138}
139
140impl Codec<'_> for u16 {
141    fn encode(&self, bytes: &mut Vec<u8>) {
142        let mut b16 = [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(2) {
149            Some(&[b1, b2]) => 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 = 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(3) {
176            Some(&[a, b, c]) => 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(4) {
189            Some(&[a, b, c, d]) => 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(v);
198}
199
200impl Codec<'_> for u64 {
201    fn encode(&self, bytes: &mut Vec<u8>) {
202        let mut b64 = [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(8) {
209            Some(&[a, b, c, d, e, f, g, h]) => 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 = T::SIZE_LEN.read(r)?;
229        let mut sub = r.sub(len)?;
230        let mut ret = Self::new();
231        while sub.any_left() {
232            ret.push(T::read(&mut sub)?);
233        }
234
235        Ok(ret)
236    }
237}
238
239/// A trait for types that can be encoded and decoded in a list.
240///
241/// This trait is used to implement `Codec` for `Vec<T>`. Lists in the TLS wire format are
242/// prefixed with a length, the size of which depends on the type of the list elements.
243/// As such, the `Codec` implementation for `Vec<T>` requires an implementation of this trait
244/// for its element type `T`.
245pub(crate) trait TlsListElement {
246    const SIZE_LEN: ListLength;
247}
248
249/// The length of the length prefix for a list.
250///
251/// The types that appear in lists are limited to three kinds of length prefixes:
252/// 1, 2, and 3 bytes. For the latter kind, we require a `TlsListElement` implementer
253/// to specify a maximum length and error if the actual length is larger.
254pub(crate) enum ListLength {
255    /// U8 but non-empty
256    NonZeroU8 { empty_error: InvalidMessage },
257
258    /// U16, perhaps empty
259    U16,
260
261    /// U16 but non-empty
262    NonZeroU16 { empty_error: InvalidMessage },
263
264    /// U24 with imposed upper bound
265    U24 { max: usize, error: InvalidMessage },
266}
267
268impl ListLength {
269    pub(crate) fn read(&self, r: &mut Reader<'_>) -> Result<usize, InvalidMessage> {
270        Ok(match self {
271            Self::NonZeroU8 { empty_error } => match usize::from(u8::read(r)?) {
272                0 => return Err(*empty_error),
273                len => len,
274            },
275            Self::U16 => usize::from(u16::read(r)?),
276            Self::NonZeroU16 { empty_error } => match usize::from(u16::read(r)?) {
277                0 => return Err(*empty_error),
278                len => len,
279            },
280            Self::U24 { max, error } => match usize::from(u24::read(r)?) {
281                len if len > *max => return Err(*error),
282                len => len,
283            },
284        })
285    }
286}
287
288/// Tracks encoding a length-delimited structure in a single pass.
289pub(crate) struct LengthPrefixedBuffer<'a> {
290    pub(crate) buf: &'a mut Vec<u8>,
291    len_offset: usize,
292    size_len: ListLength,
293}
294
295impl<'a> LengthPrefixedBuffer<'a> {
296    /// Inserts a dummy length into `buf`, and remembers where it went.
297    ///
298    /// After this, the body of the length-delimited structure should be appended to `LengthPrefixedBuffer::buf`.
299    /// The length header is corrected in `LengthPrefixedBuffer::drop`.
300    pub(crate) fn new(size_len: ListLength, buf: &'a mut Vec<u8>) -> Self {
301        let len_offset = buf.len();
302        buf.extend(match size_len {
303            ListLength::NonZeroU8 { .. } => &[0xff][..],
304            ListLength::U16 | ListLength::NonZeroU16 { .. } => &[0xff, 0xff],
305            ListLength::U24 { .. } => &[0xff, 0xff, 0xff],
306        });
307
308        Self {
309            buf,
310            len_offset,
311            size_len,
312        }
313    }
314}
315
316impl Drop for LengthPrefixedBuffer<'_> {
317    /// Goes back and corrects the length previously inserted at the start of the structure.
318    fn drop(&mut self) {
319        match self.size_len {
320            ListLength::NonZeroU8 { .. } => {
321                let len = self.buf.len() - self.len_offset - 1;
322                debug_assert!(len <= 0xff);
323                self.buf[self.len_offset] = len as u8;
324            }
325            ListLength::U16 | ListLength::NonZeroU16 { .. } => {
326                let len = self.buf.len() - self.len_offset - 2;
327                debug_assert!(len <= 0xffff);
328                let out: &mut [u8; 2] = (&mut self.buf[self.len_offset..self.len_offset + 2])
329                    .try_into()
330                    .unwrap();
331                *out = u16::to_be_bytes(len as u16);
332            }
333            ListLength::U24 { .. } => {
334                let len = self.buf.len() - self.len_offset - 3;
335                debug_assert!(len <= 0xff_ffff);
336                let len_bytes = u32::to_be_bytes(len as u32);
337                let out: &mut [u8; 3] = (&mut self.buf[self.len_offset..self.len_offset + 3])
338                    .try_into()
339                    .unwrap();
340                out.copy_from_slice(&len_bytes[1..]);
341            }
342        }
343    }
344}
345
346#[cfg(test)]
347mod tests {
348    use std::prelude::v1::*;
349    use std::vec;
350
351    use super::*;
352
353    #[test]
354    fn interrupted_length_prefixed_buffer_leaves_maximum_length() {
355        let mut buf = Vec::new();
356        let nested = LengthPrefixedBuffer::new(ListLength::U16, &mut buf);
357        nested.buf.push(0xaa);
358        assert_eq!(nested.buf, &vec![0xff, 0xff, 0xaa]);
359        // <- if the buffer is accidentally read here, there is no possibility
360        //    that the contents of the length-prefixed buffer are interpreted
361        //    as a subsequent encoding (perhaps allowing injection of a different
362        //    extension)
363        drop(nested);
364        assert_eq!(buf, vec![0x00, 0x01, 0xaa]);
365    }
366}
rustls/msgs/codec.rs

rustls/msgs/
codec.rs
