xmtp_id/key_package/
construction.rs

1use super::{WelcomePointersExtension, WrapperAlgorithm, WrapperEncryptionExtension};
2use openmls::{
3    credentials::{Credential, CredentialWithKey},
4    extensions::{
5        ApplicationIdExtension, Extension, ExtensionType, Extensions, LastResortExtension,
6    },
7    key_packages::{KeyPackage, KeyPackageBundle, Lifetime},
8    messages::proposals::ProposalType,
9    prelude::{Capabilities, HpkeKeyPair, LeafNode},
10};
11use openmls_traits::OpenMlsProvider;
12use xmtp_configuration::{
13    CIPHERSUITE, GROUP_MEMBERSHIP_EXTENSION_ID, GROUP_PERMISSIONS_EXTENSION_ID,
14    MUTABLE_METADATA_EXTENSION_ID, WELCOME_POINTEE_ENCRYPTION_AEAD_TYPES_EXTENSION_ID,
15    WELCOME_WRAPPER_ENCRYPTION_EXTENSION_ID,
16};
17use xmtp_cryptography::{
18    GeneratePostQuantumKeyError, XmtpInstallationCredential, generate_post_quantum_key,
19};
20
21#[derive(Debug, thiserror::Error)]
22pub enum KeyPackageConstructionError {
23    #[error(transparent)]
24    Generation(#[from] openmls::key_packages::errors::KeyPackageNewError),
25    #[error(transparent)]
26    InvalidExtension(#[from] openmls::prelude::InvalidExtensionError),
27    #[error(transparent)]
28    Encode(#[from] prost::EncodeError),
29    #[error(transparent)]
30    PostQuantum(#[from] GeneratePostQuantumKeyError),
31}
32
33/// Options for construction. These do not change admission policy.
34pub struct KeyPackageOptions {
35    pub include_post_quantum: bool,
36    pub welcome_pointers: bool,
37    pub app_data_dictionary: bool,
38    pub lifetime: Option<Lifetime>,
39}
40impl Default for KeyPackageOptions {
41    fn default() -> Self {
42        Self {
43            include_post_quantum: false,
44            welcome_pointers: true,
45            app_data_dictionary: true,
46            lifetime: None,
47        }
48    }
49}
50pub struct GeneratedKeyPackage {
51    pub bundle: KeyPackageBundle,
52    pub post_quantum_keypair: Option<HpkeKeyPair>,
53}
54pub fn build_post_quantum_public_key_extension(
55    public_key: &[u8],
56) -> Result<Extension, prost::EncodeError> {
57    WrapperEncryptionExtension::new(WrapperAlgorithm::XWingMLKEM768Draft6, public_key.to_vec())
58        .try_into()
59}
60
61/// Build a package with the supplied provider. Client bookkeeping is separate.
62pub fn build_key_package(
63    inbox_id: &str,
64    credential: Credential,
65    installation_keys: &XmtpInstallationCredential,
66    provider: &impl OpenMlsProvider,
67    options: KeyPackageOptions,
68) -> Result<GeneratedKeyPackage, KeyPackageConstructionError> {
69    let last_resort = Extension::LastResort(LastResortExtension::default());
70    let welcome_pointee_encryption_aead_types =
71        WelcomePointersExtension::available_types().try_into()?;
72    let mut extensions = vec![last_resort, welcome_pointee_encryption_aead_types];
73    if !options.welcome_pointers {
74        extensions.pop();
75    }
76    let mut post_quantum_keypair = None;
77    if options.include_post_quantum {
78        let keypair = generate_post_quantum_key()?;
79        extensions.push(build_post_quantum_public_key_extension(&keypair.public)?);
80        post_quantum_keypair = Some(keypair);
81    }
82    let key_package_extensions = Extensions::from_vec(extensions)?;
83
84    let application_id = Extension::ApplicationId(ApplicationIdExtension::new(inbox_id.as_bytes()));
85    let leaf_node_extensions = Extensions::<LeafNode>::single(application_id)?;
86
87    let mut capability_extensions = vec![
88        ExtensionType::LastResort,
89        ExtensionType::ApplicationId,
90        ExtensionType::ImmutableMetadata,
91        // Default capabilities let clients join groups that use AppDataUpdate.
92        ExtensionType::AppDataDictionary,
93        ExtensionType::Unknown(GROUP_PERMISSIONS_EXTENSION_ID),
94        ExtensionType::Unknown(MUTABLE_METADATA_EXTENSION_ID),
95        ExtensionType::Unknown(GROUP_MEMBERSHIP_EXTENSION_ID),
96        ExtensionType::Unknown(WELCOME_WRAPPER_ENCRYPTION_EXTENSION_ID),
97        ExtensionType::Unknown(WELCOME_POINTEE_ENCRYPTION_AEAD_TYPES_EXTENSION_ID),
98    ];
99    // Fixtures can model clients that do not support AppDataDictionary.
100    if !options.app_data_dictionary {
101        capability_extensions.retain(|e| *e != ExtensionType::AppDataDictionary);
102    }
103    // Defaults preserve both proposal capabilities advertised by clients.
104    let capabilities = Capabilities::new(
105        None,
106        Some(&[CIPHERSUITE]),
107        Some(&capability_extensions),
108        Some(&[
109            ProposalType::GroupContextExtensions,
110            ProposalType::AppDataUpdate,
111        ]),
112        None,
113    );
114
115    let kp_builder = KeyPackage::builder()
116        .leaf_node_capabilities(capabilities)
117        .leaf_node_extensions(leaf_node_extensions)
118        .key_package_extensions(key_package_extensions);
119
120    let kp_builder = if let Some(lifetime) = options.lifetime {
121        kp_builder.key_package_lifetime(lifetime)
122    } else {
123        kp_builder
124    };
125
126    let kp = kp_builder.build(
127        CIPHERSUITE,
128        provider,
129        installation_keys,
130        CredentialWithKey {
131            credential,
132            signature_key: installation_keys.public_slice().into(),
133        },
134    )?;
135
136    Ok(GeneratedKeyPackage {
137        bundle: kp,
138        post_quantum_keypair,
139    })
140}
141
142#[cfg(test)]
143mod tests {
144    use super::*;
145    use crate::key_package::{VerifiedKeyPackageV2, create_credential};
146    use openmls::prelude::tls_codec::Serialize;
147    use openmls_rust_crypto::OpenMlsRustCrypto;
148
149    #[xmtp_common::test(unwrap_try = true)]
150    fn generated_package_preserves_options_and_verifies() {
151        for include_post_quantum in [false, true] {
152            for capabilities in [false, true] {
153                let provider = OpenMlsRustCrypto::default();
154                let key = XmtpInstallationCredential::new();
155                let generated = build_key_package(
156                    "inbox",
157                    create_credential("inbox"),
158                    &key,
159                    &provider,
160                    KeyPackageOptions {
161                        include_post_quantum,
162                        welcome_pointers: capabilities,
163                        app_data_dictionary: capabilities,
164                        lifetime: Some(Lifetime::new(3600)),
165                    },
166                )?;
167                let bytes = generated.bundle.key_package().tls_serialize_detached()?;
168                let verified = VerifiedKeyPackageV2::from_bytes(provider.crypto(), &bytes)?;
169                assert_eq!(verified.credential.inbox_id, "inbox");
170                assert_eq!(verified.installation_public_key, key.public_slice());
171                let leaf = verified.inner.leaf_node();
172                assert_eq!(
173                    leaf.capabilities()
174                        .extensions()
175                        .contains(&ExtensionType::AppDataDictionary),
176                    capabilities
177                );
178                assert!(
179                    leaf.capabilities()
180                        .proposals()
181                        .contains(&ProposalType::AppDataUpdate)
182                );
183                assert!(
184                    leaf.capabilities()
185                        .proposals()
186                        .contains(&ProposalType::GroupContextExtensions)
187                );
188                let pointer = verified
189                    .inner
190                    .extensions()
191                    .unknown(WELCOME_POINTEE_ENCRYPTION_AEAD_TYPES_EXTENSION_ID);
192                assert_eq!(pointer.is_some(), capabilities);
193                let wrapper = verified.wrapper_encryption()?;
194                assert_eq!(wrapper.is_some(), include_post_quantum);
195                if let Some(pair) = generated.post_quantum_keypair {
196                    assert_eq!(wrapper.unwrap().pub_key_bytes, pair.public);
197                    assert!(!pair.private.is_empty());
198                }
199            }
200        }
201    }
202
203    #[xmtp_common::test(unwrap_try = true)]
204    fn credential_shape_is_not_a_key_package_admission_rule() {
205        let provider = OpenMlsRustCrypto::default();
206        let key = XmtpInstallationCredential::new();
207        for inbox in ["", "not-a-hex-inbox"] {
208            let generated = build_key_package(
209                inbox,
210                create_credential(inbox),
211                &key,
212                &provider,
213                Default::default(),
214            )?;
215            let bytes = generated.bundle.key_package().tls_serialize_detached()?;
216            assert_eq!(
217                VerifiedKeyPackageV2::from_bytes(provider.crypto(), &bytes)?
218                    .credential
219                    .inbox_id,
220                inbox
221            );
222        }
223    }
224
225    #[xmtp_common::test(unwrap_try = true)]
226    fn configured_expired_lifetime_still_fails_verification() {
227        let provider = OpenMlsRustCrypto::default();
228        let key = XmtpInstallationCredential::new();
229        let generated = build_key_package(
230            "inbox",
231            create_credential("inbox"),
232            &key,
233            &provider,
234            KeyPackageOptions {
235                lifetime: Some(Lifetime::init(1, 2)),
236                ..Default::default()
237            },
238        )?;
239        let bytes = generated.bundle.key_package().tls_serialize_detached()?;
240        assert!(matches!(
241            VerifiedKeyPackageV2::from_bytes(provider.crypto(), &bytes),
242            Err(crate::key_package::KeyPackageVerificationError::MlsValidation(_))
243        ));
244    }
245}
xmtp_id/key_package/construction.rs

xmtp_id/key_package/
construction.rs
