Skip to content

Commit

Permalink
[fastx] Replace KeyPair With Signer & Load Keys From File (MystenLabs…
Browse files Browse the repository at this point in the history
…#590)

* Replace KeyPair With Signer & Load Keys From File

* Fixup after rebase

* Add license header

* impl TryFrom<Vec<u8>> instead of From<Vec<u8>> for SuiAddress
  • Loading branch information
patrickkuo authored Mar 1, 2022
1 parent f53fbe9 commit 0da6369
Show file tree
Hide file tree
Showing 10 changed files with 171 additions and 65 deletions.
2 changes: 1 addition & 1 deletion sui/Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -25,7 +25,7 @@ ed25519-dalek = { version = "1.0.1", features = ["batch", "serde"] }
rocksdb = "0.17.0"
hex = "0.4.3"
async-trait = "0.1.52"
serde_with = "1.11.0"
serde_with = { version = "1.11.0", features = ["hex"] }
tracing = { version = "0.1", features = ["log"] }
tracing-subscriber = { version = "0.3", features = ["time", "env-filter"] }
serde-value = "0.7.0"
Expand Down
28 changes: 12 additions & 16 deletions sui/src/config.rs
Original file line number Diff line number Diff line change
Expand Up @@ -2,11 +2,14 @@
// Copyright (c) 2022, Mysten Labs, Inc.
// SPDX-License-Identifier: Apache-2.0

use crate::keystore::KeystoreType;
use anyhow::anyhow;
use once_cell::sync::Lazy;
use serde::de::DeserializeOwned;
use serde::{Deserialize, Serialize};
use serde_json::Value;
use serde_with::hex::Hex;
use serde_with::serde_as;
use std::collections::BTreeMap;
use std::fmt::{Display, Formatter};
use std::fs::{self, File};
Expand All @@ -21,7 +24,6 @@ use sui_network::transport;
use sui_types::base_types::*;
use sui_types::committee::Committee;
use sui_types::crypto::{get_key_pair, KeyPair};
use sui_types::error::SuiError;
use tracing::log::trace;

const DEFAULT_WEIGHT: usize = 1;
Expand All @@ -32,13 +34,6 @@ pub const DEFAULT_STARTING_PORT: u16 = 10000;
static PORT_ALLOCATOR: Lazy<Mutex<PortAllocator>> =
Lazy::new(|| Mutex::new(PortAllocator::new(DEFAULT_STARTING_PORT)));

#[derive(Serialize, Deserialize)]
pub struct AccountInfo {
#[serde(serialize_with = "bytes_as_hex", deserialize_with = "bytes_from_hex")]
pub address: SuiAddress,
pub key_pair: KeyPair,
}

#[derive(Serialize, Deserialize)]
pub struct AuthorityInfo {
#[serde(serialize_with = "bytes_as_hex", deserialize_with = "bytes_from_hex")]
Expand Down Expand Up @@ -107,28 +102,35 @@ impl<'de> Deserialize<'de> for AuthorityPrivateInfo {
}
}

#[serde_as]
#[derive(Serialize, Deserialize)]
pub struct WalletConfig {
pub accounts: Vec<AccountInfo>,
#[serde_as(as = "Vec<Hex>")]
pub accounts: Vec<SuiAddress>,
pub authorities: Vec<AuthorityInfo>,
pub send_timeout: Duration,
pub recv_timeout: Duration,
pub buffer_size: usize,
pub db_folder_path: PathBuf,
pub keystore: KeystoreType,

#[serde(skip)]
config_path: PathBuf,
}

impl Config for WalletConfig {
fn create(path: &Path) -> Result<Self, anyhow::Error> {
let working_dir = path
.parent()
.ok_or(anyhow!("Cannot determine parent directory."))?;
Ok(WalletConfig {
accounts: Vec::new(),
authorities: Vec::new(),
send_timeout: Duration::from_micros(4000000),
recv_timeout: Duration::from_micros(4000000),
buffer_size: transport::DEFAULT_MAX_DATAGRAM_SIZE.to_string().parse()?,
db_folder_path: PathBuf::from("./client_db"),
db_folder_path: working_dir.join("client_db"),
keystore: KeystoreType::File(working_dir.join("wallet.ks")),
config_path: path.to_path_buf(),
})
}
Expand Down Expand Up @@ -166,12 +168,6 @@ impl WalletConfig {
}
authority_clients
}
pub fn get_account_cfg_info(&self, address: &SuiAddress) -> Result<&AccountInfo, SuiError> {
self.accounts
.iter()
.find(|info| &info.address == address)
.ok_or(SuiError::AccountNotFound)
}
}

