From b805670c8a7f2a1888c195a29672f6b699972485 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?K=C3=A9vin=20Commaille?= <zecakeh@tedomum.fr>
Date: Tue, 25 May 2021 21:43:01 +0200
Subject: [PATCH] sdk: Add methods for media content

---
 matrix_sdk/src/client.rs | 96 +++++++++++++++++++++++++++++++++++++++-
 matrix_sdk/src/error.rs  |  7 ++-
 matrix_sdk/src/lib.rs    |  2 +-
 3 files changed, 101 insertions(+), 4 deletions(-)

diff --git a/matrix_sdk/src/client.rs b/matrix_sdk/src/client.rs
index e928b18e..6a64a18f 100644
--- a/matrix_sdk/src/client.rs
+++ b/matrix_sdk/src/client.rs
@@ -24,7 +24,7 @@ use std::{
 use std::{
     fmt::{self, Debug},
     future::Future,
-    io::Read,
+    io::{Cursor, Read},
     path::Path,
     result::Result as StdResult,
     sync::Arc,
@@ -41,7 +41,11 @@ use matrix_sdk_base::crypto::{
     OutgoingRequests, RoomMessageRequest, ToDeviceRequest,
 };
 use matrix_sdk_base::{
-    deserialized_responses::SyncResponse, events::AnyMessageEventContent, identifiers::MxcUri,
+    crypto::AttachmentDecryptor,
+    deserialized_responses::SyncResponse,
+    events::AnyMessageEventContent,
+    identifiers::MxcUri,
+    media::{MediaFormat, MediaRequest, MediaType},
     BaseClient, BaseClientConfig, SendAccessToken, Session, Store,
 };
 use mime::{self, Mime};
@@ -2465,6 +2469,94 @@ impl Client {
 
         Ok(olm.import_keys(import, |_, _| {}).await?)
     }
+
+    /// Get a media file's content.
+    ///
+    /// If the content is encrypted and encryption is enabled, the content will
+    /// be decrypted.
+    ///
+    /// # Arguments
+    ///
+    /// * `request` - The `MediaRequest` of the content.
+    ///
+    /// * `use_cache` - If we should use the media cache for this request.
+    pub async fn get_media_content(
+        &self,
+        request: &MediaRequest,
+        use_cache: bool,
+    ) -> Result<Vec<u8>> {
+        let content = if use_cache {
+            self.base_client.store().get_media_content(request).await?
+        } else {
+            None
+        };
+
+        if let Some(content) = content {
+            Ok(content)
+        } else {
+            let content: Vec<u8> = match &request.media_type {
+                MediaType::Encrypted(file) => {
+                    let content: Vec<u8> =
+                        self.send(get_content::Request::from_url(&file.url)?, None).await?.file;
+
+                    #[cfg(feature = "encryption")]
+                    let content = {
+                        let mut cursor = Cursor::new(content);
+                        let mut reader =
+                            AttachmentDecryptor::new(&mut cursor, file.clone().into())?;
+
+                        let mut decrypted = Vec::new();
+                        reader.read_to_end(&mut decrypted)?;
+
+                        decrypted
+                    };
+
+                    content
+                }
+                MediaType::Uri(uri) => {
+                    if let MediaFormat::Thumbnail(size) = &request.format {
+                        self.send(
+                            get_content_thumbnail::Request::from_url(
+                                &uri,
+                                size.width,
+                                size.height,
+                            )?,
+                            None,
+                        )
+                        .await?
+                        .file
+                    } else {
+                        self.send(get_content::Request::from_url(&uri)?, None).await?.file
+                    }
+                }
+            };
+
+            if use_cache {
+                self.base_client.store().add_media_content(request, content.clone()).await?;
+            }
+
+            Ok(content)
+        }
+    }
+
+    /// Remove a media file's content from the store.
+    ///
+    /// # Arguments
+    ///
+    /// * `request` - The `MediaRequest` of the content.
+    pub async fn remove_media_content(&self, request: &MediaRequest) -> Result<()> {
+        Ok(self.base_client.store().remove_media_content(request).await?)
+    }
+
+    /// Delete all the media content corresponding to the given
+    /// uri from the store.
+    ///
+    /// # Arguments
+    ///
+    /// * `uri` - The `MxcUri` of the files.
+    pub async fn remove_media_content_for_uri(&self, uri: &MxcUri) -> Result<()> {
+        Ok(self.base_client.store().remove_media_content_for_uri(&uri).await?)
+    }
 }
 
 #[cfg(test)]
diff --git a/matrix_sdk/src/error.rs b/matrix_sdk/src/error.rs
index e2b032ec..061446c9 100644
--- a/matrix_sdk/src/error.rs
+++ b/matrix_sdk/src/error.rs
@@ -18,7 +18,7 @@ use std::io::Error as IoError;
 
 use http::StatusCode;
 #[cfg(feature = "encryption")]
-use matrix_sdk_base::crypto::store::CryptoStoreError;
+use matrix_sdk_base::crypto::{store::CryptoStoreError, DecryptorError};
 use matrix_sdk_base::{Error as MatrixError, StoreError};
 use matrix_sdk_common::{
     api::{
@@ -122,6 +122,11 @@ pub enum Error {
     #[error(transparent)]
     CryptoStoreError(#[from] CryptoStoreError),
 
+    /// An error occurred during decryption.
+    #[cfg(feature = "encryption")]
+    #[error(transparent)]
+    DecryptorError(#[from] DecryptorError),
+
     /// An error occured in the state store.
     #[error(transparent)]
     StateStore(#[from] StoreError),
diff --git a/matrix_sdk/src/lib.rs b/matrix_sdk/src/lib.rs
index 8b5da4eb..9a105305 100644
--- a/matrix_sdk/src/lib.rs
+++ b/matrix_sdk/src/lib.rs
@@ -81,7 +81,7 @@ pub use bytes::{Bytes, BytesMut};
 #[cfg_attr(feature = "docs", doc(cfg(encryption)))]
 pub use matrix_sdk_base::crypto::{EncryptionInfo, LocalTrust};
 pub use matrix_sdk_base::{
-    Error as BaseError, Room as BaseRoom, RoomInfo, RoomMember as BaseRoomMember, RoomType,
+    media, Error as BaseError, Room as BaseRoom, RoomInfo, RoomMember as BaseRoomMember, RoomType,
     Session, StateChanges, StoreError,
 };
 pub use matrix_sdk_common::*;