From 81da73475c3aacf351490b91a7476f4bb53611ac Mon Sep 17 00:00:00 2001
From: lwshang <lwshang@users.noreply.github.com>
Date: Thu, 19 Mar 2020 17:05:49 -0400
Subject: [PATCH] Prompt error if vanity pattern is not valid base58 (#5308)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

* Prompt error if vanity pattern is not valid base58

* Update bin/utils/subkey/src/vanity.rs

* Add tests for validating pattern of subkey vanity

Co-authored-by: Bastian Köcher <bkchr@users.noreply.github.com>
---
 substrate/bin/utils/subkey/src/vanity.rs | 27 ++++++++++++++++++++++++
 1 file changed, 27 insertions(+)

diff --git a/substrate/bin/utils/subkey/src/vanity.rs b/substrate/bin/utils/subkey/src/vanity.rs
index 8a167cd062..f921470946 100644
--- a/substrate/bin/utils/subkey/src/vanity.rs
+++ b/substrate/bin/utils/subkey/src/vanity.rs
@@ -62,6 +62,12 @@ fn calculate_score(_desired: &str, key: &str) -> usize {
 	0
 }
 
+/// Validate whether the char is allowed to be used in base58.
+/// num 0, lower l, upper I and O are not allowed.
+fn validate_base58(c :char) -> bool {
+	c.is_alphanumeric() && !"0lIO".contains(c)
+}
+
 pub(super) fn generate_key<C: Crypto>(desired: &str) -> Result<KeyPair<C>, &'static str> where
 		PublicOf<C>: PublicT,
 {
@@ -69,6 +75,11 @@ pub(super) fn generate_key<C: Crypto>(desired: &str) -> Result<KeyPair<C>, &'sta
 		return Err("Pattern must not be empty");
 	}
 
+	if !desired.chars().all(validate_base58) {
+		return Err("Pattern can only contains valid characters in base58 \
+			(all alphanumeric except for 0, l, I and O)");
+	}
+
 	eprintln!("Generating key containing pattern '{}'", desired);
 
 	let top = 45 + (desired.len() * 48);
@@ -162,6 +173,22 @@ mod tests {
 		);
 	}
 
+	#[test]
+	fn test_invalid_pattern() {
+		assert!(generate_key::<Ed25519>("").is_err());
+		assert!(generate_key::<Ed25519>("0").is_err());
+		assert!(generate_key::<Ed25519>("l").is_err());
+		assert!(generate_key::<Ed25519>("I").is_err());
+		assert!(generate_key::<Ed25519>("O").is_err());
+		assert!(generate_key::<Ed25519>("!").is_err());
+	}
+
+	#[test]
+	fn test_valid_pattern() {
+		assert!(generate_key::<Ed25519>("o").is_ok());
+		assert!(generate_key::<Ed25519>("L").is_ok());
+	}
+
 	#[cfg(feature = "bench")]
 	#[bench]
 	fn bench_paranoiac(b: &mut Bencher) {