Fix BABE logic leaks

client/consensus/babe/src/authorship.rs

@@ -18,17 +18,21 @@
 
 use sp_application_crypto::AppKey;
 use sp_consensus_babe::{
-	BABE_ENGINE_ID, BABE_VRF_PREFIX,
+	BABE_VRF_PREFIX,
 	AuthorityId, BabeAuthorityWeight,
 	SlotNumber,
+	make_transcript,
 };
 use sp_consensus_babe::digests::{
 	PreDigest, PrimaryPreDigest, SecondaryPlainPreDigest, SecondaryVRFPreDigest,
 };
 use sp_consensus_vrf::schnorrkel::{VRFOutput, VRFProof};
-use sp_core::{U256, blake2_256, traits::BareCryptoStore};
+use sp_core::{U256, blake2_256, crypto::Public, traits::BareCryptoStore};
 use codec::Encode;
-use schnorrkel::vrf::VRFInOut;
+use schnorrkel::{
+	keys::PublicKey,
+	vrf::VRFInOut,
+};
 use sc_keystore::KeyStorePtr;
 use super::Epoch;
 
@@ -144,19 +148,19 @@ fn claim_secondary_slot(
 	for (authority_id, authority_index) in keys {
 		if authority_id == expected_author {
 			let pre_digest = if author_secondary_vrf {
-				let result = keystore.read().vrf_sign(
-					AuthorityId::ID,
-					authority_id.as_ref(),
-					&BABE_ENGINE_ID,
-					BABE_VRF_PREFIX,
+				let transcript = super::authorship::make_transcript(
 					randomness,
 					slot_number,
 					*epoch_index,
-					u128::MAX,
 				);
-				if let Ok(Some((output, proof)))  = result {
-					let proof = schnorrkel::vrf::VRFProof::from_bytes(&proof).ok()?;
-					let output = schnorrkel::vrf::VRFOutput::from_bytes(&output).ok()?;
+				let result = keystore.read().vrf_sign(
+					AuthorityId::ID,
+					authority_id.as_ref(),
+					transcript,
+				);
+				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,
@@ -239,6 +243,7 @@ fn claim_primary_slot(
 	let Epoch { authorities, randomness, epoch_index, .. } = epoch;
 
 	for (authority_id, authority_index) in keys {
+		let transcript = super::authorship::make_transcript(randomness, slot_number, *epoch_index);
 		// Compute the threshold we will use.
 		//
 		// We already checked that authorities contains `key.public()`, so it can't
@@ -248,25 +253,26 @@ fn claim_primary_slot(
 		let result = keystore.read().vrf_sign(
 			AuthorityId::ID,
 			authority_id.as_ref(),
-			&BABE_ENGINE_ID,
-			BABE_VRF_PREFIX,
-			randomness,
-			slot_number,
-			*epoch_index,
-			threshold,
+			transcript.clone(),
 		);
-		if let Ok(Some((output, proof)))  = result {
-			let proof = schnorrkel::vrf::VRFProof::from_bytes(&proof).ok()?;
-			let output = schnorrkel::vrf::VRFOutput::from_bytes(&output).ok()?;
-
-			let pre_digest = PreDigest::Primary(PrimaryPreDigest {
-				slot_number,
-				vrf_output: VRFOutput(output),
-				vrf_proof: VRFProof(proof),
-				authority_index: *authority_index as u32,
-			});
-
-			return Some((pre_digest, authority_id.clone()));
+		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(output),
+					vrf_proof: VRFProof(proof),
+					authority_index: *authority_index as u32,
+				});
+
+				return Some((pre_digest, authority_id.clone()));
+			}
 		}
 	}
 

client/keystore/src/lib.rs

@@ -20,14 +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, Error as TraitError},
+	traits::{BareCryptoStore, Error as TraitError, VRFSignature},
 	sr25519::{Public as Sr25519Public, Pair as Sr25519Pair},
 	Encode,
 };
 use sp_application_crypto::{AppKey, AppPublic, AppPair, ed25519, sr25519, ecdsa};
 use parking_lot::RwLock;
 use merlin::Transcript;
-use schnorrkel::vrf::VRFInOut;
 
 /// Keystore pointer
 pub type KeyStorePtr = Arc<RwLock<Store>>;
@@ -299,24 +298,6 @@ 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
-	}
-
-	fn check_primary_threshold(&self, prefix: &'static [u8], inout: &VRFInOut, threshold: u128) -> bool {
-		u128::from_le_bytes(inout.make_bytes::<[u8; 16]>(prefix)) < threshold
-	}
 }
 
 impl BareCryptoStore for Store {
@@ -464,23 +445,16 @@ impl BareCryptoStore for Store {
 		&self,
 		key_type: KeyTypeId,
 		public: &Sr25519Public,
-		label: &'static [u8],
-		prefix: &'static [u8],
-		randomness: &[u8],
-		slot_number: u64,
-		epoch: u64,
-		threshold: u128,
-	) -> std::result::Result<Option<(Vec<u8>, Vec<u8>)>, TraitError> {
-		let transcript = self.make_vrf_transcript(label, randomness, slot_number, epoch);
+		transcript: Transcript,
+	) -> std::result::Result<VRFSignature, TraitError> {
 		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);
-		if self.check_primary_threshold(prefix, &inout, threshold) {
-			Ok(Some((inout.to_output().to_bytes().to_vec(), proof.to_bytes().to_vec())))
-		} else {
-			Ok(None)
-		}
+		Ok(VRFSignature {
+			output: inout.to_output().to_bytes().to_vec(),
+			proof: proof.to_bytes().to_vec(),
+		})
 	}
 }
 

