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 { 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 { 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> { 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> + '_, > { 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), ) } }