VRF refactory (#13889)

* First iteration to encapsulate schnorrkel and merlin usage

* Remove schnorkel direct dependency from BABE pallet

* Remove schnorrkel direct dependency from BABE client

* Trivial renaming for VrfTranscript data and value

* Better errors

* Expose a function to get a schnorrkel friendly transcript

* Keep the vrf signature stuff together (preventing some clones around)

* Fix tests

* Remove vrf agnostic transcript and define it as an associated type for VrfSigner and VrfVerifier

* Fix babe pallet mock

* Inner types are required to be public for polkadot

* Update client/consensus/babe/src/verification.rs

Co-authored-by: Koute <koute@users.noreply.github.com>

* Nit

* Remove Deref implementations

* make_bytes as a method

* Trigger CI

---------

Co-authored-by: Koute <koute@users.noreply.github.com>

substrate/client/consensus/babe/src/authorship.rs

@@ -18,17 +18,19 @@
 
 //! BABE authority selection and slot claiming.
 
-use super::Epoch;
+use super::{Epoch, AUTHORING_SCORE_LENGTH, AUTHORING_SCORE_VRF_CONTEXT};
 use codec::Encode;
 use sc_consensus_epochs::Epoch as EpochT;
-use schnorrkel::{keys::PublicKey, vrf::VRFInOut};
 use sp_application_crypto::AppCrypto;
 use sp_consensus_babe::{
 	digests::{PreDigest, PrimaryPreDigest, SecondaryPlainPreDigest, SecondaryVRFPreDigest},
-	make_transcript, make_transcript_data, AuthorityId, BabeAuthorityWeight, Slot, BABE_VRF_PREFIX,
+	make_transcript, AuthorityId, BabeAuthorityWeight, Randomness, Slot,
+};
+use sp_core::{
+	blake2_256,
+	crypto::{ByteArray, Wraps},
+	U256,
 };
-use sp_consensus_vrf::schnorrkel::{VRFOutput, VRFProof};
-use sp_core::{blake2_256, crypto::ByteArray, U256};
 use sp_keystore::KeystorePtr;
 
 /// Calculates the primary selection threshold for a given authority, taking
@@ -95,19 +97,13 @@ pub(super) fn calculate_primary_threshold(
 	)
 }
 
