crypto: use SuiKeyPair enum instead of AccountKeyPair type (MystenLab…

…s#3801)
brucx · Aug 15, 2022 · 65efcc1 · 65efcc1
1 parent 3eb4227
commit 65efcc1
Show file tree

Hide file tree

Showing 16 changed files with 483 additions and 75 deletions.
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/crates/sui-benchmark/src/bin/stress.rs b/crates/sui-benchmark/src/bin/stress.rs
@@ -20,6 +20,7 @@ use sui_node::metrics;
 use sui_node::SuiNode;
 use sui_types::base_types::ObjectID;
 use sui_types::base_types::SuiAddress;
+use sui_types::crypto::AccountKeyPair;
 
 use std::collections::HashMap;
 use std::path::PathBuf;
@@ -34,9 +35,7 @@ use sui_benchmark::workloads::workload::Payload;
 use sui_benchmark::workloads::workload::Workload;
 use sui_core::epoch::epoch_store::EpochStore;
 use sui_sdk::crypto::FileBasedKeystore;
-use sui_types::crypto::AccountKeyPair;
 use sui_types::crypto::EncodeDecodeBase64;
-use sui_types::crypto::KeypairTraits;
 
 use test_utils::authority::spawn_test_authorities;
 use test_utils::authority::test_and_configure_authority_configs;
@@ -303,7 +302,7 @@ async fn main() -> Result<()> {
             .key_pairs()
             .iter()
             .find(|x| {
-                let address: SuiAddress = Into::<SuiAddress>::into(x.public());
+                let address: SuiAddress = Into::<SuiAddress>::into(&x.public());
                 address == primary_gas_account
             })
             .map(|x| x.encode_base64())

diff --git a/crates/sui-cluster-test/src/cluster.rs b/crates/sui-cluster-test/src/cluster.rs
@@ -15,6 +15,7 @@ use sui_swarm::memory::Node;
 use sui_swarm::memory::Swarm;
 use sui_types::base_types::SuiAddress;
 use sui_types::crypto::KeypairTraits;
+use sui_types::crypto::SuiKeyPair;
 use sui_types::crypto::{get_key_pair, AccountKeyPair};
 use test_utils::network::{start_rpc_test_network_with_fullnode, TestNetwork};
 use tracing::info;
@@ -246,7 +247,11 @@ pub async fn new_wallet_context_from_cluster(
     let keystore_path = temp_dir.path().join(SUI_KEYSTORE_FILENAME);
     let keystore = KeystoreType::File(keystore_path);
     let address: SuiAddress = key_pair.public().into();
-    keystore.init().unwrap().add_key(key_pair).unwrap();
+    keystore
+        .init()
+        .unwrap()
+        .add_key(SuiKeyPair::Ed25519SuiKeyPair(key_pair))
+        .unwrap();
     SuiClientConfig {
         keystore,
         gateway: ClientType::RPC(rpc_url.into()),

diff --git a/crates/sui-open-rpc/spec/openrpc.json b/crates/sui-open-rpc/spec/openrpc.json
@@ -3037,6 +3037,9 @@
           }
         }
       },
+      "Secp256k1SuiSignature": {
+        "$ref": "#/components/schemas/Base64"
+      },
       "SequenceNumber": {
         "type": "integer",
         "format": "uint64",
@@ -3055,6 +3058,18 @@
               }
             },
             "additionalProperties": false
+          },
+          {
+            "type": "object",
+            "required": [
+              "Secp256k1SuiSignature"
+            ],
+            "properties": {
+              "Secp256k1SuiSignature": {
+                "$ref": "#/components/schemas/Secp256k1SuiSignature"
+              }
+            },
+            "additionalProperties": false
           }
         ]
       },

diff --git a/crates/sui-sdk/src/crypto.rs b/crates/sui-sdk/src/crypto.rs
@@ -18,8 +18,7 @@ use rand::rngs::adapter::ReadRng;
 
 use sui_types::base_types::SuiAddress;
 use sui_types::crypto::{
-    get_key_pair_from_rng, AccountKeyPair, AccountPublicKey, EncodeDecodeBase64, KeypairTraits,
-    Signature,
+    get_key_pair_from_rng, EncodeDecodeBase64, PublicKey, Signature, SuiKeyPair,
 };
 
 #[derive(Serialize, Deserialize)]
