Skip to main content

openmls/extensions/
extension_in.rs

1use std::{collections::hash_set::HashSet, io::Read, marker::PhantomData};
2
3use tls_codec::{Deserialize, Size, VLBytes};
4
5use crate::{
6    extensions::{
7        codec::deserialize_extension_exact, Extension, ExtensionType, Extensions,
8        InvalidExtensionError, UnknownExtension,
9    },
10    prelude::ExtensionValidator,
11};
12
13/// The initial capacity of the HashSet used to aggregate which extension types are in use.
14/// We pre-allocate using with_capacity to avoid heap re-allocations. We could just use the number
15/// of extensions in the list, but that is a user-provided number, so the user could make us
16/// allocate a lot of memory. Therefore, we make this a fixed number.
17///
18/// On the number choice:
19/// - Most applications are unlikely to use many extension types in a single list (>10 or so)
20/// - The entries are small (2 bytes each), so overestimating is not very expensive.
21///   However, the allocation size doubles at 29 from ~100B to ~200B, and 28 extension types should
22///   be enough anyway.
23const EXTENSION_TYPE_SET_INITIAL_CAP: usize = 28;
24
25/// An [`Extension`] decoded in the context of `T`, i.e. its extension type
26/// was checked with `T::validate_extension_type` during deserialization.
27/// Only used as an intermediate when deserializing [`Extensions<T>`].
28///
29/// [`Extensions<T>`]: crate::extensions::Extensions
30#[repr(transparent)]
31pub struct ExtensionIn<T>(Extension, PhantomData<T>);
32
33impl<T: ExtensionValidator> ExtensionIn<T> {
34    /// Extracts the Extension
35    pub fn into_extension(self) -> Extension {
36        self.0
37    }
38
39    /// Returns the [`ExtensionType`] of the extension
40    pub fn extension_type(&self) -> ExtensionType {
41        self.0.extension_type()
42    }
43}
44
45impl<T: ExtensionValidator> TryFrom<Vec<ExtensionIn<T>>> for Extensions<T>
46where
47    InvalidExtensionError: From<T::Error>,
48{
49    type Error = InvalidExtensionError;
50
51    fn try_from(candidate: Vec<ExtensionIn<T>>) -> Result<Self, Self::Error> {
52        let mut seen = HashSet::with_capacity(EXTENSION_TYPE_SET_INITIAL_CAP);
53        for extension in candidate.iter() {
54            if !seen.insert(extension.extension_type()) {
55                return Err(InvalidExtensionError::Duplicate);
56            }
57        }
58
59        Ok(Self {
60            unique: candidate
61                .into_iter()
62                .map(ExtensionIn::into_extension)
63                .collect(),
64            _object: PhantomData,
65        })
66    }
67}
68
69impl<T: ExtensionValidator> Size for ExtensionIn<T> {
70    fn tls_serialized_len(&self) -> usize {
71        self.0.tls_serialized_len()
72    }
73}
74
75impl<T: ExtensionValidator<Error: ToString>> Deserialize for ExtensionIn<T> {
76    fn tls_deserialize<R: Read>(bytes: &mut R) -> Result<Self, tls_codec::Error>
77    where
78        Self: Sized,
79    {
80        // Read the extension type and extension data.
81        let extension_type = ExtensionType::tls_deserialize(bytes)?;
82        let extension_data = VLBytes::tls_deserialize(bytes)?;
83
84        // Ensure the extension type is valid in this context
85        T::validate_extension_type(extension_type)
86            .map_err(|e| tls_codec::Error::DecodingError(e.to_string()))?;
87
88        // Now deserialize the extension itself from the extension data.
89        let extension_data = extension_data.as_slice();
90        let extension = match extension_type {
91            ExtensionType::ApplicationId => {
92                Extension::ApplicationId(deserialize_extension_exact(extension_data)?)
93            }
94            ExtensionType::RatchetTree => {
95                Extension::RatchetTree(deserialize_extension_exact(extension_data)?)
96            }
97            ExtensionType::RequiredCapabilities => {
98                Extension::RequiredCapabilities(deserialize_extension_exact(extension_data)?)
99            }
100            ExtensionType::ExternalPub => {
101                Extension::ExternalPub(deserialize_extension_exact(extension_data)?)
102            }
103            ExtensionType::ExternalSenders => {
104                Extension::ExternalSenders(deserialize_extension_exact(extension_data)?)
105            }
106            #[cfg(feature = "extensions-draft")]
107            ExtensionType::AppDataDictionary => {
108                Extension::AppDataDictionary(deserialize_extension_exact(extension_data)?)
109            }
110            ExtensionType::LastResort => {
111                Extension::LastResort(deserialize_extension_exact(extension_data)?)
112            }
113            ExtensionType::Grease(grease) | ExtensionType::Unknown(grease) => {
114                Extension::Unknown(grease, UnknownExtension(extension_data.to_vec()))
115            }
116        };
117
118        Ok(Self(extension, PhantomData))
119    }
120}