feat: download media and thumbnails over federation

4 years ago · aa5e9e607e
parent 4e44fedbcd
commit aa5e9e607e
2 changed files with 90 additions and 7 deletions
--- a/src/client_server/media.rs
+++ b/src/client_server/media.rs
@ -1,5 +1,7 @@
 use super::State;
-use crate::{database::media::FileMeta, utils, ConduitResult, Database, Error, Ruma};
+use crate::{
+    database::media::FileMeta, server_server, utils, ConduitResult, Database, Error, Ruma,
+};
 use ruma::api::client::{
    error::ErrorKind,
    r0::media::{create_content, get_content, get_content_thumbnail, get_media_config},
@ -35,7 +37,7 @@ pub fn create_content_route(
        utils::random_string(MXC_LENGTH)
    );
    db.media
-        .create(mxc.clone(), &body.filename, &body.content_type, &body.file)?;
+        .create(mxc.clone(), &body.filename.as_deref(), &body.content_type, &body.file)?;

    Ok(create_content::Response { content_uri: mxc }.into())
 }
@ -47,19 +49,25 @@ pub fn create_content_route(
        data = "<body>"
    )
 )]
-pub fn get_content_route(
+pub async fn get_content_route(
    db: State<'_, Database>,
    body: Ruma<get_content::Request<'_>>,
    _server_name: String,
    _media_id: String,
 ) -> ConduitResult<get_content::Response> {
+        let mxc = format!(
+            "mxc://{}/{}",
+            db.globals.server_name(),
+            utils::random_string(MXC_LENGTH)
+        );
+
    if let Some(FileMeta {
        filename,
        content_type,
        file,
    }) = db
        .media
-        .get(format!("mxc://{}/{}", body.server_name, body.media_id))?
+        .get(&mxc)?
    {
        Ok(get_content::Response {
            file,
@ -67,6 +75,26 @@ pub fn get_content_route(
            content_disposition: filename.unwrap_or_default(), // TODO: Spec says this should be optional
        }
        .into())
+    } else if body.allow_remote {
+        let get_content_response = server_server::send_request(
+            &db,
+            body.server_name.as_ref(),
+            get_content::Request {
+                allow_remote: false,
+                server_name: &body.server_name,
+                media_id: &body.media_id,
+            },
+        )
+        .await?;
+
+        db.media.create(
+            mxc,
+            &Some(&get_content_response.content_disposition),
+            &get_content_response.content_type,
+            &get_content_response.file,
+        )?;
+
+        Ok(get_content_response.into())
    } else {
        Err(Error::BadRequest(ErrorKind::NotFound, "Media not found."))
    }
@ -79,7 +107,7 @@ pub fn get_content_route(
        data = "<body>"
    )
 )]
-pub fn get_content_thumbnail_route(
+pub async fn get_content_thumbnail_route(
    db: State<'_, Database>,
    body: Ruma<get_content_thumbnail::Request<'_>>,
    _server_name: String,
@ -97,6 +125,37 @@ pub fn get_content_thumbnail_route(
            .map_err(|_| Error::BadRequest(ErrorKind::InvalidParam, "Width is invalid."))?,
    )? {
        Ok(get_content_thumbnail::Response { file, content_type }.into())
+    } else if body.allow_remote {
+        let get_thumbnail_response = server_server::send_request(
+            &db,
+            body.server_name.as_ref(),
+            get_content_thumbnail::Request {
+                allow_remote: false,
+                height: body.height,
+                width: body.width,
+                method: body.method,
+                server_name: &body.server_name,
+                media_id: &body.media_id,
+            },
+        )
+        .await?;
+
+        let mxc = format!(
+            "mxc://{}/{}",
+            db.globals.server_name(),
+            utils::random_string(MXC_LENGTH)
+        );
+
+        db.media.upload_thumbnail(
+            mxc,
+            &None,
+            &get_thumbnail_response.content_type,
+            body.width.try_into().expect("all UInts are valid u32s"),
+            body.height.try_into().expect("all UInts are valid u32s"),
+            &get_thumbnail_response.file,
+        )?;
+
+        Ok(get_thumbnail_response.into())
    } else {
        Err(Error::BadRequest(ErrorKind::NotFound, "Media not found."))
    }
--- a/src/database/media.rs
+++ b/src/database/media.rs
@ -16,7 +16,7 @@ impl Media {
    pub fn create(
        &self,
        mxc: String,
-        filename: &Option<String>,
+        filename: &Option<&str>,
        content_type: &str,
        file: &[u8],
    ) -> Result<()> {
@ -34,8 +34,32 @@ impl Media {
        Ok(())
    }

+    /// Uploads or replaces a file thumbnail.
+    pub fn upload_thumbnail(
+        &self,
+        mxc: String,
+        filename: &Option<String>,
+        content_type: &str,
+        width: u32,
+        height: u32,
+        file: &[u8],
+    ) -> Result<()> {
+        let mut key = mxc.as_bytes().to_vec();
+        key.push(0xff);
+        key.extend_from_slice(&width.to_be_bytes());
+        key.extend_from_slice(&height.to_be_bytes());
+        key.push(0xff);
+        key.extend_from_slice(filename.as_ref().map(|f| f.as_bytes()).unwrap_or_default());
+        key.push(0xff);
+        key.extend_from_slice(content_type.as_bytes());
+
+        self.mediaid_file.insert(key, file)?;
+
+        Ok(())
+    }
+
    /// Downloads a file.
-    pub fn get(&self, mxc: String) -> Result<Option<FileMeta>> {
+    pub fn get(&self, mxc: &str) -> Result<Option<FileMeta>> {
        let mut prefix = mxc.as_bytes().to_vec();
        prefix.push(0xff);
        prefix.extend_from_slice(&0_u32.to_be_bytes()); // Width = 0 if it's not a thumbnail