Skip to main content

matrix_sdk_sqlite/
utils.rs

1// Copyright 2022 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
15use core::fmt;
16use std::{
17    borrow::{Borrow, Cow},
18    cmp::min,
19    iter,
20    ops::Deref,
21};
22
23use async_trait::async_trait;
24use base64::Engine as _;
25use deadpool_sync::InteractError;
26use itertools::Itertools;
27use matrix_sdk_store_encryption::{EncryptableValue, StoreCipher};
28use ruma::{OwnedEventId, OwnedRoomId, serde::Raw, time::SystemTime};
29use rusqlite::{OptionalExtension, Params, Row, Statement, Transaction, limits::Limit};
30use serde::{Serialize, de::DeserializeOwned};
31use tracing::{error, trace, warn};
32use vodozemac::base64_encode;
33use zeroize::Zeroize;
34
35use crate::{
36    OpenStoreError, RuntimeConfig, Secret,
37    connection::Connection as SqliteAsyncConn,
38    error::{Error, Result},
39};
40
41#[derive(Clone, Debug, PartialEq, Eq, PartialOrd, Ord)]
42pub(crate) enum Key {
43    Plain(Vec<u8>),
44    Hashed([u8; 32]),
45}
46
47impl Deref for Key {
48    type Target = [u8];
49
50    fn deref(&self) -> &Self::Target {
51        match self {
52            Key::Plain(slice) => slice,
53            Key::Hashed(bytes) => bytes,
54        }
55    }
56}
57
58impl Borrow<[u8]> for Key {
59    fn borrow(&self) -> &[u8] {
60        self.deref()
61    }
62}
63
64impl rusqlite::ToSql for Key {
65    fn to_sql(&self) -> rusqlite::Result<rusqlite::types::ToSqlOutput<'_>> {
66        self.deref().to_sql()
67    }
68}
69
70#[async_trait]
71pub(crate) trait SqliteAsyncConnExt {
72    async fn execute<P>(
73        &self,
74        sql: impl AsRef<str> + Send + 'static,
75        params: P,
76    ) -> rusqlite::Result<usize>
77    where
78        P: Params + Send + 'static;
79
80    async fn execute_batch(&self, sql: impl AsRef<str> + Send + 'static) -> rusqlite::Result<()>;
81
82    async fn prepare<T, F>(
83        &self,
84        sql: impl AsRef<str> + Send + 'static,
85        f: F,
86    ) -> rusqlite::Result<T>
87    where
88        T: Send + 'static,
89        F: FnOnce(Statement<'_>) -> rusqlite::Result<T> + Send + 'static;
90
91    async fn query_row<T, P, F>(
92        &self,
93        sql: impl AsRef<str> + Send + 'static,
94        params: P,
95        f: F,
96    ) -> rusqlite::Result<T>
97    where
98        T: Send + 'static,
99        P: Params + Send + 'static,
100        F: FnOnce(&Row<'_>) -> rusqlite::Result<T> + Send + 'static;
101
102    async fn query_one<T, P, F>(
103        &self,
104        sql: impl AsRef<str> + Send + 'static,
105        params: P,
106        f: F,
107    ) -> rusqlite::Result<T>
108    where
109        T: Send + 'static,
110        P: Params + Send + 'static,
111        F: FnOnce(&Row<'_>) -> rusqlite::Result<T> + Send + 'static;
112
113    async fn query_many<T, P, F>(
114        &self,
115        sql: impl AsRef<str> + Send + 'static,
116        params: P,
117        f: F,
118    ) -> rusqlite::Result<Vec<T>>
119    where
120        T: Send + 'static,
121        P: Params + Send + 'static,
122        F: FnMut(&Row<'_>) -> rusqlite::Result<T> + Send + 'static;
123
124    async fn with_transaction<T, E, F>(&self, f: F) -> Result<T, E>
125    where
126        T: Send + 'static,
127        E: From<rusqlite::Error> + Send + 'static,
128        F: FnOnce(&Transaction<'_>) -> Result<T, E> + Send + 'static;
129
130    /// Chunk a large query over some keys.
131    ///
132    /// Imagine there is a _dynamic_ query that runs potentially large number of
133    /// parameters, so much that the maximum number of parameters can be hit.
134    /// Then, this helper is for you. It will execute the query on chunks of
135    /// parameters.
136    async fn chunk_large_query_over<Query, Res>(
137        &self,
138        mut keys_to_chunk: Vec<Key>,
139        result_capacity: Option<usize>,
140        do_query: Query,
141    ) -> Result<Vec<Res>>
142    where
143        Res: Send + 'static,
144        Query: Fn(&Transaction<'_>, ChunkFromLargeQuery<Key>) -> Result<Vec<Res>> + Send + 'static;
145
146    /// Apply the [`RuntimeConfig`].
147    ///
148    /// It will call the `Self::optimize`, `Self::cache_size` or
149    /// `Self::journal_size_limit` methods automatically based on the
150    /// `RuntimeConfig` values.
151    ///
152    /// It is possible to call these methods individually though. This
153    /// `apply_runtime_config` method allows to automate this process.
154    async fn apply_runtime_config(&self, runtime_config: RuntimeConfig) -> Result<()> {
155        let RuntimeConfig { optimize, cache_size, journal_size_limit } = runtime_config;
156
157        if optimize {
158            self.optimize().await?;
159        }
160
161        self.cache_size(cache_size).await?;
162        self.journal_size_limit(journal_size_limit).await?;
163
164        Ok(())
165    }
166
167    /// Optimize the database.
168    ///
169    /// The SQLite documentation recommends to run this regularly and after any
170    /// schema change. The easiest is to do it consistently when the store is
171    /// constructed, after eventual migrations.
172    ///
173    /// See [`PRAGMA optimize`] to learn more.
174    ///
175    /// [`PRAGMA cache_size`]: https://www.sqlite.org/pragma.html#pragma_optimize
176    async fn optimize(&self) -> Result<()> {
177        self.execute_batch("PRAGMA optimize = 0x10002;").await?;
178        Ok(())
179    }
180
181    /// Define the maximum size in **bytes** the SQLite cache can use.
182    ///
183    /// See [`PRAGMA cache_size`] to learn more.
184    ///
185    /// [`PRAGMA cache_size`]: https://www.sqlite.org/pragma.html#pragma_cache_size
186    async fn cache_size(&self, cache_size: u32) -> Result<()> {
187        // `N` in `PRAGMA cache_size = -N` is expressed in kibibytes.
188        // `cache_size` is expressed in bytes. Let's convert.
189        let n = cache_size / 1024;
190
191        self.execute_batch(format!("PRAGMA cache_size = -{n};")).await?;
192        Ok(())
193    }
194
195    /// Limit the size of the WAL file, in **bytes**.
196    ///
197    /// By default, while the DB connections of the databases are open, [the
198    /// size of the WAL file can keep increasing][size_wal_file] depending on
199    /// the size needed for the transactions. A critical case is `VACUUM`
200    /// which basically writes the content of the DB file to the WAL file
201    /// before writing it back to the DB file, so we end up taking twice the
202    /// size of the database.
203    ///
204    /// By setting this limit, the WAL file is truncated after its content is
205    /// written to the database, if it is bigger than the limit.
206    ///
207    /// See [`PRAGMA journal_size_limit`] to learn more. The value `limit`
208    /// corresponds to `N` in `PRAGMA journal_size_limit = N`.
209    ///
210    /// [size_wal_file]: https://www.sqlite.org/wal.html#avoiding_excessively_large_wal_files
211    /// [`PRAGMA journal_size_limit`]: https://www.sqlite.org/pragma.html#pragma_journal_size_limit
212    async fn journal_size_limit(&self, limit: u32) -> Result<()> {
213        self.execute_batch(format!("PRAGMA journal_size_limit = {limit};")).await?;
214        Ok(())
215    }
216
217    /// Defragment the database and free space on the filesystem.
218    ///
219    /// Only returns an error in tests, otherwise the error is only logged.
220    async fn vacuum(&self) -> Result<()> {
221        // Truncate the WAL file before vacuuming so it has room to grow.
222        self.wal_checkpoint().await;
223        if let Err(error) = self.execute_batch("VACUUM").await {
224            // Since this is an optimisation step, do not propagate the error
225            // but log it.
226            #[cfg(not(any(test, debug_assertions)))]
227            tracing::warn!("Failed to vacuum database: {error}");
228
229            // We want to know if there is an error with this step during tests.
230            #[cfg(any(test, debug_assertions))]
231            return Err(error.into());
232        } else {
233            trace!("VACUUM complete");
234            // Once vacuumed, truncate the WAL file again to purge the copied DB contents.
235            self.wal_checkpoint().await;
236        }
237
238        Ok(())
239    }
240
241    /// Adds a manual [WAL checkpoint] to copy back the contents of the WAL
242    /// files into the actual database, resetting the write-ahead log.
243    ///
244    /// [WAL checkpoint]: https://sqlite.org/c3ref/wal_checkpoint.html
245    async fn wal_checkpoint(&self) {
246        match self.execute_batch("PRAGMA wal_checkpoint(TRUNCATE);").await {
247            Ok(_) => trace!("WAL checkpoint completed"),
248            Err(error) => error!(?error, "WAL checkpoint error"),
249        }
250    }
251
252    async fn get_db_size(&self) -> Result<usize> {
253        let page_size =
254            self.query_row("PRAGMA page_size;", (), |row| row.get::<_, usize>(0)).await?;
255        let total_pages =
256            self.query_row("PRAGMA page_count;", (), |row| row.get::<_, usize>(0)).await?;
257
258        Ok(total_pages * page_size)
259    }
260}
261
262#[async_trait]
263impl SqliteAsyncConnExt for SqliteAsyncConn {
264    async fn execute<P>(
265        &self,
266        sql: impl AsRef<str> + Send + 'static,
267        params: P,
268    ) -> rusqlite::Result<usize>
269    where
270        P: Params + Send + 'static,
271    {
272        self.interact(move |conn| conn.execute(sql.as_ref(), params))
273            .await
274            .map_err(map_interact_err)?
275    }
276
277    async fn execute_batch(&self, sql: impl AsRef<str> + Send + 'static) -> rusqlite::Result<()> {
278        self.interact(move |conn| conn.execute_batch(sql.as_ref()))
279            .await
280            .map_err(map_interact_err)?
281    }
282
283    async fn prepare<T, F>(
284        &self,
285        sql: impl AsRef<str> + Send + 'static,
286        f: F,
287    ) -> rusqlite::Result<T>
288    where
289        T: Send + 'static,
290        F: FnOnce(Statement<'_>) -> rusqlite::Result<T> + Send + 'static,
291    {
292        self.interact(move |conn| f(conn.prepare(sql.as_ref())?)).await.map_err(map_interact_err)?
293    }
294
295    async fn query_row<T, P, F>(
296        &self,
297        sql: impl AsRef<str> + Send + 'static,
298        params: P,
299        f: F,
300    ) -> rusqlite::Result<T>
301    where
302        T: Send + 'static,
303        P: Params + Send + 'static,
304        F: FnOnce(&Row<'_>) -> rusqlite::Result<T> + Send + 'static,
305    {
306        self.interact(move |conn| conn.query_row(sql.as_ref(), params, f))
307            .await
308            .map_err(map_interact_err)?
309    }
310
311    async fn query_one<T, P, F>(
312        &self,
313        sql: impl AsRef<str> + Send + 'static,
314        params: P,
315        f: F,
316    ) -> rusqlite::Result<T>
317    where
318        T: Send + 'static,
319        P: Params + Send + 'static,
320        F: FnOnce(&Row<'_>) -> rusqlite::Result<T> + Send + 'static,
321    {
322        self.interact(move |conn| conn.query_one(sql.as_ref(), params, f))
323            .await
324            .map_err(map_interact_err)?
325    }
326
327    async fn query_many<T, P, F>(
328        &self,
329        sql: impl AsRef<str> + Send + 'static,
330        params: P,
331        f: F,
332    ) -> rusqlite::Result<Vec<T>>
333    where
334        T: Send + 'static,
335        P: Params + Send + 'static,
336        F: FnMut(&Row<'_>) -> rusqlite::Result<T> + Send + 'static,
337    {
338        self.interact(move |conn| {
339            let mut stmt = conn.prepare(sql.as_ref())?;
340            stmt.query_and_then(params, f)?.collect()
341        })
342        .await
343        .map_err(map_interact_err)?
344    }
345
346    async fn with_transaction<T, E, F>(&self, f: F) -> Result<T, E>
347    where
348        T: Send + 'static,
349        E: From<rusqlite::Error> + Send + 'static,
350        F: FnOnce(&Transaction<'_>) -> Result<T, E> + Send + 'static,
351    {
352        self.interact(move |conn| {
353            let txn = conn.transaction()?;
354            let result = f(&txn)?;
355            txn.commit()?;
356            Ok(result)
357        })
358        .await
359        .map_err(map_interact_err)
360        .map_err(E::from)?
361    }
362
363    async fn chunk_large_query_over<Query, Res>(
364        &self,
365        keys_to_chunk: Vec<Key>,
366        result_capacity: Option<usize>,
367        do_query: Query,
368    ) -> Result<Vec<Res>>
369    where
370        Res: Send + 'static,
371        Query: Fn(&Transaction<'_>, ChunkFromLargeQuery<Key>) -> Result<Vec<Res>> + Send + 'static,
372    {
373        self.with_transaction(move |txn| {
374            txn.chunk_large_query_over(keys_to_chunk, result_capacity, do_query)
375        })
376        .await
377    }
378}
379
380/// Map an [`InteractError`] into a [`rusqlite::Error`].
381///
382/// An [`InteractError::Panic`] will panic. An [`InteractError::Cancelled`] will
383/// generate a [`rusqlite::Error::SqliteFailure`] with the
384/// [`rusqlite::ffi::SQLITE_ABORT`] code.
385fn map_interact_err(error: InteractError) -> rusqlite::Error {
386    match error {
387        InteractError::Panic(p) => panic!("{p:?}"),
388        InteractError::Cancelled => rusqlite::Error::SqliteFailure(
389            rusqlite::ffi::Error::new(rusqlite::ffi::SQLITE_ABORT),
390            None,
391        ),
392    }
393}
394
395pub(crate) trait SqliteTransactionExt {
396    /// See [`SqliteAsyncConnExt::chunk_large_query_over`].
397    fn chunk_large_query_over<Key, Query, Res>(
398        &self,
399        keys_to_chunk: Vec<Key>,
400        result_capacity: Option<usize>,
401        do_query: Query,
402    ) -> Result<Vec<Res>>
403    where
404        Res: Send + 'static,
405        Query: Fn(&Transaction<'_>, ChunkFromLargeQuery<Key>) -> Result<Vec<Res>> + Send + 'static;
406}
407
408/// Represent the new chunk prepared by
409/// [`SqliteAsyncConnExt::chunk_large_query_over`] or
410/// [`SqliteTransactionExt::chunk_large_query_over`].
411#[repr(transparent)]
412pub(crate) struct ChunkFromLargeQuery<Key>(Vec<Key>);
413
414impl<Key> IntoIterator for ChunkFromLargeQuery<Key> {
415    type Item = Key;
416    type IntoIter = std::vec::IntoIter<Key>;
417
418    fn into_iter(self) -> Self::IntoIter {
419        self.0.into_iter()
420    }
421}
422
423impl<Key> ChunkFromLargeQuery<Key> {
424    /// Return the correct number of host parameters (the `?` variable in an SQL
425    /// query) equals to the size of the chunk.
426    pub fn host_parameters(&self) -> impl fmt::Display + use<Key> {
427        host_parameters(self.0.len())
428    }
429
430    /// Iterate over the keys in the chunk, by reference.
431    ///
432    /// To get an owned iterator, see the `IntoIterator` implementation.
433    pub fn iter(&self) -> std::slice::Iter<'_, Key> {
434        self.0.iter()
435    }
436}
437
438impl SqliteTransactionExt for Transaction<'_> {
439    fn chunk_large_query_over<Key, Query, Res>(
440        &self,
441        mut keys_to_chunk: Vec<Key>,
442        result_capacity: Option<usize>,
443        do_query: Query,
444    ) -> Result<Vec<Res>>
445    where
446        Res: Send + 'static,
447        Query: Fn(&Transaction<'_>, ChunkFromLargeQuery<Key>) -> Result<Vec<Res>> + Send + 'static,
448    {
449        // Divide by 2 to allow space for more static parameters (not part of
450        // `keys_to_chunk`).
451        let maximum_chunk_size = self.limit(Limit::SQLITE_LIMIT_VARIABLE_NUMBER)? / 2;
452        let maximum_chunk_size: usize = maximum_chunk_size
453            .try_into()
454            .map_err(|_| Error::SqliteMaximumVariableNumber(maximum_chunk_size))?;
455
456        if keys_to_chunk.len() < maximum_chunk_size {
457            // Chunking isn't necessary.
458            let chunk = keys_to_chunk;
459
460            Ok(do_query(self, ChunkFromLargeQuery(chunk))?)
461        } else {
462            // Chunking _is_ necessary.
463
464            // Define the accumulator.
465            let capacity = result_capacity.unwrap_or_default();
466            let mut all_results = Vec::with_capacity(capacity);
467
468            while !keys_to_chunk.is_empty() {
469                // Chunk and run the query.
470                let tail = keys_to_chunk.split_off(min(keys_to_chunk.len(), maximum_chunk_size));
471                let chunk = keys_to_chunk;
472                keys_to_chunk = tail;
473
474                all_results.extend(do_query(self, ChunkFromLargeQuery(chunk))?);
475            }
476
477            Ok(all_results)
478        }
479    }
480}
481
482/// Extension trait for a [`rusqlite::Connection`] that contains a key-value
483/// table named `kv`.
484///
485/// The table should be created like this:
486///
487/// ```sql
488/// CREATE TABLE "kv" (
489///     "key" TEXT PRIMARY KEY NOT NULL,
490///     "value" BLOB NOT NULL
491/// );
492/// ```
493pub(crate) trait SqliteKeyValueStoreConnExt {
494    /// Store the given value for the given key.
495    fn set_kv(&self, key: &str, value: &[u8]) -> rusqlite::Result<()>;
496
497    /// Store the given value for the given key by serializing it.
498    fn set_serialized_kv<T: Serialize + Send>(&self, key: &str, value: T) -> Result<()> {
499        let serialized_value = rmp_serde::to_vec_named(&value)?;
500        self.set_kv(key, &serialized_value)?;
501
502        Ok(())
503    }
504
505    /// Removes the current key and value if exists.
506    fn clear_kv(&self, key: &str) -> rusqlite::Result<()>;
507
508    /// Set the version of the database.
509    fn set_db_version(&self, version: u8) -> rusqlite::Result<()> {
510        self.set_kv("version", &[version])
511    }
512}
513
514impl SqliteKeyValueStoreConnExt for rusqlite::Connection {
515    fn set_kv(&self, key: &str, value: &[u8]) -> rusqlite::Result<()> {
516        self.execute(
517            "INSERT INTO kv VALUES (?1, ?2) ON CONFLICT (key) DO UPDATE SET value = ?2",
518            (key, value),
519        )?;
520        Ok(())
521    }
522
523    fn clear_kv(&self, key: &str) -> rusqlite::Result<()> {
524        self.execute("DELETE FROM kv WHERE key = ?1", (key,))?;
525        Ok(())
526    }
527}
528
529/// Extension trait for an [`SqliteAsyncConn`] that contains a key-value
530/// table named `kv`.
531///
532/// The table should be created like this:
533///
534/// ```sql
535/// CREATE TABLE "kv" (
536///     "key" TEXT PRIMARY KEY NOT NULL,
537///     "value" BLOB NOT NULL
538/// );
539/// ```
540#[async_trait]
541pub(crate) trait SqliteKeyValueStoreAsyncConnExt: SqliteAsyncConnExt {
542    /// Whether the `kv` table exists in this database.
543    async fn kv_table_exists(&self) -> rusqlite::Result<bool> {
544        self.query_row(
545            "SELECT EXISTS (SELECT 1 FROM sqlite_master WHERE type = 'table' AND name = 'kv')",
546            (),
547            |row| row.get(0),
548        )
549        .await
550    }
551
552    /// Get the stored value for the given key.
553    async fn get_kv(&self, key: &str) -> rusqlite::Result<Option<Vec<u8>>> {
554        let key = key.to_owned();
555        self.query_row("SELECT value FROM kv WHERE key = ?", (key,), |row| row.get(0))
556            .await
557            .optional()
558    }
559
560    /// Get the stored serialized value for the given key.
561    async fn get_serialized_kv<T: DeserializeOwned>(&self, key: &str) -> Result<Option<T>> {
562        let Some(bytes) = self.get_kv(key).await? else {
563            return Ok(None);
564        };
565
566        Ok(Some(rmp_serde::from_slice(&bytes)?))
567    }
568
569    /// Store the given value for the given key.
570    async fn set_kv(&self, key: &str, value: Vec<u8>) -> rusqlite::Result<()>;
571
572    /// Store the given value for the given key by serializing it.
573    async fn set_serialized_kv<T: Serialize + Send + 'static>(
574        &self,
575        key: &str,
576        value: T,
577    ) -> Result<()>;
578
579    /// Clears the given value for the given key.
580    async fn clear_kv(&self, key: &str) -> rusqlite::Result<()>;
581
582    /// Get the version of the database.
583    async fn db_version(&self) -> Result<u8, OpenStoreError> {
584        let kv_exists = self.kv_table_exists().await.map_err(OpenStoreError::LoadVersion)?;
585
586        if kv_exists {
587            match self.get_kv("version").await.map_err(OpenStoreError::LoadVersion)?.as_deref() {
588                Some([v]) => Ok(*v),
589                Some(_) => Err(OpenStoreError::InvalidVersion),
590                None => Err(OpenStoreError::MissingVersion),
591            }
592        } else {
593            Ok(0)
594        }
595    }
596
597    /// Get the [`StoreCipher`] of the database or create it.
598    async fn get_or_create_store_cipher(
599        &self,
600        secret: Secret,
601    ) -> Result<StoreCipher, OpenStoreError> {
602        const STORAGE_KEY: &str = "cipher";
603
604        let encrypted_cipher =
605            self.get_kv(STORAGE_KEY).await.map_err(OpenStoreError::LoadCipher)?;
606
607        let cipher = if let Some(encrypted) = encrypted_cipher {
608            match &secret {
609                Secret::PassPhrase(passphrase) => StoreCipher::import(passphrase, &encrypted)?,
610                Secret::Key(key) => StoreCipher::import_with_key(key.as_slice(), &encrypted)?,
611                Secret::HighEntropyPassPhrase { key, base64_variant } => {
612                    // Element X apps used the passphrase-based secret variant even though the
613                    // underlying secret was a randomly generated key.
614                    //
615                    // The `HighEntropyPassPhrase` variant was introduced to migrate these cipher
616                    // exports from a passphrase-based setup to a key-based setup.
617                    //
618                    // We first attempt to decrypt the cipher using the provided high-entropy
619                    // passphrase as a key. If this results in a KDF mismatch, it indicates that
620                    // the export was originally encrypted with the high-entropy passphrase being
621                    // used as a passphrase instead.
622                    //
623                    // In that case, we re-encrypt the cipher using the key-based setup. On the next
624                    // import attempt, `import_with_key()` can then decrypt it successfully.
625                    match StoreCipher::import_with_key(key.as_slice(), &encrypted) {
626                        Ok(cipher) => cipher,
627                        Err(matrix_sdk_store_encryption::Error::KdfMismatch) => {
628                            // EX generated a byte array for a key but converted it into a string by
629                            // base64 encoding it to use it as a passphrase. So let's do that as
630                            // well.
631                            //
632                            // Funnily enough, iOS used padded base64, while Android used unpadded.
633                            let mut base64_passphrase = match base64_variant {
634                                crate::Base64Variant::Unpadded => base64_encode(key),
635                                crate::Base64Variant::Padded => {
636                                    base64::prelude::BASE64_STANDARD.encode(key)
637                                }
638                            };
639
640                            let cipher = StoreCipher::import(&base64_passphrase, &encrypted);
641                            base64_passphrase.zeroize();
642
643                            let cipher = cipher?;
644                            let export = cipher.export_with_key(key.as_slice())?;
645
646                            self.set_kv(STORAGE_KEY, export)
647                                .await
648                                .map_err(OpenStoreError::SaveCipher)?;
649
650                            cipher
651                        }
652                        Err(e) => return Err(e.into()),
653                    }
654                }
655            }
656        } else {
657            let cipher = StoreCipher::new()?;
658
659            let export = match &secret {
660                Secret::PassPhrase(passphrase) => {
661                    #[cfg(not(test))]
662                    {
663                        cipher.export(passphrase)
664                    }
665                    #[cfg(test)]
666                    {
667                        cipher._insecure_export_fast_for_testing(passphrase)
668                    }
669                }
670                Secret::Key(key) => cipher.export_with_key(key.as_slice()),
671                Secret::HighEntropyPassPhrase { key, .. } => cipher.export_with_key(key.as_slice()),
672            }?;
673
674            self.set_kv(STORAGE_KEY, export).await.map_err(OpenStoreError::SaveCipher)?;
675
676            cipher
677        };
678
679        Ok(cipher)
680    }
681}
682
683#[async_trait]
684impl SqliteKeyValueStoreAsyncConnExt for SqliteAsyncConn {
685    async fn set_kv(&self, key: &str, value: Vec<u8>) -> rusqlite::Result<()> {
686        let key = key.to_owned();
687        self.interact(move |conn| conn.set_kv(&key, &value)).await.unwrap()?;
688
689        Ok(())
690    }
691
692    async fn set_serialized_kv<T: Serialize + Send + 'static>(
693        &self,
694        key: &str,
695        value: T,
696    ) -> Result<()> {
697        let key = key.to_owned();
698        self.interact(move |conn| conn.set_serialized_kv(&key, value)).await.unwrap()?;
699
700        Ok(())
701    }
702
703    async fn clear_kv(&self, key: &str) -> rusqlite::Result<()> {
704        let key = key.to_owned();
705        self.interact(move |conn| conn.clear_kv(&key)).await.unwrap()?;
706
707        Ok(())
708    }
709}
710
711/// Repeat `?` n times, where n is defined by `count`. `?` are comma-separated.
712pub(crate) fn host_parameters(count: usize) -> impl fmt::Display {
713    assert_ne!(count, 0, "Can't generate zero host parameters");
714
715    iter::repeat_n("?", count).format(",")
716}
717
718/// Convert the given `SystemTime` to a timestamp, as the number of seconds
719/// since Unix Epoch.
720///
721/// Returns an `i64` as it is the numeric type used by SQLite.
722pub(crate) fn time_to_timestamp(time: SystemTime) -> i64 {
723    time.duration_since(SystemTime::UNIX_EPOCH)
724        .ok()
725        .and_then(|d| d.as_secs().try_into().ok())
726        // It is unlikely to happen unless the time on the system is seriously wrong, but we always
727        // need a value.
728        .unwrap_or(0)
729}
730
731/// Trait for a store that can encrypt its values, based on the presence of a
732/// cipher or not.
733///
734/// A single method must be implemented: `get_cypher`, which returns an optional
735/// cipher.
736///
737/// All the other methods come for free, based on the implementation of
738/// `get_cypher`.
739pub(crate) trait EncryptableStore {
740    fn get_cypher(&self) -> Option<&StoreCipher>;
741
742    /// If the store is using encryption, this will hash the given key. This is
743    /// useful when we need to do queries against a given key, but we don't
744    /// need to store the key in plain text (i.e. it's not both a key and a
745    /// value).
746    fn encode_key(&self, table_name: &str, key: impl AsRef<[u8]>) -> Key {
747        let bytes = key.as_ref();
748        if let Some(store_cipher) = self.get_cypher() {
749            Key::Hashed(store_cipher.hash_key(table_name, bytes))
750        } else {
751            Key::Plain(bytes.to_owned())
752        }
753    }
754
755    fn encode_value<V>(&self, value: V) -> Result<Vec<u8>>
756    where
757        V: EncryptableValue + Into<Vec<u8>>,
758    {
759        if let Some(key) = self.get_cypher() {
760            let encrypted = key.encrypt_value_data(value)?;
761            Ok(rmp_serde::to_vec_named(&encrypted)?)
762        } else {
763            Ok(value.into())
764        }
765    }
766
767    fn decode_value<'a>(&self, value: &'a [u8]) -> Result<Cow<'a, [u8]>> {
768        if let Some(key) = self.get_cypher() {
769            let encrypted = rmp_serde::from_slice(value)?;
770            let decrypted = key.decrypt_value_data(encrypted)?;
771            Ok(Cow::Owned(decrypted))
772        } else {
773            Ok(Cow::Borrowed(value))
774        }
775    }
776
777    fn serialize_value(&self, value: &impl Serialize) -> Result<Vec<u8>> {
778        let serialized = rmp_serde::to_vec_named(value)?;
779        self.encode_value(serialized)
780    }
781
782    fn deserialize_value<T: DeserializeOwned>(&self, value: &[u8]) -> Result<T> {
783        let decoded = self.decode_value(value)?;
784        Ok(rmp_serde::from_slice(&decoded)?)
785    }
786
787    fn serialize_json(&self, value: &impl Serialize) -> Result<Vec<u8>> {
788        let serialized = serde_json::to_vec(value)?;
789        self.encode_value(serialized)
790    }
791
792    fn deserialize_json<T: DeserializeOwned>(&self, data: &[u8]) -> Result<T> {
793        let decoded = self.decode_value(data)?;
794
795        let json_deserializer = &mut serde_json::Deserializer::from_slice(&decoded);
796
797        serde_path_to_error::deserialize(json_deserializer).map_err(|err| {
798            let raw_json: Option<Raw<serde_json::Value>> = serde_json::from_slice(&decoded).ok();
799
800            let target_type = std::any::type_name::<T>();
801            let serde_path = err.path().to_string();
802
803            error!(
804                sentry = true,
805                %err,
806                "Failed to deserialize {target_type} in a store: {serde_path}",
807            );
808
809            if let Some(raw) = raw_json {
810                if let Some(room_id) = raw.get_field::<OwnedRoomId>("room_id").ok().flatten() {
811                    warn!("Found a room id in the source data to deserialize: {room_id}");
812                }
813                if let Some(event_id) = raw.get_field::<OwnedEventId>("event_id").ok().flatten() {
814                    warn!("Found an event id in the source data to deserialize: {event_id}");
815                }
816            }
817
818            err.into_inner().into()
819        })
820    }
821}
822
823#[cfg(test)]
824mod unit_tests {
825    use std::time::Duration;
826
827    use super::*;
828
829    #[test]
830    fn test_can_generate_host_parameters() {
831        assert_eq!(host_parameters(1).to_string(), "?");
832        assert_eq!(host_parameters(2).to_string(), "?,?");
833        assert_eq!(host_parameters(5).to_string(), "?,?,?,?,?");
834    }
835
836    #[test]
837    #[should_panic(expected = "Can't generate zero host parameters")]
838    fn test_generating_zero_host_parameters_panics() {
839        host_parameters(0);
840    }
841
842    #[test]
843    fn test_time_to_timestamp() {
844        assert_eq!(time_to_timestamp(SystemTime::UNIX_EPOCH), 0);
845        assert_eq!(time_to_timestamp(SystemTime::UNIX_EPOCH + Duration::from_secs(60)), 60);
846
847        // Fallback value on overflow.
848        assert_eq!(time_to_timestamp(SystemTime::UNIX_EPOCH - Duration::from_secs(60)), 0);
849    }
850}