openmls/extensions/

mod.rs

//! # Extensions
//!
//! In MLS, extensions appear in the following places:
//!
//! - In [`KeyPackages`](`crate::key_packages`) and [`LeafNode`](`crate::treesync::node::leaf_node::LeafNode`),
//!   to describe client capabilities
//!   and aspects of their participation in the group.
//!
//! - In `GroupInfo`, to inform new members of the group's parameters and to
//!   provide any additional information required to join the group.
//!
//! - In the `GroupContext` object, to ensure that all members of the group have
//!   a consistent view of the parameters in use.
//!
//! Note that `GroupInfo` and `GroupContext` are not exposed via OpenMLS' public
//! API.
//!
//! OpenMLS supports the following extensions:
//!
//! - [`ApplicationIdExtension`] (KeyPackage extension)
//! - [`RatchetTreeExtension`] (GroupInfo extension)
//! - [`RequiredCapabilitiesExtension`] (GroupContext extension)
//! - [`ExternalPubExtension`] (GroupInfo extension)
//! - [`ExternalSendersExtension`] (GroupContext extension)
//! - [`LastResortExtension`] (KeyPackage extension)

use std::{
    collections::HashSet,
    convert::Infallible,
    fmt::Debug,
    io::{Read, Write},
    marker::PhantomData,
};

use serde::{Deserialize, Serialize};

// Private
#[cfg(feature = "extensions-draft")]
mod app_data_dict_extension;
mod application_id_extension;
mod codec;
mod extension_in;
mod external_pub_extension;
mod external_sender_extension;
mod last_resort;
mod ratchet_tree_extension;
mod required_capabilities;
use errors::*;

// Public
pub mod errors;

// Public re-exports
#[cfg(feature = "extensions-draft")]
pub use app_data_dict_extension::{AppDataDictionary, AppDataDictionaryExtension};
pub use application_id_extension::ApplicationIdExtension;
pub use external_pub_extension::ExternalPubExtension;
pub use external_sender_extension::{
    ExternalSender, ExternalSendersExtension, SenderExtensionIndex,
};
pub use last_resort::LastResortExtension;
pub use ratchet_tree_extension::RatchetTreeExtension;
pub use required_capabilities::RequiredCapabilitiesExtension;

use tls_codec::{
    Deserialize as TlsDeserializeTrait, DeserializeBytes, Error, Serialize as TlsSerializeTrait,
    Size, TlsDeserialize, TlsSerialize, TlsSize,
};

use crate::{
    extensions::extension_in::ExtensionIn, group::GroupContext, key_packages::KeyPackage,
    messages::group_info::GroupInfo, treesync::LeafNode,
};

#[cfg(test)]
mod tests;

/// MLS Extension Types
///
/// Copied from draft-ietf-mls-protocol-16:
///
/// | Value            | Name                     | Message(s) | Recommended | Reference |
/// |:-----------------|:-------------------------|:-----------|:------------|:----------|
/// | 0x0000           | RESERVED                 | N/A        | N/A         | RFC XXXX  |
/// | 0x0001           | application_id           | LN         | Y           | RFC XXXX  |
/// | 0x0002           | ratchet_tree             | GI         | Y           | RFC XXXX  |
/// | 0x0003           | required_capabilities    | GC         | Y           | RFC XXXX  |
/// | 0x0004           | external_pub             | GI         | Y           | RFC XXXX  |
/// | 0x0005           | external_senders         | GC         | Y           | RFC XXXX  |
/// | 0xff00  - 0xffff | Reserved for Private Use | N/A        | N/A         | RFC XXXX  |
///
/// Note: OpenMLS does not provide a `Reserved` variant in [ExtensionType].
#[derive(Debug, Copy, Clone, PartialEq, Eq, Hash, Ord, PartialOrd)]
#[cfg_attr(
    feature = "0-8-1-storage-format",
    derive(serde::Serialize, serde::Deserialize)
)]
#[cfg_attr(
    not(feature = "0-8-1-storage-format"),
    derive(
        openmls_serialization_helpers::Serialize,
        openmls_serialization_helpers::Deserialize,
    )
)]
pub enum ExtensionType {
    #[cfg_attr(not(feature = "0-8-1-storage-format"), storage_tag = 0)]
    /// The application id extension allows applications to add an explicit,
    /// application-defined identifier to a KeyPackage.
    ApplicationId,

    #[cfg_attr(not(feature = "0-8-1-storage-format"), storage_tag = 1)]
    /// The ratchet tree extensions provides the whole public state of the
    /// ratchet tree.
    RatchetTree,

    #[cfg_attr(not(feature = "0-8-1-storage-format"), storage_tag = 2)]
    /// The required capabilities extension defines the configuration of a group
    /// that imposes certain requirements on clients in the group.
    RequiredCapabilities,

    #[cfg_attr(not(feature = "0-8-1-storage-format"), storage_tag = 3)]
    /// To join a group via an External Commit, a new member needs a GroupInfo
    /// with an ExternalPub extension present in its extensions field.
    ExternalPub,

