feat(libp2p): multi signature schemes & session signing (#2599)

* feat(libp2p): multi signature schemes & session signing * clean up session sig scheme * fmt * chore * fix: tests * fix: session sig
dojoengine · Oct 31, 2024 · 66b76e6 · 66b76e6
1 parent a49b532
commit 66b76e6
Show file tree

Hide file tree

Showing 4 changed files with 77 additions and 35 deletions.
diff --git a/crates/torii/libp2p/src/errors.rs b/crates/torii/libp2p/src/errors.rs
@@ -4,6 +4,7 @@ use std::io;
 use libp2p::gossipsub::{PublishError, SubscriptionError};
 #[cfg(not(target_arch = "wasm32"))]
 use libp2p::noise;
+use starknet::providers::ProviderError;
 use thiserror::Error;
 
 #[derive(Error, Debug)]
@@ -45,4 +46,7 @@ pub enum Error {
 
     #[error("Invalid type provided: {0}")]
     InvalidTypeError(String),
+
+    #[error(transparent)]
+    ProviderError(#[from] ProviderError),
 }
diff --git a/crates/torii/libp2p/src/server/mod.rs b/crates/torii/libp2p/src/server/mod.rs
@@ -37,8 +37,8 @@ use crate::errors::Error;
 mod events;
 
 use crate::server::events::ServerEvent;
-use crate::typed_data::{parse_value_to_ty, PrimitiveType, TypedData};
-use crate::types::Message;
+use crate::typed_data::{encode_type, parse_value_to_ty, PrimitiveType, TypedData};
+use crate::types::{Message, Signature};
 
 pub(crate) const LOG_TARGET: &str = "torii::relay::server";
 
@@ -302,37 +302,15 @@ impl<P: Provider + Sync> Relay<P> {
                             // to prevent replay attacks.
 
                             // Verify the signature
-                            let message_hash =
-                                if let Ok(message) = data.message.encode(entity_identity) {
-                                    message
-                                } else {
-                                    info!(
-                                        target: LOG_TARGET,
-                                        "Encoding message."
-                                    );
-                                    continue;
-                                };
-
-                            let mut calldata = vec![message_hash];
-                            calldata.push(Felt::from(data.signature.len()));
-
-                            calldata.extend(data.signature);
-                            if !match self
-                                .provider
-                                .call(
-                                    FunctionCall {
-                                        contract_address: entity_identity,
-                                        entry_point_selector: get_selector_from_name(
-                                            "is_valid_signature",
-                                        )
-                                        .unwrap(),
-                                        calldata,
-                                    },
-                                    BlockId::Tag(BlockTag::Pending),
-                                )
-                                .await
+                            if !match validate_signature(
+                                &self.provider,
+                                entity_identity,
+                                &data.message,
+                                &data.signature,
+                            )
+                            .await
                             {
-                                Ok(res) => res[0] != Felt::ZERO,
+                                Ok(res) => res,
                                 Err(e) => {
                                     warn!(
                                         target: LOG_TARGET,
@@ -432,6 +410,57 @@ impl<P: Provider + Sync> Relay<P> {
     }
 }
 
+async fn validate_signature<P: Provider + Sync>(
+    provider: &P,
+    entity_identity: Felt,
+    message: &TypedData,
+    signature: &Signature,
+) -> Result<bool, Error> {
+    let message_hash = message.encode(entity_identity)?;
+
+    match signature {
+        Signature::Account(signature) => {
+            let mut calldata = vec![message_hash, Felt::from(signature.len())];
+            calldata.extend(signature);
+            provider
+                .call(
+                    FunctionCall {
+                        contract_address: entity_identity,
+                        entry_point_selector: get_selector_from_name("is_valid_signature").unwrap(),
+                        calldata,
+                    },
+                    BlockId::Tag(BlockTag::Pending),
+                )
+                .await
+                .map_err(Error::ProviderError)
+                .map(|res| res[0] != Felt::ZERO)
+        }
+        Signature::Session(signature) => {
+            let mut calldata = vec![
+                Felt::ONE,
+                get_selector_from_name(&encode_type(&message.primary_type, &message.types)?)
+                    .map_err(|e| Error::InvalidMessageError(e.to_string()))?,
+                message_hash,
+                signature.len().into(),
+            ];
+            calldata.extend(signature);
+            provider
+                .call(
+                    FunctionCall {
+                        contract_address: entity_identity,
+                        entry_point_selector: get_selector_from_name("is_session_sigature_valid")
+                            .unwrap(),
+                        calldata,
+                    },
+                    BlockId::Tag(BlockTag::Pending),
+                )
+                .await
+                .map_err(Error::ProviderError)
+                .map(|res| res[0] != Felt::ZERO)
+        }
+    }
+}
+
 fn ty_keys(ty: &Ty) -> Result<Vec<Felt>, Error> {
     if let Ty::Struct(s) = &ty {
         let mut keys = Vec::new();

diff --git a/crates/torii/libp2p/src/tests.rs b/crates/torii/libp2p/src/tests.rs
@@ -545,7 +545,7 @@ mod test {
 
         use crate::server::Relay;
         use crate::typed_data::{Domain, Field, SimpleField, TypedData};
-        use crate::types::Message;
+        use crate::types::{Message, Signature};
 
         let _ = tracing_subscriber::fmt()
             .with_env_filter("torii::relay::client=debug,torii::relay::server=debug")
@@ -683,7 +683,10 @@ mod test {
 
         client
             .command_sender
-            .publish(Message { message: typed_data, signature: vec![signature.r, signature.s] })
+            .publish(Message {
+                message: typed_data,
+                signature: Signature::Account(vec![signature.r, signature.s]),
+            })
             .await?;
 
         sleep(std::time::Duration::from_secs(2)).await;

diff --git a/crates/torii/libp2p/src/types.rs b/crates/torii/libp2p/src/types.rs
@@ -3,8 +3,14 @@ use starknet::core::types::Felt;
 
 use crate::typed_data::TypedData;
 
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub enum Signature {
+    Account(Vec<Felt>),
+    Session(Vec<Felt>),
+}
+
 #[derive(Debug, Clone, Serialize, Deserialize)]
 pub struct Message {
     pub message: TypedData,
-    pub signature: Vec<Felt>,
+    pub signature: Signature,
 }