Skip to content

Commit

Permalink
Add alias support for sui addresses in the CLI
Browse files Browse the repository at this point in the history
  • Loading branch information
@stefan-mysten
stefan-mysten committed Sep 11, 2023
1 parent df0b128 commit 24f360e
Show file tree
Hide file tree
Showing 12 changed files with 533 additions and 93 deletions.
4 changes: 3 additions & 1 deletion crates/sui-cluster-test/src/cluster.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -317,7 +317,9 @@ pub async fn new_wallet_context_from_cluster(
let keystore_path = config_dir.join(SUI_KEYSTORE_FILENAME);
let mut keystore = Keystore::from(FileBasedKeystore::new(&keystore_path).unwrap());
let address: SuiAddress = key_pair.public().into();
keystore.add_key(SuiKeyPair::Ed25519(key_pair)).unwrap();
keystore
.add_key(None, SuiKeyPair::Ed25519(key_pair))
.unwrap();
SuiClientConfig {
keystore,
envs: vec![SuiEnv {
Expand Down
223 changes: 180 additions & 43 deletions crates/sui-keys/src/keystore.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,7 @@
// SPDX-License-Identifier: Apache-2.0

use crate::key_derive::{derive_key_pair_from_path, generate_new_key};
use crate::random_names::random_name;
use anyhow::anyhow;
use bip32::DerivationPath;
use bip39::{Language, Mnemonic, Seed};
Expand Down Expand Up @@ -29,12 +30,14 @@ pub enum Keystore {
}
#[enum_dispatch]
pub trait AccountKeystore: Send + Sync {
fn add_key(&mut self, keypair: SuiKeyPair) -> Result<(), anyhow::Error>;
fn add_key(&mut self, alias: Option<String>, keypair: SuiKeyPair) -> Result<(), anyhow::Error>;
fn keys(&self) -> Vec<PublicKey>;
fn keys_with_aliases(&self) -> Vec<(SuiAddressAlias, PublicKey)>;
fn get_key(&self, address: &SuiAddress) -> Result<&SuiKeyPair, anyhow::Error>;

fn get_aliases(&self) -> Vec<SuiAddressAlias>;
fn sign_hashed(&self, address: &SuiAddress, msg: &[u8]) -> Result<Signature, signature::Error>;

fn alias_exists(&self, alias: &Option<String>) -> Result<(), anyhow::Error>;
fn update_alias(&mut self, old_alias: &str, new_alias: &str) -> Result<(), anyhow::Error>;
fn sign_secure<T>(
&self,
address: &SuiAddress,
Expand All @@ -47,30 +50,40 @@ pub trait AccountKeystore: Send + Sync {
self.keys().iter().map(|k| k.into()).collect()
}

/// Return the Sui addressess and their aliases
fn addresses_with_aliases(&self) -> Vec<(SuiAddressAlias, SuiAddress)> {
self.keys_with_aliases()
.iter()
.map(|(alias, pk)| (alias.into(), pk.into()))
.collect()
}

fn generate_and_add_new_key(
&mut self,
key_scheme: SignatureScheme,
alias: Option<String>,
derivation_path: Option<DerivationPath>,
word_length: Option<String>,
) -> Result<(SuiAddress, String, SignatureScheme), anyhow::Error> {
let (address, kp, scheme, phrase) =
generate_new_key(key_scheme, derivation_path, word_length)?;
self.add_key(kp)?;
self.add_key(alias, kp)?;
Ok((address, phrase, scheme))
}

fn import_from_mnemonic(
&mut self,
phrase: &str,
key_scheme: SignatureScheme,
alias: Option<String>,
derivation_path: Option<DerivationPath>,
) -> Result<SuiAddress, anyhow::Error> {
let mnemonic = Mnemonic::from_phrase(phrase, Language::English)
.map_err(|e| anyhow::anyhow!("Invalid mnemonic phrase: {:?}", e))?;
let seed = Seed::new(&mnemonic, "");
match derive_key_pair_from_path(seed.as_bytes(), derivation_path, &key_scheme) {
Ok((address, kp)) => {
self.add_key(kp)?;
self.add_key(alias, kp)?;
Ok(address)
}
Err(e) => Err(anyhow!("error getting keypair {:?}", e)),
Expand All @@ -95,9 +108,17 @@ impl Display for Keystore {
}
}

pub type SuiAddressAlias = String;

#[derive(Serialize, Deserialize)]
pub struct KeystoreAliasPrivateKey {
pub alias: String,
pub private_key_base64: String,
}

#[derive(Default)]
pub struct FileBasedKeystore {
keys: BTreeMap<SuiAddress, SuiKeyPair>,
keys: BTreeMap<SuiAddress, (SuiAddressAlias, SuiKeyPair)>,
path: Option<PathBuf>,
}

Expand Down Expand Up @@ -129,12 +150,14 @@ impl<'de> Deserialize<'de> for FileBasedKeystore {

impl AccountKeystore for FileBasedKeystore {
fn sign_hashed(&self, address: &SuiAddress, msg: &[u8]) -> Result<Signature, signature::Error> {
Ok(Signature::new_hashed(
msg,
self.keys.get(address).ok_or_else(|| {
let (_, key) = self
.keys
.get(address)
.ok_or_else(|| {
signature::Error::from_source(format!("Cannot find key for address: [{address}]"))
})?,
))
})?
.clone();
Ok(Signature::new_hashed(msg, key))
}
fn sign_secure<T>(
&self,
Expand All @@ -145,31 +168,80 @@ impl AccountKeystore for FileBasedKeystore {
where
T: Serialize,
{
Ok(Signature::new_secure(
&IntentMessage::new(intent, msg),
self.keys.get(address).ok_or_else(|| {
signature::Error::from_source(format!("Cannot find key for address: [{address}]"))
})?,
))
let (_, key) = self.keys.get(address).ok_or_else(|| {
signature::Error::from_source(format!("Cannot find key for address: [{address}]"))
})?;
Ok(Signature::new_secure(&IntentMessage::new(intent, msg), key))
}

fn add_key(&mut self, keypair: SuiKeyPair) -> Result<(), anyhow::Error> {
/// Add a key to the keystore
fn add_key(&mut self, alias: Option<String>, keypair: SuiKeyPair) -> Result<(), anyhow::Error> {
self.alias_exists(&alias)?;
let alias = alias.unwrap_or_else(|| random_name(self.get_aliases()));
let address: SuiAddress = (&keypair.public()).into();
self.keys.insert(address, keypair);
self.keys.insert(address, (alias, keypair));
self.save()?;
Ok(())
}

/// Return the public keys in the keystore
fn keys(&self) -> Vec<PublicKey> {
self.keys.values().map(|key| key.public()).collect()
self.keys.values().map(|(_, key)| key.public()).collect()
}

/// Return a list of pairs with the alias and the public key
fn keys_with_aliases(&self) -> Vec<(SuiAddressAlias, PublicKey)> {
self.keys
.values()
.map(|(alias, key)| (alias.clone(), key.public()))
.collect()
}

/// Return the corresponding SuiKeyPair of the provided sui address
fn get_key(&self, address: &SuiAddress) -> Result<&SuiKeyPair, anyhow::Error> {
match self.keys.get(address) {
Some(key) => Ok(key),
Some((_, key)) => Ok(key),
None => Err(anyhow!("Cannot find key for address: [{address}]")),
}
}

/// Return the list of aliases in the keystore
fn get_aliases(&self) -> Vec<SuiAddressAlias> {
self.keys_with_aliases()
.into_iter()
.map(|(alias, _)| alias)
.collect()
}

/// Return an error if the provided alias exists already
fn alias_exists(&self, alias: &Option<String>) -> Result<(), anyhow::Error> {
if let Some(alias) = alias {
if self.get_aliases().contains(alias) {
anyhow::bail!("Alias {alias} already exists. Please choose another alias.")
}
}
Ok(())
}

/// Update the old alias to the new alias
fn update_alias(&mut self, old_alias: &str, new_alias: &str) -> Result<(), anyhow::Error> {
if !self.get_aliases().contains(&old_alias.to_string()) {
anyhow::bail!("Alias {old_alias} does not exist in keystore.")
} else {
self.keys.values_mut().for_each(|x| {
if x.0 == old_alias {
x.0 = new_alias.to_string()
}
});
self.save()?;
Ok(())
}
}

/// Return a list of the addresses in the keystore
fn addresses(&self) -> Vec<SuiAddress> {
self.keys().iter().map(|k| k.into()).collect()
}
}

impl FileBasedKeystore {
Expand All @@ -179,15 +251,19 @@ impl FileBasedKeystore {
File::open(path)
.map_err(|e| anyhow!("Can't open FileBasedKeystore from {:?}: {e}", path))?,
);
let kp_strings: Vec<String> = serde_json::from_reader(reader)
let kp_strings: Vec<KeystoreAliasPrivateKey> = serde_json::from_reader(reader)
.map_err(|e| anyhow!("Can't deserialize FileBasedKeystore from {:?}: {e}", path))?;
kp_strings
.iter()
.map(|kpstr| {
let key = SuiKeyPair::decode_base64(kpstr);
key.map(|k| (Into::<SuiAddress>::into(&k.public()), k))
.map(|x| {
let key = SuiKeyPair::decode_base64(&x.private_key_base64);
key.map(|k| {
let sui_address = Into::<SuiAddress>::into(&k.public());
let value = (x.alias.clone(), k);
(sui_address, value)
})
})
.collect::<Result<BTreeMap<_, _>, _>>()
.collect::<Result<BTreeMap<SuiAddress, (String, SuiKeyPair)>, _>>()
.map_err(|e| anyhow::anyhow!("Invalid Keypair file {:#?} {:?}", e, path))?
} else {
BTreeMap::new()
Expand All @@ -209,7 +285,14 @@ impl FileBasedKeystore {
&self
.keys
.values()
.map(EncodeDecodeBase64::encode_base64)
.map(|(alias, key)| {
let encoded_key = EncodeDecodeBase64::encode_base64(key);

KeystoreAliasPrivateKey {
alias: alias.clone(),
private_key_base64: encoded_key,
}
})
.collect::<Vec<_>>(),
)
.unwrap();
Expand All @@ -219,23 +302,26 @@ impl FileBasedKeystore {
}

pub fn key_pairs(&self) -> Vec<&SuiKeyPair> {
self.keys.values().collect()
self.keys.values().map(|(_, key)| key).collect()
}
}

#[derive(Default, Serialize, Deserialize)]
pub struct InMemKeystore {
keys: BTreeMap<SuiAddress, SuiKeyPair>,
keys: BTreeMap<SuiAddress, (SuiAddressAlias, SuiKeyPair)>,
}

impl AccountKeystore for InMemKeystore {
fn sign_hashed(&self, address: &SuiAddress, msg: &[u8]) -> Result<Signature, signature::Error> {
Ok(Signature::new_hashed(
msg,
self.keys.get(address).ok_or_else(|| {
let (_, key) = self
.keys
.get(address)
.ok_or_else(|| {
signature::Error::from_source(format!("Cannot find key for address: [{address}]"))
})?,
))
})?
.clone();

Ok(Signature::new_hashed(msg, key))
}
fn sign_secure<T>(
&self,
Expand All @@ -246,39 +332,90 @@ impl AccountKeystore for InMemKeystore {
where
T: Serialize,
{
let key = self
.keys
.get(address)
.ok_or_else(|| {
signature::Error::from_source(format!("Cannot find key for address: [{address}]"))
})?
.clone();

Ok(Signature::new_secure(
&IntentMessage::new(intent, msg),
self.keys.get(address).ok_or_else(|| {
signature::Error::from_source(format!("Cannot find key for address: [{address}]"))
})?,
&key.1,
))
}

fn add_key(&mut self, keypair: SuiKeyPair) -> Result<(), anyhow::Error> {
/// Add a key to the keystore
fn add_key(&mut self, alias: Option<String>, keypair: SuiKeyPair) -> Result<(), anyhow::Error> {
self.alias_exists(&alias)?;
let alias = alias.unwrap_or_else(|| random_name(self.get_aliases()));
let address: SuiAddress = (&keypair.public()).into();
self.keys.insert(address, keypair);
self.keys.insert(address, (alias, keypair));
Ok(())
}

/// Return the public keys in the keystore
fn keys(&self) -> Vec<PublicKey> {
self.keys.values().map(|key| key.public()).collect()
self.keys.values().map(|(_, key)| key.public()).collect()
}

/// Return a list of pairs with the alias and the public key
fn keys_with_aliases(&self) -> Vec<(SuiAddressAlias, PublicKey)> {
self.keys
.values()
.map(|(alias, key)| (alias.clone(), key.public()))
.collect()
}

/// Return the corresponding SuiKeyPair of the provided sui address
fn get_key(&self, address: &SuiAddress) -> Result<&SuiKeyPair, anyhow::Error> {
match self.keys.get(address) {
Some(key) => Ok(key),
Some((_, key)) => Ok(key),
None => Err(anyhow!("Cannot find key for address: [{address}]")),
}
}

/// Return the list of aliases in the keystore
fn get_aliases(&self) -> Vec<SuiAddressAlias> {
self.keys_with_aliases()
.into_iter()
.map(|(alias, _)| alias)
.collect()
}

/// Return an error if the provided alias exists already
fn alias_exists(&self, alias: &Option<String>) -> Result<(), anyhow::Error> {
if let Some(alias) = alias {
if self.get_aliases().contains(alias) {
anyhow::bail!("Alias {alias} already exists. Please choose another alias.")
}
}
Ok(())
}

/// Update the old alias to the new alias
fn update_alias(&mut self, old_alias: &str, new_alias: &str) -> Result<(), anyhow::Error> {
if !self.get_aliases().contains(&old_alias.to_string()) {
anyhow::bail!("Alias {old_alias} does not exist in keystore.")
} else {
self.keys.values_mut().for_each(|x| {
if x.0 == old_alias {
x.0 = new_alias.to_string()
}
});
Ok(())
}
}
}

impl InMemKeystore {
pub fn new_insecure_for_tests(initial_key_number: usize) -> Self {
let mut rng = StdRng::from_seed([0; 32]);
let keys = (0..initial_key_number)
.map(|_| get_key_pair_from_rng(&mut rng))
.map(|(ad, k)| (ad, SuiKeyPair::Ed25519(k)))
.collect::<BTreeMap<SuiAddress, SuiKeyPair>>();
.map(|(ad, k)| (ad, (random_name(vec![]), SuiKeyPair::Ed25519(k))))
.collect::<BTreeMap<SuiAddress, (String, SuiKeyPair)>>();

Self { keys }
}
Expand Down
1 change: 1 addition & 0 deletions crates/sui-keys/src/lib.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -4,3 +4,4 @@
pub mod key_derive;
pub mod keypair_file;
pub mod keystore;
pub mod random_names;
Loading

0 comments on commit 24f360e

Please sign in to comment.