Check primary threshold in signer

client/consensus/babe/src/authorship.rs

@@ -21,15 +21,14 @@ use sp_consensus_babe::{
 	BABE_ENGINE_ID, 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::VRFSigner, crypto::Public};
+use sp_core::{U256, blake2_256, traits::VRFSigner};
 use codec::Encode;
-use schnorrkel::{vrf::VRFInOut, keys::PublicKey};
+use schnorrkel::vrf::VRFInOut;
 use sc_keystore::KeyStorePtr;
 use super::Epoch;
 
@@ -145,31 +144,39 @@ fn claim_secondary_slot(
 	for (authority_id, authority_index) in keys {
 		if authority_id == expected_author {
 			let pre_digest = if author_secondary_vrf {
-				let (output, proof) = keystore.read().vrf_sign(
+				let result = keystore.read().vrf_sign(
 					AuthorityId::ID,
 					authority_id.as_ref(),
 					&BABE_ENGINE_ID,
+					BABE_VRF_PREFIX,
 					randomness,
 					slot_number,
 					*epoch_index,
-				).ok()?;
-				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(output),
-					vrf_proof: VRFProof(proof),
-					authority_index: *authority_index as u32,
-				})
+					u128::MAX,
+				).ok();
+				if let Some((output, proof))  = result {
+					let proof = schnorrkel::vrf::VRFProof::from_bytes(&proof).ok()?;
+					let output = schnorrkel::vrf::VRFOutput::from_bytes(&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, authority_id.clone()));
+			if let Some(pre_digest) = pre_digest {
+				return Some((pre_digest, authority_id.clone()));
+			}
 		}
 	}
 
@@ -232,38 +239,33 @@ 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
 		// be empty.  Therefore, this division in `calculate_threshold` is safe.
 		let threshold = super::authorship::calculate_primary_threshold(c, authorities, *authority_index);
 
-		let (output, proof) = keystore.read().vrf_sign(
-			AuthorityId::ID, authority_id.as_ref(), &BABE_ENGINE_ID, randomness, slot_number, *epoch_index,
-		).ok()?;
-		let proof = schnorrkel::vrf::VRFProof::from_bytes(&proof).ok()?;
-		let output = schnorrkel::vrf::VRFOutput::from_bytes(&output).ok()?;
-		let public_key = PublicKey::from_bytes(authority_id.as_slice()).ok()?;
-		let inout = output.attach_input_hash(
-			&public_key,
-			transcript
-		).ok()?;
-
-		let pre_digest = if super::authorship::check_primary_threshold(&inout, threshold) {
-			Some(PreDigest::Primary(PrimaryPreDigest {
+		let result = keystore.read().vrf_sign(
+			AuthorityId::ID,
+			authority_id.as_ref(),
+			&BABE_ENGINE_ID,
+			BABE_VRF_PREFIX,
+			randomness,
+			slot_number,
+			*epoch_index,
+			threshold,
+		).ok();
+		if let 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,
-			}))
-		} else {
-			None
-		};
+			});
 
-		// early exit on first successful claim
-		if let Some(pre_digest) = pre_digest {
 			return Some((pre_digest, authority_id.clone()));
 		}
 	}

client/keystore/src/lib.rs

@@ -27,6 +27,7 @@ use sp_core::{
 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>>;
@@ -312,6 +313,10 @@ impl Store {
 		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 {
@@ -462,14 +467,20 @@ impl VRFSigner for Store {
 		key_type: KeyTypeId,
 		public: &Sr25519Public,
 		label: &'static [u8],
+		prefix: &'static [u8],
 		randomness: &[u8],
 		slot_number: u64,
-		epoch: u64
+		epoch: u64,
+		threshold: u128,
 	) -> std::result::Result<(Vec<u8>, Vec<u8>), ()> {
 		let transcript = self.make_vrf_transcript(label, randomness, slot_number, epoch);
 		let pair = self.key_pair_by_type::<Sr25519Pair>(public, key_type).map_err(|_| ())?;
 		let (inout, proof, _) = pair.as_ref().vrf_sign(transcript);
-		Ok((inout.to_output().to_bytes().to_vec(), proof.to_bytes().to_vec()))
+		if self.check_primary_threshold(prefix, &inout, threshold) {
+			Ok((inout.to_output().to_bytes().to_vec(), proof.to_bytes().to_vec()))
+		} else {
+			Err(())
+		}
 	}
 }
 

primitives/core/src/traits.rs

@@ -182,9 +182,11 @@ pub trait VRFSigner {
 		key_type: KeyTypeId,
 		public: &sr25519::Public,
 		label: &'static [u8],
+		prefix: &'static [u8],
 		randomness: &[u8],
 		slot_number: u64,
-		epoch: u64
+		epoch: u64,
+		threshold: u128,
 	) -> Result<(Vec<u8>, Vec<u8>), ()>;
 }