[cli] Support local address mapping

rooch-network · Dec 28, 2023 · bec3156 · bec3156
1 parent b2821c4
commit bec3156
Show file tree

Hide file tree

Showing 14 changed files with 493 additions and 96 deletions.
diff --git a/crates/rooch-key/src/keystore/account_keystore.rs b/crates/rooch-key/src/keystore/account_keystore.rs
@@ -1,6 +1,7 @@
 // Copyright (c) RoochNetwork
 // 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,
@@ -9,6 +10,7 @@ use crate::keystore::ImportedMnemonic;
 use bip32::DerivationPath;
 use bip39::{Language, Mnemonic, Seed};
 use fastcrypto::encoding::{Base64, Encoding};
+use rooch_types::framework::session_key::SessionKey;
 use rooch_types::key_struct::{MnemonicData, MnemonicResult};
 use rooch_types::{
     address::RoochAddress,
@@ -21,6 +23,8 @@ use rooch_types::{
 use serde::Serialize;
 
 pub trait AccountKeystore {
+    fn get_accounts(&self, password: Option<String>) -> Result<Vec<LocalAccount>, anyhow::Error>;
+
     fn add_address_encryption_data(
         &mut self,
         address: RoochAddress,
@@ -30,6 +34,7 @@ pub trait AccountKeystore {
         &self,
         password: Option<String>,
     ) -> Result<Vec<(RoochAddress, PublicKey)>, anyhow::Error>;
+
     fn get_public_key(&self, password: Option<String>) -> Result<PublicKey, anyhow::Error>;
     fn get_key_pairs(
         &self,
@@ -196,6 +201,13 @@ pub trait AccountKeystore {
         password: Option<String>,
     ) -> Result<AuthenticationKey, anyhow::Error>;
 
+    /// Binding on-chain SessionKey to LocalSessionKey
+    fn binding_session_key(
+        &mut self,
+        address: RoochAddress,
+        session_key: SessionKey,
+    ) -> Result<(), anyhow::Error>;
+
     fn sign_transaction_via_session_key(
         &self,
         address: &RoochAddress,

diff --git a/crates/rooch-key/src/keystore/base_keystore.rs b/crates/rooch-key/src/keystore/base_keystore.rs
@@ -3,11 +3,12 @@
 
 use std::collections::BTreeMap;
 
-use anyhow::anyhow;
+use super::types::{AddressMapping, LocalAccount, LocalSessionKey};
+use crate::key_derive::{decrypt_key, generate_new_key_pair, retrieve_key_pair};
+use crate::keystore::account_keystore::AccountKeystore;
+use anyhow::{anyhow, ensure};
 use fastcrypto::encoding::{Base64, Encoding};
-use serde::{Deserialize, Serialize};
-use serde_with::serde_as;
-
+use rooch_types::framework::session_key::SessionKey;
 use rooch_types::key_struct::{MnemonicData, MnemonicResult};
 use rooch_types::{
     address::RoochAddress,
@@ -20,19 +21,25 @@ use rooch_types::{
         rooch::{RoochTransaction, RoochTransactionData},
     },
 };
-
-use crate::key_derive::{decrypt_key, generate_new_key_pair, retrieve_key_pair};
-use crate::keystore::account_keystore::AccountKeystore;
+use serde::{Deserialize, Serialize};
+use serde_with::serde_as;
 
 #[derive(Default, Debug, Serialize, Deserialize)]
 #[serde_as]
 pub(crate) struct BaseKeyStore {
+    #[serde(default)]
     pub(crate) keys: BTreeMap<RoochAddress, EncryptionData>,
+    #[serde(default)]
     pub(crate) mnemonics: BTreeMap<String, MnemonicData>,
+    #[serde(default)]
     #[serde_as(as = "BTreeMap<DisplayFromStr, BTreeMap<DisplayFromStr, _>>")]
-    pub(crate) session_keys: BTreeMap<RoochAddress, BTreeMap<AuthenticationKey, EncryptionData>>,
+    pub(crate) session_keys: BTreeMap<RoochAddress, BTreeMap<AuthenticationKey, LocalSessionKey>>,
+    #[serde(default)]
     pub(crate) password_hash: Option<String>,
+    #[serde(default)]
     pub(crate) is_password_empty: bool,
+    #[serde(default)]
+    pub(crate) address_mapping: AddressMapping,
 }
 
 impl BaseKeyStore {
@@ -43,11 +50,52 @@ impl BaseKeyStore {
             session_keys: BTreeMap::new(),
             password_hash: None,
             is_password_empty: true,
+            address_mapping: AddressMapping::default(),
         }
     }
 }
 
 impl AccountKeystore for BaseKeyStore {
+    fn get_accounts(&self, password: Option<String>) -> Result<Vec<LocalAccount>, anyhow::Error> {
+        let mut accounts = BTreeMap::new();
+        for (address, encryption) in &self.keys {
+            let keypair: RoochKeyPair = retrieve_key_pair(encryption, password.clone())?;
+            let public_key = keypair.public();
+            let multichain_address = self
+                .address_mapping
+                .rooch_to_multichain
+                .get(address)
+                .cloned();
+            let has_session_key = self.session_keys.get(address).is_some();
+            let local_account = LocalAccount {
+                address: *address,
+                multichain_address,
+                public_key: Some(public_key),
+                has_session_key,
+            };
+            accounts.insert(*address, local_account);
+        }
+        for address in self.session_keys.keys() {
+            if accounts.contains_key(address) {
+                continue;
+            }
+            let multichain_address = self
+                .address_mapping
+                .rooch_to_multichain
+                .get(address)
+                .cloned();
+            let has_session_key = true;
+            let local_account = LocalAccount {
+                address: *address,
+                multichain_address,
+                public_key: None,
+                has_session_key,
+            };
+            accounts.insert(*address, local_account);
+        }
+        Ok(accounts.into_values().into_iter().collect())
+    }
+
     fn get_key_pair_with_password(
         &self,
         address: &RoochAddress,
@@ -182,21 +230,37 @@ impl AccountKeystore for BaseKeyStore {
             retrieve_key_pair(&result.key_pair_data.private_key_encryption, password)?;
         let authentication_key = kp.public().authentication_key();
         let inner_map = self.session_keys.entry(*address).or_default();
-        inner_map.insert(
-            authentication_key.clone(),
-            result.key_pair_data.private_key_encryption,
-        );
+        let local_session_key = LocalSessionKey {
+            session_key: None,
+            private_key: result.key_pair_data.private_key_encryption,
+        };
+        inner_map.insert(authentication_key.clone(), local_session_key);
         Ok(authentication_key)
     }
 
+    fn binding_session_key(
+        &mut self,
+        address: RoochAddress,
+        session_key: SessionKey,
+    ) -> Result<(), anyhow::Error> {
+        let inner_map: &mut BTreeMap<AuthenticationKey, LocalSessionKey> =
+            self.session_keys.entry(address).or_default();
+        let authentication_key = session_key.authentication_key();
+        let local_session_key = inner_map.get_mut(&authentication_key).ok_or_else(||{
+            anyhow::Error::new(RoochError::KeyConversionError(format!("Cannot find session key for address:[{address}] and authentication_key:[{authentication_key}]", address = address, authentication_key = authentication_key)))
+        })?;
+        local_session_key.session_key = Some(session_key);
+        Ok(())
+    }
+
     fn sign_transaction_via_session_key(
         &self,
         address: &RoochAddress,
         msg: RoochTransactionData,
         authentication_key: &AuthenticationKey,
         password: Option<String>,
     ) -> Result<RoochTransaction, anyhow::Error> {
-        let encryption = self
+        let local_session_key = self
             .session_keys
             .get(address)
             .ok_or_else(|| {
@@ -210,9 +274,16 @@ impl AccountKeystore for BaseKeyStore {
                     "Cannot find SessionKey for authentication_key: [{authentication_key}]"
                 ))
             })?;
-
-        let kp: RoochKeyPair =
-            retrieve_key_pair(encryption, password).map_err(signature::Error::from_source)?;
+        let session_key = local_session_key.session_key.as_ref().ok_or_else(||{
+            signature::Error::from_source(
+                format!("SessionKey for authentication_key:[{authentication_key}] do not binding to on-chain SessionKey")
+            )
+        })?;
+        ensure!(session_key.is_scope_match_with_action(&msg.action), signature::Error::from_source(
+            format!("SessionKey for authentication_key:[{authentication_key}] scope do not match with transaction")
+        ));
+        let kp: RoochKeyPair = retrieve_key_pair(&local_session_key.private_key, password)
+            .map_err(signature::Error::from_source)?;
 
         let signature = Signature::new_hashed(msg.hash().as_bytes(), &kp);
 

diff --git a/crates/rooch-key/src/keystore/file_keystore.rs b/crates/rooch-key/src/keystore/file_keystore.rs
@@ -1,6 +1,7 @@
 // Copyright (c) RoochNetwork
 // SPDX-License-Identifier: Apache-2.0
 
+use super::types::LocalAccount;
 use crate::key_derive::retrieve_key_pair;
 use crate::keystore::account_keystore::AccountKeystore;
 use crate::keystore::base_keystore::BaseKeyStore;
@@ -27,6 +28,10 @@ pub struct FileBasedKeystore {
 }
 
 impl AccountKeystore for FileBasedKeystore {
+    fn get_accounts(&self, password: Option<String>) -> Result<Vec<LocalAccount>, anyhow::Error> {
+        self.keystore.get_accounts(password)
+    }
+
     fn add_address_encryption_data(
         &mut self,
         address: RoochAddress,
@@ -138,6 +143,16 @@ impl AccountKeystore for FileBasedKeystore {
         Ok(auth_key)
     }
 
+    fn binding_session_key(
+        &mut self,
+        address: RoochAddress,
+        session_key: rooch_types::framework::session_key::SessionKey,
+    ) -> Result<(), anyhow::Error> {
+        self.keystore.binding_session_key(address, session_key)?;
+        self.save()?;
+        Ok(())
+    }
+
     fn sign_transaction_via_session_key(
         &self,
         address: &RoochAddress,

diff --git a/crates/rooch-key/src/keystore/memory_keystore.rs b/crates/rooch-key/src/keystore/memory_keystore.rs
@@ -1,6 +1,7 @@
 // Copyright (c) RoochNetwork
 // SPDX-License-Identifier: Apache-2.0
 
+use super::types::LocalAccount;
 use crate::key_derive::get_key_pair_from_red;
 use crate::keystore::account_keystore::AccountKeystore;
 use crate::keystore::base_keystore::BaseKeyStore;
@@ -21,6 +22,10 @@ pub struct InMemKeystore {
 }
 
 impl AccountKeystore for InMemKeystore {
+    fn get_accounts(&self, password: Option<String>) -> Result<Vec<LocalAccount>, anyhow::Error> {
+        self.keystore.get_accounts(password)
+    }
+
     fn add_address_encryption_data(
         &mut self,
         address: RoochAddress,
@@ -124,6 +129,14 @@ impl AccountKeystore for InMemKeystore {
         self.keystore.generate_session_key(address, password)
     }
 
+    fn binding_session_key(
+        &mut self,
+        address: RoochAddress,
+        session_key: rooch_types::framework::session_key::SessionKey,
+    ) -> Result<(), anyhow::Error> {
+        self.keystore.binding_session_key(address, session_key)
+    }
+
     fn sign_transaction_via_session_key(
         &self,
         address: &RoochAddress,

diff --git a/crates/rooch-key/src/keystore/mod.rs b/crates/rooch-key/src/keystore/mod.rs
@@ -21,6 +21,7 @@ pub mod account_keystore;
 pub mod base_keystore;
 pub mod file_keystore;
 pub mod memory_keystore;
+pub mod types;
 
 pub struct ImportedMnemonic {
     pub address: RoochAddress,
@@ -35,6 +36,16 @@ pub enum Keystore {
 }
 
 impl AccountKeystore for Keystore {
+    fn get_accounts(
+        &self,
+        password: Option<String>,
+    ) -> Result<Vec<types::LocalAccount>, anyhow::Error> {
+        match self {
+            Keystore::File(file_keystore) => file_keystore.get_accounts(password),
+            Keystore::InMem(inmem_keystore) => inmem_keystore.get_accounts(password),
+        }
+    }
+
     fn sign_transaction_via_session_key(
         &self,
         address: &RoochAddress,
@@ -204,6 +215,21 @@ impl AccountKeystore for Keystore {
         }
     }
 
+    fn binding_session_key(
+        &mut self,
+        address: RoochAddress,
+        session_key: rooch_types::framework::session_key::SessionKey,
+    ) -> Result<(), anyhow::Error> {
+        match self {
+            Keystore::File(file_keystore) => {
+                file_keystore.binding_session_key(address, session_key)
+            }
+            Keystore::InMem(inmem_keystore) => {
+                inmem_keystore.binding_session_key(address, session_key)
+            }
+        }
+    }
+
     fn addresses(&self) -> Vec<RoochAddress> {
         match self {
             Keystore::File(file_keystore) => file_keystore.addresses(),

diff --git a/crates/rooch-key/src/keystore/types.rs b/crates/rooch-key/src/keystore/types.rs
@@ -0,0 +1,40 @@
+// Copyright (c) RoochNetwork
+// SPDX-License-Identifier: Apache-2.0
+
+use rooch_types::{
+    address::{MultiChainAddress, RoochAddress},
+    crypto::PublicKey,
+    framework::session_key::SessionKey,
+    key_struct::EncryptionData,
+};
+use serde::{Deserialize, Serialize};
+use std::collections::BTreeMap;
+
+#[derive(Serialize, Deserialize, Debug)]
+pub struct LocalSessionKey {
+    pub session_key: Option<SessionKey>,
+    pub private_key: EncryptionData,
+}
+
+#[derive(Serialize, Deserialize, Debug)]
+pub struct LocalAccount {
+    pub address: RoochAddress,
+    pub multichain_address: Option<MultiChainAddress>,
+    pub public_key: Option<PublicKey>,
+    pub has_session_key: bool,
+}
+
+#[derive(Serialize, Deserialize, Debug)]
+pub struct AddressMapping {
+    pub rooch_to_multichain: BTreeMap<RoochAddress, MultiChainAddress>,
+    pub multichain_to_rooch: BTreeMap<MultiChainAddress, RoochAddress>,
+}
+
+impl Default for AddressMapping {
+    fn default() -> Self {
+        Self {
+            rooch_to_multichain: BTreeMap::new(),
+            multichain_to_rooch: BTreeMap::new(),
+        }
+    }
+}
diff --git a/crates/rooch-rpc-server/src/server/btc_server.rs b/crates/rooch-rpc-server/src/server/btc_server.rs
@@ -50,7 +50,7 @@ impl BtcAPIServer for BtcServer {
         let resolve_address = match filter.clone() {
             Some(UTXOFilterView::Owner(address)) => {
                 let multi_chain_address = MultiChainAddress::try_from_str_with_multichain_id(
-                    RoochMultiChainID::Bitcoin.multichain_id().id(),
+                    RoochMultiChainID::Bitcoin,
                     address.to_string().as_str(),
                 )?;
                 self.rpc_service
@@ -104,7 +104,7 @@ impl BtcAPIServer for BtcServer {
         let resolve_address = match filter.clone() {
             Some(InscriptionFilterView::Owner(address)) => {
                 let multi_chain_address = MultiChainAddress::try_from_str_with_multichain_id(
-                    RoochMultiChainID::Bitcoin.multichain_id().id(),
+                    RoochMultiChainID::Bitcoin,
                     address.to_string().as_str(),
                 )?;
                 self.rpc_service