Babe VRF Signing in keystore

client/consensus/babe/src/authorship.rs

@@ -17,14 +17,16 @@
 //! BABE authority selection and slot claiming.
 
 use sp_consensus_babe::{
-	make_transcript, AuthorityId, BabeAuthorityWeight, BABE_VRF_PREFIX,
+	BABE_ENGINE_ID, BABE_VRF_PREFIX,
+	AuthorityId, BabeAuthorityWeight,
 	SlotNumber, AuthorityPair,
+	make_transcript,
 };
 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, traits::VRFSigner};
 use codec::Encode;
 use schnorrkel::vrf::VRFInOut;
 use sp_core::Pair;
@@ -125,6 +127,7 @@ fn claim_secondary_slot(
 	slot_number: SlotNumber,
 	epoch: &Epoch,
 	key_pairs: &[(AuthorityPair, usize)],
+	keystore: &KeyStorePtr,
 	author_secondary_vrf: bool,
 ) -> Option<(PreDigest, AuthorityId)> {
 	let Epoch { authorities, randomness, epoch_index, .. } = epoch;
@@ -142,18 +145,16 @@ fn claim_secondary_slot(
 	for (pair, authority_index) in key_pairs {
 		if pair.public() == *expected_author {
 			let pre_digest = if author_secondary_vrf {
-				let transcript = super::authorship::make_transcript(
-					randomness,
-					slot_number,
-					*epoch_index,
+				let (output, proof) = keystore.read().vrf_sign(
+					pair.as_ref(), &BABE_ENGINE_ID, randomness, slot_number, *epoch_index,
 				);
-
-				let s = get_keypair(&pair).vrf_sign(transcript);
+				let proof = schnorrkel::vrf::VRFProof::from_bytes(&proof).ok()?;
+				let output = schnorrkel::vrf::VRFOutput::from_bytes(&output).ok()?;
 
 				PreDigest::SecondaryVRF(SecondaryVRFPreDigest {
 					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,
 				})
 			} else {
@@ -188,17 +189,18 @@ pub fn claim_slot(
 			})
 			.collect::<Vec<_>>()
 	};
-	claim_slot_using_key_pairs(slot_number, epoch, &key_pairs)
+	claim_slot_using_key_pairs(slot_number, epoch, keystore, &key_pairs)
 }
 
 /// 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(
 	slot_number: SlotNumber,
 	epoch: &Epoch,
+	keystore: &KeyStorePtr,
 	key_pairs: &[(AuthorityPair, 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, &key_pairs)
 		.or_else(|| {
 			if epoch.config.allowed_slots.is_secondary_plain_slots_allowed() ||
 				epoch.config.allowed_slots.is_secondary_vrf_slots_allowed()
@@ -207,6 +209,7 @@ pub fn claim_slot_using_key_pairs(
 					slot_number,
 					&epoch,
 					&key_pairs,
+					keystore,
 					epoch.config.allowed_slots.is_secondary_vrf_slots_allowed(),
 				)
 			} else {
@@ -215,11 +218,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,6 +226,7 @@ fn claim_primary_slot(
 	slot_number: SlotNumber,
 	epoch: &Epoch,
 	c: (u64, u64),
+	keystore: &KeyStorePtr,
 	key_pairs: &[(AuthorityPair, usize)],
 ) -> Option<(PreDigest, AuthorityId)> {
 	let Epoch { authorities, randomness, epoch_index, .. } = epoch;
@@ -241,16 +240,26 @@ fn claim_primary_slot(
 		// 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 {
-					slot_number,
-					vrf_output: VRFOutput(s.0.to_output()),
-					vrf_proof: VRFProof(s.1),
-					authority_index: *authority_index as u32,
-				})
-			});
+		let (output, proof) = keystore.read().vrf_sign(
+			pair.as_ref(), &BABE_ENGINE_ID, randomness, slot_number, *epoch_index,
+		);
+		let proof = schnorrkel::vrf::VRFProof::from_bytes(&proof).ok()?;
+		let output = schnorrkel::vrf::VRFOutput::from_bytes(&output).ok()?;
+		let inout = output.attach_input_hash(
+			&pair.as_ref().as_ref().public,
+			transcript
+		).ok()?;
+
+		let pre_digest = if super::authorship::check_primary_threshold(&inout, threshold) {
+			Some(PreDigest::Primary(PrimaryPreDigest {
+				slot_number,
+				vrf_output: VRFOutput(output),
+				vrf_proof: VRFProof(proof),
+				authority_index: *authority_index as u32,
+			}))
+		} else {
+			None
+		};
 
 		// early exit on first successful claim
 		if let Some(pre_digest) = pre_digest {

client/keystore/Cargo.toml

@@ -18,10 +18,12 @@ 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"
+schnorrkel = { version = "0.9.1", features = ["preaudit_deprecated", "u64_backend"], default-features = false}
 
 [dev-dependencies]
 tempfile = "3.1.0"

client/keystore/src/lib.rs

@@ -20,11 +20,13 @@
 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, BareCryptoStoreError as TraitError, VRFSigner},
+	sr25519::Pair as Sr25519Pair,
 	Encode,
 };
 use sp_application_crypto::{AppKey, AppPublic, AppPair, ed25519, sr25519, ecdsa};
 use parking_lot::RwLock;
+use merlin::Transcript;
 
 /// Keystore pointer
 pub type KeyStorePtr = Arc<RwLock<Store>>;
@@ -296,6 +298,20 @@ impl Store {
 
 		Ok(public_keys)
 	}
+
+	fn make_vrf_transcript(
+		&self,
+		label: &'static [u8],
+		randomness: &[u8],
+		slot_number: u64,
+		epoch: u64
+	) -> Transcript {
+		let mut transcript = Transcript::new(label.clone());
+		transcript.append_u64(b"slot number", slot_number);
+		transcript.append_u64(b"current epoch", epoch);
+		transcript.append_message(b"chain randomness", &randomness[..]);
+		transcript
+	}
 }
 
 impl BareCryptoStore for Store {
@@ -440,6 +456,22 @@ impl BareCryptoStore for Store {
 	}
 }
 
+impl VRFSigner for Store {
+	fn vrf_sign(
+		&self,
+		pair: &Sr25519Pair,
+		label: &'static [u8],
+		randomness: &[u8],
+		slot_number: u64,
+		epoch: u64
+	) -> (Vec<u8>, Vec<u8>) {
+		let transcript = self.make_vrf_transcript(label, randomness, slot_number, epoch);
+		let pair = pair.as_ref();
+		let (inout, proof, _) = pair.vrf_sign(transcript);
+		(inout.to_output().to_bytes().to_vec(), proof.to_bytes().to_vec())
+	}
+}
+
 #[cfg(test)]
 mod tests {
 	use super::*;

primitives/core/src/traits.rs

@@ -174,6 +174,19 @@ pub trait BareCryptoStore: Send + Sync {
 	}
 }
 
+/// VRF-signing capability.
+pub trait VRFSigner {
+	/// VRF-sign round
+	fn vrf_sign(
+		&self,
+		pair: &sr25519::Pair,
+		label: &'static [u8],
+		randomness: &[u8],
+		slot_number: u64,
+		epoch: u64
+	) -> (Vec<u8>, Vec<u8>);
+}
+
 /// A pointer to the key store.
 pub type BareCryptoStorePtr = Arc<parking_lot::RwLock<dyn BareCryptoStore>>;