Add functions to retrieve avatars (whisperfish#234)

* Add functions to retrieve avatars --------- Co-authored-by: Gabriel Féron <[email protected]>
Schmiddiii · Feb 24, 2024 · 6c49f6f · 6c49f6f
1 parent 34c85e5
commit 6c49f6f
Show file tree

Hide file tree

Showing 6 changed files with 170 additions and 6 deletions.
diff --git a/presage-store-sled/src/lib.rs b/presage-store-sled/src/lib.rs
@@ -27,7 +27,7 @@ use presage::libsignal_service::{
     Profile, ServiceAddress,
 };
 use presage::store::{ContentExt, ContentsStore, StateStore, Store, Thread};
-use presage::{manager::RegistrationData, proto::verified};
+use presage::{manager::RegistrationData, proto::verified, AvatarBytes};
 use prost::Message;
 use serde::{de::DeserializeOwned, Deserialize, Serialize};
 use sha2::{Digest, Sha256};
@@ -42,6 +42,7 @@ pub use error::SledStoreError;
 
 const SLED_TREE_CONTACTS: &str = "contacts";
 const SLED_TREE_GROUPS: &str = "groups";
+const SLED_TREE_GROUP_AVATARS: &str = "group_avatars";
 const SLED_TREE_IDENTITIES: &str = "identities";
 const SLED_TREE_PRE_KEYS: &str = "pre_keys";
 const SLED_TREE_SENDER_KEYS: &str = "sender_keys";
@@ -51,6 +52,7 @@ const SLED_TREE_KYBER_PRE_KEYS: &str = "kyber_pre_keys";
 const SLED_TREE_STATE: &str = "state";
 const SLED_TREE_THREADS_PREFIX: &str = "threads";
 const SLED_TREE_PROFILES: &str = "profiles";
+const SLED_TREE_PROFILE_AVATARS: &str = "profile_avatars";
 const SLED_TREE_PROFILE_KEYS: &str = "profile_keys";
 
 const SLED_KEY_NEXT_SIGNED_PRE_KEY_ID: &str = "next_signed_pre_key_id";
@@ -93,11 +95,13 @@ pub enum SchemaVersion {
     V1 = 1,
     V2 = 2,
     V3 = 3,
+    // Introduction of avatars, requires dropping all profiles from the cache
+    V4 = 4,
 }
 
 impl SchemaVersion {
     fn current() -> SchemaVersion {
-        Self::V3
+        Self::V4
     }
 
     /// return an iterator on all the necessary migration steps from another version
@@ -110,6 +114,7 @@ impl SchemaVersion {
             1 => SchemaVersion::V1,
             2 => SchemaVersion::V2,
             3 => SchemaVersion::V3,
+            4 => SchemaVersion::V4,
             _ => unreachable!("oops, this not supposed to happen!"),
         })
     }
@@ -347,6 +352,12 @@ fn migrate(
                     db.drop_tree(SLED_TREE_GROUPS)?;
                     db.flush()?;
                 }
+                SchemaVersion::V4 => {
+                    debug!("migrating from schema v3 to v4: dropping profile cache");
+                    let db = store.write();
+                    db.drop_tree(SLED_TREE_PROFILES)?;
+                    db.flush()?;
+                }
                 _ => return Err(SledStoreError::MigrationConflict),
             }
 
@@ -482,6 +493,22 @@ impl ContentsStore for SledStore {
         Ok(())
     }
 
+    fn group_avatar(
+        &self,
+        master_key_bytes: GroupMasterKeyBytes,
+    ) -> Result<Option<AvatarBytes>, SledStoreError> {
+        self.get(SLED_TREE_GROUP_AVATARS, master_key_bytes)
+    }
+
+    fn save_group_avatar(
+        &self,
+        master_key: GroupMasterKeyBytes,
+        avatar: &AvatarBytes,
+    ) -> Result<(), SledStoreError> {
+        self.insert(SLED_TREE_GROUP_AVATARS, master_key, avatar)?;
+        Ok(())
+    }
+
     /// Messages
 
     fn clear_messages(&mut self) -> Result<(), SledStoreError> {
@@ -604,6 +631,26 @@ impl ContentsStore for SledStore {
         let key = self.profile_key_for_uuid(uuid, key);
         self.get(SLED_TREE_PROFILES, key)
     }
+
+    fn save_profile_avatar(
+        &mut self,
+        uuid: Uuid,
+        key: ProfileKey,
+        avatar: &AvatarBytes,
+    ) -> Result<(), SledStoreError> {
+        let key = self.profile_key_for_uuid(uuid, key);
+        self.insert(SLED_TREE_PROFILE_AVATARS, key, avatar)?;
+        Ok(())
+    }
+
+    fn profile_avatar(
+        &self,
+        uuid: Uuid,
+        key: ProfileKey,
+    ) -> Result<Option<AvatarBytes>, SledStoreError> {
+        let key = self.profile_key_for_uuid(uuid, key);
+        self.get(SLED_TREE_PROFILE_AVATARS, key)
+    }
 }
 
 #[async_trait(?Send)]
