[rooch-networkgh-933] refactor key store and type (1.2/2).

feliciss · Oct 11, 2023 · edac265 · edac265
1 parent 8d54117
commit edac265
Show file tree

Hide file tree

Showing 19 changed files with 394 additions and 408 deletions.
diff --git a/crates/rooch-framework-tests/src/tests/ethereum_light_client_test.rs b/crates/rooch-framework-tests/src/tests/ethereum_light_client_test.rs
@@ -5,17 +5,15 @@ use crate::binding_test;
 use ethers::prelude::*;
 use moveos_types::transaction::MoveAction;
 use rooch_key::keystore::{AccountKeystore, InMemKeystore};
-use rooch_types::address::RoochAddress;
 use rooch_types::framework::ethereum_light_client::BlockHeader;
-use rooch_types::keypair_type::KeyPairType;
 use rooch_types::transaction::rooch::RoochTransactionData;
 
 #[test]
 fn test_submit_block() {
     tracing_subscriber::fmt::init();
     let mut binding_test = binding_test::RustBindingTest::new().unwrap();
 
-    let keystore = InMemKeystore::<RoochAddress>::new_insecure_for_tests(1);
+    let keystore = InMemKeystore::new_insecure_for_tests(1);
     let sender = keystore.addresses()[0];
     let sequence_number = 0;
 
@@ -55,14 +53,7 @@ fn test_submit_block() {
     let block_header = BlockHeader::try_from(&ethereum_block).unwrap();
     let action = MoveAction::Function(rooch_types::framework::ethereum_light_client::EthereumLightClientModule::create_submit_new_block_call(&block_header));
     let tx_data = RoochTransactionData::new_for_test(sender, sequence_number, action);
-    let tx = keystore
-        .sign_transaction(
-            &sender,
-            tx_data,
-            KeyPairType::RoochKeyPairType,
-            Some("".to_owned()),
-        )
-        .unwrap();
+    let tx = keystore.sign_transaction(&sender, tx_data, None).unwrap();
     binding_test.execute(tx).unwrap();
 
     let timestamp_module =

diff --git a/crates/rooch-framework-tests/src/tests/ethereum_validator_tests.rs b/crates/rooch-framework-tests/src/tests/ethereum_validator_tests.rs
@@ -4,9 +4,8 @@
 use ethers::types::{Bytes, U256};
 use moveos_types::transaction::MoveAction;
 use rooch_key::keystore::{AccountKeystore, InMemKeystore};
-use rooch_types::address::{EthereumAddress, MultiChainAddress};
+use rooch_types::address::MultiChainAddress;
 use rooch_types::framework::empty::Empty;
-use rooch_types::keypair_type::KeyPairType;
 use rooch_types::transaction::ethereum::EthereumTransactionData;
 use rooch_types::transaction::AbstractTransaction;
 
@@ -21,20 +20,15 @@ fn test_validate() {
     let address_mapping =
         binding_test.as_module_bundle::<rooch_types::framework::address_mapping::AddressMapping>();
 
-    let keystore = InMemKeystore::<EthereumAddress>::new_insecure_for_tests(1);
+    let keystore = InMemKeystore::new_insecure_for_tests(1);
     let sender = keystore.addresses()[0];
     let sequence_number = U256::zero();
     let action = MoveAction::new_function_call(Empty::empty_function_id(), vec![], vec![]);
     let action_bytes =
         Bytes::try_from(bcs::to_bytes(&action).unwrap()).expect("Convert action to bytes failed.");
     let tx_data = EthereumTransactionData::new_for_test(sender, sequence_number, action_bytes);
     keystore
-        .sign_transaction(
-            &sender,
-            tx_data.clone(),
-            KeyPairType::EthereumKeyPairType,
-            Some("".to_owned()),
-        )
+        .sign_transaction(&sender, tx_data.clone(), None)
         .unwrap();
     let auth_info = tx_data.authenticator_info().unwrap();
     let multichain_address = MultiChainAddress::from(sender);

diff --git a/crates/rooch-framework-tests/src/tests/native_validator_tests.rs b/crates/rooch-framework-tests/src/tests/native_validator_tests.rs
@@ -3,9 +3,7 @@
 
 use moveos_types::transaction::MoveAction;
 use rooch_key::keystore::{AccountKeystore, InMemKeystore};
-use rooch_types::address::RoochAddress;
 use rooch_types::framework::empty::Empty;
-use rooch_types::keypair_type::KeyPairType;
 use rooch_types::transaction::{rooch::RoochTransactionData, AbstractTransaction};
 
 use crate::binding_test;
@@ -17,19 +15,12 @@ fn test_validate() {
         .as_module_bundle::<rooch_types::framework::native_validator::NativeValidatorModule>(
     );
 
-    let keystore = InMemKeystore::<RoochAddress>::new_insecure_for_tests(1);
+    let keystore = InMemKeystore::new_insecure_for_tests(1);
     let sender = keystore.addresses()[0];
     let sequence_number = 0;
     let action = MoveAction::new_function_call(Empty::empty_function_id(), vec![], vec![]);
     let tx_data = RoochTransactionData::new_for_test(sender, sequence_number, action);
-    let tx = keystore
-        .sign_transaction(
-            &sender,
-            tx_data,
-            KeyPairType::RoochKeyPairType,
-            Some("".to_owned()),
-        )
-        .unwrap();
+    let tx = keystore.sign_transaction(&sender, tx_data, None).unwrap();
     let auth_info = tx.authenticator_info().unwrap();
     let move_tx = tx.construct_moveos_transaction(sender.into()).unwrap();
 

diff --git a/crates/rooch-framework-tests/src/tests/transaction_validator_tests.rs b/crates/rooch-framework-tests/src/tests/transaction_validator_tests.rs
@@ -10,10 +10,9 @@ use move_core_types::vm_status::{AbortLocation, VMStatus};
 use moveos_types::move_types::FunctionId;
 use moveos_types::{module_binding::ModuleBinding, transaction::MoveAction};
 use rooch_key::keystore::{AccountKeystore, InMemKeystore};
-use rooch_types::address::{EthereumAddress, MultiChainAddress, RoochAddress};
+use rooch_types::address::MultiChainAddress;
 use rooch_types::framework::session_key::SessionKeyModule;
 use rooch_types::framework::timestamp::TimestampModule;
-use rooch_types::keypair_type::KeyPairType;
 use rooch_types::transaction::ethereum::EthereumTransactionData;
 use rooch_types::{addresses::ROOCH_FRAMEWORK_ADDRESS, framework::empty::Empty};
 use rooch_types::{
@@ -30,19 +29,12 @@ fn test_validate_rooch() {
         .as_module_bundle::<rooch_types::framework::transaction_validator::TransactionValidator>(
     );
 
-    let keystore = InMemKeystore::<RoochAddress>::new_insecure_for_tests(1);
+    let keystore = InMemKeystore::new_insecure_for_tests(1);
     let sender = keystore.addresses()[0];
     let sequence_number = 0;
     let action = MoveAction::new_function_call(Empty::empty_function_id(), vec![], vec![]);
     let tx_data = RoochTransactionData::new_for_test(sender, sequence_number, action);
-    let tx = keystore
-        .sign_transaction(
-            &sender,
-            tx_data,
-            KeyPairType::RoochKeyPairType,
-            Some("".to_owned()),
-        )
-        .unwrap();
+    let tx = keystore.sign_transaction(&sender, tx_data, None).unwrap();
     let auth_info = tx.authenticator_info().unwrap();
     let move_tx = tx.construct_moveos_transaction(sender.into()).unwrap();
 
@@ -62,20 +54,15 @@ fn test_validate_ethereum() {
     let address_mapping =
         binding_test.as_module_bundle::<rooch_types::framework::address_mapping::AddressMapping>();
 
-    let keystore = InMemKeystore::<EthereumAddress>::new_insecure_for_tests(1);
+    let keystore = InMemKeystore::new_insecure_for_tests(1);
     let sender = keystore.addresses()[0];
     let sequence_number = U256::zero();
     let action = MoveAction::new_function_call(Empty::empty_function_id(), vec![], vec![]);
     let action_bytes =
         Bytes::try_from(bcs::to_bytes(&action).unwrap()).expect("Convert action to bytes failed.");
     let tx_data = EthereumTransactionData::new_for_test(sender, sequence_number, action_bytes);
-    let (_, _sig) = keystore
-        .sign_transaction(
-            &sender,
-            tx_data.clone(),
-            KeyPairType::EthereumKeyPairType,
-            Some("".to_owned()),
-        )
+    let tx = keystore
+        .sign_transaction(&sender, tx_data.clone(), None)
         .unwrap();
     let auth_info = tx_data.authenticator_info().unwrap();
     let multichain_address = MultiChainAddress::from(sender);
@@ -98,7 +85,7 @@ fn test_session_key_rooch() {
     // tracing_subscriber::fmt::init();
     let mut binding_test = binding_test::RustBindingTest::new().unwrap();
 
-    let mut keystore = InMemKeystore::<RoochAddress>::new_insecure_for_tests(1);
+    let mut keystore = InMemKeystore::new_insecure_for_tests(1);
     let sender = keystore.addresses()[0];
     let sequence_number = 0;
 
@@ -116,14 +103,7 @@ fn test_session_key_rooch() {
         max_inactive_interval,
     );
     let tx_data = RoochTransactionData::new_for_test(sender, sequence_number, action);
-    let tx = keystore
-        .sign_transaction(
-            &sender,
-            tx_data,
-            KeyPairType::RoochKeyPairType,
-            Some("".to_owned()),
-        )
-        .unwrap();
+    let tx = keystore.sign_transaction(&sender, tx_data, None).unwrap();
     binding_test.execute(tx).unwrap();
 
     let session_key_module =
@@ -194,14 +174,7 @@ fn test_session_key_rooch() {
     // because previous transaction is failed, so the sequence number is not increased.
     let tx_data =
         RoochTransactionData::new_for_test(sender, sequence_number + 2, update_time_action);
-    let tx = keystore
-        .sign_transaction(
-            &sender,
-            tx_data,
-            KeyPairType::RoochKeyPairType,
-            Some("".to_owned()),
-        )
-        .unwrap();
+    let tx = keystore.sign_transaction(&sender, tx_data, None).unwrap();
     binding_test.execute(tx).unwrap();
 
     let action = MoveAction::new_function_call(Empty::empty_function_id(), vec![], vec![]);

diff --git a/crates/rooch-key/src/key_derive.rs b/crates/rooch-key/src/key_derive.rs
@@ -4,17 +4,14 @@
 use argon2::password_hash::{PasswordHash, PasswordHasher, SaltString};
 use argon2::Argon2;
 use argon2::PasswordVerifier;
-use bip32::{DerivationPath, XPrv};
+use bip32::DerivationPath;
 use bip39::{Language, Mnemonic, MnemonicType, Seed};
 use chacha20poly1305::aead::Aead;
 use chacha20poly1305::{AeadCore, ChaCha20Poly1305, KeyInit};
 use fastcrypto::ed25519::{Ed25519KeyPair, Ed25519PrivateKey};
-use fastcrypto::secp256k1::recoverable::{
-    Secp256k1RecoverableKeyPair, Secp256k1RecoverablePrivateKey,
-};
 use fastcrypto::traits::{KeyPair, ToFromBytes};
 use rand::rngs::OsRng;
-use rooch_types::address::{EthereumAddress, RoochAddress};
+use rooch_types::address::RoochAddress;
 use rooch_types::crypto::RoochKeyPair;
 use rooch_types::error::RoochError;
 use rooch_types::key_struct::{EncryptionData, GenerateNewKeyPair, GeneratedKeyPair};
@@ -28,11 +25,9 @@ pub const DERIVATION_PATH_PURPOSE_SCHNORR: u32 = 44;
 pub const DERIVATION_PATH_PURPOSE_ECDSA: u32 = 54;
 pub const DERIVATION_PATH_PURPOSE_SECP256R1: u32 = 74;
 
-type EncryptionKeyResult = Result<(Vec<u8>, Vec<u8>, Vec<u8>), RoochError>;
-
 pub fn encrypt_private_key(
     private_key: &[u8],
-    password: Option<&str>,
+    password: Option<String>,
 ) -> Result<EncryptionData, RoochError> {
     let nonce = ChaCha20Poly1305::generate_nonce(&mut OsRng);
     let mut output_key_material = [0u8; 32];
@@ -47,7 +42,7 @@ pub fn encrypt_private_key(
     let cipher = ChaCha20Poly1305::new_from_slice(&output_key_material)
         .map_err(|e| RoochError::KeyConversionError(e.to_string()))?;
 
-    let ciphertext_with_tag = match cipher.encrypt(&nonce, &*private_key) {
+    let ciphertext_with_tag = match cipher.encrypt(&nonce, private_key) {
         Ok(ciphertext) => ciphertext,
         Err(_) => {
             return Err(RoochError::KeyConversionError(
@@ -70,7 +65,7 @@ pub fn decrypt_private_key(
     nonce: &[u8],
     ciphertext: &[u8],
     tag: &[u8],
-    password: Option<&str>,
+    password: Option<String>,
 ) -> Result<Vec<u8>, RoochError> {
     let mut output_key_material = [0u8; 32];
     Argon2::default()
@@ -97,7 +92,7 @@ pub fn decrypt_private_key(
 }
 
 pub fn verify_password(
-    password: Option<&str>,
+    password: Option<String>,
     password_hash: String,
 ) -> Result<bool, argon2::password_hash::Error> {
     let parsed_hash = PasswordHash::new(&password_hash)?;
@@ -106,9 +101,9 @@ pub fn verify_password(
         .is_ok())
 }
 
-pub fn hash_password(nonce: &[u8], password: Option<&str>) -> Result<String, RoochError> {
-    let salt = SaltString::encode_b64(&nonce)
-        .map_err(|e| RoochError::KeyConversionError(e.to_string()))?;
+pub fn hash_password(nonce: &[u8], password: Option<String>) -> Result<String, RoochError> {
+    let salt =
+        SaltString::encode_b64(nonce).map_err(|e| RoochError::KeyConversionError(e.to_string()))?;
     let argon2 = Argon2::default();
     let password_hash = argon2
         .hash_password(password.unwrap_or_default().as_bytes(), &salt)
@@ -195,17 +190,10 @@ pub fn generate_new_key_pair(
 
     let sk = derive_private_key_from_path(seed.as_bytes(), derivation_path)?;
 
-    let (nonce, ciphertext, tag) = encrypt_private_key(&sk.clone(), password.clone())
-        .expect("Encryption failed for private key");
+    let encryption = encrypt_private_key(&sk, password).expect("Encryption failed for private key");
 
     let address = derive_address_from_private_key(sk)?;
 
-    let encryption = EncryptionData {
-        nonce,
-        ciphertext,
-        tag,
-    };
-
     let result = GenerateNewKeyPair {
         encryption,
         mnemonic: mnemonic.phrase().to_string(),
@@ -229,8 +217,7 @@ fn parse_word_length(s: Option<String>) -> Result<MnemonicType, anyhow::Error> {
 /// Get a keypair from a random encryption data
 pub fn get_key_pair_from_red() -> (RoochAddress, EncryptionData) {
     let random_encryption_data = EncryptionData::new_for_test();
-    let kp: RoochKeyPair =
-        retrieve_key_pair(&random_encryption_data, Some("".to_owned()), None).unwrap();
+    let kp: RoochKeyPair = retrieve_key_pair(&random_encryption_data, None).unwrap();
 
     ((&kp.public()).into(), random_encryption_data)
 }