    #[cfg_attr(not(feature = "0-8-1-storage-format"), storage_tag = 4)]
    /// Group context extension that contains the credentials and signature keys
    /// of senders that are permitted to send external proposals to the group.
    ExternalSenders,

    #[cfg_attr(not(feature = "0-8-1-storage-format"), storage_tag = 5)]
    /// KeyPackage extension that marks a KeyPackage for use in a last resort
    /// scenario.
    LastResort,

    #[cfg(feature = "extensions-draft")]
    #[cfg_attr(not(feature = "0-8-1-storage-format"), storage_tag = 8)]
    /// AppDataDictionary extension
    AppDataDictionary,

    #[cfg_attr(not(feature = "0-8-1-storage-format"), storage_tag = 7)]
    /// A GREASE extension type for ensuring extensibility.
    Grease(u16),

    #[cfg_attr(not(feature = "0-8-1-storage-format"), storage_tag = 6)]
    /// A currently unknown extension type.
    Unknown(u16),
}

impl ExtensionType {
    /// Returns true for all extension types that are considered "default" by the spec.
    pub(crate) fn is_default(self) -> bool {
        match self {
            ExtensionType::ApplicationId
            | ExtensionType::RatchetTree
            | ExtensionType::RequiredCapabilities
            | ExtensionType::ExternalPub
            | ExtensionType::ExternalSenders => true,
            ExtensionType::LastResort | ExtensionType::Grease(_) | ExtensionType::Unknown(_) => {
                false
            }
            #[cfg(feature = "extensions-draft")]
            ExtensionType::AppDataDictionary => false,
        }
    }

    /// Returns whether an extension type is valid when used in leaf nodes.
    /// Returns [`true`] for unknown extensions.
    //  https://validation.openmls.tech/#valn1601
    pub(crate) fn is_valid_in_leaf_node(self) -> bool {
        match self {
            ExtensionType::LastResort
            | ExtensionType::RatchetTree
            | ExtensionType::RequiredCapabilities
            | ExtensionType::ExternalPub
            | ExtensionType::ExternalSenders => false,
            // GREASE may appear as an extension type in `leaf_node.extensions`
            // and must be tolerated there. It is still subject to the normal rule
            // that it must be declared in `capabilities` (checked separately),
            // so this only permits the type, it does not exempt it from that check.
            ExtensionType::Grease(_) | ExtensionType::Unknown(_) | ExtensionType::ApplicationId => {
                true
            }
            #[cfg(feature = "extensions-draft")]
            ExtensionType::AppDataDictionary => true,
        }
    }
    pub(crate) fn is_valid_in_group_info(self) -> Option<bool> {
        match self {
            ExtensionType::LastResort
            | ExtensionType::RequiredCapabilities
            | ExtensionType::ExternalSenders
            | ExtensionType::ApplicationId => Some(false),
            ExtensionType::RatchetTree | ExtensionType::ExternalPub => Some(true),
            // GREASE is treated like an unknown extension type (tolerated): a
            // GREASE-valued extension used to be reported as `Unknown` here, and
            // must not become stricter now that it maps to `Grease`.
            ExtensionType::Grease(_) | ExtensionType::Unknown(_) => None,
            #[cfg(feature = "extensions-draft")]
            ExtensionType::AppDataDictionary => Some(true),
        }
    }

    pub(crate) fn is_valid_in_key_package(self) -> bool {
        match self {
            ExtensionType::RatchetTree
            | ExtensionType::RequiredCapabilities
            | ExtensionType::ExternalPub
            | ExtensionType::ExternalSenders
            | ExtensionType::ApplicationId => false,
            // GREASE may appear as an extension type in `key_package.extensions`
            // and must be tolerated there. It is still subject to the normal rule
            // that it be declared in `capabilities` (checked separately), so this
            // only permits the type, it does not exempt it from that check.
            ExtensionType::Grease(_) | ExtensionType::Unknown(_) | ExtensionType::LastResort => {
                true
            }
            #[cfg(feature = "extensions-draft")]
            ExtensionType::AppDataDictionary => true,
        }
    }

    pub(crate) fn is_valid_in_group_context(self) -> bool {
        match self {
            ExtensionType::RequiredCapabilities
            | ExtensionType::ExternalSenders
            | ExtensionType::Unknown(_) => true,
            // GREASE is treated like an unknown extension type: structurally
            // allowed to appear here, but (like any unknown extension in the
            // GroupContext) still subject to the per-member support check, which
            // enforces the consensus rule that every member support it. GREASE
            // does not bypass that check.
            ExtensionType::Grease(_) => true,
            #[cfg(feature = "extensions-draft")]
            ExtensionType::AppDataDictionary => true,
            _ => false,
        }
    }

    /// Returns true if this is a GREASE extension type.
    ///
    /// GREASE values are used to ensure implementations properly handle unknown
    /// extension types. See [RFC 9420 Section 13.5](https://www.rfc-editor.org/rfc/rfc9420.html#section-13.5).
    pub fn is_grease(&self) -> bool {
        matches!(self, ExtensionType::Grease(_))
    }
}

