Reapply keystore (#4999)

* [2/n] crypto: simplify keystore (#4909) * do not panic on none address
MystenLabs · Oct 5, 2022 · 79b27eb · 79b27eb · github-actions · Oct 5, 2022
1 parent 8eb9d03
commit 79b27eb
Show file tree

Hide file tree

Showing 28 changed files with 283 additions and 238 deletions.
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/crates/sui-benchmark/src/util.rs b/crates/sui-benchmark/src/util.rs
@@ -2,10 +2,10 @@
 // SPDX-License-Identifier: Apache-2.0
 
 use anyhow::Result;
-use sui_sdk::crypto::FileBasedKeystore;
+use sui_sdk::crypto::{AccountKeystore, FileBasedKeystore};
 use sui_types::{
     base_types::SuiAddress,
-    crypto::{AccountKeyPair, EncodeDecodeBase64, SuiKeyPair},
+    crypto::{AccountKeyPair, KeypairTraits, SuiKeyPair},
 };
 
 use std::path::PathBuf;
@@ -14,20 +14,9 @@ pub fn get_ed25519_keypair_from_keystore(
     keystore_path: PathBuf,
     requested_address: &SuiAddress,
 ) -> Result<AccountKeyPair> {
-    let keystore = FileBasedKeystore::load_or_create(&keystore_path)?;
-    let keypair = keystore
-        .key_pairs()
-        .iter()
-        .find(|x| {
-            let address: SuiAddress = Into::<SuiAddress>::into(&x.public());
-            address == *requested_address
-        })
-        .map(|x| x.encode_base64())
-        .unwrap();
-    // TODO(joyqvq): This is a hack to decode base64 keypair with added flag, ok for now since it is for benchmark use.
-    // Rework to get the typed keypair directly from above.
-    Ok(match SuiKeyPair::decode_base64(&keypair).unwrap() {
-        SuiKeyPair::Ed25519SuiKeyPair(x) => x,
-        _ => panic!("Unexpected keypair type"),
-    })
+    let keystore = FileBasedKeystore::new(&keystore_path)?;
+    match keystore.get_key(requested_address) {
+        Ok(SuiKeyPair::Ed25519SuiKeyPair(kp)) => Ok(kp.copy()),
+        _ => Err(anyhow::anyhow!("Unsupported key type")),
+    }
 }
diff --git a/crates/sui-cluster-test/src/cluster.rs b/crates/sui-cluster-test/src/cluster.rs
@@ -9,7 +9,9 @@ use sui::config::SuiClientConfig;
 use sui_config::genesis_config::GenesisConfig;
 use sui_config::Config;
 use sui_config::SUI_KEYSTORE_FILENAME;
-use sui_sdk::crypto::KeystoreType;
+use sui_sdk::crypto::AccountKeystore;
+use sui_sdk::crypto::FileBasedKeystore;
+use sui_sdk::crypto::Keystore;
 use sui_sdk::ClientType;
 use sui_swarm::memory::Node;
 use sui_swarm::memory::Swarm;
@@ -285,11 +287,9 @@ pub async fn new_wallet_context_from_cluster(
     let rpc_url = cluster.rpc_url();
     info!("Use gateway: {}", &rpc_url);
     let keystore_path = temp_dir.path().join(SUI_KEYSTORE_FILENAME);
-    let keystore = KeystoreType::File(keystore_path);
+    let mut keystore = Keystore::from(FileBasedKeystore::new(&keystore_path).unwrap());
     let address: SuiAddress = key_pair.public().into();
     keystore
-        .init()
-        .unwrap()
         .add_key(SuiKeyPair::Ed25519SuiKeyPair(key_pair))
         .unwrap();
     SuiClientConfig {

diff --git a/crates/sui-cluster-test/src/wallet_client.rs b/crates/sui-cluster-test/src/wallet_client.rs
@@ -5,6 +5,7 @@ use crate::cluster::new_wallet_context_from_cluster;
 
 use super::Cluster;
 use sui::client_commands::WalletContext;
+use sui_sdk::crypto::AccountKeystore;
 use sui_sdk::SuiClient;
 use sui_types::base_types::SuiAddress;
 use sui_types::crypto::{KeypairTraits, Signature};
@@ -68,6 +69,7 @@ impl WalletClient {
 
     pub fn sign(&self, txn_data: &TransactionData, desc: &str) -> Signature {
         self.get_wallet()
+            .config
             .keystore
             .sign(&self.address, &txn_data.to_bytes())
             .unwrap_or_else(|e| panic!("Failed to sign transaction for {}. {}", desc, e))

diff --git a/crates/sui-faucet/Cargo.toml b/crates/sui-faucet/Cargo.toml
@@ -28,6 +28,7 @@ sui-node = { path = "../sui-node" }
 sui-json-rpc-types= { path = "../sui-json-rpc-types" }
 sui-types = { path = "../sui-types" }
 sui-config = { path = "../sui-config" }
+sui-sdk = { path = "../sui-sdk" }
 telemetry-subscribers.workspace = true
 workspace-hack.workspace = true
 

diff --git a/crates/sui-faucet/src/faucet/simple_faucet.rs b/crates/sui-faucet/src/faucet/simple_faucet.rs
@@ -1,11 +1,11 @@
 // Copyright (c) 2022, Mysten Labs, Inc.
 // SPDX-License-Identifier: Apache-2.0
 
+use crate::metrics::FaucetMetrics;
 use anyhow::anyhow;
 use async_trait::async_trait;
-
-use crate::metrics::FaucetMetrics;
 use prometheus::Registry;
+use sui_sdk::crypto::AccountKeystore;
 
 // HashSet is in fact used but linter does not think so
 #[allow(unused_imports)]
@@ -259,7 +259,7 @@ impl SimpleFaucet {
             .construct_transfer_sui_txn_with_retry(coin_id, signer, recipient, budget, amount, uuid)
             .await?;
 
-        let signature = context.keystore.sign(&signer, &data.to_bytes())?;
+        let signature = context.config.keystore.sign(&signer, &data.to_bytes())?;
 
         let tx = Transaction::new(data, signature);
         info!(tx_digest = ?tx.digest(), ?recipient, ?coin_id, ?uuid, "Broadcasting transfer obj txn");

diff --git a/crates/sui-gateway/src/unit_tests/rpc_server_tests.rs b/crates/sui-gateway/src/unit_tests/rpc_server_tests.rs
@@ -11,7 +11,9 @@ use sui_json_rpc::api::{
     RpcGatewayApiClient, RpcReadApiClient, RpcTransactionBuilderClient, WalletSyncApiClient,
 };
 use sui_json_rpc_types::{GetObjectDataResponse, SuiTransactionResponse, TransactionBytes};
-use sui_sdk::crypto::KeystoreType;
+use sui_sdk::crypto::AccountKeystore;
+use sui_sdk::crypto::FileBasedKeystore;
+use sui_sdk::crypto::Keystore;
 use sui_types::base_types::ObjectID;
 use sui_types::base_types::TransactionDigest;
 use sui_types::gas_coin::GAS;
@@ -52,7 +54,7 @@ async fn test_public_transfer_object() -> Result<(), anyhow::Error> {
         .await?;
 
     let keystore_path = test_network.network.dir().join(SUI_KEYSTORE_FILENAME);
-    let keystore = KeystoreType::File(keystore_path).init()?;
+    let keystore = Keystore::from(FileBasedKeystore::new(&keystore_path)?);
 
     let signature = keystore.sign(address, &transaction_bytes.tx_bytes.to_vec()?)?;
     let tx = Transaction::new(transaction_bytes.to_data().unwrap(), signature);
@@ -91,7 +93,7 @@ async fn test_publish() -> Result<(), anyhow::Error> {
         .await?;
 
     let keystore_path = test_network.network.dir().join(SUI_KEYSTORE_FILENAME);
-    let keystore = KeystoreType::File(keystore_path).init()?;
+    let keystore = Keystore::from(FileBasedKeystore::new(&keystore_path)?);
     let signature = keystore.sign(address, &transaction_bytes.tx_bytes.to_vec()?)?;
 
     let tx = Transaction::new(transaction_bytes.to_data().unwrap(), signature);
@@ -139,7 +141,7 @@ async fn test_move_call() -> Result<(), anyhow::Error> {
         .await?;
 
     let keystore_path = test_network.network.dir().join(SUI_KEYSTORE_FILENAME);
-    let keystore = KeystoreType::File(keystore_path).init()?;
+    let keystore = Keystore::from(FileBasedKeystore::new(&keystore_path)?);
 
     let signature = keystore.sign(address, &transaction_bytes.tx_bytes.to_vec()?)?;
 
@@ -192,7 +194,7 @@ async fn test_get_transaction() -> Result<(), anyhow::Error> {
             .await?;
 
         let keystore_path = test_network.network.dir().join(SUI_KEYSTORE_FILENAME);
-        let keystore = KeystoreType::File(keystore_path).init()?;
+        let keystore = Keystore::from(FileBasedKeystore::new(&keystore_path)?);
 
         let signature = keystore.sign(address, &transaction_bytes.tx_bytes.to_vec()?)?;
 

diff --git a/crates/sui-open-rpc/Cargo.toml b/crates/sui-open-rpc/Cargo.toml
@@ -25,6 +25,7 @@ sui-json-rpc = { path = "../sui-json-rpc" }
 sui-json-rpc-types = { path = "../sui-json-rpc-types" }
 sui-json = { path = "../sui-json" }
 sui-types = { path = "../sui-types" }
+sui-sdk = { path = "../sui-sdk" }
 sui-config = { path = "../sui-config" }
 test-utils = { path = "../test-utils" }
 rand = "0.8.5"

diff --git a/crates/sui-open-rpc/src/generate_json_rpc_spec.rs b/crates/sui-open-rpc/src/generate_json_rpc_spec.rs
@@ -1,10 +1,6 @@
 // Copyright (c) 2022, Mysten Labs, Inc.
 // SPDX-License-Identifier: Apache-2.0
 
-use std::collections::BTreeMap;
-use std::fs::File;
-use std::io::Write;
-
 use clap::ArgEnum;
 use clap::Parser;
 use hyper::body::Buf;
@@ -13,7 +9,11 @@ use move_package::BuildConfig;
 use pretty_assertions::assert_str_eq;
 use serde::Serialize;
 use serde_json::{json, Map, Value};
+use std::collections::BTreeMap;
+use std::fs::File;
+use std::io::Write;
 use sui_json_rpc::api::EventReadApiOpenRpc;
+use sui_sdk::crypto::AccountKeystore;
 use sui_types::messages::Transaction;
 
 use crate::examples::RpcExampleProvider;
@@ -136,7 +136,13 @@ async fn create_response_sample() -> Result<
     let config = working_dir.join(SUI_CLIENT_CONFIG);
 
     let mut context = WalletContext::new(&config).await?;
-    let address = context.keystore.addresses().first().cloned().unwrap();
+    let address = context
+        .config
+        .keystore
+        .addresses()
+        .first()
+        .cloned()
+        .unwrap();
 
     context
         .client
@@ -403,6 +409,7 @@ async fn create_error_response(
         .await?;
 
     let signature = context
+        .config
         .keystore
         .sign(&address, &response.tx_bytes.to_vec()?)?;
 

diff --git a/crates/sui-rosetta/Cargo.toml b/crates/sui-rosetta/Cargo.toml
@@ -30,6 +30,7 @@ sui-types = { path = "../sui-types" }
 sui-core = { path = "../sui-core" }
 sui-node = { path = "../sui-node" }
 sui-config = { path = "../sui-config" }
+sui-sdk = { path = "../sui-sdk" }
 
 move-core-types.workspace = true
 

diff --git a/crates/sui-rosetta/src/main.rs b/crates/sui-rosetta/src/main.rs
@@ -170,12 +170,12 @@ fn read_prefunded_account(path: &Path) -> Result<Vec<PrefundedAccount>, anyhow::
 
 #[test]
 fn test_read_keystore() {
-    use sui_sdk::crypto::KeystoreType;
+    use sui_sdk::crypto::{AccountKeystore, FileBasedKeystore, Keystore};
     use sui_types::crypto::SignatureScheme;
 
     let temp_dir = tempfile::tempdir().unwrap();
     let path = temp_dir.path().join("sui.keystore");
-    let mut ks = KeystoreType::File(path.clone()).init().unwrap();
+    let mut ks = Keystore::from(FileBasedKeystore::new(&path).unwrap());
     let key1 = ks.generate_new_key(SignatureScheme::ED25519, None).unwrap();
     let key2 = ks
         .generate_new_key(SignatureScheme::Secp256k1, None)

diff --git a/crates/sui-sdk/examples/tic_tac_toe.rs b/crates/sui-sdk/examples/tic_tac_toe.rs
@@ -13,8 +13,9 @@ use clap::Parser;
 use clap::Subcommand;
 use serde::Deserialize;
 
+use sui_sdk::crypto::{AccountKeystore, FileBasedKeystore};
 use sui_sdk::{
-    crypto::{KeystoreType, SuiKeystore},
+    crypto::Keystore,
     json::SuiJsonValue,
     rpc_types::SuiData,
     types::{
@@ -29,7 +30,7 @@ use sui_sdk::{
 async fn main() -> Result<(), anyhow::Error> {
     let opts: TicTacToeOpts = TicTacToeOpts::parse();
     let keystore_path = opts.keystore_path.unwrap_or_else(default_keystore_path);
-    let keystore = KeystoreType::File(keystore_path).init()?;
+    let keystore = Keystore::File(FileBasedKeystore::new(&keystore_path)?);
 
     let game = TicTacToe {
         game_package_id: opts.game_package_id,
@@ -55,7 +56,7 @@ async fn main() -> Result<(), anyhow::Error> {
 struct TicTacToe {
     game_package_id: ObjectID,
     client: SuiClient,
-    keystore: SuiKeystore,
+    keystore: Keystore,
 }
 
 impl TicTacToe {

diff --git a/crates/sui-sdk/examples/transfer_coins.rs b/crates/sui-sdk/examples/transfer_coins.rs
@@ -3,7 +3,7 @@
 
 use std::str::FromStr;
 use sui_sdk::{
-    crypto::KeystoreType,
+    crypto::{AccountKeystore, FileBasedKeystore, Keystore},
     types::{
         base_types::{ObjectID, SuiAddress},
         messages::Transaction,
@@ -30,8 +30,8 @@ async fn main() -> Result<(), anyhow::Error> {
         .transfer_sui(my_address, gas_object_id, 1000, recipient, Some(1000))
         .await?;
 
-    // Get signer from keystore
-    let keystore = KeystoreType::File(keystore_path).init()?;
+    // Sign transaction
+    let keystore = Keystore::from(FileBasedKeystore::new(&keystore_path)?);
     let signature = keystore.sign(&my_address, &transfer_tx.to_bytes())?;
 
     // Execute the transaction