media: add MediaFileKey wrapper

One more win in the fight against the Vec<u8>s

src/database.rs

@@ -25,9 +25,10 @@ use ruma::{
 use tracing::{debug, error, info, info_span, warn, Instrument};
 
 use crate::{
-    config::DatabaseBackend, observability::FilterReloadHandles,
+    config::DatabaseBackend,
-    service::rooms::timeline::PduCount, services, utils, Config, Error,
+    observability::FilterReloadHandles,
-    PduEvent, Result, Services, SERVICES,
+    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)?;

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<Vec<u8>> {
+    ) -> Result<MediaFileKey> {
         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<String>, Option<String>, Vec<u8>)> {
+    ) -> Result<(Option<String>, Option<String>, 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()

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
     }
 

src/service/media.rs

@@ -24,6 +24,19 @@ pub(crate) struct FileMeta {
     pub(crate) content_type: Option<String>,
 }
 
+#[derive(Clone, Debug, PartialEq, Eq, Hash)]
+pub(crate) struct MediaFileKey(Vec<u8>);
+
+impl MediaFileKey {
+    pub(crate) fn new(key: Vec<u8>) -> Self {
+        Self(key)
+    }
+
+    pub(crate) fn as_bytes(&self) -> &[u8] {
+        &self.0
+    }
+}
+
 pub(crate) struct Service {
     pub(crate) db: &'static dyn Data,
 }

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<Vec<u8>>;
+    ) -> Result<MediaFileKey>;
 
     /// 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<String>, Option<String>, Vec<u8>)>;
+    ) -> Result<(Option<String>, Option<String>, MediaFileKey)>;
 }