impl Size for ExtensionType {
    fn tls_serialized_len(&self) -> usize {
        2
    }
}

impl TlsDeserializeTrait for ExtensionType {
    fn tls_deserialize<R: Read>(bytes: &mut R) -> Result<Self, Error>
    where
        Self: Sized,
    {
        let mut extension_type = [0u8; 2];
        bytes.read_exact(&mut extension_type)?;

        Ok(ExtensionType::from(u16::from_be_bytes(extension_type)))
    }
}

impl DeserializeBytes for ExtensionType {
    fn tls_deserialize_bytes(bytes: &[u8]) -> Result<(Self, &[u8]), Error>
    where
        Self: Sized,
    {
        let mut bytes_ref = bytes;
        let extension_type = ExtensionType::tls_deserialize(&mut bytes_ref)?;
        Ok((extension_type, bytes_ref))
    }
}

impl TlsSerializeTrait for ExtensionType {
    fn tls_serialize<W: Write>(&self, writer: &mut W) -> Result<usize, Error> {
        writer.write_all(&u16::from(*self).to_be_bytes())?;

        Ok(2)
    }
}

impl From<u16> for ExtensionType {
    fn from(a: u16) -> Self {
        match a {
            1 => ExtensionType::ApplicationId,
            2 => ExtensionType::RatchetTree,
            3 => ExtensionType::RequiredCapabilities,
            4 => ExtensionType::ExternalPub,
            5 => ExtensionType::ExternalSenders,
            #[cfg(feature = "extensions-draft")]
            6 => ExtensionType::AppDataDictionary,
            10 => ExtensionType::LastResort,
            unknown if crate::grease::is_grease_value(unknown) => ExtensionType::Grease(unknown),
            unknown => ExtensionType::Unknown(unknown),
        }
    }
}

impl From<ExtensionType> for u16 {
    fn from(value: ExtensionType) -> Self {
        match value {
            ExtensionType::ApplicationId => 1,
            ExtensionType::RatchetTree => 2,
            ExtensionType::RequiredCapabilities => 3,
            ExtensionType::ExternalPub => 4,
            ExtensionType::ExternalSenders => 5,
            #[cfg(feature = "extensions-draft")]
            ExtensionType::AppDataDictionary => 6,
            ExtensionType::LastResort => 10,
            ExtensionType::Grease(value) => value,
            ExtensionType::Unknown(unknown) => unknown,
        }
    }
}

/// # Extension
///
/// An extension is one of the [`Extension`] enum values.
/// The enum provides a set of common functionality for all extensions.
///
/// See the individual extensions for more details on each extension.
///
/// ```c
/// // draft-ietf-mls-protocol-16
/// struct {
///     ExtensionType extension_type;
///     opaque extension_data<V>;
/// } Extension;
/// ```
#[derive(Debug, Clone, PartialEq, Eq)]
#[cfg_attr(
    feature = "0-8-1-storage-format",
    derive(serde::Serialize, serde::Deserialize)
)]
#[cfg_attr(
    not(feature = "0-8-1-storage-format"),
    derive(
        openmls_serialization_helpers::Serialize,
        openmls_serialization_helpers::Deserialize,
    )
)]
pub enum Extension {
    #[cfg_attr(not(feature = "0-8-1-storage-format"), storage_tag = 0)]
    /// An [`ApplicationIdExtension`]
    ApplicationId(ApplicationIdExtension),

    #[cfg_attr(not(feature = "0-8-1-storage-format"), storage_tag = 1)]
    /// A [`RatchetTreeExtension`]
    RatchetTree(RatchetTreeExtension),

    #[cfg_attr(not(feature = "0-8-1-storage-format"), storage_tag = 2)]
    /// A [`RequiredCapabilitiesExtension`]
    RequiredCapabilities(RequiredCapabilitiesExtension),

    #[cfg_attr(not(feature = "0-8-1-storage-format"), storage_tag = 3)]
    /// An [`ExternalPubExtension`]
    ExternalPub(ExternalPubExtension),

    #[cfg_attr(not(feature = "0-8-1-storage-format"), storage_tag = 4)]
    /// An [`ExternalSendersExtension`]
    ExternalSenders(ExternalSendersExtension),

    #[cfg_attr(not(feature = "0-8-1-storage-format"), storage_tag = 7)]
    /// An [`AppDataDictionaryExtension`]
    #[cfg(feature = "extensions-draft")]
    AppDataDictionary(AppDataDictionaryExtension),

    #[cfg_attr(not(feature = "0-8-1-storage-format"), storage_tag = 5)]
    /// A [`LastResortExtension`]
    LastResort(LastResortExtension),

    #[cfg_attr(not(feature = "0-8-1-storage-format"), storage_tag = 6)]
    /// A currently unknown extension.
    Unknown(u16, UnknownExtension),
}

