Add database models (#289)

* Add key_store table and generate

* Implement store/fetch/delete traits

* Add SQL key store

* Use SQL key store for provider, remove in-memory key store

* Use reference to encrypted store to allow multiple consumers

* Tidy identity

* Add identity table, queries, and unit tests

* Persist and retrieve identity from storage inside builder

* Tidy up names and types

* Add tests and fix bugs

* Fix lints

* Update xmtp_mls/src/storage/encrypted_store/mod.rs

Co-authored-by: Andrew Plaza <[email protected]>

* Dont use the word Error in error names

* Move identity initialization to identity strategy

* Refactor DB stuff into multiple files

* Delete models.rs

* Cleanup

* Add kind to group_messages

* Update schema

* Update schema

* Update schema

* Add ciphersuite back in as import

---------

Co-authored-by: Richard Hua <[email protected]>
Co-authored-by: Richard Hua <[email protected]>
Co-authored-by: Andrew Plaza <[email protected]>

xmtp/migrations/2023-06-30-063813_create_users_and_conversations/up.sql

@@ -1,8 +1,9 @@
 CREATE TABLE users (
-    user_address TEXT PRIMARY KEY NOT NULL,
-    created_at BIGINT NOT NULL,
-    last_refreshed BIGINT NOT NULL
+  user_address TEXT PRIMARY KEY NOT NULL,
+  created_at BIGINT NOT NULL,
+  last_refreshed BIGINT NOT NULL
 );
+
 CREATE TABLE conversations (
   convo_id TEXT PRIMARY KEY NOT NULL,
   peer_address TEXT NOT NULL,

xmtp_mls/migrations/2023-10-29-205333_state_machine_init/down.sql

@@ -0,0 +1,9 @@
+DROP TABLE IF EXISTS groups;
+
+DROP TABLE IF EXISTS group_messages;
+
+DROP TABLE IF EXISTS topic_refresh_state;
+
+DROP TABLE IF EXISTS group_intents;
+
+DROP TABLE IF EXISTS outbound_welcome_messages;

xmtp_mls/migrations/2023-10-29-205333_state_machine_init/up.sql

@@ -0,0 +1,58 @@
+CREATE TABLE groups (
+    -- Random ID generated by group creator
+    "id" BLOB PRIMARY KEY NOT NULL,
+    -- Based on the timestamp of the welcome message
+    "created_at_ns" BIGINT NOT NULL,
+    -- Enum of GROUP_MEMBERSHIP_STATE
+    "membership_state" INT NOT NULL
+);
+
+-- Allow for efficient sorting of groups
+CREATE INDEX groups_created_at_idx ON groups(created_at_ns);
+
+CREATE INDEX groups_membership_state ON groups(membership_state);
+
+-- Successfully processed messages meant to be returned to the user
+CREATE TABLE group_messages (
+    -- Derived via SHA256(CONCAT(decrypted_message_bytes, conversation_id, timestamp))
+    "id" BLOB PRIMARY KEY NOT NULL,
+    "group_id" BLOB NOT NULL,
+    -- Message contents after decryption
+    "decrypted_message_bytes" BLOB NOT NULL,
+    -- Based on the timestamp of the message
+    "sent_at_ns" BIGINT NOT NULL,
+    -- Enum GROUP_MESSAGE_KIND
+    "kind" INT NOT NULL,
+    -- Could remove this if we added a table mapping installation_ids to wallet addresses
+    "sender_installation_id" BLOB NOT NULL,
+    "sender_wallet_address" TEXT NOT NULL,
+    FOREIGN KEY (group_id) REFERENCES groups(id)
+);
+
+CREATE INDEX group_messages_group_id_sort_idx ON group_messages(group_id, sent_at_ns);
+
+-- Used to keep track of the last seen message timestamp in a topic
+CREATE TABLE topic_refresh_state (
+    "topic" TEXT PRIMARY KEY NOT NULL,
+    "last_message_timestamp_ns" BIGINT NOT NULL
+);
+
+-- This table is required to retry messages that do not send successfully due to epoch conflicts
+CREATE TABLE group_intents (
+    -- Serial ID auto-generated by the DB
+    "id" INTEGER PRIMARY KEY AUTOINCREMENT NOT NULL,
+    -- Enum INTENT_KIND
+    "kind" INT NOT NULL,
+    "group_id" BLOB NOT NULL,
+    -- Some sort of serializable blob that can be used to re-try the message if the first attempt failed due to conflict
+    "data" BLOB NOT NULL,
+    -- INTENT_STATE,
+    "state" INT NOT NULL,
+    -- The hash of the encrypted, concrete, form of the message if it was published.
+    "payload_hash" BLOB,
+    -- (Optional) data needed for the post-commit flow. For example, welcome messages
+    "post_commit_data" BLOB,
+    FOREIGN KEY (group_id) REFERENCES groups(id)
+);
+
+CREATE INDEX group_intents_group_id_id ON group_intents(group_id, id);

xmtp_mls/src/builder.rs

@@ -1,4 +1,5 @@
-use crate::storage::StoredIdentity;
+use crate::configuration::CIPHERSUITE;
+use crate::storage::identity::StoredIdentity;
 use crate::xmtp_openmls_provider::XmtpOpenMlsProvider;
 use crate::{
     client::{Client, Network},

xmtp_mls/src/identity.rs

@@ -13,7 +13,7 @@ use xmtp_cryptography::signature::SignatureError;
 use crate::{
     association::{AssociationError, AssociationText, Eip191Association},
     configuration::CIPHERSUITE,
-    storage::{EncryptedMessageStore, StorageError, StoredIdentity},
+    storage::{identity::StoredIdentity, EncryptedMessageStore, StorageError},
     types::Address,
     xmtp_openmls_provider::XmtpOpenMlsProvider,
     InboxOwner, Store,

xmtp_mls/src/storage/encrypted_store/group.rs

@@ -0,0 +1,11 @@
+use super::schema::groups;
+use diesel::prelude::*;
+
+#[derive(Insertable, Identifiable, Queryable, Debug, Clone)]
+#[diesel(table_name = groups)]
+#[diesel(primary_key(id))]
+pub struct StoredGroup {
+    pub id: Vec<u8>,
+    pub created_at_ns: i64,
+    pub membership_state: i32,
+}

xmtp_mls/src/storage/encrypted_store/group_intent.rs

@@ -0,0 +1,24 @@
+use super::schema::group_intents;
+use diesel::prelude::*;
+
+#[derive(Queryable, Identifiable, Debug, Clone)]
+#[diesel(table_name = group_intents)]
+#[diesel(primary_key(id))]
+pub struct StoredGroupIntent {
+    pub id: i32,
+    pub kind: i32,
+    pub state: i32,
+    pub group_id: Vec<u8>,
+    pub data: Vec<u8>,
+    pub payload_hash: Option<Vec<u8>>,
+    pub post_commit_data: Option<Vec<u8>>,
+}
+
+#[derive(Insertable, Debug, Clone)]
+#[diesel(table_name = group_intents)]
+pub struct NewGroupIntent {
+    pub kind: i32,
+    pub state: i32,
+    pub group_id: Vec<u8>,
+    pub data: Vec<u8>,
+}

xmtp_mls/src/storage/encrypted_store/group_message.rs

@@ -0,0 +1,14 @@
+use super::schema::group_messages;
+use diesel::prelude::*;
+
+#[derive(Insertable, Identifiable, Queryable, Debug, Clone)]
+#[diesel(table_name = group_messages)]
+#[diesel(primary_key(id))]
+pub struct StoredGroupMessage {
+    pub id: Vec<u8>,
+    pub group_id: Vec<u8>,
+    pub decrypted_message_bytes: Vec<u8>,
+    pub sent_at_ns: i64,
+    pub sender_installation_id: Vec<u8>,
+    pub sender_wallet_address: String,
+}

xmtp_mls/src/storage/encrypted_store/models.rs → xmtp_mls/src/storage/encrypted_store/identity.rs

@@ -1,19 +1,10 @@
-use diesel::prelude::*;
-
+use super::{schema::identity, DbConnection, StorageError};
 use crate::{
     identity::Identity,
     storage::serialization::{db_deserialize, db_serialize},
+    Fetch, Store,
 };
-
-use super::schema::*;
-
-#[derive(Insertable, Queryable, Debug, Clone)]
-#[diesel(table_name = openmls_key_store)]
-#[diesel(primary_key(key_bytes))]
-pub struct StoredKeyStoreEntry {
-    pub key_bytes: Vec<u8>,
-    pub value_bytes: Vec<u8>,
-}
+use diesel::prelude::*;
 
 #[derive(Insertable, Queryable, Debug, Clone)]
 #[diesel(table_name = identity)]
@@ -59,3 +50,47 @@ impl From<StoredIdentity> for Identity {
         }
     }
 }
+
+impl Store<DbConnection> for StoredIdentity {
+    fn store(&self, into: &mut DbConnection) -> Result<(), StorageError> {
+        diesel::insert_into(identity::table)
+            .values(self)
+            .execute(into)?;
+        Ok(())
+    }
+}
+
+impl Fetch<StoredIdentity> for DbConnection {
+    type Key = ();
+    fn fetch(&mut self, _key: ()) -> Result<Option<StoredIdentity>, StorageError> where {
+        use super::schema::identity::dsl::*;
+        Ok(identity.first(self).optional()?)
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::{
+        super::{tests::rand_vec, EncryptedMessageStore, StorageOption},
+        StoredIdentity,
+    };
+    use crate::Store;
+
+    #[test]
+    fn can_only_store_one_identity() {
+        let store = EncryptedMessageStore::new(
+            StorageOption::Ephemeral,
+            EncryptedMessageStore::generate_enc_key(),
+        )
+        .unwrap();
+        let conn = &mut store.conn().unwrap();
+
+        StoredIdentity::new("".to_string(), rand_vec(), rand_vec())
+            .store(conn)
+            .unwrap();
+
+        let duplicate_insertion =
+            StoredIdentity::new("".to_string(), rand_vec(), rand_vec()).store(conn);
+        assert!(duplicate_insertion.is_err());
+    }
+}

xmtp_mls/src/storage/encrypted_store/key_store_entry.rs

@@ -0,0 +1,38 @@
+use super::DbConnection;
+use super::{schema::openmls_key_store, StorageError};
+use crate::{Delete, Fetch, Store};
+use diesel::prelude::*;
+
+#[derive(Insertable, Queryable, Debug, Clone)]
+#[diesel(table_name = openmls_key_store)]
+#[diesel(primary_key(key_bytes))]
+pub struct StoredKeyStoreEntry {
+    pub key_bytes: Vec<u8>,
+    pub value_bytes: Vec<u8>,
+}
+
+impl Store<DbConnection> for StoredKeyStoreEntry {
+    fn store(&self, into: &mut DbConnection) -> Result<(), StorageError> {
+        diesel::insert_into(openmls_key_store::table)
+            .values(self)
+            .execute(into)?;
+
+        Ok(())
+    }
+}
+
+impl Fetch<StoredKeyStoreEntry> for DbConnection {
+    type Key = Vec<u8>;
+    fn fetch(&mut self, key: Vec<u8>) -> Result<Option<StoredKeyStoreEntry>, StorageError> where {
+        use super::schema::openmls_key_store::dsl::*;
+        Ok(openmls_key_store.find(key).first(self).optional()?)
+    }
+}
+
+impl Delete<StoredKeyStoreEntry> for DbConnection {
+    type Key = Vec<u8>;
+    fn delete(&mut self, key: Vec<u8>) -> Result<usize, StorageError> where {
+        use super::schema::openmls_key_store::dsl::*;
+        Ok(diesel::delete(openmls_key_store.filter(key_bytes.eq(key))).execute(self)?)
+    }
+}

xmtp_mls/src/storage/encrypted_store/mod.rs

@@ -11,12 +11,13 @@
 //! `diesel print-schema` or use `cargo run update-schema` which will update the files for you.      
 //!
 
-pub mod models;
+pub mod group;
+pub mod group_intent;
+pub mod group_message;
+pub mod identity;
+pub mod key_store_entry;
 pub mod schema;
-
-use crate::{Delete, Fetch, Store};
-
-use self::{models::*, schema::*};
+pub mod topic_refresh_state;
 
 use super::StorageError;
 use diesel::{
@@ -165,52 +166,9 @@ fn warn_length<T>(list: &Vec<T>, str_id: &str, max_length: usize) {
     }
 }
 
-impl Store<DbConnection> for StoredKeyStoreEntry {
-    fn store(&self, into: &mut DbConnection) -> Result<(), StorageError> {
-        diesel::insert_into(openmls_key_store::table)
-            .values(self)
-            .execute(into)?;
-
-        Ok(())
-    }
-}
-
-impl Fetch<StoredKeyStoreEntry> for DbConnection {
-    type Key = Vec<u8>;
-    fn fetch(&mut self, key: Vec<u8>) -> Result<Option<StoredKeyStoreEntry>, StorageError> where {
-        use self::schema::openmls_key_store::dsl::*;
-        Ok(openmls_key_store.find(key).first(self).optional()?)
-    }
-}
-
-impl Delete<StoredKeyStoreEntry> for DbConnection {
-    type Key = Vec<u8>;
-    fn delete(&mut self, key: Vec<u8>) -> Result<usize, StorageError> where {
-        use self::schema::openmls_key_store::dsl::*;
-        Ok(diesel::delete(openmls_key_store.filter(key_bytes.eq(key))).execute(self)?)
-    }
-}
-
-impl Store<DbConnection> for StoredIdentity {
-    fn store(&self, into: &mut DbConnection) -> Result<(), StorageError> {
-        diesel::insert_into(identity::table)
-            .values(self)
-            .execute(into)?;
-        Ok(())
-    }
-}
-
-impl Fetch<StoredIdentity> for DbConnection {
-    type Key = ();
-    fn fetch(&mut self, _key: ()) -> Result<Option<StoredIdentity>, StorageError> where {
-        use self::schema::identity::dsl::*;
-        Ok(identity.first(self).optional()?)
-    }
-}
-
 #[cfg(test)]
 mod tests {
-    use super::{models::*, EncryptedMessageStore, StorageError, StorageOption};
+    use super::{identity::StoredIdentity, EncryptedMessageStore, StorageError, StorageOption};
     use crate::{Fetch, Store};
     use rand::{
         distributions::{Alphanumeric, DistString},
@@ -219,11 +177,11 @@ mod tests {
     use std::boxed::Box;
     use std::fs;
 
-    fn rand_string() -> String {
+    pub(crate) fn rand_string() -> String {
         Alphanumeric.sample_string(&mut rand::thread_rng(), 16)
     }
 
-    fn rand_vec() -> Vec<u8> {
+    pub(crate) fn rand_vec() -> Vec<u8> {
         rand::thread_rng().gen::<[u8; 16]>().to_vec()
     }
 
@@ -296,24 +254,6 @@ mod tests {
         fs::remove_file(db_path).unwrap();
     }
 
-    #[test]
-    fn can_only_store_one_identity() {
-        let store = EncryptedMessageStore::new(
-            StorageOption::Ephemeral,
-            EncryptedMessageStore::generate_enc_key(),
-        )
-        .unwrap();
-        let conn = &mut store.conn().unwrap();
-
-        StoredIdentity::new("".to_string(), rand_vec(), rand_vec())
-            .store(conn)
-            .unwrap();
-
-        let duplicate_insertion =
-            StoredIdentity::new("".to_string(), rand_vec(), rand_vec()).store(conn);
-        assert!(duplicate_insertion.is_err());
-    }
-
     #[test]
     fn it_returns_ok_when_given_ok_result() {
         let result: Result<(), diesel::result::Error> = Ok(());