Skip to main content

xmtp_id/key_package/
verified_key_package_v2.rs

1use super::WrapperEncryptionExtension;
2use openmls::{
3    credentials::{BasicCredential, errors::BasicCredentialError},
4    key_packages::Lifetime,
5    prelude::{
6        KeyPackage, KeyPackageIn, KeyPackageVerifyError,
7        tls_codec::{Deserialize, Error as TlsCodecError},
8    },
9};
10use openmls_rust_crypto::RustCrypto;
11use prost::Message;
12use std::panic::{self, AssertUnwindSafe};
13use thiserror::Error;
14use xmtp_common::ErrorCode;
15use xmtp_configuration::MLS_PROTOCOL_VERSION;
16use xmtp_configuration::WELCOME_WRAPPER_ENCRYPTION_EXTENSION_ID;
17use xmtp_proto::xmtp::identity::MlsCredential;
18
19#[derive(Debug, Error, ErrorCode)]
20pub enum KeyPackageVerificationError {
21    /// TLS codec error.
22    ///
23    /// MLS TLS encoding/decoding failed. Not retryable.
24    #[error("TLS Codec error: {0}")]
25    TlsError(#[from] TlsCodecError),
26    /// MLS validation error.
27    ///
28    /// Key package verification failed. Not retryable.
29    #[error("mls validation: {0}")]
30    MlsValidation(#[from] KeyPackageVerifyError),
31    /// Wrong credential type.
32    ///
33    /// Unexpected MLS credential type. Not retryable.
34    #[error("wrong credential type")]
35    WrongCredentialType(#[from] BasicCredentialError),
36    #[error(transparent)]
37    #[error_code(inherit)]
38    ConversionError(#[from] xmtp_proto::ConversionError),
39}
40
41impl From<prost::DecodeError> for KeyPackageVerificationError {
42    fn from(value: prost::DecodeError) -> Self {
43        Self::ConversionError(value.into())
44    }
45}
46
47pub struct VerifiedLifetime {
48    pub not_before: u64,
49    pub not_after: u64,
50}
51
52impl From<&Lifetime> for VerifiedLifetime {
53    fn from(value: &Lifetime) -> Self {
54        Self {
55            not_before: value.not_before(),
56            not_after: value.not_after(),
57        }
58    }
59}
60/// A wrapper around the MLS key package struct with some additional fields
61#[derive(Clone, Debug)]
62pub struct VerifiedKeyPackageV2 {
63    pub inner: KeyPackage,
64    pub credential: MlsCredential,
65    pub installation_public_key: Vec<u8>,
66}
67
68impl VerifiedKeyPackageV2 {
69    /// Create a new verified key package from its raw parts.
70    pub fn new(
71        kp: KeyPackage,
72        credential: MlsCredential,
73        installation_public_key: Vec<u8>,
74    ) -> Self {
75        Self {
76            inner: kp,
77            credential,
78            installation_public_key,
79        }
80    }
81
82    /// Create a verified key package from TLS-Serialized bytes.
83    pub fn from_bytes(
84        crypto_provider: &RustCrypto,
85        data: &[u8],
86    ) -> Result<Self, KeyPackageVerificationError> {
87        let kp_in: KeyPackageIn = KeyPackageIn::tls_deserialize_exact(data)?;
88        let kp = kp_in.validate(
89            crypto_provider,
90            MLS_PROTOCOL_VERSION,
91            openmls::prelude::LeafNodeLifetimePolicy::Verify,
92        )?;
93
94        kp.try_into()
95    }
96
97    pub fn installation_id(&self) -> Vec<u8> {
98        self.inner.leaf_node().signature_key().as_slice().to_vec()
99    }
100
101    pub fn hpke_init_key(&self) -> Vec<u8> {
102        self.inner.hpke_init_key().as_slice().to_vec()
103    }
104
105    pub fn wrapper_encryption(
106        &self,
107    ) -> Result<Option<WrapperEncryptionExtension>, KeyPackageVerificationError> {
108        self.inner
109            .extensions()
110            .unknown(WELCOME_WRAPPER_ENCRYPTION_EXTENSION_ID)
111            .map(|ext| ext.try_into().map_err(Into::into))
112            .transpose()
113    }
114
115    pub fn life_time(&self) -> Option<VerifiedLifetime> {
116        let lifetime_result = panic::catch_unwind(AssertUnwindSafe(|| {
117            self.inner.life_time() // This might panic
118        }));
119
120        match lifetime_result {
121            Ok(lifetime) => Some(lifetime.into()),
122            Err(_) => None,
123        }
124    }
125}
126
127impl TryFrom<KeyPackage> for VerifiedKeyPackageV2 {
128    type Error = KeyPackageVerificationError;
129
130    fn try_from(kp: KeyPackage) -> Result<Self, Self::Error> {
131        let leaf_node = kp.leaf_node();
132        let basic_credential = BasicCredential::try_from(leaf_node.credential().clone())?;
133        let pub_key_bytes = leaf_node.signature_key().as_slice().to_vec();
134        let credential = MlsCredential::decode(basic_credential.identity())?;
135
136        Ok(Self::new(kp, credential, pub_key_bytes))
137    }
138}