/// A unknown/unparsed extension represented by raw bytes.
#[derive(
    PartialEq, Eq, Clone, Debug, Serialize, Deserialize, TlsSize, TlsSerialize, TlsDeserialize,
)]
pub struct UnknownExtension(pub Vec<u8>);

/// A Extension for Object of type T
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct Extensions<T> {
    unique: Vec<Extension>,
    #[serde(skip)]
    _object: core::marker::PhantomData<T>,
}

#[derive(Clone, Copy, PartialEq, Eq, Debug, Default, TlsSize, TlsSerialize, TlsDeserialize)]
/// Any object
pub struct AnyObject;

impl<T> Default for Extensions<T> {
    fn default() -> Self {
        Self {
            unique: vec![],
            _object: PhantomData,
        }
    }
}

impl<T> Size for Extensions<T> {
    fn tls_serialized_len(&self) -> usize {
        Vec::tls_serialized_len(&self.unique)
    }
}

impl<T> TlsSerializeTrait for Extensions<T> {
    fn tls_serialize<W: Write>(&self, writer: &mut W) -> Result<usize, Error> {
        self.unique.tls_serialize(writer)
    }
}

impl<T: ExtensionValidator<Error: ToString>> TlsDeserializeTrait for Extensions<T>
where
    InvalidExtensionError: From<T::Error>,
{
    fn tls_deserialize<R: Read>(bytes: &mut R) -> Result<Self, Error>
    where
        Self: Sized,
    {
        let candidate: Vec<ExtensionIn<T>> = Vec::tls_deserialize(bytes)?;
        Extensions::<T>::try_from(candidate)
            .map_err(|_| Error::DecodingError("Found duplicate extensions".into()))
    }
}

impl<T: ExtensionValidator<Error: ToString>> DeserializeBytes for Extensions<T>
where
    InvalidExtensionError: From<T::Error>,
{
    fn tls_deserialize_bytes(bytes: &[u8]) -> Result<(Self, &[u8]), Error>
    where
        Self: Sized,
    {
        let mut bytes_ref = bytes;
        let extensions = Extensions::<T>::tls_deserialize(&mut bytes_ref)?;
        Ok((extensions, bytes_ref))
    }
}

impl<T: ExtensionValidator> Extensions<T> {
    /// Create an empty extension list.
    pub fn empty() -> Self {
        Self {
            unique: vec![],
            _object: PhantomData,
        }
    }

    /// Returns an iterator over the extension list.
    pub fn iter(&self) -> impl Iterator<Item = &Extension> {
        self.unique.iter()
    }

    /// Remove an extension from the extension list.
    ///
    /// Returns the removed extension or `None` when there is no extension with
    /// the given extension type.
    pub fn remove(&mut self, extension_type: ExtensionType) -> Option<Extension> {
        if let Some(pos) = self
            .unique
            .iter()
            .position(|ext| ext.extension_type() == extension_type)
        {
            Some(self.unique.remove(pos))
        } else {
            None
        }
    }

    /// Returns `true` iff the extension list contains an extension with the
    /// given extension type.
    pub fn contains(&self, extension_type: ExtensionType) -> bool {
        self.unique
            .iter()
            .any(|ext| ext.extension_type() == extension_type)
    }
}

impl<T> Extensions<T>
where
    T: ExtensionValidator,
    InvalidExtensionError: From<T::Error>,
{
    /// Create an extension list with a single extension.
    pub fn single(extension: Extension) -> Result<Self, InvalidExtensionError> {
        T::validate_extension_type(extension.extension_type())?;
        Ok(Self {
            unique: vec![extension],
            _object: PhantomData,
        })
    }

    /// Create an extension list with multiple extensions.
    ///
    /// This function will fail when the list of extensions contains duplicate
    /// extension types.
    pub fn from_vec(extensions: Vec<Extension>) -> Result<Self, InvalidExtensionError> {
        extensions.try_into()
    }

    /// Validate if the extensions are valid for this context
    pub fn validate<'a>(
        extensions: impl Iterator<Item = &'a Extension>,
    ) -> Result<(), InvalidExtensionError> {
        for ext in extensions {
            T::validate_extension_type(ext.extension_type())?;
        }
        Ok(())
    }

    /// Add an extension to the extension list.
    ///
    /// Returns an error when there already is an extension with the same
    /// extension type.
    pub fn add(&mut self, extension: Extension) -> Result<(), InvalidExtensionError> {
        T::validate_extension_type(extension.extension_type())?;
        if self.contains(extension.extension_type()) {
            return Err(InvalidExtensionError::Duplicate);
        }

        self.unique.push(extension);

        Ok(())
    }

    /// Add an extension to the extension list (or replace an existing one.)
    ///
    /// Returns the replaced extension (if any).
    pub fn add_or_replace(
        &mut self,
        extension: Extension,
    ) -> Result<Option<Extension>, InvalidExtensionError> {
        T::validate_extension_type(extension.extension_type())?;
        let replaced = self.remove(extension.extension_type());
        self.unique.push(extension);
        Ok(replaced)
    }
}

