1use base64::prelude::*;
2use phonenumber::PhoneNumber;
3use rand::{CryptoRng, Rng};
4use reqwest::Method;
5use std::collections::HashMap;
6use std::convert::{TryFrom, TryInto};
7
8use aes::cipher::{KeyIvInit, StreamCipher as _};
9use hmac::digest::Output;
10use hmac::{Hmac, Mac};
11use libsignal_protocol::{
12 kem, Aci, GenericSignedPreKey, IdentityKey, IdentityKeyPair,
13 IdentityKeyStore, KeyPair, KyberPreKeyRecord, PrivateKey, ProtocolStore,
14 PublicKey, SenderKeyStore, ServiceIdKind, SignedPreKeyRecord, Timestamp,
15};
16use prost::Message;
17use serde::{Deserialize, Serialize};
18use sha2::{Digest, Sha256};
19use tracing_futures::Instrument;
20use zkgroup::profiles::ProfileKey;
21
22use crate::content::ContentBody;
23use crate::master_key::MasterKey;
24use crate::pre_keys::{
25 KyberPreKeyEntity, PreKeyEntity, PreKeysStore, SignedPreKeyEntity,
26 PRE_KEY_BATCH_SIZE, PRE_KEY_MINIMUM,
27};
28use crate::prelude::{MessageSender, MessageSenderError};
29use crate::proto::sync_message::PniChangeNumber;
30use crate::proto::{DeviceName, SyncMessage};
31use crate::provisioning::generate_registration_id;
32use crate::push_service::{
33 AvatarWrite, CaptchaAttributes, DeviceActivationRequest, DeviceInfo,
34 HttpAuthOverride, RegistrationMethod, ReqwestExt, VerifyAccountResponse,
35 DEFAULT_DEVICE_ID,
36};
37use crate::sender::OutgoingPushMessage;
38use crate::service_address::ServiceIdExt;
39use crate::session_store::SessionStoreExt;
40use crate::timestamp::TimestampExt as _;
41use crate::utils::{random_length_padding, BASE64_RELAXED};
42use crate::{
43 configuration::{Endpoint, ServiceCredentials},
44 pre_keys::PreKeyState,
45 profile_cipher::{ProfileCipher, ProfileCipherError},
46 profile_name::ProfileName,
47 proto::{ProvisionEnvelope, ProvisionMessage, ProvisioningVersion},
48 provisioning::{ProvisioningCipher, ProvisioningError},
49 push_service::{AccountAttributes, PushService, ServiceError},
50 utils::serde_base64,
51};
52
53type Aes256Ctr128BE = ctr::Ctr128BE<aes::Aes256>;
54
55pub struct AccountManager {
56 service: PushService,
57 profile_key: Option<ProfileKey>,
58}
59
60#[derive(thiserror::Error, Debug)]
61pub enum ProfileManagerError {
62 #[error(transparent)]
63 ServiceError(#[from] ServiceError),
64 #[error(transparent)]
65 ProfileCipherError(#[from] ProfileCipherError),
66}
67
68#[derive(Debug, Default, Serialize, Deserialize, Clone)]
69pub struct Profile {
70 pub name: Option<ProfileName<String>>,
71 pub about: Option<String>,
72 pub about_emoji: Option<String>,
73 pub avatar: Option<String>,
74 pub unrestricted_unidentified_access: bool,
75}
76
77impl AccountManager {
78 pub fn new(service: PushService, profile_key: Option<ProfileKey>) -> Self {
79 Self {
80 service,
81 profile_key,
82 }
83 }
84
85 #[allow(clippy::too_many_arguments)]
86 #[tracing::instrument(skip(self, protocol_store))]
87 pub async fn check_pre_keys<P: PreKeysStore>(
88 &mut self,
89 protocol_store: &mut P,
90 service_id_kind: ServiceIdKind,
91 ) -> Result<bool, ServiceError> {
92 let Some(signed_prekey_id) = protocol_store.signed_prekey_id().await?
93 else {
94 tracing::warn!("No signed prekey found");
95 return Ok(false);
96 };
97 let Some(kyber_prekey_id) =
100 protocol_store.last_resort_kyber_prekey_id().await?
101 else {
102 tracing::warn!("No last resort kyber prekey found");
103 return Ok(false);
104 };
105
106 let signed_prekey =
107 protocol_store.get_signed_pre_key(signed_prekey_id).await?;
108 let kyber_prekey =
109 protocol_store.get_kyber_pre_key(kyber_prekey_id).await?;
110
111 let mut hash = Sha256::default();
113 hash.update(
114 protocol_store
115 .get_identity_key_pair()
116 .await?
117 .public_key()
118 .serialize(),
119 );
120 hash.update((u32::from(signed_prekey_id) as u64).to_be_bytes());
121 hash.update(signed_prekey.public_key()?.serialize());
122 hash.update((u32::from(kyber_prekey_id) as u64).to_be_bytes());
123 hash.update(kyber_prekey.public_key()?.serialize());
124
125 self.service
126 .check_pre_keys(service_id_kind, hash.finalize().as_ref())
127 .await
128 }
129
130 #[allow(clippy::too_many_arguments)]
137 #[tracing::instrument(skip(self, csprng, protocol_store))]
138 pub async fn update_pre_key_bundle<R: Rng + CryptoRng, P: PreKeysStore>(
139 &mut self,
140 protocol_store: &mut P,
141 service_id_kind: ServiceIdKind,
142 use_last_resort_key: bool,
143 csprng: &mut R,
144 ) -> Result<(), ServiceError> {
145 let prekey_status = match self
146 .service
147 .get_pre_key_status(service_id_kind)
148 .instrument(tracing::span!(
149 tracing::Level::DEBUG,
150 "Fetching pre key status"
151 ))
152 .await
153 {
154 Ok(status) => status,
155 Err(ServiceError::Unauthorized) => {
156 tracing::info!("Got Unauthorized when fetching pre-key status. Assuming first installment.");
157 crate::push_service::PreKeyStatus {
160 count: 0,
161 pq_count: 0,
162 }
163 },
164 Err(e) => return Err(e),
165 };
166 tracing::trace!("Remaining pre-keys on server: {:?}", prekey_status);
167
168 let check_pre_keys = self
169 .check_pre_keys(protocol_store, service_id_kind)
170 .instrument(tracing::span!(
171 tracing::Level::DEBUG,
172 "Checking pre keys"
173 ))
174 .await?;
175 if !check_pre_keys {
176 tracing::info!(
177 "Last resort pre-keys are not up to date; refreshing."
178 );
179 } else {
180 tracing::debug!("Last resort pre-keys are up to date.");
181 }
182
183 if check_pre_keys
188 && (prekey_status.count >= PRE_KEY_MINIMUM
189 && prekey_status.pq_count >= PRE_KEY_MINIMUM)
190 {
191 if protocol_store.signed_pre_keys_count().await? > 0
192 && protocol_store.kyber_pre_keys_count(true).await? > 0
193 && protocol_store.signed_prekey_id().await?.is_some()
194 && protocol_store
195 .last_resort_kyber_prekey_id()
196 .await?
197 .is_some()
198 {
199 tracing::debug!("Available keys sufficient");
200 return Ok(());
201 }
202 tracing::info!("Available keys sufficient; forcing refresh.");
203 }
204
205 let identity_key_pair = protocol_store
206 .get_identity_key_pair()
207 .instrument(tracing::trace_span!("get identity key pair"))
208 .await?;
209
210 let last_resort_keys = protocol_store
211 .load_last_resort_kyber_pre_keys()
212 .instrument(tracing::trace_span!("fetch last resort key"))
213 .await?;
214
215 let has_last_resort_key = !last_resort_keys.is_empty();
217
218 let (pre_keys, signed_pre_key, pq_pre_keys, pq_last_resort_key) =
219 crate::pre_keys::replenish_pre_keys(
220 protocol_store,
221 csprng,
222 &identity_key_pair,
223 use_last_resort_key && !has_last_resort_key,
224 PRE_KEY_BATCH_SIZE,
225 PRE_KEY_BATCH_SIZE,
226 )
227 .await?;
228
229 let pq_last_resort_key = if has_last_resort_key {
230 if last_resort_keys.len() > 1 {
231 tracing::warn!(
232 "More than one last resort key found; only uploading first"
233 );
234 }
235 Some(KyberPreKeyEntity::try_from(last_resort_keys[0].clone())?)
236 } else {
237 pq_last_resort_key
238 .map(KyberPreKeyEntity::try_from)
239 .transpose()?
240 };
241
242 let identity_key = *identity_key_pair.identity_key();
243
244 let pre_keys: Vec<_> = pre_keys
245 .into_iter()
246 .map(PreKeyEntity::try_from)
247 .collect::<Result<_, _>>()?;
248 let signed_pre_key = signed_pre_key.try_into()?;
249 let pq_pre_keys: Vec<_> = pq_pre_keys
250 .into_iter()
251 .map(KyberPreKeyEntity::try_from)
252 .collect::<Result<_, _>>()?;
253
254 tracing::info!(
255 "Uploading pre-keys: {} one-time, {} PQ, {} PQ last resort",
256 pre_keys.len(),
257 pq_pre_keys.len(),
258 if pq_last_resort_key.is_some() { 1 } else { 0 }
259 );
260
261 let pre_key_state = PreKeyState {
262 pre_keys,
263 signed_pre_key,
264 identity_key,
265 pq_pre_keys,
266 pq_last_resort_key,
267 };
268
269 self.service
270 .register_pre_keys(service_id_kind, pre_key_state)
271 .instrument(tracing::span!(
272 tracing::Level::DEBUG,
273 "Uploading pre keys"
274 ))
275 .await?;
276
277 Ok(())
278 }
279
280 async fn new_device_provisioning_code(
281 &mut self,
282 ) -> Result<String, ServiceError> {
283 #[derive(serde::Deserialize)]
284 #[serde(rename_all = "camelCase")]
285 struct DeviceCode {
286 verification_code: String,
287 }
288
289 let dc: DeviceCode = self
290 .service
291 .request(
292 Method::GET,
293 Endpoint::service("/v1/devices/provisioning/code"),
294 HttpAuthOverride::NoOverride,
295 )?
296 .send()
297 .await?
298 .service_error_for_status()
299 .await?
300 .json()
301 .await?;
302
303 Ok(dc.verification_code)
304 }
305
306 async fn send_provisioning_message(
307 &mut self,
308 destination: &str,
309 env: ProvisionEnvelope,
310 ) -> Result<(), ServiceError> {
311 #[derive(serde::Serialize)]
312 struct ProvisioningMessage {
313 body: String,
314 }
315
316 let body = env.encode_to_vec();
317
318 self.service
319 .request(
320 Method::PUT,
321 Endpoint::service(format!("/v1/provisioning/{destination}")),
322 HttpAuthOverride::NoOverride,
323 )?
324 .json(&ProvisioningMessage {
325 body: BASE64_RELAXED.encode(body),
326 })
327 .send()
328 .await?
329 .service_error_for_status()
330 .await?;
331
332 Ok(())
333 }
334
335 pub async fn link_device<R: Rng + CryptoRng>(
349 &mut self,
350 csprng: &mut R,
351 url: url::Url,
352 aci_identity_store: &dyn IdentityKeyStore,
353 pni_identity_store: &dyn IdentityKeyStore,
354 credentials: ServiceCredentials,
355 master_key: Option<MasterKey>,
356 ) -> Result<(), ProvisioningError> {
357 let query: HashMap<_, _> = url.query_pairs().collect();
358 let ephemeral_id =
359 query.get("uuid").ok_or(ProvisioningError::MissingUuid)?;
360 let pub_key = query
361 .get("pub_key")
362 .ok_or(ProvisioningError::MissingPublicKey)?;
363
364 let pub_key = BASE64_RELAXED
365 .decode(&**pub_key)
366 .map_err(|e| ProvisioningError::InvalidPublicKey(e.into()))?;
367 let pub_key = PublicKey::deserialize(&pub_key)
368 .map_err(|e| ProvisioningError::InvalidPublicKey(e.into()))?;
369
370 let aci_identity_key_pair =
371 aci_identity_store.get_identity_key_pair().await?;
372 let pni_identity_key_pair =
373 pni_identity_store.get_identity_key_pair().await?;
374
375 if credentials.aci.is_none() {
376 tracing::warn!("No local ACI set");
377 }
378 if credentials.pni.is_none() {
379 tracing::warn!("No local PNI set");
380 }
381
382 let provisioning_code = self.new_device_provisioning_code().await?;
383
384 let msg = ProvisionMessage {
385 aci: credentials.aci.as_ref().map(|u| u.to_string()),
386 aci_binary: credentials.aci.map(|u| u.into_bytes().into()),
387 aci_identity_key_public: Some(
388 aci_identity_key_pair.public_key().serialize().into_vec(),
389 ),
390 aci_identity_key_private: Some(
391 aci_identity_key_pair.private_key().serialize(),
392 ),
393 number: Some(credentials.e164()),
394 pni_identity_key_public: Some(
395 pni_identity_key_pair.public_key().serialize().into_vec(),
396 ),
397 pni_identity_key_private: Some(
398 pni_identity_key_pair.private_key().serialize(),
399 ),
400 pni: credentials.pni.as_ref().map(uuid::Uuid::to_string),
401 pni_binary: credentials.pni.map(|u| u.into_bytes().into()),
402 profile_key: self.profile_key.as_ref().map(|x| x.bytes.to_vec()),
403 provisioning_version: Some(i32::from(
405 ProvisioningVersion::TabletSupport,
406 ) as _),
407 provisioning_code: Some(provisioning_code),
408 read_receipts: None,
409 user_agent: None,
410 master_key: master_key.map(|x| x.into()),
411 ephemeral_backup_key: None,
412 account_entropy_pool: None,
413 media_root_backup_key: None,
414 };
415
416 let cipher = ProvisioningCipher::from_public(pub_key);
417
418 let encrypted = cipher.encrypt(csprng, msg)?;
419 self.send_provisioning_message(ephemeral_id, encrypted)
420 .await?;
421 Ok(())
422 }
423
424 pub async fn linked_devices(
425 &mut self,
426 aci_identity_store: &dyn IdentityKeyStore,
427 ) -> Result<Vec<DeviceInfo>, ServiceError> {
428 let device_infos = self.service.devices().await?;
429 let aci_identity_keypair =
430 aci_identity_store.get_identity_key_pair().await?;
431
432 device_infos
433 .into_iter()
434 .map(|i| {
435 Ok(DeviceInfo {
436 id: i.id,
437 name: i.name.and_then(|s| {
438 match decrypt_device_name_from_device_info(
439 &s,
440 &aci_identity_keypair,
441 ) {
442 Ok(name) => Some(name),
443 Err(e) => {
444 tracing::error!("{e}");
445 None
446 },
447 }
448 }),
449 created: i.created,
450 last_seen: i.last_seen,
451 })
452 })
453 .collect()
454 }
455
456 pub async fn register_account<
457 R: Rng + CryptoRng,
458 Aci: PreKeysStore + IdentityKeyStore,
459 Pni: PreKeysStore + IdentityKeyStore,
460 >(
461 &mut self,
462 csprng: &mut R,
463 registration_method: RegistrationMethod<'_>,
464 account_attributes: AccountAttributes,
465 aci_protocol_store: &mut Aci,
466 pni_protocol_store: &mut Pni,
467 skip_device_transfer: bool,
468 ) -> Result<VerifyAccountResponse, ProvisioningError> {
469 let aci_identity_key_pair = aci_protocol_store
470 .get_identity_key_pair()
471 .instrument(tracing::trace_span!("get ACI identity key pair"))
472 .await?;
473 let pni_identity_key_pair = pni_protocol_store
474 .get_identity_key_pair()
475 .instrument(tracing::trace_span!("get PNI identity key pair"))
476 .await?;
477
478 let (
479 _aci_pre_keys,
480 aci_signed_pre_key,
481 _aci_kyber_pre_keys,
482 aci_last_resort_kyber_prekey,
483 ) = crate::pre_keys::replenish_pre_keys(
484 aci_protocol_store,
485 csprng,
486 &aci_identity_key_pair,
487 true,
488 0,
489 0,
490 )
491 .await?;
492
493 let (
494 _pni_pre_keys,
495 pni_signed_pre_key,
496 _pni_kyber_pre_keys,
497 pni_last_resort_kyber_prekey,
498 ) = crate::pre_keys::replenish_pre_keys(
499 pni_protocol_store,
500 csprng,
501 &pni_identity_key_pair,
502 true,
503 0,
504 0,
505 )
506 .await?;
507
508 let aci_identity_key = aci_identity_key_pair.identity_key();
509 let pni_identity_key = pni_identity_key_pair.identity_key();
510
511 let dar = DeviceActivationRequest {
512 aci_signed_pre_key: aci_signed_pre_key.try_into()?,
513 pni_signed_pre_key: pni_signed_pre_key.try_into()?,
514 aci_pq_last_resort_pre_key: aci_last_resort_kyber_prekey
515 .expect("requested last resort prekey")
516 .try_into()?,
517 pni_pq_last_resort_pre_key: pni_last_resort_kyber_prekey
518 .expect("requested last resort prekey")
519 .try_into()?,
520 };
521
522 let result = self
523 .service
524 .submit_registration_request(
525 registration_method,
526 account_attributes,
527 skip_device_transfer,
528 aci_identity_key,
529 pni_identity_key,
530 dar,
531 )
532 .await?;
533
534 Ok(result)
535 }
536
537 pub async fn upload_versioned_profile_without_avatar<
548 R: Rng + CryptoRng,
549 S: AsRef<str>,
550 >(
551 &mut self,
552 aci: libsignal_protocol::Aci,
553 name: ProfileName<S>,
554 about: Option<String>,
555 about_emoji: Option<String>,
556 retain_avatar: bool,
557 csprng: &mut R,
558 ) -> Result<(), ProfileManagerError> {
559 self.upload_versioned_profile::<std::io::Cursor<Vec<u8>>, _, _>(
560 aci,
561 name,
562 about,
563 about_emoji,
564 if retain_avatar {
565 AvatarWrite::RetainAvatar
566 } else {
567 AvatarWrite::NoAvatar
568 },
569 csprng,
570 )
571 .await?;
572 Ok(())
573 }
574
575 pub async fn retrieve_profile(
576 &mut self,
577 address: Aci,
578 ) -> Result<Profile, ProfileManagerError> {
579 let profile_key =
580 self.profile_key.expect("set profile key in AccountManager");
581
582 let encrypted_profile = self
583 .service
584 .retrieve_profile_by_id(address, Some(profile_key))
585 .await?;
586
587 let profile_cipher = ProfileCipher::new(profile_key);
588 Ok(profile_cipher.decrypt(encrypted_profile)?)
589 }
590
591 pub async fn upload_versioned_profile<
597 's,
598 C: std::io::Read + Send + 's,
599 R: Rng + CryptoRng,
600 S: AsRef<str>,
601 >(
602 &mut self,
603 aci: libsignal_protocol::Aci,
604 name: ProfileName<S>,
605 about: Option<String>,
606 about_emoji: Option<String>,
607 avatar: AvatarWrite<&'s mut C>,
608 csprng: &mut R,
609 ) -> Result<Option<String>, ProfileManagerError> {
610 let profile_key =
611 self.profile_key.expect("set profile key in AccountManager");
612 let profile_cipher = ProfileCipher::new(profile_key);
613
614 let name = profile_cipher.encrypt_name(name.as_ref(), csprng)?;
616 let about = about.unwrap_or_default();
617 let about = profile_cipher.encrypt_about(about, csprng)?;
618 let about_emoji = about_emoji.unwrap_or_default();
619 let about_emoji = profile_cipher.encrypt_emoji(about_emoji, csprng)?;
620
621 if matches!(avatar, AvatarWrite::NewAvatar(_)) {
623 unimplemented!("Setting avatar requires ProfileCipherStream")
626 }
627
628 let profile_key = profile_cipher.into_inner();
629 let commitment = profile_key.get_commitment(aci);
630 let profile_key_version = profile_key.get_profile_key_version(aci);
631
632 Ok(self
633 .service
634 .write_profile::<C, S>(
635 &profile_key_version,
636 &name,
637 &about,
638 &about_emoji,
639 &commitment,
640 avatar,
641 )
642 .await?)
643 }
644
645 pub async fn set_account_attributes(
649 &mut self,
650 attributes: AccountAttributes,
651 ) -> Result<(), ServiceError> {
652 self.service.set_account_attributes(attributes).await
653 }
654
655 pub async fn update_device_name<R: Rng + CryptoRng>(
657 &mut self,
658 device_id: libsignal_core::DeviceId,
659 device_name: &str,
660 aci: Aci,
661 aci_identity_store: &dyn IdentityKeyStore,
662 csprng: &mut R,
663 ) -> Result<(), ServiceError> {
664 let addr = aci.to_protocol_address(device_id).unwrap();
665 let public_key = aci_identity_store.get_identity(&addr).await?;
666 let Some(public_key) = public_key else {
667 return Err(ServiceError::SendError {
668 reason: format!("public key for device {addr:?} not found"),
669 });
670 };
671 let encrypted_device_name =
672 encrypt_device_name(csprng, device_name, &public_key)?;
673
674 #[derive(Serialize)]
675 #[serde(rename_all = "camelCase")]
676 struct Data {
677 #[serde(with = "serde_base64")]
678 device_name: Vec<u8>,
679 }
680
681 self.service
682 .request(
683 Method::PUT,
684 Endpoint::service(format!(
685 "/v1/accounts/name?deviceId={}",
686 device_id
687 )),
688 HttpAuthOverride::NoOverride,
689 )?
690 .json(&Data {
691 device_name: encrypted_device_name.encode_to_vec(),
692 })
693 .send()
694 .await?
695 .service_error_for_status()
696 .await?;
697
698 Ok(())
699 }
700
701 pub async fn submit_recaptcha_challenge(
708 &mut self,
709 token: &str,
710 captcha: &str,
711 ) -> Result<(), ServiceError> {
712 self.service
713 .request(
714 Method::PUT,
715 Endpoint::service("/v1/challenge"),
716 HttpAuthOverride::NoOverride,
717 )?
718 .json(&CaptchaAttributes {
719 challenge_type: "captcha",
720 token,
721 captcha,
722 })
723 .send()
724 .await?
725 .service_error_for_status()
726 .await?;
727
728 Ok(())
729 }
730
731 #[tracing::instrument(skip(self, aci_protocol_store, pni_protocol_store, sender, local_aci, csprng), fields(local_aci = local_aci.service_id_string()))]
737 pub async fn pnp_initialize_devices<
738 R: Rng + CryptoRng,
739 AciStore: PreKeysStore + SessionStoreExt,
740 PniStore: PreKeysStore,
741 AciOrPni: ProtocolStore + SenderKeyStore + SessionStoreExt + Sync + Clone,
742 >(
743 &mut self,
744 aci_protocol_store: &mut AciStore,
745 pni_protocol_store: &mut PniStore,
746 mut sender: MessageSender<AciOrPni>,
747 local_aci: Aci,
748 e164: PhoneNumber,
749 csprng: &mut R,
750 ) -> Result<(), MessageSenderError> {
751 let pni_identity_key_pair =
752 pni_protocol_store.get_identity_key_pair().await?;
753
754 let pni_identity_key = pni_identity_key_pair.identity_key();
755
756 let local_device_ids = aci_protocol_store
758 .get_sub_device_sessions(&local_aci.into())
759 .await?;
760
761 let mut device_messages =
762 Vec::<OutgoingPushMessage>::with_capacity(local_device_ids.len());
763 let mut device_pni_signed_prekeys =
764 HashMap::<String, SignedPreKeyEntity>::with_capacity(
765 local_device_ids.len(),
766 );
767 let mut device_pni_last_resort_kyber_prekeys =
768 HashMap::<String, KyberPreKeyEntity>::with_capacity(
769 local_device_ids.len(),
770 );
771 let mut pni_registration_ids =
772 HashMap::<String, u32>::with_capacity(local_device_ids.len());
773
774 let signature_valid_on_each_signed_pre_key = true;
775 for local_device_id in
776 std::iter::once(*DEFAULT_DEVICE_ID).chain(local_device_ids)
777 {
778 let local_protocol_address =
779 local_aci.to_protocol_address(local_device_id)?;
780 let span = tracing::trace_span!(
781 "filtering devices",
782 address = %local_protocol_address
783 );
784 if (local_device_id != *DEFAULT_DEVICE_ID)
786 && aci_protocol_store
787 .load_session(&local_protocol_address)
788 .instrument(span)
789 .await?
790 .is_none()
791 {
792 tracing::warn!(
793 "No session with device {}, skipping PNI provisioning",
794 local_device_id
795 );
796 continue;
797 }
798 let (
799 _pre_keys,
800 signed_pre_key,
801 _kyber_pre_keys,
802 last_resort_kyber_prekey,
803 ) = if local_device_id == *DEFAULT_DEVICE_ID {
804 crate::pre_keys::replenish_pre_keys(
805 pni_protocol_store,
806 csprng,
807 &pni_identity_key_pair,
808 true,
809 0,
810 0,
811 )
812 .await?
813 } else {
814 let signed_pre_key_pair = KeyPair::generate(csprng);
816 let signed_pre_key_public = signed_pre_key_pair.public_key;
817 let signed_pre_key_signature = pni_identity_key_pair
818 .private_key()
819 .calculate_signature(
820 &signed_pre_key_public.serialize(),
821 csprng,
822 )
823 .map_err(MessageSenderError::InvalidPrivateKey)?;
824
825 let signed_prekey_record = SignedPreKeyRecord::new(
826 csprng.random_range::<u32, _>(0..0xFFFFFF).into(),
827 Timestamp::now(),
828 &signed_pre_key_pair,
829 &signed_pre_key_signature,
830 );
831
832 let kyber_pre_key_record = KyberPreKeyRecord::generate(
834 kem::KeyType::Kyber1024,
835 csprng.random_range::<u32, _>(0..0xFFFFFF).into(),
836 pni_identity_key_pair.private_key(),
837 )?;
838 (
839 vec![],
840 signed_prekey_record,
841 vec![],
842 Some(kyber_pre_key_record),
843 )
844 };
845
846 let registration_id = if local_device_id == *DEFAULT_DEVICE_ID {
847 pni_protocol_store.get_local_registration_id().await?
848 } else {
849 loop {
850 let regid = generate_registration_id(csprng);
851 if !pni_registration_ids.iter().any(|(_k, v)| *v == regid) {
852 break regid;
853 }
854 }
855 };
856
857 let local_device_id_s = local_device_id.to_string();
858 device_pni_signed_prekeys.insert(
859 local_device_id_s.clone(),
860 SignedPreKeyEntity::try_from(&signed_pre_key)?,
861 );
862 device_pni_last_resort_kyber_prekeys.insert(
863 local_device_id_s.clone(),
864 KyberPreKeyEntity::try_from(
865 last_resort_kyber_prekey
866 .as_ref()
867 .expect("requested last resort key"),
868 )?,
869 );
870 pni_registration_ids
871 .insert(local_device_id_s.clone(), registration_id);
872
873 assert!(_pre_keys.is_empty());
874 assert!(_kyber_pre_keys.is_empty());
875
876 if local_device_id == *DEFAULT_DEVICE_ID {
877 continue;
880 }
881 let msg = SyncMessage {
883 pni_change_number: Some(PniChangeNumber {
884 identity_key_pair: Some(
885 pni_identity_key_pair.serialize().to_vec(),
886 ),
887 signed_pre_key: Some(signed_pre_key.serialize()?),
888 last_resort_kyber_pre_key: Some(
889 last_resort_kyber_prekey
890 .expect("requested last resort key")
891 .serialize()?,
892 ),
893 registration_id: Some(registration_id),
894 new_e164: Some(
895 e164.format().mode(phonenumber::Mode::E164).to_string(),
896 ),
897 }),
898 padding: Some(random_length_padding(csprng, 512)),
899 ..SyncMessage::default()
900 };
901 let content: ContentBody = msg.into();
902 let msg = sender
903 .create_encrypted_message(
904 &local_aci.into(),
905 None,
906 local_device_id,
907 &content.into_proto().encode_to_vec(),
908 )
909 .await?;
910 device_messages.push(msg);
911 }
912
913 self.service
914 .distribute_pni_keys(
915 pni_identity_key,
916 device_messages,
917 device_pni_signed_prekeys,
918 device_pni_last_resort_kyber_prekeys,
919 pni_registration_ids,
920 signature_valid_on_each_signed_pre_key,
921 )
922 .await?;
923
924 Ok(())
925 }
926}
927
928#[expect(clippy::result_large_err)]
929fn calculate_hmac256(
930 mac_key: &[u8],
931 ciphertext: &[u8],
932) -> Result<Output<Hmac<Sha256>>, ServiceError> {
933 let mut mac = Hmac::<Sha256>::new_from_slice(mac_key)
934 .map_err(|_| ServiceError::MacError)?;
935 mac.update(ciphertext);
936 Ok(mac.finalize().into_bytes())
937}
938
939#[expect(clippy::result_large_err)]
940pub fn encrypt_device_name<R: rand::Rng + rand::CryptoRng>(
941 csprng: &mut R,
942 device_name: &str,
943 identity_public: &IdentityKey,
944) -> Result<DeviceName, ServiceError> {
945 let plaintext = device_name.as_bytes().to_vec();
946 let ephemeral_key_pair = KeyPair::generate(csprng);
947
948 let master_secret = ephemeral_key_pair
949 .private_key
950 .calculate_agreement(identity_public.public_key())?;
951
952 let key1 = calculate_hmac256(&master_secret, b"auth")?;
953 let synthetic_iv = calculate_hmac256(&key1, &plaintext)?;
954 let synthetic_iv = &synthetic_iv[..16];
955
956 let key2 = calculate_hmac256(&master_secret, b"cipher")?;
957 let cipher_key = calculate_hmac256(&key2, synthetic_iv)?;
958
959 let mut ciphertext = plaintext;
960
961 const IV: [u8; 16] = [0; 16];
962 let mut cipher = Aes256Ctr128BE::new(&cipher_key, &IV.into());
963 cipher.apply_keystream(&mut ciphertext);
964
965 let device_name = DeviceName {
966 ephemeral_public: Some(
967 ephemeral_key_pair.public_key.serialize().to_vec(),
968 ),
969 synthetic_iv: Some(synthetic_iv.to_vec()),
970 ciphertext: Some(ciphertext),
971 };
972
973 Ok(device_name)
974}
975
976#[expect(clippy::result_large_err)]
977fn decrypt_device_name_from_device_info(
978 string: &str,
979 aci: &IdentityKeyPair,
980) -> Result<String, ServiceError> {
981 let data = BASE64_RELAXED.decode(string)?;
982 let name = DeviceName::decode(&*data)?;
983 crate::decrypt_device_name(aci.private_key(), &name)
984}
985
986#[expect(clippy::result_large_err)]
987pub fn decrypt_device_name(
988 private_key: &PrivateKey,
989 device_name: &DeviceName,
990) -> Result<String, ServiceError> {
991 let DeviceName {
992 ephemeral_public: Some(ephemeral_public),
993 synthetic_iv: Some(synthetic_iv),
994 ciphertext: Some(ciphertext),
995 } = device_name
996 else {
997 return Err(ServiceError::InvalidDeviceName);
998 };
999
1000 let synthetic_iv: [u8; 16] = synthetic_iv[..synthetic_iv.len().min(16)]
1001 .try_into()
1002 .map_err(|_| ServiceError::MacError)?;
1003
1004 let ephemeral_public = PublicKey::deserialize(ephemeral_public)?;
1005
1006 let master_secret = private_key.calculate_agreement(&ephemeral_public)?;
1007 let key2 = calculate_hmac256(&master_secret, b"cipher")?;
1008 let cipher_key = calculate_hmac256(&key2, &synthetic_iv)?;
1009
1010 let mut plaintext = ciphertext.to_vec();
1011 const IV: [u8; 16] = [0; 16];
1012 let mut cipher =
1013 Aes256Ctr128BE::new(cipher_key.as_slice().into(), &IV.into());
1014 cipher.apply_keystream(&mut plaintext);
1015
1016 let key1 = calculate_hmac256(&master_secret, b"auth")?;
1017 let our_synthetic_iv = calculate_hmac256(&key1, &plaintext)?;
1018 let our_synthetic_iv = &our_synthetic_iv[..16];
1019
1020 if synthetic_iv != our_synthetic_iv {
1021 Err(ServiceError::MacError)
1022 } else {
1023 Ok(String::from_utf8_lossy(&plaintext).to_string())
1024 }
1025}
1026
1027#[cfg(test)]
1028mod tests {
1029 use crate::utils::BASE64_RELAXED;
1030 use base64::Engine;
1031 use libsignal_protocol::{IdentityKeyPair, PrivateKey, PublicKey};
1032
1033 use super::DeviceName;
1034
1035 #[test]
1036 fn encrypt_device_name() -> anyhow::Result<()> {
1037 let input_device_name = "Nokia 3310 Millenial Edition";
1038 let mut csprng = rand::rng();
1039 let identity = IdentityKeyPair::generate(&mut csprng);
1040
1041 let device_name = super::encrypt_device_name(
1042 &mut csprng,
1043 input_device_name,
1044 identity.identity_key(),
1045 )?;
1046
1047 let decrypted_device_name =
1048 super::decrypt_device_name(identity.private_key(), &device_name)?;
1049
1050 assert_eq!(input_device_name, decrypted_device_name);
1051
1052 Ok(())
1053 }
1054
1055 #[test]
1056 fn decrypt_device_name() -> anyhow::Result<()> {
1057 let ephemeral_private_key = PrivateKey::deserialize(
1058 &BASE64_RELAXED
1059 .decode("0CgxHjwwblXjvX8sD5wZDWdYToMRf+CZSlgaUrxCGVo=")?,
1060 )?;
1061 let ephemeral_public_key = PublicKey::deserialize(
1062 &BASE64_RELAXED
1063 .decode("BcZS+Lt6yAKbEpXnRX+I5wHqesuvu93Q2V+fjidwW8R6")?,
1064 )?;
1065
1066 let device_name = DeviceName {
1067 ephemeral_public: Some(ephemeral_public_key.serialize().to_vec()),
1068 synthetic_iv: Some(
1069 BASE64_RELAXED.decode("86gekHGmltnnZ9QARhiFcg==")?,
1070 ),
1071 ciphertext: Some(
1072 BASE64_RELAXED
1073 .decode("MtJ9/9KBWLBVAxfZJD4pLKzP4q+iodRJeCc+/A==")?,
1074 ),
1075 };
1076
1077 let decrypted_device_name =
1078 super::decrypt_device_name(&ephemeral_private_key, &device_name)?;
1079
1080 assert_eq!(decrypted_device_name, "Nokia 3310 Millenial Edition");
1081
1082 Ok(())
1083 }
1084}