-/// Returns true if the given VRF output is lower than the given threshold,
-/// false otherwise.
-pub(super) fn check_primary_threshold(inout: &VRFInOut, threshold: u128) -> bool {
-	u128::from_le_bytes(inout.make_bytes::<[u8; 16]>(BABE_VRF_PREFIX)) < threshold
-}
-
 /// Get the expected secondary author for the given slot and with given
 /// authorities. This should always assign the slot to some authority unless the
 /// authorities list is empty.
 pub(super) fn secondary_slot_author(
 	slot: Slot,
 	authorities: &[(AuthorityId, BabeAuthorityWeight)],
-	randomness: [u8; 32],
+	randomness: Randomness,
 ) -> Option<&AuthorityId> {
 	if authorities.is_empty() {
 		return None
@@ -152,18 +148,14 @@ fn claim_secondary_slot(
 	for (authority_id, authority_index) in keys {
 		if authority_id == expected_author {
 			let pre_digest = if author_secondary_vrf {
-				let transcript_data = make_transcript_data(randomness, slot, epoch_index);
-				let result = keystore.sr25519_vrf_sign(
-					AuthorityId::ID,
-					authority_id.as_ref(),
-					transcript_data,
-				);
-				if let Ok(Some(signature)) = result {
+				let transcript = make_transcript(randomness, slot, epoch_index);
+				let result =
+					keystore.sr25519_vrf_sign(AuthorityId::ID, authority_id.as_ref(), &transcript);
+				if let Ok(Some(vrf_signature)) = result {
 					Some(PreDigest::SecondaryVRF(SecondaryVRFPreDigest {
 						slot,
-						vrf_output: VRFOutput(signature.output),
-						vrf_proof: VRFProof(signature.proof),
 						authority_index: *authority_index as u32,
+						vrf_signature,
 					}))
 				} else {
 					None
@@ -247,25 +239,28 @@ fn claim_primary_slot(
 		epoch_index = epoch.clone_for_slot(slot).epoch_index;
 	}
 
-	for (authority_id, authority_index) in keys {
-		let transcript = make_transcript(randomness, slot, epoch_index);
-		let transcript_data = make_transcript_data(randomness, slot, epoch_index);
-		let result =
-			keystore.sr25519_vrf_sign(AuthorityId::ID, authority_id.as_ref(), transcript_data);
-		if let Ok(Some(signature)) = result {
-			let public = PublicKey::from_bytes(&authority_id.to_raw_vec()).ok()?;
-			let inout = match signature.output.attach_input_hash(&public, transcript) {
-				Ok(inout) => inout,
-				Err(_) => continue,
-			};
+	let transcript = make_transcript(randomness, slot, epoch_index);
 
+	for (authority_id, authority_index) in keys {
+		let result = keystore.sr25519_vrf_sign(AuthorityId::ID, authority_id.as_ref(), &transcript);
+		if let Ok(Some(vrf_signature)) = result {
 			let threshold = calculate_primary_threshold(c, authorities, *authority_index);
-			if check_primary_threshold(&inout, threshold) {
+
+			let can_claim = authority_id
+				.as_inner_ref()
+				.make_bytes::<[u8; AUTHORING_SCORE_LENGTH]>(
+					AUTHORING_SCORE_VRF_CONTEXT,
+					&transcript,
+					&vrf_signature.output,
+				)
+				.map(|bytes| u128::from_le_bytes(bytes) < threshold)
+				.unwrap_or_default();
+
+			if can_claim {
 				let pre_digest = PreDigest::Primary(PrimaryPreDigest {
 					slot,
-					vrf_output: VRFOutput(signature.output),
-					vrf_proof: VRFProof(signature.proof),
 					authority_index: *authority_index as u32,
+					vrf_signature,
 				});
 
 				return Some((pre_digest, authority_id.clone()))

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

@@ -86,7 +86,6 @@ use futures::{
 use log::{debug, info, log, trace, warn};
 use parking_lot::Mutex;
 use prometheus_endpoint::Registry;
-use schnorrkel::SignatureError;
 
 use sc_client_api::{
 	backend::AuxStore, AuxDataOperations, Backend as BackendT, FinalityNotification,
@@ -134,7 +133,7 @@ pub use sp_consensus_babe::{
 		PrimaryPreDigest, SecondaryPlainPreDigest,
 	},
 	AuthorityId, AuthorityPair, AuthoritySignature, BabeApi, BabeAuthorityWeight, BabeBlockWeight,
-	BabeConfiguration, BabeEpochConfiguration, ConsensusLog, BABE_ENGINE_ID, VRF_OUTPUT_LENGTH,
+	BabeConfiguration, BabeEpochConfiguration, ConsensusLog, Randomness, BABE_ENGINE_ID,
 };
 
 pub use aux_schema::load_block_weight as block_weight;
@@ -149,6 +148,12 @@ mod tests;
 
 const LOG_TARGET: &str = "babe";
 
+/// VRF context used for slots claiming lottery.
+const AUTHORING_SCORE_VRF_CONTEXT: &[u8] = b"substrate-babe-vrf";
+
+/// VRF output length for slots claiming lottery.
+const AUTHORING_SCORE_LENGTH: usize = 16;
+
 /// BABE epoch information
 #[derive(Decode, Encode, PartialEq, Eq, Clone, Debug, scale_info::TypeInfo)]
 pub struct Epoch {
@@ -161,7 +166,7 @@ pub struct Epoch {
 	/// The authorities and their weights.
 	pub authorities: Vec<(AuthorityId, BabeAuthorityWeight)>,
 	/// Randomness for this epoch.
-	pub randomness: [u8; VRF_OUTPUT_LENGTH],
+	pub randomness: Randomness,
 	/// Configuration of the epoch.
 	pub config: BabeEpochConfiguration,
 }
@@ -308,12 +313,12 @@ pub enum Error<B: BlockT> {
 	/// No secondary author expected.
 	#[error("No secondary author expected.")]
 	NoSecondaryAuthorExpected,
-	/// VRF verification of block by author failed
-	#[error("VRF verification of block by author {0:?} failed: threshold {1} exceeded")]
-	VRFVerificationOfBlockFailed(AuthorityId, u128),
 	/// VRF verification failed
-	#[error("VRF verification failed: {0:?}")]
-	VRFVerificationFailed(SignatureError),
+	#[error("VRF verification failed")]
+	VrfVerificationFailed,
+	/// Primary slot threshold too low
+	#[error("VRF output rejected, threshold {0} exceeded")]
+	VrfThresholdExceeded(u128),
 	/// Could not fetch parent header
 	#[error("Could not fetch parent header: {0}")]
 	FetchParentHeader(sp_blockchain::Error),

substrate/client/consensus/babe/src/migration.rs

@@ -18,7 +18,7 @@
 
 use crate::{
 	AuthorityId, BabeAuthorityWeight, BabeConfiguration, BabeEpochConfiguration, Epoch,
-	NextEpochDescriptor, VRF_OUTPUT_LENGTH,
+	NextEpochDescriptor, Randomness,
 };
 use codec::{Decode, Encode};
 use sc_consensus_epochs::Epoch as EpochT;
@@ -36,7 +36,7 @@ pub struct EpochV0 {
 	/// The authorities and their weights.
 	pub authorities: Vec<(AuthorityId, BabeAuthorityWeight)>,
 	/// Randomness for this epoch.
-	pub randomness: [u8; VRF_OUTPUT_LENGTH],
+	pub randomness: Randomness,
 }
 
 impl EpochT for EpochV0 {

substrate/client/consensus/babe/src/tests.rs

@@ -20,10 +20,6 @@
 
 use super::*;
 use authorship::claim_slot;
-use rand_chacha::{
-	rand_core::{RngCore, SeedableRng},
-	ChaChaRng,
-};
 use sc_block_builder::{BlockBuilder, BlockBuilderProvider};
 use sc_client_api::{backend::TransactionFor, BlockchainEvents, Finalizer};
 use sc_consensus::{BoxBlockImport, BoxJustificationImport};
@@ -33,16 +29,13 @@ use sc_network_test::{Block as TestBlock, *};
 use sp_application_crypto::key_types::BABE;
 use sp_consensus::{DisableProofRecording, NoNetwork as DummyOracle, Proposal};
 use sp_consensus_babe::{
-	inherents::InherentDataProvider, make_transcript, make_transcript_data, AllowedSlots,
-	AuthorityId, AuthorityPair, Slot,
+	inherents::InherentDataProvider, make_transcript, AllowedSlots, AuthorityId, AuthorityPair,
+	Slot,
 };
 use sp_consensus_slots::SlotDuration;
-use sp_consensus_vrf::schnorrkel::VRFOutput;
 use sp_core::crypto::Pair;
 use sp_keyring::Sr25519Keyring;
-use sp_keystore::{
-	testing::MemoryKeystore, vrf::make_transcript as transcript_from_data, Keystore,
-};
+use sp_keystore::{testing::MemoryKeystore, Keystore};
 use sp_runtime::{
 	generic::{Digest, DigestItem},
 	traits::Block as BlockT,
@@ -637,24 +630,24 @@ fn claim_vrf_check() {
 		PreDigest::Primary(d) => d,
 		v => panic!("Unexpected pre-digest variant {:?}", v),
 	};
-	let transcript = make_transcript_data(&epoch.randomness.clone(), 0.into(), epoch.epoch_index);
+	let transcript = make_transcript(&epoch.randomness.clone(), 0.into(), epoch.epoch_index);
 	let sign = keystore
-		.sr25519_vrf_sign(AuthorityId::ID, &public, transcript)
+		.sr25519_vrf_sign(AuthorityId::ID, &public, &transcript)
 		.unwrap()
 		.unwrap();
-	assert_eq!(pre_digest.vrf_output, VRFOutput(sign.output));
+	assert_eq!(pre_digest.vrf_signature.output, sign.output);
 
 	// We expect a SecondaryVRF claim for slot 1
 	let pre_digest = match claim_slot(1.into(), &epoch, &keystore).unwrap().0 {
 		PreDigest::SecondaryVRF(d) => d,
 		v => panic!("Unexpected pre-digest variant {:?}", v),
 	};
-	let transcript = make_transcript_data(&epoch.randomness.clone(), 1.into(), epoch.epoch_index);
+	let transcript = make_transcript(&epoch.randomness.clone(), 1.into(), epoch.epoch_index);
 	let sign = keystore
-		.sr25519_vrf_sign(AuthorityId::ID, &public, transcript)
+		.sr25519_vrf_sign(AuthorityId::ID, &public, &transcript)
 		.unwrap()
 		.unwrap();
-	assert_eq!(pre_digest.vrf_output, VRFOutput(sign.output));
+	assert_eq!(pre_digest.vrf_signature.output, sign.output);
 
 	// Check that correct epoch index has been used if epochs are skipped (primary VRF)
 	let slot = Slot::from(103);
@@ -663,13 +656,13 @@ fn claim_vrf_check() {
 		v => panic!("Unexpected claim variant {:?}", v),
 	};
 	let fixed_epoch = epoch.clone_for_slot(slot);
-	let transcript = make_transcript_data(&epoch.randomness.clone(), slot, fixed_epoch.epoch_index);
+	let transcript = make_transcript(&epoch.randomness.clone(), slot, fixed_epoch.epoch_index);
 	let sign = keystore
-		.sr25519_vrf_sign(AuthorityId::ID, &public, transcript)
+		.sr25519_vrf_sign(AuthorityId::ID, &public, &transcript)
 		.unwrap()
 		.unwrap();
 	assert_eq!(fixed_epoch.epoch_index, 11);
-	assert_eq!(claim.vrf_output, VRFOutput(sign.output));
+	assert_eq!(claim.vrf_signature.output, sign.output);
 
 	// Check that correct epoch index has been used if epochs are skipped (secondary VRF)
 	let slot = Slot::from(100);
@@ -678,13 +671,13 @@ fn claim_vrf_check() {
 		v => panic!("Unexpected claim variant {:?}", v),
 	};
 	let fixed_epoch = epoch.clone_for_slot(slot);
-	let transcript = make_transcript_data(&epoch.randomness.clone(), slot, fixed_epoch.epoch_index);
+	let transcript = make_transcript(&epoch.randomness.clone(), slot, fixed_epoch.epoch_index);
 	let sign = keystore
-		.sr25519_vrf_sign(AuthorityId::ID, &public, transcript)
+		.sr25519_vrf_sign(AuthorityId::ID, &public, &transcript)
 		.unwrap()
 		.unwrap();
 	assert_eq!(fixed_epoch.epoch_index, 11);
-	assert_eq!(pre_digest.vrf_output, VRFOutput(sign.output));
+	assert_eq!(pre_digest.vrf_signature.output, sign.output);
 }
 
 // Propose and import a new BABE block on top of the given parent.
@@ -1084,36 +1077,6 @@ async fn verify_slots_are_strictly_increasing() {
 	propose_and_import_block(&b1, Some(999.into()), &mut proposer_factory, &mut block_import).await;
 }
 
-#[test]
-fn babe_transcript_generation_match() {
-	sp_tracing::try_init_simple();
-
-	let authority = Sr25519Keyring::Alice;
-	let _keystore = create_keystore(authority);
-
-	let epoch = Epoch {
-		start_slot: 0.into(),
-		authorities: vec![(authority.public().into(), 1)],
-		randomness: [0; 32],
-		epoch_index: 1,
-		duration: 100,
-		config: BabeEpochConfiguration {
-			c: (3, 10),
-			allowed_slots: AllowedSlots::PrimaryAndSecondaryPlainSlots,
-		},
-	};
-
-	let orig_transcript = make_transcript(&epoch.randomness.clone(), 1.into(), epoch.epoch_index);
-	let new_transcript = make_transcript_data(&epoch.randomness, 1.into(), epoch.epoch_index);
-
-	let test = |t: merlin::Transcript| -> [u8; 16] {
-		let mut b = [0u8; 16];
-		t.build_rng().finalize(&mut ChaChaRng::from_seed([0u8; 32])).fill_bytes(&mut b);
-		b
-	};
-	debug_assert!(test(orig_transcript) == test(transcript_from_data(new_transcript)));
-}
-
 #[tokio::test]
 async fn obsolete_blocks_aux_data_cleanup() {
 	let mut net = BabeTestNet::new(1);

substrate/client/consensus/babe/src/verification.rs

@@ -18,8 +18,9 @@
 
 //! Verification for BABE headers.
 use crate::{
-	authorship::{calculate_primary_threshold, check_primary_threshold, secondary_slot_author},
-	babe_err, find_pre_digest, BlockT, Epoch, Error, LOG_TARGET,
+	authorship::{calculate_primary_threshold, secondary_slot_author},
+	babe_err, find_pre_digest, BlockT, Epoch, Error, AUTHORING_SCORE_LENGTH,
+	AUTHORING_SCORE_VRF_CONTEXT, LOG_TARGET,
 };
 use log::{debug, trace};
 use sc_consensus_epochs::Epoch as EpochT;
@@ -32,7 +33,10 @@ use sp_consensus_babe::{
 	make_transcript, AuthorityId, AuthorityPair, AuthoritySignature,
 };
 use sp_consensus_slots::Slot;
-use sp_core::{ByteArray, Pair};
+use sp_core::{
+	crypto::{VrfVerifier, Wraps},
+	Pair,
+};
 use sp_runtime::{traits::Header, DigestItem};
 
 /// BABE verification parameters
@@ -155,7 +159,7 @@ fn check_primary_header<B: BlockT + Sized>(
 	epoch: &Epoch,
 	c: (u64, u64),
 ) -> Result<(), Error<B>> {
-	let author = &epoch.authorities[pre_digest.authority_index as usize].0;
+	let authority_id = &epoch.authorities[pre_digest.authority_index as usize].0;
 	let mut epoch_index = epoch.epoch_index;
 
 	if epoch.end_slot() <= pre_digest.slot {
@@ -163,28 +167,34 @@ fn check_primary_header<B: BlockT + Sized>(
 		epoch_index = epoch.clone_for_slot(pre_digest.slot).epoch_index;
 	}
 
-	if AuthorityPair::verify(&signature, pre_hash, author) {
-		let (inout, _) = {
-			let transcript = make_transcript(&epoch.randomness, pre_digest.slot, epoch_index);
-
-			schnorrkel::PublicKey::from_bytes(author.as_slice())
-				.and_then(|p| {
-					p.vrf_verify(transcript, &pre_digest.vrf_output, &pre_digest.vrf_proof)
-				})
-				.map_err(|s| babe_err(Error::VRFVerificationFailed(s)))?
-		};
-
-		let threshold =
-			calculate_primary_threshold(c, &epoch.authorities, pre_digest.authority_index as usize);
-
-		if !check_primary_threshold(&inout, threshold) {
-			return Err(babe_err(Error::VRFVerificationOfBlockFailed(author.clone(), threshold)))
-		}
-
-		Ok(())
-	} else {
-		Err(babe_err(Error::BadSignature(pre_hash)))
+	if !AuthorityPair::verify(&signature, pre_hash, authority_id) {
+		return Err(babe_err(Error::BadSignature(pre_hash)))
 	}
+
+	let transcript = make_transcript(&epoch.randomness, pre_digest.slot, epoch_index);
+
+	if !authority_id.as_inner_ref().vrf_verify(&transcript, &pre_digest.vrf_signature) {
+		return Err(babe_err(Error::VrfVerificationFailed))
+	}
+
+	let threshold =
+		calculate_primary_threshold(c, &epoch.authorities, pre_digest.authority_index as usize);
+
+	let score = authority_id
+		.as_inner_ref()
+		.make_bytes::<[u8; AUTHORING_SCORE_LENGTH]>(
+			AUTHORING_SCORE_VRF_CONTEXT,
+			&transcript,
+			&pre_digest.vrf_signature.output,
+		)
+		.map(u128::from_le_bytes)
+		.map_err(|_| babe_err(Error::VrfVerificationFailed))?;
+
+	if score >= threshold {
+		return Err(babe_err(Error::VrfThresholdExceeded(threshold)))
+	}
+
+	Ok(())
 }
 
 /// Check a secondary slot proposal header. We validate that the given header is
@@ -197,8 +207,7 @@ fn check_secondary_plain_header<B: BlockT>(
 	signature: AuthoritySignature,
 	epoch: &Epoch,
 ) -> Result<(), Error<B>> {
-	// check the signature is valid under the expected authority and
-	// chain state.
+	// check the signature is valid under the expected authority and chain state.
 	let expected_author =
 		secondary_slot_author(pre_digest.slot, &epoch.authorities, epoch.randomness)
 			.ok_or(Error::NoSecondaryAuthorExpected)?;
@@ -209,11 +218,11 @@ fn check_secondary_plain_header<B: BlockT>(
 		return Err(Error::InvalidAuthor(expected_author.clone(), author.clone()))
 	}
 
-	if AuthorityPair::verify(&signature, pre_hash.as_ref(), author) {
-		Ok(())
-	} else {
-		Err(Error::BadSignature(pre_hash))
+	if !AuthorityPair::verify(&signature, pre_hash.as_ref(), author) {
+		return Err(Error::BadSignature(pre_hash))
 	}
+
+	Ok(())
 }
 
 /// Check a secondary VRF slot proposal header.
@@ -223,8 +232,7 @@ fn check_secondary_vrf_header<B: BlockT>(
 	signature: AuthoritySignature,
 	epoch: &Epoch,
 ) -> Result<(), Error<B>> {
-	// check the signature is valid under the expected authority and
-	// chain state.
+	// check the signature is valid under the expected authority and chain state.
 	let expected_author =
 		secondary_slot_author(pre_digest.slot, &epoch.authorities, epoch.randomness)
 			.ok_or(Error::NoSecondaryAuthorExpected)?;
@@ -241,15 +249,15 @@ fn check_secondary_vrf_header<B: BlockT>(
 		epoch_index = epoch.clone_for_slot(pre_digest.slot).epoch_index;
 	}
 
-	if AuthorityPair::verify(&signature, pre_hash.as_ref(), author) {
-		let transcript = make_transcript(&epoch.randomness, pre_digest.slot, epoch_index);
-
-		schnorrkel::PublicKey::from_bytes(author.as_slice())
-			.and_then(|p| p.vrf_verify(transcript, &pre_digest.vrf_output, &pre_digest.vrf_proof))
-			.map_err(|s| babe_err(Error::VRFVerificationFailed(s)))?;
-
-		Ok(())
-	} else {
-		Err(Error::BadSignature(pre_hash))
+	if !AuthorityPair::verify(&signature, pre_hash.as_ref(), author) {
+		return Err(Error::BadSignature(pre_hash))
 	}
+
+	let transcript = make_transcript(&epoch.randomness, pre_digest.slot, epoch_index);
+
+	if !author.as_inner_ref().vrf_verify(&transcript, &pre_digest.vrf_signature) {
+		return Err(Error::VrfVerificationFailed)
+	}
+
+	Ok(())
 }