VRF refactory (paritytech#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>
nathanwhit · Jul 19, 2023 · 95f4229 · 95f4229
1 parent 5ff3de0
commit 95f4229
Show file tree

Hide file tree

Showing 28 changed files with 468 additions and 712 deletions.
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/Cargo.toml b/Cargo.toml
@@ -190,7 +190,6 @@ members = [
 	"primitives/consensus/grandpa",
 	"primitives/consensus/pow",
 	"primitives/consensus/slots",
-	"primitives/consensus/vrf",
 	"primitives/core",
 	"primitives/core/hashing",
 	"primitives/core/hashing/proc-macro",

diff --git a/client/consensus/babe/Cargo.toml b/client/consensus/babe/Cargo.toml
@@ -19,12 +19,10 @@ scale-info = { version = "2.5.0", features = ["derive"] }
 codec = { package = "parity-scale-codec", version = "3.2.2", features = ["derive"] }
 futures = "0.3.21"
 log = "0.4.17"
-merlin = "2.0"
 num-bigint = "0.4.3"
 num-rational = "0.4.1"
 num-traits = "0.2.8"
 parking_lot = "0.12.1"
-schnorrkel = { version = "0.9.1", features = ["preaudit_deprecated"] }
 thiserror = "1.0"
 fork-tree = { version = "3.0.0", path = "../../../utils/fork-tree" }
 prometheus-endpoint = { package = "substrate-prometheus-endpoint", version = "0.10.0-dev", path = "../../../utils/prometheus" }
@@ -41,7 +39,6 @@ sp-blockchain = { version = "4.0.0-dev", path = "../../../primitives/blockchain"
 sp-consensus = { version = "0.10.0-dev", path = "../../../primitives/consensus/common" }
 sp-consensus-babe = { version = "0.10.0-dev", path = "../../../primitives/consensus/babe" }
 sp-consensus-slots = { version = "0.10.0-dev", path = "../../../primitives/consensus/slots" }
-sp-consensus-vrf = { version = "0.10.0-dev", path = "../../../primitives/consensus/vrf" }
 sp-core = { version = "7.0.0", path = "../../../primitives/core" }
 sp-inherents = { version = "4.0.0-dev", path = "../../../primitives/inherents" }
 sp-keystore = { version = "0.13.0", path = "../../../primitives/keystore" }

diff --git a/client/consensus/babe/src/authorship.rs b/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()))

diff --git a/client/consensus/babe/src/lib.rs b/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),

diff --git a/client/consensus/babe/src/migration.rs b/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 {