From 22d2c9c08da7aad4b41ab1adab80f88b33304d0b Mon Sep 17 00:00:00 2001
From: Philip Vu <pvu2002@outlook.com>
Date: Mon, 28 Oct 2024 13:34:34 -0700
Subject: [PATCH] Add ProfileConfig deserialization

---
 crates/aptos/src/common/types.rs |  9 +++++----
 crates/aptos/src/common/utils.rs | 34 +++++++++++++++++++++++++++++++-
 2 files changed, 38 insertions(+), 5 deletions(-)

diff --git a/crates/aptos/src/common/types.rs b/crates/aptos/src/common/types.rs
index 6c96e01e3e4b3f..8f940f2c6dfbbc 100644
--- a/crates/aptos/src/common/types.rs
+++ b/crates/aptos/src/common/types.rs
@@ -7,10 +7,10 @@ use crate::{
         init::Network,
         local_simulation,
         utils::{
-            check_if_file_exists, create_dir_if_not_exist, dir_default_to_current,
-            get_account_with_state, get_auth_key, get_sequence_number, parse_json_file,
-            prompt_yes_with_override, read_from_file, start_logger, to_common_result,
-            to_common_success_result, write_to_file, write_to_file_with_opts,
+            check_if_file_exists, create_dir_if_not_exist, deserialize_private_key_with_prefix,
+            dir_default_to_current, get_account_with_state, get_auth_key, get_sequence_number,
+            parse_json_file, prompt_yes_with_override, read_from_file, start_logger,
+            to_common_result, to_common_success_result, write_to_file, write_to_file_with_opts,
             write_to_user_only_file,
         },
     },
@@ -246,6 +246,7 @@ pub struct ProfileConfig {
     pub network: Option<Network>,
     /// Private key for commands.
     #[serde(skip_serializing_if = "Option::is_none")]
+    #[serde(deserialize_with = "deserialize_private_key_with_prefix")]
     pub private_key: Option<Ed25519PrivateKey>,
     /// Public key for commands
     #[serde(skip_serializing_if = "Option::is_none")]
diff --git a/crates/aptos/src/common/utils.rs b/crates/aptos/src/common/utils.rs
index 94c12570b25bea..8c8ec58858d88e 100644
--- a/crates/aptos/src/common/utils.rs
+++ b/crates/aptos/src/common/utils.rs
@@ -27,7 +27,7 @@ use aptos_types::{
 use itertools::Itertools;
 use move_core_types::{account_address::AccountAddress, language_storage::CORE_CODE_ADDRESS};
 use reqwest::Url;
-use serde::{Deserialize, Serialize};
+use serde::{Deserialize, Deserializer, Serialize};
 #[cfg(unix)]
 use std::os::unix::fs::OpenOptionsExt;
 use std::{
@@ -609,3 +609,35 @@ pub fn strip_private_key_prefix(key: &String) -> String {
         .unwrap_or(key) // If no prefix is found, return the original key
         .to_string()
 }
+
+/// Deserializes an Ed25519 private key with a prefix AIP-80 prefix if present.
+///
+/// [Read about AIP-80](https://github.com/aptos-foundation/AIPs/blob/main/aips/aip-80.md)
+pub fn deserialize_private_key_with_prefix<'de, D>(
+    deserializer: D,
+) -> Result<Option<Ed25519PrivateKey>, D::Error>
+where
+    D: Deserializer<'de>,
+{
+    use serde::de::Error;
+
+    // Deserialize the field as an Option<String>
+    let opt: Option<String> = Option::deserialize(deserializer)?;
+
+    // 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)
+        }
+    })
+}