diff --git a/src/database.rs b/src/database.rs index 13f989a2..445c8a6a 100644 --- a/src/database.rs +++ b/src/database.rs @@ -25,9 +25,10 @@ use ruma::{ use tracing::{debug, error, info, info_span, warn, Instrument}; use crate::{ - config::DatabaseBackend, observability::FilterReloadHandles, - service::rooms::timeline::PduCount, services, utils, Config, Error, - PduEvent, Result, Services, SERVICES, + config::DatabaseBackend, + observability::FilterReloadHandles, + service::{media::MediaFileKey, rooms::timeline::PduCount}, + services, utils, Config, Error, PduEvent, Result, Services, SERVICES, }; pub(crate) struct KeyValueDatabase { @@ -606,6 +607,7 @@ impl KeyValueDatabase { if services().globals.database_version()? < 3 { // Move media to filesystem for (key, content) in db.mediaid_file.iter() { + let key = MediaFileKey::new(key); if content.is_empty() { continue; } @@ -613,7 +615,7 @@ impl KeyValueDatabase { let path = services().globals.get_media_file(&key); let mut file = fs::File::create(path)?; file.write_all(&content)?; - db.mediaid_file.insert(&key, &[])?; + db.mediaid_file.insert(key.as_bytes(), &[])?; } services().globals.bump_database_version(3)?; diff --git a/src/database/key_value/media.rs b/src/database/key_value/media.rs index 57154a51..868842ba 100644 --- a/src/database/key_value/media.rs +++ b/src/database/key_value/media.rs @@ -1,6 +1,10 @@ use ruma::api::client::error::ErrorKind; -use crate::{database::KeyValueDatabase, service, utils, Error, Result}; +use crate::{ + database::KeyValueDatabase, + service::{self, media::MediaFileKey}, + utils, Error, Result, +}; impl service::media::Data for KeyValueDatabase { fn create_file_metadata( @@ -10,7 +14,7 @@ impl service::media::Data for KeyValueDatabase { height: u32, content_disposition: Option<&str>, content_type: Option<&str>, - ) -> Result> { + ) -> Result { let mut key = mxc.as_bytes().to_vec(); key.push(0xFF); key.extend_from_slice(&width.to_be_bytes()); @@ -27,7 +31,9 @@ impl service::media::Data for KeyValueDatabase { content_type.as_ref().map(|c| c.as_bytes()).unwrap_or_default(), ); - self.mediaid_file.insert(&key, &[])?; + let key = MediaFileKey::new(key); + + self.mediaid_file.insert(key.as_bytes(), &[])?; Ok(key) } @@ -37,7 +43,7 @@ impl service::media::Data for KeyValueDatabase { mxc: String, width: u32, height: u32, - ) -> Result<(Option, Option, Vec)> { + ) -> Result<(Option, Option, MediaFileKey)> { let mut prefix = mxc.as_bytes().to_vec(); prefix.push(0xFF); prefix.extend_from_slice(&width.to_be_bytes()); @@ -49,7 +55,9 @@ impl service::media::Data for KeyValueDatabase { Error::BadRequest(ErrorKind::NotFound, "Media not found"), )?; - let mut parts = key.rsplit(|&b| b == 0xFF); + let key = MediaFileKey::new(key); + + let mut parts = key.as_bytes().rsplit(|&b| b == 0xFF); let content_type = parts .next() diff --git a/src/service/globals.rs b/src/service/globals.rs index 2bdad38c..75b545fa 100644 --- a/src/service/globals.rs +++ b/src/service/globals.rs @@ -35,6 +35,7 @@ use trust_dns_resolver::TokioAsyncResolver; use crate::{ api::server_server::FedDest, observability::FilterReloadHandles, + service::media::MediaFileKey, services, utils::on_demand_hashmap::{OnDemandHashMap, TokenSet}, Config, Error, Result, @@ -506,11 +507,11 @@ impl Service { r } - pub(crate) fn get_media_file(&self, key: &[u8]) -> PathBuf { + pub(crate) fn get_media_file(&self, key: &MediaFileKey) -> PathBuf { let mut r = PathBuf::new(); r.push(self.config.database.path.clone()); r.push("media"); - r.push(general_purpose::URL_SAFE_NO_PAD.encode(key)); + r.push(general_purpose::URL_SAFE_NO_PAD.encode(key.as_bytes())); r } diff --git a/src/service/media.rs b/src/service/media.rs index 21b178a2..12cb2538 100644 --- a/src/service/media.rs +++ b/src/service/media.rs @@ -24,6 +24,19 @@ pub(crate) struct FileMeta { pub(crate) content_type: Option, } +#[derive(Clone, Debug, PartialEq, Eq, Hash)] +pub(crate) struct MediaFileKey(Vec); + +impl MediaFileKey { + pub(crate) fn new(key: Vec) -> Self { + Self(key) + } + + pub(crate) fn as_bytes(&self) -> &[u8] { + &self.0 + } +} + pub(crate) struct Service { pub(crate) db: &'static dyn Data, } diff --git a/src/service/media/data.rs b/src/service/media/data.rs index 0aea3f9d..0f337342 100644 --- a/src/service/media/data.rs +++ b/src/service/media/data.rs @@ -1,3 +1,4 @@ +use super::MediaFileKey; use crate::Result; pub(crate) trait Data: Send + Sync { @@ -8,7 +9,7 @@ pub(crate) trait Data: Send + Sync { height: u32, content_disposition: Option<&str>, content_type: Option<&str>, - ) -> Result>; + ) -> Result; /// Returns `content_disposition`, `content_type` and the `metadata` key. fn search_file_metadata( @@ -16,5 +17,5 @@ pub(crate) trait Data: Send + Sync { mxc: String, width: u32, height: u32, - ) -> Result<(Option, Option, Vec)>; + ) -> Result<(Option, Option, MediaFileKey)>; }