From 3de2a92598474828b0c170856d76cfe7b8f7b45d Mon Sep 17 00:00:00 2001
From: Serban Iorga <serban@parity.io>
Date: Tue, 1 Oct 2024 19:38:36 +0300
Subject: [PATCH] Beefy equivocation: check all the MMR roots (#5857)

Normally, the BEEFY protocol only accepts a single MMR Root entry in a
commitment's payload. But to be extra careful, when validating
equivocation reports, let's check all the MMR roots, if there are more.

---------

Co-authored-by: Adrian Catangiu <adrian@parity.io>
---
 prdoc/pr_5857.prdoc                           | 19 +++++++++++
 substrate/frame/beefy-mmr/src/lib.rs          | 32 ++++++++++++++-----
 substrate/frame/beefy-mmr/src/tests.rs        | 24 ++++++++++++--
 .../primitives/consensus/beefy/src/payload.rs | 18 +++++++++++
 4 files changed, 83 insertions(+), 10 deletions(-)
 create mode 100644 prdoc/pr_5857.prdoc

diff --git a/prdoc/pr_5857.prdoc b/prdoc/pr_5857.prdoc
new file mode 100644
index 000000000000..00ee0a8cc704
--- /dev/null
+++ b/prdoc/pr_5857.prdoc
@@ -0,0 +1,19 @@
+# Schema: Polkadot SDK PRDoc Schema (prdoc) v1.0.0
+# See doc at https://raw.githubusercontent.com/paritytech/polkadot-sdk/master/prdoc/schema_user.json
+
+title: "Beefy equivocation: check all the MMR roots"
+
+doc:
+  - audience:
+    - Runtime Dev
+    - Runtime User
+    description: |
+      This PR adjusts the logic for `report_fork_voting` exposed by `pallet-beefy`.
+      Normally, the BEEFY protocol only accepts a single MMR Root entry in a commitment's payload. But, in order to
+      be extra careful, now, when validating equivocation reports, we check all the MMR roots, if there are more.
+
+crates:
+  - name: sp-consensus-beefy
+    bump: patch
+  - name: pallet-beefy-mmr
+    bump: patch
diff --git a/substrate/frame/beefy-mmr/src/lib.rs b/substrate/frame/beefy-mmr/src/lib.rs
index 73119c3faa9b..ef99bc1e9cf1 100644
--- a/substrate/frame/beefy-mmr/src/lib.rs
+++ b/substrate/frame/beefy-mmr/src/lib.rs
@@ -258,17 +258,33 @@ where
 				},
 			};
 
-		let commitment_root =
-			match commitment.payload.get_decoded::<MerkleRootOf<T>>(&known_payloads::MMR_ROOT_ID) {
-				Some(commitment_root) => commitment_root,
+		let mut found_commitment_root = false;
+		let commitment_roots = commitment
+			.payload
+			.get_all_decoded::<MerkleRootOf<T>>(&known_payloads::MMR_ROOT_ID);
+		for maybe_commitment_root in commitment_roots {
+			match maybe_commitment_root {
+				Some(commitment_root) => {
+					found_commitment_root = true;
+					if canonical_prev_root != commitment_root {
+						// If the commitment contains an MMR root, that is not equal to
+						// `canonical_prev_root`, the commitment is invalid
+						return true;
+					}
+				},
 				None => {
-					// If the commitment doesn't contain any MMR root, while the proof is valid,
-					// the commitment is invalid
-					return true
+					// If the commitment contains an MMR root, that can't be decoded, it is invalid.
+					return true;
 				},
-			};
+			}
+		}
+		if !found_commitment_root {
+			// If the commitment doesn't contain any MMR root, while the proof is valid,
+			// the commitment is invalid
+			return true;
+		}
 
-		canonical_prev_root != commitment_root
+		false
 	}
 }
 
diff --git a/substrate/frame/beefy-mmr/src/tests.rs b/substrate/frame/beefy-mmr/src/tests.rs
index b126a01012b4..297fb28647ac 100644
--- a/substrate/frame/beefy-mmr/src/tests.rs
+++ b/substrate/frame/beefy-mmr/src/tests.rs
@@ -278,8 +278,28 @@ fn is_non_canonical_should_work_correctly() {
 				&Commitment {
 					payload: Payload::from_single_entry(
 						known_payloads::MMR_ROOT_ID,
-						H256::repeat_byte(0).encode(),
-					),
+						prev_roots[250 - 1].encode()
+					)
+					.push_raw(known_payloads::MMR_ROOT_ID, H256::repeat_byte(0).encode(),),
+					block_number: 250,
+					validator_set_id: 0,
+				},
+				valid_proof.clone(),
+				Mmr::mmr_root(),
+			),
+			true
+		);
+
+		// If the `commitment.payload` contains an MMR root that can't be decoded,
+		// it's non-canonical.
+		assert_eq!(
+			BeefyMmr::is_non_canonical(
+				&Commitment {
+					payload: Payload::from_single_entry(
+						known_payloads::MMR_ROOT_ID,
+						prev_roots[250 - 1].encode()
+					)
+					.push_raw(known_payloads::MMR_ROOT_ID, vec![],),
 					block_number: 250,
 					validator_set_id: 0,
 				},
diff --git a/substrate/primitives/consensus/beefy/src/payload.rs b/substrate/primitives/consensus/beefy/src/payload.rs
index d22255c384bc..9e792670fef5 100644
--- a/substrate/primitives/consensus/beefy/src/payload.rs
+++ b/substrate/primitives/consensus/beefy/src/payload.rs
@@ -56,6 +56,16 @@ impl Payload {
 		Some(&self.0[index].1)
 	}
 
+	/// Returns all the raw payloads under given `id`.
+	pub fn get_all_raw<'a>(
+		&'a self,
+		id: &'a BeefyPayloadId,
+	) -> impl Iterator<Item = &Vec<u8>> + 'a {
+		self.0
+			.iter()
+			.filter_map(move |probe| if &probe.0 != id { return None } else { Some(&probe.1) })
+	}
+
 	/// Returns a decoded payload value under given `id`.
 	///
 	/// In case the value is not there, or it cannot be decoded `None` is returned.
@@ -63,6 +73,14 @@ impl Payload {
 		self.get_raw(id).and_then(|raw| T::decode(&mut &raw[..]).ok())
 	}
 
+	/// Returns all decoded payload values under given `id`.
+	pub fn get_all_decoded<'a, T: Decode>(
+		&'a self,
+		id: &'a BeefyPayloadId,
+	) -> impl Iterator<Item = Option<T>> + 'a {
+		self.get_all_raw(id).map(|raw| T::decode(&mut &raw[..]).ok())
+	}
+
 	/// Push a `Vec<u8>` with a given id into the payload vec.
 	/// This method will internally sort the payload vec after every push.
 	///