Skip to main content

xmtp_mls_identity/identity_lifecycle/
key_package.rs

1//! Building and storing an installation's key packages.
2
3use derive_builder::Builder;
4use openmls::{key_packages::KeyPackage, prelude::Credential as OpenMlsCredential};
5#[cfg(any(test, feature = "test-hooks"))]
6use std::sync::atomic::Ordering;
7use xmtp_cryptography::XmtpInstallationCredential;
8use xmtp_db::MlsProviderExt;
9use xmtp_mls_core::identity::{IdentityError, store_key_package_references};
10
11/// A new key package and, if requested, its post-quantum public key.
12pub struct NewKeyPackageResult {
13    /// The key package.
14    pub key_package: KeyPackage,
15    /// The post-quantum public key, if one was generated.
16    pub pq_pub_key: Option<Vec<u8>>,
17}
18
19/// Whether key packages built in its scope advertise welcome pointees, with a
20/// count of the packages that read it so a test can tell that it ran.
21#[cfg(any(test, feature = "test-hooks"))]
22#[doc(hidden)]
23#[derive(Clone, Debug)]
24pub struct WelcomePointersHook {
25    enabled: bool,
26    reads: std::sync::Arc<std::sync::atomic::AtomicUsize>,
27}
28
29#[cfg(any(test, feature = "test-hooks"))]
30impl WelcomePointersHook {
31    pub fn new(enabled: bool) -> Self {
32        Self {
33            enabled,
34            reads: Default::default(),
35        }
36    }
37
38    /// Record a package build that reads this hook, and return its setting.
39    pub fn read(&self) -> bool {
40        self.reads.fetch_add(1, Ordering::SeqCst);
41        self.enabled
42    }
43
44    /// How many package builds read this hook.
45    pub fn reads(&self) -> usize {
46        self.reads.load(Ordering::SeqCst)
47    }
48}
49
50#[cfg(any(test, feature = "test-hooks"))]
51tokio::task_local! {
52    pub static ENABLE_WELCOME_POINTERS: WelcomePointersHook;
53}
54
55#[derive(Builder, Debug)]
56#[builder(build_fn(error = "IdentityError", name = "inner_build", private))]
57pub struct XmtpKeyPackage {
58    #[builder(setter(into))]
59    inbox_id: String,
60    #[builder(setter(into))]
61    credential: OpenMlsCredential,
62    #[builder(setter(into))]
63    installation_keys: XmtpInstallationCredential,
64}
65
66impl XmtpKeyPackage {
67    /// Starts building a key package.
68    pub fn builder() -> XmtpKeyPackageBuilder {
69        XmtpKeyPackageBuilder::default()
70    }
71}
72
73impl XmtpKeyPackageBuilder {
74    /// Builds the key package and stores its private keys through `provider`.
75    pub fn build(
76        &mut self,
77        provider: &impl MlsProviderExt,
78        include_post_quantum: bool,
79    ) -> Result<NewKeyPackageResult, IdentityError> {
80        let this = self.inner_build()?;
81        #[allow(unused_mut)]
82        let mut options = xmtp_id::key_package::KeyPackageOptions {
83            include_post_quantum,
84            ..Default::default()
85        };
86        #[cfg(any(test, feature = "test-hooks"))]
87        {
88            options.welcome_pointers = ENABLE_WELCOME_POINTERS
89                .try_with(WelcomePointersHook::read)
90                .unwrap_or(true);
91            options.lifetime =
92                Some(crate::identity_lifecycle::test_mocks::maybe_mock_package_lifetime());
93        }
94        let generated = xmtp_id::key_package::build_key_package(
95            &this.inbox_id,
96            this.credential,
97            &this.installation_keys,
98            provider,
99            options,
100        )
101        .map_err(|error| match error {
102            xmtp_id::key_package::KeyPackageConstructionError::Generation(e) => {
103                IdentityError::KeyPackageGenerationError(e)
104            }
105            xmtp_id::key_package::KeyPackageConstructionError::InvalidExtension(e) => {
106                IdentityError::InvalidExtension(e)
107            }
108            xmtp_id::key_package::KeyPackageConstructionError::Encode(e) => {
109                IdentityError::CredentialSerialization(e)
110            }
111            xmtp_id::key_package::KeyPackageConstructionError::PostQuantum(e) => {
112                IdentityError::GeneratePostQuantumKey(e)
113            }
114        })?;
115        store_key_package_references(
116            provider,
117            generated.bundle.key_package(),
118            &generated.post_quantum_keypair,
119        )?;
120        Ok(NewKeyPackageResult {
121            key_package: generated.bundle.key_package().clone(),
122            pq_pub_key: generated.post_quantum_keypair.map(|kp| kp.public),
123        })
124    }
125}