grapevine/src/database/key_value/media.rs

use ruma::OwnedMxcUri;

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

#[derive(Debug, Eq, PartialEq)]
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> {
        // Extract parts

        // Extracting mxc url and thumbnail size is a bit fiddly, because the
        // thumbnail size may contain 0xFF bytes.
        let mut parts = key.as_bytes().splitn(2, |&b| b == 0xFF);

        let mxc_bytes = parts
            .next()
            .ok_or_else(|| Error::BadDatabase("Missing MXC URI bytes"))?;

        let rest = parts.next().ok_or_else(|| {
            Error::BadDatabase("Missing thumbnail size bytes")
        })?;

        // Thumbnail size is always 8 bytes
        let (thumbnail_size_bytes, rest) =
            rest.split_at_checked(8).ok_or_else(|| {
                Error::BadDatabase("Missing thumbnail size bytes")
            })?;

        // And always followed immediately by a 0xFF separator
        let mut parts = rest.split(|&b| b == 0xFF);

        let thumbnail_size_rest = parts.next().ok_or_else(|| {
            Error::BadDatabase("Missing Content-Disposition bytes")
        })?;
        if !thumbnail_size_rest.is_empty() {
            return Err(Error::BadDatabase("Thumbnail size part is too long"));
        }

        // The remaining parts are straightforward 0xFF-delimited fields
        let content_disposition_bytes = parts.next().ok_or_else(|| {
            Error::BadDatabase("Missing Content-Disposition bytes")
        })?;

        let content_type_bytes = parts
            .next()
            .ok_or_else(|| Error::BadDatabase("Missing Content-Type bytes"))?;

        if parts.next().is_some() {
            return Err(Error::BadDatabase("Too many parts"));
        }

        // Parse parts

        let mxc = utils::string_from_bytes(mxc_bytes)
            .map_err(|_| Error::BadDatabase("Invalid unicode in MXC URI"))?
            .into();

        let (width, height) = <&[u8; 8]>::try_from(thumbnail_size_bytes)
            .map(|eight_bytes| {
                let width = u32::from_be_bytes(
                    eight_bytes[..4].try_into().expect("should be 4 bytes"),
                );
                let height = u32::from_be_bytes(
                    eight_bytes[4..].try_into().expect("should be 4 bytes"),
                );

                (width, height)
            })
            .map_err(|_| {
                Error::BadDatabase("Wrong number of thumbnail size bytes")
            })?;

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

        let content_type = if content_type_bytes.is_empty() {
            None
        } else {
            Some(utils::string_from_bytes(content_type_bytes).map_err(
                |_| Error::BadDatabase("Invalid unicode in Content-Type"),
            )?)
        };

        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<Option<(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 Some((key, _)) = self.mediaid_file.scan_prefix(prefix).next()
        else {
            return Ok(None);
        };

        let key = MediaFileKey::new(key);
        let parts = MediaFileKeyParts::try_from(&key)?;
        Ok(Some((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),
        )
    }
}

#[cfg(test)]
mod test {
    use super::{FileMeta, MediaFileKey, MediaFileKeyParts};

    #[test]
    fn parse_key_basic() {
        let mut key = b"mxc://example.com/someid".to_vec();
        key.push(0xFF);
        key.extend_from_slice(&640_u32.to_be_bytes());
        key.extend_from_slice(&480_u32.to_be_bytes());
        key.push(0xFF);
        key.extend_from_slice(b"inline");
        key.push(0xFF);
        key.extend_from_slice(b"image/png");
        let key = MediaFileKey::new(key);

        assert_eq!(
            MediaFileKeyParts::try_from(&key).unwrap(),
            MediaFileKeyParts {
                mxc: "mxc://example.com/someid".into(),
                width: 640,
                height: 480,
                meta: FileMeta {
                    content_disposition: Some("inline".to_owned()),
                    content_type: Some("image/png".to_owned()),
                }
            }
        );
    }

    #[test]
    fn parse_key_no_content_type() {
        let mut key = b"mxc://example.com/someid".to_vec();
        key.push(0xFF);
        key.extend_from_slice(&640_u32.to_be_bytes());
        key.extend_from_slice(&480_u32.to_be_bytes());
        key.push(0xFF);
        key.extend_from_slice(b"inline");
        key.push(0xFF);
        // No content type
        let key = MediaFileKey::new(key);

        assert_eq!(
            MediaFileKeyParts::try_from(&key).unwrap(),
            MediaFileKeyParts {
                mxc: "mxc://example.com/someid".into(),
                width: 640,
                height: 480,
                meta: FileMeta {
                    content_disposition: Some("inline".to_owned()),
                    content_type: None,
                }
            }
        );
    }

    #[test]
    fn parse_key_no_content_disposition() {
        let mut key = b"mxc://example.com/someid".to_vec();
        key.push(0xFF);
        key.extend_from_slice(&640_u32.to_be_bytes());
        key.extend_from_slice(&480_u32.to_be_bytes());
        key.push(0xFF);
        // No content disposition
        key.push(0xFF);
        key.extend_from_slice(b"image/png");
        let key = MediaFileKey::new(key);

        assert_eq!(
            MediaFileKeyParts::try_from(&key).unwrap(),
            MediaFileKeyParts {
                mxc: "mxc://example.com/someid".into(),
                width: 640,
                height: 480,
                meta: FileMeta {
                    content_disposition: None,
                    content_type: Some("image/png".to_owned()),
                }
            }
        );
    }

    #[test]
    fn parse_key_no_content_disposition_or_type() {
        let mut key = b"mxc://example.com/someid".to_vec();
        key.push(0xFF);
        key.extend_from_slice(&640_u32.to_be_bytes());
        key.extend_from_slice(&480_u32.to_be_bytes());
        key.push(0xFF);
        // No content disposition
        key.push(0xFF);
        // No content type
        let key = MediaFileKey::new(key);

        assert_eq!(
            MediaFileKeyParts::try_from(&key).unwrap(),
            MediaFileKeyParts {
                mxc: "mxc://example.com/someid".into(),
                width: 640,
                height: 480,
                meta: FileMeta {
                    content_disposition: None,
                    content_type: None,
                }
            }
        );
    }

    // Our current media service code has an allowlist of thumbnail dimensions,
    // and so we don't expect to create new thumbnails with dimensions
    // containing a 0xFF byte. Thumbnails with a 0xFF in the dimensions may
    // have been created previously, so we need to be able to read them.
    #[test]
    fn parse_key_separator_in_thumbnail_dims() {
        let mut key = b"mxc://example.com/someid".to_vec();
        key.push(0xFF);
        key.extend_from_slice(&[0x0, 0x0, 0xFF, 0xFF]);
        key.extend_from_slice(&[0x0, 0x0, 0x10, 0xFF]);
        key.push(0xFF);
        key.extend_from_slice(b"inline");
        key.push(0xFF);
        key.extend_from_slice(b"image/png");
        let key = MediaFileKey::new(key);

        assert_eq!(
            MediaFileKeyParts::try_from(&key).unwrap(),
            MediaFileKeyParts {
                mxc: "mxc://example.com/someid".into(),
                width: 0xFFFF,
                height: 0x10FF,
                meta: FileMeta {
                    content_disposition: Some("inline".to_owned()),
                    content_type: Some("image/png".to_owned()),
                }
            }
        );
    }
}