libsignal_service/
account_manager.rs

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        // XXX: should we instead use the `load_last_resort_kyber_pre_keys` method? Or refactor
98        //      those whole traits?
99        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        // `SHA256(identityKeyBytes || signedEcPreKeyId || signedEcPreKeyIdBytes || lastResortKeyId || lastResortKeyBytes)`
112        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    /// Checks the availability of pre-keys, and updates them as necessary.
131    ///
132    /// Parameters are the protocol's `StoreContext`, and the offsets for the next pre-key and
133    /// signed pre-keys.
134    ///
135    /// Equivalent to Java's RefreshPreKeysJob
136    #[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                // Additionally, the second PUT request will fail if this really comes down to an
158                // authorization failure.
159                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        // XXX We should honestly compare the pre-key count with the number of pre-keys we have
184        // locally. If we have more than the server, we should upload them.
185        // Currently the trait doesn't allow us to do that, so we just upload the batch size and
186        // pray.
187        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        // XXX: Maybe this check should be done in the generate_pre_keys function?
216        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    /// Link a new device, given a tsurl.
336    ///
337    /// Equivalent of Java's `AccountManager::addDevice()`
338    ///
339    /// When calling this, make sure that UnidentifiedDelivery is disabled, ie., that your
340    /// application does not send any unidentified messages before linking is complete.
341    /// Cfr.:
342    /// - `app/src/main/java/org/thoughtcrime/securesms/migrations/LegacyMigrationJob.java`:250 and;
343    /// - `app/src/main/java/org/thoughtcrime/securesms/DeviceActivity.java`:195
344    ///
345    /// ```java
346    /// TextSecurePreferences.setIsUnidentifiedDeliveryEnabled(context, false);
347    /// ```
348    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            // CURRENT is not exposed by prost :(
404            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    /// Upload a profile
538    ///
539    /// Panics if no `profile_key` was set.
540    ///
541    /// Convenience method for
542    /// ```ignore
543    /// manager.upload_versioned_profile::<std::io::Cursor<Vec<u8>>, _>(uuid, name, about, about_emoji, _)
544    /// ```
545    /// in which the `retain_avatar` parameter sets whether to remove (`false`) or retain (`true`) the
546    /// currently set avatar.
547    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    /// Upload a profile
592    ///
593    /// Panics if no `profile_key` was set.
594    ///
595    /// Returns the avatar url path.
596    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        // Profile encryption
615        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 avatar -> upload
622        if matches!(avatar, AvatarWrite::NewAvatar(_)) {
623            // FIXME ProfileCipherOutputStream.java
624            // It's just AES GCM, but a bit of work to decently implement it with a stream.
625            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    /// Set profile attributes
646    ///
647    /// Signal Android does not allow unsetting voice/video.
648    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    /// Update (encrypted) device name
656    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    /// Upload a proof-required reCaptcha token and response.
702    ///
703    /// Token gotten originally with HTTP status 428 response to sending a message.
704    /// Captcha gotten from user completing the challenge captcha.
705    ///
706    /// It's either a silent OK, or throws a ServiceError.
707    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    /// Initialize PNI on linked devices.
732    ///
733    /// Should be called as the primary device to migrate from pre-PNI to PNI.
734    ///
735    /// This is the equivalent of Android's PnpInitializeDevicesJob or iOS' PniHelloWorldManager.
736    #[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        // For every linked device, we generate a new set of pre-keys, and send them to the device.
757        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            // Skip if we don't have a session with the device
785            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                // Generate a signed prekey
815                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                // Generate a last-resort Kyber prekey
833                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                // This is the primary device
878                // We don't need to send a message to the primary device
879                continue;
880            }
881            // cfr. SignalServiceMessageSender::getEncryptedSyncPniInitializeDeviceMessage
882            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}