[CLI] Extend vanity address support to multisig (aptos-labs#7694)

crates/aptos/src/common/types.rs

@@ -502,13 +502,6 @@ impl EncodingType {
     }
 }
 
-#[derive(Clone, Debug, Parser)]
-pub struct VanityPrefix {
-    /// Vanity prefix that resultant account address should start with, e.g. 0xaceface or d00d
-    #[clap(long)]
-    pub vanity_prefix: Option<String>,
-}
-
 #[derive(Clone, Debug, Parser)]
 pub struct RngArgs {
     /// The seed used for key generation, should be a 64 character hex string and only used for testing

crates/aptos/src/common/utils.rs

@@ -16,6 +16,7 @@ use aptos_logger::{debug, Level};
 use aptos_rest_client::{aptos_api_types::HashValue, Account, Client, State};
 use aptos_telemetry::service::telemetry_is_disabled;
 use aptos_types::{
+    account_address::create_multisig_account_address,
     chain_id::ChainId,
     transaction::{authenticator::AuthenticationKey, TransactionPayload},
 };
@@ -314,24 +315,27 @@ where
     Ok(map)
 }
 
-/// Generate a vanity account for Ed25519 single signer scheme.
+/// Generate a vanity account for Ed25519 single signer scheme, either standard or multisig.
 ///
 /// The default authentication key for an Ed25519 account is the same as the account address. Hence
