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

xmtp_id/key_package/
construction.rs
