Merge pull request #292 from logos-blockchain/fryorcraken/labels

add `wallet account label` feature

wallet/src/chain_storage.rs

@@ -11,17 +11,19 @@ use key_protocol::{
 use log::debug;
 use nssa::program::Program;
 
-use crate::config::{InitialAccountData, PersistentAccountData, WalletConfig};
+use crate::config::{InitialAccountData, Label, PersistentAccountData, WalletConfig};
 
 pub struct WalletChainStore {
     pub user_data: NSSAUserData,
     pub wallet_config: WalletConfig,
+    pub labels: HashMap<String, Label>,
 }
 
 impl WalletChainStore {
     pub fn new(
         config: WalletConfig,
         persistent_accounts: Vec<PersistentAccountData>,
+        labels: HashMap<String, Label>,
     ) -> Result<Self> {
         if persistent_accounts.is_empty() {
             anyhow::bail!("Roots not found; please run setup beforehand");
@@ -85,6 +87,7 @@ impl WalletChainStore {
                 private_tree,
             )?,
             wallet_config: config,
+            labels,
         })
     }
 
@@ -120,6 +123,7 @@ impl WalletChainStore {
                 private_tree,
             )?,
             wallet_config: config,
+            labels: HashMap::new(),
         })
     }
 
@@ -291,6 +295,6 @@ mod tests {
         let config = create_sample_wallet_config();
         let accs = create_sample_persistent_accounts();
 
-        let _ = WalletChainStore::new(config.clone(), accs).unwrap();
+        let _ = WalletChainStore::new(config.clone(), accs, HashMap::new()).unwrap();
     }
 }

wallet/src/cli/account.rs

@@ -9,6 +9,7 @@ use serde::Serialize;
 use crate::{
     TokenDefinition, TokenHolding, WalletCore,
     cli::{SubcommandReturnValue, WalletSubcommand},
+    config::Label,
     helperfunctions::{AccountPrivacyKind, HumanReadableAccount, parse_addr_with_privacy_prefix},
 };
 
@@ -39,6 +40,15 @@ pub enum AccountSubcommand {
         #[arg(short, long)]
         long: bool,
     },
+    /// Set a label for an account
+    Label {
+        /// Valid 32 byte base58 string with privacy prefix
+        #[arg(short, long)]
+        account_id: String,
+        /// The label to assign to the account
+        #[arg(short, long)]
+        label: String,
+    },
 }
 
 /// Represents generic register CLI subcommand
