chore: rewrite platform/generic/proteus in terms of the new traits

Author: coriolinus

keystore/src/entities/platform/generic/proteus/identity.rs

@@ -1,140 +1,120 @@
 use rusqlite::OptionalExtension;
 
 use crate::{
-    CryptoKeystoreError, MissingKeyErrorKind,
+    CryptoKeystoreError, CryptoKeystoreResult, Entity, EntityBase, EntityTransactionExt, UniqueEntity,
     connection::{KeystoreDatabaseConnection, TransactionWrapper},
-    entities::{Entity, EntityBase, EntityFindParams, EntityTransactionExt, ProteusIdentity, StringEntityId},
+    entities::ProteusIdentity,
 };
 
-#[async_trait::async_trait]
-impl Entity for ProteusIdentity {
-    fn id_raw(&self) -> &[u8] {
-        Self::ID
-    }
-
-    async fn find_all(
-        conn: &mut Self::ConnectionType,
-        _params: EntityFindParams,
-    ) -> crate::CryptoKeystoreResult<Vec<Self>> {
-        let mut res = vec![];
-        if let Some(identity) = Self::find_one(conn, &StringEntityId::default()).await? {
-            res.push(identity);
-        }
+impl EntityBase for ProteusIdentity {
+    type ConnectionType = KeystoreDatabaseConnection;
+    type AutoGeneratedFields = ();
+    const COLLECTION_NAME: &'static str = "proteus_identities";
 
-        Ok(res)
+    fn to_transaction_entity(self) -> crate::transaction::dynamic_dispatch::Entity {
+        crate::transaction::dynamic_dispatch::Entity::ProteusIdentity(self)
     }
+}
 
-    async fn find_one(
-        conn: &mut Self::ConnectionType,
-        _id: &StringEntityId,
-    ) -> crate::CryptoKeystoreResult<Option<Self>> {
-        let mut conn = conn.conn().await;
-        let transaction = conn.transaction()?;
+impl UniqueEntity for ProteusIdentity {
+    const KEY: () = ();
+}
 
-        let mut row_id: Option<i64> = transaction
-            .query_row(
-                "SELECT rowid FROM proteus_identities ORDER BY rowid ASC LIMIT 1",
-                [],
-                |r| r.get(0),
-            )
-            .optional()?;
+#[async_trait::async_trait]
+impl Entity for ProteusIdentity {
+    /// Each distinct [`PrimaryKey`] uniquely identifies either 0 or 1 instance.
+    ///
+    /// This constraint should be enforced at the DB level.
+    type PrimaryKey = ();
+
+    /// Get this entity's primary key.
+    fn primary_key(&self) -> Self::PrimaryKey {
+        ()
+    }
 
-        let row_id = if let Some(rowid) = row_id.take() {
-            rowid
-        } else {
+    /// Get an entity by its primary key.
+    ///
+    /// The type signature here is somewhat complicated, but it breaks down simply: if our primary key is something
+    /// like `Vec<u8>`, we want to be able to use this method even if what we have on hand is `&[u8]`.
+    ///
+    /// Actually ignoring the key here is unusual, but acceptable for a unique entity.
+    async fn get(conn: &mut Self::ConnectionType, _key: &()) -> CryptoKeystoreResult<Option<Self>> {
+        let conn = conn.conn().await;
+        // this gets the oldest identity, and is retained behavior from a previous implementation:
+        // it used to save additional bonus identities without error
+        let mut statement = conn.prepare_cached("SELECT pk, sk FROM proteus_identities ORDER BY rowid ASC LIMIT 1")?;
+        let Some(identity) = statement
+            .query_row([], |row| {
+                let identity = ProteusIdentity {
+                    pk: row.get("pk")?,
+                    sk: row.get("sk")?,
+                };
+                Ok(identity)
+            })
+            .optional()?
+        else {
             return Ok(None);
         };
 
-        use std::io::Read as _;
-        let mut blob = transaction.blob_open(rusqlite::MAIN_DB, "proteus_identities", "pk", row_id, true)?;
-        if blob.len() != Self::PK_KEY_SIZE {
+        if identity.pk.len() != Self::PK_KEY_SIZE {
             return Err(CryptoKeystoreError::InvalidKeySize {
                 expected: Self::PK_KEY_SIZE,
-                actual: blob.len(),
+                actual: identity.pk.len(),
                 key: "pk",
             });
         }
-        let mut pk = Vec::with_capacity(blob.len());
-        blob.read_to_end(&mut pk)?;
-        blob.close()?;
-
-        let mut blob = transaction.blob_open(rusqlite::MAIN_DB, "proteus_identities", "sk", row_id, true)?;
-        if blob.len() != Self::SK_KEY_SIZE {
+        if identity.sk.len() != Self::SK_KEY_SIZE {
             return Err(CryptoKeystoreError::InvalidKeySize {
                 expected: Self::SK_KEY_SIZE,
-                actual: blob.len(),
+                actual: identity.sk.len(),
                 key: "sk",
             });
         }
-        let mut sk = Vec::with_capacity(blob.len());
-        blob.read_to_end(&mut sk)?;
-        blob.close()?;
 
-        Ok(Some(Self { pk, sk }))
+        Ok(Some(identity))
     }
 
-    async fn count(conn: &mut Self::ConnectionType) -> crate::CryptoKeystoreResult<usize> {
+    /// Count the number of entities of this type in the database.
+    async fn count(conn: &mut Self::ConnectionType) -> CryptoKeystoreResult<u32> {
         let conn = conn.conn().await;
-        let count = conn.query_row("SELECT COUNT(*) FROM proteus_identities", [], |r| r.get(0))?;
+        let mut statement = conn.prepare_cached("SELECT COUNT(*) FROM proteus_identities")?;
+        let count = statement.query_one([], |row| row.get(0))?;
         // This should always be less or equal 1
         debug_assert!(count <= 1);
         Ok(count)
     }
-}
 
-#[async_trait::async_trait]
-impl EntityBase for ProteusIdentity {
-    type ConnectionType = KeystoreDatabaseConnection;
-    type AutoGeneratedFields = ();
-    const COLLECTION_NAME: &'static str = "proteus_identities";
-
-    fn to_missing_key_err_kind() -> MissingKeyErrorKind {
-        MissingKeyErrorKind::ProteusIdentity
-    }
-
-    fn to_transaction_entity(self) -> crate::transaction::dynamic_dispatch::Entity {
-        crate::transaction::dynamic_dispatch::Entity::ProteusIdentity(self)
+    /// Retrieve all entities of this type from the database.
+    async fn load_all(conn: &mut Self::ConnectionType) -> CryptoKeystoreResult<Vec<Self>> {
+        Self::get(conn, &())
+            .await
+            .map(|optional| optional.into_iter().collect())
     }
 }
 
 #[async_trait::async_trait]
-impl EntityTransactionExt for ProteusIdentity {
-    async fn save(&self, transaction: &TransactionWrapper<'_>) -> crate::CryptoKeystoreResult<()> {
-        use rusqlite::ToSql as _;
-        transaction.execute(
-            "INSERT INTO proteus_identities (sk, pk) VALUES (?, ?)",
-            [
-                rusqlite::blob::ZeroBlob(self.sk.len() as i32).to_sql()?,
-                rusqlite::blob::ZeroBlob(self.pk.len() as i32).to_sql()?,
-            ],
-        )?;
-
-        let row_id = transaction.last_insert_rowid();
-
-        use std::io::Write as _;
-        let mut blob = transaction.blob_open(rusqlite::MAIN_DB, "proteus_identities", "sk", row_id, false)?;
-        blob.write_all(&self.sk)?;
-        blob.close()?;
-
-        let mut blob = transaction.blob_open(rusqlite::MAIN_DB, "proteus_identities", "pk", row_id, false)?;
-        blob.write_all(&self.pk)?;
-        blob.close()?;
+impl<'a> EntityTransactionExt<'a> for ProteusIdentity {
+    type Transaction = TransactionWrapper<'a>;
+
+    /// Use the transaction's interface to save this entity to the database
+    async fn save(&self, tx: &Self::Transaction) -> CryptoKeystoreResult<()> {
+        let mut statement = tx.prepare_cached("SELECT COUNT(*) FROM proteus_identities")?;
+        let count = statement.query_one([], |row| row.get::<_, u32>(0))?;
+        if count > 0 {
+            return Err(CryptoKeystoreError::AlreadyExists("ProteusEntity"));
+        }
 
+        let mut statement = tx.prepare_cached("INSERT INTO proteus_identities (sk, pk) values (?, ?)")?;
+        statement.execute([&self.sk, &self.pk])?;
         Ok(())
     }
 
-    async fn delete_fail_on_missing_id(
-        transaction: &TransactionWrapper<'_>,
-        _id: StringEntityId<'_>,
-    ) -> crate::CryptoKeystoreResult<()> {
-        let row_id = transaction.query_row(
-            "SELECT rowid FROM proteus_identities ORDER BY rowid ASC LIMIT 1",
-            [],
-            |r| r.get::<_, i64>(0),
-        )?;
-        use rusqlite::ToSql as _;
-        transaction.execute("DELETE FROM proteus_identities WHERE rowid = ?", [row_id.to_sql()?])?;
-
-        Ok(())
+    /// Use the transaction's inteface to delete this entity from the database.
+    ///
+    /// Returns `true` if at least one entity was deleted, or `false` if the id was not found in the database.
+    async fn delete(tx: &Self::Transaction, _id: &()) -> CryptoKeystoreResult<bool> {
+        let mut statement = tx.prepare_cached("DELETE FROM proteus_identities")?;
+        let affected = statement.execute([])?;
+        Ok(affected > 0)
     }
 }