feat(integrtity): introduce zero cost wrapper Verified

server/crates/arbiter-server/src/crypto/integrity/v1.rs

@@ -1,9 +1,9 @@
-use crate::{
-    actors::vault::{self, GetState},
-    crypto::integrity::hashing::Hashable,
-};
+use crate::{actors::vault::{self, GetState}, crypto::integrity::hashing::Hashable};
 use hmac::Hmac;
 use sha2::Sha256;
+use std::future::Future;
+use std::ops::Deref;
+use std::pin::Pin;
 
 use diesel::{ExpressionMethods as _, QueryDsl, dsl::insert_into, sqlite::Sqlite};
 use diesel_async::{AsyncConnection, RunQueryDsl};
@@ -11,16 +11,23 @@ use kameo::{actor::ActorRef, error::SendError};
 use sha2::Digest as _;
 
 pub mod hashing;
+pub mod verified;
 
 use crate::{
     actors::vault::{SignIntegrity, Vault, VerifyIntegrity},
     db::{
         self,
-        models::{IntegrityEnvelope, NewIntegrityEnvelope},
+        models::{IntegrityEnvelope as IntegrityEnvelopeRow, NewIntegrityEnvelope},
         schema::integrity_envelope,
     },
 };
 
+pub const CURRENT_PAYLOAD_VERSION: i32 = 1;
+pub const INTEGRITY_SUBKEY_TAG: &[u8] = b"arbiter/db-integrity-key/v1";
+
+pub type HmacSha256 = Hmac<Sha256>;
+pub use self::verified::{Nested, VerificationOrigin, Verified};
+
 #[derive(Debug, thiserror::Error)]
 pub enum Error {
     #[error("Database error: {0}")]
@@ -49,71 +56,90 @@ pub enum Error {
 }
 
 #[derive(Debug, Clone, Copy, PartialEq, Eq)]
+#[must_use]
 pub enum AttestationStatus {
     Attested,
     Unavailable,
 }
 
-pub const CURRENT_PAYLOAD_VERSION: i32 = 1;
-pub const INTEGRITY_SUBKEY_TAG: &[u8] = b"arbiter/db-integrity-key/v1";
-
-pub type HmacSha256 = Hmac<Sha256>;
-
 pub trait Integrable: Hashable {
     const KIND: &'static str;
     const VERSION: i32 = 1;
 }
 
-fn payload_hash(payload: &impl Hashable) -> [u8; 32] {
-    let mut hasher = Sha256::new();
-    payload.hash(&mut hasher);
-    hasher.finalize().into()
+impl<T: Integrable> Integrable for &T {
+    const KIND: &'static str = T::KIND;
+    const VERSION: i32 = T::VERSION;
 }
 
-fn push_len_prefixed(out: &mut Vec<u8>, bytes: &[u8]) {
-    out.extend_from_slice(&(bytes.len() as u32).to_be_bytes());
-    out.extend_from_slice(bytes);
-}
+#[derive(Debug, Clone)]
+pub struct EntityId(Vec<u8>);
 
-fn build_mac_input(
-    entity_kind: &str,
-    entity_id: &[u8],
-    payload_version: i32,
-    payload_hash: &[u8; 32],
-) -> Vec<u8> {
-    let mut out = Vec::with_capacity(8 + entity_kind.len() + entity_id.len() + 32);
-    push_len_prefixed(&mut out, entity_kind.as_bytes());
-    push_len_prefixed(&mut out, entity_id);
-    out.extend_from_slice(&payload_version.to_be_bytes());
-    out.extend_from_slice(payload_hash);
-    out
-}
+impl Deref for EntityId {
+    type Target = [u8];
 
-pub trait IntoId {
-    fn into_id(self) -> Vec<u8>;
-}
-
-impl IntoId for i32 {
-    fn into_id(self) -> Vec<u8> {
-        self.to_be_bytes().to_vec()
+    fn deref(&self) -> &Self::Target {
+        &self.0
     }
 }
 
-impl IntoId for &'_ [u8] {
-    fn into_id(self) -> Vec<u8> {
-        self.to_vec()
+impl From<i32> for EntityId {
+    fn from(value: i32) -> Self {
+        Self(value.to_be_bytes().to_vec())
     }
 }
 
-pub async fn sign_entity<E: Integrable>(
+impl From<&'_ [u8]> for EntityId {
+    fn from(bytes: &'_ [u8]) -> Self {
+        Self(bytes.to_vec())
+    }
+}
+
+pub async fn lookup_verified<E, Id, C, F, Fut>(
+    conn: &mut C,
+    vault: &ActorRef<Vault>,
+    entity_id: Id,
+    load: F,
+) -> Result<VerifiedEntity<E, Id>, Error>
+where
+    C: AsyncConnection<Backend = Sqlite>,
+    E: Integrable,
+    Id: Into<EntityId> + Clone,
+    F: FnOnce(&mut C) -> Fut,
+    Fut: Future<Output = Result<E, db::DatabaseError>>,
+{
+    let entity = load(conn).await?;
+    verify_entity(conn, vault, entity, entity_id).await
+}
+
+pub async fn lookup_verified_from_query<E, Id, C, F>(
+    conn: &mut C,
+    vault: &ActorRef<Vault>,
+    load: F,
+) -> Result<VerifiedEntity<E, Id>, Error>
+where
+    C: AsyncConnection<Backend = Sqlite> + Send,
+    E: Integrable,
+    Id: Into<EntityId> + Clone,
+    F: for<'a> FnOnce(
+        &'a mut C,
+    ) -> Pin<
+        Box<dyn Future<Output = Result<(Id, E), db::DatabaseError>> + Send + 'a>,
+    >,
+{
+    let (entity_id, entity) = load(conn).await?;
+    verify_entity(conn, vault, entity, entity_id).await
+}
+
+pub async fn sign_entity<E: Integrable, Id: Into<EntityId> + Clone>(
     conn: &mut impl AsyncConnection<Backend = Sqlite>,
     vault: &ActorRef<Vault>,
     entity: &E,
-    entity_id: impl IntoId,
-) -> Result<(), Error> {
-    let payload_hash = payload_hash(&entity);
+    as_entity_id: Id,
+) -> Result<Verified<Id, Nested<E>>, Error> {
+    let payload_hash = payload_hash(entity);
 
-    let entity_id = entity_id.into_id();
+    let entity_id = as_entity_id.clone().into();
 
     let mac_input = build_mac_input(E::KIND, &entity_id, E::VERSION, &payload_hash);
 
@@ -129,7 +155,7 @@ pub async fn sign_entity<E: Integrable>(
     insert_into(integrity_envelope::table)
         .values(NewIntegrityEnvelope {
             entity_kind: E::KIND.to_owned(),
-            entity_id,
+            entity_id: entity_id.to_vec(),
             payload_version: E::VERSION,
             key_version,
             mac: mac.to_vec(),
@@ -148,19 +174,19 @@ pub async fn sign_entity<E: Integrable>(
         .await
         .map_err(db::DatabaseError::from)?;
 
-    Ok(())
+    Ok(Verified::<Id, Nested<E>>::new(as_entity_id))
 }
 
-pub async fn verify_entity<E: Integrable>(
+pub async fn check_entity_attestation<E: Integrable>(
     conn: &mut impl AsyncConnection<Backend = Sqlite>,
     vault: &ActorRef<Vault>,
     entity: &E,
-    entity_id: impl IntoId,
+    entity_id: impl Into<EntityId>,
 ) -> Result<AttestationStatus, Error> {
-    let entity_id = entity_id.into_id();
-    let envelope: IntegrityEnvelope = integrity_envelope::table
+    let entity_id = entity_id.into();
+    let envelope: IntegrityEnvelopeRow = integrity_envelope::table
         .filter(integrity_envelope::entity_kind.eq(E::KIND))
-        .filter(integrity_envelope::entity_id.eq(&entity_id))
+        .filter(integrity_envelope::entity_id.eq(&*entity_id))
         .first(conn)
         .await
         .map_err(|err| match err {
@@ -178,7 +204,7 @@ pub async fn verify_entity<E: Integrable>(
         });
     }
 
-    let payload_hash = payload_hash(&entity);
+    let payload_hash = payload_hash(entity);
     let mac_input = build_mac_input(E::KIND, &entity_id, envelope.payload_version, &payload_hash);
 
     let result = vault
@@ -194,24 +220,111 @@ pub async fn verify_entity<E: Integrable>(
         Ok(false) => Err(Error::MacMismatch {
             entity_kind: E::KIND,
         }),
-        Err(SendError::HandlerError(vault::Error::Sealed)) => {
-            Ok(AttestationStatus::Unavailable)
-        }
+        Err(SendError::HandlerError(vault::Error::Sealed)) => Ok(AttestationStatus::Unavailable),
         Err(_) => Err(Error::VaultSend),
     }
 }
 
+#[derive(Debug, Clone)]
+#[repr(C)]
+pub struct VerifiedEntity<E, Id> {
+    pub entity: Verified<E>,
+    pub entity_id: Verified<Id, Nested<E>>,
+}
+
+impl<E, Id> Deref for VerifiedEntity<E, Id> {
+    type Target = Verified<E>;
+
+    fn deref(&self) -> &Self::Target {
+        &self.entity
+    }
+}
+
+pub async fn verify_entity<E: Integrable, Id: Into<EntityId> + Clone>(
+    conn: &mut impl AsyncConnection<Backend = Sqlite>,
+    vault: &ActorRef<Vault>,
+    entity: E,
+    entity_id: Id,
+) -> Result<VerifiedEntity<E, Id>, Error> {
+    match check_entity_attestation(conn, vault, &entity, entity_id.clone()).await? {
+        AttestationStatus::Attested => Ok(VerifiedEntity {
+            entity: Verified::new(entity),
+            entity_id: Verified::new(entity_id),
+        }),
+        AttestationStatus::Unavailable => Err(Error::Vault(vault::Error::Sealed)),
+    }
+}
+
+pub async fn verify_entity_ref<'e, E: Integrable, Id: Into<EntityId> + Clone>(
+    conn: &mut impl AsyncConnection<Backend = Sqlite>,
+    vault: &ActorRef<Vault>,
+    entity: &'e E,
+    entity_id: Id,
+) -> Result<Verified<VerifiedEntity<&'e E, Id>, Nested<E>>, Error> {
+    match check_entity_attestation(conn, vault, entity, entity_id.clone()).await? {
+        AttestationStatus::Attested => Ok(Verified::<VerifiedEntity<&'e E, Id>, Nested<E>>::new(
+            VerifiedEntity {
+                entity: Verified::new(entity),
+                entity_id: Verified::new(entity_id),
+            },
+        )),
+        AttestationStatus::Unavailable => Err(Error::Vault(vault::Error::Sealed)),
+    }
+}
+
+pub async fn delete_envelope<E: Integrable>(
+    conn: &mut impl AsyncConnection<Backend = Sqlite>,
+    entity_id: impl Into<EntityId>,
+) -> Result<usize, Error> {
+    let entity_id = entity_id.into();
+
+    let affected = diesel::delete(
+        integrity_envelope::table
+            .filter(integrity_envelope::entity_kind.eq(E::KIND))
+            .filter(integrity_envelope::entity_id.eq(&*entity_id)),
+    )
+    .execute(conn)
+    .await
+    .map_err(db::DatabaseError::from)?;
+
+    Ok(affected)
+}
+
 pub async fn is_signing_available(vault: &ActorRef<Vault>) -> Result<bool, Error> {
     let state = vault.ask(GetState).await.map_err(|_| Error::VaultSend)?;
     Ok(matches!(state, vault::VaultState::Unsealed))
 }
 
+fn payload_hash(payload: &impl Hashable) -> [u8; 32] {
+    let mut hasher = Sha256::new();
+    payload.hash(&mut hasher);
+    hasher.finalize().into()
+}
+
+fn build_mac_input(
+    entity_kind: &str,
+    entity_id: &[u8],
+    payload_version: i32,
+    payload_hash: &[u8; 32],
+) -> Vec<u8> {
+    let mut out = Vec::with_capacity(8 + entity_kind.len() + entity_id.len() + 32);
+    push_len_prefixed(&mut out, entity_kind.as_bytes());
+    push_len_prefixed(&mut out, entity_id);
+    out.extend_from_slice(&payload_version.to_be_bytes());
+    out.extend_from_slice(payload_hash);
+    out
+}
+
+fn push_len_prefixed(out: &mut Vec<u8>, bytes: &[u8]) {
+    out.extend_from_slice(&(bytes.len() as u32).to_be_bytes());
+    out.extend_from_slice(bytes);
+}
+
 #[cfg(test)]
 mod tests {
     use diesel::{ExpressionMethods as _, QueryDsl};
     use diesel_async::RunQueryDsl;
     use kameo::{actor::ActorRef, prelude::Spawn};
-
     use sha2::Digest;
 
     use crate::{
@@ -224,7 +337,7 @@ mod tests {
     use arbiter_crypto::safecell::{SafeCell, SafeCellHandle as _};
 
     use super::hashing::Hashable;
-    use super::{Error, Integrable, sign_entity, verify_entity};
+    use super::{Error, Integrable, check_entity_attestation, sign_entity};
 
     #[derive(Clone)]
     struct DummyEntity {
@@ -272,7 +385,8 @@ mod tests {
 
         sign_entity(&mut conn, &vault, &entity, ENTITY_ID)
             .await
-            .unwrap();
+            .unwrap()
+            .drop_verification_provenance();
 
         let count: i64 = schema::integrity_envelope::table
             .filter(schema::integrity_envelope::entity_kind.eq("dummy_entity"))
@@ -283,9 +397,11 @@ mod tests {
             .unwrap();
 
         assert_eq!(count, 1, "envelope row must be created exactly once");
-        verify_entity(&mut conn, &vault, &entity, ENTITY_ID)
+
+        let status = check_entity_attestation(&mut conn, &vault, &entity, ENTITY_ID)
             .await
             .unwrap();
+        assert!(matches!(status, super::AttestationStatus::Attested));
     }
 
     #[tokio::test]
@@ -303,7 +419,8 @@ mod tests {
 
         sign_entity(&mut conn, &vault, &entity, ENTITY_ID)
             .await
-            .unwrap();
+            .unwrap()
+            .drop_verification_provenance();
 
         diesel::update(schema::integrity_envelope::table)
             .filter(schema::integrity_envelope::entity_kind.eq("dummy_entity"))
@@ -313,35 +430,7 @@ mod tests {
             .await
             .unwrap();
 
-        let err = verify_entity(&mut conn, &vault, &entity, ENTITY_ID)
-            .await
-            .unwrap_err();
-        assert!(matches!(err, Error::MacMismatch { .. }));
-    }
-
-    #[tokio::test]
-    async fn changed_payload_fails_verification() {
-        let db = db::create_test_pool().await;
-        let vault = bootstrapped_vault(&db).await;
-        let mut conn = db.get().await.unwrap();
-
-        const ENTITY_ID: &[u8] = b"entity-id-21";
-
-        let entity = DummyEntity {
-            payload_version: 1,
-            payload: b"payload-v1".to_vec(),
-        };
-
-        sign_entity(&mut conn, &vault, &entity, ENTITY_ID)
-            .await
-            .unwrap();
-
-        let tampered = DummyEntity {
-            payload: b"payload-v1-but-tampered".to_vec(),
-            ..entity
-        };
-
-        let err = verify_entity(&mut conn, &vault, &tampered, ENTITY_ID)
+        let err = check_entity_attestation(&mut conn, &vault, &entity, ENTITY_ID)
             .await
             .unwrap_err();
         assert!(matches!(err, Error::MacMismatch { .. }));