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
102    /// given 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
111    /// key to encrypt private data.
112    pub async fn open_with_key(
113        path: impl AsRef<Path>,
114        key: Option<&[u8; 32]>,
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
145    /// pool. 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 fails with
276        // the error message: "cannot change into wal mode from within a transaction".
277        conn.execute_batch("PRAGMA journal_mode = wal;").await?;
278        conn.with_transaction(|txn| {
279            txn.execute_batch(include_str!("../migrations/media_store/001_init.sql"))?;
280            txn.set_db_version(1)
281        })
282        .await?;
283    }
284
285    if version < 2 {
286        debug!("Upgrading database to version 2");
287        conn.with_transaction(|txn| {
288            txn.execute_batch(include_str!(
289                "../migrations/media_store/002_lease_locks_with_generation.sql"
290            ))?;
291            txn.set_db_version(2)
292        })
293        .await?;
294    }
295
296    Ok(())
297}
298
299#[async_trait]
300impl MediaStore for SqliteMediaStore {
301    type Error = Error;
302
303    #[instrument(skip(self))]
304    async fn try_take_leased_lock(
305        &self,
306        lease_duration_ms: u32,
307        key: &str,
308        holder: &str,
309    ) -> Result<Option<CrossProcessLockGeneration>> {
310        let key = key.to_owned();
311        let holder = holder.to_owned();
312
313        let now: u64 = MilliSecondsSinceUnixEpoch::now().get().into();
314        let expiration = now + lease_duration_ms as u64;
315
316        // Learn about the `excluded` keyword in https://sqlite.org/lang_upsert.html.
317        let generation = self
318            .write()
319            .await?
320            .with_transaction(move |txn| {
321                txn.query_one(
322                    "INSERT INTO lease_locks (key, holder, expiration)
323                    VALUES (?1, ?2, ?3)
324                    ON CONFLICT (key)
325                    DO
326                        UPDATE SET
327                            holder = excluded.holder,
328                            expiration = excluded.expiration,
329                            generation =
330                                CASE holder
331                                    WHEN excluded.holder THEN generation
332                                    ELSE generation + 1
333                                END
334                        WHERE
335                            holder = excluded.holder
336                            OR expiration < ?4
337                    RETURNING generation
338                    ",
339                    (key, holder, expiration, now),
340                    |row| row.get(0),
341                )
342                .optional()
343            })
344            .await?;
345
346        Ok(generation)
347    }
348
349    async fn add_media_content(
350        &self,
351        request: &MediaRequestParameters,
352        content: Vec<u8>,
353        ignore_policy: IgnoreMediaRetentionPolicy,
354    ) -> Result<()> {
355        let _timer = timer!("method");
356
357        self.media_service.add_media_content(self, request, content, ignore_policy).await
358    }
359
360    #[instrument(skip_all)]
361    async fn replace_media_key(
362        &self,
363        from: &MediaRequestParameters,
364        to: &MediaRequestParameters,
365    ) -> Result<(), Self::Error> {
366        let _timer = timer!("method");
367
368        let prev_uri = self.encode_key(keys::MEDIA, from.source.unique_key());
369        let prev_format = self.encode_key(keys::MEDIA, from.format.unique_key());
370
371        let new_uri = self.encode_key(keys::MEDIA, to.source.unique_key());
372        let new_format = self.encode_key(keys::MEDIA, to.format.unique_key());
373
374        let conn = self.write().await?;
375        conn.execute(
376            r#"UPDATE media SET uri = ?, format = ? WHERE uri = ? AND format = ?"#,
377            (new_uri, new_format, prev_uri, prev_format),
378        )
379        .await?;
380
381        Ok(())
382    }
383
384    #[instrument(skip_all)]
385    async fn get_media_content(&self, request: &MediaRequestParameters) -> Result<Option<Vec<u8>>> {
386        let _timer = timer!("method");
387
388        self.media_service.get_media_content(self, request).await
389    }
390
391    #[instrument(skip_all)]
392    async fn remove_media_content(&self, request: &MediaRequestParameters) -> Result<()> {
393        let _timer = timer!("method");
394
395        let uri = self.encode_key(keys::MEDIA, request.source.unique_key());
396        let format = self.encode_key(keys::MEDIA, request.format.unique_key());
397
398        let conn = self.write().await?;
399        conn.execute("DELETE FROM media WHERE uri = ? AND format = ?", (uri, format)).await?;
400
401        Ok(())
402    }
403
404    #[instrument(skip(self))]
405    async fn remove_media_content_for_uri(&self, uri: &MxcUri) -> Result<()> {
406        let _timer = timer!("method");
407
408        let uri = self.encode_key(keys::MEDIA, uri);
409
410        let conn = self.write().await?;
411        conn.execute("DELETE FROM media WHERE uri = ?", (uri,)).await?;
412
413        Ok(())
414    }
415
416    #[instrument(skip_all)]
417    async fn set_media_retention_policy(
418        &self,
419        policy: MediaRetentionPolicy,
420    ) -> Result<(), Self::Error> {
421        let _timer = timer!("method");
422
423        self.media_service.set_media_retention_policy(self, policy).await
424    }
425
426    #[instrument(skip_all)]
427    fn media_retention_policy(&self) -> MediaRetentionPolicy {
428        let _timer = timer!("method");
429
430        self.media_service.media_retention_policy()
431    }
432
433    #[instrument(skip_all)]
434    async fn set_ignore_media_retention_policy(
435        &self,
436        request: &MediaRequestParameters,
437        ignore_policy: IgnoreMediaRetentionPolicy,
438    ) -> Result<(), Self::Error> {
439        let _timer = timer!("method");
440
441        self.media_service.set_ignore_media_retention_policy(self, request, ignore_policy).await
442    }
443
444    #[instrument(skip_all)]
445    async fn clean(&self) -> Result<(), Self::Error> {
446        let _timer = timer!("method");
447
448        self.media_service.clean(self).await
449    }
450
451    async fn close(&self) -> Result<(), Self::Error> {
452        SqliteMediaStore::close(self).await
453    }
454
455    async fn reopen(&self) -> Result<(), Self::Error> {
456        SqliteMediaStore::reopen(self).await
457    }
458
459    async fn optimize(&self) -> Result<(), Self::Error> {
460        Ok(self.vacuum().await?)
461    }
462
463    async fn get_size(&self) -> Result<Option<usize>, Self::Error> {
464        self.get_db_size().await
465    }
466}
467
468#[cfg_attr(target_family = "wasm", async_trait(?Send))]
469#[cfg_attr(not(target_family = "wasm"), async_trait)]
470impl MediaStoreInner for SqliteMediaStore {
471    type Error = Error;
472
473    async fn media_retention_policy_inner(
474        &self,
475    ) -> Result<Option<MediaRetentionPolicy>, Self::Error> {
476        let conn = self.read().await?;
477        conn.get_serialized_kv(keys::MEDIA_RETENTION_POLICY).await
478    }
479
480    async fn set_media_retention_policy_inner(
481        &self,
482        policy: MediaRetentionPolicy,
483    ) -> Result<(), Self::Error> {
484        let conn = self.write().await?;
485        conn.set_serialized_kv(keys::MEDIA_RETENTION_POLICY, policy).await?;
486        Ok(())
487    }
488
489    async fn add_media_content_inner(
490        &self,
491        request: &MediaRequestParameters,
492        data: Vec<u8>,
493        last_access: SystemTime,
494        policy: MediaRetentionPolicy,
495        ignore_policy: IgnoreMediaRetentionPolicy,
496    ) -> Result<(), Self::Error> {
497        let ignore_policy = ignore_policy.is_yes();
498        let data = self.encode_value(data)?;
499
500        if !ignore_policy && policy.exceeds_max_file_size(data.len() as u64) {
501            return Ok(());
502        }
503
504        let uri = self.encode_key(keys::MEDIA, request.source.unique_key());
505        let format = self.encode_key(keys::MEDIA, request.format.unique_key());
506        let timestamp = time_to_timestamp(last_access);
507
508        let conn = self.write().await?;
509        conn.execute(
510            "INSERT OR REPLACE INTO media (uri, format, data, last_access, ignore_policy) VALUES (?, ?, ?, ?, ?)",
511            (uri, format, data, timestamp, ignore_policy),
512        )
513        .await?;
514
515        Ok(())
516    }
517
518    async fn set_ignore_media_retention_policy_inner(
519        &self,
520        request: &MediaRequestParameters,
521        ignore_policy: IgnoreMediaRetentionPolicy,
522    ) -> Result<(), Self::Error> {
523        let uri = self.encode_key(keys::MEDIA, request.source.unique_key());
524        let format = self.encode_key(keys::MEDIA, request.format.unique_key());
525        let ignore_policy = ignore_policy.is_yes();
526
527        let conn = self.write().await?;
528        conn.execute(
529            r#"UPDATE media SET ignore_policy = ? WHERE uri = ? AND format = ?"#,
530            (ignore_policy, uri, format),
531        )
532        .await?;
533
534        Ok(())
535    }
536
537    async fn get_media_content_inner(
538        &self,
539        request: &MediaRequestParameters,
540        current_time: SystemTime,
541    ) -> Result<Option<Vec<u8>>, Self::Error> {
542        let uri = self.encode_key(keys::MEDIA, request.source.unique_key());
543        let format = self.encode_key(keys::MEDIA, request.format.unique_key());
544        let timestamp = time_to_timestamp(current_time);
545
546        let conn = self.write().await?;
547        let data = conn
548            .with_transaction::<_, rusqlite::Error, _>(move |txn| {
549                // Update the last access.
550                // We need to do this first so the transaction is in write mode right away.
551                // See: https://sqlite.org/lang_transaction.html#read_transactions_versus_write_transactions
552                txn.execute(
553                    "UPDATE media SET last_access = ? WHERE uri = ? AND format = ?",
554                    (timestamp, &uri, &format),
555                )?;
556
557                txn.query_row::<Vec<u8>, _, _>(
558                    "SELECT data FROM media WHERE uri = ? AND format = ?",
559                    (&uri, &format),
560                    |row| row.get(0),
561                )
562                .optional()
563            })
564            .await?;
565
566        data.map(|v| self.decode_value(&v).map(Into::into)).transpose()
567    }
568
569    async fn clean_inner(
570        &self,
571        policy: MediaRetentionPolicy,
572        current_time: SystemTime,
573    ) -> Result<(), Self::Error> {
574        if !policy.has_limitations() {
575            // We can safely skip all the checks.
576            return Ok(());
577        }
578
579        let conn = self.write().await?;
580        let removed = conn
581            .with_transaction::<_, Error, _>(move |txn| {
582                let mut removed = false;
583
584                // First, check media content that exceed the max filesize.
585                if let Some(max_file_size) = policy.computed_max_file_size() {
586                    let count = txn.execute(
587                        "DELETE FROM media WHERE ignore_policy IS FALSE AND length(data) > ?",
588                        (max_file_size,),
589                    )?;
590
591                    if count > 0 {
592                        removed = true;
593                    }
594                }
595
596                // Then, clean up expired media content.
597                if let Some(last_access_expiry) = policy.last_access_expiry {
598                    let current_timestamp = time_to_timestamp(current_time);
599                    let expiry_secs = last_access_expiry.as_secs();
600                    let count = txn.execute(
601                        "DELETE FROM media WHERE ignore_policy IS FALSE AND (? - last_access) >= ?",
602                        (current_timestamp, expiry_secs),
603                    )?;
604
605                    if count > 0 {
606                        removed = true;
607                    }
608                }
609
610                // Finally, if the cache size is too big, remove old items until it fits.
611                if let Some(max_cache_size) = policy.max_cache_size {
612                    // i64 is the integer type used by SQLite, use it here to avoid usize overflow
613                    // during the conversion of the result.
614                    let cache_size = txn
615                        .query_row(
616                            "SELECT sum(length(data)) FROM media WHERE ignore_policy IS FALSE",
617                            (),
618                            |row| {
619                                // `sum()` returns `NULL` if there are no rows.
620                                row.get::<_, Option<u64>>(0)
621                            },
622                        )?
623                        .unwrap_or_default();
624
625                    // If the cache size is overflowing or bigger than max cache size, clean up.
626                    if cache_size > max_cache_size {
627                        // Get the sizes of the media contents ordered by last access.
628                        let mut cached_stmt = txn.prepare_cached(
629                            "SELECT rowid, length(data) FROM media \
630                             WHERE ignore_policy IS FALSE ORDER BY last_access DESC",
631                        )?;
632                        let content_sizes = cached_stmt
633                            .query(())?
634                            .mapped(|row| Ok((row.get::<_, i64>(0)?, row.get::<_, u64>(1)?)));
635
636                        let mut accumulated_items_size = 0u64;
637                        let mut limit_reached = false;
638                        let mut rows_to_remove = Vec::new();
639
640                        for result in content_sizes {
641                            let (row_id, size) = match result {
642                                Ok(content_size) => content_size,
643                                Err(error) => {
644                                    return Err(error.into());
645                                }
646                            };
647
648                            if limit_reached {
649                                rows_to_remove.push(row_id);
650                                continue;
651                            }
652
653                            match accumulated_items_size.checked_add(size) {
654                                Some(acc) if acc > max_cache_size => {
655                                    // We can stop accumulating.
656                                    limit_reached = true;
657                                    rows_to_remove.push(row_id);
658                                }
659                                Some(acc) => accumulated_items_size = acc,
660                                None => {
661                                    // The accumulated size is overflowing but the setting cannot be
662                                    // bigger than usize::MAX, we can stop accumulating.
663                                    limit_reached = true;
664                                    rows_to_remove.push(row_id);
665                                }
666                            }
667                        }
668
669                        if !rows_to_remove.is_empty() {
670                            removed = true;
671                        }
672
673                        txn.chunk_large_query_over(rows_to_remove, None, |txn, row_ids| {
674                            let sql_params = row_ids.host_parameters();
675                            let query = format!("DELETE FROM media WHERE rowid IN ({sql_params})");
676                            txn.prepare(&query)?.execute(params_from_iter(row_ids))?;
677                            Ok(Vec::<()>::new())
678                        })?;
679                    }
680                }
681
682                txn.set_serialized_kv(keys::LAST_MEDIA_CLEANUP_TIME, current_time)?;
683
684                Ok(removed)
685            })
686            .await?;
687
688        // If we removed media, defragment the database and free space on the
689        // filesystem.
690        if removed {
691            conn.vacuum().await?;
692        }
693
694        Ok(())
695    }
696
697    async fn last_media_cleanup_time_inner(&self) -> Result<Option<SystemTime>, Self::Error> {
698        let conn = self.read().await?;
699        conn.get_serialized_kv(keys::LAST_MEDIA_CLEANUP_TIME).await
700    }
701}
702
703#[cfg(test)]
704mod tests {
705    use std::{
706        path::PathBuf,
707        sync::{
708            LazyLock,
709            atomic::{AtomicU32, Ordering::SeqCst},
710        },
711        time::Duration,
712    };
713
714    use matrix_sdk_base::{
715        media::{
716            MediaFormat, MediaRequestParameters, MediaThumbnailSettings,
717            store::{IgnoreMediaRetentionPolicy, MediaStore, MediaStoreError},
718        },
719        media_store_inner_integration_tests, media_store_integration_tests,
720        media_store_integration_tests_time,
721    };
722    use matrix_sdk_test::async_test;
723    use ruma::{events::room::MediaSource, media::Method, mxc_uri, uint};
724    use tempfile::{TempDir, tempdir};
725
726    use super::SqliteMediaStore;
727    use crate::{SqliteStoreConfig, utils::SqliteAsyncConnExt};
728
729    static TMP_DIR: LazyLock<TempDir> = LazyLock::new(|| tempdir().unwrap());
730    static NUM: AtomicU32 = AtomicU32::new(0);
731
732    fn new_media_store_workspace() -> PathBuf {
733        let name = NUM.fetch_add(1, SeqCst).to_string();
734        TMP_DIR.path().join(name)
735    }
736
737    async fn get_media_store() -> Result<SqliteMediaStore, MediaStoreError> {
738        let tmpdir_path = new_media_store_workspace();
739
740        tracing::info!("using media store @ {}", tmpdir_path.to_str().unwrap());
741
742        Ok(SqliteMediaStore::open(tmpdir_path.to_str().unwrap(), None).await.unwrap())
743    }
744
745    media_store_integration_tests!();
746    media_store_integration_tests_time!();
747    media_store_inner_integration_tests!();
748
749    async fn get_media_store_content_sorted_by_last_access(
750        media_store: &SqliteMediaStore,
751    ) -> Vec<Vec<u8>> {
752        let sqlite_db = media_store.read().await.expect("accessing sqlite db failed");
753        sqlite_db
754            .prepare("SELECT data FROM media ORDER BY last_access DESC", |mut stmt| {
755                stmt.query(())?.mapped(|row| row.get(0)).collect()
756            })
757            .await
758            .expect("querying media cache content by last access failed")
759    }
760
761    #[async_test]
762    async fn test_pool_size() {
763        let tmpdir_path = new_media_store_workspace();
764        let store_open_config = SqliteStoreConfig::new(tmpdir_path).pool_max_size(42);
765
766        let store = SqliteMediaStore::open_with_config(&store_open_config).await.unwrap();
767
768        assert_eq!(store.pool_max_size().await.unwrap(), 42);
769    }
770
771    #[async_test]
772    async fn test_last_access() {
773        let media_store = get_media_store().await.expect("creating media cache failed");
774        let uri = mxc_uri!("mxc://localhost/media");
775        let file_request = MediaRequestParameters {
776            source: MediaSource::Plain(uri.to_owned()),
777            format: MediaFormat::File,
778        };
779        let thumbnail_request = MediaRequestParameters {
780            source: MediaSource::Plain(uri.to_owned()),
781            format: MediaFormat::Thumbnail(MediaThumbnailSettings::with_method(
782                Method::Crop,
783                uint!(100),
784                uint!(100),
785            )),
786        };
787
788        let content: Vec<u8> = "hello world".into();
789        let thumbnail_content: Vec<u8> = "hello…".into();
790
791        // Add the media.
792        media_store
793            .add_media_content(&file_request, content.clone(), IgnoreMediaRetentionPolicy::No)
794            .await
795            .expect("adding file failed");
796
797        // Since the precision of the timestamp is in seconds, wait so the timestamps
798        // differ.
799        tokio::time::sleep(Duration::from_secs(3)).await;
800
801        media_store
802            .add_media_content(
803                &thumbnail_request,
804                thumbnail_content.clone(),
805                IgnoreMediaRetentionPolicy::No,
806            )
807            .await
808            .expect("adding thumbnail failed");
809
810        // File's last access is older than thumbnail.
811        let contents = get_media_store_content_sorted_by_last_access(&media_store).await;
812
813        assert_eq!(contents.len(), 2, "media cache contents length is wrong");
814        assert_eq!(contents[0], thumbnail_content, "thumbnail is not last access");
815        assert_eq!(contents[1], content, "file is not second-to-last access");
816
817        // Since the precision of the timestamp is in seconds, wait so the timestamps
818        // differ.
819        tokio::time::sleep(Duration::from_secs(3)).await;
820
821        // Access the file so its last access is more recent.
822        let _ = media_store
823            .get_media_content(&file_request)
824            .await
825            .expect("getting file failed")
826            .expect("file is missing");
827
828        // File's last access is more recent than thumbnail.
829        let contents = get_media_store_content_sorted_by_last_access(&media_store).await;
830
831        assert_eq!(contents.len(), 2, "media cache contents length is wrong");
832        assert_eq!(contents[0], content, "file is not last access");
833        assert_eq!(contents[1], thumbnail_content, "thumbnail is not second-to-last access");
834    }
835}
836
837#[cfg(test)]
838mod close_reopen_tests {
839    use std::sync::{
840        LazyLock,
841        atomic::{AtomicU32, Ordering::SeqCst},
842    };
843
844    use matrix_sdk_base::media::{
845        MediaFormat, MediaRequestParameters,
846        store::{IgnoreMediaRetentionPolicy, MediaStore},
847    };
848    use matrix_sdk_test::async_test;
849    use ruma::{events::room::MediaSource, mxc_uri};
850    use tempfile::{TempDir, tempdir};
851
852    use super::SqliteMediaStore;
853
854    static TMP_DIR: LazyLock<TempDir> = LazyLock::new(|| tempdir().unwrap());
855    static NUM: AtomicU32 = AtomicU32::new(0);
856
857    async fn new_store() -> SqliteMediaStore {
858        let name = NUM.fetch_add(1, SeqCst).to_string();
859        let tmpdir_path = TMP_DIR.path().join(name);
860        SqliteMediaStore::open(tmpdir_path, None).await.unwrap()
861    }
862
863    fn test_request() -> MediaRequestParameters {
864        MediaRequestParameters {
865            source: MediaSource::Plain(mxc_uri!("mxc://localhost/test_media").to_owned()),
866            format: MediaFormat::File,
867        }
868    }
869
870    #[async_test]
871    async fn test_close_completes_without_timeout() {
872        let store = new_store().await;
873
874        // Close should complete quickly without hitting any timeout.
875        let start = std::time::Instant::now();
876        store.close().await.unwrap();
877        let elapsed = start.elapsed();
878
879        assert!(
880            elapsed < std::time::Duration::from_secs(2),
881            "close() took {elapsed:?}, expected < 2s (no timeout)"
882        );
883
884        // Connections should be None after close.
885        let guard = store.connections.lock().await;
886        assert!(guard.is_none(), "connections should be None after close");
887    }
888
889    #[async_test]
890    async fn test_reopen_restores_connections() {
891        let store = new_store().await;
892
893        store.close().await.unwrap();
894
895        {
896            let guard = store.connections.lock().await;
897            assert!(guard.is_none());
898        }
899
900        store.reopen().await.unwrap();
901
902        {
903            let guard = store.connections.lock().await;
904            assert!(guard.is_some(), "connections should be Some after reopen");
905        }
906    }
907
908    #[async_test]
909    async fn test_close_is_idempotent() {
910        let store = new_store().await;
911
912        store.close().await.unwrap();
913        // Second close should be a no-op.
914        store.close().await.unwrap();
915
916        let guard = store.connections.lock().await;
917        assert!(guard.is_none());
918    }
919
920    #[async_test]
921    async fn test_reopen_is_idempotent() {
922        let store = new_store().await;
923
924        // Reopen on an active store should be a no-op.
925        store.reopen().await.unwrap();
926
927        let guard = store.connections.lock().await;
928        assert!(guard.is_some());
929    }
930
931    #[async_test]
932    async fn test_read_fails_when_closed() {
933        let store = new_store().await;
934        store.close().await.unwrap();
935
936        let err = store.get_media_content(&test_request()).await;
937        assert!(err.is_err(), "read should fail when closed");
938
939        let err_msg = err.unwrap_err().to_string();
940        assert!(err_msg.contains("closed"), "error should mention 'closed', got: {err_msg}");
941    }
942
943    #[async_test]
944    async fn test_write_fails_when_closed() {
945        let store = new_store().await;
946        store.close().await.unwrap();
947
948        let err = store
949            .add_media_content(&test_request(), b"data".to_vec(), IgnoreMediaRetentionPolicy::No)
950            .await;
951        assert!(err.is_err(), "write should fail when closed");
952
953        let err_msg = err.unwrap_err().to_string();
954        assert!(err_msg.contains("closed"), "error should mention 'closed', got: {err_msg}");
955    }
956
957    #[async_test]
958    async fn test_data_persists_across_close_reopen() {
959        let store = new_store().await;
960
961        // Write some media content.
962        store
963            .add_media_content(
964                &test_request(),
965                b"hello world".to_vec(),
966                IgnoreMediaRetentionPolicy::Yes,
967            )
968            .await
969            .unwrap();
970
971        // Verify it's there.
972        let content = store.get_media_content(&test_request()).await.unwrap();
973        assert_eq!(content.as_deref(), Some(b"hello world".as_slice()));
974
975        // Close and reopen.
976        store.close().await.unwrap();
977        store.reopen().await.unwrap();
978
979        // Content should still be there after reopen.
980        let content = store.get_media_content(&test_request()).await.unwrap();
981        assert_eq!(
982            content.as_deref(),
983            Some(b"hello world".as_slice()),
984            "media content should persist across close/reopen"
985        );
986    }
987
988    #[async_test]
989    async fn test_multiple_close_reopen_cycles() {
990        let store = new_store().await;
991
992        for _ in 0..5 {
993            store.close().await.unwrap();
994            store.reopen().await.unwrap();
995
996            // After each cycle, the store should be fully operational.
997            let result = store.get_media_content(&test_request()).await;
998            assert!(result.is_ok(), "store should work after close/reopen cycle");
999        }
1000    }
1001
1002    #[async_test]
1003    async fn test_pool_is_fully_drained_after_close() {
1004        let store = new_store().await;
1005
1006        // Do a few reads to exercise the pool.
1007        let _ = store.get_media_content(&test_request()).await;
1008        let _ = store.get_media_content(&test_request()).await;
1009
1010        store.close().await.unwrap();
1011
1012        // After close, the connections field should be None (pool and write
1013        // connection have been fully torn down).
1014        let guard = store.connections.lock().await;
1015        assert!(guard.is_none(), "all connections should be released after close");
1016    }
1017
1018    #[async_test]
1019    async fn test_operations_work_immediately_after_reopen() {
1020        let store = new_store().await;
1021
1022        store.close().await.unwrap();
1023        store.reopen().await.unwrap();
1024
1025        // Read should work immediately after reopen.
1026        let result = store.get_media_content(&test_request()).await;
1027        assert!(result.is_ok(), "read should succeed immediately after reopen");
1028
1029        // Write should work immediately after reopen.
1030        let result = store
1031            .add_media_content(
1032                &test_request(),
1033                b"after_reopen".to_vec(),
1034                IgnoreMediaRetentionPolicy::No,
1035            )
1036            .await;
1037        assert!(result.is_ok(), "write should succeed immediately after reopen");
1038    }
1039
1040    #[async_test]
1041    async fn test_close_waits_for_held_read_connection_to_drain() {
1042        let store = new_store().await;
1043
1044        // Acquire a read connection and hold it, simulating an in-flight read.
1045        let held_conn = store.read().await.unwrap();
1046
1047        // Spawn close in a background task — it will close the pool and then
1048        // poll-wait for pool.status().size == 0 in the drain loop.
1049        let store_clone = store.clone();
1050        let close_handle = tokio::spawn(async move {
1051            store_clone.close().await.unwrap();
1052        });
1053
1054        // Give close() a moment to close the pool and enter the drain loop.
1055        tokio::time::sleep(std::time::Duration::from_millis(200)).await;
1056
1057        // The close task should still be running because we hold a connection.
1058        assert!(!close_handle.is_finished(), "close should be waiting for the held connection");
1059
1060        // Release the held connection — this lets pool.status().size drop to 0.
1061        drop(held_conn);
1062
1063        // Now close should complete promptly (well within the 5s timeout).
1064        let timeout = tokio::time::timeout(std::time::Duration::from_secs(3), close_handle).await;
1065        assert!(timeout.is_ok(), "close should complete after the held connection is released");
1066        timeout.unwrap().unwrap();
1067
1068        // Verify the store is fully closed.
1069        let guard = store.connections.lock().await;
1070        assert!(guard.is_none(), "connections should be None after close");
1071    }
1072}
1073
1074#[cfg(test)]
1075mod encrypted_tests {
1076    use std::sync::{
1077        LazyLock,
1078        atomic::{AtomicU32, Ordering::SeqCst},
1079    };
1080
1081    use matrix_sdk_base::{
1082        media::store::MediaStoreError, media_store_inner_integration_tests,
1083        media_store_integration_tests, media_store_integration_tests_time,
1084    };
1085    use tempfile::{TempDir, tempdir};
1086
1087    use super::SqliteMediaStore;
1088
1089    static TMP_DIR: LazyLock<TempDir> = LazyLock::new(|| tempdir().unwrap());
1090    static NUM: AtomicU32 = AtomicU32::new(0);
1091
1092    async fn get_media_store() -> Result<SqliteMediaStore, MediaStoreError> {
1093        let name = NUM.fetch_add(1, SeqCst).to_string();
1094        let tmpdir_path = TMP_DIR.path().join(name);
1095
1096        tracing::info!("using media store @ {}", tmpdir_path.to_str().unwrap());
1097
1098        Ok(SqliteMediaStore::open(tmpdir_path.to_str().unwrap(), Some("default_test_password"))
1099            .await
1100            .unwrap())
1101    }
1102
1103    media_store_integration_tests!();
1104    media_store_integration_tests_time!();
1105    media_store_inner_integration_tests!();
1106}