rustls 0.24.0-dev.1

use alloc::vec::Vec;
use core::fmt::{self, Debug};
use core::marker::PhantomData;
use core::mem;

use pki_types::{CertificateDer, SubjectPublicKeyInfoDer};
use zeroize::Zeroize;

use crate::crypto::cipher::Payload;
use crate::error::InvalidMessage;

/// An arbitrary, unknown-content, u24-length-prefixed payload
#[derive(Clone, Eq, PartialEq)]
pub(crate) struct SizedPayload<'a, L, C: Cardinality = MaybeEmpty> {
    pub(crate) inner: Payload<'a>,
    pub(crate) _marker: PhantomData<(L, C)>,
}

impl<'a, L, C: Cardinality> SizedPayload<'a, L, C> {
    pub(crate) fn into_owned(self) -> SizedPayload<'static, L, C> {
        SizedPayload {
            inner: self.inner.into_owned(),
            _marker: PhantomData,
        }
    }

    pub(crate) fn into_vec(self) -> Vec<u8> {
        self.inner.into_owned().into_vec()
    }

    pub(crate) fn as_mut(&mut self) -> Option<&mut [u8]> {
        match &mut self.inner {
            Payload::Owned(vec) => Some(vec.as_mut_slice()),
            Payload::Borrowed(_) => None,
        }
    }

    pub(crate) fn to_vec(&self) -> Vec<u8> {
        self.inner.bytes().to_vec()
    }

    pub(crate) fn bytes(&'a self) -> &'a [u8] {
        self.inner.bytes()
    }

    pub(crate) fn is_empty(&self) -> bool {
        self.inner.bytes().is_empty()
    }
}

impl<'a, L: PayloadSize<'a>> SizedPayload<'a, L, MaybeEmpty> {
    #[cfg(test)]
    pub(crate) fn empty() -> Self {
        Self {
            inner: Payload::Borrowed(&[]),
            _marker: PhantomData,
        }
    }
}

impl<'a, L: PayloadSize<'a>, C: Cardinality> Codec<'a> for SizedPayload<'a, L, C> {
    fn encode(&self, bytes: &mut Vec<u8>) {
        let inner = self.inner.bytes();
        debug_assert!(inner.len() >= C::MIN);
        debug_assert!(inner.len() <= L::MAX);
        L::length(inner).encode(bytes);
        bytes.extend_from_slice(inner);
    }

    fn read(r: &mut Reader<'a>) -> Result<Self, InvalidMessage> {
        let len = L::read(r)?.into();
        if len < C::MIN {
            return Err(InvalidMessage::IllegalEmptyList("SizedPayload"));
        }
        let mut sub = r.sub(len)?;
        Ok(Self {
            inner: Payload::read(&mut sub),
            _marker: PhantomData,
        })
    }
}

impl<C: Cardinality> Zeroize for SizedPayload<'_, u8, C> {
    #[inline(never)]
    fn zeroize(&mut self) {
        if let Payload::Owned(buf) = &mut self.inner {
            buf.zeroize();
        }
    }
}

impl<'a, L: PayloadSize<'a>, C: Cardinality> From<Payload<'a>> for SizedPayload<'a, L, C> {
    fn from(inner: Payload<'a>) -> Self {
        debug_assert!(inner.bytes().len() >= C::MIN);
        debug_assert!(inner.bytes().len() <= L::MAX);
        Self {
            inner,
            _marker: PhantomData,
        }
    }
}

impl<'a, L: PayloadSize<'a>, C: Cardinality> From<Vec<u8>> for SizedPayload<'a, L, C> {
    fn from(inner: Vec<u8>) -> Self {
        debug_assert!(inner.len() >= C::MIN);
        debug_assert!(inner.len() <= L::MAX);
        Self {
            inner: Payload::Owned(inner),
            _marker: PhantomData,
        }
    }
}

impl<'a, L: PayloadSize<'a>, C: Cardinality> Debug for SizedPayload<'a, L, C> {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        self.inner.fmt(f)
    }
}

impl<'a> PayloadSize<'a> for U24 {
    fn length(bytes: &[u8]) -> Self {
        Self(bytes.len() as u32)
    }

    const MAX: usize = 0xFFFFFF;
}

impl<'a> PayloadSize<'a> for u16 {
    fn length(bytes: &[u8]) -> Self {
        bytes.len() as Self
    }

