Skip to main content

matrix_sdk_sqlite/
media_store.rs

1// Copyright 2024 The Matrix.org Foundation C.I.C.
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15//! An SQLite-based backend for the [`MediaStore`].
16
17use std::{
18    fmt,
19    path::{Path, PathBuf},
20    sync::Arc,
21};
22
23use async_trait::async_trait;
24use deadpool::managed::PoolConfig;
25use matrix_sdk_base::{
26    cross_process_lock::CrossProcessLockGeneration,
27    media::{
28        MediaRequestParameters, UniqueKey,
29        store::{
30            IgnoreMediaRetentionPolicy, MediaRetentionPolicy, MediaService, MediaStore,
31            MediaStoreInner,
32        },
33    },
34    timer,
35};
36use matrix_sdk_store_encryption::StoreCipher;
37use ruma::{MilliSecondsSinceUnixEpoch, MxcUri, time::SystemTime};
38use rusqlite::{OptionalExtension, params_from_iter};
39use tokio::{
40    fs,
41    sync::{Mutex, OwnedMutexGuard},
42};
43use tracing::{debug, instrument};
44
45use crate::{
46    OpenStoreError, RuntimeConfig, Secret, SqliteStoreConfig,
47    connection::{self, Connection as SqliteAsyncConn, Pool as SqlitePool, SqliteConnections},
48    error::{Error, Result},
49    utils::{
50        EncryptableStore, SqliteAsyncConnExt, SqliteKeyValueStoreAsyncConnExt,
51        SqliteKeyValueStoreConnExt, SqliteTransactionExt, time_to_timestamp,
52    },
53};
54
55mod keys {
56    // Entries in Key-value store
57    pub const MEDIA_RETENTION_POLICY: &str = "media_retention_policy";
58    pub const LAST_MEDIA_CLEANUP_TIME: &str = "last_media_cleanup_time";
59
60    // Tables
61    pub const MEDIA: &str = "media";
62}
63
64/// The database name.
65const DATABASE_NAME: &str = "matrix-sdk-media.sqlite3";
66
67/// An SQLite-based media store.
68#[derive(Clone)]
69pub struct SqliteMediaStore {
70    store_cipher: Option<Arc<StoreCipher>>,
71
72    /// `Some` when active, `None` when closed.
73    connections: Arc<Mutex<Option<SqliteConnections>>>,
74
75    /// Retained so we can rebuild the pool on reopen.
76    db_path: PathBuf,
77
78    /// Retained so we can rebuild the pool on reopen.
79    pool_config: PoolConfig,
80
81    /// Retained so we can re-apply runtime config on reopen.
82    runtime_config: RuntimeConfig,
83
84    media_service: MediaService,
85}
86
87#[cfg(not(tarpaulin_include))]
88impl fmt::Debug for SqliteMediaStore {
89    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
90        f.debug_struct("SqliteMediaStore").finish_non_exhaustive()
91    }
92}
93
94impl EncryptableStore for SqliteMediaStore {
95    fn get_cypher(&self) -> Option<&StoreCipher> {
96        self.store_cipher.as_deref()
97    }
98}
99
100impl SqliteMediaStore {
101    /// Open the SQLite-based media store at the given path using the given
102    /// passphrase to encrypt private data.
103    pub async fn open(
104        path: impl AsRef<Path>,
105        passphrase: Option<&str>,
106    ) -> Result<Self, OpenStoreError> {
107        Self::open_with_config(&SqliteStoreConfig::new(path).passphrase(passphrase)).await
108    }
109
110    /// Open the SQLite-based media store at the given path using the given key
111    /// to encrypt private data.
112    pub async fn open_with_key(
113        path: impl AsRef<Path>,
114        key: Option<&[u8]>,
115    ) -> Result<Self, OpenStoreError> {
116        Self::open_with_config(&SqliteStoreConfig::new(path).key(key)).await
117    }
118
119    /// Open the SQLite-based media store with the config open config.
120    #[instrument(skip(config), fields(path = ?config.path))]
121    pub async fn open_with_config(config: &SqliteStoreConfig) -> Result<Self, OpenStoreError> {
122        debug!(?config);
123
124        let _timer = timer!("open_with_config");
125
126        fs::create_dir_all(&config.path).await.map_err(OpenStoreError::CreateDir)?;
127
128        let db_path = config.path.join(DATABASE_NAME);
129        let pool_config = config.pool_config();
130        let runtime_config = config.runtime_config();
131
132        let pool = config.build_pool_of_connections(DATABASE_NAME)?;
133
134        let this =
135            Self::open_with_pool(pool, db_path, pool_config, runtime_config, config.secret.clone())
136                .await?;
137
138        // Apply runtime config on the write connection.
139        this.write().await?.apply_runtime_config(runtime_config).await?;
140
141        Ok(this)
142    }
143
144    /// Open an SQLite-based media store using the given SQLite database pool.
145    /// The given passphrase will be used to encrypt private data.
146    async fn open_with_pool(
147        pool: SqlitePool,
148        db_path: PathBuf,
149        pool_config: PoolConfig,
150        runtime_config: RuntimeConfig,
151        secret: Option<Secret>,
152    ) -> Result<Self, OpenStoreError> {
153        let conn = pool.get().await?;
154
155        let version = conn.db_version().await?;
156        run_migrations(&conn, version).await?;
157
158        conn.wal_checkpoint().await;
159
160        let store_cipher = match &secret {
161            Some(s) => Some(Arc::new(conn.get_or_create_store_cipher(s.clone()).await?)),
162            None => None,
163        };
164
165        let media_service = MediaService::new();
166        let media_retention_policy = conn.get_serialized_kv(keys::MEDIA_RETENTION_POLICY).await?;
167        let last_media_cleanup_time = conn.get_serialized_kv(keys::LAST_MEDIA_CLEANUP_TIME).await?;
168        media_service.restore(media_retention_policy, last_media_cleanup_time);
169
170        let connections = SqliteConnections {
171            pool,
172            // Use `conn` as our selected write connection.
173            write_connection: Arc::new(Mutex::new(conn)),
174        };
175
176        Ok(Self {
177            store_cipher,
178            connections: Arc::new(Mutex::new(Some(connections))),
179            db_path,
180            pool_config,
181            runtime_config,
182            media_service,
183        })
184    }
185
186    // Acquire a connection for executing read operations.
187    #[instrument(skip_all)]
188    async fn read(&self) -> Result<SqliteAsyncConn> {
189        let pool = {
190            let guard = self.connections.lock().await;
191            let conns = guard.as_ref().ok_or(Error::StoreClosed)?;
192            conns.pool.clone()
193        };
194
195        let connection = pool.get().await?;
196
197        // Per https://www.sqlite.org/foreignkeys.html#fk_enable, foreign key
198        // support must be enabled on a per-connection basis. Execute it every
199        // time we try to get a connection, since we can't guarantee a previous
200        // connection did enable it before.
201        connection.execute_batch("PRAGMA foreign_keys = ON;").await?;
202
203        Ok(connection)
204    }
205
206    // Acquire a connection for executing write operations.
207    #[instrument(skip_all)]
208    async fn write(&self) -> Result<OwnedMutexGuard<SqliteAsyncConn>> {
209        let write_connection = {
210            let guard = self.connections.lock().await;
211            let conns = guard.as_ref().ok_or(Error::StoreClosed)?;
212            conns.write_connection.clone()
213        };
214
215        let connection = write_connection.lock_owned().await;
216
217        // Per https://www.sqlite.org/foreignkeys.html#fk_enable, foreign key
218        // support must be enabled on a per-connection basis. Execute it every
219        // time we try to get a connection, since we can't guarantee a previous
220        // connection did enable it before.
221        connection.execute_batch("PRAGMA foreign_keys = ON;").await?;
222
223        Ok(connection)
224    }
225
226    pub async fn vacuum(&self) -> Result<()> {
227        let write_connection = {
228            let guard = self.connections.lock().await;
229            let conns = guard.as_ref().ok_or(Error::StoreClosed)?;
230            conns.write_connection.clone()
231        };
232        write_connection.lock().await.vacuum().await
233    }
234
235    async fn get_db_size(&self) -> Result<Option<usize>> {
236        let pool = {
237            let guard = self.connections.lock().await;
238            let conns = guard.as_ref().ok_or(Error::StoreClosed)?;
239            conns.pool.clone()
240        };
241        Ok(Some(pool.get().await?.get_db_size().await?))
242    }
243
244    pub async fn close(&self) -> Result<()> {
245        connection::close_connections(&self.connections, "Media store").await;
246        Ok(())
247    }
248
249    pub async fn reopen(&self) -> Result<()> {
250        connection::reopen_connections(
251            &self.connections,
252            self.db_path.clone(),
253            self.pool_config,
254            self.runtime_config,
255        )
256        .await?;
257        Ok(())
258    }
259
260    /// Returns the pool size status, for testing purposes.
261    #[cfg(test)]
262    async fn pool_max_size(&self) -> Option<usize> {
263        let guard = self.connections.lock().await;
264        guard.as_ref().map(|conns| conns.pool.status().max_size)
265    }
266}
267
268/// Run migrations for the given version of the database.
269async fn run_migrations(conn: &SqliteAsyncConn, version: u8) -> Result<()> {
270    // Always enable foreign keys for the current connection.
271    conn.execute_batch("PRAGMA foreign_keys = ON;").await?;
272
273    if version < 1 {
274        debug!("Creating database");
275        // First turn on WAL mode, this can't be done in the transaction, it
276        // fails with the error message: "cannot change into wal mode from
277        // within a transaction".
278        conn.execute_batch("PRAGMA journal_mode = wal;").await?;
279        conn.with_transaction(|txn| {
280            txn.execute_batch(include_str!("../migrations/media_store/001_init.sql"))?;
281            txn.set_db_version(1)
282        })
283        .await?;
284    }
285
286    if version < 2 {
287        debug!("Upgrading database to version 2");
288        conn.with_transaction(|txn| {
289            txn.execute_batch(include_str!(
290                "../migrations/media_store/002_lease_locks_with_generation.sql"
291            ))?;
292            txn.set_db_version(2)
293        })
294        .await?;
295    }
296
297    Ok(())
298}
299
300#[async_trait]
301impl MediaStore for SqliteMediaStore {
302    type Error = Error;
303
304    #[instrument(skip(self))]
305    async fn try_take_leased_lock(
306        &self,
307        lease_duration_ms: u32,
308        key: &str,
309        holder: &str,
310    ) -> Result<Option<CrossProcessLockGeneration>> {
311        let key = key.to_owned();
312        let holder = holder.to_owned();
313
314        let now: u64 = MilliSecondsSinceUnixEpoch::now().get().into();
315        let expiration = now + lease_duration_ms as u64;
316
317        // Learn about the `excluded` keyword in https://sqlite.org/lang_upsert.html.
318        let generation = self
319            .write()
320            .await?
321            .with_transaction(move |txn| {
322                txn.query_one(
323                    "INSERT INTO lease_locks (key, holder, expiration)
324                    VALUES (?1, ?2, ?3)
325                    ON CONFLICT (key)
326                    DO
327                        UPDATE SET
328                            holder = excluded.holder,
329                            expiration = excluded.expiration,
330                            generation =
331                                CASE holder
332                                    WHEN excluded.holder THEN generation
333                                    ELSE generation + 1
334                                END
335                        WHERE
336                            holder = excluded.holder
337                            OR expiration < ?4
338                    RETURNING generation
339                    ",
340                    (key, holder, expiration, now),
341                    |row| row.get(0),
342                )
343                .optional()
344            })
345            .await?;
346
347        Ok(generation)
348    }
349
350    async fn add_media_content(
351        &self,
352        request: &MediaRequestParameters,
353        content: Vec<u8>,
354        ignore_policy: IgnoreMediaRetentionPolicy,
355    ) -> Result<()> {
356        let _timer = timer!("method");
357
358        self.media_service.add_media_content(self, request, content, ignore_policy).await
359    }
360
361    #[instrument(skip_all)]
362    async fn replace_media_key(
363        &self,
364        from: &MediaRequestParameters,
365        to: &MediaRequestParameters,
366    ) -> Result<(), Self::Error> {
367        let _timer = timer!("method");
368
369        let prev_uri = self.encode_key(keys::MEDIA, from.source.unique_key());
370        let prev_format = self.encode_key(keys::MEDIA, from.format.unique_key());
371
372        let new_uri = self.encode_key(keys::MEDIA, to.source.unique_key());
373        let new_format = self.encode_key(keys::MEDIA, to.format.unique_key());
374
375        let conn = self.write().await?;
376        conn.execute(
377            r#"UPDATE media SET uri = ?, format = ? WHERE uri = ? AND format = ?"#,
378            (new_uri, new_format, prev_uri, prev_format),
379        )
380        .await?;
381
382        Ok(())
383    }
384
385    #[instrument(skip_all)]
386    async fn get_media_content(&self, request: &MediaRequestParameters) -> Result<Option<Vec<u8>>> {
387        let _timer = timer!("method");
388
389        self.media_service.get_media_content(self, request).await
390    }
391
392    #[instrument(skip_all)]
393    async fn remove_media_content(&self, request: &MediaRequestParameters) -> Result<()> {
394        let _timer = timer!("method");
395
396        let uri = self.encode_key(keys::MEDIA, request.source.unique_key());
397        let format = self.encode_key(keys::MEDIA, request.format.unique_key());
398
399        let conn = self.write().await?;
400        conn.execute("DELETE FROM media WHERE uri = ? AND format = ?", (uri, format)).await?;
401
402        Ok(())
403    }
404
405    #[instrument(skip(self))]
406    async fn remove_media_content_for_uri(&self, uri: &MxcUri) -> Result<()> {
407        let _timer = timer!("method");
408
409        let uri = self.encode_key(keys::MEDIA, uri);
410
411        let conn = self.write().await?;
412        conn.execute("DELETE FROM media WHERE uri = ?", (uri,)).await?;
413
414        Ok(())
415    }
416
417    #[instrument(skip_all)]
418    async fn set_media_retention_policy(
419        &self,
420        policy: MediaRetentionPolicy,
421    ) -> Result<(), Self::Error> {
422        let _timer = timer!("method");
423
424        self.media_service.set_media_retention_policy(self, policy).await
425    }
426
427    #[instrument(skip_all)]
428    fn media_retention_policy(&self) -> MediaRetentionPolicy {
429        let _timer = timer!("method");
430
431        self.media_service.media_retention_policy()
432    }
433
434    #[instrument(skip_all)]
435    async fn set_ignore_media_retention_policy(
436        &self,
437        request: &MediaRequestParameters,
438        ignore_policy: IgnoreMediaRetentionPolicy,
439    ) -> Result<(), Self::Error> {
440        let _timer = timer!("method");
441
442        self.media_service.set_ignore_media_retention_policy(self, request, ignore_policy).await
443    }
444
445    #[instrument(skip_all)]
446    async fn clean(&self) -> Result<(), Self::Error> {
447        let _timer = timer!("method");
448
449        self.media_service.clean(self).await
450    }
451
452    async fn close(&self) -> Result<(), Self::Error> {
453        SqliteMediaStore::close(self).await
454    }
455
456    async fn reopen(&self) -> Result<(), Self::Error> {
457        SqliteMediaStore::reopen(self).await
458    }
459
460    async fn optimize(&self) -> Result<(), Self::Error> {
461        Ok(self.vacuum().await?)
462    }
463
464    async fn get_size(&self) -> Result<Option<usize>, Self::Error> {
465        self.get_db_size().await
466    }
467}
468
469#[cfg_attr(target_family = "wasm", async_trait(?Send))]
470#[cfg_attr(not(target_family = "wasm"), async_trait)]
471impl MediaStoreInner for SqliteMediaStore {
472    type Error = Error;
473
474    async fn media_retention_policy_inner(
475        &self,
476    ) -> Result<Option<MediaRetentionPolicy>, Self::Error> {
477        let conn = self.read().await?;
478        conn.get_serialized_kv(keys::MEDIA_RETENTION_POLICY).await
479    }
480
481    async fn set_media_retention_policy_inner(
482        &self,
483        policy: MediaRetentionPolicy,
484    ) -> Result<(), Self::Error> {
485        let conn = self.write().await?;
486        conn.set_serialized_kv(keys::MEDIA_RETENTION_POLICY, policy).await?;
487        Ok(())
488    }
489
490    async fn add_media_content_inner(
491        &self,
492        request: &MediaRequestParameters,
493        data: Vec<u8>,
494        last_access: SystemTime,
495        policy: MediaRetentionPolicy,
496        ignore_policy: IgnoreMediaRetentionPolicy,
497    ) -> Result<(), Self::Error> {
498        let ignore_policy = ignore_policy.is_yes();
499        let data = self.encode_value(data)?;
500
501        if !ignore_policy && policy.exceeds_max_file_size(data.len() as u64) {
502            return Ok(());
503        }
504
505        let uri = self.encode_key(keys::MEDIA, request.source.unique_key());
506        let format = self.encode_key(keys::MEDIA, request.format.unique_key());
507        let timestamp = time_to_timestamp(last_access);
508
509        let conn = self.write().await?;
510        conn.execute(
511            "INSERT OR REPLACE INTO media (uri, format, data, last_access, ignore_policy) VALUES (?, ?, ?, ?, ?)",
512            (uri, format, data, timestamp, ignore_policy),
513        )
514        .await?;
515
516        Ok(())
517    }
518
519    async fn set_ignore_media_retention_policy_inner(
520        &self,
521        request: &MediaRequestParameters,
522        ignore_policy: IgnoreMediaRetentionPolicy,
523    ) -> Result<(), Self::Error> {
524        let uri = self.encode_key(keys::MEDIA, request.source.unique_key());
525        let format = self.encode_key(keys::MEDIA, request.format.unique_key());
526        let ignore_policy = ignore_policy.is_yes();
527
528        let conn = self.write().await?;
529        conn.execute(
530            r#"UPDATE media SET ignore_policy = ? WHERE uri = ? AND format = ?"#,
531            (ignore_policy, uri, format),
532        )
533        .await?;
534
535        Ok(())
536    }
537
538    async fn get_media_content_inner(
539        &self,
540        request: &MediaRequestParameters,
541        current_time: SystemTime,
542    ) -> Result<Option<Vec<u8>>, Self::Error> {
543        let uri = self.encode_key(keys::MEDIA, request.source.unique_key());
544        let format = self.encode_key(keys::MEDIA, request.format.unique_key());
545        let timestamp = time_to_timestamp(current_time);
546
547        let conn = self.write().await?;
548        let data = conn
549            .with_transaction::<_, rusqlite::Error, _>(move |txn| {
550                // Update the last access. We need to do this first so the
551                // transaction is in write mode right away. See:
552                // https://sqlite.org/lang_transaction.html#read_transactions_versus_write_transactions
553                txn.execute(
554                    "UPDATE media SET last_access = ? WHERE uri = ? AND format = ?",
555                    (timestamp, &uri, &format),
556                )?;
557
558                txn.query_row::<Vec<u8>, _, _>(
559                    "SELECT data FROM media WHERE uri = ? AND format = ?",
560                    (&uri, &format),
561                    |row| row.get(0),
562                )
563                .optional()
564            })
565            .await?;
566
567        data.map(|v| self.decode_value(&v).map(Into::into)).transpose()
568    }
569
570    async fn clean_inner(
571        &self,
572        policy: MediaRetentionPolicy,
573        current_time: SystemTime,
574    ) -> Result<(), Self::Error> {
575        if !policy.has_limitations() {
576            // We can safely skip all the checks.
577            return Ok(());
578        }
579
580        let conn = self.write().await?;
581        let removed = conn
582            .with_transaction::<_, Error, _>(move |txn| {
583                let mut removed = false;
584
585                // First, check media content that exceed the max filesize.
586                if let Some(max_file_size) = policy.computed_max_file_size() {
587                    let count = txn.execute(
588                        "DELETE FROM media WHERE ignore_policy IS FALSE AND length(data) > ?",
589                        (max_file_size,),
590                    )?;
591
592                    if count > 0 {
593                        removed = true;
594                    }
595                }
596
597                // Then, clean up expired media content.
598                if let Some(last_access_expiry) = policy.last_access_expiry {
599                    let current_timestamp = time_to_timestamp(current_time);
600                    let expiry_secs = last_access_expiry.as_secs();
601                    let count = txn.execute(
602                        "DELETE FROM media WHERE ignore_policy IS FALSE AND (? - last_access) >= ?",
603                        (current_timestamp, expiry_secs),
604                    )?;
605
606                    if count > 0 {
607                        removed = true;
608                    }
609                }
610
611                // Finally, if the cache size is too big, remove old items until it fits.
612                if let Some(max_cache_size) = policy.max_cache_size {
613                    // i64 is the integer type used by SQLite, use it here to
614                    // avoid usize overflow during the conversion of the result.
615                    let cache_size = txn
616                        .query_row(
617                            "SELECT sum(length(data)) FROM media WHERE ignore_policy IS FALSE",
618                            (),
619                            |row| {
620                                // `sum()` returns `NULL` if there are no rows.
621                                row.get::<_, Option<u64>>(0)
622                            },
623                        )?
624                        .unwrap_or_default();
625
626                    // If the cache size is overflowing or bigger than max cache size, clean up.
627                    if cache_size > max_cache_size {
628                        // Get the sizes of the media contents ordered by last access.
629                        let mut cached_stmt = txn.prepare_cached(
630                            "SELECT rowid, length(data) FROM media \
631                             WHERE ignore_policy IS FALSE ORDER BY last_access DESC",
632                        )?;
633                        let content_sizes = cached_stmt
634                            .query(())?
635                            .mapped(|row| Ok((row.get::<_, i64>(0)?, row.get::<_, u64>(1)?)));
636
637                        let mut accumulated_items_size = 0u64;
638                        let mut limit_reached = false;
639                        let mut rows_to_remove = Vec::new();
640
641                        for result in content_sizes {
642                            let (row_id, size) = match result {
643                                Ok(content_size) => content_size,
644                                Err(error) => {
645                                    return Err(error.into());
646                                }
647                            };
648
649                            if limit_reached {
650                                rows_to_remove.push(row_id);
651                                continue;
652                            }
653
654                            match accumulated_items_size.checked_add(size) {
655                                Some(acc) if acc > max_cache_size => {
656                                    // We can stop accumulating.
657                                    limit_reached = true;
658                                    rows_to_remove.push(row_id);
659                                }
660                                Some(acc) => accumulated_items_size = acc,
661                                None => {
662                                    // The accumulated size is overflowing but
663                                    // the setting cannot be bigger than
664                                    // usize::MAX, we can stop accumulating.
665                                    limit_reached = true;
666                                    rows_to_remove.push(row_id);
667                                }
668                            }
669                        }
670
671                        if !rows_to_remove.is_empty() {
672                            removed = true;
673                        }
674
675                        txn.chunk_large_query_over(rows_to_remove, None, |txn, row_ids| {
676                            let sql_params = row_ids.host_parameters();
677                            let query = format!("DELETE FROM media WHERE rowid IN ({sql_params})");
678                            txn.prepare(&query)?.execute(params_from_iter(row_ids))?;
679                            Ok(Vec::<()>::new())
680                        })?;
681                    }
682                }
683
684                txn.set_serialized_kv(keys::LAST_MEDIA_CLEANUP_TIME, current_time)?;
685
686                Ok(removed)
687            })
688            .await?;
689
690        // If we removed media, defragment the database and free space on the
691        // filesystem.
692        if removed {
693            conn.vacuum().await?;
694        }
695
696        Ok(())
697    }
698
699    async fn last_media_cleanup_time_inner(&self) -> Result<Option<SystemTime>, Self::Error> {
700        let conn = self.read().await?;
701        conn.get_serialized_kv(keys::LAST_MEDIA_CLEANUP_TIME).await
702    }
703}
704
705#[cfg(test)]
706mod tests {
707    use std::{
708        path::PathBuf,
709        sync::{
710            LazyLock,
711            atomic::{AtomicU32, Ordering::SeqCst},
712        },
713        time::Duration,
714    };
715
716    use matrix_sdk_base::{
717        media::{
718            MediaFormat, MediaRequestParameters, MediaThumbnailSettings,
719            store::{IgnoreMediaRetentionPolicy, MediaStore, MediaStoreError},
720        },
721        media_store_inner_integration_tests, media_store_integration_tests,
722        media_store_integration_tests_time,
723    };
724    use matrix_sdk_test::async_test;
725    use ruma::{events::room::MediaSource, media::Method, mxc_uri, uint};
726    use tempfile::{TempDir, tempdir};
727
728    use super::SqliteMediaStore;
729    use crate::{SqliteStoreConfig, utils::SqliteAsyncConnExt};
730
731    static TMP_DIR: LazyLock<TempDir> = LazyLock::new(|| tempdir().unwrap());
732    static NUM: AtomicU32 = AtomicU32::new(0);
733
734    fn new_media_store_workspace() -> PathBuf {
735        let name = NUM.fetch_add(1, SeqCst).to_string();
736        TMP_DIR.path().join(name)
737    }
738
739    async fn get_media_store() -> Result<SqliteMediaStore, MediaStoreError> {
740        let tmpdir_path = new_media_store_workspace();
741
742        tracing::info!("using media store @ {}", tmpdir_path.to_str().unwrap());
743
744        Ok(SqliteMediaStore::open(tmpdir_path.to_str().unwrap(), None).await.unwrap())
745    }
746
747    media_store_integration_tests!();
748    media_store_integration_tests_time!();
749    media_store_inner_integration_tests!();
750
751    async fn get_media_store_content_sorted_by_last_access(
752        media_store: &SqliteMediaStore,
753    ) -> Vec<Vec<u8>> {
754        let sqlite_db = media_store.read().await.expect("accessing sqlite db failed");
755        sqlite_db
756            .prepare("SELECT data FROM media ORDER BY last_access DESC", |mut stmt| {
757                stmt.query(())?.mapped(|row| row.get(0)).collect()
758            })
759            .await
760            .expect("querying media cache content by last access failed")
761    }
762
763    #[async_test]
764    async fn test_pool_size() {
765        let tmpdir_path = new_media_store_workspace();
766        let store_open_config = SqliteStoreConfig::new(tmpdir_path).pool_max_size(42);
767
768        let store = SqliteMediaStore::open_with_config(&store_open_config).await.unwrap();
769
770        assert_eq!(store.pool_max_size().await.unwrap(), 42);
771    }
772
773    #[async_test]
774    async fn test_last_access() {
775        let media_store = get_media_store().await.expect("creating media cache failed");
776        let uri = mxc_uri!("mxc://localhost/media");
777        let file_request = MediaRequestParameters {
778            source: MediaSource::Plain(uri.to_owned()),
779            format: MediaFormat::File,
780        };
781        let thumbnail_request = MediaRequestParameters {
782            source: MediaSource::Plain(uri.to_owned()),
783            format: MediaFormat::Thumbnail(MediaThumbnailSettings::with_method(
784                Method::Crop,
785                uint!(100),
786                uint!(100),
787            )),
788        };
789
790        let content: Vec<u8> = "hello world".into();
791        let thumbnail_content: Vec<u8> = "hello…".into();
792
793        // Add the media.
794        media_store
795            .add_media_content(&file_request, content.clone(), IgnoreMediaRetentionPolicy::No)
796            .await
797            .expect("adding file failed");
798
799        // Since the precision of the timestamp is in seconds, wait so the
800        // timestamps differ.
801        tokio::time::sleep(Duration::from_secs(3)).await;
802
803        media_store
804            .add_media_content(
805                &thumbnail_request,
806                thumbnail_content.clone(),
807                IgnoreMediaRetentionPolicy::No,
808            )
809            .await
810            .expect("adding thumbnail failed");
811
812        // File's last access is older than thumbnail.
813        let contents = get_media_store_content_sorted_by_last_access(&media_store).await;
814
815        assert_eq!(contents.len(), 2, "media cache contents length is wrong");
816        assert_eq!(contents[0], thumbnail_content, "thumbnail is not last access");
817        assert_eq!(contents[1], content, "file is not second-to-last access");
818
819        // Since the precision of the timestamp is in seconds, wait so the
820        // timestamps differ.
821        tokio::time::sleep(Duration::from_secs(3)).await;
822
823        // Access the file so its last access is more recent.
824        let _ = media_store
825            .get_media_content(&file_request)
826            .await
827            .expect("getting file failed")
828            .expect("file is missing");
829
830        // File's last access is more recent than thumbnail.
831        let contents = get_media_store_content_sorted_by_last_access(&media_store).await;
832
833        assert_eq!(contents.len(), 2, "media cache contents length is wrong");
834        assert_eq!(contents[0], content, "file is not last access");
835        assert_eq!(contents[1], thumbnail_content, "thumbnail is not second-to-last access");
836    }
837}
838
839#[cfg(test)]
840mod close_reopen_tests {
841    use std::sync::{
842        LazyLock,
843        atomic::{AtomicU32, Ordering::SeqCst},
844    };
845
846    use matrix_sdk_base::media::{
847        MediaFormat, MediaRequestParameters,
848        store::{IgnoreMediaRetentionPolicy, MediaStore},
849    };
850    use matrix_sdk_test::async_test;
851    use ruma::{events::room::MediaSource, mxc_uri};
852    use tempfile::{TempDir, tempdir};
853
854    use super::SqliteMediaStore;
855
856    static TMP_DIR: LazyLock<TempDir> = LazyLock::new(|| tempdir().unwrap());
857    static NUM: AtomicU32 = AtomicU32::new(0);
858
859    async fn new_store() -> SqliteMediaStore {
860        let name = NUM.fetch_add(1, SeqCst).to_string();
861        let tmpdir_path = TMP_DIR.path().join(name);
862        SqliteMediaStore::open(tmpdir_path, None).await.unwrap()
863    }
864
865    fn test_request() -> MediaRequestParameters {
866        MediaRequestParameters {
867            source: MediaSource::Plain(mxc_uri!("mxc://localhost/test_media").to_owned()),
868            format: MediaFormat::File,
869        }
870    }
871
872    #[async_test]
873    async fn test_close_completes_without_timeout() {
874        let store = new_store().await;
875
876        // Close should complete quickly without hitting any timeout.
877        let start = std::time::Instant::now();
878        store.close().await.unwrap();
879        let elapsed = start.elapsed();
880
881        assert!(
882            elapsed < std::time::Duration::from_secs(2),
883            "close() took {elapsed:?}, expected < 2s (no timeout)"
884        );
885
886        // Connections should be None after close.
887        let guard = store.connections.lock().await;
888        assert!(guard.is_none(), "connections should be None after close");
889    }
890
891    #[async_test]
892    async fn test_reopen_restores_connections() {
893        let store = new_store().await;
894
895        store.close().await.unwrap();
896
897        {
898            let guard = store.connections.lock().await;
899            assert!(guard.is_none());
900        }
901
902        store.reopen().await.unwrap();
903
904        {
905            let guard = store.connections.lock().await;
906            assert!(guard.is_some(), "connections should be Some after reopen");
907        }
908    }
909
910    #[async_test]
911    async fn test_close_is_idempotent() {
912        let store = new_store().await;
913
914        store.close().await.unwrap();
915        // Second close should be a no-op.
916        store.close().await.unwrap();
917
918        let guard = store.connections.lock().await;
919        assert!(guard.is_none());
920    }
921
922    #[async_test]
923    async fn test_reopen_is_idempotent() {
924        let store = new_store().await;
925
926        // Reopen on an active store should be a no-op.
927        store.reopen().await.unwrap();
928
929        let guard = store.connections.lock().await;
930        assert!(guard.is_some());
931    }
932
933    #[async_test]
934    async fn test_read_fails_when_closed() {
935        let store = new_store().await;
936        store.close().await.unwrap();
937
938        let err = store.get_media_content(&test_request()).await;
939        assert!(err.is_err(), "read should fail when closed");
940
941        let err_msg = err.unwrap_err().to_string();
942        assert!(err_msg.contains("closed"), "error should mention 'closed', got: {err_msg}");
943    }
944
945    #[async_test]
946    async fn test_write_fails_when_closed() {
947        let store = new_store().await;
948        store.close().await.unwrap();
949
950        let err = store
951            .add_media_content(&test_request(), b"data".to_vec(), IgnoreMediaRetentionPolicy::No)
952            .await;
953        assert!(err.is_err(), "write should fail when closed");
954
955        let err_msg = err.unwrap_err().to_string();
956        assert!(err_msg.contains("closed"), "error should mention 'closed', got: {err_msg}");
957    }
958
959    #[async_test]
960    async fn test_data_persists_across_close_reopen() {
961        let store = new_store().await;
962
963        // Write some media content.
964        store
965            .add_media_content(
966                &test_request(),
967                b"hello world".to_vec(),
968                IgnoreMediaRetentionPolicy::Yes,
969            )
970            .await
971            .unwrap();
972
973        // Verify it's there.
974        let content = store.get_media_content(&test_request()).await.unwrap();
975        assert_eq!(content.as_deref(), Some(b"hello world".as_slice()));
976
977        // Close and reopen.
978        store.close().await.unwrap();
979        store.reopen().await.unwrap();
980
981        // Content should still be there after reopen.
982        let content = store.get_media_content(&test_request()).await.unwrap();
983        assert_eq!(
984            content.as_deref(),
985            Some(b"hello world".as_slice()),
986            "media content should persist across close/reopen"
987        );
988    }
989
990    #[async_test]
991    async fn test_multiple_close_reopen_cycles() {
992        let store = new_store().await;
993
994        for _ in 0..5 {
995            store.close().await.unwrap();
996            store.reopen().await.unwrap();
997
998            // After each cycle, the store should be fully operational.
999            let result = store.get_media_content(&test_request()).await;
1000            assert!(result.is_ok(), "store should work after close/reopen cycle");
1001        }
1002    }
1003
1004    #[async_test]
1005    async fn test_pool_is_fully_drained_after_close() {
1006        let store = new_store().await;
1007
1008        // Do a few reads to exercise the pool.
1009        let _ = store.get_media_content(&test_request()).await;
1010        let _ = store.get_media_content(&test_request()).await;
1011
1012        store.close().await.unwrap();
1013
1014        // After close, the connections field should be None (pool and write
1015        // connection have been fully torn down).
1016        let guard = store.connections.lock().await;
1017        assert!(guard.is_none(), "all connections should be released after close");
1018    }
1019
1020    #[async_test]
1021    async fn test_operations_work_immediately_after_reopen() {
1022        let store = new_store().await;
1023
1024        store.close().await.unwrap();
1025        store.reopen().await.unwrap();
1026
1027        // Read should work immediately after reopen.
1028        let result = store.get_media_content(&test_request()).await;
1029        assert!(result.is_ok(), "read should succeed immediately after reopen");
1030
1031        // Write should work immediately after reopen.
1032        let result = store
1033            .add_media_content(
1034                &test_request(),
1035                b"after_reopen".to_vec(),
1036                IgnoreMediaRetentionPolicy::No,
1037            )
1038            .await;
1039        assert!(result.is_ok(), "write should succeed immediately after reopen");
1040    }
1041
1042    #[async_test]
1043    async fn test_close_waits_for_held_read_connection_to_drain() {
1044        let store = new_store().await;
1045
1046        // Acquire a read connection and hold it, simulating an in-flight read.
1047        let held_conn = store.read().await.unwrap();
1048
1049        // Spawn close in a background task — it will close the pool and then
1050        // poll-wait for pool.status().size == 0 in the drain loop.
1051        let store_clone = store.clone();
1052        let close_handle = tokio::spawn(async move {
1053            store_clone.close().await.unwrap();
1054        });
1055
1056        // Give close() a moment to close the pool and enter the drain loop.
1057        tokio::time::sleep(std::time::Duration::from_millis(200)).await;
1058
1059        // The close task should still be running because we hold a connection.
1060        assert!(!close_handle.is_finished(), "close should be waiting for the held connection");
1061
1062        // Release the held connection — this lets pool.status().size drop to 0.
1063        drop(held_conn);
1064
1065        // Now close should complete promptly (well within the 5s timeout).
1066        let timeout = tokio::time::timeout(std::time::Duration::from_secs(3), close_handle).await;
1067        assert!(timeout.is_ok(), "close should complete after the held connection is released");
1068        timeout.unwrap().unwrap();
1069
1070        // Verify the store is fully closed.
1071        let guard = store.connections.lock().await;
1072        assert!(guard.is_none(), "connections should be None after close");
1073    }
1074}
1075
1076#[cfg(test)]
1077mod encrypted_tests {
1078    use std::sync::{
1079        LazyLock,
1080        atomic::{AtomicU32, Ordering::SeqCst},
1081    };
1082
1083    use matrix_sdk_base::{
1084        media::store::MediaStoreError, media_store_inner_integration_tests,
1085        media_store_integration_tests, media_store_integration_tests_time,
1086    };
1087    use tempfile::{TempDir, tempdir};
1088
1089    use super::SqliteMediaStore;
1090
1091    static TMP_DIR: LazyLock<TempDir> = LazyLock::new(|| tempdir().unwrap());
1092    static NUM: AtomicU32 = AtomicU32::new(0);
1093
1094    async fn get_media_store() -> Result<SqliteMediaStore, MediaStoreError> {
1095        let name = NUM.fetch_add(1, SeqCst).to_string();
1096        let tmpdir_path = TMP_DIR.path().join(name);
1097
1098        tracing::info!("using media store @ {}", tmpdir_path.to_str().unwrap());
1099
1100        Ok(SqliteMediaStore::open(tmpdir_path.to_str().unwrap(), Some("default_test_password"))
1101            .await
1102            .unwrap())
1103    }
1104
1105    media_store_integration_tests!();
1106    media_store_integration_tests_time!();
1107    media_store_inner_integration_tests!();
1108}