diff --git a/crates/aptos/src/common/init.rs b/crates/aptos/src/common/init.rs index cb3e526d2713ff..9a416e274b5b91 100644 --- a/crates/aptos/src/common/init.rs +++ b/crates/aptos/src/common/init.rs @@ -221,10 +221,10 @@ impl CliCommand<()> for InitTool { .generate_ed25519_private_key() } } else { - Ed25519PrivateKey::from_encoded_string(&strip_private_key_prefix( - &input.to_string(), - )) - .map_err(|err| CliError::UnableToParse("Ed25519PrivateKey", err.to_string()))? + let stripped = strip_private_key_prefix(&input.to_string())?; + Ed25519PrivateKey::from_encoded_string(&stripped).map_err(|err| { + CliError::UnableToParse("Ed25519PrivateKey", err.to_string()) + })? } }; diff --git a/crates/aptos/src/common/types.rs b/crates/aptos/src/common/types.rs index 0028cb78f147ba..57805330e01b0c 100644 --- a/crates/aptos/src/common/types.rs +++ b/crates/aptos/src/common/types.rs @@ -689,7 +689,7 @@ pub trait ParsePrivateKey { encoding.load_key("--private-key-file", file.as_path())?, )) } else if let Some(ref key) = private_key { - let key = strip_private_key_prefix(key).as_bytes().to_vec(); + let key = strip_private_key_prefix(key)?.as_bytes().to_vec(); Ok(Some(encoding.decode_key("--private-key", key)?)) } else { Ok(None) diff --git a/crates/aptos/src/common/utils.rs b/crates/aptos/src/common/utils.rs index 4819b819fd35c8..b822d0836a9b01 100644 --- a/crates/aptos/src/common/utils.rs +++ b/crates/aptos/src/common/utils.rs @@ -601,10 +601,29 @@ pub fn explorer_transaction_link( /// Strips the private key prefix for a given key string if it is AIP-80 compliant. /// /// [Read about AIP-80](https://github.com/aptos-foundation/AIPs/blob/main/aips/aip-80.md) -pub fn strip_private_key_prefix(key: &String) -> String { - key.strip_prefix("ed25519-priv-") - .unwrap_or(key) // If the prefix is not found, return the original key - .to_string() +pub fn strip_private_key_prefix(key: &String) -> CliTypedResult { + let disabled_prefixes = ["secp256k1-priv-"]; + let enabled_prefixes = ["ed25519-priv-"]; + + // Check for disabled prefixes first + for prefix in disabled_prefixes { + if key.starts_with(prefix) { + return Err(CliError::UnexpectedError(format!( + "Private key not supported. Cannot parse private key with '{}' prefix.", + prefix + ))); + } + } + + // Try to strip enabled prefixes + for prefix in enabled_prefixes { + if key.starts_with(prefix) { + return Ok(key.strip_prefix(prefix).unwrap().to_string()); + } + } + + // If no prefix is found, return the original key + Ok(key.to_string()) } /// Deserializes an Ed25519 private key with a prefix AIP-80 prefix if present. @@ -623,18 +642,14 @@ where // Transform Option into Option opt.map_or(Ok(None), |s| { - if let Some(stripped) = s.strip_prefix("ed25519-priv-") { - // Deserialize using Ed25519PrivateKey's DeserializeKey implementation - Ed25519PrivateKey::deserialize(serde::de::value::StrDeserializer::::new( - stripped, - )) - .map(Some) - .map_err(D::Error::custom) - } else { - // Attempt normal deserialization - Ed25519PrivateKey::deserialize(serde::de::value::StrDeserializer::::new(&s)) - .map(Some) - .map_err(D::Error::custom) - } + // Use strip_private_key_prefix to handle the AIP-80 prefix + let stripped = strip_private_key_prefix(&s).map_err(D::Error::custom)?; + + // Attempt deserialization with the stripped key + Ed25519PrivateKey::deserialize(serde::de::value::StrDeserializer::::new( + &stripped, + )) + .map(Some) + .map_err(D::Error::custom) }) }