    const MAX: usize = 0xFFFF;
}

impl<'a> PayloadSize<'a> for u8 {
    fn length(bytes: &[u8]) -> Self {
        bytes.len() as Self
    }

    const MAX: usize = 0xFF;
}

pub(crate) trait PayloadSize<'a>: Codec<'a> + Into<usize> {
    fn length(bytes: &[u8]) -> Self;

    const MAX: usize;
}

pub(crate) trait Cardinality: Clone + Eq + PartialEq {
    const MIN: usize;
}

#[derive(Clone, Eq, PartialEq)]
pub(crate) struct MaybeEmpty;

impl Cardinality for MaybeEmpty {
    const MIN: usize = 0;
}

#[derive(Clone, Eq, PartialEq)]
pub(crate) struct NonEmpty;

impl Cardinality for NonEmpty {
    const MIN: usize = 1;
}

impl<'a> Codec<'a> for CertificateDer<'a> {
    fn encode(&self, bytes: &mut Vec<u8>) {
        let nest = LengthPrefixedBuffer::new(Self::SIZE_LEN, bytes);
        nest.buf.extend(self.as_ref());
    }

    fn read(r: &mut Reader<'a>) -> Result<Self, InvalidMessage> {
        let len = ListLength::NonZeroU24 {
            max: CERTIFICATE_MAX_SIZE_LIMIT,
            empty_error: InvalidMessage::IllegalEmptyList("CertificateDer"),
            too_many_error: InvalidMessage::CertificatePayloadTooLarge,
        }
        .read(r)?;

        let mut sub = r.sub(len)?;
        let body = sub.rest();
        Ok(Self::from(body))
    }
}

impl TlsListElement for CertificateDer<'_> {
    const SIZE_LEN: ListLength = ListLength::U24 {
        max: CERTIFICATE_MAX_SIZE_LIMIT,
        error: InvalidMessage::CertificatePayloadTooLarge,
    };
}

impl<'a> Codec<'a> for SubjectPublicKeyInfoDer<'a> {
    fn encode(&self, bytes: &mut Vec<u8>) {
        let nest = LengthPrefixedBuffer::new(Self::SIZE_LEN, bytes);
        nest.buf.extend(self.as_ref());
    }

    fn read(r: &mut Reader<'a>) -> Result<Self, InvalidMessage> {
        let len = Self::SIZE_LEN.read(r)?;
        let mut sub = r.sub(len)?;
        let body = sub.rest();
        Ok(Self::from(body))
    }
}

impl TlsListElement for SubjectPublicKeyInfoDer<'_> {
    const SIZE_LEN: ListLength = CertificateDer::SIZE_LEN;
}

/// An iterator over a vector of `TlsListElements`.
///
/// All uses _MUST_ exhaust the iterator, as errors may be delayed
/// until the last element.
pub(crate) struct TlsListIter<'a, T: Codec<'a> + TlsListElement + Debug> {
    sub: Reader<'a>,
    _t: PhantomData<T>,
}

impl<'a, T: Codec<'a> + TlsListElement + Debug> TlsListIter<'a, T> {
    pub(crate) fn new(r: &mut Reader<'a>) -> Result<Self, InvalidMessage> {
        let len = T::SIZE_LEN.read(r)?;
        let sub = r.sub(len)?;
        Ok(Self {
            sub,
            _t: PhantomData,
        })
    }
}

impl<'a, T: Codec<'a> + TlsListElement + Debug> Iterator for TlsListIter<'a, T> {
    type Item = Result<T, InvalidMessage>;

    fn next(&mut self) -> Option<Self::Item> {
        match self.sub.any_left() {
            true => Some(T::read(&mut self.sub)),
            false => None,
        }
    }
}

/// Implement `Codec` for lists of elements that implement `TlsListElement`.
///
/// `TlsListElement` provides the size of the length prefix for the list.
impl<'a, T: Codec<'a> + TlsListElement + Debug> Codec<'a> for Vec<T> {
    fn encode(&self, bytes: &mut Vec<u8>) {
        let nest = LengthPrefixedBuffer::new(T::SIZE_LEN, bytes);

        for i in self {
            i.encode(nest.buf);
        }
    }

    fn read(r: &mut Reader<'a>) -> Result<Self, InvalidMessage> {
        let mut ret = Self::new();
        for item in TlsListIter::<T>::new(r)? {
            ret.push(item?);
        }

        Ok(ret)
    }
}

