From 1284a491c21a09dac37a004473911c25c8aa3673 Mon Sep 17 00:00:00 2001
From: Aten <jincxmain@gmail.com>
Date: Thu, 30 Jul 2020 16:41:32 +0800
Subject: [PATCH] support custom ss58addressformat in
 from_ss58check_with_version (#5526)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

* support custom ss58addressformat in from_ss58check_with_version

* fix str parse

1. if can parse with u8, use u8 into.
2. if u8 can't parse, convert to str then parse

* add a test

* typo

* add error description in test

* fix the `TryFrom<u8>` for `Ss58AddressFormat`

change check logic in TryFrom<u8> to replace modified code in `from_ss58check_with_version`

* use Ss58AddressFormat::default() replace DEFAULT_VERSION

* Apply suggestions from code review

* Update primitives/core/src/crypto.rs

* Update primitives/core/src/crypto.rs

* Update primitives/core/src/crypto.rs

* Update primitives/core/src/crypto.rs

Co-authored-by: Bastian Köcher <bkchr@users.noreply.github.com>
---
 substrate/primitives/core/src/crypto.rs | 13 +++++++++++--
 substrate/primitives/core/src/ecdsa.rs  | 18 +++++++++++++++++-
 2 files changed, 28 insertions(+), 3 deletions(-)

diff --git a/substrate/primitives/core/src/crypto.rs b/substrate/primitives/core/src/crypto.rs
index c34115ac8c..0cdbebde9f 100644
--- a/substrate/primitives/core/src/crypto.rs
+++ b/substrate/primitives/core/src/crypto.rs
@@ -366,7 +366,16 @@ macro_rules! ss58_address_format {
 			fn try_from(x: u8) -> Result<Ss58AddressFormat, ()> {
 				match x {
 					$($number => Ok(Ss58AddressFormat::$identifier)),*,
-					_ => Err(()),
+					_ => {
+						#[cfg(feature = "std")]
+						match Ss58AddressFormat::default() {
+							Ss58AddressFormat::Custom(n) if n == x => Ok(Ss58AddressFormat::Custom(x)),
+							_ => Err(()),
+						}
+						
+						#[cfg(not(feature = "std"))]
+						Err(())
+					},
 				}
 			}
 		}
@@ -377,7 +386,7 @@ macro_rules! ss58_address_format {
 			fn try_from(x: &'a str) -> Result<Ss58AddressFormat, ()> {
 				match x {
 					$($name => Ok(Ss58AddressFormat::$identifier)),*,
-					a => a.parse::<u8>().map(Ss58AddressFormat::Custom).map_err(|_| ()),
+					a => a.parse::<u8>().map_err(|_| ()).and_then(TryFrom::try_from),
 				}
 			}
 		}
diff --git a/substrate/primitives/core/src/ecdsa.rs b/substrate/primitives/core/src/ecdsa.rs
index 29fa9a9c5c..da6b7614c7 100644
--- a/substrate/primitives/core/src/ecdsa.rs
+++ b/substrate/primitives/core/src/ecdsa.rs
@@ -552,7 +552,7 @@ impl CryptoType for Pair {
 mod test {
 	use super::*;
 	use hex_literal::hex;
-	use crate::crypto::DEV_PHRASE;
+	use crate::crypto::{DEV_PHRASE, set_default_ss58_version};
 	use serde_json;
 
 	#[test]
@@ -676,6 +676,22 @@ mod test {
 		assert_eq!(cmp, public);
 	}
 
+	#[test]
+	fn ss58check_custom_format_works() {
+		use crate::crypto::Ss58AddressFormat;
+		// temp save default format version
+		let default_format = Ss58AddressFormat::default();
+		// set current ss58 version is custom "200" `Ss58AddressFormat::Custom(200)`
+		set_default_ss58_version(Ss58AddressFormat::Custom(200));
+		// custom addr encoded by version 200
+		let addr = "2X64kMNEWAW5KLZMSKcGKEc96MyuaRsRUku7vomuYxKgqjVCRj";
+		Public::from_ss58check(&addr).unwrap();
+		set_default_ss58_version(default_format);
+		// set  current ss58 version to default version
+		let addr = "KWAfgC2aRG5UVD6CpbPQXCx4YZZUhvWqqAJE6qcYc9Rtr6g5C";
+		Public::from_ss58check(&addr).unwrap();
+	}
+
 	#[test]
 	fn signature_serialization_works() {
 		let pair = Pair::from_seed(b"12345678901234567890123456789012");