primitives/core/src/testing.rs

@@ -22,10 +22,12 @@ use crate::crypto::KeyTypeId;
 use crate::{
 	crypto::{Pair, Public, CryptoTypePublicPair},
 	ed25519, sr25519, ecdsa,
-	traits::Error
+	traits::{Error, VRFSignature},
 };
 #[cfg(feature = "std")]
 use std::collections::HashSet;
+#[cfg(feature = "std")]
+use merlin::Transcript;
 /// Key type for generic Ed25519 key.
 pub const ED25519: KeyTypeId = KeyTypeId(*b"ed25");
 /// Key type for generic Sr 25519 key.
@@ -243,13 +245,8 @@ impl crate::traits::BareCryptoStore for KeyStore {
 		&self,
 		_key_type: KeyTypeId,
 		_public: &sr25519::Public,
-		_label: &'static [u8],
-		_prefix: &'static [u8],
-		_randomness: &[u8],
-		_slot_number: u64,
-		_epoch: u64,
-		_threshold: u128
-	) -> Result<Option<(Vec<u8>, Vec<u8>)>, Error> {
+		_transcript: Transcript,
+	) -> Result<VRFSignature, Error> {
 		Err(Error::Unavailable)
 	}
 }

primitives/core/src/traits.rs

@@ -28,6 +28,7 @@ use std::{
 	panic::UnwindSafe,
 	sync::Arc,
 };
+use merlin::Transcript;
 
 pub use sp_externalities::{Externalities, ExternalitiesExt};
 
@@ -51,6 +52,13 @@ pub enum Error {
 	Other(String)
 }
 
+/// VRF signature data
+pub struct VRFSignature {
+	/// The VRFOutput serialized
+	pub output: Vec<u8>,
+	/// The calculated VRFProof
+	pub proof: Vec<u8>,
+}
 
 /// Something that generates, stores and provides access to keys.
 pub trait BareCryptoStore: Send + Sync {
@@ -174,18 +182,26 @@ pub trait BareCryptoStore: Send + Sync {
 		Ok(keys.iter().map(|k| self.sign_with(id, k, msg)).collect())
 	}
 
-	/// Generate VRF proof for given transacript data.
+	/// Generate VRF signature for given transcript data.
+	///
+	/// Receives KeyTypeId and Public key to be able to map
+	/// them to a private key that exists in the keystore which
+	/// is, in turn, used for signing the provided transcript.
+	///
+	/// Returns a result containing the signature data.
+	/// Namely, VRFOutput and VRFProof which are returned
+	/// inside the `VRFSignature` container struct.
+	///
+	/// This function will return an error in the cases where
+	/// the public key and key type provided do not match a private
+	/// key in the keystore. Or, in the context of remote signing
+	/// an error could be a network one.
 	fn vrf_sign(
 		&self,
 		key_type: KeyTypeId,
 		public: &sr25519::Public,
-		label: &'static [u8],
-		prefix: &'static [u8],
-		randomness: &[u8],
-		slot_number: u64,
-		epoch: u64,
-		threshold: u128,
-	) -> Result<Option<(Vec<u8>, Vec<u8>)>, Error>;
+		transcript: Transcript,
+	) -> Result<VRFSignature, Error>;
 }
 
 /// A pointer to the key store.