From 2c6508198122091dfb350a997e8ad195991eb11c Mon Sep 17 00:00:00 2001
From: Ruben De Smet <ruben.de.smet@rubdos.be>
Date: Fri, 11 Nov 2022 09:41:19 +0100
Subject: [PATCH] Implement profile avatar retention

---
 libsignal-service/src/account_manager.rs | 20 +++++++++++++------
 libsignal-service/src/push_service.rs    | 25 +++++++++++++++++-------
 2 files changed, 32 insertions(+), 13 deletions(-)

diff --git a/libsignal-service/src/account_manager.rs b/libsignal-service/src/account_manager.rs
index 784c9e4ba..bf68d0355 100644
--- a/libsignal-service/src/account_manager.rs
+++ b/libsignal-service/src/account_manager.rs
@@ -14,6 +14,7 @@ use serde::{Deserialize, Serialize};
 use sha2::Sha256;
 use zkgroup::profiles::ProfileKey;
 
+use crate::push_service::AvatarWrite;
 use crate::ServiceAddress;
 use crate::{
     configuration::{Endpoint, ServiceCredentials},
@@ -292,21 +293,28 @@ impl<Service: PushService> AccountManager<Service> {
     ///
     /// Convenience method for
     /// ```ignore
-    /// manager.upload_versioned_profile::<std::io::Cursor<Vec<u8>>, _>(uuid, name, about, about_emoji, None)
+    /// manager.upload_versioned_profile::<std::io::Cursor<Vec<u8>>, _>(uuid, name, about, about_emoji, _)
     /// ```
+    /// in which the `retain_avatar` parameter sets whether to remove (`false`) or retain (`true`) the
+    /// currently set avatar.
     pub async fn upload_versioned_profile_without_avatar<S: AsRef<str>>(
         &mut self,
         uuid: uuid::Uuid,
         name: ProfileName<S>,
         about: Option<String>,
         about_emoji: Option<String>,
+        retain_avatar: bool,
     ) -> Result<(), ProfileManagerError> {
         self.upload_versioned_profile::<std::io::Cursor<Vec<u8>>, _>(
             uuid,
             name,
             about,
             about_emoji,
-            None,
+            if retain_avatar {
+                AvatarWrite::RetainAvatar
+            } else {
+                AvatarWrite::NoAvatar
+            },
         )
         .await?;
         Ok(())
@@ -344,7 +352,7 @@ impl<Service: PushService> AccountManager<Service> {
         name: ProfileName<S>,
         about: Option<String>,
         about_emoji: Option<String>,
-        avatar: Option<&'s mut C>,
+        avatar: AvatarWrite<&'s mut C>,
     ) -> Result<Option<String>, ProfileManagerError> {
         let profile_key =
             self.profile_key.expect("set profile key in AccountManager");
@@ -359,7 +367,7 @@ impl<Service: PushService> AccountManager<Service> {
         let about_emoji = profile_cipher.encrypt_emoji(about_emoji)?;
 
         // If avatar -> upload
-        if let Some(_avatar) = avatar {
+        if matches!(avatar, AvatarWrite::NewAvatar(_)) {
             // FIXME ProfileCipherOutputStream.java
             // It's just AES GCM, but a bit of work to decently implement it with a stream.
             unimplemented!("Setting avatar requires ProfileCipherStream")
@@ -372,13 +380,13 @@ impl<Service: PushService> AccountManager<Service> {
 
         Ok(self
             .service
-            .write_profile(
+            .write_profile::<C, S>(
                 &profile_key_version,
                 &name,
                 &about,
                 &about_emoji,
                 &commitment,
-                None, // FIXME avatar
+                avatar,
             )
             .await?)
     }
diff --git a/libsignal-service/src/push_service.rs b/libsignal-service/src/push_service.rs
index 1c19f9718..ec57af14f 100644
--- a/libsignal-service/src/push_service.rs
+++ b/libsignal-service/src/push_service.rs
@@ -139,6 +139,13 @@ pub enum HttpAuthOverride {
     Identified(HttpAuth),
 }
 
+#[derive(Debug, Clone, Eq, PartialEq)]
+pub enum AvatarWrite<C> {
+    NewAvatar(C),
+    RetainAvatar,
+    NoAvatar,
+}
+
 impl fmt::Debug for HttpAuth {
     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
         write!(f, "HTTP auth with username {}", self.username)
@@ -766,15 +773,14 @@ pub trait PushService: MaybeSend {
     /// See [`AccountManager`][struct@crate::AccountManager] for a convenience method.
     ///
     /// Java equivalent: `writeProfile`
-    async fn write_profile(
+    async fn write_profile<'s, C: std::io::Read + Send + 's, S: AsRef<str>>(
         &mut self,
         version: &ProfileKeyVersion,
         name: &[u8],
         about: &[u8],
         emoji: &[u8],
         commitment: &ProfileKeyCommitment,
-        // FIXME cfr also account manager
-        avatar: Option<()>,
+        avatar: AvatarWrite<&mut C>,
     ) -> Result<Option<String>, ServiceError> {
         #[derive(Debug, Serialize)]
         #[serde(rename_all = "camelCase")]
@@ -788,6 +794,7 @@ pub trait PushService: MaybeSend {
             #[serde(with = "serde_base64")]
             about_emoji: &'s [u8],
             avatar: bool,
+            same_avatar: bool,
             #[serde(with = "serde_base64")]
             commitment: &'s [u8],
         }
@@ -803,7 +810,8 @@ pub trait PushService: MaybeSend {
             name,
             about,
             about_emoji: emoji,
-            avatar: avatar.is_some(),
+            avatar: matches!(avatar, AvatarWrite::NewAvatar(_)),
+            same_avatar: matches!(avatar, AvatarWrite::RetainAvatar),
             commitment: &commitment,
         };
 
@@ -817,19 +825,22 @@ pub trait PushService: MaybeSend {
             )
             .await;
         match (response, avatar) {
-            (Ok(_url), Some(_avatar)) => {
+            (Ok(_url), AvatarWrite::NewAvatar(_avatar)) => {
                 // FIXME
                 unreachable!("Uploading avatar unimplemented");
             },
             // FIXME cleanup when #54883 is stable and MSRV:
             // or-patterns syntax is experimental
             // see issue #54883 <https://github.com/rust-lang/rust/issues/54883> for more information
-            (Err(ServiceError::JsonDecodeError { .. }), None) => {
+            (
+                Err(ServiceError::JsonDecodeError { .. }),
+                AvatarWrite::RetainAvatar | AvatarWrite::NoAvatar,
+            ) => {
                 // OWS sends an empty string when there's no attachment
                 Ok(None)
             },
             (Err(e), _) => Err(e),
-            (Ok(_resp), None) => {
+            (Ok(_resp), AvatarWrite::RetainAvatar | AvatarWrite::NoAvatar) => {
                 log::warn!(
                     "No avatar supplied but got avatar upload URL. Ignoring"
                 );