Make BEEFY payload extensible (#10307)

* make BEEFY payload extensible * cargo fmt * cargo fmt * remove generic payload param in beefy-primitives * cargo fmt * Apply suggestions from code review Co-authored-by: Tomasz Drwięga <tomusdrw@users.noreply.github.com> * adds Paylaod Type * remove hex * fix tests * Apply suggestions from code review Co-authored-by: Tomasz Drwięga <tomusdrw@users.noreply.github.com> * use binary_search_by to sort * Payload::new() * fix tests * Apply suggestions from code review Co-authored-by: Tomasz Drwięga <tomusdrw@users.noreply.github.com> * fix tests * cargo fmt * fix get_decoded * fix test Co-authored-by: Tomasz Drwięga <tomusdrw@users.noreply.github.com>
paritytech · Dec 1, 2021 · 8bf9836 · 8bf9836
1 parent b6a3ecc
commit 8bf9836
Show file tree

Hide file tree

Showing 9 changed files with 152 additions and 99 deletions.
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/client/beefy/src/gossip.rs b/client/beefy/src/gossip.rs
@@ -30,7 +30,7 @@ use wasm_timer::Instant;
 
 use beefy_primitives::{
 	crypto::{Public, Signature},
-	MmrRootHash, VoteMessage,
+	VoteMessage,
 };
 
 use crate::keystore::BeefyKeystore;
@@ -142,9 +142,7 @@ where
 		sender: &PeerId,
 		mut data: &[u8],
 	) -> ValidationResult<B::Hash> {
-		if let Ok(msg) =
-			VoteMessage::<MmrRootHash, NumberFor<B>, Public, Signature>::decode(&mut data)
-		{
+		if let Ok(msg) = VoteMessage::<NumberFor<B>, Public, Signature>::decode(&mut data) {
 			let msg_hash = twox_64(data);
 			let round = msg.commitment.block_number;
 
@@ -178,9 +176,7 @@ where
 	fn message_expired<'a>(&'a self) -> Box<dyn FnMut(B::Hash, &[u8]) -> bool + 'a> {
 		let known_votes = self.known_votes.read();
 		Box::new(move |_topic, mut data| {
-			let msg = match VoteMessage::<MmrRootHash, NumberFor<B>, Public, Signature>::decode(
-				&mut data,
-			) {
+			let msg = match VoteMessage::<NumberFor<B>, Public, Signature>::decode(&mut data) {
 				Ok(vote) => vote,
 				Err(_) => return true,
 			};
@@ -214,9 +210,7 @@ where
 				return do_rebroadcast
 			}
 
-			let msg = match VoteMessage::<MmrRootHash, NumberFor<B>, Public, Signature>::decode(
-				&mut data,
-			) {
+			let msg = match VoteMessage::<NumberFor<B>, Public, Signature>::decode(&mut data) {
 				Ok(vote) => vote,
 				Err(_) => return true,
 			};
@@ -237,9 +231,11 @@ mod tests {
 	use sc_network_test::Block;
 	use sp_keystore::{SyncCryptoStore, SyncCryptoStorePtr};
 
-	use beefy_primitives::{crypto::Signature, Commitment, MmrRootHash, VoteMessage, KEY_TYPE};
-
 	use crate::keystore::{tests::Keyring, BeefyKeystore};
+	use beefy_primitives::{
+		crypto::Signature, known_payload_ids, Commitment, MmrRootHash, Payload, VoteMessage,
+		KEY_TYPE,
+	};
 
 	use super::*;
 
@@ -345,10 +341,7 @@ mod tests {
 		}
 	}
 
-	fn sign_commitment<BN: Encode, P: Encode>(
-		who: &Keyring,
-		commitment: &Commitment<BN, P>,
-	) -> Signature {
+	fn sign_commitment<BN: Encode>(who: &Keyring, commitment: &Commitment<BN>) -> Signature {
 		let store: SyncCryptoStorePtr = std::sync::Arc::new(LocalKeystore::in_memory());
 		SyncCryptoStore::ecdsa_generate_new(&*store, KEY_TYPE, Some(&who.to_seed())).unwrap();
 		let beefy_keystore: BeefyKeystore = Some(store).into();
@@ -362,11 +355,8 @@ mod tests {
 		let sender = sc_network::PeerId::random();
 		let mut context = TestContext;
 
-		let commitment = Commitment {
-			payload: MmrRootHash::default(),
-			block_number: 3_u64,
-			validator_set_id: 0,
-		};
+		let payload = Payload::new(known_payload_ids::MMR_ROOT_ID, MmrRootHash::default().encode());
+		let commitment = Commitment { payload, block_number: 3_u64, validator_set_id: 0 };
 
 		let signature = sign_commitment(&Keyring::Alice, &commitment);
 

diff --git a/client/beefy/src/notification.rs b/client/beefy/src/notification.rs
@@ -24,8 +24,7 @@ use sp_runtime::traits::{Block, NumberFor};
 use parking_lot::Mutex;
 
 /// Stream of signed commitments returned when subscribing.
-pub type SignedCommitment<Block> =
-	beefy_primitives::SignedCommitment<NumberFor<Block>, beefy_primitives::MmrRootHash>;
+pub type SignedCommitment<Block> = beefy_primitives::SignedCommitment<NumberFor<Block>>;
 
 /// Stream of signed commitments returned when subscribing.
 type SignedCommitmentStream<Block> = TracingUnboundedReceiver<SignedCommitment<Block>>;

diff --git a/client/beefy/src/round.rs b/client/beefy/src/round.rs
@@ -53,14 +53,14 @@ fn threshold(authorities: usize) -> usize {
 	authorities - faulty
 }
 
-pub(crate) struct Rounds<Hash, Number> {
-	rounds: BTreeMap<(Hash, Number), RoundTracker>,
+pub(crate) struct Rounds<Payload, Number> {
+	rounds: BTreeMap<(Payload, Number), RoundTracker>,
 	validator_set: ValidatorSet<Public>,
 }
 
-impl<H, N> Rounds<H, N>
+impl<P, N> Rounds<P, N>
 where
-	H: Ord + Hash,
+	P: Ord + Hash,
 	N: Ord + AtLeast32BitUnsigned + MaybeDisplay,
 {
 	pub(crate) fn new(validator_set: ValidatorSet<Public>) -> Self {
@@ -70,8 +70,8 @@ where
 
 impl<H, N> Rounds<H, N>
 where
-	H: Ord + Hash,
-	N: Ord + AtLeast32BitUnsigned + MaybeDisplay,
+	H: Ord + Hash + Clone,
+	N: Ord + AtLeast32BitUnsigned + MaybeDisplay + Clone,
 {
 	pub(crate) fn validator_set_id(&self) -> ValidatorSetId {
 		self.validator_set.id
@@ -81,9 +81,9 @@ where
 		self.validator_set.validators.clone()
 	}
 
-	pub(crate) fn add_vote(&mut self, round: (H, N), vote: (Public, Signature)) -> bool {
+	pub(crate) fn add_vote(&mut self, round: &(H, N), vote: (Public, Signature)) -> bool {
 		if self.validator_set.validators.iter().any(|id| vote.0 == *id) {
-			self.rounds.entry(round).or_default().add_vote(vote)
+			self.rounds.entry(round.clone()).or_default().add_vote(vote)
 		} else {
 			false
 		}
@@ -179,29 +179,29 @@ mod tests {
 		let mut rounds = Rounds::<H256, NumberFor<Block>>::new(validators);
 
 		assert!(rounds.add_vote(
-			(H256::from_low_u64_le(1), 1),
+			&(H256::from_low_u64_le(1), 1),
 			(Keyring::Alice.public(), Keyring::Alice.sign(b"I am committed"))
 		));
 
 		assert!(!rounds.is_done(&(H256::from_low_u64_le(1), 1)));
 
 		// invalid vote
 		assert!(!rounds.add_vote(
-			(H256::from_low_u64_le(1), 1),
+			&(H256::from_low_u64_le(1), 1),
 			(Keyring::Dave.public(), Keyring::Dave.sign(b"I am committed"))
 		));
 
 		assert!(!rounds.is_done(&(H256::from_low_u64_le(1), 1)));
 
 		assert!(rounds.add_vote(
-			(H256::from_low_u64_le(1), 1),
+			&(H256::from_low_u64_le(1), 1),
 			(Keyring::Bob.public(), Keyring::Bob.sign(b"I am committed"))
 		));
 
 		assert!(!rounds.is_done(&(H256::from_low_u64_le(1), 1)));
 
 		assert!(rounds.add_vote(
-			(H256::from_low_u64_le(1), 1),
+			&(H256::from_low_u64_le(1), 1),
 			(Keyring::Charlie.public(), Keyring::Charlie.sign(b"I am committed"))
 		));
 
@@ -225,31 +225,31 @@ mod tests {
 
 		// round 1
 		rounds.add_vote(
-			(H256::from_low_u64_le(1), 1),
+			&(H256::from_low_u64_le(1), 1),
 			(Keyring::Alice.public(), Keyring::Alice.sign(b"I am committed")),
 		);
 		rounds.add_vote(
-			(H256::from_low_u64_le(1), 1),
+			&(H256::from_low_u64_le(1), 1),
 			(Keyring::Bob.public(), Keyring::Bob.sign(b"I am committed")),
 		);
 
 		// round 2
 		rounds.add_vote(
-			(H256::from_low_u64_le(2), 2),
+			&(H256::from_low_u64_le(2), 2),
 			(Keyring::Alice.public(), Keyring::Alice.sign(b"I am again committed")),
 		);
 		rounds.add_vote(
-			(H256::from_low_u64_le(2), 2),
+			&(H256::from_low_u64_le(2), 2),
 			(Keyring::Bob.public(), Keyring::Bob.sign(b"I am again committed")),
 		);
 
 		// round 3
 		rounds.add_vote(
-			(H256::from_low_u64_le(3), 3),
+			&(H256::from_low_u64_le(3), 3),
 			(Keyring::Alice.public(), Keyring::Alice.sign(b"I am still committed")),
 		);
 		rounds.add_vote(
-			(H256::from_low_u64_le(3), 3),
+			&(H256::from_low_u64_le(3), 3),
 			(Keyring::Bob.public(), Keyring::Bob.sign(b"I am still committed")),
 		);
 

diff --git a/client/beefy/src/worker.rs b/client/beefy/src/worker.rs
@@ -36,8 +36,8 @@ use sp_runtime::{
 
 use beefy_primitives::{
 	crypto::{AuthorityId, Public, Signature},
-	BeefyApi, Commitment, ConsensusLog, MmrRootHash, SignedCommitment, ValidatorSet,
-	VersionedCommitment, VoteMessage, BEEFY_ENGINE_ID, GENESIS_AUTHORITY_SET_ID,
+	known_payload_ids, BeefyApi, Commitment, ConsensusLog, MmrRootHash, Payload, SignedCommitment,
+	ValidatorSet, VersionedCommitment, VoteMessage, BEEFY_ENGINE_ID, GENESIS_AUTHORITY_SET_ID,
 };
 
 use crate::{
@@ -79,7 +79,7 @@ where
 	/// Min delta in block numbers between two blocks, BEEFY should vote on
 	min_block_delta: u32,
 	metrics: Option<Metrics>,
-	rounds: round::Rounds<MmrRootHash, NumberFor<B>>,
+	rounds: round::Rounds<Payload, NumberFor<B>>,
 	finality_notifications: FinalityNotifications<B>,
 	/// Best block we received a GRANDPA notification for
 	best_grandpa_block: NumberFor<B>,
@@ -262,8 +262,9 @@ where
 					return
 				};
 
+			let payload = Payload::new(known_payload_ids::MMR_ROOT_ID, mmr_root.encode());
 			let commitment = Commitment {
-				payload: mmr_root,
+				payload,
 				block_number: notification.header.number(),
 				validator_set_id: self.rounds.validator_set_id(),
 			};
@@ -301,10 +302,10 @@ where
 		}
 	}
 
-	fn handle_vote(&mut self, round: (MmrRootHash, NumberFor<B>), vote: (Public, Signature)) {
+	fn handle_vote(&mut self, round: (Payload, NumberFor<B>), vote: (Public, Signature)) {
 		self.gossip_validator.note_round(round.1);
 
-		let vote_added = self.rounds.add_vote(round, vote);
+		let vote_added = self.rounds.add_vote(&round, vote);
 
 		if vote_added && self.rounds.is_done(&round) {
 			if let Some(signatures) = self.rounds.drop(&round) {
@@ -352,7 +353,7 @@ where
 			|notification| async move {
 				debug!(target: "beefy", "🥩 Got vote message: {:?}", notification);
 
-				VoteMessage::<MmrRootHash, NumberFor<B>, Public, Signature>::decode(
+				VoteMessage::<NumberFor<B>, Public, Signature>::decode(
 					&mut &notification.message[..],
 				)
 				.ok()

diff --git a/primitives/beefy/Cargo.toml b/primitives/beefy/Cargo.toml
@@ -18,8 +18,8 @@ sp-runtime = { version = "4.0.0-dev", path = "../runtime", default-features = fa
 sp-std = { version = "4.0.0-dev", path = "../std", default-features = false }
 
 [dev-dependencies]
+hex = "0.4.3"
 hex-literal = "0.3"
-
 sp-keystore = { version = "0.10.0-dev", path = "../keystore" }
 
 [features]