Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

[cli] Support local address mapping #1277

Merged
merged 4 commits into from
Dec 29, 2023
Merged
Show file tree
Hide file tree
Changes from 3 commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
12 changes: 12 additions & 0 deletions crates/rooch-key/src/keystore/account_keystore.rs
Original file line number Diff line number Diff line change
@@ -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,
Expand All @@ -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,
Expand All @@ -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,
Expand All @@ -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,
Expand Down Expand Up @@ -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,
Expand Down
103 changes: 87 additions & 16 deletions crates/rooch-key/src/keystore/base_keystore.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand All @@ -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 {
Expand All @@ -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().collect())
}

fn get_key_pair_with_password(
&self,
address: &RoochAddress,
Expand Down Expand Up @@ -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(|| {
Expand All @@ -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);

Expand Down
15 changes: 15 additions & 0 deletions crates/rooch-key/src/keystore/file_keystore.rs
Original file line number Diff line number Diff line change
@@ -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;
Expand All @@ -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,
Expand Down Expand Up @@ -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,
Expand Down
13 changes: 13 additions & 0 deletions crates/rooch-key/src/keystore/memory_keystore.rs
Original file line number Diff line number Diff line change
@@ -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;
Expand All @@ -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,
Expand Down Expand Up @@ -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,
Expand Down
26 changes: 26 additions & 0 deletions crates/rooch-key/src/keystore/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand All @@ -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,
Expand Down Expand Up @@ -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(),
Expand Down
31 changes: 31 additions & 0 deletions crates/rooch-key/src/keystore/types.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,31 @@
// 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, Default, Debug)]
pub struct AddressMapping {
pub rooch_to_multichain: BTreeMap<RoochAddress, MultiChainAddress>,
pub multichain_to_rooch: BTreeMap<MultiChainAddress, RoochAddress>,
}
4 changes: 2 additions & 2 deletions crates/rooch-rpc-server/src/server/btc_server.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand Down
4 changes: 2 additions & 2 deletions crates/rooch-rpc-server/src/server/rooch_server.rs
Original file line number Diff line number Diff line change
Expand Up @@ -33,11 +33,11 @@ use rooch_rpc_api::{
api::{MAX_RESULT_LIMIT, MAX_RESULT_LIMIT_USIZE},
jsonrpc_types::BytesView,
};
use rooch_types::address::MultiChainAddress;
use rooch_types::indexer::event_filter::IndexerEventID;
use rooch_types::indexer::state::IndexerStateID;
use rooch_types::transaction::rooch::RoochTransaction;
use rooch_types::transaction::{AbstractTransaction, TypedTransaction};
use rooch_types::{address::MultiChainAddress, multichain_id::RoochMultiChainID};
use std::cmp::min;
use tracing::info;

Expand Down Expand Up @@ -438,7 +438,7 @@ impl RoochAPIServer for RoochServer {
address,
} => {
let multi_chain_address = MultiChainAddress::try_from_str_with_multichain_id(
multichain_id,
RoochMultiChainID::try_from(multichain_id)?,
address.as_str(),
)?;
self.rpc_service
Expand Down
Loading