conduit/src/database/media.rs

use crate::{utils, Error, Result};
use std::mem;

pub struct Media {
    pub(super) mediaid_file: sled::Tree, // MediaId = MXC + WidthHeight + Filename + ContentType
}

impl Media {
    /// Uploads or replaces a file.
    pub fn create(
        &self,
        mxc: String,
        filename: Option<&String>,
        content_type: &str,
        file: &[u8],
    ) -> 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());

        self.mediaid_file.insert(key, file)?;

        Ok(())
    }

    /// Downloads a file.
    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.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 {
            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)
        }
    }
}