Skip to main content

xmtp_mls_identity/
identity_lifecycle.rs

1//! Creating an identity and keeping its key packages: the identity-owned
2//! constructors and key package rotation over the [`Identity`] struct.
3
4use 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/// The identity strategy of a test client: nonce 1 of the owner's identifier.
45#[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/// Create a new [Identity] instance.
58///
59/// If the address is already associated with an inbox_id, the existing inbox_id will be used.
60/// If it is not, a new inbox_id is generated from the identifier and nonce.
61/// Either way, the caller must sign with their wallet.
62#[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    // check if address is already associated with an inbox_id
72    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 an inbox is associated with address, we'd use it to create Identity and ignore the nonce.
80        // We would need a signature from user's wallet.
81        if *associated_inbox_id != inbox_id {
82            return Err(IdentityError::NewIdentity("Inbox ID mismatch".to_string()));
83        }
84
85        // get sequence_id from identity updates and loaded into the DB
86        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        // The deployment sets the ceiling; the wrapper carries the
106        // snapshot resolved before any identity work ran.
107        let max_installations = api_client.configuration().mls.max_installations_per_inbox;
108        let current_installation_count = state.installation_ids().len();
109        // implements: CONF-044
110        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
164/// An unregistered identity that creates the inbox of `identifier` and
165/// `nonce` with `installation_keys`. The request is pre-signed with the
166/// installation key; the caller adds the wallet signature.
167async 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    // We can pre-sign the request with an installation key signature, since we have access to the key
185    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/// An unregistered identity for a new inbox that uses the given
204/// installation key. Use it with [`IdentityStrategy::ExternalIdentity`],
205/// then register with the owner of `identifier`.
206#[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/// Generate a new key package and store the associated keys in the database.
222#[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/// Store fresh key material before upload, then record its publication receipt.
236/// Receipt order controls retirement. A read settles unknown publications first.
237/// Receipt bookkeeping and the next rotation deadline commit together.
238#[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    // implements: JOIN-074
248    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    // Generate and store key package locally
253    let (kp_bytes, history_id) =
254        generate_and_store_key_package(identity, mls_storage, include_post_quantum)?;
255
256    // Upload to network
257    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            // Backend publication order can differ from local generation order.
263            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        // INVALID_ARGUMENT stored nothing, so no sender can address the package.
278        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
286/// Settle receipt-less keys against the package the backend serves, before a new one exists.
287async 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/// Store key material and its history row in one state transaction.
313/// Return bytes for upload and the history ID for its later receipt.
314#[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}