From 5fcb1cc226abfe5176922f8d30696039714c414c Mon Sep 17 00:00:00 2001
From: Philip Vu <pvu2002@outlook.com>
Date: Wed, 6 Nov 2024 16:20:55 -0800
Subject: [PATCH] Restrict prefixes to ed25519

---
 crates/aptos/src/common/init.rs  |  8 +++---
 crates/aptos/src/common/types.rs |  2 +-
 crates/aptos/src/common/utils.rs | 49 +++++++++++++++++++++-----------
 3 files changed, 37 insertions(+), 22 deletions(-)

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<String> {
+    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<String> into Option<Ed25519PrivateKey>
     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::<D::Error>::new(
-                stripped,
-            ))
-            .map(Some)
-            .map_err(D::Error::custom)
-        } else {
-            // Attempt normal deserialization
-            Ed25519PrivateKey::deserialize(serde::de::value::StrDeserializer::<D::Error>::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::<D::Error>::new(
+            &stripped,
+        ))
+        .map(Some)
+        .map_err(D::Error::custom)
     })
 }