Babe VRF Signing in keystore

client/consensus/babe/rpc/Cargo.toml

@@ -27,12 +27,12 @@ derive_more = "0.99.2"
 sp-api = { version = "2.0.0-rc2", path = "../../../../primitives/api" }
 sp-consensus = { version = "0.8.0-rc2", path = "../../../../primitives/consensus/common" }
 sp-core = { version = "2.0.0-rc2", path = "../../../../primitives/core" }
+sp-application-crypto = { version = "2.0.0-rc2", path = "../../../../primitives/application-crypto" }
 sc-keystore = { version = "2.0.0-rc2", path = "../../../keystore" }
 
 [dev-dependencies]
 sc-consensus = { version = "0.8.0-rc2", path = "../../../consensus/common" }
 serde_json = "1.0.50"
-sp-application-crypto = { version = "2.0.0-rc2", path = "../../../../primitives/application-crypto" }
 sp-keyring = { version = "2.0.0-rc2", path = "../../../../primitives/keyring" }
 substrate-test-runtime-client = { version = "2.0.0-rc2", path = "../../../../test-utils/runtime/client" }
 tempfile = "3.1.0"

client/consensus/babe/rpc/src/lib.rs

@@ -32,6 +32,11 @@ use sp_consensus_babe::{
 	digests::PreDigest,
 };
 use serde::{Deserialize, Serialize};
+use sp_core::{
+	crypto::Public,
+	traits::BareCryptoStore,
+};
+use sp_application_crypto::AppKey;
 use sc_keystore::KeyStorePtr;
 use sc_rpc_api::DenyUnsafe;
 use sp_api::{ProvideRuntimeApi, BlockId};
@@ -126,22 +131,21 @@ impl<B, C, SC> BabeApi for BabeRpcHandler<B, C, SC>
 
 			let mut claims: HashMap<AuthorityId, EpochAuthorship> = HashMap::new();
 