/// Tracks encoding a length-delimited structure in a single pass.
pub(crate) struct LengthPrefixedBuffer<'a> {
    pub(crate) buf: &'a mut Vec<u8>,
    len_offset: usize,
    size_len: ListLength,
}

impl<'a> LengthPrefixedBuffer<'a> {
    /// Inserts a dummy length into `buf`, and remembers where it went.
    ///
    /// After this, the body of the length-delimited structure should be appended to `LengthPrefixedBuffer::buf`.
    /// The length header is corrected in `LengthPrefixedBuffer::drop`.
    pub(crate) fn new(size_len: ListLength, buf: &'a mut Vec<u8>) -> Self {
        let len_offset = buf.len();
        buf.extend(match size_len {
            ListLength::NonZeroU8 { .. } => &[0xff][..],
            ListLength::U16 | ListLength::NonZeroU16 { .. } => &[0xff, 0xff],
            ListLength::U24 { .. } | ListLength::NonZeroU24 { .. } => &[0xff, 0xff, 0xff],
        });

        Self {
            buf,
            len_offset,
            size_len,
        }
    }
}

impl Drop for LengthPrefixedBuffer<'_> {
    /// Goes back and corrects the length previously inserted at the start of the structure.
    fn drop(&mut self) {
        match self.size_len {
            ListLength::NonZeroU8 { .. } => {
                let len = self.buf.len() - self.len_offset - 1;
                debug_assert!(len <= 0xff);
                self.buf[self.len_offset] = len as u8;
            }
            ListLength::U16 | ListLength::NonZeroU16 { .. } => {
                let len = self.buf.len() - self.len_offset - 2;
                debug_assert!(len <= 0xffff);
                let out: &mut [u8; 2] = (&mut self.buf[self.len_offset..self.len_offset + 2])
                    .try_into()
                    .unwrap();
                *out = u16::to_be_bytes(len as u16);
            }
            ListLength::U24 { .. } | ListLength::NonZeroU24 { .. } => {
                let len = self.buf.len() - self.len_offset - 3;
                debug_assert!(len <= 0xff_ffff);
                let len_bytes = u32::to_be_bytes(len as u32);
                let out: &mut [u8; 3] = (&mut self.buf[self.len_offset..self.len_offset + 3])
                    .try_into()
                    .unwrap();
                out.copy_from_slice(&len_bytes[1..]);
            }
        }
    }
}

impl Codec<'_> for u8 {
    fn encode(&self, bytes: &mut Vec<u8>) {
        bytes.push(*self);
    }

    fn read(r: &mut Reader<'_>) -> Result<Self, InvalidMessage> {
        r.take_array("u8").map(|&[byte]| byte)
    }
}

pub(crate) fn put_u16(v: u16, out: &mut [u8]) {
    let out: &mut [u8; 2] = (&mut out[..2]).try_into().unwrap();
    *out = u16::to_be_bytes(v);
}

impl Codec<'_> for u16 {
    fn encode(&self, bytes: &mut Vec<u8>) {
        let mut b16 = [0u8; 2];
        put_u16(*self, &mut b16);
        bytes.extend_from_slice(&b16);
    }

    fn read(r: &mut Reader<'_>) -> Result<Self, InvalidMessage> {
        r.take_array("u16")
            .map(|&[b1, b2]| Self::from_be_bytes([b1, b2]))
    }
}

// Make a distinct type for u24, even though it's a u32 underneath
#[derive(Debug, Copy, Clone)]
pub struct U24(pub u32);

#[cfg(any(target_pointer_width = "32", target_pointer_width = "64"))]
impl From<U24> for usize {
    #[inline]
    fn from(v: U24) -> Self {
        v.0 as Self
    }
}

impl Codec<'_> for U24 {
    fn encode(&self, bytes: &mut Vec<u8>) {
        let be_bytes = u32::to_be_bytes(self.0);
        bytes.extend_from_slice(&be_bytes[1..]);
    }

    fn read(r: &mut Reader<'_>) -> Result<Self, InvalidMessage> {
        r.take_array("u24")
            .map(|&[a, b, c]| Self(u32::from_be_bytes([0, a, b, c])))
    }
}

