Get an Association State from a list of IdentityUpdates in Validation…

… Service (#656)

* protos + get_association_state

* test for get_state

* panic tests for IdentityUpdates

Cargo.toml

@@ -50,3 +50,4 @@ tracing-subscriber = { version = "0.3.18", features = ["fmt", "env-filter"] }
 xmtp_cryptography = { path = "xmtp_cryptography" }
 xmtp_mls = { path = "xmtp_mls" }
 xmtp_proto = { path = "xmtp_proto" }
+xmtp_id = { path = "xmtp_id" }

mls_validation_service/Cargo.toml

@@ -16,18 +16,22 @@ openmls = { workspace = true }
 openmls_basic_credential = { workspace = true }
 openmls_rust_crypto = { workspace = true }
 openmls_traits = { workspace = true }
-prost = { version = "^0.12", features = ["prost-derive"] }
+prost = { workspace = true, features = ["prost-derive"] }
 serde = { workspace = true }
 tokio = { workspace = true, features = ["full"] }
 tonic = { workspace = true }
 warp = "0.3.6"
-xmtp_mls = { path = "../xmtp_mls" }
 xmtp_proto = { path = "../xmtp_proto", features = [
     "proto_full",
     "grpc",
     "tonic",
 ] }
+xmtp_id.workspace = true
+xmtp_mls.workspace = true 
+thiserror.workspace = true
 
 [dev-dependencies]
-ethers = "2.0.13"
+ethers.workspace = true 
 rand = { workspace = true }
+xmtp_id = { workspace = true, features = ["test-utils"]}
+anyhow.workspace = true

mls_validation_service/src/handlers.rs

@@ -5,14 +5,33 @@ use openmls::{
 use openmls_rust_crypto::RustCrypto;
 use tonic::{Request, Response, Status};
 
+use xmtp_id::associations::{self, try_map_vec, AssociationError, DeserializationError};
 use xmtp_mls::{utils::id::serialize_group_id, verified_key_package::VerifiedKeyPackage};
-use xmtp_proto::xmtp::mls_validation::v1::{
-    validate_group_messages_response::ValidationResponse as ValidateGroupMessageValidationResponse,
-    validate_key_packages_response::ValidationResponse as ValidateKeyPackageValidationResponse,
-    validation_api_server::ValidationApi, ValidateGroupMessagesRequest,
-    ValidateGroupMessagesResponse, ValidateKeyPackagesRequest, ValidateKeyPackagesResponse,
+use xmtp_proto::xmtp::{
+    identity::associations::IdentityUpdate as IdentityUpdateProto,
+    mls_validation::v1::{
+        validate_group_messages_response::ValidationResponse as ValidateGroupMessageValidationResponse,
+        validate_key_packages_response::ValidationResponse as ValidateKeyPackageValidationResponse,
+        validation_api_server::ValidationApi, GetAssociationStateRequest,
+        GetAssociationStateResponse, ValidateGroupMessagesRequest, ValidateGroupMessagesResponse,
+        ValidateKeyPackagesRequest, ValidateKeyPackagesResponse,
+    },
 };
 
+#[derive(Debug, thiserror::Error)]
+pub enum GrpcServerError {
+    #[error(transparent)]
+    Deserialization(#[from] DeserializationError),
+    #[error(transparent)]
+    Association(#[from] AssociationError),
+}
+
+impl From<GrpcServerError> for Status {
+    fn from(err: GrpcServerError) -> Self {
+        Status::invalid_argument(err.to_string())
+    }
+}
+
 #[derive(Debug, Default)]
 pub struct ValidationService {}
 
@@ -81,6 +100,49 @@ impl ValidationApi for ValidationService {
             responses: out,
         }))
     }
+
+    async fn get_association_state(
+        &self,
+        request: Request<GetAssociationStateRequest>,
+    ) -> Result<Response<GetAssociationStateResponse>, Status> {
+        let GetAssociationStateRequest {
+            old_updates,
+            new_updates,
+        } = request.into_inner();
+
+        get_association_state(old_updates, new_updates)
+            .map(Response::new)
+            .map_err(Into::into)
+    }
+}
+
+fn get_association_state(
+    old_updates: Vec<IdentityUpdateProto>,
+    new_updates: Vec<IdentityUpdateProto>,
+) -> Result<GetAssociationStateResponse, GrpcServerError> {
+    let (old_updates, new_updates) = (try_map_vec(old_updates)?, try_map_vec(new_updates)?);
+
+    if old_updates.is_empty() {
+        let new_state = associations::get_state(&new_updates)?;
+        return Ok(GetAssociationStateResponse {
+            association_state: Some(new_state.clone().into()),
+            state_diff: Some(new_state.as_diff().into()),
+        });
+    }
+
+    let old_state = associations::get_state(&old_updates)?;
+    let new_state = new_updates
+        .into_iter()
+        .try_fold(old_state.clone(), |state, update| {
+            associations::apply_update(state, update)
+        })?;
+
+    let state_diff = old_state.diff(&new_state);
+
+    Ok(GetAssociationStateResponse {
+        association_state: Some(new_state.into()),
+        state_diff: Some(state_diff.into()),
+    })
 }
 
 struct ValidateGroupMessageResult {
@@ -140,8 +202,10 @@ mod tests {
     use openmls_basic_credential::SignatureKeyPair;
     use openmls_rust_crypto::OpenMlsRustCrypto;
     use prost::Message;
+    use xmtp_id::associations::{generate_inbox_id, Action, CreateInbox, IdentityUpdate};
     use xmtp_mls::{credential::Credential, InboxOwner};
     use xmtp_proto::xmtp::{
+        identity::associations::IdentityUpdate as IdentityUpdateProto,
         mls::message_contents::MlsCredential as CredentialProto,
         mls_validation::v1::validate_key_packages_request::KeyPackage as KeyPackageProtoWrapper,
     };
@@ -252,4 +316,29 @@ mod tests {
         assert!(!first_response.is_ok);
         assert_eq!(first_response.account_address, "".to_string());
     }
+
+    // this test will panic until signature recovery is added
+    // and `MockSignature` is updated with signatures that can be recovered
+    #[tokio::test]
+    #[should_panic]
+    async fn test_get_association_state() {
+        let create_request = CreateInbox::default();
+        let inbox_id = generate_inbox_id(&create_request.account_address, &create_request.nonce);
+
+        let updates = vec![IdentityUpdate::new_test(
+            vec![Action::CreateInbox(create_request)],
+            inbox_id.clone(),
+        )];
+
+        ValidationService::default()
+            .get_association_state(Request::new(GetAssociationStateRequest {
+                old_updates: vec![],
+                new_updates: updates
+                    .into_iter()
+                    .map(IdentityUpdateProto::from)
+                    .collect::<Vec<_>>(),
+            }))
+            .await
+            .unwrap();
+    }
 }

xmtp_id/Cargo.toml

@@ -6,7 +6,6 @@ version = "0.1.0"
 # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
 
 [dependencies]
-anyhow.workspace = true
 async-trait.workspace = true
 chrono.workspace = true
 ethers.workspace = true
@@ -27,17 +26,11 @@ xmtp_cryptography.workspace = true
 xmtp_proto = { workspace = true, features = ["proto_full"] }
 
 [dev-dependencies]
-anyhow.workspace = true
 ctor = "0.2.5"
 ethers = { workspace = true, features = ["ws"] }
 futures = "0.3"
-jsonrpsee = { workspace = true, features = ["macros", "ws-client"] }
 regex = "1.10"
-serde_json.workspace = true
-surf = "2.3"
 tokio = { workspace = true, features = ["time"] }
-tokio-test = "0.4"
-tracing-subscriber.workspace = true
 
 [features]
 test-utils = []

xmtp_id/src/associations/association_log.rs

@@ -59,6 +59,7 @@ pub trait IdentityAction {
 }
 
 /// CreateInbox Action
+#[derive(Debug)]
 pub struct CreateInbox {
     pub nonce: u64,
     pub account_address: String,
@@ -100,6 +101,7 @@ impl IdentityAction for CreateInbox {
 }
 
 /// AddAssociation Action
+#[derive(Debug)]
 pub struct AddAssociation {
     pub new_member_signature: Box<dyn Signature>,
     pub new_member_identifier: MemberIdentifier,
@@ -196,6 +198,7 @@ impl IdentityAction for AddAssociation {
 }
 
 /// RevokeAssociation Action
+#[derive(Debug)]
 pub struct RevokeAssociation {
     pub recovery_address_signature: Box<dyn Signature>,
     pub revoked_member: MemberIdentifier,
@@ -248,6 +251,7 @@ impl IdentityAction for RevokeAssociation {
 }
 
 /// ChangeRecoveryAddress Action
+#[derive(Debug)]
 pub struct ChangeRecoveryAddress {
     pub recovery_address_signature: Box<dyn Signature>,
     pub new_recovery_address: String,
@@ -282,6 +286,7 @@ impl IdentityAction for ChangeRecoveryAddress {
 }
 
 /// All possible Action types that can be used inside an `IdentityUpdate`
+#[derive(Debug)]
 pub enum Action {
     CreateInbox(CreateInbox),
     AddAssociation(AddAssociation),
@@ -313,6 +318,7 @@ impl IdentityAction for Action {
 }
 
 /// An `IdentityUpdate` contains one or more Actions that can be applied to the AssociationState
+#[derive(Debug)]
 pub struct IdentityUpdate {
     pub inbox_id: String,
     pub client_timestamp_ns: u64,

xmtp_id/src/associations/hashes.rs

@@ -1,6 +1,6 @@
 use sha2::{Digest, Sha256};
 
-pub fn sha256_string(input: String) -> String {
+fn sha256_string(input: String) -> String {
     let mut hasher = Sha256::new();
     hasher.update(input.as_bytes());
     let result = hasher.finalize();

xmtp_id/src/associations/mod.rs

@@ -12,7 +12,7 @@ mod unsigned_actions;
 pub use self::association_log::*;
 pub use self::hashes::generate_inbox_id;
 pub use self::member::{Member, MemberIdentifier, MemberKind};
-pub use self::serialization::DeserializationError;
+pub use self::serialization::{map_vec, try_map_vec, DeserializationError};
 pub use self::signature::{Signature, SignatureError, SignatureKind};
 pub use self::state::{AssociationState, AssociationStateDiff};
 
@@ -25,8 +25,10 @@ pub fn apply_update(
 }
 
 // Get the current state from an array of `IdentityUpdate`s. Entire operation fails if any operation fails
-pub fn get_state(updates: Vec<IdentityUpdate>) -> Result<AssociationState, AssociationError> {
-    let new_state = updates.iter().try_fold(
+pub fn get_state<Updates: AsRef<[IdentityUpdate]>>(
+    updates: Updates,
+) -> Result<AssociationState, AssociationError> {
+    let new_state = updates.as_ref().iter().try_fold(
         None,
         |state, update| -> Result<Option<AssociationState>, AssociationError> {
             let updated_state = update.update_state(state)?;
@@ -52,7 +54,6 @@ pub mod test_defaults {
         fn default() -> Self {
             let existing_member = rand_string();
             let new_member = rand_vec();
-
             Self {
                 existing_member_signature: MockSignature::new_boxed(
                     true,
@@ -75,7 +76,6 @@ pub mod test_defaults {
     impl Default for CreateInbox {
         fn default() -> Self {
             let signer = rand_string();
-
             Self {
                 nonce: rand_u64(),
                 account_address: signer.clone(),
@@ -92,7 +92,6 @@ pub mod test_defaults {
     impl Default for RevokeAssociation {
         fn default() -> Self {
             let signer = rand_string();
-
             Self {
                 recovery_address_signature: MockSignature::new_boxed(
                     true,