@@ -657,6 +704,8 @@ impl Store for SledStore {
         let db = self.write();
         db.drop_tree(SLED_TREE_CONTACTS)?;
         db.drop_tree(SLED_TREE_GROUPS)?;
+        db.drop_tree(SLED_TREE_PROFILES)?;
+        db.drop_tree(SLED_TREE_PROFILE_AVATARS)?;
 
         for tree in db
             .tree_names()

diff --git a/presage/Cargo.toml b/presage/Cargo.toml
@@ -7,8 +7,8 @@ edition = "2021"
 license = "AGPL-3.0-only"
 
 [dependencies]
-libsignal-service = { git = "https://github.com/whisperfish/libsignal-service-rs", rev = "a2e7540a71866a62028ad0205574a5feb0e717ec" }
-libsignal-service-hyper = { git = "https://github.com/whisperfish/libsignal-service-rs", rev = "a2e7540a71866a62028ad0205574a5feb0e717ec" }
+libsignal-service = { git = "https://github.com/whisperfish/libsignal-service-rs", rev = "1a8cdf81" }
+libsignal-service-hyper = { git = "https://github.com/whisperfish/libsignal-service-rs", rev = "1a8cdf81" }
 
 base64 = "0.21"
 futures = "0.3"

diff --git a/presage/src/errors.rs b/presage/src/errors.rs
@@ -70,6 +70,8 @@ pub enum Error<S: std::error::Error> {
     UnexpectedAttachmentChecksum,
     #[error("Unverified registration session (i.e. wrong verification code)")]
     UnverifiedRegistrationSession,
+    #[error("profile cipher error")]
+    ProfileCipherError(#[from] libsignal_service::profile_cipher::ProfileCipherError),
 }
 
 impl<S: StoreError> From<S> for Error<S> {

diff --git a/presage/src/lib.rs b/presage/src/lib.rs
@@ -12,3 +12,5 @@ pub use errors::Error;
 pub use manager::Manager;
 
 const USER_AGENT: &str = concat!(env!("CARGO_PKG_NAME"), "-rs-", env!("CARGO_PKG_VERSION"));
+
+pub type AvatarBytes = Vec<u8>;
diff --git a/presage/src/manager/registered.rs b/presage/src/manager/registered.rs
@@ -34,6 +34,7 @@ use libsignal_service::utils::{
     serde_signaling_key,
 };
 use libsignal_service::websocket::SignalWebSocket;
+use libsignal_service::zkgroup::groups::{GroupMasterKey, GroupSecretParams};
 use libsignal_service::zkgroup::profiles::ProfileKey;
 use libsignal_service::{cipher, AccountManager, Profile, ServiceAddress};
 use libsignal_service_hyper::push_service::HyperPushService;
@@ -47,7 +48,7 @@ use tokio::sync::Mutex;
 use crate::cache::CacheCell;
 use crate::serde::serde_profile_key;
 use crate::store::{Store, Thread};
-use crate::{Error, Manager};
+use crate::{AvatarBytes, Error, Manager};
 
 type ServiceCipher<S> = cipher::ServiceCipher<S, StdRng>;
 type MessageSender<S> = libsignal_service::prelude::MessageSender<HyperPushService, S, StdRng>;
@@ -455,6 +456,8 @@ impl<S: Store> Manager<S, Registered> {
         profile_key: ProfileKey,
     ) -> Result<Profile, Error<S::Error>> {
         // Check if profile is cached.
+        // TODO: Create a migration in the store removing all profiles.
+        // TODO: Is there some way to know if this is outdated?
         if let Some(profile) = self.store.profile(uuid, profile_key).ok().flatten() {
             return Ok(profile);
         }
@@ -468,6 +471,86 @@ impl<S: Store> Manager<S, Registered> {
         Ok(profile)
     }
 
+    pub async fn retrieve_group_avatar(
+        &mut self,
+        context: GroupContextV2,
+    ) -> Result<Option<AvatarBytes>, Error<S::Error>> {
+        let master_key_bytes = context
+            .master_key()
+            .try_into()
+            .expect("Master key bytes to be of size 32.");
+
+        // Check if group avatar is cached.
+        // TODO: Is there some way to know if this is outdated?
+        if let Some(avatar) = self.store.group_avatar(master_key_bytes).ok().flatten() {
+            return Ok(Some(avatar));
+        }
+
+        let mut gm = self.groups_manager()?;
+        let Some(group) = upsert_group(
+            self.store(),
+            &mut gm,
+            context.master_key(),
+            &context.revision(),
+        )
+        .await?
+        else {
+            return Ok(None);
+        };
+
+        // Empty path means no avatar was set.
+        if group.avatar.is_empty() {
+            return Ok(None);
+        }
+
+        let avatar = gm
+            .retrieve_avatar(
+                &group.avatar,
+                GroupSecretParams::derive_from_master_key(GroupMasterKey::new(master_key_bytes)),
+            )
+            .await?;
+        if let Some(avatar) = &avatar {
+            let _ = self.store.save_group_avatar(master_key_bytes, avatar);
+        }
+        Ok(avatar)
+    }
+
+    pub async fn retrieve_profile_avatar_by_uuid(
+        &mut self,
+        uuid: Uuid,
+        profile_key: ProfileKey,
+    ) -> Result<Option<AvatarBytes>, Error<S::Error>> {
+        // Check if profile avatar is cached.
+        // TODO: Is there some way to know if this is outdated?
+        if let Some(avatar) = self.store.profile_avatar(uuid, profile_key).ok().flatten() {
+            return Ok(Some(avatar));
+        }
+
+        let profile = if let Some(profile) = self.store.profile(uuid, profile_key).ok().flatten() {
+            profile
+        } else {
+            self.retrieve_profile_by_uuid(uuid, profile_key).await?
+        };
+
+        let Some(avatar) = profile.avatar.as_ref() else {
+            return Ok(None);
+        };
+
+        let mut service = self.unidentified_push_service();
+
+        let mut avatar_stream = service.retrieve_profile_avatar(avatar).await?;
+        // 10MB is what Signal Android allocates
+        let mut contents = Vec::with_capacity(10 * 1024 * 1024);
+        let len = avatar_stream.read_to_end(&mut contents).await?;
+        contents.truncate(len);
+
+        let cipher = ProfileCipher::from(profile_key);
+
+        let avatar = cipher.decrypt_avatar(&contents)?;
+        let _ = self.store.save_profile_avatar(uuid, profile_key, &avatar);
+        Ok(Some(avatar))
+    }
+
     /// Get an iterator of messages in a thread, optionally starting from a point in time.
     pub fn messages(
         &self,

diff --git a/presage/src/store.rs b/presage/src/store.rs
@@ -20,7 +20,7 @@ use libsignal_service::{
 use log::error;
 use serde::{Deserialize, Serialize};
 
-use crate::manager::RegistrationData;
+use crate::{manager::RegistrationData, AvatarBytes};
 
 /// An error trait implemented by store error types
 pub trait StoreError: std::error::Error + Sync + Send + 'static {}
@@ -214,6 +214,19 @@ pub trait ContentsStore {
         master_key: GroupMasterKeyBytes,
     ) -> Result<Option<Group>, Self::ContentsStoreError>;
 
+    /// Save a group avatar in the cache
+    fn save_group_avatar(
+        &self,
+        master_key: GroupMasterKeyBytes,
+        avatar: &AvatarBytes,
+    ) -> Result<(), Self::ContentsStoreError>;
+
+    /// Retrieve a group avatar from the cache.
+    fn group_avatar(
+        &self,
+        master_key: GroupMasterKeyBytes,
+    ) -> Result<Option<AvatarBytes>, Self::ContentsStoreError>;
+
     // Profiles
 
     /// Insert or update the profile key of a contact
@@ -240,6 +253,21 @@ pub trait ContentsStore {
         uuid: Uuid,
         key: ProfileKey,
     ) -> Result<Option<Profile>, Self::ContentsStoreError>;
+
+    /// Save a profile avatar by [Uuid] and [ProfileKey].
+    fn save_profile_avatar(
+        &mut self,
+        uuid: Uuid,
+        key: ProfileKey,
+        profile: &AvatarBytes,
+    ) -> Result<(), Self::ContentsStoreError>;
+
+    /// Retrieve a profile avatar by [Uuid] and [ProfileKey].
+    fn profile_avatar(
+        &self,
+        uuid: Uuid,
+        key: ProfileKey,
+    ) -> Result<Option<AvatarBytes>, Self::ContentsStoreError>;
 }
 
 /// The manager store trait combining all other stores into a single one