-			let key_pairs = {
-				let keystore = keystore.read();
-				epoch.authorities.iter()
-					.enumerate()
-					.flat_map(|(i, a)| {
-						keystore
-							.key_pair::<sp_consensus_babe::AuthorityPair>(&a.0)
-							.ok()
-							.map(|kp| (kp, i))
-					})
-					.collect::<Vec<_>>()
-			};
+			let ks = keystore.read();
+			let keys = epoch.authorities.iter()
+				.enumerate()
+				.filter_map(|(i, a)| {
+					if ks.has_keys(&[(a.0.to_raw_vec(), AuthorityId::ID)]) {
+						Some((a.0.clone(), i))
+					} else {
+						None
+					}
+				})
+				.collect::<Vec<_>>();
 
 			for slot_number in epoch_start..epoch_end {
 				if let Some((claim, key)) =
-					authorship::claim_slot_using_key_pairs(slot_number, &epoch, &key_pairs)
+					authorship::claim_slot_using_keys(slot_number, &epoch, &keystore, &keys)
 				{
 					match claim {
 						PreDigest::Primary { .. } => {

client/consensus/babe/src/authorship.rs

@@ -16,18 +16,24 @@
 
 //! BABE authority selection and slot claiming.
 
+use sp_application_crypto::AppKey;
 use sp_consensus_babe::{
-	make_transcript, AuthorityId, BabeAuthorityWeight, BABE_VRF_PREFIX,
-	SlotNumber, AuthorityPair,
+	BABE_VRF_PREFIX,
+	AuthorityId, BabeAuthorityWeight,
+	SlotNumber,
+	make_transcript,
+	make_transcript_data,
 };
 use sp_consensus_babe::digests::{
 	PreDigest, PrimaryPreDigest, SecondaryPlainPreDigest, SecondaryVRFPreDigest,
 };
 use sp_consensus_vrf::schnorrkel::{VRFOutput, VRFProof};
-use sp_core::{U256, blake2_256};
+use sp_core::{U256, blake2_256, crypto::Public, traits::BareCryptoStore};
 use codec::Encode;
-use schnorrkel::vrf::VRFInOut;
-use sp_core::Pair;
+use schnorrkel::{
+	keys::PublicKey,
+	vrf::VRFInOut,
+};
 use sc_keystore::KeyStorePtr;
 use super::Epoch;
 
@@ -124,7 +130,8 @@ pub(super) fn secondary_slot_author(
 fn claim_secondary_slot(
 	slot_number: SlotNumber,
 	epoch: &Epoch,
-	key_pairs: &[(AuthorityPair, usize)],
+	keys: &[(AuthorityId, usize)],
+	keystore: &KeyStorePtr,
 	author_secondary_vrf: bool,
 ) -> Option<(PreDigest, AuthorityId)> {
 	let Epoch { authorities, randomness, epoch_index, .. } = epoch;
@@ -139,31 +146,43 @@ fn claim_secondary_slot(
 		*randomness,
 	)?;
 
-	for (pair, authority_index) in key_pairs {
-		if pair.public() == *expected_author {
+	let ks = keystore.read();
+	for (authority_id, authority_index) in keys {
+		if authority_id == expected_author {
 			let pre_digest = if author_secondary_vrf {
-				let transcript = super::authorship::make_transcript(
+				let transcript_data = super::authorship::make_transcript_data(
 					randomness,
 					slot_number,
-					*epoch_index,
+					*epoch_index
 				);
-
-				let s = get_keypair(&pair).vrf_sign(transcript);
-
-				PreDigest::SecondaryVRF(SecondaryVRFPreDigest {
-					slot_number,
-					vrf_output: VRFOutput(s.0.to_output()),
-					vrf_proof: VRFProof(s.1),
-					authority_index: *authority_index as u32,
-				})
+				let result = ks.sr25519_vrf_sign(
+					AuthorityId::ID,
+					authority_id.as_ref(),
+					transcript_data,
+				);
+				if let Ok(signature)  = result {
+					let proof = schnorrkel::vrf::VRFProof::from_bytes(&signature.proof).ok()?;
+					let output = schnorrkel::vrf::VRFOutput::from_bytes(&signature.output).ok()?;
+
+					Some(PreDigest::SecondaryVRF(SecondaryVRFPreDigest {
+						slot_number,
+						vrf_output: VRFOutput(output),
+						vrf_proof: VRFProof(proof),
+						authority_index: *authority_index as u32,
+					}))
+				} else {
+					None
+				}
 			} else {
-				PreDigest::SecondaryPlain(SecondaryPlainPreDigest {
+				Some(PreDigest::SecondaryPlain(SecondaryPlainPreDigest {
 					slot_number,
 					authority_index: *authority_index as u32,
-				})
+				}))
 			};
 
-			return Some((pre_digest, pair.public()));
+			if let Some(pre_digest) = pre_digest {
+				return Some((pre_digest, authority_id.clone()));
+			}
 		}
 	}
 
@@ -179,34 +198,31 @@ pub fn claim_slot(
 	epoch: &Epoch,
 	keystore: &KeyStorePtr,
 ) -> Option<(PreDigest, AuthorityId)> {
-	let key_pairs = {
-		let keystore = keystore.read();
-		epoch.authorities.iter()
-			.enumerate()
-			.flat_map(|(i, a)| {
-				keystore.key_pair::<AuthorityPair>(&a.0).ok().map(|kp| (kp, i))
-			})
-			.collect::<Vec<_>>()
-	};
-	claim_slot_using_key_pairs(slot_number, epoch, &key_pairs)
+	let authorities = epoch.authorities.iter()
+		.enumerate()
+		.map(|(index, a)| (a.0.clone(), index))
+		.collect::<Vec<_>>();
+	claim_slot_using_keys(slot_number, epoch, keystore, &authorities)
 }
 
 /// Like `claim_slot`, but allows passing an explicit set of key pairs. Useful if we intend
 /// to make repeated calls for different slots using the same key pairs.
-pub fn claim_slot_using_key_pairs(
+pub fn claim_slot_using_keys(
 	slot_number: SlotNumber,
 	epoch: &Epoch,
-	key_pairs: &[(AuthorityPair, usize)],
+	keystore: &KeyStorePtr,
+	keys: &[(AuthorityId, usize)],
 ) -> Option<(PreDigest, AuthorityId)> {
-	claim_primary_slot(slot_number, epoch, epoch.config.c, &key_pairs)
+	claim_primary_slot(slot_number, epoch, epoch.config.c, keystore, &keys)
 		.or_else(|| {
 			if epoch.config.allowed_slots.is_secondary_plain_slots_allowed() ||
 				epoch.config.allowed_slots.is_secondary_vrf_slots_allowed()
 			{
 				claim_secondary_slot(
 					slot_number,
 					&epoch,
-					&key_pairs,
+					keys,
+					keystore,
 					epoch.config.allowed_slots.is_secondary_vrf_slots_allowed(),
 				)
 			} else {
@@ -215,11 +231,6 @@ pub fn claim_slot_using_key_pairs(
 		})
 }
 
-fn get_keypair(q: &AuthorityPair) -> &schnorrkel::Keypair {
-	use sp_core::crypto::IsWrappedBy;
-	sp_core::sr25519::Pair::from_ref(q).as_ref()
-}
-
 /// Claim a primary slot if it is our turn.  Returns `None` if it is not our turn.
 /// This hashes the slot number, epoch, genesis hash, and chain randomness into
 /// the VRF.  If the VRF produces a value less than `threshold`, it is our turn,
@@ -228,33 +239,52 @@ fn claim_primary_slot(
 	slot_number: SlotNumber,
 	epoch: &Epoch,
 	c: (u64, u64),
-	key_pairs: &[(AuthorityPair, usize)],
+	keystore: &KeyStorePtr,
+	keys: &[(AuthorityId, usize)],
 ) -> Option<(PreDigest, AuthorityId)> {
 	let Epoch { authorities, randomness, epoch_index, .. } = epoch;
 
-	for (pair, authority_index) in key_pairs {
-		let transcript = super::authorship::make_transcript(randomness, slot_number, *epoch_index);
-
+	let ks = keystore.read();
+	for (authority_id, authority_index) in keys {
+		let transcript = super::authorship::make_transcript(
+			randomness,
+			slot_number,
+			*epoch_index
+		);
+		let transcript_data = super::authorship::make_transcript_data(
+			randomness,
+			slot_number,
+			*epoch_index
+		);
 		// Compute the threshold we will use.
 		//
 		// We already checked that authorities contains `key.public()`, so it can't
 		// be empty.  Therefore, this division in `calculate_threshold` is safe.
 		let threshold = super::authorship::calculate_primary_threshold(c, authorities, *authority_index);
 
-		let pre_digest = get_keypair(pair)
-			.vrf_sign_after_check(transcript, |inout| super::authorship::check_primary_threshold(inout, threshold))
-			.map(|s| {
-				PreDigest::Primary(PrimaryPreDigest {
+		let result = ks.sr25519_vrf_sign(
+			AuthorityId::ID,
+			authority_id.as_ref(),
+			transcript_data,
+		);
+		if let Ok(signature)  = result {
+			let proof = schnorrkel::vrf::VRFProof::from_bytes(&signature.proof).ok()?;
+			let output = schnorrkel::vrf::VRFOutput::from_bytes(&signature.output).ok()?;
+			let public = PublicKey::from_bytes(&authority_id.to_raw_vec()).ok()?;
+			let inout = match output.attach_input_hash(&public, transcript) {
+				Ok(inout) => inout,
+				Err(_) => continue,
+			};
+			if super::authorship::check_primary_threshold(&inout, threshold) {
+				let pre_digest = PreDigest::Primary(PrimaryPreDigest {
 					slot_number,
-					vrf_output: VRFOutput(s.0.to_output()),
-					vrf_proof: VRFProof(s.1),
+					vrf_output: VRFOutput(output),
+					vrf_proof: VRFProof(proof),
 					authority_index: *authority_index as u32,
-				})
-			});
+				});
 
-		// early exit on first successful claim
-		if let Some(pre_digest) = pre_digest {
-			return Some((pre_digest, pair.public()));
+				return Some((pre_digest, authority_id.clone()));
+			}
 		}
 	}
 

client/keystore/Cargo.toml

@@ -18,10 +18,11 @@ derive_more = "0.99.2"
 sp-core = { version = "2.0.0-rc2", path = "../../primitives/core" }
 sp-application-crypto = { version = "2.0.0-rc2", path = "../../primitives/application-crypto" }
 hex = "0.4.0"
+merlin = { version = "2.0", default-features = false }
+parking_lot = "0.10.0"
 rand = "0.7.2"
 serde_json = "1.0.41"
 subtle = "2.1.1"
-parking_lot = "0.10.0"
 
 [dev-dependencies]
 tempfile = "3.1.0"

client/keystore/src/lib.rs

@@ -20,7 +20,9 @@
 use std::{collections::{HashMap, HashSet}, path::PathBuf, fs::{self, File}, io::{self, Write}, sync::Arc};
 use sp_core::{
 	crypto::{IsWrappedBy, CryptoTypePublicPair, KeyTypeId, Pair as PairT, Protected, Public},
-	traits::{BareCryptoStore, BareCryptoStoreError as TraitError},
+	traits::{BareCryptoStore, Error as TraitError},
+	sr25519::{Public as Sr25519Public, Pair as Sr25519Pair},
+	vrf::{VRFTranscriptData, VRFSignature, make_transcript},
 	Encode,
 };
 use sp_application_crypto::{AppKey, AppPublic, AppPair, ed25519, sr25519, ecdsa};
@@ -438,6 +440,23 @@ impl BareCryptoStore for Store {
 	fn has_keys(&self, public_keys: &[(Vec<u8>, KeyTypeId)]) -> bool {
 		public_keys.iter().all(|(p, t)| self.key_phrase_by_type(&p, *t).is_ok())
 	}
+
+	fn sr25519_vrf_sign(
+		&self,
+		key_type: KeyTypeId,
+		public: &Sr25519Public,
+		transcript_data: VRFTranscriptData,
+	) -> std::result::Result<VRFSignature, TraitError> {
+		let transcript = make_transcript(transcript_data);
+		let pair = self.key_pair_by_type::<Sr25519Pair>(public, key_type)
+			.map_err(|e| TraitError::PairNotFound(e.to_string()))?;
+
+		let (inout, proof, _) = pair.as_ref().vrf_sign(transcript);
+		Ok(VRFSignature {
+			output: inout.to_output().to_bytes().to_vec(),
+			proof: proof.to_bytes().to_vec(),
+		})
+	}
 }
 
 #[cfg(test)]

primitives/consensus/babe/Cargo.toml

@@ -17,6 +17,7 @@ codec = { package = "parity-scale-codec", version = "1.3.0", default-features =
 merlin = { version = "2.0", default-features = false }
 sp-std = { version = "2.0.0-rc2", default-features = false, path = "../../std" }
 sp-api = { version = "2.0.0-rc2", default-features = false, path = "../../api" }
+sp-core = { version = "2.0.0-rc2", default-features = false, path = "../../core" }
 sp-consensus = { version = "0.8.0-rc2", optional = true, path = "../common" }
 sp-consensus-vrf = { version = "0.8.0-rc2", path = "../vrf", default-features = false }
 sp-inherents = { version = "2.0.0-rc2", default-features = false, path = "../../inherents" }
@@ -26,6 +27,7 @@ sp-timestamp = { version = "2.0.0-rc2", default-features = false, path = "../../
 [features]
 default = ["std"]
 std = [
+	"sp-core/std",
 	"sp-application-crypto/std",
 	"codec/std",
 	"merlin/std",

primitives/consensus/babe/src/lib.rs

@@ -31,6 +31,8 @@ pub use merlin::Transcript;
 use codec::{Encode, Decode};
 use sp_std::vec::Vec;
 use sp_runtime::{ConsensusEngineId, RuntimeDebug};
+#[cfg(feature = "std")]
+use sp_core::vrf::{VRFTranscriptData, VRFTranscriptValue};
 use crate::digests::{NextEpochDescriptor, NextConfigDescriptor};
 
 mod app {
@@ -94,6 +96,23 @@ pub fn make_transcript(
 	transcript
 }
 
+/// Make a VRF transcript data container
+#[cfg(feature = "std")]
+pub fn make_transcript_data(
+	randomness: &Randomness,
+	slot_number: u64,
+	epoch: u64,
+) -> VRFTranscriptData {
+	VRFTranscriptData {
+		label: &BABE_ENGINE_ID,
+		items: vec![
+			("slot_number", VRFTranscriptValue::U64(slot_number)),
+			("current epoch", VRFTranscriptValue::U64(epoch)),
+			("chain randomness", VRFTranscriptValue::Bytes(&randomness[..])),
+		]
+	}
+}
+
 /// An consensus log item for BABE.
 #[derive(Decode, Encode, Clone, PartialEq, Eq)]
 pub enum ConsensusLog {