refactor EncryptionData

crates/rooch-benchmarks/benches/bench_tx_sequence.rs

@@ -21,9 +21,7 @@ pub fn tx_sequence_benchmark(c: &mut Criterion) {
 
     let rooch_account = keystore.addresses()[0];
     let rooch_key_pair = keystore
-        .get_key_pairs(&rooch_account, None)
-        .unwrap()
-        .pop()
+        .get_key_pair(&rooch_account, None)
         .expect("key pair should have value");
     let sequencer_keypair = rooch_key_pair.copy();
     let mut sequencer =

crates/rooch-benchmarks/src/tx.rs

@@ -83,8 +83,7 @@ pub async fn setup_service(
     // init keystore
     let rooch_account = keystore.addresses()[0];
     let rooch_key_pair = keystore
-        .get_key_pairs(&rooch_account, None)?
-        .pop()
+        .get_key_pair(&rooch_account, None)
         .expect("Key pair should have value");
 
     let sequencer_keypair = rooch_key_pair.copy();

crates/rooch-key/src/key_derive.rs

@@ -5,18 +5,9 @@ use argon2::password_hash::{PasswordHash, PasswordHasher, SaltString};
 use argon2::Argon2;
 use argon2::PasswordVerifier;
 use bip39::{Language, Mnemonic, MnemonicType, Seed};
-use bitcoin::address::Address;
-use bitcoin::bip32::{self, ChildNumber, DerivationPath, Xpriv, Xpub};
-use bitcoin::hex::FromHex;
-use bitcoin::secp256k1::ffi::types::AlignedType;
-use bitcoin::secp256k1::Secp256k1;
-use bitcoin::Network;
-use chacha20poly1305::aead::Aead;
-use chacha20poly1305::{AeadCore, ChaCha20Poly1305, KeyInit};
+use bitcoin::bip32::{self, DerivationPath};
 use fastcrypto::ed25519::{Ed25519KeyPair, Ed25519PrivateKey};
-use fastcrypto::encoding::{Base64, Encoding};
 use fastcrypto::traits::{KeyPair, ToFromBytes};
-use rand::rngs::OsRng;
 use rooch_types::address::RoochAddress;
 use rooch_types::crypto::RoochKeyPair;
 use rooch_types::error::RoochError;
@@ -33,69 +24,6 @@ pub const DERIVATION_PATH_PURPOSE_ECDSA: u32 = 54;
 pub const DERIVATION_PATH_PURPOSE_SECP256R1: u32 = 74;
 pub const DERIVATION_PATH_PURPOSE_BIP84: u32 = 84;
 
-pub fn encrypt_key(key: &[u8], password: Option<String>) -> Result<EncryptionData, RoochError> {
-    let nonce = ChaCha20Poly1305::generate_nonce(&mut OsRng);
-    let mut output_key_material = [0u8; 32];
-    Argon2::default()
-        .hash_password_into(
-            password.unwrap_or_default().as_bytes(),
-            &nonce,
-            &mut output_key_material,
-        )
-        .map_err(|e| RoochError::KeyConversionError(e.to_string()))?;
-
-    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, key) {
-        Ok(ciphertext) => ciphertext,
-        Err(_) => {
-            return Err(RoochError::KeyConversionError(
-                "Encryption failed".to_owned(),
-            ))
-        }
-    };
-
-    let ciphertext = ciphertext_with_tag[..ciphertext_with_tag.len() - 16].to_vec();
-    let tag = ciphertext_with_tag[ciphertext_with_tag.len() - 16..].to_vec();
-
-    Ok(EncryptionData {
-        nonce: Base64::encode(nonce),
-        ciphertext: Base64::encode(ciphertext),
-        tag: Base64::encode(tag),
-    })
-}
-
-pub fn decrypt_key(
-    nonce: &[u8],
-    ciphertext: &[u8],
-    tag: &[u8],
-    password: Option<String>,
-) -> Result<Vec<u8>, RoochError> {
-    let mut output_key_material = [0u8; 32];
-    Argon2::default()
-        .hash_password_into(
-            password.unwrap_or_default().as_bytes(),
-            nonce,
-            &mut output_key_material,
-        )
-        .map_err(|e| RoochError::KeyConversionError(e.to_string()))?;
-
-    let cipher = ChaCha20Poly1305::new_from_slice(&output_key_material)
-        .map_err(|e| RoochError::KeyConversionError(e.to_string()))?;
-
-    let mut ciphertext_with_tag = Vec::with_capacity(tag.len() + ciphertext.len());
-    ciphertext_with_tag.extend_from_slice(ciphertext);
-    ciphertext_with_tag.extend_from_slice(tag);
-
-    match cipher.decrypt(nonce.into(), &*ciphertext_with_tag) {
-        Ok(pk) => Ok(pk),
-        Err(_) => Err(RoochError::KeyConversionError(
-            "Decryption failed".to_owned(),
-        )),
-    }
-}
-
 pub fn verify_password(
     password: Option<String>,
     password_hash: String,
@@ -127,9 +55,8 @@ pub fn hash_password(nonce: &[u8], password: Option<String>) -> Result<String, R
 
 pub(crate) fn derive_private_key_from_path(
     seed: &[u8],
-    derivation_path: Option<DerivationPath>,
+    path: DerivationPath,
 ) -> Result<Vec<u8>, anyhow::Error> {
-    let path = validate_derivation_path(derivation_path)?;
     let indexes = path.to_u32_vec();
     let derived = derive_ed25519_private_key(seed, &indexes);
     let sk = Ed25519PrivateKey::from_bytes(&derived)
@@ -168,15 +95,8 @@ pub fn derive_address_from_private_key(private_key: Vec<u8>) -> Result<RoochAddr
 pub fn retrieve_key_pair(
     encryption: &EncryptionData,
     password: Option<String>,
-) -> Result<RoochKeyPair, RoochError> {
-    let nonce = Base64::decode(&encryption.nonce)
-        .map_err(|e| RoochError::KeyConversionError(e.to_string()))?;
-    let ciphertext = Base64::decode(&encryption.ciphertext)
-        .map_err(|e| RoochError::KeyConversionError(e.to_string()))?;
-    let tag = Base64::decode(&encryption.tag)
-        .map_err(|e| RoochError::KeyConversionError(e.to_string()))?;
-
-    let private_key = decrypt_key(&nonce, &ciphertext, &tag, password)?;
+) -> Result<RoochKeyPair, anyhow::Error> {
+    let private_key = encryption.decrypt(password)?;
 
     let kp = Ed25519KeyPair::from(
         Ed25519PrivateKey::from_bytes(&private_key)
@@ -186,6 +106,8 @@ pub fn retrieve_key_pair(
     Ok(kp.into())
 }
 
+//In the future, we may support for custom derivation path to recover the keypair
+#[allow(dead_code)]
 fn validate_derivation_path(path: Option<DerivationPath>) -> Result<DerivationPath, anyhow::Error> {
     let (purpose, coin_type) = (
         DERIVATION_PATH_PURPOSE_ED25519,
@@ -243,7 +165,7 @@ pub(crate) fn generate_derivation_path(account_index: u32) -> Result<DerivationP
 
 pub fn generate_new_key_pair(
     mnemonic_phrase: Option<String>,
-    derivation_path: Option<DerivationPath>,
+    derivation_path: DerivationPath,
     word_length: Option<String>,
     password: Option<String>,
 ) -> Result<GeneratedKeyPair, anyhow::Error> {
@@ -259,10 +181,9 @@ pub fn generate_new_key_pair(
 
     let sk = derive_private_key_from_path(seed.as_bytes(), derivation_path)?;
 
-    let private_key_encryption =
-        encrypt_key(&sk, password.clone()).expect("Encryption failed for private key");
-    let mnemonic_phrase_encryption = encrypt_key(mnemonic.phrase().as_bytes(), password)
-        .expect("Encryption failed for mnemonic phrase");
+    let private_key_encryption = EncryptionData::encrypt(&sk, password.clone())?;
+    let mnemonic_phrase_encryption =
+        EncryptionData::encrypt(mnemonic.phrase().as_bytes(), password)?;
 
     let address = derive_address_from_private_key(sk)?;
 
@@ -289,11 +210,3 @@ fn parse_word_length(s: Option<String>) -> Result<MnemonicType, anyhow::Error> {
         _ => Err(anyhow::anyhow!("Invalid word length")),
     }
 }
-
-/// Get a keypair from a random encryption data
-pub(crate) fn generate_key_pair_for_tests() -> (RoochAddress, EncryptionData) {
-    let random_encryption_data = EncryptionData::new_for_test();
-    let kp: RoochKeyPair = retrieve_key_pair(&random_encryption_data, None).unwrap();
-
-    ((&kp.public()).into(), random_encryption_data)
-}

crates/rooch-key/src/keystore/account_keystore.rs

@@ -2,21 +2,13 @@
 // SPDX-License-Identifier: Apache-2.0
 
 use super::types::LocalAccount;
-use crate::key_derive::{
-    derive_address_from_private_key, derive_private_key_from_path, encrypt_key,
-    generate_derivation_path, generate_new_key_pair, hash_password,
-};
-use crate::keystore::ImportedMnemonic;
-use bip39::{Language, Mnemonic, Seed};
-use bitcoin::bip32::DerivationPath;
-use fastcrypto::encoding::{Base64, Encoding};
+use crate::key_derive::{generate_derivation_path, generate_new_key_pair};
 use rooch_types::framework::session_key::SessionKey;
 use rooch_types::key_struct::{MnemonicData, MnemonicResult};
 use rooch_types::{
     address::RoochAddress,
     authentication_key::AuthenticationKey,
-    crypto::{PublicKey, RoochKeyPair, Signature},
-    error::RoochError,
+    crypto::{RoochKeyPair, Signature},
     key_struct::{EncryptionData, GeneratedKeyPair},
     transaction::rooch::{RoochTransaction, RoochTransactionData},
 };
@@ -30,12 +22,8 @@ pub trait AccountKeystore {
         password: Option<String>,
     ) -> Result<GeneratedKeyPair, anyhow::Error> {
         let derivation_path = generate_derivation_path(0)?;
-        let result = generate_new_key_pair(
-            mnemonic_phrase,
-            Some(derivation_path),
-            word_length,
-            password,
-        )?;
+        let result =
+            generate_new_key_pair(mnemonic_phrase, derivation_path, word_length, password)?;
         let new_address = result.address;
         self.add_address_encryption_data(
             new_address,
@@ -64,7 +52,7 @@ pub trait AccountKeystore {
 
         let result = generate_new_key_pair(
             Some(mnemonic.mnemonic_phrase),
-            Some(derivation_path),
+            derivation_path,
             None,
             password,
         )?;
@@ -84,7 +72,7 @@ pub trait AccountKeystore {
         encryption: EncryptionData,
     ) -> Result<(), anyhow::Error>;
 
-    fn get_key_pair_with_password(
+    fn get_key_pair(
         &self,
         address: &RoochAddress,
         password: Option<String>,

crates/rooch-key/src/keystore/base_keystore.rs

@@ -4,16 +4,15 @@
 use std::collections::BTreeMap;
 
 use super::types::{AddressMapping, LocalAccount, LocalSessionKey};
-use crate::key_derive::{decrypt_key, generate_new_key_pair, retrieve_key_pair};
+use crate::key_derive::retrieve_key_pair;
 use crate::keystore::account_keystore::AccountKeystore;
-use anyhow::{anyhow, ensure};
-use fastcrypto::encoding::{Base64, Encoding};
+use anyhow::ensure;
 use rooch_types::framework::session_key::SessionKey;
 use rooch_types::key_struct::{MnemonicData, MnemonicResult};
 use rooch_types::{
     address::RoochAddress,
     authentication_key::AuthenticationKey,
-    crypto::{PublicKey, RoochKeyPair, Signature},
+    crypto::{RoochKeyPair, Signature},
     error::RoochError,
     key_struct::EncryptionData,
     transaction::{
@@ -102,7 +101,7 @@ impl AccountKeystore for BaseKeyStore {
         Ok(accounts.into_values().collect())
     }
 
-    fn get_key_pair_with_password(
+    fn get_key_pair(
         &self,
         address: &RoochAddress,
         password: Option<String>,
@@ -126,7 +125,7 @@ impl AccountKeystore for BaseKeyStore {
     ) -> Result<Signature, anyhow::Error> {
         Ok(Signature::new_hashed(
             msg,
-            &self.get_key_pair_with_password(address, password)?,
+            &self.get_key_pair(address, password)?,
         ))
     }
 
@@ -141,7 +140,7 @@ impl AccountKeystore for BaseKeyStore {
     {
         Ok(Signature::new_secure(
             msg,
-            &self.get_key_pair_with_password(address, password)?,
+            &self.get_key_pair(address, password)?,
         ))
     }
 
@@ -151,12 +150,9 @@ impl AccountKeystore for BaseKeyStore {
         msg: RoochTransactionData,
         password: Option<String>,
     ) -> Result<RoochTransaction, anyhow::Error> {
-        let kp = self
-            .get_key_pair_with_password(address, password)
-            .ok()
-            .ok_or_else(|| {
-                RoochError::SignMessageError(format!("Cannot find key for address: [{address}]"))
-            })?;
+        let kp = self.get_key_pair(address, password).ok().ok_or_else(|| {
+            RoochError::SignMessageError(format!("Cannot find key for address: [{address}]"))
+        })?;
 
         let signature = Signature::new_hashed(msg.tx_hash().as_bytes(), &kp);
 
@@ -187,9 +183,10 @@ impl AccountKeystore for BaseKeyStore {
         let kp: RoochKeyPair = RoochKeyPair::generate();
         let authentication_key = kp.public().authentication_key();
         let inner_map = self.session_keys.entry(*address).or_default();
+        let private_key_encryption = EncryptionData::encrypt(kp.private(), password)?;
         let local_session_key = LocalSessionKey {
             session_key: None,
-            private_key: result.key_pair_data.private_key_encryption,
+            private_key: private_key_encryption,
         };
         inner_map.insert(authentication_key.clone(), local_session_key);
         Ok(authentication_key)
@@ -278,25 +275,7 @@ impl AccountKeystore for BaseKeyStore {
     fn get_mnemonic(&self, password: Option<String>) -> Result<MnemonicResult, anyhow::Error> {
         match &self.mnemonic {
             Some(mnemonic_data) => {
-                let nonce = Base64::decode(&mnemonic_data.mnemonic_phrase_encryption.nonce)
-                    .map_err(|e| {
-                        anyhow::Error::new(RoochError::KeyConversionError(e.to_string()))
-                    })?;
-                let ciphertext = Base64::decode(
-                    &mnemonic_data.mnemonic_phrase_encryption.ciphertext,
-                )
-                .map_err(|e| anyhow::Error::new(RoochError::KeyConversionError(e.to_string())))?;
-                let tag =
-                    Base64::decode(&mnemonic_data.mnemonic_phrase_encryption.tag).map_err(|e| {
-                        anyhow::Error::new(RoochError::KeyConversionError(e.to_string()))
-                    })?;
-
-                let mnemonic_phrase = decrypt_key(
-                    nonce.as_slice(),
-                    ciphertext.as_slice(),
-                    tag.as_slice(),
-                    password,
-                )?;
+                let mnemonic_phrase = mnemonic_data.mnemonic_phrase_encryption.decrypt(password)?;
 
                 let mnemonic_phrase = String::from_utf8(mnemonic_phrase)
                     .map_err(|e| anyhow::anyhow!("Parse mnemonic phrase error:{}", e))?;

crates/rooch-key/src/keystore/file_keystore.rs

@@ -10,12 +10,11 @@ use rooch_types::key_struct::{MnemonicData, MnemonicResult};
 use rooch_types::{
     address::RoochAddress,
     authentication_key::AuthenticationKey,
-    crypto::{PublicKey, RoochKeyPair, Signature},
+    crypto::{RoochKeyPair, Signature},
     key_struct::EncryptionData,
     transaction::rooch::{RoochTransaction, RoochTransactionData},
 };
 use serde::{Deserialize, Serialize};
-use std::collections::BTreeMap;
 use std::fs;
 use std::fs::File;
 use std::io::BufReader;
@@ -49,12 +48,12 @@ impl AccountKeystore for FileBasedKeystore {
         Ok(())
     }
 
-    fn get_key_pair_with_password(
+    fn get_key_pair(
         &self,
         address: &RoochAddress,
         password: Option<String>,
     ) -> Result<RoochKeyPair, anyhow::Error> {
-        self.keystore.get_key_pair_with_password(address, password)
+        self.keystore.get_key_pair(address, password)
     }
 
     fn nullify(&mut self, address: &RoochAddress) -> Result<(), anyhow::Error> {