babe: make secondary slot randomness available on-chain (#7053)

* babe: make secondary slot randomness available on-chain

* babe: extract out vrf_output function

* babe: add missing comment

* babe: fix incorrectly storing primary randomness

* babe: add test for onchain author vrf

* babe: fix reviewer nits

* runtime: bump spec_version

* babe: remove outer Option for AuthorVrfRandomness

* babe: fix reviewer nits on doc strings

* babe: move make_vrf_output to mock.rs

* babe: cleanup docs

* babe: kill ephemeral entry instead of take

* babe: use type alias for maybe randomness

Co-authored-by: André Silva <andrerfosilva@gmail.com>

substrate/frame/babe/src/lib.rs

@@ -210,6 +210,11 @@ decl_storage! {
 		/// if per-block initialization has already been called for current block.
 		Initialized get(fn initialized): Option<MaybeRandomness>;
 
+		/// Temporary value (cleared at block finalization) that includes the VRF output generated
+		/// at this block. This field should always be populated during block processing unless
+		/// secondary plain slots are enabled (which don't contain a VRF output).
+		AuthorVrfRandomness get(fn author_vrf_randomness): MaybeRandomness;
+
 		/// How late the current block is compared to its parent.
 		///
 		/// This entry is populated as part of block execution and is cleaned up
@@ -255,6 +260,9 @@ decl_module! {
 				Self::deposit_randomness(&randomness);
 			}
 
+			// The stored author generated VRF output is ephemeral.
+			AuthorVrfRandomness::kill();
+
 			// remove temporary "environment" entry from storage
 			Lateness::<T>::kill();
 		}
@@ -517,7 +525,9 @@ impl<T: Trait> Module<T> {
 			})
 			.next();
 
-		let maybe_randomness: Option<schnorrkel::Randomness> = maybe_pre_digest.and_then(|digest| {
+		let is_primary = matches!(maybe_pre_digest, Some(PreDigest::Primary(..)));
+
+		let maybe_randomness: MaybeRandomness = maybe_pre_digest.and_then(|digest| {
 			// on the first non-zero block (i.e. block #1)
 			// this is where the first epoch (epoch #0) actually starts.
 			// we need to adjust internal storage accordingly.
@@ -546,38 +556,44 @@ impl<T: Trait> Module<T> {
 			Lateness::<T>::put(lateness);
 			CurrentSlot::put(current_slot);
 
-			if let PreDigest::Primary(primary) = digest {
-				// place the VRF output into the `Initialized` storage item
-				// and it'll be put onto the under-construction randomness
-				// later, once we've decided which epoch this block is in.
-				//
-				// Reconstruct the bytes of VRFInOut using the authority id.
-				Authorities::get()
-					.get(primary.authority_index as usize)
-					.and_then(|author| {
-						schnorrkel::PublicKey::from_bytes(author.0.as_slice()).ok()
-					})
-					.and_then(|pubkey| {
-						let transcript = sp_consensus_babe::make_transcript(
-							&Self::randomness(),
-							current_slot,
-							EpochIndex::get(),
-						);
+			let authority_index = digest.authority_index();
 
-						primary.vrf_output.0.attach_input_hash(
-							&pubkey,
-							transcript
-						).ok()
-					})
-					.map(|inout| {
-						inout.make_bytes(&sp_consensus_babe::BABE_VRF_INOUT_CONTEXT)
-					})
-			} else {
-				None
-			}
+			// Extract out the VRF output if we have it
+			digest
+				.vrf_output()
+				.and_then(|vrf_output| {
+					// Reconstruct the bytes of VRFInOut using the authority id.
+					Authorities::get()
+						.get(authority_index as usize)
+						.and_then(|author| {
+							schnorrkel::PublicKey::from_bytes(author.0.as_slice()).ok()
+						})
+						.and_then(|pubkey| {
+							let transcript = sp_consensus_babe::make_transcript(
+								&Self::randomness(),
+								current_slot,
+								EpochIndex::get(),
+							);
+
+							vrf_output.0.attach_input_hash(
+								&pubkey,
+								transcript
+							).ok()
+						})
+						.map(|inout| {
+							inout.make_bytes(&sp_consensus_babe::BABE_VRF_INOUT_CONTEXT)
+						})
+				})
 		});
 
-		Initialized::put(maybe_randomness);
+		// For primary VRF output we place it in the `Initialized` storage
+		// item and it'll be put onto the under-construction randomness later,
+		// once we've decided which epoch this block is in.
+		Initialized::put(if is_primary { maybe_randomness } else { None });
+
+		// Place either the primary or secondary VRF output into the
+		// `AuthorVrfRandomness` storage item.
+		AuthorVrfRandomness::put(maybe_randomness);
 
 		// enact epoch change, if necessary.
 		T::EpochChangeTrigger::trigger::<T>(now)