@@ -32,8 +31,8 @@ pub enum KeystoreType {
 
 pub trait AccountKeystore: Send + Sync {
     fn sign(&self, address: &SuiAddress, msg: &[u8]) -> Result<Signature, signature::Error>;
-    fn add_key(&mut self, keypair: AccountKeyPair) -> Result<(), anyhow::Error>;
-    fn keys(&self) -> Vec<AccountPublicKey>;
+    fn add_key(&mut self, keypair: SuiKeyPair) -> Result<(), anyhow::Error>;
+    fn keys(&self) -> Vec<PublicKey>;
 }
 
 impl KeystoreType {
@@ -64,9 +63,9 @@ impl Display for KeystoreType {
     }
 }
 
-#[derive(Serialize, Deserialize, Default)]
+#[derive(Default)]
 pub struct FileBasedKeystore {
-    keys: BTreeMap<SuiAddress, AccountKeyPair>,
+    keys: BTreeMap<SuiAddress, SuiKeyPair>,
     path: Option<PathBuf>,
 }
 
@@ -80,15 +79,15 @@ impl AccountKeystore for FileBasedKeystore {
             .try_sign(msg)
     }
 
-    fn add_key(&mut self, keypair: AccountKeyPair) -> Result<(), anyhow::Error> {
-        let address: SuiAddress = keypair.public().into();
+    fn add_key(&mut self, keypair: SuiKeyPair) -> Result<(), anyhow::Error> {
+        let address: SuiAddress = (&keypair.public()).into();
         self.keys.insert(address, keypair);
         self.save()?;
         Ok(())
     }
 
-    fn keys(&self) -> Vec<AccountPublicKey> {
-        self.keys.values().map(|key| key.public().clone()).collect()
+    fn keys(&self) -> Vec<PublicKey> {
+        self.keys.values().map(|key| key.public()).collect()
     }
 }
 
@@ -100,8 +99,8 @@ impl FileBasedKeystore {
             kp_strings
                 .iter()
                 .map(|kpstr| {
-                    let key = AccountKeyPair::decode_base64(kpstr);
-                    key.map(|k| (k.public().into(), k))
+                    let key = SuiKeyPair::decode_base64(kpstr);
+                    key.map(|k| (Into::<SuiAddress>::into(&k.public()), k))
                 })
                 .collect::<Result<BTreeMap<_, _>, _>>()
                 .map_err(|_| anyhow::anyhow!("Invalid Keypair file"))?
@@ -125,7 +124,7 @@ impl FileBasedKeystore {
                 &self
                     .keys
                     .values()
-                    .map(|k| k.encode_base64())
+                    .map(EncodeDecodeBase64::encode_base64)
                     .collect::<Vec<_>>(),
             )
             .unwrap();
@@ -134,7 +133,7 @@ impl FileBasedKeystore {
         Ok(())
     }
 
-    pub fn key_pairs(&self) -> Vec<&AccountKeyPair> {
+    pub fn key_pairs(&self) -> Vec<&SuiKeyPair> {
         self.keys.values().collect()
     }
 }
@@ -146,20 +145,21 @@ impl SuiKeystore {
         Self(Box::new(keystore))
     }
 
-    pub fn add_key(&mut self, keypair: AccountKeyPair) -> Result<(), anyhow::Error> {
+    pub fn add_key(&mut self, keypair: SuiKeyPair) -> Result<(), anyhow::Error> {
         self.0.add_key(keypair)
     }
 
     pub fn generate_new_key(&mut self) -> Result<(SuiAddress, String), anyhow::Error> {
         let mnemonic = Mnemonic::generate(12)?;
         let seed = mnemonic.to_seed("");
         let mut rng = RngWrapper(ReadRng::new(&seed));
+        // TODO(joyqvq): add ability to generate as Secp256k1 keypair
         let (address, kp) = get_key_pair_from_rng(&mut rng);
-        self.0.add_key(kp)?;
+        self.0.add_key(SuiKeyPair::Ed25519SuiKeyPair(kp))?;
         Ok((address, mnemonic.to_string()))
     }
 
-    pub fn keys(&self) -> Vec<AccountPublicKey> {
+    pub fn keys(&self) -> Vec<PublicKey> {
         self.0.keys()
     }
 
@@ -175,7 +175,8 @@ impl SuiKeystore {
         let seed = &Mnemonic::from_str(phrase).unwrap().to_seed("");
         let mut rng = RngWrapper(ReadRng::new(seed));
         let (address, kp) = get_key_pair_from_rng(&mut rng);
-        self.0.add_key(kp)?;
+        // TODO(joyqvq): add ability to import as Secp256k1 keypair
+        self.0.add_key(SuiKeyPair::Ed25519SuiKeyPair(kp))?;
         Ok(address)
     }
 
@@ -226,9 +227,9 @@ impl Signer<Signature> for KeystoreSigner<'_> {
     }
 }
 
-#[derive(Serialize, Deserialize, Default)]
+#[derive(Default)]
 struct InMemKeystore {
-    keys: BTreeMap<SuiAddress, AccountKeyPair>,
+    keys: BTreeMap<SuiAddress, SuiKeyPair>,
 }
 
 impl AccountKeystore for InMemKeystore {
@@ -241,14 +242,14 @@ impl AccountKeystore for InMemKeystore {
             .try_sign(msg)
     }
 
-    fn add_key(&mut self, keypair: AccountKeyPair) -> Result<(), anyhow::Error> {
-        let address: SuiAddress = keypair.public().into();
+    fn add_key(&mut self, keypair: SuiKeyPair) -> Result<(), anyhow::Error> {
+        let address: SuiAddress = (&keypair.public()).into();
         self.keys.insert(address, keypair);
         Ok(())
     }
 
-    fn keys(&self) -> Vec<AccountPublicKey> {
-        self.keys.values().map(|key| key.public().clone()).collect()
+    fn keys(&self) -> Vec<PublicKey> {
+        self.keys.values().map(|key| key.public()).collect()
     }
 }
 
@@ -257,7 +258,8 @@ impl InMemKeystore {
         let mut rng = StdRng::from_seed([0; 32]);
         let keys = (0..initial_key_number)
             .map(|_| get_key_pair_from_rng(&mut rng))
-            .collect::<BTreeMap<SuiAddress, AccountKeyPair>>();
+            .map(|(ad, k)| (ad, SuiKeyPair::Ed25519SuiKeyPair(k)))
+            .collect::<BTreeMap<SuiAddress, SuiKeyPair>>();
 
         Self { keys }
     }
@@ -268,11 +270,11 @@ impl AccountKeystore for Box<dyn AccountKeystore> {
         (**self).sign(address, msg)
     }
 
-    fn add_key(&mut self, keypair: AccountKeyPair) -> Result<(), anyhow::Error> {
+    fn add_key(&mut self, keypair: SuiKeyPair) -> Result<(), anyhow::Error> {
         (**self).add_key(keypair)
     }
 
-    fn keys(&self) -> Vec<AccountPublicKey> {
+    fn keys(&self) -> Vec<PublicKey> {
         (**self).keys()
     }
 }
diff --git a/crates/sui-types/Cargo.toml b/crates/sui-types/Cargo.toml
@@ -36,6 +36,7 @@ strum = "^0.24"
 strum_macros = "^0.24"
 roaring = "0.9.0"
 enum_dispatch = "^0.3"
+eyre = "0.6.8"
 
 # This version is incompatible with ed25519-dalek
 rand_latest = { version = "0.8.5", package = "rand" }

diff --git a/crates/sui-types/src/base_types.rs b/crates/sui-types/src/base_types.rs
@@ -26,7 +26,9 @@ use sha2::Sha512;
 use sha3::Sha3_256;
 
 use crate::committee::EpochId;
-use crate::crypto::{AuthorityPublicKey, AuthorityPublicKeyBytes, KeypairTraits, SuiPublicKey};
+use crate::crypto::{
+    AuthorityPublicKey, AuthorityPublicKeyBytes, KeypairTraits, PublicKey, SuiPublicKey,
+};
 use crate::error::ExecutionError;
 use crate::error::ExecutionErrorKind;
 use crate::error::SuiError;
@@ -208,6 +210,19 @@ impl<T: SuiPublicKey> From<&T> for SuiAddress {
     }
 }
 
+impl From<&PublicKey> for SuiAddress {
+    fn from(pk: &PublicKey) -> Self {
+        let mut hasher = Sha3_256::default();
+        hasher.update(&[pk.flag()]);
+        hasher.update(pk);
+        let g_arr = hasher.finalize();
+
+        let mut res = [0u8; SUI_ADDRESS_LENGTH];
+        res.copy_from_slice(&AsRef::<[u8]>::as_ref(&g_arr)[..SUI_ADDRESS_LENGTH]);
+        SuiAddress(res)
+    }
+}
+
 impl TryFrom<&[u8]> for SuiAddress {
     type Error = SuiError;