xmtp_mls_identity/identity_lifecycle/
key_package.rs1use 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
11pub struct NewKeyPackageResult {
13 pub key_package: KeyPackage,
15 pub pq_pub_key: Option<Vec<u8>>,
17}
18
19#[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 pub fn read(&self) -> bool {
40 self.reads.fetch_add(1, Ordering::SeqCst);
41 self.enabled
42 }
43
44 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 pub fn builder() -> XmtpKeyPackageBuilder {
69 XmtpKeyPackageBuilder::default()
70 }
71}
72
73impl XmtpKeyPackageBuilder {
74 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}