feat(event cache store): Implement renaming media keys

Author: bnjbvr

crates/matrix-sdk-base/src/event_cache_store/integration_tests.rs

@@ -31,6 +31,9 @@ use crate::media::{MediaFormat, MediaRequest, MediaThumbnailSettings};
 pub trait EventCacheStoreIntegrationTests {
     /// Test media content storage.
     async fn test_media_content(&self);
+
+    /// Test replacing a MXID.
+    async fn test_replace_media_key(&self);
 }
 
 #[cfg_attr(target_arch = "wasm32", async_trait(?Send))]
@@ -139,6 +142,42 @@ impl EventCacheStoreIntegrationTests for DynEventCacheStore {
             "other media was removed"
         );
     }
+
+    async fn test_replace_media_key(&self) {
+        let uri = mxc_uri!("mxc://sendqueue.local/tr4n-s4ct-10n1-d");
+        let req =
+            MediaRequest { source: MediaSource::Plain(uri.to_owned()), format: MediaFormat::File };
+
+        let content = "hello".as_bytes().to_owned();
+
+        // Media isn't present in the cache.
+        assert!(self.get_media_content(&req).await.unwrap().is_none(), "unexpected media found");
+
+        // Add the media.
+        self.add_media_content(&req, content.clone()).await.expect("adding media failed");
+
+        // Sanity-check: media is found after adding it.
+        assert_eq!(self.get_media_content(&req).await.unwrap().unwrap(), b"hello");
+
+        // Replacing a media request works.
+        let new_uri = mxc_uri!("mxc://matrix.org/tr4n-s4ct-10n1-d");
+        let new_req = MediaRequest {
+            source: MediaSource::Plain(new_uri.to_owned()),
+            format: MediaFormat::File,
+        };
+        self.replace_media_key(&req, &new_req)
+            .await
+            .expect("replacing the media request key failed");
+
+        // Finding with the previous request doesn't work anymore.
+        assert!(
+            self.get_media_content(&req).await.unwrap().is_none(),
+            "unexpected media found with the old key"
+        );
+
+        // Finding with the new request does work.
+        assert_eq!(self.get_media_content(&new_req).await.unwrap().unwrap(), b"hello");
+    }
 }
 
 /// Macro building to allow your `EventCacheStore` implementation to run the
@@ -184,6 +223,13 @@ macro_rules! event_cache_store_integration_tests {
                     get_event_cache_store().await.unwrap().into_event_cache_store();
                 event_cache_store.test_media_content().await;
             }
+
+            #[async_test]
+            async fn test_replace_media_key() {
+                let event_cache_store =
+                    get_event_cache_store().await.unwrap().into_event_cache_store();
+                event_cache_store.test_replace_media_key().await;
+            }
         }
     };
 }

crates/matrix-sdk-base/src/event_cache_store/memory_store.rs

@@ -60,18 +60,35 @@ impl EventCacheStore for MemoryStore {
         Ok(())
     }
 
+    async fn replace_media_key(
+        &self,
+        from: &MediaRequest,
+        to: &MediaRequest,
+    ) -> Result<(), Self::Error> {
+        let expected_key = from.unique_key();
+
+        let mut medias = self.media.write().unwrap();
+        if let Some((mxc, key, _)) = medias.iter_mut().find(|(_, key, _)| *key == expected_key) {
+            *mxc = to.uri().to_owned();
+            *key = to.unique_key();
+        }
+
+        Ok(())
+    }
+
     async fn get_media_content(&self, request: &MediaRequest) -> Result<Option<Vec<u8>>> {
-        let media = self.media.read().unwrap();
         let expected_key = request.unique_key();
 
+        let media = self.media.read().unwrap();
         Ok(media.iter().find_map(|(_media_uri, media_key, media_content)| {
             (media_key == &expected_key).then(|| media_content.to_owned())
         }))
     }
 
     async fn remove_media_content(&self, request: &MediaRequest) -> Result<()> {
-        let mut media = self.media.write().unwrap();
         let expected_key = request.unique_key();
+
+        let mut media = self.media.write().unwrap();
         let Some(index) = media
             .iter()
             .position(|(_media_uri, media_key, _media_content)| media_key == &expected_key)

crates/matrix-sdk-base/src/event_cache_store/traits.rs

@@ -42,6 +42,28 @@ pub trait EventCacheStore: AsyncTraitDeps {
         content: Vec<u8>,
     ) -> Result<(), Self::Error>;
 
+    /// Replaces the given media's content key with another one.
+    ///
+    /// This should be used whenever a temporary (local) MXID has been used, and
+    /// it must now be replaced with its actual remote counterpart (after
+    /// uploading some content, or creating an empty MXC URI).
+    ///
+    /// ⚠ No check is performed to ensure that the media formats are consistent,
+    /// i.e. it's possible to update with a thumbnail key a media that was
+    /// keyed as a file before. The caller is responsible of ensuring that
+    /// the replacement makes sense, according to their use case.
+    ///
+    /// # Arguments
+    ///
+    /// * `from` - The previous `MediaRequest` of the file.
+    ///
+    /// * `to` - The new `MediaRequest` of the file.
+    async fn replace_media_key(
+        &self,
+        from: &MediaRequest,
+        to: &MediaRequest,
+    ) -> Result<(), Self::Error>;
+
     /// Get a media file's content out of the media store.
     ///
     /// # Arguments
@@ -91,6 +113,14 @@ impl<T: EventCacheStore> EventCacheStore for EraseEventCacheStoreError<T> {
         self.0.add_media_content(request, content).await.map_err(Into::into)
     }
 
+    async fn replace_media_key(
+        &self,
+        from: &MediaRequest,
+        to: &MediaRequest,
+    ) -> Result<(), Self::Error> {
+        self.0.replace_media_key(from, to).await.map_err(Into::into)
+    }
+
     async fn get_media_content(
         &self,
         request: &MediaRequest,

crates/matrix-sdk-sqlite/src/event_cache_store.rs

@@ -155,6 +155,28 @@ impl EventCacheStore for SqliteEventCacheStore {
         Ok(())
     }
 
+    async fn replace_media_key(
+        &self,
+        from: &MediaRequest,
+        to: &MediaRequest,
+    ) -> Result<(), Self::Error> {
+        let prev_uri = self.encode_key(keys::MEDIA, from.source.unique_key());
+        let prev_format = self.encode_key(keys::MEDIA, from.format.unique_key());
+
+        let new_uri = self.encode_key(keys::MEDIA, to.source.unique_key());
+        let new_format = self.encode_key(keys::MEDIA, to.format.unique_key());
+
+        let conn = self.acquire().await?;
+        conn.execute(
+            r#"UPDATE media SET uri = ?, format = ?, last_access = CAST(strftime('%s') as INT)
+               WHERE uri = ? AND format = ?"#,
+            (new_uri, new_format, prev_uri, prev_format),
+        )
+        .await?;
+
+        Ok(())
+    }
+
     async fn get_media_content(&self, request: &MediaRequest) -> Result<Option<Vec<u8>>> {
         let uri = self.encode_key(keys::MEDIA, request.source.unique_key());
         let format = self.encode_key(keys::MEDIA, request.format.unique_key());