-/// this function generates Ed25519 private keys until finding one that has an authentication key
-/// that begins with the given vanity prefix. Note that while a valid hex string must have an even
-/// number of characters, a vanity prefix can have an odd number of characters since
-/// account addresses are human-readable.
+/// for a standard account, this function generates Ed25519 private keys until finding one that has
+/// an authentication key (account address) that begins with the given vanity prefix.
+///
+/// For a multisig account, this function generates private keys until finding one that can create
+/// a multisig account with the given vanity prefix as its first transaction (sequence number 0).
+///
+/// Note that while a valid hex string must have an even number of characters, a vanity prefix can
+/// have an odd number of characters since account addresses are human-readable.
 ///
 /// `vanity_prefix_ref` is a reference to a hex string vanity prefix, optionally prefixed with "0x".
 /// For example "0xaceface" or "d00d".
 pub fn generate_vanity_account_ed25519(
     vanity_prefix_ref: &str,
+    multisig: bool,
 ) -> CliTypedResult<Ed25519PrivateKey> {
-    // Optionally strip leading 0x from input string.
     let vanity_prefix_ref = vanity_prefix_ref
         .strip_prefix("0x")
-        .unwrap_or(vanity_prefix_ref);
-    // Get string instance to check if vanity prefix is valid hex (may need to add a 0 at end).
+        .unwrap_or(vanity_prefix_ref); // Optionally strip leading 0x from input string.
     let mut to_check_if_is_hex = String::from(vanity_prefix_ref);
     // If an odd number of characters append a 0 for verifying that prefix contains valid hex.
     if to_check_if_is_hex.len() % 2 != 0 {
@@ -340,19 +344,22 @@ pub fn generate_vanity_account_ed25519(
     hex::decode(to_check_if_is_hex).  // Check that the vanity prefix can be decoded into hex.
         map_err(|error| CliError::CommandArgumentError(format!(
             "The vanity prefix could not be decoded to hex: {}", error)))?;
-    // Create a random key generator based on OS random number generator.
-    let mut key_generator = KeyGen::from_os_rng();
-    // Randomly generate a new Ed25519 private key.
-    let mut private_key = key_generator.generate_ed25519_private_key();
-    // While the account address does not start with the vanity prefix:
-    while !account_address_from_public_key(&Ed25519PublicKey::from(&private_key))
-        .short_str_lossless()
-        .starts_with(vanity_prefix_ref)
-    {
-        // Generate a new account.
-        private_key = key_generator.generate_ed25519_private_key();
+    let mut key_generator = KeyGen::from_os_rng(); // Get random key generator.
+    loop {
+        // Generate new keys until finding a match against the vanity prefix.
+        let private_key = key_generator.generate_ed25519_private_key();
+        let mut account_address =
+            account_address_from_public_key(&Ed25519PublicKey::from(&private_key));
+        if multisig {
+            account_address = create_multisig_account_address(account_address, 0)
+        };
+        if account_address
+            .short_str_lossless()
+            .starts_with(vanity_prefix_ref)
+        {
+            return Ok(private_key);
+        };
     }
-    Ok(private_key) // Return the resulting private key.
 }
 
 pub fn current_dir() -> CliTypedResult<PathBuf> {

crates/aptos/src/op/key.rs

@@ -5,7 +5,7 @@ use crate::{
     common::{
         types::{
             account_address_from_public_key, CliError, CliTypedResult, EncodingOptions,
-            EncodingType, KeyType, RngArgs, SaveFile, VanityPrefix,
+            EncodingType, KeyType, RngArgs, SaveFile,
         },
         utils::{
             append_file_extension, check_if_file_exists, generate_vanity_account_ed25519,
@@ -17,7 +17,9 @@ use crate::{
 use aptos_config::config::{Peer, PeerRole};
 use aptos_crypto::{bls12381, ed25519, x25519, PrivateKey, ValidCryptoMaterial};
 use aptos_genesis::config::HostAndPort;
-use aptos_types::account_address::{from_identity_public_key, AccountAddress};
+use aptos_types::account_address::{
+    create_multisig_account_address, from_identity_public_key, AccountAddress,
+};
 use async_trait::async_trait;
 use clap::{Parser, Subcommand};
 use std::{
@@ -170,13 +172,21 @@ pub struct GenerateKey {
     /// Key type to generate. Must be one of [x25519, ed25519, bls12381]
     #[clap(long, default_value_t = KeyType::Ed25519)]
     pub(crate) key_type: KeyType,
-
+    /// Vanity prefix that resultant account address should start with, e.g. 0xaceface or d00d. Each
+    /// additional character multiplies by a factor of 16 the computational difficulty associated
+    /// with generating an address, so try out shorter prefixes first and be prepared to wait for
+    /// longer ones
+    #[clap(long)]
+    pub vanity_prefix: Option<String>,
+    /// Use this flag when vanity prefix is for a multisig account. This mines a private key for
+    /// a single signer account that can, as its first transaction, create a multisig account with
+    /// the given vanity prefix
+    #[clap(long)]
+    pub vanity_multisig: bool,
     #[clap(flatten)]
     pub rng_args: RngArgs,
     #[clap(flatten)]
     pub(crate) save_params: SaveKey,
-    #[clap(flatten)]
-    pub vanity_prefix: VanityPrefix,
 }
 
 #[async_trait]
@@ -186,16 +196,19 @@ impl CliCommand<HashMap<&'static str, PathBuf>> for GenerateKey {
     }
 
     async fn execute(self) -> CliTypedResult<HashMap<&'static str, PathBuf>> {
-        self.save_params.check_key_file()?;
-        let mut keygen = self.rng_args.key_generator()?;
-        // Verify key type is Ed25519 if a vanity prefix is specified.
-        if self.vanity_prefix.vanity_prefix.is_some() && !matches!(self.key_type, KeyType::Ed25519)
-        {
+        if self.vanity_prefix.is_some() && !matches!(self.key_type, KeyType::Ed25519) {
             return Err(CliError::CommandArgumentError(format!(
                 "Vanity prefixes are only accepted for {} keys",
                 KeyType::Ed25519
             )));
         }
+        if self.vanity_multisig && self.vanity_prefix.is_none() {
+            return Err(CliError::CommandArgumentError(
+                "No vanity prefix provided".to_string(),
+            ));
+        }
+        self.save_params.check_key_file()?;
+        let mut keygen = self.rng_args.key_generator()?;
         match self.key_type {
             KeyType::X25519 => {
                 let private_key = keygen.generate_x25519_private_key().map_err(|err| {
@@ -208,25 +221,34 @@ impl CliCommand<HashMap<&'static str, PathBuf>> for GenerateKey {
             },
             KeyType::Ed25519 => {
                 // If no vanity prefix specified, generate a standard Ed25519 private key.
-                let private_key = if self.vanity_prefix.vanity_prefix.is_none() {
+                let private_key = if self.vanity_prefix.is_none() {
                     keygen.generate_ed25519_private_key()
                 } else {
                     // If a vanity prefix is specified, generate vanity Ed25519 account from it.
                     generate_vanity_account_ed25519(
-                        self.vanity_prefix.vanity_prefix.clone().unwrap().as_str(),
+                        self.vanity_prefix.clone().unwrap().as_str(),
+                        self.vanity_multisig,
                     )?
                 };
-                // Store CLI result map from key save operation, to append vanity address if needed.
+                // Store CLI result from key save operation, to append vanity address(es) if needed.
                 let mut result_map = self.save_params.save_key(&private_key, "ed25519").unwrap();
-                // If a vanity prefix was specified:
-                if self.vanity_prefix.vanity_prefix.is_some() {
-                    // Get account address.
+                if self.vanity_prefix.is_some() {
                     let account_address = account_address_from_public_key(
                         &ed25519::Ed25519PublicKey::from(&private_key),
-                    )
-                    .to_hex_literal();
-                    // Put account address in PathBuf so it can be displayed from typed CLI result.
-                    result_map.insert("Account Address:", PathBuf::from(account_address));
+                    );
+                    // Store account address in a PathBuf so it can be displayed in CLI result.
+                    result_map.insert(
+                        "Account Address:",
+                        PathBuf::from(account_address.to_hex_literal()),
+                    );
+                    if self.vanity_multisig {
+                        let multisig_account_address =
+                            create_multisig_account_address(account_address, 0);
+                        result_map.insert(
+                            "Multisig Account Address:",
+                            PathBuf::from(multisig_account_address.to_hex_literal()),
+                        );
+                    }
                 }
                 return Ok(result_map);
             },

crates/aptos/src/test/mod.rs

@@ -16,7 +16,7 @@ use crate::{
             EncodingOptions, FaucetOptions, GasOptions, KeyType, MoveManifestAccountWrapper,
             MovePackageDir, OptionalPoolAddressArgs, PoolAddressArgs, PrivateKeyInputOptions,
             PromptOptions, PublicKeyInputOptions, RestOptions, RngArgs, SaveFile,
-            TransactionOptions, TransactionSummary, VanityPrefix,
+            TransactionOptions, TransactionSummary,
         },
         utils::write_to_file,
     },
@@ -720,9 +720,8 @@ impl CliTestFramework {
                 },
                 encoding_options: Default::default(),
             },
-            vanity_prefix: VanityPrefix {
-                vanity_prefix: None,
-            },
+            vanity_prefix: None,
+            vanity_multisig: false,
         }
         .execute()
         .await