impl Extensions<AnyObject> {
    /// Assume that the extensions contain the given extension type.
    ///
    /// # Safety
    ///
    /// The caller must guarantee that the extensions are of the correct type.
    #[cfg(feature = "unchecked-conversions")]
    pub fn into_unchecked<T>(self) -> Extensions<T> {
        Extensions {
            unique: self.unique,
            _object: PhantomData,
        }
    }
}

mod private {
    /// Used to seal other traits
    pub trait Sealed {}
}

/// Can be implemented by a type to validate extensions.
pub trait ExtensionValidator: private::Sealed {
    /// The error returned by the validator
    type Error;

    /// Check if the extension is valid.
    fn validate_extension_type(ext: ExtensionType) -> Result<(), Self::Error>;
}

impl private::Sealed for AnyObject {}

impl ExtensionValidator for AnyObject {
    type Error = Infallible;

    fn validate_extension_type(_ext: ExtensionType) -> Result<(), Infallible> {
        Ok(())
    }
}

impl<T: ExtensionValidator> TryFrom<Vec<Extension>> for Extensions<T>
where
    InvalidExtensionError: From<T::Error>,
{
    type Error = InvalidExtensionError;

    fn try_from(candidate: Vec<Extension>) -> Result<Self, Self::Error> {
        let mut seen = HashSet::with_capacity(candidate.len());
        for extension in candidate.iter() {
            T::validate_extension_type(extension.extension_type())?;

            if !seen.insert(extension.extension_type()) {
                return Err(InvalidExtensionError::Duplicate);
            }
        }

        Ok(Self {
            unique: candidate,
            _object: PhantomData,
        })
    }
}

impl private::Sealed for GroupInfo {}

// https://validation.openmls.tech/#valn1602
impl ExtensionValidator for GroupInfo {
    type Error = ExtensionTypeNotValidInGroupInfoError;

    fn validate_extension_type(
        extension_type: ExtensionType,
    ) -> Result<(), ExtensionTypeNotValidInGroupInfoError> {
        if extension_type.is_valid_in_group_info() == Some(true)
            || extension_type.is_valid_in_group_info().is_none()
        {
            Ok(())
        } else {
            Err(ExtensionTypeNotValidInGroupInfoError(extension_type))
        }
    }
}

impl private::Sealed for GroupContext {}

// https://validation.openmls.tech/#valn1603
impl ExtensionValidator for GroupContext {
    type Error = ExtensionTypeNotValidInGroupContextError;

    fn validate_extension_type(
        extension_type: ExtensionType,
    ) -> Result<(), ExtensionTypeNotValidInGroupContextError> {
        if extension_type.is_valid_in_group_context() {
            Ok(())
        } else {
            Err(ExtensionTypeNotValidInGroupContextError(extension_type))
        }
    }
}

impl private::Sealed for KeyPackage {}

// https://validation.openmls.tech/#valn1604
impl ExtensionValidator for KeyPackage {
    type Error = ExtensionTypeNotValidInKeyPackageError;

    fn validate_extension_type(
        extension_type: ExtensionType,
    ) -> Result<(), ExtensionTypeNotValidInKeyPackageError> {
        if extension_type.is_valid_in_key_package() {
            Ok(())
        } else {
            Err(ExtensionTypeNotValidInKeyPackageError(extension_type))
        }
    }
}

impl private::Sealed for LeafNode {}

// https://validation.openmls.tech/#valn1601
impl ExtensionValidator for LeafNode {
    type Error = ExtensionTypeNotValidInLeafNodeError;

    fn validate_extension_type(
        extension_type: ExtensionType,
    ) -> Result<(), ExtensionTypeNotValidInLeafNodeError> {
        if extension_type.is_valid_in_leaf_node() {
            Ok(())
        } else {
            Err(ExtensionTypeNotValidInLeafNodeError(extension_type))
        }
    }
}

impl<T> Extensions<T> {
    fn find_by_type(&self, extension_type: ExtensionType) -> Option<&Extension> {
        self.unique
            .iter()
            .find(|ext| ext.extension_type() == extension_type)
    }

    /// Get a reference to the [`ApplicationIdExtension`] if there is any.
    pub fn application_id(&self) -> Option<&ApplicationIdExtension> {
        self.find_by_type(ExtensionType::ApplicationId)
            .and_then(|e| match e {
                Extension::ApplicationId(e) => Some(e),
                _ => None,
            })
    }

    /// Get a reference to the [`RatchetTreeExtension`] if there is any.
    pub fn ratchet_tree(&self) -> Option<&RatchetTreeExtension> {
        self.find_by_type(ExtensionType::RatchetTree)
            .and_then(|e| match e {
                Extension::RatchetTree(e) => Some(e),
                _ => None,
            })
    }

    /// Get a reference to the [`RequiredCapabilitiesExtension`] if there is
    /// any.
    pub fn required_capabilities(&self) -> Option<&RequiredCapabilitiesExtension> {
        self.find_by_type(ExtensionType::RequiredCapabilities)
            .and_then(|e| match e {
                Extension::RequiredCapabilities(e) => Some(e),
                _ => None,
            })
    }

