Generic keystore (paritytech#3008)

* Add KeyTypeId.

* Implement clone for sr25519::Pair.

* Extend Pair with to_raw_vec.

* Implement TypedKey for Signature and Pair.

* Add trait Public.

* Make keystore generic.

* Fixup clone.

* Fix tests.

* Update service.

* Fix imports.

* Fix build.

* Fix babe build.

* Fix subkey build.

* Make authority setup generic.

* Update node-template.

* Fix build.

* Remove unsafe code.

* Fix tests.

core/consensus/babe/src/lib.rs

@@ -45,7 +45,7 @@ use std::{sync::Arc, u64, fmt::{Debug, Display}, time::{Instant, Duration}};
 use runtime_support::serde::{Serialize, Deserialize};
 use parity_codec::{Decode, Encode};
 use parking_lot::Mutex;
-use primitives::{crypto::Pair, sr25519};
+use primitives::{Pair, Public, sr25519};
 use merlin::Transcript;
 use inherents::{InherentDataProviders, InherentData};
 use substrate_telemetry::{

core/keyring/src/ed25519.rs

@@ -18,7 +18,7 @@
 
 use std::{collections::HashMap, ops::Deref};
 use lazy_static::lazy_static;
-use substrate_primitives::{ed25519::{Pair, Public, Signature}, Pair as PairT, H256};
+use substrate_primitives::{ed25519::{Pair, Public, Signature}, Pair as PairT, Public as PublicT, H256};
 pub use substrate_primitives::ed25519;
 
 /// Set of test accounts.

core/keyring/src/sr25519.rs

@@ -19,7 +19,7 @@
 use std::collections::HashMap;
 use std::ops::Deref;
 use lazy_static::lazy_static;
-use substrate_primitives::{sr25519::{Pair, Public, Signature}, Pair as PairT, H256};
+use substrate_primitives::{sr25519::{Pair, Public, Signature}, Pair as PairT, Public as PublicT, H256};
 pub use substrate_primitives::sr25519;
 
 /// Set of test accounts.
@@ -68,7 +68,7 @@ impl Keyring {
 	}
 
 	pub fn to_raw_public_vec(self) -> Vec<u8> {
-		Public::from(self).into_raw_vec()
+		Public::from(self).to_raw_vec()
 	}
 
 	pub fn sign(self, msg: &[u8]) -> Signature {

core/keystore/src/lib.rs

@@ -23,7 +23,7 @@ use std::path::PathBuf;
 use std::fs::{self, File};
 use std::io::{self, Write};
 
-use substrate_primitives::{ed25519::{Pair, Public}, Pair as PairT};
+use substrate_primitives::crypto::{KeyTypeId, Pair, Public, TypedKey};
 
 /// Keystore error.
 #[derive(Debug, derive_more::Display, derive_more::From)]
@@ -59,7 +59,7 @@ impl std::error::Error for Error {
 /// Key store.
 pub struct Store {
 	path: PathBuf,
-	additional: HashMap<Public, Pair>,
+	additional: HashMap<(KeyTypeId, Vec<u8>), Vec<u8>>,
 }
 
 impl Store {
@@ -69,33 +69,49 @@ impl Store {
 		Ok(Store { path, additional: HashMap::new() })
 	}
 
+	fn get_pair<TPair: Pair + TypedKey>(&self, public: &TPair::Public) -> Result<Option<TPair>> {
+		let key = (TPair::KEY_TYPE, public.to_raw_vec());
+		if let Some(bytes) = self.additional.get(&key) {
+			let pair = TPair::from_seed_slice(bytes)
+				.map_err(|_| Error::InvalidSeed)?;
+			return Ok(Some(pair));
+		}
+		Ok(None)
+	}
+
+	fn insert_pair<TPair: Pair + TypedKey>(&mut self, pair: &TPair) {
+		let key = (TPair::KEY_TYPE, pair.public().to_raw_vec());
+		self.additional.insert(key, pair.to_raw_vec());
+	}
+
 	/// Generate a new key, placing it into the store.
-	pub fn generate(&self, password: &str) -> Result<Pair> {
-		let (pair, phrase, _) = Pair::generate_with_phrase(Some(password));
-		let mut file = File::create(self.key_file_path(&pair.public()))?;
+	pub fn generate<TPair: Pair + TypedKey>(&self, password: &str) -> Result<TPair> {
+		let (pair, phrase, _) = TPair::generate_with_phrase(Some(password));
+		let mut file = File::create(self.key_file_path::<TPair>(&pair.public()))?;
 		::serde_json::to_writer(&file, &phrase)?;
 		file.flush()?;
 		Ok(pair)
 	}
 
 	/// Create a new key from seed. Do not place it into the store.
-	pub fn generate_from_seed(&mut self, seed: &str) -> Result<Pair> {
-		let pair = Pair::from_string(seed, None)
+	pub fn generate_from_seed<TPair: Pair + TypedKey>(&mut self, seed: &str) -> Result<TPair> {
+		let pair = TPair::from_string(seed, None)
 			.ok().ok_or(Error::InvalidSeed)?;
-		self.additional.insert(pair.public(), pair.clone());
+		self.insert_pair(&pair);
 		Ok(pair)
 	}
 
 	/// Load a key file with given public key.
-	pub fn load(&self, public: &Public, password: &str) -> Result<Pair> {
-		if let Some(pair) = self.additional.get(public) {
-			return Ok(pair.clone());
+	pub fn load<TPair: Pair + TypedKey>(&self, public: &TPair::Public, password: &str) -> Result<TPair> {
+		if let Some(pair) = self.get_pair(public)? {
+			return Ok(pair)
 		}
-		let path = self.key_file_path(public);
+
+		let path = self.key_file_path::<TPair>(public);
 		let file = File::open(path)?;
 
 		let phrase: String = ::serde_json::from_reader(&file)?;
-		let (pair, _) = Pair::from_phrase(&phrase, Some(password))
+		let (pair, _) = TPair::from_phrase(&phrase, Some(password))
 			.ok().ok_or(Error::InvalidPhrase)?;
 		if &pair.public() != public {
 			return Err(Error::InvalidPassword);
@@ -104,22 +120,28 @@ impl Store {
 	}
 
 	/// Get public keys of all stored keys.
-	pub fn contents(&self) -> Result<Vec<Public>> {
-		let mut public_keys: Vec<Public> = self.additional.keys().cloned().collect();
+	pub fn contents<TPublic: Public + TypedKey>(&self) -> Result<Vec<TPublic>> {
+		let mut public_keys: Vec<TPublic> = self.additional.keys()
+			.filter_map(|(ty, public)| {
+				if *ty != TPublic::KEY_TYPE {
+					return None
+				}
+				Some(TPublic::from_slice(public))
+			})
+			.collect();
+
+		let key_type: [u8; 4] = TPublic::KEY_TYPE.to_le_bytes();
 		for entry in fs::read_dir(&self.path)? {
 			let entry = entry?;
 			let path = entry.path();
 
 			// skip directories and non-unicode file names (hex is unicode)
 			if let Some(name) = path.file_name().and_then(|n| n.to_str()) {
-				if name.len() != 64 { continue }
-
 				match hex::decode(name) {
-					Ok(ref hex) if hex.len() == 32 => {
-						let mut buf = [0; 32];
-						buf.copy_from_slice(&hex[..]);
-
-						public_keys.push(Public(buf));
+					Ok(ref hex) => {
+						if hex[0..4] != key_type { continue	}
+						let public = TPublic::from_slice(&hex[4..]);
+						public_keys.push(public);
 					}
 					_ => continue,
 				}
@@ -129,9 +151,12 @@ impl Store {
 		Ok(public_keys)
 	}
 
-	fn key_file_path(&self, public: &Public) -> PathBuf {
+	fn key_file_path<TPair: Pair + TypedKey>(&self, public: &TPair::Public) -> PathBuf {
 		let mut buf = self.path.clone();
-		buf.push(hex::encode(public.as_slice()));
+		let bytes: [u8; 4] = TPair::KEY_TYPE.to_le_bytes();
+		let key_type = hex::encode(bytes);
+		let key = hex::encode(public.as_slice());
+		buf.push(key_type + key.as_str());
 		buf
 	}
 }
@@ -140,31 +165,34 @@ impl Store {
 mod tests {
 	use super::*;
 	use tempdir::TempDir;
+	use substrate_primitives::ed25519;
 	use substrate_primitives::crypto::Ss58Codec;
 
 	#[test]
 	fn basic_store() {
 		let temp_dir = TempDir::new("keystore").unwrap();
 		let store = Store::open(temp_dir.path().to_owned()).unwrap();
 
-		assert!(store.contents().unwrap().is_empty());
+		assert!(store.contents::<ed25519::Public>().unwrap().is_empty());
 
-		let key = store.generate("thepassword").unwrap();
-		let key2 = store.load(&key.public(), "thepassword").unwrap();
+		let key: ed25519::Pair = store.generate("thepassword").unwrap();
+		let key2: ed25519::Pair = store.load(&key.public(), "thepassword").unwrap();
 
-		assert!(store.load(&key.public(), "notthepassword").is_err());
+		assert!(store.load::<ed25519::Pair>(&key.public(), "notthepassword").is_err());
 
 		assert_eq!(key.public(), key2.public());
 
-		assert_eq!(store.contents().unwrap()[0], key.public());
+		assert_eq!(store.contents::<ed25519::Public>().unwrap()[0], key.public());
 	}
 
 	#[test]
 	fn test_generate_from_seed() {
 		let temp_dir = TempDir::new("keystore").unwrap();
 		let mut store = Store::open(temp_dir.path().to_owned()).unwrap();
 
-		let pair = store.generate_from_seed("0x3d97c819d68f9bafa7d6e79cb991eebcd77d966c5334c0b94d9e1fa7ad0869dc").unwrap();
+		let pair: ed25519::Pair = store
+			.generate_from_seed("0x3d97c819d68f9bafa7d6e79cb991eebcd77d966c5334c0b94d9e1fa7ad0869dc")
+			.unwrap();
 		assert_eq!("5DKUrgFqCPV8iAXx9sjy1nyBygQCeiUYRFWurZGhnrn3HJCA", pair.public().to_ss58check());
 	}
 }

core/primitives/src/crypto.rs

@@ -26,6 +26,8 @@ use parity_codec::{Encode, Decode};
 use regex::Regex;
 #[cfg(feature = "std")]
 use base58::{FromBase58, ToBase58};
+#[cfg(feature = "std")]
+use std::hash::Hash;
 
 /// The root phrase for our publicly known keys.
 pub const DEV_PHRASE: &str = "bottom drive obey lake curtain smoke basket hold race lonely fit walk";
@@ -286,13 +288,30 @@ impl<T: AsMut<[u8]> + AsRef<[u8]> + Default + Derive> Ss58Codec for T {
 	}
 }
 
+/// Trait suitable for typical cryptographic PKI key public type.
+pub trait Public: PartialEq + Eq {
+	/// A new instance from the given slice that should be 32 bytes long.
+	///
+	/// NOTE: No checking goes on to ensure this is a real public key. Only use it if
+	/// you are certain that the array actually is a pubkey. GIGO!
+	fn from_slice(data: &[u8]) -> Self;
+
+	/// Return a `Vec<u8>` filled with raw data.
+	#[cfg(feature = "std")]
+	fn to_raw_vec(&self) -> Vec<u8>;
+
+	/// Return a slice filled with raw data.
+	fn as_slice(&self) -> &[u8];
+}
+
 /// Trait suitable for typical cryptographic PKI key pair type.
 ///
 /// For now it just specifies how to create a key from a phrase and derivation path.
 #[cfg(feature = "std")]
-pub trait Pair: Sized + 'static {
+pub trait Pair: Sized + 'static
+{
 	/// TThe type which is used to encode a public key.
-	type Public;
+	type Public: Public + Hash;
 
 	/// The type used to (minimally) encode the data required to securely create
 	/// a new key pair.
@@ -412,6 +431,31 @@ pub trait Pair: Sized + 'static {
 			path,
 		)
 	}
+
+	/// Return a vec filled with raw data.
+	fn to_raw_vec(&self) -> Vec<u8>;
+}
+
+/// An identifier for a type of cryptographic key.
+///
+/// 0-1024 are reserved.
+pub type KeyTypeId = u32;
+
+/// Constant key types.
+pub mod key_types {
+	use super::KeyTypeId;
+
+	/// ED25519 public key.
+	pub const ED25519: KeyTypeId = 10;
+
+	/// SR25519 public key.
+	pub const SR25519: KeyTypeId = 20;
+}
+
+/// A trait for something that has a key type ID.
+pub trait TypedKey {
+	/// The type ID of this key.
+	const KEY_TYPE: KeyTypeId;
 }
 
 #[cfg(test)]
@@ -429,8 +473,21 @@ mod tests {
 		Seed(Vec<u8>),
 	}
 
+	#[derive(PartialEq, Eq, Hash)]
+	struct TestPublic;
+	impl Public for TestPublic {
+		fn from_slice(bytes: &[u8]) -> Self {
+			Self
+		}
+		fn as_slice(&self) -> &[u8] {
+			&[]
+		}
+		fn to_raw_vec(&self) -> Vec<u8> {
+			vec![]
+		}
+	}
 	impl Pair for TestPair {
-		type Public = ();
+		type Public = TestPublic;
 		type Seed = [u8; 0];
 		type Signature = ();
 		type DeriveError = ();
@@ -464,7 +521,7 @@ mod tests {
 			_message: M,
 			_pubkey: P
 		) -> bool { true }
-		fn public(&self) -> Self::Public { () }
+		fn public(&self) -> Self::Public { TestPublic }
 		fn from_standard_components<I: Iterator<Item=DeriveJunction>>(
 			phrase: &str,
 			password: Option<&str>,
@@ -481,6 +538,9 @@ mod tests {
 		{
 			Ok(TestPair::Seed(seed.to_owned()))
 		}
+		fn to_raw_vec(&self) -> Vec<u8> {
+			vec![]
+		}
 	}
 
 	#[test]