Migrate PushService methods to the emulated HTTP request/response inside of the websocket (#366)

Author: gferon

src/account_manager.rs

@@ -30,30 +30,37 @@ use crate::proto::sync_message::PniChangeNumber;
 use crate::proto::{DeviceName, SyncMessage};
 use crate::provisioning::generate_registration_id;
 use crate::push_service::{
-    AvatarWrite, CaptchaAttributes, DeviceActivationRequest, DeviceInfo,
-    HttpAuthOverride, RegistrationMethod, ReqwestExt, VerifyAccountResponse,
-    DEFAULT_DEVICE_ID,
+    AvatarWrite, HttpAuthOverride, ReqwestExt, DEFAULT_DEVICE_ID,
 };
 use crate::sender::OutgoingPushMessage;
 use crate::service_address::ServiceIdExt;
 use crate::session_store::SessionStoreExt;
 use crate::timestamp::TimestampExt as _;
 use crate::utils::{random_length_padding, BASE64_RELAXED};
+use crate::websocket::account::DeviceInfo;
+use crate::websocket::keys::PreKeyStatus;
+use crate::websocket::registration::{
+    CaptchaAttributes, DeviceActivationRequest, RegistrationMethod,
+    VerifyAccountResponse,
+};
+use crate::websocket::{self, SignalWebSocket};
 use crate::{
     configuration::{Endpoint, ServiceCredentials},
     pre_keys::PreKeyState,
     profile_cipher::{ProfileCipher, ProfileCipherError},
     profile_name::ProfileName,
     proto::{ProvisionEnvelope, ProvisionMessage, ProvisioningVersion},
     provisioning::{ProvisioningCipher, ProvisioningError},
-    push_service::{AccountAttributes, PushService, ServiceError},
+    push_service::{PushService, ServiceError},
     utils::serde_base64,
+    websocket::account::AccountAttributes,
 };
 
 type Aes256Ctr128BE = ctr::Ctr128BE<aes::Aes256>;
 
 pub struct AccountManager {
     service: PushService,
+    websocket: SignalWebSocket<websocket::Identified>,
     profile_key: Option<ProfileKey>,
 }
 
@@ -75,9 +82,14 @@ pub struct Profile {
 }
 
 impl AccountManager {
-    pub fn new(service: PushService, profile_key: Option<ProfileKey>) -> Self {
+    pub fn new(
+        service: PushService,
+        websocket: SignalWebSocket<websocket::Identified>,
+        profile_key: Option<ProfileKey>,
+    ) -> Self {
         Self {
             service,
+            websocket,
             profile_key,
         }
     }
@@ -122,7 +134,7 @@ impl AccountManager {
         hash.update((u32::from(kyber_prekey_id) as u64).to_be_bytes());
         hash.update(kyber_prekey.public_key()?.serialize());
 
-        self.service
+        self.websocket
             .check_pre_keys(service_id_kind, hash.finalize().as_ref())
             .await
     }
@@ -134,16 +146,15 @@ impl AccountManager {
     ///
     /// Equivalent to Java's RefreshPreKeysJob
     #[allow(clippy::too_many_arguments)]
-    #[tracing::instrument(skip(self, csprng, protocol_store))]
-    pub async fn update_pre_key_bundle<R: Rng + CryptoRng, P: PreKeysStore>(
+    #[tracing::instrument(skip(self, protocol_store))]
+    pub async fn update_pre_key_bundle<P: PreKeysStore>(
         &mut self,
         protocol_store: &mut P,
         service_id_kind: ServiceIdKind,
         use_last_resort_key: bool,
-        csprng: &mut R,
     ) -> Result<(), ServiceError> {
         let prekey_status = match self
-            .service
+            .websocket
             .get_pre_key_status(service_id_kind)
             .instrument(tracing::span!(
                 tracing::Level::DEBUG,
@@ -156,7 +167,7 @@ impl AccountManager {
                 tracing::info!("Got Unauthorized when fetching pre-key status. Assuming first installment.");
                 // Additionally, the second PUT request will fail if this really comes down to an
                 // authorization failure.
-                crate::push_service::PreKeyStatus {
+                PreKeyStatus {
                     count: 0,
                     pq_count: 0,
                 }
@@ -218,7 +229,7 @@ impl AccountManager {
         let (pre_keys, signed_pre_key, pq_pre_keys, pq_last_resort_key) =
             crate::pre_keys::replenish_pre_keys(
                 protocol_store,
-                csprng,
+                &mut rand::rng(),
                 &identity_key_pair,
                 use_last_resort_key && !has_last_resort_key,
                 PRE_KEY_BATCH_SIZE,
@@ -266,7 +277,7 @@ impl AccountManager {
             pq_last_resort_key,
         };
 
-        self.service
+        self.websocket
             .register_pre_keys(service_id_kind, pre_key_state)
             .instrument(tracing::span!(
                 tracing::Level::DEBUG,
@@ -425,7 +436,7 @@ impl AccountManager {
         &mut self,
         aci_identity_store: &dyn IdentityKeyStore,
     ) -> Result<Vec<DeviceInfo>, ServiceError> {
-        let device_infos = self.service.devices().await?;
+        let device_infos = self.websocket.devices().await?;
         let aci_identity_keypair =
             aci_identity_store.get_identity_key_pair().await?;
 
@@ -520,7 +531,7 @@ impl AccountManager {
         };
 
         let result = self
-            .service
+            .websocket
             .submit_registration_request(
                 registration_method,
                 account_attributes,
@@ -580,7 +591,7 @@ impl AccountManager {
             self.profile_key.expect("set profile key in AccountManager");
 
         let encrypted_profile = self
-            .service
+            .websocket
             .retrieve_profile_by_id(address, Some(profile_key))
             .await?;
 
@@ -630,7 +641,7 @@ impl AccountManager {
         let profile_key_version = profile_key.get_profile_key_version(aci);
 
         Ok(self
-            .service
+            .websocket
             .write_profile::<C, S>(
                 &profile_key_version,
                 &name,
@@ -649,7 +660,7 @@ impl AccountManager {
         &mut self,
         attributes: AccountAttributes,
     ) -> Result<(), ServiceError> {
-        self.service.set_account_attributes(attributes).await
+        self.websocket.set_account_attributes(attributes).await
     }
 
     /// Update (encrypted) device name
@@ -910,7 +921,7 @@ impl AccountManager {
             device_messages.push(msg);
         }
 
-        self.service
+        self.websocket
             .distribute_pni_keys(
                 pni_identity_key,
                 device_messages,
@@ -925,7 +936,6 @@ impl AccountManager {
     }
 }
 
-#[expect(clippy::result_large_err)]
 fn calculate_hmac256(
     mac_key: &[u8],
     ciphertext: &[u8],
@@ -936,7 +946,6 @@ fn calculate_hmac256(
     Ok(mac.finalize().into_bytes())
 }
 
-#[expect(clippy::result_large_err)]
 pub fn encrypt_device_name<R: rand::Rng + rand::CryptoRng>(
     csprng: &mut R,
     device_name: &str,
@@ -973,7 +982,6 @@ pub fn encrypt_device_name<R: rand::Rng + rand::CryptoRng>(
     Ok(device_name)
 }
 
-#[expect(clippy::result_large_err)]
 fn decrypt_device_name_from_device_info(
     string: &str,
     aci: &IdentityKeyPair,
@@ -983,7 +991,6 @@ fn decrypt_device_name_from_device_info(
     crate::decrypt_device_name(aci.private_key(), &name)
 }
 
-#[expect(clippy::result_large_err)]
 pub fn decrypt_device_name(
     private_key: &PrivateKey,
     device_name: &DeviceName,

src/cipher.rs

@@ -457,7 +457,7 @@ struct Plaintext {
     data: Vec<u8>,
 }
 
-#[expect(clippy::comparison_chain, clippy::result_large_err)]
+#[expect(clippy::comparison_chain)]
 fn add_padding(version: u32, contents: &[u8]) -> Result<Vec<u8>, ServiceError> {
     if version < 2 {
         Err(ServiceError::PaddingVersion(version))
@@ -480,7 +480,7 @@ fn add_padding(version: u32, contents: &[u8]) -> Result<Vec<u8>, ServiceError> {
     }
 }
 
-#[expect(clippy::comparison_chain, clippy::result_large_err)]
+#[expect(clippy::comparison_chain)]
 fn strip_padding_version(
     version: u32,
     contents: &mut Vec<u8>,
@@ -497,7 +497,6 @@ fn strip_padding_version(
     }
 }
 
-#[expect(clippy::result_large_err)]
 fn strip_padding(contents: &mut Vec<u8>) -> Result<(), ServiceError> {
     let new_length = Iso7816::raw_unpad(contents)?.len();
     contents.resize(new_length, 0);

src/content.rs

@@ -72,7 +72,6 @@ impl Content {
     }
 
     /// Converts a proto::Content into a public Content, including metadata.
-    #[expect(clippy::result_large_err)]
     pub fn from_proto(
         p: crate::proto::Content,
         metadata: Metadata,

src/envelope.rs

@@ -8,7 +8,6 @@ use crate::{configuration::SignalingKey, push_service::ServiceError};
 pub use crate::proto::Envelope;
 
 impl Envelope {
-    #[expect(clippy::result_large_err)]
     #[tracing::instrument(skip(input, signaling_key), fields(signaling_key_present = signaling_key.is_some(), input_size = input.len()))]
     pub fn decrypt(
         input: &[u8],

src/groups_v2/manager.rs

@@ -10,6 +10,7 @@ use crate::{
     proto::GroupContextV2,
     push_service::{HttpAuth, HttpAuthOverride, ReqwestExt, ServiceIds},
     utils::BASE64_RELAXED,
+    websocket::{self, SignalWebSocket},
 };
 
 use base64::prelude::*;
@@ -38,7 +39,6 @@ pub struct CredentialResponse {
     credentials: Vec<TemporalCredential>,
 }
 
-#[expect(clippy::result_large_err)]
 impl CredentialResponse {
     pub fn parse(
         self,
@@ -133,21 +133,24 @@ impl<T: CredentialsCache> CredentialsCache for &mut T {
 
 pub struct GroupsManager<C: CredentialsCache> {
     service_ids: ServiceIds,
-    push_service: PushService,
+    identified_push_service: PushService,
+    unidentified_websocket: SignalWebSocket<websocket::Unidentified>,
     credentials_cache: C,
     server_public_params: ServerPublicParams,
 }
 
 impl<C: CredentialsCache> GroupsManager<C> {
     pub fn new(
         service_ids: ServiceIds,
-        push_service: PushService,
+        identified_push_service: PushService,
+        unidentified_websocket: SignalWebSocket<websocket::Unidentified>,
         credentials_cache: C,
         server_public_params: ServerPublicParams,
     ) -> Self {
         Self {
             service_ids,
-            push_service,
+            identified_push_service,
+            unidentified_websocket,
             credentials_cache,
             server_public_params,
         }
@@ -169,7 +172,7 @@ impl<C: CredentialsCache> GroupsManager<C> {
             format!("/v1/certificate/auth/group?redemptionStartSeconds={}&redemptionEndSeconds={}&pniAsServiceId=true", today, today_plus_7_days);
 
             let credentials_response: CredentialResponse = self
-                .push_service
+                .identified_push_service
                 .request(
                     Method::GET,
                     Endpoint::service(path),
@@ -199,7 +202,6 @@ impl<C: CredentialsCache> GroupsManager<C> {
         )
     }
 
-    #[expect(clippy::result_large_err)]
     fn get_authorization_string<R: Rng + CryptoRng>(
         &self,
         csprng: &mut R,
@@ -257,7 +259,7 @@ impl<C: CredentialsCache> GroupsManager<C> {
         let authorization = self
             .get_authorization_for_today(csprng, group_secret_params)
             .await?;
-        self.push_service.get_group(authorization).await
+        self.identified_push_service.get_group(authorization).await
     }
 
     #[tracing::instrument(
@@ -270,7 +272,7 @@ impl<C: CredentialsCache> GroupsManager<C> {
         group_secret_params: GroupSecretParams,
     ) -> Result<Option<Vec<u8>>, ServiceError> {
         let mut encrypted_avatar = self
-            .push_service
+            .unidentified_websocket
             .retrieve_groups_v2_profile_avatar(path)
             .await?;
         let mut result = Vec::with_capacity(10 * 1024 * 1024);
@@ -301,7 +303,6 @@ impl<C: CredentialsCache> GroupsManager<C> {
     }
 }
 
-#[expect(clippy::result_large_err)]
 pub fn decrypt_group(
     master_key_bytes: &[u8],
     encrypted_group: crate::proto::Group,

src/messagepipe.rs

@@ -14,7 +14,10 @@ pub use crate::{
     },
 };
 
-use crate::{push_service::ServiceError, websocket::SignalWebSocket};
+use crate::{
+    push_service::ServiceError,
+    websocket::{self, SignalWebSocket},
+};
 
 #[derive(Debug)]
 pub enum Incoming {
@@ -23,20 +26,20 @@ pub enum Incoming {
 }
 
 pub struct MessagePipe {
-    ws: SignalWebSocket,
+    ws: SignalWebSocket<websocket::Identified>,
     credentials: ServiceCredentials,
 }
 
 impl MessagePipe {
     pub fn from_socket(
-        ws: SignalWebSocket,
+        ws: SignalWebSocket<websocket::Identified>,
         credentials: ServiceCredentials,
     ) -> Self {
         MessagePipe { ws, credentials }
     }
 
     /// Return a SignalWebSocket for sending messages and other purposes beyond receiving messages.
-    pub fn ws(&self) -> SignalWebSocket {
+    pub fn ws(&self) -> SignalWebSocket<websocket::Identified> {
         self.ws.clone()
     }
 

src/profile_cipher.rs

@@ -5,7 +5,8 @@ use rand::{rand_core, CryptoRng, RngCore};
 use zkgroup::profiles::ProfileKey;
 
 use crate::{
-    profile_name::ProfileName, push_service::SignalServiceProfile, Profile,
+    profile_name::ProfileName, websocket::profile::SignalServiceProfile,
+    Profile,
 };
 
 /// Encrypt and decrypt a [`ProfileName`] and other profile information.

src/proto.rs

@@ -1,6 +1,7 @@
 #![allow(clippy::all)]
 
 use rand::{CryptoRng, Rng};
+use reqwest::StatusCode;
 include!(concat!(env!("OUT_DIR"), "/signalservice.rs"));
 include!(concat!(env!("OUT_DIR"), "/signal.rs"));
 
@@ -66,6 +67,10 @@ impl WebSocketResponseMessage {
             }
         }
     }
+
+    pub fn status_code(&self) -> Option<reqwest::StatusCode> {
+        StatusCode::from_u16(self.status().try_into().ok()?).ok()
+    }
 }
 
 impl SyncMessage {