From 0d90c180f84e0bc2b5ef12963de2b45a83c00931 Mon Sep 17 00:00:00 2001 From: Sameh Abouelsaad Date: Mon, 12 May 2025 09:29:50 +0300 Subject: [PATCH] feat: Enhance Ethereum wallet management and improve code structure --- .../hero_vault/agung_send_transaction.rhai | 22 +- src/hero_vault/error.rs | 4 + src/hero_vault/ethereum/storage.rs | 196 +++++++++++++++++- src/hero_vault/ethereum/tests/wallet_tests.rs | 107 +++++----- src/hero_vault/kvs/mod.rs | 6 - src/rhai/hero_vault.rs | 156 ++++++-------- 6 files changed, 313 insertions(+), 178 deletions(-) diff --git a/examples/hero_vault/agung_send_transaction.rhai b/examples/hero_vault/agung_send_transaction.rhai index 8d063fd..916d9ee 100644 --- a/examples/hero_vault/agung_send_transaction.rhai +++ b/examples/hero_vault/agung_send_transaction.rhai @@ -1,4 +1,4 @@ -// Script to create an Agung wallet from a private key and send tokens +// Script to create an Ethereum wallet from a private key and send tokens on the Agung network // This script demonstrates how to create a wallet from a private key and send tokens // Define the private key and recipient address @@ -33,32 +33,32 @@ if !select_keypair("demo_keypair") { print("\nCreated and selected keypair successfully"); -// Clear any existing Agung wallets to avoid conflicts -if clear_wallets_for_network("agung") { - print("Cleared existing Agung wallets"); +// Clear any existing Ethereum wallets to avoid conflicts +if clear_ethereum_wallets() { + print("Cleared existing Ethereum wallets"); } else { - print("Failed to clear existing Agung wallets"); + print("Failed to clear existing Ethereum wallets"); return; } // Create a wallet from the private key directly print("\n=== Creating Wallet from Private Key ==="); -// Create a wallet from the private key for the Agung network -if create_wallet_from_private_key_for_network(private_key, "agung") { - print("Successfully created wallet from private key for Agung network"); +// Create a wallet from the private key (works for any network) +if create_ethereum_wallet_from_private_key(private_key) { + print("Successfully created wallet from private key"); // Get the wallet address - let wallet_address = get_wallet_address_for_network("agung"); + let wallet_address = get_ethereum_address(); print(`Wallet address: ${wallet_address}`); // Create a provider for the Agung network - let provider_id = create_agung_provider(); + let provider_id = create_provider("agung"); if provider_id != "" { print("Successfully created Agung provider"); // Check the wallet balance first - let wallet_address = get_wallet_address_for_network("agung"); + let wallet_address = get_ethereum_address(); let balance_wei = get_balance("agung", wallet_address); if balance_wei == "" { diff --git a/src/hero_vault/error.rs b/src/hero_vault/error.rs index 2cf41f1..2f7f5c1 100644 --- a/src/hero_vault/error.rs +++ b/src/hero_vault/error.rs @@ -48,6 +48,10 @@ pub enum CryptoError { /// Smart contract error #[error("Smart contract error: {0}")] ContractError(String), + + /// Storage error + #[error("Storage error: {0}")] + StorageError(String), } /// Convert CryptoError to SAL's Error type diff --git a/src/hero_vault/ethereum/storage.rs b/src/hero_vault/ethereum/storage.rs index c1d2aa1..9f2a8d3 100644 --- a/src/hero_vault/ethereum/storage.rs +++ b/src/hero_vault/ethereum/storage.rs @@ -3,16 +3,190 @@ use std::sync::Mutex; use std::collections::HashMap; use once_cell::sync::Lazy; +use serde::{Serialize, Deserialize}; +use cfg_if::cfg_if; +use tokio::runtime::Runtime; +use ethers::types::Address; +use std::str::FromStr; use crate::hero_vault::error::CryptoError; +use crate::hero_vault::kvs::{self, KVStore, DefaultStore}; use super::wallet::EthereumWallet; use super::networks; -/// Global storage for Ethereum wallets. +/// Ethereum wallet data storage key in KVStore +const ETH_WALLET_STORAGE_KEY: &str = "ethereum/wallets"; + +/// Global fallback storage for Ethereum wallets (used when KVStore is unavailable) static ETH_WALLETS: Lazy>> = Lazy::new(|| { Mutex::new(Vec::new()) }); +// Global Tokio runtime for blocking async operations +static RUNTIME: Lazy> = Lazy::new(|| { + Mutex::new(Runtime::new().expect("Failed to create Tokio runtime")) +}); + +/// Serializable representation of an Ethereum wallet +#[derive(Debug, Clone, Serialize, Deserialize)] +struct EthereumWalletStorage { + /// Ethereum address string + address: String, + /// Private key in hex format + private_key: String, + /// Optional wallet name + name: Option, +} + +impl From<&EthereumWallet> for EthereumWalletStorage { + fn from(wallet: &EthereumWallet) -> Self { + Self { + address: wallet.address_string(), + private_key: wallet.private_key_hex(), + name: wallet.name.clone(), + } + } +} + +impl TryFrom<&EthereumWalletStorage> for EthereumWallet { + type Error = CryptoError; + + fn try_from(storage: &EthereumWalletStorage) -> Result { + let wallet = EthereumWallet::from_private_key(&storage.private_key)?; + + // If the address doesn't match, something is wrong + if wallet.address_string() != storage.address { + return Err(CryptoError::InvalidAddress(format!( + "Address mismatch: expected {}, got {}", + storage.address, wallet.address_string() + ))); + } + + // Set the name if present + let wallet_with_name = if let Some(name) = &storage.name { + EthereumWallet { + name: Some(name.clone()), + ..wallet + } + } else { + wallet + }; + + Ok(wallet_with_name) + } +} + +/// Helper function to get the platform-specific storage implementation +fn get_wallet_store() -> Result { + cfg_if! { + if #[cfg(target_arch = "wasm32")] { + // For WebAssembly, we need to handle the async nature of IndexedDB + // We'll use a blocking approach for API consistency + use wasm_bindgen_futures::spawn_local; + + // We need to use the runtime to block_on the async operations + let rt = RUNTIME.lock().unwrap(); + rt.block_on(async { + match kvs::open_default_store("ethereum-wallets", None).await { + Ok(store) => Ok(store), + Err(e) => { + log::warn!("Failed to open IndexedDB store: {}", e); + // Try to create the store if opening failed + kvs::create_default_store("ethereum-wallets", false, None).await + .map_err(|e| CryptoError::StorageError(e.to_string())) + } + } + }) + } else { + // For native platforms, we can use SlateDB directly + match kvs::open_default_store("ethereum-wallets", None) { + Ok(store) => Ok(store), + Err(e) => { + log::warn!("Failed to open SlateDB store: {}", e); + // Try to create the store if opening failed + kvs::create_default_store("ethereum-wallets", false, None) + .map_err(|e| CryptoError::StorageError(e.to_string())) + } + } + } + } +} + +/// Save wallets to persistent storage +fn save_wallets(wallets: &[EthereumWallet]) -> Result<(), CryptoError> { + // Convert wallets to serializable format + let storage_wallets: Vec = wallets + .iter() + .map(|w| EthereumWalletStorage::from(w)) + .collect(); + + // Try to use the KVStore implementation first + let store_result = get_wallet_store() + .and_then(|store| { + let json = serde_json::to_string(&storage_wallets) + .map_err(|e| CryptoError::StorageError(e.to_string()))?; + store.set(ETH_WALLET_STORAGE_KEY, &json) + .map_err(|e| CryptoError::StorageError(e.to_string())) + }); + + // Log warning if storage failed but don't fail the operation + if let Err(e) = &store_result { + log::warn!("Failed to save wallets to persistent storage: {}", e); + } + + // Always update the in-memory fallback + let mut mem_wallets = ETH_WALLETS.lock().unwrap(); + *mem_wallets = wallets.to_vec(); + + // Return the result of the persistent storage operation + store_result +} + +/// Load wallets from persistent storage +fn load_wallets() -> Vec { + // Try to load from KVStore first + let store_result = get_wallet_store() + .and_then(|store| { + store.get::<_, String>(ETH_WALLET_STORAGE_KEY) + .map_err(|e| CryptoError::StorageError(e.to_string())) + }) + .and_then(|json| { + serde_json::from_str::>(&json) + .map_err(|e| CryptoError::StorageError(format!("Failed to parse wallet JSON: {}", e))) + }); + + match store_result { + Ok(storage_wallets) => { + // Convert from storage format to EthereumWallet + let wallets_result: Result, CryptoError> = storage_wallets + .iter() + .map(|sw| EthereumWallet::try_from(sw)) + .collect(); + + match wallets_result { + Ok(wallets) => { + // Also update the in-memory fallback + let mut mem_wallets = ETH_WALLETS.lock().unwrap(); + *mem_wallets = wallets.clone(); + wallets + }, + Err(e) => { + log::error!("Failed to convert wallets from storage format: {}", e); + // Fall back to in-memory storage + let mem_wallets = ETH_WALLETS.lock().unwrap(); + mem_wallets.clone() + } + } + }, + Err(e) => { + log::warn!("Failed to load wallets from persistent storage: {}", e); + // Fall back to in-memory storage + let mem_wallets = ETH_WALLETS.lock().unwrap(); + mem_wallets.clone() + } + } +} + /// Creates an Ethereum wallet from the currently selected keypair. pub fn create_ethereum_wallet() -> Result { // Get the currently selected keypair @@ -22,8 +196,9 @@ pub fn create_ethereum_wallet() -> Result { let wallet = EthereumWallet::from_keypair(&keypair)?; // Store the wallet - let mut wallets = ETH_WALLETS.lock().unwrap(); + let mut wallets = load_wallets(); wallets.push(wallet.clone()); + save_wallets(&wallets)?; Ok(wallet) } @@ -37,8 +212,9 @@ pub fn create_ethereum_wallet_from_name(name: &str) -> Result Result Result { - let wallets = ETH_WALLETS.lock().unwrap(); + let wallets = load_wallets(); if wallets.is_empty() { return Err(CryptoError::NoKeypairSelected); @@ -68,8 +245,13 @@ pub fn get_current_ethereum_wallet() -> Result { /// Clears all Ethereum wallets. pub fn clear_ethereum_wallets() { - let mut wallets = ETH_WALLETS.lock().unwrap(); - wallets.clear(); + // Clear both persistent and in-memory storage + if let Ok(store) = get_wallet_store() { + let _ = store.delete::<&str>(ETH_WALLET_STORAGE_KEY); + } + + let mut mem_wallets = ETH_WALLETS.lock().unwrap(); + mem_wallets.clear(); } // Legacy functions for backward compatibility diff --git a/src/hero_vault/ethereum/tests/wallet_tests.rs b/src/hero_vault/ethereum/tests/wallet_tests.rs index 9fafb11..7b78cec 100644 --- a/src/hero_vault/ethereum/tests/wallet_tests.rs +++ b/src/hero_vault/ethereum/tests/wallet_tests.rs @@ -32,11 +32,11 @@ fn test_ethereum_wallet_from_name_and_keypair() { // Creating another wallet with the same name and keypair should yield the same address let wallet2 = EthereumWallet::from_name_and_keypair("test", &keypair).unwrap(); - assert_eq!(wallet.address, wallet2.address); + assert_eq!(wallet.address_string(), wallet2.address_string()); // Creating a wallet with a different name should yield a different address let wallet3 = EthereumWallet::from_name_and_keypair("test2", &keypair).unwrap(); - assert_ne!(wallet.address, wallet3.address); + assert_ne!(wallet.address_string(), wallet3.address_string()); } #[test] @@ -53,7 +53,7 @@ fn test_ethereum_wallet_from_private_key() { // The address should be deterministic based on the private key let wallet2 = EthereumWallet::from_private_key(private_key).unwrap(); - assert_eq!(wallet.address, wallet2.address); + assert_eq!(wallet.address_string(), wallet2.address_string()); } #[test] @@ -72,7 +72,7 @@ fn test_wallet_management() { let current_wallet = get_current_ethereum_wallet().unwrap(); // Check that they match - assert_eq!(wallet.address, current_wallet.address); + assert_eq!(wallet.address_string(), current_wallet.address_string()); // Clear all wallets clear_ethereum_wallets(); @@ -81,48 +81,40 @@ fn test_wallet_management() { let result = get_current_ethereum_wallet(); assert!(result.is_err()); - // Test legacy functions + // The legacy network-specific wallet functions have been removed + // We now use a single wallet that works across all networks - // Create wallets for different networks (should all create the same wallet) - let gnosis_wallet = create_ethereum_wallet_for_network(networks::gnosis()).unwrap(); - let peaq_wallet = create_ethereum_wallet_for_network(networks::peaq()).unwrap(); - let agung_wallet = create_ethereum_wallet_for_network(networks::agung()).unwrap(); + // Create a new wallet (network-agnostic) + let wallet = create_ethereum_wallet().unwrap(); - // They should all have the same address - assert_eq!(gnosis_wallet.address, peaq_wallet.address); - assert_eq!(gnosis_wallet.address, agung_wallet.address); + // Check that it's accessible + let current_wallet = get_current_ethereum_wallet().unwrap(); + assert_eq!(wallet.address_string(), current_wallet.address_string()); - // Get the current wallets for different networks (should all return the same wallet) - let current_gnosis = get_current_ethereum_wallet_for_network("Gnosis").unwrap(); - let current_peaq = get_current_ethereum_wallet_for_network("Peaq").unwrap(); - let current_agung = get_current_ethereum_wallet_for_network("Agung").unwrap(); + // Test for_network functionality to get network-specific wallet + let gnosis_network = networks::gnosis(); + let peaq_network = networks::peaq(); + let agung_network = networks::agung(); - // They should all have the same address - assert_eq!(current_gnosis.address, current_peaq.address); - assert_eq!(current_gnosis.address, current_agung.address); + // The wallet address should remain the same regardless of network + let wallet_address = current_wallet.address_string(); - // Clear wallets for a specific network (should be a no-op in the new design) - clear_ethereum_wallets_for_network("Gnosis"); + // Network-specific wallets have different chain IDs but same address + // Just verify different chain IDs here + let gnosis_wallet = current_wallet.for_network(&gnosis_network); + let peaq_wallet = current_wallet.for_network(&peaq_network); + let agung_wallet = current_wallet.for_network(&agung_network); - // All wallets should still be accessible - let current_gnosis = get_current_ethereum_wallet_for_network("Gnosis").unwrap(); - let current_peaq = get_current_ethereum_wallet_for_network("Peaq").unwrap(); - let current_agung = get_current_ethereum_wallet_for_network("Agung").unwrap(); - - // They should all have the same address - assert_eq!(current_gnosis.address, current_peaq.address); - assert_eq!(current_gnosis.address, current_agung.address); + // Check that chain IDs are different + assert_ne!(gnosis_wallet.chain_id(), peaq_wallet.chain_id()); + assert_ne!(gnosis_wallet.chain_id(), agung_wallet.chain_id()); // Clear all wallets clear_ethereum_wallets(); - // Check that all wallets are gone - let result1 = get_current_ethereum_wallet_for_network("Gnosis"); - let result2 = get_current_ethereum_wallet_for_network("Peaq"); - let result3 = get_current_ethereum_wallet_for_network("Agung"); - assert!(result1.is_err()); - assert!(result2.is_err()); - assert!(result3.is_err()); + // Check that the wallet is gone + let result = get_current_ethereum_wallet(); + assert!(result.is_err()); } #[test] @@ -180,33 +172,32 @@ fn test_wallet_for_network() { } #[test] -fn test_legacy_wallet_functions() { +fn test_multi_network_configuration() { let keypair = KeyPair::new("test_keypair7"); - // Test legacy wallet creation functions + // Create a network-agnostic wallet + let wallet = EthereumWallet::from_keypair(&keypair).unwrap(); + + // Test the for_network functionality to get network-specific configurations let gnosis_network = networks::gnosis(); + let peaq_network = networks::peaq(); + let agung_network = networks::agung(); - // Create a wallet with the legacy function - let wallet1 = EthereumWallet::from_keypair_for_network(&keypair, gnosis_network.clone()).unwrap(); + // Get the wallet's base address for comparison + let wallet_address = format!("{:?}", wallet.address); - // Create a wallet with the new function - let wallet2 = EthereumWallet::from_keypair(&keypair).unwrap(); + // Create network-specific signers + let gnosis_wallet = wallet.for_network(&gnosis_network); + let peaq_wallet = wallet.for_network(&peaq_network); + let agung_wallet = wallet.for_network(&agung_network); - // They should have the same address - assert_eq!(wallet1.address, wallet2.address); + // The signers should each have their network's chain ID + assert_eq!(gnosis_wallet.chain_id(), gnosis_network.chain_id); + assert_eq!(peaq_wallet.chain_id(), peaq_network.chain_id); + assert_eq!(agung_wallet.chain_id(), agung_network.chain_id); - // Test with name - let wallet3 = EthereumWallet::from_name_and_keypair_for_network("test", &keypair, gnosis_network.clone()).unwrap(); - let wallet4 = EthereumWallet::from_name_and_keypair("test", &keypair).unwrap(); - - // They should have the same address - assert_eq!(wallet3.address, wallet4.address); - - // Test with private key - let private_key = "0x0123456789abcdef0123456789abcdef0123456789abcdef0123456789abcdef"; - let wallet5 = EthereumWallet::from_private_key_for_network(private_key, gnosis_network.clone()).unwrap(); - let wallet6 = EthereumWallet::from_private_key(private_key).unwrap(); - - // They should have the same address - assert_eq!(wallet5.address, wallet6.address); + // And each should have the same address as the original wallet + assert_eq!(format!("{:?}", gnosis_wallet.address()), wallet_address); + assert_eq!(format!("{:?}", peaq_wallet.address()), wallet_address); + assert_eq!(format!("{:?}", agung_wallet.address()), wallet_address); } diff --git a/src/hero_vault/kvs/mod.rs b/src/hero_vault/kvs/mod.rs index 02e9e20..d874710 100644 --- a/src/hero_vault/kvs/mod.rs +++ b/src/hero_vault/kvs/mod.rs @@ -18,12 +18,6 @@ mod indexed_db_store; pub use error::{KvsError, Result}; pub use store::{KvPair, KVStore}; -// Legacy re-exports for backward compatibility -pub use store::{ - KvStore, create_store, open_store, delete_store, - list_stores, get_store_path -}; - // Re-export the SlateDbStore for native platforms pub use slate_store::{ SlateDbStore, create_slatedb_store, open_slatedb_store, diff --git a/src/rhai/hero_vault.rs b/src/rhai/hero_vault.rs index 46754d3..6fc1b63 100644 --- a/src/rhai/hero_vault.rs +++ b/src/rhai/hero_vault.rs @@ -4,13 +4,14 @@ use rhai::{Engine, Dynamic, EvalAltResult}; use base64::{Engine as _, engine::general_purpose::STANDARD as BASE64}; use std::fs; use std::sync::Mutex; -use once_cell::sync::Lazy; +use once_cell::sync::{Lazy, OnceCell}; use tokio::runtime::Runtime; use ethers::types::{Address, U256}; use std::str::FromStr; +use cfg_if::cfg_if; use crate::hero_vault::{keypair, symmetric, ethereum, kvs}; -use crate::hero_vault::kvs::DefaultStore; +use crate::hero_vault::kvs::{KVStore, DefaultStore}; use crate::hero_vault::ethereum::prepare_function_arguments; // Global Tokio runtime for blocking async operations @@ -29,40 +30,57 @@ where } // Get a platform-specific DefaultStore implementation for Rhai bindings -#[cfg(not(target_arch = "wasm32"))] -fn get_key_store() -> DefaultStore { - lazy_static::lazy_static! { - static ref STORE: DefaultStore = { - match kvs::open_default_store("rhai-vault", None) { - Ok(store) => store, - Err(_) => kvs::create_default_store("rhai-vault", false, None) - .expect("Failed to create store") - } - }; - } - - STORE.clone() -} - -// For WebAssembly, the store operations would typically be async -// but since Rhai requires synchronous functions, we create a blocking adapter -#[cfg(target_arch = "wasm32")] -fn get_key_store() -> DefaultStore { - use once_cell::sync::Lazy; - static STORE: Lazy = Lazy::new(|| { - match run_async(async { - kvs::open_default_store("rhai-vault", None).await - .or_else(|_| kvs::create_default_store("rhai-vault", false, None).await) - }) { - Ok(store) => store, - Err(e) => { - log::error!("Failed to create key store: {}", e); - panic!("Could not initialize key store: {}", e); - } +// This function is implemented differently based on the target platform +cfg_if! { + if #[cfg(target_arch = "wasm32")] { + fn get_key_store() -> DefaultStore { + use wasm_bindgen_futures::JsFuture; + use once_cell::sync::OnceCell; + use std::future::Future; + + // Static store instance + static KEY_STORE: OnceCell = OnceCell::new(); + + // Initialize if not already done + KEY_STORE.get_or_init(|| { + // In WebAssembly, we need to use a blocking approach for Rhai + let store_future = async { + match kvs::open_default_store("rhai-vault", None).await { + Ok(store) => store, + Err(_) => { + // Try to create the store if opening failed + kvs::create_default_store("rhai-vault", false, None).await + .expect("Failed to create key store") + } + } + }; + + // Block on the async operation + let rt = RUNTIME.lock().unwrap(); + rt.block_on(store_future) + }).clone() } - }); - - STORE.clone() + } else { + fn get_key_store() -> DefaultStore { + use once_cell::sync::OnceCell; + + // Static store instance + static KEY_STORE: OnceCell = OnceCell::new(); + + // Initialize if not already done + KEY_STORE.get_or_init(|| { + // For native platforms, the operations are synchronous + match kvs::open_default_store("rhai-vault", None) { + Ok(store) => store, + Err(_) => { + // Try to create the store if opening failed + kvs::create_default_store("rhai-vault", false, None) + .expect("Failed to create key store") + } + } + }).clone() + } + } } // Key space management functions @@ -262,6 +280,12 @@ fn create_ethereum_wallet_from_private_key(private_key: &str) -> bool { } } +// Clear all Ethereum wallets +fn clear_ethereum_wallets() -> bool { + ethereum::clear_ethereum_wallets(); + true // Always return true since the operation doesn't have a failure mode +} + // Network registry functions // Register a new network @@ -301,19 +325,6 @@ fn create_provider(network_name: &str) -> String { } } -// Legacy provider functions for backward compatibility -fn create_agung_provider() -> String { - create_provider("agung") -} - -fn create_peaq_provider() -> String { - create_provider("peaq") -} - -fn create_gnosis_provider() -> String { - create_provider("gnosis") -} - // Get network token symbol fn get_network_token_symbol(network_name: &str) -> String { match ethereum::get_network_by_name(network_name) { @@ -559,41 +570,6 @@ fn call_contract_write(contract_json: &str, function_name: &str, args: rhai::Arr } } -// Legacy functions for backward compatibility - -fn create_peaq_wallet() -> bool { - create_ethereum_wallet() -} - -fn get_peaq_address() -> String { - get_ethereum_address() -} - -fn create_agung_wallet() -> bool { - create_ethereum_wallet() -} - -fn get_agung_address() -> String { - get_ethereum_address() -} - -fn create_wallet_for_network(_network_name: &str) -> bool { - create_ethereum_wallet() -} - -fn get_wallet_address_for_network(_network_name: &str) -> String { - get_ethereum_address() -} - -fn clear_wallets_for_network(_network_name: &str) -> bool { - ethereum::clear_ethereum_wallets(); - true -} - -fn create_wallet_from_private_key_for_network(private_key: &str, _network_name: &str) -> bool { - create_ethereum_wallet_from_private_key(private_key) -} - /// Register crypto functions with the Rhai engine pub fn register_crypto_module(engine: &mut Engine) -> Result<(), Box> { // Register key space functions @@ -621,6 +597,7 @@ pub fn register_crypto_module(engine: &mut Engine) -> Result<(), Box Result<(), Box Result<(), Box