impl Codec<'_> for u32 {
    fn encode(&self, bytes: &mut Vec<u8>) {
        bytes.extend(Self::to_be_bytes(*self));
    }

    fn read(r: &mut Reader<'_>) -> Result<Self, InvalidMessage> {
        r.take_array("u32")
            .map(|&[a, b, c, d]| Self::from_be_bytes([a, b, c, d]))
    }
}

pub(crate) fn put_u64(v: u64, bytes: &mut [u8]) {
    let bytes: &mut [u8; 8] = (&mut bytes[..8]).try_into().unwrap();
    *bytes = u64::to_be_bytes(v);
}

impl Codec<'_> for u64 {
    fn encode(&self, bytes: &mut Vec<u8>) {
        let mut b64 = [0u8; 8];
        put_u64(*self, &mut b64);
        bytes.extend_from_slice(&b64);
    }

    fn read(r: &mut Reader<'_>) -> Result<Self, InvalidMessage> {
        r.take_array("u64")
            .map(|&[a, b, c, d, e, f, g, h]| Self::from_be_bytes([a, b, c, d, e, f, g, h]))
    }
}

impl Codec<'_> for () {
    fn encode(&self, _: &mut Vec<u8>) {}

    fn read(r: &mut Reader<'_>) -> Result<Self, InvalidMessage> {
        r.all("Empty", |_| Ok(()))
    }
}

/// Trait for implementing encoding and decoding functionality
/// on something.
pub(crate) trait Codec<'a>: Debug + Sized {
    /// Function for encoding itself by appending itself to
    /// the provided vec of bytes.
    fn encode(&self, bytes: &mut Vec<u8>);

    /// Function for decoding itself from the provided reader
    /// will return `Ok` if the decoding was successful or
    /// `Err(InvalidMessage)` if it was not.
    fn read(_: &mut Reader<'a>) -> Result<Self, InvalidMessage>;

    /// Convenience function for encoding the implementation
    /// into a vec and returning it
    fn get_encoding(&self) -> Vec<u8> {
        let mut bytes = Vec::new();
        self.encode(&mut bytes);
        bytes
    }

    /// Function for wrapping a call to the read function in
    /// a Reader for the slice of bytes provided
    ///
    /// Returns `Err(InvalidMessage::TrailingData(_))` if
    /// `Self::read` does not read the entirety of `bytes`.
    fn read_bytes(bytes: &'a [u8]) -> Result<Self, InvalidMessage> {
        Reader::new(bytes).all("read_bytes", Self::read)
    }
}

/// Wrapper over a slice of bytes that allows reading chunks from
/// with the current position state held using a cursor.
///
/// A new reader for a sub section of the buffer can be created
/// using the `sub` function or a section of a certain length can
/// be obtained using the `take` function
pub(crate) struct Reader<'a> {
    /// The underlying buffer storing the readers content
    buffer: &'a [u8],
}

impl<'a> Reader<'a> {
    /// Creates a new Reader of the provided `bytes` slice.
    pub(crate) fn new(buffer: &'a [u8]) -> Self {
        Self { buffer }
    }

    /// Reads all of `buffer` into a type of `T`, checking for trailing data.
    pub(crate) fn all<T, E: From<InvalidMessage>, F: FnOnce(&mut Self) -> Result<T, E>>(
        &mut self,
        type_name: &'static str,
        f: F,
    ) -> Result<T, E> {
        let value = f(self)?;
        match self.any_left() {
            true => Err(InvalidMessage::TrailingData(type_name).into()),
            false => Ok(value),
        }
    }

    /// Attempts to create a new Reader on a sub section of this
    /// readers bytes by taking a slice of the provided `length`
    /// will return `Err(InvalidMessage::MessageTooShort)` if there is not enough bytes
    pub(crate) fn sub(&mut self, length: usize) -> Result<Self, InvalidMessage> {
        match self.take(length) {
            Some(bytes) => Ok(Reader::new(bytes)),
            None => Err(InvalidMessage::MessageTooShort),
        }
    }

