feat: image thumbnails

src/database/media.rs

@@ -1,7 +1,8 @@
 use crate::{utils, Error, Result};
+use std::mem;
 
 pub struct Media {
-    pub(super) mediaid_file: sled::Tree, // MediaId = MXC + Filename + ContentType
+    pub(super) mediaid_file: sled::Tree, // MediaId = MXC + WidthHeight + Filename + ContentType
 }
 
 impl Media {
@@ -15,6 +16,9 @@ impl Media {
     ) -> Result<()> {
         let mut key = mxc.as_bytes().to_vec();
         key.push(0xff);
+        key.extend_from_slice(&0_u32.to_be_bytes()); // Width = 0 if it's not a thumbnail
+        key.extend_from_slice(&0_u32.to_be_bytes()); // Height = 0 if it's not a thumbnail
+        key.push(0xff);
         key.extend_from_slice(filename.map(|f| f.as_bytes()).unwrap_or_default());
         key.push(0xff);
         key.extend_from_slice(content_type.as_bytes());
@@ -28,10 +32,19 @@ impl Media {
     pub fn get(&self, mxc: String) -> Result<Option<(Option<String>, String, Vec<u8>)>> {
         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
+        prefix.extend_from_slice(&0_u32.to_be_bytes()); // Height = 0 if it's not a thumbnail
+        prefix.push(0xff);
 
         if let Some(r) = self.mediaid_file.scan_prefix(&prefix).next() {
             let (key, file) = r?;
-            let mut parts = key.split(|&b| b == 0xff).skip(1);
+            let mut parts = key.rsplit(|&b| b == 0xff);
+
+            let content_type = utils::string_from_bytes(
+                parts
+                    .next()
+                    .ok_or(Error::BadDatabase("mediaid is invalid"))?,
+            )?;
 
             let filename_bytes = parts
                 .next()
@@ -42,13 +55,99 @@ impl Media {
                 Some(utils::string_from_bytes(filename_bytes)?)
             };
 
+            Ok(Some((filename, content_type, file.to_vec())))
+        } else {
+            Ok(None)
+        }
+    }
+
+    /// Downloads a file's thumbnail.
+    pub fn get_thumbnail(
+        &self,
+        mxc: String,
+        width: u32,
+        height: u32,
+    ) -> Result<Option<(Option<String>, String, Vec<u8>)>> {
+        let mut main_prefix = mxc.as_bytes().to_vec();
+        main_prefix.push(0xff);
+
+        let mut thumbnail_prefix = main_prefix.clone();
+        thumbnail_prefix.extend_from_slice(&width.to_be_bytes());
+        thumbnail_prefix.extend_from_slice(&height.to_be_bytes());
+        thumbnail_prefix.push(0xff);
+
+        let mut original_prefix = main_prefix;
+        original_prefix.extend_from_slice(&0_u32.to_be_bytes()); // Width = 0 if it's not a thumbnail
+        original_prefix.extend_from_slice(&0_u32.to_be_bytes()); // Height = 0 if it's not a thumbnail
+        original_prefix.push(0xff);
+
+        if let Some(r) = self.mediaid_file.scan_prefix(&thumbnail_prefix).next() {
+            // Using saved thumbnail
+            let (key, file) = r?;
+            let mut parts = key.rsplit(|&b| b == 0xff);
+
             let content_type = utils::string_from_bytes(
                 parts
                     .next()
                     .ok_or(Error::BadDatabase("mediaid is invalid"))?,
             )?;
 
+            let filename_bytes = parts
+                .next()
+                .ok_or(Error::BadDatabase("mediaid is invalid"))?;
+            let filename = if filename_bytes.is_empty() {
+                None
+            } else {
+                Some(utils::string_from_bytes(filename_bytes)?)
+            };
+
             Ok(Some((filename, content_type, file.to_vec())))
+        } else if let Some(r) = self.mediaid_file.scan_prefix(&original_prefix).next() {
+            // Generate a thumbnail
+            let (key, file) = r?;
+            let mut parts = key.rsplit(|&b| b == 0xff);
+
+            let content_type = utils::string_from_bytes(
+                parts
+                    .next()
+                    .ok_or(Error::BadDatabase("mediaid is invalid"))?,
+            )?;
+
+            let filename_bytes = parts
+                .next()
+                .ok_or(Error::BadDatabase("mediaid is invalid"))?;
+            let filename = if filename_bytes.is_empty() {
+                None
+            } else {
+                Some(utils::string_from_bytes(filename_bytes)?)
+            };
+
+            if let Ok(image) = image::load_from_memory(&file) {
+                let thumbnail = image.thumbnail(width, height);
+                let mut thumbnail_bytes = Vec::new();
+                thumbnail.write_to(&mut thumbnail_bytes, image::ImageOutputFormat::Jpeg(75))?;
+
+                // Save thumbnail in database so we don't have to generate it again next time
+                let mut thumbnail_key = key.to_vec();
+                let width_index = thumbnail_key
+                    .iter()
+                    .position(|&b| b == 0xff)
+                    .ok_or(Error::BadDatabase("mediaid is invalid"))?
+                    + 1;
+                let mut widthheight = width.to_be_bytes().to_vec();
+                widthheight.extend_from_slice(&height.to_be_bytes());
+
+                thumbnail_key.splice(
+                    width_index..width_index + 2 * mem::size_of::<u32>(),
+                    widthheight,
+                );
+
+                self.mediaid_file.insert(thumbnail_key, &*thumbnail_bytes)?;
+
+                Ok(Some((filename, content_type, thumbnail_bytes)))
+            } else {
+                Ok(None)
+            }
         } else {
             Ok(None)
         }