From fcb717de430e85c616c7b57e92ad066a469054c4 Mon Sep 17 00:00:00 2001
From: Andrew Plaza <github@andrewplaza.dev>
Date: Thu, 28 Mar 2024 13:32:08 -0400
Subject: [PATCH] verification trait

---
 xmtp_id/src/credential.rs | 91 ++++++++++-----------------------------
 xmtp_id/src/lib.rs        | 14 +++---
 2 files changed, 29 insertions(+), 76 deletions(-)

diff --git a/xmtp_id/src/credential.rs b/xmtp_id/src/credential.rs
index be43862d2..d3c9bed7d 100644
--- a/xmtp_id/src/credential.rs
+++ b/xmtp_id/src/credential.rs
@@ -1,10 +1,10 @@
 use prost::Message;
 use xmtp_mls::{
-    credential::{GrantMessagingAccessAssociation, LegacyCreateIdentityAssociation},
+    credential::{Credential, GrantMessagingAccessAssociation, LegacyCreateIdentityAssociation},
     types::Address,
 };
 use xmtp_proto::xmtp::mls::message_contents::MlsCredential as CredentialProto;
-
+#[derive(Debug, Clone)]
 pub enum AssociationType {
     ExternallyOwned,
     SmartContract,
@@ -40,19 +40,21 @@ impl VerifiedCredential {
 }
 
 pub struct VerificationRequest {
-    expected_account_address: String,
     installation_public_key: Vec<u8>,
     credential: Vec<u8>,
 }
 
-type VerificationResult = Result<VerifiedCredential, VerificationError>;
-
-pub trait Credential {
-    fn address(&self) -> String;
-    fn installation_public_key(&self) -> Vec<u8>;
-    fn created_ns(&self) -> u64;
+impl VerificationRequest {
+    pub fn new(installation_public_key: Vec<u8>, credential: Vec<u8>) -> Self {
+        Self {
+            installation_public_key,
+            credential,
+        }
+    }
 }
 
+type VerificationResult = Result<VerifiedCredential, VerificationError>;
+
 #[async_trait::async_trait]
 pub trait CredentialVerifier {
     async fn verify_credential(request: VerificationRequest) -> VerificationResult;
@@ -68,68 +70,21 @@ pub trait CredentialVerifier {
     }
 }
 
-impl<'a> Credential for &'a GrantMessagingAccessAssociation {
-    fn address(&self) -> String {
-        self.account_address().clone()
-    }
-
-    fn installation_public_key(&self) -> Vec<u8> {
-        self.installation_public_key().clone()
-    }
-
-    fn created_ns(&self) -> u64 {
-        GrantMessagingAccessAssociation::created_ns(self)
-    }
-}
-
-fn validate_credential(
-    credential: impl Credential,
-    request: VerificationRequest,
-) -> Result<(), VerificationError> {
-    if credential.address() != request.expected_account_address {
-        return Err(VerificationError::AddressMismatch {
-            provided_addr: request.expected_account_address.to_string(),
-            signing_addr: credential.address(),
-        });
-    }
-
-    if credential.installation_public_key() != request.installation_public_key {
-        return Err(VerificationError::InstallationPublicKeyMismatch);
-    }
-
-    Ok(())
-}
-
-#[async_trait::async_trait]
-impl CredentialVerifier for GrantMessagingAccessAssociation {
-    async fn verify_credential(request: VerificationRequest) -> VerificationResult {
-        let proto = CredentialProto::decode(request.credential);
-        let credential = GrantMessagingAccessAssociation::from_proto_validated(
-            proto,
-            Some(request.installation_public_key),
-        );
-        validate_credential(&credential, request)?;
-
-        Ok(VerifiedCredential {
-            account_address: credential.account_address(),
-            account_type: AssociationType::EOA,
-        })
-    }
-}
-
 #[async_trait::async_trait]
-impl CredentialVerifier for LegacyCreateIdentityAssociation {
+impl CredentialVerifier for Credential {
     async fn verify_credential(request: VerificationRequest) -> VerificationResult {
         let proto = CredentialProto::decode(request.credential);
-        let credential = LegacyCreateIdentityAssociation::from_proto_validated(
-            proto,
-            Some(request.installation_public_key),
-        );
-        validate_credential(&credential, request)?;
-
-        Ok(VerifiedCredential {
-            account_address: credential.account_address(),
-            account_type: AssociationType::Legacy,
+        let credential =
+            Credential::from_proto_validated(proto, None, Some(request.installation_public_key));
+        Ok(match credential {
+            Credential::GrantMessagingAccess(cred) => VerifiedCredential {
+                account_address: cred.account_address(),
+                account_type: AssociationType::ExternallyOwned,
+            },
+            Credential::LegacyCreateIdentity(cred) => VerifiedCredential {
+                account_address: cred.account_address(),
+                account_type: AssociationType::Legacy,
+            },
         })
     }
 }
diff --git a/xmtp_id/src/lib.rs b/xmtp_id/src/lib.rs
index 7cf25a69d..3c4708c84 100644
--- a/xmtp_id/src/lib.rs
+++ b/xmtp_id/src/lib.rs
@@ -13,7 +13,10 @@ use xmtp_mls::{
 };
 use xmtp_proto::xmtp::mls::message_contents::MlsCredential as CredentialProto;
 
-use crate::error::IdentityError;
+use crate::{
+    credential::{CredentialVerifier, VerificationRequest, VerifiedCredential},
+    error::IdentityError,
+};
 
 pub struct Identity {
     #[allow(dead_code)]
@@ -66,13 +69,8 @@ impl Identity {
         credential: &[u8],
         installation_public_key: &[u8],
     ) -> Result<String, IdentityError> {
-        let proto = CredentialProto::decode(credential)?;
-        let credential = Credential::from_proto_validated(
-            proto,
-            None, // expected_account_address
-            Some(installation_public_key),
-        )?;
-
+        let request = VerificationRequest::new(credential, installation_public_key);
+        let credential = CredentialVerifier::verify_credential(request)?;
         Ok(credential.address())
     }
 }