xmtp_db/encrypted_store/
mod.rs

1//! A durable object store powered by Sqlite and Diesel.
2//!
3//! Provides mechanism to store objects between sessions. The behavior of the store can be tailored
4//! by choosing an appropriate `StoreOption`.
5//!
6//! ## Migrations
7//!
8//! Table definitions are located `<PackageRoot>/migrations/`. On initialization the store will see
9//! if there are any outstanding database migrations and perform them as needed. When updating the
10//! table definitions `schema.rs` must also be updated. To generate the correct schemas you can run
11//! `diesel print-schema` or use `cargo run update-schema` which will update the files for you.
12
13pub mod association_state;
14pub mod consent_record;
15pub mod conversation_list;
16pub mod database;
17pub mod db_connection;
18pub mod delivery;
19pub mod group;
20pub mod group_intent;
21pub mod group_message;
22pub mod identity;
23pub mod identity_cache;
24pub mod identity_update;
25pub mod incoming_envelope;
26pub mod key_package_history;
27pub mod key_store_entry;
28pub mod local_commit_log;
29pub mod message_deletion;
30pub mod migrations;
31pub mod notifications;
32pub mod pending_remove;
33pub mod pragmas;
34pub mod processed_device_sync_messages;
35pub mod readd_status;
36pub mod refresh_state;
37pub mod remote_commit_log;
38pub mod schema;
39mod schema_gen;
40pub mod server_configuration;
41pub mod store;
42pub mod stream_storage;
43pub mod tasks;
44pub mod user_preferences;
45
46pub use self::db_connection::DbConnection;
47use diesel::{migration::Migration, result::DatabaseErrorKind};
48pub use diesel::{
49    migration::MigrationSource,
50    sqlite::{Sqlite, SqliteConnection},
51};
52use openmls::storage::OpenMlsProvider;
53use prost::DecodeError;
54use xmtp_common::{ErrorCode, MaybeSend, MaybeSync, RetryableError};
55use xmtp_proto::ConversionError;
56use zeroize::ZeroizeOnDrop;
57
58use super::StorageError;
59use crate::sql_key_store::SqlKeyStoreError;
60use crate::{Store, XmtpMlsStorageProvider};
61
62pub use database::*;
63pub use store::*;
64
65use diesel::{prelude::*, sql_query};
66use diesel_migrations::{EmbeddedMigrations, MigrationHarness, embed_migrations};
67use std::{ops::Deref, sync::Arc};
68pub const MIGRATIONS: EmbeddedMigrations = embed_migrations!("./migrations/");
69
70#[derive(ZeroizeOnDrop, Clone)]
71pub struct EncryptionKey([u8; 32]);
72impl std::fmt::Debug for EncryptionKey {
73    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
74        f.debug_tuple("EncryptionKey").field(&"xxxx").finish()
75    }
76}
77
78impl Deref for EncryptionKey {
79    type Target = [u8; 32];
80    fn deref(&self) -> &Self::Target {
81        &self.0
82    }
83}
84
85impl<T> AsRef<T> for EncryptionKey
86where
87    T: ?Sized,
88    <EncryptionKey as Deref>::Target: AsRef<T>,
89{
90    fn as_ref(&self) -> &T {
91        self.deref().as_ref()
92    }
93}
94
95impl TryFrom<Vec<u8>> for EncryptionKey {
96    type Error = ConversionError;
97    fn try_from(v: Vec<u8>) -> Result<EncryptionKey, Self::Error> {
98        Ok(EncryptionKey(v.as_slice().try_into()?))
99    }
100}
101
102impl From<[u8; 32]> for EncryptionKey {
103    fn from(v: [u8; 32]) -> Self {
104        EncryptionKey(v)
105    }
106}
107
108impl TryFrom<&[u8]> for EncryptionKey {
109    type Error = ConversionError;
110    fn try_from(v: &[u8]) -> Result<EncryptionKey, Self::Error> {
111        let bytes: [u8; 32] = v.try_into()?;
112        Ok(EncryptionKey(bytes))
113    }
114}
115
116// For PRAGMA query log statements
117#[derive(QueryableByName, Debug)]
118struct SqliteVersion {
119    #[diesel(sql_type = diesel::sql_types::Text)]
120    version: String,
121}
122
123#[derive(Default, Clone, Debug, zeroize::ZeroizeOnDrop)]
124pub enum StorageOption {
125    #[default]
126    Ephemeral,
127    Persistent(String),
128}
129
130impl std::fmt::Display for StorageOption {
131    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
132        match self {
133            StorageOption::Ephemeral => write!(f, "Ephemeral"),
134            StorageOption::Persistent(path) => write!(f, "Persistent({})", path),
135        }
136    }
137}
138
139#[derive(thiserror::Error, Debug, ErrorCode)]
140pub enum ConnectionError {
141    /// Database error.
142    ///
143    /// Diesel database query error. May be retryable.
144    #[error(transparent)]
145    Database(#[from] diesel::result::Error),
146    #[error(transparent)]
147    #[error_code(inherit)]
148    Platform(#[from] PlatformStorageError),
149    /// Decode error.
150    ///
151    /// Protobuf decode failed within DB layer. Not retryable.
152    #[error(transparent)]
153    DecodeError(#[from] DecodeError),
154    /// Disconnect in transaction.
155    ///
156    /// Cannot disconnect while transaction is active. Retryable.
157    #[error("disconnect not possible in transaction")]
158    DisconnectInTransaction,
159    /// Reconnect in transaction.
160    ///
161    /// Cannot reconnect while transaction is active. Retryable.
162    #[error("reconnect not possible in transaction")]
163    ReconnectInTransaction,
164    /// Invalid query.
165    ///
166    /// Invalid query parameters or configuration. Not retryable.
167    #[error("invalid query: {0}")]
168    InvalidQuery(String),
169    /// Invalid version.
170    ///
171    /// DB migration version mismatch -- running a newer DB on older LibXMTP. Not retryable.
172    #[error(
173        "Applied migrations does not match available migrations.\n\
174    This is likely due to running a database that is newer than this version of libxmtp.\n\
175    Expected: {expected}, found: {found}"
176    )]
177    InvalidVersion { expected: String, found: String },
178}
179
180impl RetryableError for ConnectionError {
181    fn is_retryable(&self) -> bool {
182        match self {
183            Self::Database(d) => d.is_retryable(),
184            Self::Platform(n) => n.is_retryable(),
185            Self::DecodeError(_) => false,
186            Self::DisconnectInTransaction => true,
187            Self::ReconnectInTransaction => true,
188            Self::InvalidQuery(_) => false,
189            Self::InvalidVersion { .. } => false,
190        }
191    }
192}
193
194impl ConnectionError {
195    /// True when the pool can't currently hand out a connection. Mirrors
196    /// [`StorageError::db_needs_connection`].
197    #[cfg(not(target_arch = "wasm32"))]
198    pub fn db_needs_connection(&self) -> bool {
199        use PlatformStorageError::{Pool, PoolNeedsConnection};
200        matches!(self, Self::Platform(PoolNeedsConnection | Pool(_)))
201    }
202
203    #[cfg(target_arch = "wasm32")]
204    pub fn db_needs_connection(&self) -> bool {
205        matches!(self, Self::Platform(PlatformStorageError::Disconnected))
206    }
207}
208
209pub trait ConnectionExt: MaybeSend + MaybeSync {
210    /// Run a scoped query against the underlying SQLite connection.
211    fn raw_query<T, F>(&self, fun: F) -> Result<T, crate::ConnectionError>
212    where
213        F: FnOnce(&mut SqliteConnection) -> Result<T, diesel::result::Error>,
214        Self: Sized;
215
216    fn disconnect(&self) -> Result<(), ConnectionError>;
217    fn reconnect(&self) -> Result<(), ConnectionError>;
218}
219
220impl<C> ConnectionExt for &C
221where
222    C: ConnectionExt,
223{
224    fn raw_query<T, F>(&self, fun: F) -> Result<T, crate::ConnectionError>
225    where
226        F: FnOnce(&mut SqliteConnection) -> Result<T, diesel::result::Error>,
227        Self: Sized,
228    {
229        <C as ConnectionExt>::raw_query(self, fun)
230    }
231
232    fn disconnect(&self) -> Result<(), ConnectionError> {
233        <C as ConnectionExt>::disconnect(self)
234    }
235
236    fn reconnect(&self) -> Result<(), ConnectionError> {
237        <C as ConnectionExt>::reconnect(self)
238    }
239}
240
241impl<C> ConnectionExt for &mut C
242where
243    C: ConnectionExt,
244{
245    fn raw_query<T, F>(&self, fun: F) -> Result<T, crate::ConnectionError>
246    where
247        F: FnOnce(&mut SqliteConnection) -> Result<T, diesel::result::Error>,
248        Self: Sized,
249    {
250        <C as ConnectionExt>::raw_query(self, fun)
251    }
252
253    fn disconnect(&self) -> Result<(), ConnectionError> {
254        <C as ConnectionExt>::disconnect(self)
255    }
256
257    fn reconnect(&self) -> Result<(), ConnectionError> {
258        <C as ConnectionExt>::reconnect(self)
259    }
260}
261
262impl<C> ConnectionExt for Arc<C>
263where
264    C: ConnectionExt,
265{
266    fn raw_query<T, F>(&self, fun: F) -> Result<T, crate::ConnectionError>
267    where
268        F: FnOnce(&mut SqliteConnection) -> Result<T, diesel::result::Error>,
269        Self: Sized,
270    {
271        <C as ConnectionExt>::raw_query(self, fun)
272    }
273
274    fn disconnect(&self) -> Result<(), ConnectionError> {
275        <C as ConnectionExt>::disconnect(self)
276    }
277
278    fn reconnect(&self) -> Result<(), ConnectionError> {
279        <C as ConnectionExt>::reconnect(self)
280    }
281}
282
283pub type BoxedDatabase = Box<
284    dyn XmtpDb<
285            Connection = diesel::SqliteConnection,
286            DbQuery = DbConnection<diesel::SqliteConnection>,
287        >,
288>;
289
290#[cfg_attr(any(feature = "test-utils", test), mockall::automock(type Connection = crate::mock::MockConnection; type DbQuery = crate::mock::MockDbQuery;))]
291pub trait XmtpDb: MaybeSend + MaybeSync {
292    /// The Connection type for this database
293    type Connection: ConnectionExt + MaybeSend + MaybeSync;
294
295    type DbQuery: crate::DbQuery + MaybeSend + MaybeSync;
296
297    /// Reject incompatible formats before migration; never infer F from old processed progress.
298    fn init(&self) -> Result<(), StorageError> {
299        self.conn().raw_query(|conn| {
300            self.validate(conn).map_err(|e| {
301                diesel::result::Error::DatabaseError(
302                    DatabaseErrorKind::Unknown,
303                    Box::new(e.to_string()),
304                )
305            })?;
306            #[derive(QueryableByName)]
307            struct MigrationTable {
308                #[diesel(sql_type = diesel::sql_types::Text)]
309                name: String,
310            }
311            let migration_table = sql_query(
312                "SELECT name FROM sqlite_master WHERE type = 'table' AND name = '__diesel_schema_migrations'",
313            ).get_result::<MigrationTable>(conn).optional()?;
314            if let Some(table) = migration_table {
315                debug_assert_eq!(table.name, "__diesel_schema_migrations");
316                let baseline = MIGRATIONS.final_migration();
317                let applied = conn.applied_migrations()
318                    .map_err(diesel::result::Error::QueryBuilderError)?;
319                if applied.iter().any(|version| version.to_string() != baseline) {
320                    return Ok(Err(StorageError::PreTransitionDatabase));
321                }
322                if !applied.is_empty() {
323                    let current_format = sql_query(
324                        "SELECT name FROM pragma_table_info('refresh_state') WHERE name = 'received_sequence_id'",
325                    ).get_result::<MigrationTable>(conn).optional()?;
326                    if current_format.is_none() {
327                        return Ok(Err(StorageError::OldStreamDatabase));
328                    }
329                    // Spec 006 CFG-044: the baseline gained `server_configuration`.
330                    // Diesel records one version for the whole baseline, so a
331                    // database created by an earlier self-hosted build is never
332                    // re-migrated and would meet the configuration queries with
333                    // "no such table". Reject it here, with the same instruction
334                    // every other baseline change gives, instead of failing later.
335                    let has_configuration = sql_query(
336                        "SELECT name FROM sqlite_master WHERE type = 'table' AND name = 'server_configuration'",
337                    ).get_result::<MigrationTable>(conn).optional()?;
338                    if has_configuration.is_none() {
339                        return Ok(Err(StorageError::OldStreamDatabase));
340                    }
341                }
342            }
343            conn.run_pending_migrations(MIGRATIONS)
344                .map_err(diesel::result::Error::QueryBuilderError)?;
345
346            // Ensure the database version is what we expect
347            let db_version = conn.final_migration()?;
348            let last_migration = MIGRATIONS.final_migration();
349            if db_version != last_migration {
350                return Ok(Err(ConnectionError::InvalidVersion {
351                    expected: last_migration,
352                    found: db_version,
353                }.into()));
354            }
355
356            let sqlite_version =
357                sql_query("SELECT sqlite_version() AS version").load::<SqliteVersion>(conn)?;
358            tracing::info!("sqlite_version={}", sqlite_version[0].version);
359
360            tracing::info!("Migrations successful");
361            Ok(Ok(()))
362        })??;
363
364        Ok(())
365    }
366
367    /// The Options this database was created with
368    fn opts(&self) -> &StorageOption;
369
370    /// Validate a connection is as expected
371    fn validate(&self, _conn: &mut SqliteConnection) -> Result<(), ConnectionError> {
372        Ok(())
373    }
374
375    /// Returns the Connection implementation for this Database
376    fn conn(&self) -> Self::Connection;
377
378    /// Returns a higher-level wrapeped DbConnection from which high-level queries may be
379    /// accessed.
380    fn db(&self) -> Self::DbQuery;
381
382    /// Reconnect to the database
383    fn reconnect(&self) -> Result<(), ConnectionError>;
384
385    /// Release connection to the database, closing it
386    fn disconnect(&self) -> Result<(), ConnectionError>;
387}
388
389#[macro_export]
390macro_rules! impl_fetch {
391    ($model:ty, $table:ident, $key:ty, select) => {
392        impl<C: $crate::ConnectionExt> $crate::Fetch<$model> for C {
393            type Key = $key;
394            fn fetch(&self, key: &Self::Key) -> Result<Option<$model>, $crate::StorageError> {
395                use $crate::diesel::{OptionalExtension, QueryDsl, RunQueryDsl, SelectableHelper};
396                self.raw_query(|conn| {
397                    $crate::encrypted_store::schema::$table::table
398                        .find(key.clone())
399                        .select(<$model>::as_select())
400                        .first(conn)
401                        .optional()
402                })
403                .map_err(Into::into)
404            }
405        }
406    };
407    ($model:ty, $table:ident) => {
408        impl<C> $crate::Fetch<$model> for C
409        where
410            C: $crate::ConnectionExt,
411        {
412            type Key = ();
413            fn fetch(&self, _key: &Self::Key) -> Result<Option<$model>, $crate::StorageError> {
414                use $crate::encrypted_store::schema::$table::dsl::*;
415                self.raw_query(|conn| $table.first(conn).optional())
416                    .map_err(Into::into)
417            }
418        }
419    };
420
421    ($model:ty, $table:ident, $key:ty) => {
422        impl<C> $crate::Fetch<$model> for C
423        where
424            C: $crate::ConnectionExt,
425        {
426            type Key = $key;
427            fn fetch(&self, key: &Self::Key) -> Result<Option<$model>, $crate::StorageError> {
428                use $crate::encrypted_store::schema::$table::dsl::*;
429                self.raw_query::<_, _>(|conn| $table.find(key.clone()).first(conn).optional())
430                    .map_err(Into::into)
431            }
432        }
433    };
434}
435
436#[macro_export]
437macro_rules! impl_fetch_list {
438    ($model:ty, $table:ident) => {
439        impl<C> $crate::FetchList<$model> for C
440        where
441            C: $crate::ConnectionExt,
442        {
443            fn fetch_list(&self) -> Result<Vec<$model>, $crate::StorageError> {
444                use $crate::encrypted_store::schema::$table::dsl::*;
445                self.raw_query(|conn| $table.load::<$model>(conn))
446                    .map_err(Into::into)
447            }
448        }
449    };
450}
451
452// Inserts the model into the database by primary key, erroring if the model already exists
453#[macro_export]
454macro_rules! impl_store {
455    ($model:ty, $table:ident) => {
456        impl<C> $crate::Store<C> for $model
457        where
458            C: $crate::ConnectionExt,
459        {
460            type Output = ();
461            fn store(&self, into: &C) -> Result<(), $crate::StorageError> {
462                into.raw_query::<_, _>(|conn| {
463                    diesel::insert_into($table::table)
464                        .values(self)
465                        .execute(conn)
466                        .map_err(Into::into)
467                        .map(|_| ())
468                })
469                .map_err(Into::into)
470            }
471        }
472    };
473}
474
475#[macro_export]
476macro_rules! impl_store_or_ignore {
477    // Original variant without return type parameter (defaults to returning ())
478    ($model:ty, $table:ident) => {
479        impl<C> $crate::StoreOrIgnore<C> for $model
480        where
481            C: $crate::ConnectionExt,
482        {
483            type Output = ();
484
485            fn store_or_ignore(&self, into: &C) -> Result<(), $crate::StorageError> {
486                into.raw_query(|conn| {
487                    diesel::insert_or_ignore_into($table::table)
488                        .values(self)
489                        .execute(conn)
490                        .map_err(Into::into)
491                        .map(|_| ())
492                })
493                .map_err(Into::into)
494            }
495        }
496    };
497}
498
499impl<T, C> Store<DbConnection<C>> for Vec<T>
500where
501    T: Store<DbConnection<C>>,
502{
503    type Output = ();
504    fn store(&self, into: &DbConnection<C>) -> Result<Self::Output, StorageError> {
505        for item in self {
506            item.store(into)?;
507        }
508        Ok(())
509    }
510}
511
512pub trait MlsProviderExt: OpenMlsProvider<StorageError = SqlKeyStoreError> {
513    type XmtpStorage: XmtpMlsStorageProvider;
514
515    fn key_store(&self) -> &Self::XmtpStorage;
516}
517
518trait EmbeddedMigrationsExt {
519    fn final_migration(&self) -> String;
520}
521impl EmbeddedMigrationsExt for EmbeddedMigrations {
522    fn final_migration(&self) -> String {
523        let migrations: Vec<Box<dyn Migration<Sqlite>>> = self
524            .migrations()
525            .expect("Migrations are directly embedded, so this cannot error");
526        migrations
527            .first()
528            .expect("There is at least one migration")
529            .name()
530            .to_string()
531            .chars()
532            .filter(|c| c.is_numeric())
533            .collect()
534    }
535}
536
537trait MigrationHarnessExt {
538    fn final_migration(&mut self) -> Result<String, diesel::result::Error>;
539}
540
541impl MigrationHarnessExt for SqliteConnection {
542    fn final_migration(&mut self) -> Result<String, diesel::result::Error> {
543        let migration: String = self
544            .applied_migrations()
545            .map_err(diesel::result::Error::QueryBuilderError)?
546            .pop()
547            .expect("This function should be run after migrations are applied")
548            .to_string();
549
550        Ok(migration)
551    }
552}
553
554#[cfg(test)]
555pub(crate) mod tests {
556    use super::*;
557    use crate::{Fetch, Store, XmtpTestDb, identity::StoredIdentity};
558    use xmtp_common::{rand_vec, tmp_path};
559
560    #[xmtp_common::test]
561    async fn ephemeral_store() {
562        let store = crate::TestDb::create_ephemeral_store().await;
563        let conn = store.conn();
564
565        let inbox_id = "inbox_id";
566        StoredIdentity::new(inbox_id.to_string(), rand_vec::<24>(), rand_vec::<24>())
567            .store(&conn)
568            .unwrap();
569
570        let fetched_identity: StoredIdentity = conn.fetch(&()).unwrap().unwrap();
571        assert_eq!(fetched_identity.inbox_id, inbox_id);
572    }
573
574    #[xmtp_common::test]
575    async fn persistent_store() {
576        let db_path = tmp_path();
577        {
578            let store = crate::TestDb::create_persistent_store(Some(db_path.clone())).await;
579            let conn = &store.conn();
580
581            let inbox_id = "inbox_id";
582            StoredIdentity::new(inbox_id.to_string(), rand_vec::<24>(), rand_vec::<24>())
583                .store(conn)
584                .unwrap();
585
586            let fetched_identity: StoredIdentity = conn.fetch(&()).unwrap().unwrap();
587            assert_eq!(fetched_identity.inbox_id, inbox_id);
588        }
589        EncryptedMessageStore::<()>::remove_db_files(db_path)
590    }
591
592    #[xmtp_common::test]
593    async fn encrypted_db_with_multiple_connections() {
594        let db_path = tmp_path();
595        {
596            let store = crate::TestDb::create_persistent_store(Some(db_path.clone())).await;
597            let conn1 = &store.conn();
598            let inbox_id = "inbox_id";
599            StoredIdentity::new(inbox_id.to_string(), rand_vec::<24>(), rand_vec::<24>())
600                .store(conn1)
601                .unwrap();
602
603            let conn2 = &store.conn();
604            tracing::info!("Getting conn 2");
605            let fetched_identity: StoredIdentity = conn2.fetch(&()).unwrap().unwrap();
606            assert_eq!(fetched_identity.inbox_id, inbox_id);
607        }
608        EncryptedMessageStore::<()>::remove_db_files(db_path)
609    }
610
611    /// A query failing because the pool can't hand out a connection must report
612    /// `db_needs_connection()`. Uses a persistent store since only it has a pool to drop.
613    #[cfg(not(target_arch = "wasm32"))]
614    #[xmtp_common::test]
615    async fn pool_failure_needs_connection() {
616        let db_path = tmp_path();
617        {
618            let store = crate::TestDb::create_persistent_store(Some(db_path.clone())).await;
619            let conn = store.conn();
620
621            // Healthy pool: a query succeeds.
622            let ok: Result<Option<StoredIdentity>, _> = conn.fetch(&());
623            assert!(ok.is_ok());
624
625            // Drop the pool, then run a real query against it.
626            conn.disconnect().unwrap();
627            let res: Result<Option<StoredIdentity>, _> = conn.fetch(&());
628            let err = res.expect_err("query against a disconnected pool should fail");
629
630            assert!(
631                err.db_needs_connection(),
632                "expected db_needs_connection() for a pool failure, got: {err:?}"
633            );
634        }
635        EncryptedMessageStore::<()>::remove_db_files(db_path)
636    }
637}
638
639#[cfg(all(test, not(target_arch = "wasm32")))]
640mod db_needs_connection_tests {
641    use crate::{ConnectionError, PlatformStorageError, StorageError};
642
643    /// Pool errors must classify as "needs reconnect" (direct and `Connection`-wrapped);
644    /// non-pool errors must not. `Pool(_)` is covered e2e in `pool_failure_needs_connection`.
645    #[test]
646    fn pool_errors_need_connection() {
647        assert!(
648            StorageError::Platform(PlatformStorageError::PoolNeedsConnection).db_needs_connection()
649        );
650        assert!(
651            StorageError::Connection(ConnectionError::Platform(
652                PlatformStorageError::PoolNeedsConnection
653            ))
654            .db_needs_connection()
655        );
656
657        // A non-pool error must NOT be classified as needing a reconnect.
658        assert!(!StorageError::DbSerialize.db_needs_connection());
659    }
660}
xmtp_db/encrypted_store/mod.rs

xmtp_db/encrypted_store/
mod.rs