    /// Get a reference to the [`ExternalPubExtension`] if there is any.
    pub fn external_pub(&self) -> Option<&ExternalPubExtension> {
        self.find_by_type(ExtensionType::ExternalPub)
            .and_then(|e| match e {
                Extension::ExternalPub(e) => Some(e),
                _ => None,
            })
    }

    /// Get a reference to the [`ExternalSendersExtension`] if there is any.
    pub fn external_senders(&self) -> Option<&ExternalSendersExtension> {
        self.find_by_type(ExtensionType::ExternalSenders)
            .and_then(|e| match e {
                Extension::ExternalSenders(e) => Some(e),
                _ => None,
            })
    }

    #[cfg(feature = "extensions-draft")]
    /// Get a reference to the [`AppDataDictionaryExtension`] if there is any.
    pub fn app_data_dictionary(&self) -> Option<&AppDataDictionaryExtension> {
        self.find_by_type(ExtensionType::AppDataDictionary)
            .and_then(|e| match e {
                Extension::AppDataDictionary(e) => Some(e),
                _ => None,
            })
    }

    /// Get a reference to the [`UnknownExtension`] with the given type id, if there is any.
    pub fn unknown(&self, extension_type_id: u16) -> Option<&UnknownExtension> {
        let extension_type: ExtensionType = extension_type_id.into();

        match extension_type {
            ExtensionType::Grease(_) | ExtensionType::Unknown(_) => {
                self.find_by_type(extension_type).and_then(|e| match e {
                    Extension::Unknown(_, e) => Some(e),
                    _ => None,
                })
            }
            _ => None,
        }
    }
}

impl Extension {
    /// Get a reference to this extension as [`ApplicationIdExtension`].
    /// Returns an [`ExtensionError::InvalidExtensionType`] if called on an
    /// [`Extension`] that's not an [`ApplicationIdExtension`].
    pub fn as_application_id_extension(&self) -> Result<&ApplicationIdExtension, ExtensionError> {
        match self {
            Self::ApplicationId(e) => Ok(e),
            _ => Err(ExtensionError::InvalidExtensionType(
                "This is not an ApplicationIdExtension".into(),
            )),
        }
    }
    #[cfg(feature = "extensions-draft")]
    /// Get a reference to this extension as [`AppDataDictionaryExtension`].
    /// Returns an [`ExtensionError::InvalidExtensionType`] if called on an
    /// [`Extension`] that's not an [`AppDataDictionaryExtension`].
    pub fn as_app_data_dictionary_extension(
        &self,
    ) -> Result<&AppDataDictionaryExtension, ExtensionError> {
        match self {
            Self::AppDataDictionary(e) => Ok(e),
            _ => Err(ExtensionError::InvalidExtensionType(
                "This is not an AppDataDictionaryExtension".into(),
            )),
        }
    }

    /// Get a reference to this extension as [`RatchetTreeExtension`].
    /// Returns an [`ExtensionError::InvalidExtensionType`] if called on
    /// an [`Extension`] that's not a [`RatchetTreeExtension`].
    pub fn as_ratchet_tree_extension(&self) -> Result<&RatchetTreeExtension, ExtensionError> {
        match self {
            Self::RatchetTree(rte) => Ok(rte),
            _ => Err(ExtensionError::InvalidExtensionType(
                "This is not a RatchetTreeExtension".into(),
            )),
        }
    }

    /// Get a reference to this extension as [`RequiredCapabilitiesExtension`].
    /// Returns an [`ExtensionError::InvalidExtensionType`] error if called on
    /// an [`Extension`] that's not a [`RequiredCapabilitiesExtension`].
    pub fn as_required_capabilities_extension(
        &self,
    ) -> Result<&RequiredCapabilitiesExtension, ExtensionError> {
        match self {
            Self::RequiredCapabilities(e) => Ok(e),
            _ => Err(ExtensionError::InvalidExtensionType(
                "This is not a RequiredCapabilitiesExtension".into(),
            )),
        }
    }

    /// Get a reference to this extension as [`ExternalPubExtension`].
    /// Returns an [`ExtensionError::InvalidExtensionType`] error if called on
    /// an [`Extension`] that's not a [`ExternalPubExtension`].
    pub fn as_external_pub_extension(&self) -> Result<&ExternalPubExtension, ExtensionError> {
        match self {
            Self::ExternalPub(e) => Ok(e),
            _ => Err(ExtensionError::InvalidExtensionType(
                "This is not an ExternalPubExtension".into(),
            )),
        }
    }

    /// Get a reference to this extension as [`ExternalSendersExtension`].
    /// Returns an [`ExtensionError::InvalidExtensionType`] error if called on
    /// an [`Extension`] that's not a [`ExternalSendersExtension`].
    pub fn as_external_senders_extension(
        &self,
    ) -> Result<&ExternalSendersExtension, ExtensionError> {
        match self {
            Self::ExternalSenders(e) => Ok(e),
            _ => Err(ExtensionError::InvalidExtensionType(
                "This is not an ExternalSendersExtension".into(),
            )),
        }
    }

