grapevine/src/database/key_value/media.rs

use ruma::{api::client::error::ErrorKind, OwnedMxcUri};

use crate::{
    database::KeyValueDatabase,
    service::{
        self,
        media::{FileMeta, MediaFileKey},
    },
    utils, Error, Result,
};

struct MediaFileKeyParts {
    mxc: OwnedMxcUri,
    width: u32,
    height: u32,
    meta: FileMeta,
}

impl TryFrom<&MediaFileKey> for MediaFileKeyParts {
    type Error = Error;

    #[tracing::instrument(
        err,
        fields(key = utils::u8_slice_to_hex(key.as_bytes())),
    )]
    fn try_from(key: &MediaFileKey) -> Result<MediaFileKeyParts> {
        let mut parts = key.as_bytes().split(|&b| b == 0xFF);

        let mxc_bytes = parts
            .next()
            .ok_or_else(|| Error::BadDatabase("Media ID in db is invalid."))?;
        let mxc = utils::string_from_bytes(mxc_bytes)
            .map_err(|_| {
                Error::BadDatabase("Media MXC URI in db is invalid unicode.")
            })?
            .into();

        let thumbnail_size_bytes = parts
            .next()
            .ok_or_else(|| Error::BadDatabase("Media ID in db is invalid."))?;
        let thumbnail_size_bytes: &[u8; 8] =
            thumbnail_size_bytes.try_into().map_err(|_| {
                Error::BadDatabase("Media ID thumbnail size in db is invalid")
            })?;
        let width = u32::from_be_bytes(
            thumbnail_size_bytes[..4].try_into().expect("should be 4 bytes"),
        );
        let height = u32::from_be_bytes(
            thumbnail_size_bytes[4..].try_into().expect("should be 4 bytes"),
        );

        let content_type = parts
            .next()
            .map(|bytes| {
                utils::string_from_bytes(bytes).map_err(|_| {
                    Error::BadDatabase(
                        "Content type in mediaid_file is invalid unicode.",
                    )
                })
            })
            .transpose()?;

        let content_disposition_bytes = parts
            .next()
            .ok_or_else(|| Error::BadDatabase("Media ID in db is invalid."))?;

        let content_disposition = if content_disposition_bytes.is_empty() {
            None
        } else {
            Some(utils::string_from_bytes(content_disposition_bytes).map_err(
                |_| {
                    Error::BadDatabase(
                        "Content Disposition in mediaid_file is invalid \
                         unicode.",
                    )
                },
            )?)
        };

        Ok(MediaFileKeyParts {
            mxc,
            width,
            height,
            meta: FileMeta {
                content_disposition,
                content_type,
            },
        })
    }
}

impl service::media::Data for KeyValueDatabase {
    fn create_file_metadata(
        &self,
        mxc: OwnedMxcUri,
        width: u32,
        height: u32,
        meta: &FileMeta,
    ) -> Result<MediaFileKey> {
        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(
            meta.content_disposition
                .as_ref()
                .map(String::as_bytes)
                .unwrap_or_default(),
        );
        key.push(0xFF);
        key.extend_from_slice(
            meta.content_type
                .as_ref()
                .map(String::as_bytes)
                .unwrap_or_default(),
        );

        let key = MediaFileKey::new(key);

        self.mediaid_file.insert(key.as_bytes(), &[])?;

        Ok(key)
    }

    fn search_file_metadata(
        &self,
        mxc: OwnedMxcUri,
        width: u32,
        height: u32,
    ) -> Result<(FileMeta, MediaFileKey)> {
        let mut prefix = mxc.as_bytes().to_vec();
        prefix.push(0xFF);
        prefix.extend_from_slice(&width.to_be_bytes());
        prefix.extend_from_slice(&height.to_be_bytes());
        prefix.push(0xFF);

        let (key, _) =
            self.mediaid_file.scan_prefix(prefix).next().ok_or(
                Error::BadRequest(ErrorKind::NotFound, "Media not found"),
            )?;

        let key = MediaFileKey::new(key);
        let parts = MediaFileKeyParts::try_from(&key)?;
        Ok((parts.meta, key))
    }

    fn delete_file_metadata(&self, key: MediaFileKey) -> Result<()> {
        self.mediaid_file.remove(key.as_bytes())
    }

    fn search_thumbnails_metadata(
        &self,
        mxc: OwnedMxcUri,
    ) -> Result<Vec<(FileMeta, MediaFileKey)>> {
        let mut prefix = mxc.as_bytes().to_vec();
        prefix.push(0xFF);

        self.mediaid_file
            .scan_prefix(prefix)
            .map(|(key, _)| {
                let key = MediaFileKey::new(key);
                let parts = MediaFileKeyParts::try_from(&key)?;
                Ok((parts.meta, key))
            })
            .collect()
    }

    fn all_file_metadata(
        &self,
    ) -> Box<
        dyn Iterator<Item = Result<(OwnedMxcUri, FileMeta, MediaFileKey)>> + '_,
    > {
        Box::new(
            self.mediaid_file
                .iter()
                .map(|(key, _)| {
                    let key = MediaFileKey::new(key);

                    let parts = MediaFileKeyParts::try_from(&key)?;
                    if parts.width != 0 && parts.height != 0 {
                        // Skip thumbnails
                        return Ok(None);
                    };

                    Ok(Some((parts.mxc, parts.meta, key)))
                })
                .filter_map(Result::transpose),
        )
    }
}