Add functions to retrieve avatars

presage-store-sled/src/lib.rs

@@ -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";
@@ -482,6 +484,22 @@ impl ContentsStore for SledStore {
         Ok(())
     }
 
+    fn group_avatar(
+        &self,
+        master_key_bytes: GroupMasterKeyBytes,
+    ) -> Result<Option<Vec<u8>>, SledStoreError> {
+        self.get(SLED_TREE_GROUP_AVATARS, master_key_bytes)
+    }
+
+    fn save_group_avatar(
+        &self,
+        master_key: GroupMasterKeyBytes,
+        avatar: Vec<u8>,
+    ) -> Result<(), SledStoreError> {
+        self.insert(SLED_TREE_GROUP_AVATARS, master_key, avatar)?;
+        Ok(())
+    }
+
     /// Messages
 
     fn clear_messages(&mut self) -> Result<(), SledStoreError> {
@@ -604,6 +622,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: Vec<u8>,
+    ) -> 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<Vec<u8>>, SledStoreError> {
+        let key = self.profile_key_for_uuid(uuid, key);
+        self.get(SLED_TREE_PROFILE_AVATARS, key)
+    }
 }
 
 #[async_trait(?Send)]
@@ -657,6 +695,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()

presage/Cargo.toml

@@ -6,8 +6,8 @@ authors = ["Gabriel Féron <[email protected]>"]
 edition = "2021"
 
 [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 = "3b51a6f" }
+libsignal-service-hyper = { git = "https://github.com/whisperfish/libsignal-service-rs", rev = "3b51a6f" }
 
 base64 = "0.21"
 futures = "0.3"

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> {

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;
@@ -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,88 @@ impl<S: Store> Manager<S, Registered> {
         Ok(profile)
     }
 
+    pub async fn retrieve_group_avatar(
+        &mut self,
+        context: GroupContextV2,
+    ) -> Result<Option<Vec<u8>>, 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);
+        };
+
+        let avatar = gm
+            .retrieve_avatar(
+                &group.avatar,
+                GroupSecretParams::derive_from_master_key(GroupMasterKey::new(
+                    master_key_bytes.clone(),
+                )),
+            )
+            .await?;
+        if let Some(avatar) = &avatar {
+            let _ = self
+                .store
+                .save_group_avatar(master_key_bytes, avatar.clone());
+        }
+        Ok(avatar)
+    }
+
+    pub async fn retrieve_profile_avatar_by_uuid(
+        &mut self,
+        uuid: Uuid,
+        profile_key: ProfileKey,
+    ) -> Result<Option<Vec<u8>>, 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);
+        };
+
+        // TODO: Identified or unidentified? Identified seems to work.
+        let mut service = self.identified_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.clone());
+        Ok(Some(avatar))
+    }
+
     /// Get an iterator of messages in a thread, optionally starting from a point in time.
     pub fn messages(
         &self,

presage/src/store.rs

@@ -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,
+        avatat: Vec<u8>,
+    ) -> Result<(), Self::ContentsStoreError>;
+
+    /// Retrieve a group avatar from the cache.
+    fn group_avatar(
+        &self,
+        master_key: GroupMasterKeyBytes,
+    ) -> Result<Option<Vec<u8>>, 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: Vec<u8>,
+    ) -> Result<(), Self::ContentsStoreError>;
+
+    /// Retrieve a profile avatar by [Uuid] and [ProfileKey].
+    fn profile_avatar(
+        &self,
+        uuid: Uuid,
+        key: ProfileKey,
+    ) -> Result<Option<Vec<u8>>, Self::ContentsStoreError>;
 }
 
 /// The manager store trait combining all other stores into a single one