    /// Returns the [`ExtensionType`]
    #[inline]
    pub const fn extension_type(&self) -> ExtensionType {
        match self {
            Extension::ApplicationId(_) => ExtensionType::ApplicationId,
            Extension::RatchetTree(_) => ExtensionType::RatchetTree,
            Extension::RequiredCapabilities(_) => ExtensionType::RequiredCapabilities,
            Extension::ExternalPub(_) => ExtensionType::ExternalPub,
            Extension::ExternalSenders(_) => ExtensionType::ExternalSenders,
            #[cfg(feature = "extensions-draft")]
            Extension::AppDataDictionary(_) => ExtensionType::AppDataDictionary,
            Extension::LastResort(_) => ExtensionType::LastResort,
            // Map GREASE-valued extension types to `Grease`, consistent with
            // `ExtensionType::from(u16)`. Without this an extension carrying a
            // GREASE value would be reported as `Unknown`, so GREASE-aware
            // validation (which ignores `Grease(_)`) would not recognize it.
            Extension::Unknown(kind, _) if crate::grease::is_grease_value(*kind) => {
                ExtensionType::Grease(*kind)
            }
            Extension::Unknown(kind, _) => ExtensionType::Unknown(*kind),
        }
    }
}

macro_rules! impl_from_extensions_validator {
    ($validator:ty, $error:ty) => {
        impl From<Extensions<$validator>> for Extensions<AnyObject> {
            fn from(value: Extensions<$validator>) -> Self {
                Extensions {
                    unique: value.unique,
                    _object: PhantomData,
                }
            }
        }

        impl TryFrom<Extensions<AnyObject>> for Extensions<$validator> {
            type Error = $error;

            fn try_from(value: Extensions<AnyObject>) -> Result<Self, $error> {
                value
                    .unique
                    .iter()
                    .map(Extension::extension_type)
                    .try_for_each(<$validator as ExtensionValidator>::validate_extension_type)?;

                Ok(Extensions {
                    unique: value.unique,
                    _object: PhantomData,
                })
            }
        }
    };
}

impl_from_extensions_validator!(GroupContext, ExtensionTypeNotValidInGroupContextError);
impl_from_extensions_validator!(LeafNode, ExtensionTypeNotValidInLeafNodeError);
impl_from_extensions_validator!(KeyPackage, ExtensionTypeNotValidInKeyPackageError);

#[cfg(any(feature = "test-utils", test))]
impl Extensions<AnyObject> {
    /// Coerces the extensions to an Extensions with the given validator. Unsafe.
    pub(crate) fn coerce<T: ExtensionValidator>(self) -> Extensions<T> {
        Extensions {
            unique: self.unique,
            _object: PhantomData,
        }
    }
}
#[cfg(test)]
mod test {
    use itertools::Itertools;
    use tls_codec::{Deserialize, Serialize, VLBytes};

    use crate::{ciphersuite::HpkePublicKey, extensions::*};

    #[test]
    fn add() {
        let mut extensions: Extensions<AnyObject> = Extensions::default();
        extensions
            .add(Extension::RequiredCapabilities(
                RequiredCapabilitiesExtension::default(),
            ))
            .unwrap();
        assert!(extensions
            .add(Extension::RequiredCapabilities(
                RequiredCapabilitiesExtension::default()
            ))
            .is_err());
    }

    #[test]
    fn grease_extension_type_mapping() {
        // A GREASE-valued extension must report a `Grease` extension type
        // (consistent with `ExtensionType::from(u16)`), not `Unknown`. Otherwise
        // GREASE-aware validation would fail to recognize it and reject peers
        // (e.g. MLS++) that decorate leaf/key-package extensions with GREASE.
        let grease = Extension::Unknown(0x5A5A, UnknownExtension(vec![1, 2, 3]));
        assert_eq!(grease.extension_type(), ExtensionType::Grease(0x5A5A));
        assert!(grease.extension_type().is_grease());

        // A non-GREASE unknown value stays `Unknown`.
        let unknown = Extension::Unknown(0xABCD, UnknownExtension(vec![]));
        assert_eq!(unknown.extension_type(), ExtensionType::Unknown(0xABCD));
    }

