atagen

matrix-rust-sdk

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