From ef9b58b617e33eba326a219a4dfb705b72db0dde Mon Sep 17 00:00:00 2001
From: Thomas Coratger <60488569+tcoratger@users.noreply.github.com>
Date: Thu, 10 Oct 2024 13:15:59 +0200
Subject: [PATCH] GenesisAccount : implement `deserialize_private_key` (#1447)

* GenesisAccount : implement deserialize_private_key

* fix no std
---
 crates/genesis/src/lib.rs | 97 +++++++++++++++++++++++++++++++++++++--
 1 file changed, 94 insertions(+), 3 deletions(-)

diff --git a/crates/genesis/src/lib.rs b/crates/genesis/src/lib.rs
index 85733936101..506f702c0cd 100644
--- a/crates/genesis/src/lib.rs
+++ b/crates/genesis/src/lib.rs
@@ -11,10 +11,11 @@
 
 extern crate alloc;
 
-use alloc::collections::BTreeMap;
+use alloc::{collections::BTreeMap, string::String};
 use alloy_primitives::{Address, Bytes, B256, U256};
 use alloy_serde::{storage::deserialize_storage_map, ttd::deserialize_json_ttd_opt, OtherFields};
-use serde::{Deserialize, Serialize};
+use core::str::FromStr;
+use serde::{de::Error as DeError, Deserialize, Deserializer, Serialize};
 
 /// The genesis block specification.
 #[derive(Clone, Debug, Default, PartialEq, Eq, Serialize, Deserialize)]
@@ -216,7 +217,12 @@ pub struct GenesisAccount {
     )]
     pub storage: Option<BTreeMap<B256, B256>>,
     /// The account's private key. Should only be used for testing.
-    #[serde(rename = "secretKey", default, skip_serializing_if = "Option::is_none")]
+    #[serde(
+        rename = "secretKey",
+        default,
+        skip_serializing_if = "Option::is_none",
+        deserialize_with = "deserialize_private_key"
+    )]
     pub private_key: Option<B256>,
 }
 
@@ -246,6 +252,28 @@ impl GenesisAccount {
     }
 }
 
+/// Custom deserialization function for the private key.
+///
+/// This function allows the private key to be deserialized from a string or a `null` value.
+///
+/// We need a custom function here especially to handle the case where the private key is `0x` and
+/// should be deserialized as `None`.
+fn deserialize_private_key<'de, D>(deserializer: D) -> Result<Option<B256>, D::Error>
+where
+    D: Deserializer<'de>,
+{
+    let opt_str: Option<String> = Option::deserialize(deserializer)?;
+
+    if let Some(ref s) = opt_str {
+        if s == "0x" {
+            return Ok(None);
+        }
+        B256::from_str(s).map(Some).map_err(D::Error::custom)
+    } else {
+        Ok(None)
+    }
+}
+
 /// Defines core blockchain settings per block.
 ///
 /// Tailors unique settings for each network based on its genesis block.
@@ -588,6 +616,7 @@ mod tests {
     use alloc::vec;
     use alloy_primitives::hex;
     use core::str::FromStr;
+    use serde_json::json;
 
     #[test]
     fn genesis_defaults_config() {
@@ -1530,4 +1559,66 @@ mod tests {
         let actual_object_value = genesis.config.extra_fields.get("object_field").unwrap();
         assert_eq!(actual_object_value, &serde_json::json!({"sub_field": "sub_value"}));
     }
+
+    #[test]
+    fn deserialize_private_key_as_none_when_0x() {
+        // Test case where "secretKey" is "0x", expecting None
+        let json_data = json!({
+            "balance": "0x0",
+            "secretKey": "0x"
+        });
+
+        let account: GenesisAccount = serde_json::from_value(json_data).unwrap();
+        assert_eq!(account.private_key, None);
+    }
+
+    #[test]
+    fn deserialize_private_key_with_valid_hex() {
+        // Test case where "secretKey" is a valid hex string
+        let json_data = json!({
+            "balance": "0x0",
+            "secretKey": "0x123456789abcdef123456789abcdef123456789abcdef123456789abcdef1234"
+        });
+
+        let account: GenesisAccount = serde_json::from_value(json_data).unwrap();
+        let expected_key =
+            B256::from_str("123456789abcdef123456789abcdef123456789abcdef123456789abcdef1234")
+                .unwrap();
+        assert_eq!(account.private_key, Some(expected_key));
+    }
+
+    #[test]
+    fn deserialize_private_key_as_none_when_null() {
+        // Test case where "secretKey" is null, expecting None
+        let json_data = json!({
+            "balance": "0x0",
+            "secretKey": null
+        });
+
+        let account: GenesisAccount = serde_json::from_value(json_data).unwrap();
+        assert_eq!(account.private_key, None);
+    }
+
+    #[test]
+    fn deserialize_private_key_with_invalid_hex_fails() {
+        // Test case where "secretKey" is an invalid hex string, expecting an error
+        let json_data = json!({
+            "balance": "0x0",
+            "secretKey": "0xINVALIDHEX"
+        });
+
+        let result: Result<GenesisAccount, _> = serde_json::from_value(json_data);
+        assert!(result.is_err()); // The deserialization should fail due to invalid hex
+    }
+
+    #[test]
+    fn deserialize_private_key_with_empty_string_fails() {
+        // Test case where "secretKey" is an empty string, expecting an error
+        let json_data = json!({
+            "secretKey": ""
+        });
+
+        let result: Result<GenesisAccount, _> = serde_json::from_value(json_data);
+        assert!(result.is_err()); // The deserialization should fail due to an empty string
+    }
 }