From 5837465121cbaeb3fec455e7ac1f5797159e9522 Mon Sep 17 00:00:00 2001
From: Ruben De Smet <ruben.de.smet@rubdos.be>
Date: Sat, 20 Jun 2020 20:16:35 +0200
Subject: [PATCH] EnvelopeEntityList wrapping type.

---
 libsignal-service/Cargo.toml          |  1 +
 libsignal-service/src/envelope.rs     |  9 +++++++++
 libsignal-service/src/lib.rs          |  2 ++
 libsignal-service/src/push_service.rs |  3 ++-
 libsignal-service/src/utils.rs        | 22 ++++++++++++++++++++++
 5 files changed, 36 insertions(+), 1 deletion(-)
 create mode 100644 libsignal-service/src/utils.rs

diff --git a/libsignal-service/Cargo.toml b/libsignal-service/Cargo.toml
index 33e5dab63..8caacf8f9 100644
--- a/libsignal-service/Cargo.toml
+++ b/libsignal-service/Cargo.toml
@@ -11,6 +11,7 @@ libsignal-protocol = { git = "https://github.com/Michael-F-Bryan/libsignal-proto
 failure = "0.1.5"
 async-trait = "0.1.30"
 url = "2.1.1"
+base64 = "0.12"
 thiserror = "1.0"
 serde = {version = "1.0", features=["derive"]}
 prost = "0.6"
diff --git a/libsignal-service/src/envelope.rs b/libsignal-service/src/envelope.rs
index b3ab3f50e..17280562e 100644
--- a/libsignal-service/src/envelope.rs
+++ b/libsignal-service/src/envelope.rs
@@ -1,3 +1,5 @@
+use crate::utils::serde_base64;
+
 pub struct Envelope {
     inner: crate::proto::Envelope,
 }
@@ -11,12 +13,19 @@ pub struct EnvelopeEntity {
     pub source: String,
     pub source_uuid: String,
     pub source_device: i32,
+    #[serde(with = "serde_base64")]
     pub message: Vec<u8>,
+    #[serde(with = "serde_base64")]
     pub content: Vec<u8>,
     pub server_timestamp: i64,
     pub guid: String,
 }
 
+#[derive(serde::Serialize, serde::Deserialize)]
+pub(crate) struct EnvelopeEntityList {
+    pub messages: Vec<EnvelopeEntity>,
+}
+
 const SUPPORTED_VERSION: usize = 1;
 const CIPHER_KEY_SIZE: usize = 32;
 const MAC_KEY_SIZE: usize = 20;
diff --git a/libsignal-service/src/lib.rs b/libsignal-service/src/lib.rs
index 78914e851..4fb0d5706 100644
--- a/libsignal-service/src/lib.rs
+++ b/libsignal-service/src/lib.rs
@@ -6,6 +6,8 @@ pub mod receiver;
 
 mod proto;
 
+mod utils;
+
 pub use crate::account_manager::AccountManager;
 
 pub const USER_AGENT: &'static str =
diff --git a/libsignal-service/src/push_service.rs b/libsignal-service/src/push_service.rs
index 4545c4a92..4c2140189 100644
--- a/libsignal-service/src/push_service.rs
+++ b/libsignal-service/src/push_service.rs
@@ -102,7 +102,8 @@ pub trait PushService {
     async fn get_messages(
         &mut self,
     ) -> Result<Vec<EnvelopeEntity>, ServiceError> {
-        Ok(self.get(MESSAGE_PATH).await?)
+        let entity_list: EnvelopeEntityList = self.get(MESSAGE_PATH).await?;
+        Ok(entity_list.messages)
     }
 }
 
diff --git a/libsignal-service/src/utils.rs b/libsignal-service/src/utils.rs
new file mode 100644
index 000000000..085bfa8cd
--- /dev/null
+++ b/libsignal-service/src/utils.rs
@@ -0,0 +1,22 @@
+pub mod serde_base64 {
+    use serde::{Deserialize, Deserializer, Serializer};
+
+    pub fn serialize<T, S>(key: &T, serializer: S) -> Result<S::Ok, S::Error>
+    where
+        T: AsRef<[u8]>,
+        S: Serializer,
+    {
+        serializer.serialize_str(&base64::encode(key.as_ref()))
+    }
+
+    pub fn deserialize<'de, D>(deserializer: D) -> Result<Vec<u8>, D::Error>
+    where
+        D: Deserializer<'de>,
+    {
+        use serde::de::Error;
+        String::deserialize(deserializer).and_then(|string| {
+            base64::decode(&string)
+                .map_err(|err| Error::custom(err.to_string()))
+        })
+    }
+}