@@ -218,9 +228,13 @@ impl WalletSubcommand for AccountSubcommand {
                 keys,
                 account_id,
             } => {
-                let (account_id, addr_kind) = parse_addr_with_privacy_prefix(&account_id)?;
+                let (account_id_str, addr_kind) = parse_addr_with_privacy_prefix(&account_id)?;
 
-                let account_id = account_id.parse()?;
+                let account_id: nssa::AccountId = account_id_str.parse()?;
+
+                if let Some(label) = wallet_core.storage.labels.get(&account_id_str) {
+                    println!("Label: {label}");
+                }
 
                 let account = match addr_kind {
                     AccountPrivacyKind::Public => {
@@ -316,32 +330,35 @@ impl WalletSubcommand for AccountSubcommand {
             }
             AccountSubcommand::List { long } => {
                 let user_data = &wallet_core.storage.user_data;
+                let labels = &wallet_core.storage.labels;
+
+                let format_with_label = |prefix: &str, id: &nssa::AccountId| {
+                    let id_str = id.to_string();
+                    if let Some(label) = labels.get(&id_str) {
+                        format!("{prefix} [{label}]")
+                    } else {
+                        prefix.to_string()
+                    }
+                };
 
                 if !long {
                     let accounts = user_data
                         .default_pub_account_signing_keys
                         .keys()
-                        .map(|id| format!("Preconfigured Public/{id}"))
-                        .chain(
-                            user_data
-                                .default_user_private_accounts
-                                .keys()
-                                .map(|id| format!("Preconfigured Private/{id}")),
-                        )
-                        .chain(
-                            user_data
-                                .public_key_tree
-                                .account_id_map
-                                .iter()
-                                .map(|(id, chain_index)| format!("{chain_index} Public/{id}")),
-                        )
-                        .chain(
-                            user_data
-                                .private_key_tree
-                                .account_id_map
-                                .iter()
-                                .map(|(id, chain_index)| format!("{chain_index} Private/{id}")),
-                        )
+                        .map(|id| format_with_label(&format!("Preconfigured Public/{id}"), id))
+                        .chain(user_data.default_user_private_accounts.keys().map(|id| {
+                            format_with_label(&format!("Preconfigured Private/{id}"), id)
+                        }))
+                        .chain(user_data.public_key_tree.account_id_map.iter().map(
+                            |(id, chain_index)| {
+                                format_with_label(&format!("{chain_index} Public/{id}"), id)
+                            },
+                        ))
+                        .chain(user_data.private_key_tree.account_id_map.iter().map(
+                            |(id, chain_index)| {
+                                format_with_label(&format!("{chain_index} Private/{id}"), id)
+                            },
+                        ))
                         .format("\n");
 
                     println!("{accounts}");
@@ -351,7 +368,10 @@ impl WalletSubcommand for AccountSubcommand {
                 // Detailed listing with --long flag
                 // Preconfigured public accounts
                 for id in user_data.default_pub_account_signing_keys.keys() {
-                    println!("Preconfigured Public/{id}");
+                    println!(
+                        "{}",
+                        format_with_label(&format!("Preconfigured Public/{id}"), id)
+                    );
                     match wallet_core.get_account_public(*id).await {
                         Ok(account) if account != Account::default() => {
                             let (description, json_view) = format_account_details(&account);
@@ -365,7 +385,10 @@ impl WalletSubcommand for AccountSubcommand {
 
                 // Preconfigured private accounts
                 for id in user_data.default_user_private_accounts.keys() {
-                    println!("Preconfigured Private/{id}");
+                    println!(
+                        "{}",
+                        format_with_label(&format!("Preconfigured Private/{id}"), id)
+                    );
                     match wallet_core.get_account_private(id) {
                         Some(account) if account != Account::default() => {
                             let (description, json_view) = format_account_details(&account);
@@ -379,7 +402,10 @@ impl WalletSubcommand for AccountSubcommand {
 
                 // Public key tree accounts
                 for (id, chain_index) in user_data.public_key_tree.account_id_map.iter() {
-                    println!("{chain_index} Public/{id}");
+                    println!(
+                        "{}",
+                        format_with_label(&format!("{chain_index} Public/{id}"), id)
+                    );
                     match wallet_core.get_account_public(*id).await {
                         Ok(account) if account != Account::default() => {
                             let (description, json_view) = format_account_details(&account);
@@ -393,7 +419,10 @@ impl WalletSubcommand for AccountSubcommand {
 
                 // Private key tree accounts
                 for (id, chain_index) in user_data.private_key_tree.account_id_map.iter() {
-                    println!("{chain_index} Private/{id}");
+                    println!(
+                        "{}",
+                        format_with_label(&format!("{chain_index} Private/{id}"), id)
+                    );
                     match wallet_core.get_account_private(id) {
                         Some(account) if account != Account::default() => {
                             let (description, json_view) = format_account_details(&account);
@@ -405,6 +434,23 @@ impl WalletSubcommand for AccountSubcommand {
                     }
                 }
 
+                Ok(SubcommandReturnValue::Empty)
+            }
+            AccountSubcommand::Label { account_id, label } => {
+                let (account_id_str, _) = parse_addr_with_privacy_prefix(&account_id)?;
+
+                let old_label = wallet_core
+                    .storage
+                    .labels
+                    .insert(account_id_str.clone(), Label::new(label.clone()));
+
+                wallet_core.store_persistent_data().await?;
+
+                if let Some(old) = old_label {
+                    eprintln!("Warning: overriding existing label '{old}'");
+                }
+                println!("Label '{label}' set for account {account_id_str}");
+
                 Ok(SubcommandReturnValue::Empty)
             }
         }

wallet/src/config.rs

@@ -1,4 +1,5 @@
 use std::{
+    collections::HashMap,
     io::{BufReader, Write as _},
     path::Path,
 };
@@ -62,10 +63,30 @@ pub enum PersistentAccountData {
     Preconfigured(InitialAccountData),
 }
 
+/// A human-readable label for an account.
+#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
+pub struct Label(String);
+
+impl Label {
+    pub fn new(label: String) -> Self {
+        Self(label)
+    }
+}
+
+impl std::fmt::Display for Label {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        write!(f, "{}", self.0)
+    }
+}
+
 #[derive(Debug, Clone, Serialize, Deserialize)]
 pub struct PersistentStorage {
     pub accounts: Vec<PersistentAccountData>,
     pub last_synced_block: u64,
+    /// Account labels keyed by account ID string (e.g.,
+    /// "2rnKprXqWGWJTkDZKsQbFXa4ctKRbapsdoTKQFnaVGG8")
+    #[serde(default)]
+    pub labels: HashMap<String, Label>,
 }
 
 impl PersistentStorage {

wallet/src/helperfunctions.rs

@@ -1,4 +1,4 @@
-use std::{path::PathBuf, str::FromStr};
+use std::{collections::HashMap, path::PathBuf, str::FromStr};
 
 use anyhow::Result;
 use base64::{Engine, engine::general_purpose::STANDARD as BASE64};
@@ -11,7 +11,7 @@ use serde::Serialize;
 use crate::{
     HOME_DIR_ENV_VAR,
     config::{
-        InitialAccountData, InitialAccountDataPrivate, InitialAccountDataPublic,
+        InitialAccountData, InitialAccountDataPrivate, InitialAccountDataPublic, Label,
         PersistentAccountDataPrivate, PersistentAccountDataPublic, PersistentStorage,
     },
 };
@@ -57,6 +57,7 @@ pub fn fetch_persistent_storage_path() -> Result<PathBuf> {
 pub fn produce_data_for_storage(
     user_data: &NSSAUserData,
     last_synced_block: u64,
+    labels: HashMap<String, Label>,
 ) -> PersistentStorage {
     let mut vec_for_storage = vec![];
 
@@ -110,6 +111,7 @@ pub fn produce_data_for_storage(
     PersistentStorage {
         accounts: vec_for_storage,
         last_synced_block,
+        labels,
     }
 }
 

wallet/src/lib.rs

@@ -149,6 +149,7 @@ impl WalletCore {
         let PersistentStorage {
             accounts: persistent_accounts,
             last_synced_block,
+            labels,
         } = PersistentStorage::from_path(&storage_path)
             .with_context(|| format!("Failed to read persistent storage at {storage_path:#?}"))?;
 
@@ -156,7 +157,7 @@ impl WalletCore {
             config_path,
             storage_path,
             config_overrides,
-            |config| WalletChainStore::new(config, persistent_accounts),
+            |config| WalletChainStore::new(config, persistent_accounts, labels),
             last_synced_block,
         )
     }
@@ -225,7 +226,11 @@ impl WalletCore {
 
     /// Store persistent data at home
     pub async fn store_persistent_data(&self) -> Result<()> {
-        let data = produce_data_for_storage(&self.storage.user_data, self.last_synced_block);
+        let data = produce_data_for_storage(
+            &self.storage.user_data,
+            self.last_synced_block,
+            self.storage.labels.clone(),
+        );
         let storage = serde_json::to_vec_pretty(&data)?;
 
         let mut storage_file = tokio::fs::File::create(&self.storage_path).await?;