    #[test]
    fn grease_extension_must_be_declared_in_capabilities() {
        // GREASE does NOT bypass the capability check: a GREASE extension type is
        // "contained" only if it is advertised in the capabilities (RFC 9420:
        // extensions in leaf_node.extensions/key_package.extensions MUST be in
        // capabilities). The fix that makes this work is that a GREASE-valued
        // extension now reports a `Grease(_)` type that matches the `Grease(_)`
        // parsed into the capabilities list.
        let advertised = crate::treesync::node::leaf_node::Capabilities::new(
            None,
            None,
            Some(&[ExtensionType::Grease(0x5A5A)]),
            None,
            None,
        );
        assert!(advertised.contains_extension_type(&ExtensionType::Grease(0x5A5A)));
        // A GREASE value that is not advertised is not contained.
        assert!(!advertised.contains_extension_type(&ExtensionType::Grease(0xAAAA)));

        // Empty capabilities contain no (non-default) extension, GREASE included.
        let empty =
            crate::treesync::node::leaf_node::Capabilities::new(None, None, None, None, None);
        assert!(!empty.contains_extension_type(&ExtensionType::Grease(0x5A5A)));
        assert!(!empty.contains_extension_type(&ExtensionType::Unknown(0xABCD)));
    }

    #[test]
    fn add_try_from() {
        // Create some extensions with different extension types and test that
        // duplicates are rejected. The extension content does not matter in this test.
        let ext_x = Extension::ApplicationId(ApplicationIdExtension::new(b"Test"));
        let ext_y = Extension::RequiredCapabilities(RequiredCapabilitiesExtension::default());

        let tests = [
            (vec![], true),
            (vec![ext_x.clone()], true),
            (vec![ext_x.clone(), ext_x.clone()], false),
            (vec![ext_x.clone(), ext_x.clone(), ext_x.clone()], false),
            (vec![ext_y.clone()], true),
            (vec![ext_y.clone(), ext_y.clone()], false),
            (vec![ext_y.clone(), ext_y.clone(), ext_y.clone()], false),
            (vec![ext_x.clone(), ext_y.clone()], true),
            (vec![ext_y.clone(), ext_x.clone()], true),
            (vec![ext_x.clone(), ext_x.clone(), ext_y.clone()], false),
            (vec![ext_y.clone(), ext_y.clone(), ext_x.clone()], false),
            (vec![ext_x.clone(), ext_y.clone(), ext_y.clone()], false),
            (vec![ext_y.clone(), ext_x.clone(), ext_x.clone()], false),
            (vec![ext_x.clone(), ext_y.clone(), ext_x.clone()], false),
            (vec![ext_y.clone(), ext_x, ext_y], false),
        ];

        for (test, should_work) in tests.into_iter() {
            // Test `add`.
            {
                let mut extensions: Extensions<AnyObject> = Extensions::default();

                let mut works = true;
                for ext in test.iter() {
                    match extensions.add(ext.clone()) {
                        Ok(_) => {}
                        Err(InvalidExtensionError::Duplicate) => {
                            works = false;
                        }
                        _ => panic!("This should have never happened."),
                    }
                }

                println!("{:?}, {:?}", test.clone(), should_work);
                assert_eq!(works, should_work);
            }

            // Test `try_from`.
            if should_work {
                assert!(Extensions::<AnyObject>::try_from(test).is_ok());
            } else {
                assert!(Extensions::<AnyObject>::try_from(test).is_err());
            }
        }
    }

    #[test]
    fn ensure_ordering() {
        // Create some extensions with different extension types and test
        // that all permutations keep their order after being (de)serialized.
        // The extension content does not matter in this test.
        let ext_x = Extension::ApplicationId(ApplicationIdExtension::new(b"Test"));
        let ext_y = Extension::ExternalPub(ExternalPubExtension::new(HpkePublicKey::new(vec![])));
        let ext_z = Extension::RequiredCapabilities(RequiredCapabilitiesExtension::default());

        for candidate in [ext_x, ext_y, ext_z]
            .into_iter()
            .permutations(3)
            .collect::<Vec<_>>()
        {
            let candidate: Extensions<AnyObject> = Extensions::try_from(candidate).unwrap();
            let bytes = candidate.tls_serialize_detached().unwrap();
            let got = Extensions::tls_deserialize(&mut bytes.as_slice()).unwrap();
            assert_eq!(candidate, got);
        }
    }

    #[test]
    fn that_unknown_extensions_are_de_serialized_correctly() {
        let extension_types = [0x0000u16, 0x0A0A, 0x7A7A, 0xF100, 0xFFFF];
        let extension_datas = [vec![], vec![0], vec![1, 2, 3]];

        for extension_type in extension_types.into_iter() {
            for extension_data in extension_datas.iter() {
                // Construct an unknown extension manually.
                let test = {
                    let mut buf = extension_type.to_be_bytes().to_vec();
                    buf.append(
                        &mut VLBytes::new(extension_data.clone())
                            .tls_serialize_detached()
                            .unwrap(),
                    );
                    buf
                };

                // Test deserialization.
                let got = Extension::tls_deserialize_exact(&test).unwrap();

                match got {
                    Extension::Unknown(got_extension_type, ref got_extension_data) => {
                        assert_eq!(extension_type, got_extension_type);
                        assert_eq!(extension_data, &got_extension_data.0);
                    }
                    other => panic!("Expected `Extension::Unknown`, got {other:?}"),
                }

                // Test serialization.
                let got_serialized = got.tls_serialize_detached().unwrap();
                assert_eq!(test, got_serialized);
            }
        }
    }
}