chore: Safe Swift refactor (#21)

reown-com · Sep 12, 2024 · aa97347 · aa97347
1 parent 5831cba
commit aa97347
Show file tree

Hide file tree

Showing 8 changed files with 433 additions and 385 deletions.
diff --git a/crates/ffi/src/account_client.rs b/crates/ffi/src/account_client.rs
@@ -83,6 +83,7 @@ impl FFIAccountClient {
                         config.chain_id.clone(),
                         config.config.into(),
                         service,
+                        config.safe,
                     );
                 account_client
             }

diff --git a/crates/ffi/src/lib.rs b/crates/ffi/src/lib.rs
@@ -47,6 +47,7 @@ mod ffi {
         pub chain_id: u64,
         pub config: FFIConfig,
         pub signer_type: String,
+        pub safe: bool,
     }
 
     enum FFIStringResult {

diff --git a/crates/yttrium/src/account_client.rs b/crates/yttrium/src/account_client.rs
@@ -1,6 +1,7 @@
 use crate::config::Config;
 use crate::private_key_service::PrivateKeyService;
 use crate::sign_service::SignService;
+use crate::transaction::send::safe_test;
 use crate::transaction::{send::send_transaction, Transaction};
 use alloy::primitives::Address;
 use alloy::signers::local::PrivateKeySigner;
@@ -47,6 +48,7 @@ pub struct AccountClient {
     chain_id: u64,
     config: Config,
     signer: Signer,
+    safe: bool,
 }
 
 impl AccountClient {
@@ -61,6 +63,7 @@ impl AccountClient {
             chain_id,
             config: config.clone(),
             signer: Signer::Native(Arc::new(Mutex::new(sign_service))),
+            safe: false,
         }
     }
 
@@ -69,6 +72,7 @@ impl AccountClient {
         chain_id: u64,
         config: Config,
         private_key_service: PrivateKeyService,
+        safe: bool,
     ) -> Self {
         Self {
             owner,
@@ -77,6 +81,7 @@ impl AccountClient {
             signer: Signer::PrivateKey(Arc::new(Mutex::new(
                 private_key_service,
             ))),
+            safe,
         }
     }
 
@@ -98,6 +103,7 @@ impl AccountClient {
             signer: Signer::PrivateKey(Arc::new(Mutex::new(
                 private_key_service,
             ))),
+            safe: false,
         }
     }
 
@@ -111,6 +117,7 @@ impl AccountClient {
             self.chain_id.clone(),
             self.config.clone(),
             self.signer.clone(),
+            self.safe,
         )
         .await
     }
@@ -136,6 +143,7 @@ impl AccountClient {
             self.chain_id.clone(),
             self.config.clone(),
             self.signer.clone(),
+            self.safe,
         )
         .await
     }
@@ -162,6 +170,7 @@ impl AccountClient {
             chain_id: 0,
             config: Config::local(),
             signer: Signer::Native(Arc::new(Mutex::new(SignService::mock()))),
+            safe: false,
         }
     }
 }
@@ -171,6 +180,7 @@ pub async fn get_address_with_signer(
     chain_id: u64,
     config: Config,
     signer: Signer,
+    safe: bool,
 ) -> eyre::Result<String> {
     match signer {
         Signer::PrivateKey(private_key_service) => {
@@ -185,6 +195,7 @@ pub async fn get_address_with_signer(
                 chain_id,
                 config,
                 private_key_signer,
+                safe,
             )
             .await
         }
@@ -199,11 +210,15 @@ pub async fn get_address_with_private_key_signer(
     chain_id: u64,
     config: Config,
     signer: PrivateKeySigner,
+    safe: bool,
 ) -> eyre::Result<String> {
     use crate::smart_accounts::simple_account::sender_address::get_sender_address_with_signer;
 
-    let sender_address =
-        get_sender_address_with_signer(config, chain_id, signer).await?;
+    let sender_address = if safe {
+        safe_test::get_address(signer, config).await?
+    } else {
+        get_sender_address_with_signer(config, chain_id, signer).await?
+    };
 
     Ok(sender_address.to_string())
 }
@@ -239,6 +254,7 @@ mod tests {
             chain_id,
             config,
             private_key_service,
+            false,
         );
 
         let transaction = Transaction::new_from_strings(
@@ -278,6 +294,7 @@ mod tests {
             chain_id,
             config,
             private_key_service,
+            false,
         );
 
         let sender_address = account_client.get_address().await?;

diff --git a/crates/yttrium/src/transaction/send.rs b/crates/yttrium/src/transaction/send.rs
@@ -1,11 +1,12 @@
+use crate::smart_accounts::safe::Execution;
 use crate::transaction::send::simple_account_test::send_transaction_with_signer;
 use crate::{
     config::Config, transaction::Transaction, user_operation::UserOperationV07,
 };
 use alloy::signers::local::PrivateKeySigner;
 use core::fmt;
 
-mod safe_test;
+pub mod safe_test;
 pub mod send_tests;
 pub mod simple_account_test;
 
@@ -50,6 +51,7 @@ pub async fn send_transaction(
     chain_id: u64,
     config: Config,
     signer: Signer,
+    safe: bool,
 ) -> eyre::Result<String> {
     match signer {
         Signer::PrivateKey(private_key_service) => {
@@ -65,6 +67,7 @@ pub async fn send_transaction(
                 chain_id,
                 config,
                 private_key_signer,
+                safe,
             )
             .await
         }
@@ -80,12 +83,27 @@ pub async fn send_transaction_with_private_key_signer(
     chain_id: u64,
     config: Config,
     private_key_signer: PrivateKeySigner,
+    safe: bool,
 ) -> eyre::Result<String> {
     let signer = private_key_signer;
 
-    let user_operation_hash =
+    let user_operation_hash = if safe {
+        safe_test::send_transaction(
+            vec![Execution {
+                target: transaction.to,
+                value: transaction.value,
+                callData: transaction.data,
+            }],
+            signer,
+            None,
+            None,
+            config,
+        )
+        .await?
+    } else {
         send_transaction_with_signer(transaction, config, chain_id, signer)
-            .await?;
+            .await?
+    };
 
     println!("user_operation_hash: {:?}", user_operation_hash);