rustls/msgs/
codec.rs

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

rustls/msgs/
codec.rs
