1pub use xmtp_cryptography::GeneratePostQuantumKeyError;
2
3use crate::worker::NeedsDbReconnect;
4use openmls::prelude::HpkeKeyPair;
5use openmls::prelude::hash_ref::HashReference;
6use openmls::{
7 credentials::errors::BasicCredentialError,
8 extensions::Extension,
9 key_packages::KeyPackage,
10 prelude::{
11 Credential as OpenMlsCredential,
12 tls_codec::{Deserialize as _, Serialize, VLBytes},
13 },
14};
15use openmls_traits::types::CryptoError;
16use std::sync::atomic::{AtomicBool, Ordering};
17use thiserror::Error;
18use xmtp_common::ErrorCode;
19use xmtp_common::time::now_ns;
20use xmtp_common::{RetryableError, retryable};
21use xmtp_configuration::KEY_PACKAGE_ROTATION_INTERVAL_NS;
22use xmtp_cryptography::signature::IdentifierValidationError;
23use xmtp_cryptography::{CredentialSign, XmtpInstallationCredential};
24use xmtp_db::db_connection::DbConnection;
25use xmtp_db::identity::StoredIdentity;
26use xmtp_db::prelude::*;
27use xmtp_db::sql_key_store::{
28 KEY_PACKAGE_REFERENCES, KEY_PACKAGE_WRAPPER_PRIVATE_KEY, SqlKeyStoreError, codec as kv_codec,
29};
30use xmtp_db::{ConnectionExt, MlsProviderExt};
31use xmtp_db::{StorageError, StoredSecret};
32use xmtp_id::associations::{AssociationError, InstallationKeyContext, PublicContext};
33use xmtp_id::key_package::KeyPackageVerificationError;
34use xmtp_id::{
35 InboxId, InboxIdRef,
36 associations::builder::{SignatureRequest, SignatureRequestError},
37};
38use xmtp_proto::types::InstallationId;
39
40#[derive(Debug, Error, ErrorCode)]
41pub enum IdentityError {
42 #[error(transparent)]
46 CredentialSerialization(#[from] prost::EncodeError),
47 #[error(transparent)]
51 Decode(#[from] prost::DecodeError),
52 #[error("installation not found: {0}")]
56 InstallationIdNotFound(String),
57 #[error(transparent)]
58 #[error_code(inherit)]
59 SignatureRequestBuilder(#[from] SignatureRequestError),
60 #[error(transparent)]
61 #[error_code(inherit)]
62 Signature(#[from] xmtp_id::associations::SignatureError),
63 #[error(transparent)]
67 BasicCredential(#[from] BasicCredentialError),
68 #[error("Uninitialized identity")]
72 UninitializedIdentity,
73 #[error("Installation key {0}")]
77 InstallationKey(String),
78 #[error(transparent)]
82 Crypto(#[from] CryptoError),
83 #[error(transparent)]
87 OpenMls(#[from] openmls::prelude::Error),
88 #[error(transparent)]
89 #[error_code(inherit)]
90 StorageError(#[from] xmtp_db::StorageError),
91 #[error(transparent)]
92 #[error_code(inherit)]
93 OpenMlsStorageError(#[from] SqlKeyStoreError),
94 #[error(transparent)]
98 KeyPackageGenerationError(#[from] openmls::key_packages::errors::KeyPackageNewError),
99 #[error(transparent)]
100 #[error_code(inherit)]
101 KeyPackageVerificationError(#[from] KeyPackageVerificationError),
102 #[error("The InboxID {id}, associated does not match the stored InboxId {stored}.")]
106 InboxIdMismatch { id: InboxId, stored: InboxId },
107 #[error("The identifier does not belong to the stored InboxId {inbox_id}.")]
111 IdentifierNotInInbox { inbox_id: InboxId },
112 #[error("The address {0} has no associated InboxID")]
116 NoAssociatedInboxId(String),
117 #[error("Required identity was not found in cache.")]
121 RequiredIdentityNotFound,
122 #[error("error creating new identity: {0}")]
126 NewIdentity(String),
127 #[error(transparent)]
128 #[error_code(inherit)]
129 Association(#[from] AssociationError),
130 #[error(transparent)]
134 Signer(#[from] xmtp_cryptography::SignerError),
135 #[error(transparent)]
136 #[error_code(inherit)]
137 ApiClient(#[from] xmtp_api::ApiError),
138 #[error(transparent)]
139 #[error_code(inherit)]
140 AddressValidation(#[from] IdentifierValidationError),
141 #[error(transparent)]
142 #[error_code(inherit)]
143 Db(#[from] xmtp_db::ConnectionError),
144 #[error(
148 "Cannot register a new installation because the InboxID {inbox_id} has already registered {count}/{max} installations. Please revoke existing installations first."
149 )]
150 TooManyInstallations {
151 inbox_id: String,
152 count: usize,
153 max: usize,
154 },
155 #[error(transparent)]
156 #[error_code(inherit)]
157 GeneratePostQuantumKey(#[from] GeneratePostQuantumKeyError),
158 #[error(transparent)]
162 InvalidExtension(#[from] openmls::prelude::InvalidExtensionError),
163 #[error("Missing post quantum public key")]
167 MissingPostQuantumPublicKey,
168 #[error(transparent)]
172 UninitializedField(#[from] derive_builder::UninitializedFieldError),
173}
174
175impl From<xmtp_db::diesel::result::Error> for IdentityError {
176 fn from(error: xmtp_db::diesel::result::Error) -> Self {
177 Self::StorageError(error.into())
178 }
179}
180
181impl NeedsDbReconnect for IdentityError {
182 #[deny(clippy::wildcard_enum_match_arm)]
183 fn needs_db_reconnect(&self) -> bool {
184 match self {
185 Self::StorageError(s) => s.db_needs_connection(),
186 Self::Db(c) => c.db_needs_connection(),
187 Self::OpenMlsStorageError(SqlKeyStoreError::Connection(c)) => c.db_needs_connection(),
189 Self::OpenMlsStorageError(_)
190 | Self::CredentialSerialization(..)
191 | Self::Decode(..)
192 | Self::InstallationIdNotFound(..)
193 | Self::SignatureRequestBuilder(..)
194 | Self::Signature(..)
195 | Self::BasicCredential(..)
196 | Self::UninitializedIdentity
197 | Self::InstallationKey(..)
198 | Self::Crypto(..)
199 | Self::OpenMls(..)
200 | Self::KeyPackageGenerationError(..)
201 | Self::KeyPackageVerificationError(..)
202 | Self::InboxIdMismatch { .. }
203 | Self::IdentifierNotInInbox { .. }
204 | Self::NoAssociatedInboxId(..)
205 | Self::RequiredIdentityNotFound
206 | Self::NewIdentity(..)
207 | Self::Association(..)
208 | Self::Signer(..)
209 | Self::ApiClient(..)
210 | Self::AddressValidation(..)
211 | Self::TooManyInstallations { .. }
212 | Self::GeneratePostQuantumKey(..)
213 | Self::InvalidExtension(..)
214 | Self::MissingPostQuantumPublicKey
215 | Self::UninitializedField(..) => false,
216 }
217 }
218}
219
220impl RetryableError for IdentityError {
221 fn is_retryable(&self) -> bool {
222 match self {
223 Self::ApiClient(err) => retryable!(err),
224 Self::StorageError(err) => retryable!(err),
225 Self::OpenMlsStorageError(err) => retryable!(err),
226 _ => false,
227 }
228 }
229
230 fn rate_limit_delay(&self) -> Option<xmtp_common::time::Duration> {
231 match self {
232 Self::ApiClient(err) => err.rate_limit_delay(),
233 _ => None,
234 }
235 }
236}
237
238#[derive(Debug)]
239pub struct Identity {
240 pub inbox_id: InboxId,
242 pub installation_keys: XmtpInstallationCredential,
244 pub credential: OpenMlsCredential,
246 pub signature_request: Option<SignatureRequest>,
249 pub is_ready: AtomicBool,
251}
252
253impl Clone for Identity {
254 fn clone(&self) -> Self {
255 Self {
256 inbox_id: self.inbox_id.clone(),
257 installation_keys: self.installation_keys.clone(),
258 credential: self.credential.clone(),
259 signature_request: self.signature_request(),
260 is_ready: AtomicBool::new(self.is_ready.load(Ordering::SeqCst)),
261 }
262 }
263}
264
265impl TryFrom<&Identity> for StoredIdentity {
266 type Error = StorageError;
267
268 fn try_from(identity: &Identity) -> Result<Self, Self::Error> {
269 StoredIdentity::builder()
270 .inbox_id(identity.inbox_id.clone())
271 .installation_keys(StoredSecret::from(xmtp_db::db_serialize(
272 &identity.installation_keys,
273 )?))
274 .credential_bytes(xmtp_db::db_serialize(&identity.credential())?)
275 .next_key_package_rotation_ns(now_ns() + KEY_PACKAGE_ROTATION_INTERVAL_NS)
276 .build()
277 }
278}
279
280impl TryFrom<StoredIdentity> for Identity {
281 type Error = StorageError;
282
283 fn try_from(identity: StoredIdentity) -> Result<Self, Self::Error> {
284 Ok(Identity {
285 inbox_id: identity.inbox_id.clone(),
286 installation_keys: xmtp_db::db_deserialize(identity.installation_keys.as_slice())?,
287 credential: xmtp_db::db_deserialize(&identity.credential_bytes)?,
288 signature_request: None,
289 is_ready: AtomicBool::new(true),
290 })
291 }
292}
293
294impl Identity {
295 pub fn inbox_id(&self) -> InboxIdRef<'_> {
296 &self.inbox_id
297 }
298
299 pub fn installation_id(&self) -> InstallationId {
300 (*self.installation_keys.public_bytes()).into()
301 }
302
303 pub fn sequence_id<C>(&self, conn: &DbConnection<C>) -> Result<i64, xmtp_db::ConnectionError>
304 where
305 C: ConnectionExt,
306 {
307 conn.get_latest_sequence_id_for_inbox(self.inbox_id.as_str())
308 }
309
310 pub fn is_ready(&self) -> bool {
311 self.is_ready.load(Ordering::SeqCst)
312 }
313
314 pub fn set_ready(&self) {
316 self.is_ready.store(true, Ordering::SeqCst)
317 }
318
319 pub fn signature_request(&self) -> Option<SignatureRequest> {
320 self.signature_request.clone()
321 }
322
323 pub fn credential(&self) -> OpenMlsCredential {
324 self.credential.clone()
325 }
326
327 pub fn sign_identity_update<Text: AsRef<str>>(
331 &self,
332 text: Text,
333 ) -> Result<Vec<u8>, IdentityError> {
334 self.installation_keys
335 .credential_sign::<InstallationKeyContext>(text)
336 .map_err(Into::into)
337 }
338
339 pub fn sign_with_public_context(
340 &self,
341 text: impl AsRef<str>,
342 ) -> Result<Vec<u8>, IdentityError> {
343 self.installation_keys
344 .credential_sign::<PublicContext>(text)
345 .map_err(Into::into)
346 }
347}
348
349pub fn serialize_key_package_hash_ref(
355 kp: &KeyPackage,
356 provider: &impl MlsProviderExt,
357) -> Result<Vec<u8>, IdentityError> {
358 let key_package_hash_ref = kp
359 .hash_ref(provider.crypto())
360 .map_err(|_| IdentityError::UninitializedIdentity)?;
361 hash_ref_key(&key_package_hash_ref)
362}
363
364fn hash_ref_key(hash_ref: &HashReference) -> Result<Vec<u8>, IdentityError> {
366 kv_codec::encode_key(KEY_PACKAGE_WRAPPER_PRIVATE_KEY, hash_ref)
367 .map_err(|_| IdentityError::UninitializedIdentity)
368}
369
370pub fn pq_key_package_references_key(raw_pub_key: &[u8]) -> Result<Vec<u8>, IdentityError> {
373 Ok(raw_pub_key.tls_serialize_detached()?)
374}
375
376pub fn deserialize_key_package_hash_ref(hash_ref: &[u8]) -> Result<HashReference, IdentityError> {
384 let raw = kv_codec::key::decode_bytes(hash_ref).ok_or(IdentityError::UninitializedIdentity)?;
385 let tls = VLBytes::new(raw.to_vec()).tls_serialize_detached()?;
388 HashReference::tls_deserialize_exact(tls).map_err(|_| IdentityError::UninitializedIdentity)
389}
390
391pub fn create_credential(inbox_id: impl AsRef<str>) -> Result<OpenMlsCredential, IdentityError> {
393 Ok(xmtp_id::key_package::create_credential(inbox_id))
394}
395
396pub fn parse_credential(credential_bytes: &[u8]) -> Result<InboxId, IdentityError> {
397 Ok(xmtp_id::key_package::parse_credential(credential_bytes)?)
398}
399
400pub fn build_post_quantum_public_key_extension(
401 public_key: &[u8],
402) -> Result<Extension, IdentityError> {
403 Ok(xmtp_id::key_package::build_post_quantum_public_key_extension(public_key)?)
404}
405
406pub fn store_key_package_references(
412 provider: &impl MlsProviderExt,
413 kp: &KeyPackage,
414 post_quantum_keypair: &Option<HpkeKeyPair>,
416) -> Result<(), IdentityError> {
417 let public_init_key = kp.hpke_init_key().tls_serialize_detached()?;
420
421 let hash_ref = kp
424 .hash_ref(provider.crypto())
425 .map_err(|_| IdentityError::UninitializedIdentity)?;
426 let hash_ref_value = kv_codec::encode_value(KEY_PACKAGE_REFERENCES, &hash_ref)?;
427 let storage = provider.key_store();
428 storage.write(KEY_PACKAGE_REFERENCES, &public_init_key, &hash_ref_value)?;
430
431 if let Some(post_quantum_keypair) = post_quantum_keypair {
432 let post_quantum_public_key = pq_key_package_references_key(&post_quantum_keypair.public)?;
433
434 storage.write(
436 KEY_PACKAGE_REFERENCES,
437 &post_quantum_public_key,
438 &hash_ref_value,
439 )?;
440
441 storage
443 .write_wrapper_private_key(&hash_ref_key(&hash_ref)?, &post_quantum_keypair.private)?;
444 }
445
446 Ok(())
447}
448
449#[cfg(test)]
450mod tests {
451 use crate::identity::{pq_key_package_references_key, serialize_key_package_hash_ref};
452 use openmls_traits::OpenMlsProvider;
453 use tls_codec::Serialize;
454 use xmtp_db::XmtpMlsStorageProvider;
455 use xmtp_db::{MlsProviderExt, sql_key_store::KEY_PACKAGE_REFERENCES};
456
457 #[xmtp_common::test(unwrap_try = true)]
463 async fn key_package_hash_ref_uses_the_key_layout_and_cbor() {
464 use super::{deserialize_key_package_hash_ref, store_key_package_references};
465 use diesel::prelude::*;
466 use xmtp_db::{
467 ConnectionExt, XmtpTestDb, schema::openmls_key_value::dsl, sql_key_store::SqlKeyStore,
468 xmtp_openmls_provider::XmtpOpenMlsProvider,
469 };
470
471 let store = xmtp_db::TestDb::create_persistent_store(None).await;
472 let provider = XmtpOpenMlsProvider::new(SqlKeyStore::new(store.conn()));
473 let bundle = xmtp_db::test_utils::generate_key_package_bundle();
474 let post_quantum = xmtp_cryptography::post_quantum::generate_post_quantum_key()?;
475 store_key_package_references(&provider, bundle.key_package(), &Some(post_quantum.clone()))?;
476
477 let hash_ref = bundle.key_package().hash_ref(provider.crypto())?;
478 let raw = hash_ref.as_slice();
479 let key_layout = [&(raw.len() as u64).to_le_bytes()[..], raw].concat();
481 let mut cbor = [
483 &[0xa1, 0x65][..],
484 b"value",
485 &[0xa1, 0x63],
486 b"vec",
487 &[0x98, u8::try_from(raw.len())?],
488 ]
489 .concat();
490 for &byte in raw {
491 if byte >= 0x18 {
492 cbor.push(0x18);
493 }
494 cbor.push(byte);
495 }
496 assert_ne!(key_layout, cbor);
497
498 let column = serialize_key_package_hash_ref(bundle.key_package(), &provider)?;
499 assert_eq!(column, key_layout, "key_package_history column");
500 assert_eq!(deserialize_key_package_hash_ref(&column)?, hash_ref);
501 assert!(deserialize_key_package_hash_ref(&column[..column.len() - 1]).is_err());
502 assert!(deserialize_key_package_hash_ref(&[column.as_slice(), &[0]].concat()).is_err());
504 assert!(deserialize_key_package_hash_ref(&cbor).is_err());
505
506 let stored = |key: &[u8]| -> Option<Vec<u8>> {
507 let storage_key = [KEY_PACKAGE_REFERENCES, key, &1u16.to_be_bytes()].concat();
508 store
509 .conn()
510 .raw_query(|conn| {
511 dsl::openmls_key_value
512 .filter(dsl::key_bytes.eq(storage_key))
513 .select(dsl::value_bytes)
514 .first::<Vec<u8>>(conn)
515 .optional()
516 })
517 .unwrap()
518 };
519 let init_key = bundle
520 .key_package()
521 .hpke_init_key()
522 .tls_serialize_detached()?;
523 let post_quantum_key = pq_key_package_references_key(&post_quantum.public)?;
524 assert_eq!(stored(&init_key), Some(cbor.clone()), "init key reference");
525 assert_eq!(
526 stored(&post_quantum_key),
527 Some(cbor.clone()),
528 "post-quantum key reference"
529 );
530 let wrapper_key = provider.key_store().wrapper_private_key(&key_layout)?;
531 assert_eq!(
532 wrapper_key.as_deref().map(Vec::as_slice),
533 Some(&*post_quantum.private)
534 );
535 assert!(provider.key_store().wrapper_private_key(&cbor)?.is_none());
536 }
537
538 #[test]
539 fn test_generate_post_quantum_key_error_codes() {
540 use super::GeneratePostQuantumKeyError;
541 use openmls_traits::types::CryptoError;
542 use xmtp_common::ErrorCode;
543
544 let crypto_err = GeneratePostQuantumKeyError::Crypto(CryptoError::CryptoLibraryError);
546 assert_eq!(
547 crypto_err.error_code(),
548 "GeneratePostQuantumKeyError::Crypto"
549 );
550 }
551
552 #[test]
553 fn test_identity_error_codes() {
554 use super::IdentityError;
555 use xmtp_common::ErrorCode;
556
557 let err = IdentityError::UninitializedIdentity;
559 assert_eq!(err.error_code(), "IdentityError::UninitializedIdentity");
560
561 let err = IdentityError::RequiredIdentityNotFound;
562 assert_eq!(err.error_code(), "IdentityError::RequiredIdentityNotFound");
563
564 let err = IdentityError::MissingPostQuantumPublicKey;
565 assert_eq!(
566 err.error_code(),
567 "IdentityError::MissingPostQuantumPublicKey"
568 );
569
570 let err = IdentityError::InstallationIdNotFound("test".to_string());
572 assert_eq!(err.error_code(), "IdentityError::InstallationIdNotFound");
573
574 let err = IdentityError::InstallationKey("test".to_string());
575 assert_eq!(err.error_code(), "IdentityError::InstallationKey");
576
577 let err = IdentityError::NewIdentity("test".to_string());
578 assert_eq!(err.error_code(), "IdentityError::NewIdentity");
579
580 let err = IdentityError::TooManyInstallations {
581 inbox_id: "test".to_string(),
582 count: 10,
583 max: 5,
584 };
585 assert_eq!(err.error_code(), "IdentityError::TooManyInstallations");
586
587 let err = IdentityError::InboxIdMismatch {
588 id: "id1".to_string(),
589 stored: "id2".to_string(),
590 };
591 assert_eq!(err.error_code(), "IdentityError::InboxIdMismatch");
592
593 let err = IdentityError::NoAssociatedInboxId("addr".to_string());
594 assert_eq!(err.error_code(), "IdentityError::NoAssociatedInboxId");
595 }
596
597 #[test]
598 fn test_identity_error_inherited_codes() {
599 use super::IdentityError;
600 use xmtp_common::ErrorCode;
601 use xmtp_db::{NotFound, StorageError};
602
603 let storage_err = StorageError::NotFound(NotFound::MessageById(vec![1, 2, 3]));
605 let err = IdentityError::StorageError(storage_err);
606 assert_eq!(err.error_code(), "StorageError::NotFound");
607 }
608}