1use crate::identity_updates::{get_association_state_with_verifier, load_identity_updates};
5use crate::key_package_maintenance::discard_key_package;
6use openmls::prelude::tls_codec::Serialize;
7use std::sync::atomic::AtomicBool;
8use xmtp_api::ApiClientWrapper;
9use xmtp_configuration::KEY_PACKAGE_ROTATION_INTERVAL_NS;
10use xmtp_cryptography::{CredentialSign, XmtpInstallationCredential};
11use xmtp_db::TransactionOutcome::Continue;
12use xmtp_db::{MlsProviderExt, StorageError, TransactionOutcome};
13use xmtp_db::{XmtpMlsStorageProvider, XmtpOpenMlsProviderRef, prelude::*};
14use xmtp_id::associations::unverified::UnverifiedSignature;
15use xmtp_id::associations::{Identifier, InstallationKeyContext};
16use xmtp_id::scw_verifier::SmartContractSignatureVerifier;
17use xmtp_id::{
18 InboxId,
19 associations::{MemberIdentifier, builder::SignatureRequestBuilder},
20};
21use xmtp_mls_core::api_status::definite_refusal;
22use xmtp_mls_core::identity::{
23 Identity, IdentityError, create_credential, serialize_key_package_hash_ref,
24};
25use xmtp_mls_core::mls_store::verify_served_key_package;
26use xmtp_proto::api_client::XmtpApi;
27
28mod key_package;
29mod strategy;
30#[cfg(all(test, not(target_arch = "wasm32")))]
31mod strategy_tests;
32#[cfg(any(test, feature = "test-hooks"))]
33pub mod test_mocks;
34#[cfg(any(test, feature = "test-utils"))]
35#[doc(hidden)]
36pub mod test_support;
37
38pub(crate) use key_package::NewKeyPackageResult;
39#[cfg(any(test, feature = "test-hooks"))]
40pub use key_package::{ENABLE_WELCOME_POINTERS, WelcomePointersHook};
41pub use key_package::{XmtpKeyPackage, XmtpKeyPackageBuilder};
42pub use strategy::IdentityStrategy;
43
44#[cfg(any(test, feature = "test-utils"))]
46pub fn identity_setup(
47 owner: impl xmtp_id::InboxOwner,
48) -> crate::identity_lifecycle::IdentityStrategy {
49 let nonce = 1;
50 let ident = owner.get_identifier().expect("the owner's identifier");
51 let inbox_id = ident
52 .inbox_id(nonce)
53 .expect("an inbox id for the identifier");
54 crate::identity_lifecycle::IdentityStrategy::new(inbox_id, ident, nonce)
55}
56
57#[tracing::instrument(name = "new", level = "trace", skip_all)]
63pub async fn new_identity<ApiClient: XmtpApi, S: XmtpMlsStorageProvider>(
64 inbox_id: InboxId,
65 identifier: Identifier,
66 nonce: u64,
67 api_client: &ApiClientWrapper<ApiClient>,
68 mls_storage: &S,
69 scw_signature_verifier: impl SmartContractSignatureVerifier,
70) -> Result<Identity, IdentityError> {
71 let inbox_ids = api_client
73 .get_inbox_ids(vec![identifier.clone().into()])
74 .await?;
75 let associated_inbox_id = inbox_ids.first().and_then(Option::as_ref);
76 let installation_keys = XmtpInstallationCredential::new();
77
78 if let Some(associated_inbox_id) = associated_inbox_id {
79 if *associated_inbox_id != inbox_id {
82 return Err(IdentityError::NewIdentity("Inbox ID mismatch".to_string()));
83 }
84
85 load_identity_updates(
87 api_client,
88 &mls_storage.db(),
89 &[associated_inbox_id.as_str()],
90 )
91 .await
92 .map_err(|e| IdentityError::NewIdentity(format!("Failed to load identity updates: {e}")))?;
93
94 let state = get_association_state_with_verifier(
95 &mls_storage.db(),
96 &inbox_id,
97 None,
98 &scw_signature_verifier,
99 )
100 .await
101 .map_err(|err| {
102 IdentityError::NewIdentity(format!("Error resolving identity state: {}", err))
103 })?;
104
105 let max_installations = api_client.configuration().mls.max_installations_per_inbox;
108 let current_installation_count = state.installation_ids().len();
109 if current_installation_count >= max_installations {
111 return Err(IdentityError::TooManyInstallations {
112 inbox_id: associated_inbox_id.clone(),
113 count: current_installation_count,
114 max: max_installations,
115 });
116 }
117
118 let builder = SignatureRequestBuilder::new(associated_inbox_id.clone());
119 let mut signature_request = builder
120 .add_association(
121 MemberIdentifier::installation(installation_keys.public_slice().to_vec()),
122 identifier.clone().into(),
123 )
124 .build();
125
126 let signature = installation_keys
127 .credential_sign::<InstallationKeyContext>(signature_request.signature_text())?;
128 signature_request
129 .add_signature(
130 UnverifiedSignature::new_installation_key(
131 signature,
132 installation_keys.verifying_key(),
133 ),
134 &scw_signature_verifier,
135 )
136 .await?;
137
138 let identity = Identity {
139 inbox_id: associated_inbox_id.clone(),
140 installation_keys,
141 credential: create_credential(associated_inbox_id.clone())?,
142 signature_request: Some(signature_request),
143 is_ready: AtomicBool::new(false),
144 };
145
146 Ok(identity)
147 } else {
148 let generated_inbox_id = identifier.inbox_id(nonce)?;
149 if inbox_id != generated_inbox_id {
150 return Err(IdentityError::NewIdentity(
151 "Inbox ID doesn't match nonce & address".to_string(),
152 ));
153 }
154 create_inbox(
155 identifier,
156 nonce,
157 installation_keys,
158 &scw_signature_verifier,
159 )
160 .await
161 }
162}
163
164async fn create_inbox(
168 identifier: Identifier,
169 nonce: u64,
170 installation_keys: XmtpInstallationCredential,
171 scw_signature_verifier: impl SmartContractSignatureVerifier,
172) -> Result<Identity, IdentityError> {
173 let inbox_id = identifier.inbox_id(nonce)?;
174 let mut builder = SignatureRequestBuilder::new(inbox_id.clone());
175 builder = builder.create_inbox(identifier.clone(), nonce);
176
177 let mut signature_request = builder
178 .add_association(
179 MemberIdentifier::installation(installation_keys.public_slice().to_vec()),
180 identifier.into(),
181 )
182 .build();
183
184 let sig = installation_keys
186 .credential_sign::<InstallationKeyContext>(signature_request.signature_text())?;
187 signature_request
188 .add_signature(
189 UnverifiedSignature::new_installation_key(sig, installation_keys.verifying_key()),
190 &scw_signature_verifier,
191 )
192 .await?;
193
194 Ok(Identity {
195 credential: create_credential(inbox_id.clone())?,
196 inbox_id,
197 installation_keys,
198 signature_request: Some(signature_request),
199 is_ready: AtomicBool::new(false),
200 })
201}
202
203#[cfg(any(test, feature = "test-utils"))]
207pub async fn new_inbox_with_installation_keys(
208 identifier: Identifier,
209 nonce: u64,
210 installation_keys: XmtpInstallationCredential,
211) -> Result<Identity, IdentityError> {
212 create_inbox(
213 identifier,
214 nonce,
215 installation_keys,
216 xmtp_id::associations::test_utils::MockSmartContractSignatureVerifier::new(true),
217 )
218 .await
219}
220
221#[tracing::instrument(level = "trace", skip_all)]
223pub fn new_key_package(
224 identity: &Identity,
225 provider: &impl MlsProviderExt,
226 include_post_quantum: bool,
227) -> Result<NewKeyPackageResult, IdentityError> {
228 XmtpKeyPackage::builder()
229 .inbox_id(identity.inbox_id())
230 .credential(identity.credential())
231 .installation_keys(identity.installation_keys.clone())
232 .build(provider, include_post_quantum)
233}
234
235#[tracing::instrument(level = "trace", skip_all)]
239pub async fn rotate_and_upload_key_package<ApiClient: XmtpApi, S: XmtpMlsStorageProvider>(
240 identity: &Identity,
241 api_client: &ApiClientWrapper<ApiClient>,
242 mls_storage: &S,
243 include_post_quantum: bool,
244) -> Result<(), IdentityError> {
245 tracing::info!("Start rotating keys and uploading the new key package");
246
247 if let Err(error) = settle_unconfirmed_key_packages(identity, api_client, mls_storage).await {
249 tracing::warn!(%error, "key package settlement deferred");
250 }
251
252 let (kp_bytes, history_id) =
254 generate_and_store_key_package(identity, mls_storage, include_post_quantum)?;
255
256 match api_client.upload_key_package(kp_bytes).await {
258 Ok(meta) => {
259 let published_cursor = meta.cursor.filter(|cursor| cursor.sequence_id > 0).ok_or(
260 xmtp_api::ApiError::InvalidResponse("key package publish cursor"),
261 )?;
262 xmtp_mls_core::state_tx::state_write(mls_storage, |tx| {
264 let storage = tx.storage();
265 storage.db().record_key_package_publication(
266 history_id,
267 xmtp_proto::types::Cursor(published_cursor.sequence_id),
268 )?;
269 storage
270 .db()
271 .reset_key_package_rotation_queue(KEY_PACKAGE_ROTATION_INTERVAL_NS)?;
272 Ok::<_, StorageError>(Continue(()))
273 })
274 .map(TransactionOutcome::into_continued)?;
275 Ok(())
276 }
277 Err(err) if definite_refusal(&err) => {
279 discard_key_package(mls_storage, history_id)?;
280 Err(IdentityError::ApiClient(err))
281 }
282 Err(err) => Err(IdentityError::ApiClient(err)),
283 }
284}
285
286async fn settle_unconfirmed_key_packages<ApiClient: XmtpApi, S: XmtpMlsStorageProvider>(
288 identity: &Identity,
289 api_client: &ApiClientWrapper<ApiClient>,
290 mls_storage: &S,
291) -> Result<(), IdentityError> {
292 let mut history = mls_storage
293 .db()
294 .find_key_package_history_entries_before_id(i32::MAX)?
295 .into_iter();
296 if !history.any(|kp| kp.published_sequence_id.is_none() && kp.delete_at_ns.is_none()) {
297 return Ok(());
298 }
299 let id = identity.installation_id();
300 let Some((package, served)) = api_client.newest_key_package(&id).await? else {
301 return Ok(());
302 };
303 let crypto = xmtp_db::XmtpOpenMlsProvider::<()>::new_crypto();
304 let Some(Ok(verified)) = verify_served_key_package(&crypto, &id, &package) else {
305 return Ok(());
306 };
307 let provider = XmtpOpenMlsProviderRef::new(mls_storage);
308 let hash_ref = serialize_key_package_hash_ref(&verified.inner, &provider)?;
309 Ok(mls_storage.db().settle_key_packages(hash_ref, served)?)
310}
311
312#[tracing::instrument(level = "trace", skip_all)]
315pub fn generate_and_store_key_package<S: XmtpMlsStorageProvider>(
316 identity: &Identity,
317 mls_storage: &S,
318 include_post_quantum: bool,
319) -> Result<(Vec<u8>, i32), IdentityError> {
320 xmtp_mls_core::state_tx::state_write_durable(mls_storage, |tx| {
321 let storage = tx.storage();
322 let provider = XmtpOpenMlsProviderRef::new(&storage);
323 let NewKeyPackageResult {
324 key_package: kp,
325 pq_pub_key,
326 } = new_key_package(identity, &provider, include_post_quantum)?;
327 let hash_ref = serialize_key_package_hash_ref(&kp, &provider)?;
328 let history_id = storage
329 .db()
330 .store_key_package_history_entry(hash_ref, pq_pub_key)?
331 .id;
332 let kp_bytes = kp.tls_serialize_detached()?;
333 Ok::<_, IdentityError>(Continue((kp_bytes, history_id)))
334 })
335 .map(TransactionOutcome::into_continued)
336}