    /// Borrow an array of `N` bytes from the buffer.
    ///
    /// If there are not enough bytes remaining `Err(InvalidMessage::MissingData)` is returned instead
    pub(crate) fn take_array<const N: usize>(
        &mut self,
        ty: &'static str,
    ) -> Result<&'a [u8; N], InvalidMessage> {
        match self.buffer.split_first_chunk() {
            Some((chunk, rest)) => {
                self.buffer = rest;
                Ok(chunk)
            }
            _ => Err(InvalidMessage::MissingData(ty)),
        }
    }

    /// Borrow a slice of `length` bytes from the buffer.
    ///
    /// If there are not enough bytes remaining to take the length `None` is returned instead.
    pub(crate) fn take(&mut self, length: usize) -> Option<&'a [u8]> {
        let (out, rest) = self.buffer.split_at_checked(length)?;
        self.buffer = rest;
        Some(out)
    }

    /// Borrows a slice of all the remaining bytes.
    ///
    /// Moves the cursor to the end of the buffer length.
    pub(crate) fn rest(&mut self) -> &'a [u8] {
        mem::take(&mut self.buffer)
    }

    /// Whether the reader has any content left.
    pub(crate) fn any_left(&self) -> bool {
        !self.buffer.is_empty()
    }

    /// Number of bytes that are still able to be read.
    pub(crate) fn left(&self) -> usize {
        self.buffer.len()
    }
}

/// A trait for types that can be encoded and decoded in a list.
///
/// This trait is used to implement `Codec` for `Vec<T>`. Lists in the TLS wire format are
/// prefixed with a length, the size of which depends on the type of the list elements.
/// As such, the `Codec` implementation for `Vec<T>` requires an implementation of this trait
/// for its element type `T`.
pub(crate) trait TlsListElement {
    const SIZE_LEN: ListLength;
}

/// The length of the length prefix for a list.
///
/// The types that appear in lists are limited to three kinds of length prefixes:
/// 1, 2, and 3 bytes. For the latter kind, we require a `TlsListElement` implementer
/// to specify a maximum length and error if the actual length is larger.
pub(crate) enum ListLength {
    /// U8 but non-empty
    NonZeroU8 { empty_error: InvalidMessage },

    /// U16, perhaps empty
    U16,

    /// U16 but non-empty
    NonZeroU16 { empty_error: InvalidMessage },

    /// U24 with imposed upper bound
    U24 { max: usize, error: InvalidMessage },

    /// U24 but non-empty, with imposed upper bound
    NonZeroU24 {
        max: usize,
        empty_error: InvalidMessage,
        too_many_error: InvalidMessage,
    },
}

impl ListLength {
    pub(crate) fn read(&self, r: &mut Reader<'_>) -> Result<usize, InvalidMessage> {
        Ok(match self {
            Self::NonZeroU8 { empty_error } => match usize::from(u8::read(r)?) {
                0 => return Err(*empty_error),
                len => len,
            },
            Self::U16 => usize::from(u16::read(r)?),
            Self::NonZeroU16 { empty_error } => match usize::from(u16::read(r)?) {
                0 => return Err(*empty_error),
                len => len,
            },
            Self::U24 { max, error } => match usize::from(U24::read(r)?) {
                len if len > *max => return Err(*error),
                len => len,
            },
            Self::NonZeroU24 {
                max,
                empty_error,
                too_many_error,
            } => match usize::from(U24::read(r)?) {
                0 => return Err(*empty_error),
                len if len > *max => return Err(*too_many_error),
                len => len,
            },
        })
    }
}

// Format an iterator of u8 into a hex string
pub(crate) fn hex<'a>(
    f: &mut fmt::Formatter<'_>,
    payload: impl IntoIterator<Item = &'a u8>,
) -> fmt::Result {
    for b in payload {
        write!(f, "{b:02x}")?;
    }
    Ok(())
}

/// TLS has a 16MB size limit on any handshake message,
/// plus a 16MB limit on any given certificate.
///
/// We contract that to 64KB to limit the amount of memory allocation
/// that is directly controllable by the peer.
pub(crate) const CERTIFICATE_MAX_SIZE_LIMIT: usize = 0x1_0000;

#[cfg(test)]
mod tests {
    use std::vec;

    use super::*;

    #[test]
    fn interrupted_length_prefixed_buffer_leaves_maximum_length() {
        let mut buf = Vec::new();
        let nested = LengthPrefixedBuffer::new(ListLength::U16, &mut buf);
        nested.buf.push(0xaa);
        assert_eq!(nested.buf, &vec![0xff, 0xff, 0xaa]);
        // <- if the buffer is accidentally read here, there is no possibility
        //    that the contents of the length-prefixed buffer are interpreted
        //    as a subsequent encoding (perhaps allowing injection of a different
        //    extension)
        drop(nested);
        assert_eq!(buf, vec![0x00, 0x01, 0xaa]);
    }
}