Feature/client manage multi addresses: Part 1 (#428)

* Add address manager

sui/src/wallet.rs

@@ -56,7 +56,7 @@ async fn main() -> Result<(), anyhow::Error> {
         .iter()
         .map(|info| info.address)
         .collect::<Vec<_>>();
-    let mut context = WalletContext::new(config);
+    let mut context = WalletContext::new(config)?;
 
     // Sync all accounts on start up.
     for address in addresses {

sui/src/wallet_commands.rs

@@ -2,7 +2,7 @@
 // SPDX-License-Identifier: Apache-2.0
 use crate::config::{AccountInfo, AuthorityInfo, WalletConfig};
 use sui_core::authority_client::AuthorityClient;
-use sui_core::client::{Client, ClientState};
+use sui_core::client::{Client, ClientAddressManager, ClientState};
 use sui_network::network::NetworkClient;
 use sui_types::base_types::{
     decode_address_hex, encode_address_hex, get_key_pair, AuthorityName, ObjectID, PublicKeyBytes,
@@ -18,6 +18,7 @@ use move_core_types::parser::{parse_transaction_argument, parse_type_tag};
 use move_core_types::transaction_argument::{convert_txn_args, TransactionArgument};
 use std::collections::BTreeMap;
 use std::path::PathBuf;
+use std::str::FromStr;
 use std::time::{Duration, Instant};
 use structopt::clap::AppSettings;
 use structopt::StructOpt;
@@ -35,6 +36,10 @@ pub enum WalletCommands {
     /// Get obj info
     #[structopt(name = "object")]
     Object {
+        /// Owner address
+        #[structopt(long, parse(try_from_str = decode_address_hex))]
+        owner: PublicKeyBytes,
+
         /// Object ID of the object to fetch
         #[structopt(long)]
         id: ObjectID,
@@ -47,6 +52,10 @@ pub enum WalletCommands {
     /// Publish Move modules
     #[structopt(name = "publish")]
     Publish {
+        /// Sender address
+        #[structopt(long, parse(try_from_str = decode_address_hex))]
+        sender: PublicKeyBytes,
+
         /// Path to directory containing a Move package
         #[structopt(long)]
         path: String,
@@ -63,6 +72,9 @@ pub enum WalletCommands {
     /// Call Move
     #[structopt(name = "call")]
     Call {
+        /// Sender address
+        #[structopt(long, parse(try_from_str = decode_address_hex))]
+        sender: PublicKeyBytes,
         /// Object ID of the package, which contains the module
         #[structopt(long)]
         package: ObjectID,
@@ -95,6 +107,10 @@ pub enum WalletCommands {
     /// Transfer funds
     #[structopt(name = "transfer")]
     Transfer {
+        /// Sender address
+        #[structopt(long, parse(try_from_str = decode_address_hex))]
+        from: PublicKeyBytes,
+
         /// Recipient address
         #[structopt(long, parse(try_from_str = decode_address_hex))]
         to: PublicKeyBytes,
@@ -135,21 +151,19 @@ impl WalletCommands {
     pub async fn execute(&mut self, context: &mut WalletContext) -> Result<(), anyhow::Error> {
         match self {
             WalletCommands::Publish {
+                sender,
                 path,
                 gas,
                 gas_budget,
             } => {
                 // Find owner of gas object
-                let owner = context.find_owner(gas)?;
-                let client_state = context.get_or_create_client_state(&owner)?;
+                let client_state = context.get_or_create_client_state(sender)?;
                 publish(client_state, path.clone(), *gas, *gas_budget).await;
             }
 
-            WalletCommands::Object { id, deep } => {
-                // Pick the first (or any) account for use in finding obj info
-                let account = context.find_owner(id)?;
+            WalletCommands::Object { id, deep, owner } => {
                 // Fetch the object ref
-                let client_state = context.get_or_create_client_state(&account)?;
+                let client_state = context.get_or_create_client_state(owner)?;
                 let object_read = client_state.get_object_info(*id).await?;
                 let object = object_read.object()?;
                 println!("{}", object);
@@ -158,6 +172,7 @@ impl WalletCommands {
                 }
             }
             WalletCommands::Call {
+                sender,
                 package,
                 module,
                 function,
@@ -167,9 +182,7 @@ impl WalletCommands {
                 gas,
                 gas_budget,
             } => {
-                // Find owner of gas object
-                let sender = context.find_owner(gas)?;
-                let client_state = context.get_or_create_client_state(&sender)?;
+                let client_state = context.get_or_create_client_state(sender)?;
 
                 let package_obj_info = client_state.get_object_info(*package).await?;
                 let package_obj_ref = package_obj_info.object().unwrap().to_object_reference();
@@ -203,10 +216,13 @@ impl WalletCommands {
                 show_object_effects(effects);
             }
 
-            WalletCommands::Transfer { to, object_id, gas } => {
-                let owner = context.find_owner(gas)?;
-
-                let client_state = context.get_or_create_client_state(&owner)?;
+            WalletCommands::Transfer {
+                to,
+                object_id,
+                gas,
+                from,
+            } => {
+                let client_state = context.get_or_create_client_state(from)?;
                 info!("Starting transfer");
                 let time_start = Instant::now();
                 let cert = client_state
@@ -219,14 +235,15 @@ impl WalletCommands {
             }
 
             WalletCommands::Addresses => {
-                let addr_strings: Vec<_> = context
-                    .config
-                    .accounts
+                let addr_strings = context
+                    .address_manager
+                    .get_managed_address_states()
                     .iter()
-                    .map(|account| encode_address_hex(&account.address))
-                    .collect();
+                    .map(|(a, _)| encode_address_hex(a))
+                    .collect::<Vec<_>>();
+
                 let addr_text = addr_strings.join("\n");
-                println!("Showing {} results.", context.config.accounts.len());
+                println!("Showing {} results.", addr_strings.len());
                 println!("{}", addr_text);
             }
 
@@ -250,6 +267,8 @@ impl WalletCommands {
                     key_pair: key,
                 });
                 context.config.save()?;
+                // Create an address to be managed
+                let _ = context.get_or_create_client_state(&address)?;
                 println!(
                     "Created new keypair for address : {}",
                     encode_address_hex(&address)
@@ -327,51 +346,23 @@ fn make_authority_clients(
 
 pub struct WalletContext {
     pub config: WalletConfig,
-    pub client_states: BTreeMap<SuiAddress, ClientState<AuthorityClient>>,
+    pub address_manager: ClientAddressManager<AuthorityClient>,
 }
 
 impl WalletContext {
-    pub fn new(config: WalletConfig) -> Self {
-        Self {
+    pub fn new(config: WalletConfig) -> Result<Self, anyhow::Error> {
+        let path = config.db_folder_path.clone();
+        Ok(Self {
             config,
-            client_states: Default::default(),
-        }
-    }
-
-    pub fn find_owner(&self, object_id: &ObjectID) -> Result<SuiAddress, SuiError> {
-        self.client_states
-            .iter()
-            .find_map(|(owner, client_state)| {
-                if client_state.get_owned_objects().contains(object_id) {
-                    Some(owner)
-                } else {
-                    None
-                }
-            })
-            .copied()
-            .ok_or(SuiError::ObjectNotFound {
-                object_id: *object_id,
-            })
-    }
-
-    pub fn get_or_create_client_state(
-        &mut self,
-        owner: &SuiAddress,
-    ) -> Result<&mut ClientState<AuthorityClient>, anyhow::Error> {
-        Ok(if !self.client_states.contains_key(owner) {
-            let new_client = self.create_client_state(owner)?;
-            self.client_states.entry(*owner).or_insert(new_client)
-        } else {
-            self.client_states.get_mut(owner).unwrap()
+            address_manager: ClientAddressManager::new(PathBuf::from_str(&path)?)?,
         })
     }
 
-    fn create_client_state(
-        &self,
+    fn get_or_create_client_state(
+        &mut self,
         owner: &SuiAddress,
-    ) -> Result<ClientState<AuthorityClient>, SuiError> {
-        let client_info = self.get_account_info(owner)?;
-
+    ) -> Result<&mut ClientState<AuthorityClient>, SuiError> {
+        let kp = Box::pin(self.get_account_cfg_info(owner)?.key_pair.copy());
         let voting_rights = self
             .config
             .authorities
@@ -385,17 +376,11 @@ impl WalletContext {
             self.config.send_timeout,
             self.config.recv_timeout,
         );
-        let path = PathBuf::from(format!("{}/{:?}", self.config.db_folder_path, owner));
-        ClientState::new(
-            path,
-            client_info.address,
-            Box::pin(client_info.key_pair.copy()),
-            committee,
-            authority_clients,
-        )
+        self.address_manager
+            .get_or_create_state_mut(*owner, kp, committee, authority_clients)
     }
 
-    pub fn get_account_info(&self, address: &SuiAddress) -> Result<&AccountInfo, SuiError> {
+    pub fn get_account_cfg_info(&self, address: &SuiAddress) -> Result<&AccountInfo, SuiError> {
         self.config
             .accounts
             .iter()

sui_core/src/client.rs

@@ -31,13 +31,59 @@ use std::{
 pub type StableSyncSigner = Pin<Box<dyn signature::Signer<ed25519_dalek::Signature> + Send + Sync>>;
 
 pub mod client_store;
-use self::client_store::ClientStore;
 
 #[cfg(test)]
 use sui_types::SUI_FRAMEWORK_ADDRESS;
 
 pub type AsyncResult<'a, T, E> = future::BoxFuture<'a, Result<T, E>>;
 
+pub struct ClientAddressManager<A> {
+    store: client_store::ClientAddressManagerStore,
+    address_states: BTreeMap<SuiAddress, ClientState<A>>,
+}
+impl<A> ClientAddressManager<A> {
+    /// Create a new manager which stores its managed addresses at `path`
+    pub fn new(path: PathBuf) -> Result<Self, SuiError> {
+        Ok(Self {
+            store: client_store::ClientAddressManagerStore::open(path),
+            address_states: BTreeMap::new(),
+        })
+    }
+
+    /// Get (if exists) or create a new managed address state
+    pub fn get_or_create_state_mut(
+        &mut self,
+        address: SuiAddress,
+        secret: StableSyncSigner,
+        committee: Committee,
+        authority_clients: BTreeMap<AuthorityName, A>,
+    ) -> Result<&mut ClientState<A>, SuiError> {
+        if let std::collections::btree_map::Entry::Vacant(e) = self.address_states.entry(address) {
+            // Load the records if available
+            let single_store = if self.store.is_managed_address(address)? {
+                // Unwrap is okay since we checked cond
+                self.store.get_managed_address(address)
+            } else {
+                self.store.manage_new_address(address)
+            }?;
+            e.insert(ClientState::new_for_manager(
+                address,
+                secret,
+                committee,
+                authority_clients,
+                single_store,
+            )?);
+        }
+
+        return Ok(self.address_states.get_mut(&address).unwrap());
+    }
+
+    /// Get all the states
+    pub fn get_managed_address_states(&self) -> &BTreeMap<PublicKeyBytes, ClientState<A>> {
+        &self.address_states
+    }
+}
+
 pub struct ClientState<AuthorityAPI> {
     /// Our Sui address.
     address: SuiAddress,
@@ -46,7 +92,7 @@ pub struct ClientState<AuthorityAPI> {
     /// Authority entry point.
     authorities: AuthorityAggregator<AuthorityAPI>,
     /// Persistent store for client
-    store: ClientStore,
+    store: client_store::ClientSingleAddressStore,
 }
 
 // Operations are considered successful when they successfully reach a quorum of authorities.
@@ -102,6 +148,7 @@ impl<A> ClientState<A> {
     /// right after constructor to fetch missing info form authorities
     /// TODO: client should manage multiple addresses instead of each addr having DBs
     /// https://github.com/MystenLabs/fastnft/issues/332
+    #[cfg(test)]
     pub fn new(
         path: PathBuf,
         address: SuiAddress,
@@ -113,7 +160,22 @@ impl<A> ClientState<A> {
             address,
             secret,
             authorities: AuthorityAggregator::new(committee, authority_clients),
-            store: ClientStore::new(path),
+            store: client_store::ClientSingleAddressStore::new(path),
+        })
+    }
+
+    pub fn new_for_manager(
+        address: SuiAddress,
+        secret: StableSyncSigner,
+        committee: Committee,
+        authority_clients: BTreeMap<AuthorityName, A>,
+        store: client_store::ClientSingleAddressStore,
+    ) -> Result<Self, SuiError> {
+        Ok(ClientState {
+            address,
+            secret,
+            authorities: AuthorityAggregator::new(committee, authority_clients),
+            store,
         })
     }
 
@@ -216,7 +278,7 @@ impl<A> ClientState<A> {
     }
 
     #[cfg(test)]
-    pub fn store(&self) -> &ClientStore {
+    pub fn store(&self) -> &client_store::ClientSingleAddressStore {
         &self.store
     }