impl Display for WalletConfig {
Expand Down
104 changes: 104 additions & 0 deletions sui/src/keystore.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,104 @@
// Copyright (c) 2022, Mysten Labs, Inc.
// SPDX-License-Identifier: Apache-2.0

use anyhow::anyhow;
use ed25519_dalek::ed25519::signature;
use ed25519_dalek::{ed25519, Signer};
use serde::{Deserialize, Serialize};
use std::collections::BTreeMap;
use std::fs;
use std::fs::File;
use std::io::BufReader;
use std::path::{Path, PathBuf};
use std::sync::{Arc, Mutex};
use sui_types::base_types::SuiAddress;
use sui_types::crypto::{KeyPair, Signature};

#[derive(Serialize, Deserialize)]
pub enum KeystoreType {
File(PathBuf),
}

pub trait Keystore: Send + Sync {
fn sign(&self, address: &SuiAddress, msg: &[u8]) -> Result<Signature, anyhow::Error>;
fn add_key(&mut self, keypair: KeyPair) -> Result<(), anyhow::Error>;
}

impl KeystoreType {
pub fn init(&self) -> Result<Box<dyn Keystore>, anyhow::Error> {
Ok(match self {
KeystoreType::File(path) => Box::new(SuiKeystore::load_or_create(path)?),
})
}
}

#[derive(Serialize, Deserialize)]
pub struct SuiKeystore {
keys: BTreeMap<SuiAddress, KeyPair>,
path: PathBuf,
}

impl Keystore for SuiKeystore {
fn sign(&self, address: &SuiAddress, msg: &[u8]) -> Result<Signature, anyhow::Error> {
Ok(self
.keys
.get(address)
.ok_or_else(|| anyhow!("Cannot find key for address: [{}]", address))?
.sign(msg))
}
fn add_key(&mut self, keypair: KeyPair) -> Result<(), anyhow::Error> {
self.keys
.insert(SuiAddress::from(keypair.public_key_bytes()), keypair);
self.save()
}
}

impl SuiKeystore {
pub fn load_or_create(path: &Path) -> Result<Self, anyhow::Error> {
let keys: Vec<KeyPair> = if path.exists() {
let reader = BufReader::new(File::open(path)?);
serde_json::from_reader(reader)?
} else {
Vec::new()
};

let keys = keys
.into_iter()
.map(|key| (SuiAddress::from(key.public_key_bytes()), key))
.collect();

Ok(Self {
keys,
path: path.to_path_buf(),
})
}

pub fn save(&self) -> Result<(), anyhow::Error> {
let store = serde_json::to_string_pretty(&self.keys.values().collect::<Vec<_>>()).unwrap();
Ok(fs::write(&self.path, store)?)
}
}

pub struct SuiKeystoreSigner {
keystore: Arc<Mutex<Box<dyn Keystore>>>,
address: SuiAddress,
}

impl SuiKeystoreSigner {
pub fn new(keystore: Arc<Mutex<Box<dyn Keystore>>>, account: SuiAddress) -> Self {
Self {
keystore,
address: account,
}
}
}

impl signature::Signer<Signature> for SuiKeystoreSigner {
fn try_sign(&self, msg: &[u8]) -> Result<Signature, ed25519::Error> {
self.keystore
.lock()
.unwrap()
.sign(&self.address, msg)
.map_err(ed25519::Error::from_source)
}
}
1 change: 1 addition & 0 deletions sui/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,7 @@
// SPDX-License-Identifier: Apache-2.0

pub mod config;
pub mod keystore;
pub mod shell;
pub mod sui_commands;
pub mod sui_json;
Expand Down
6 changes: 1 addition & 5 deletions sui/src/rest_server.rs
Original file line number Diff line number Diff line change
Expand Up @@ -311,11 +311,7 @@ async fn sui_start(
)
})?;

let addresses = wallet_config
.accounts
.iter()
.map(|info| info.address)
.collect::<Vec<_>>();
let addresses = wallet_config.accounts.clone();
let mut wallet_context = WalletContext::new(wallet_config).map_err(|error| {
custom_http_error(
StatusCode::CONFLICT,
Expand Down
11 changes: 8 additions & 3 deletions sui/src/sui_commands.rs
Original file line number Diff line number Diff line change
@@ -1,8 +1,7 @@
// Copyright (c) 2022, Mysten Labs, Inc.
// SPDX-License-Identifier: Apache-2.0
use crate::config::{
AccountInfo, AuthorityInfo, AuthorityPrivateInfo, Config, GenesisConfig, NetworkConfig,
WalletConfig,
AuthorityInfo, AuthorityPrivateInfo, Config, GenesisConfig, NetworkConfig, WalletConfig,
};
use anyhow::anyhow;
use futures::future::join_all;
Expand All @@ -17,6 +16,7 @@ use sui_core::authority::{AuthorityState, AuthorityStore};
use sui_core::authority_server::AuthorityServer;
use sui_types::base_types::{SequenceNumber, TxContext};

use crate::keystore::KeystoreType;
use sui_adapter::adapter::generate_package_id;
use sui_types::committee::Committee;
use sui_types::crypto::get_key_pair;
Expand Down Expand Up @@ -52,6 +52,7 @@ impl SuiCommand {
let wallet_path = working_dir.join("wallet.conf");
let mut wallet_config = WalletConfig::create(&wallet_path)?;
wallet_config.db_folder_path = working_dir.join("client_db");
wallet_config.keystore = KeystoreType::File(working_dir.join("wallet.key"));
genesis(config, genesis_conf, &mut wallet_config).await
}
}
Expand Down Expand Up @@ -140,12 +141,16 @@ pub async fn genesis(
"Creating {} account(s) and gas objects...",
new_account_count
);

let mut keystore = wallet_config.keystore.init()?;

for account in genesis_conf.accounts {
let address = if let Some(address) = account.address {
address
} else {
let (address, key_pair) = get_key_pair();
new_addresses.push(AccountInfo { address, key_pair });
new_addresses.push(address);
keystore.add_key(key_pair)?;
address
};
for object_conf in account.gas_objects {
Expand Down
Loading

0 comments on commit 0da6369

Please sign in to comment.