From eebc7cae7b367153e38d6fe5f75b2e21057e3f88 Mon Sep 17 00:00:00 2001
From: Alistair Singh <alistair.singh7@gmail.com>
Date: Mon, 14 Oct 2024 10:08:52 +0200
Subject: [PATCH] Fix mmr generation (#1307)

* fix mmr generation

* updated test command

* fixed test command

* remove redundant logging
---
 relayer/cmd/parachain_head_proof.go        | 40 ++++++++++++++++------
 relayer/relays/parachain/beefy-listener.go |  5 +--
 2 files changed, 32 insertions(+), 13 deletions(-)

diff --git a/relayer/cmd/parachain_head_proof.go b/relayer/cmd/parachain_head_proof.go
index 1052d30837..70bde4def6 100644
--- a/relayer/cmd/parachain_head_proof.go
+++ b/relayer/cmd/parachain_head_proof.go
@@ -43,12 +43,6 @@ func parachainHeadProofCmd() *cobra.Command {
 	)
 	cmd.MarkFlagRequired("parachain-id")
 
-	cmd.Flags().Uint64(
-		"parachain-block",
-		0,
-		"The parachain block you are trying to prove. i.e. The block containing the message.",
-	)
-	cmd.MarkFlagRequired("parachain-block")
 	return cmd
 }
 
@@ -72,7 +66,8 @@ func ParachainHeadProofFn(cmd *cobra.Command, _ []string) error {
 	copy(beefyBlockHash[:], beefyBlockHashHex[0:32])
 
 	relayChainBlock, _ := cmd.Flags().GetUint64("relaychain-block")
-	mmrProof, err := conn.GenerateProofForBlock(relayChainBlock, beefyBlockHash)
+	// magic plus 1 to make the block a leaf index
+	mmrProof, err := conn.GenerateProofForBlock(relayChainBlock+1, beefyBlockHash)
 	if err != nil {
 		log.WithError(err).Error("Cannot connect.")
 		return err
@@ -84,7 +79,6 @@ func ParachainHeadProofFn(cmd *cobra.Command, _ []string) error {
 	}).Info("conn.GenerateProofForBlock")
 
 	paraID, _ := cmd.Flags().GetUint32("parachain-id")
-	parachainBlock, _ := cmd.Flags().GetUint64("parachain-block")
 
 	relayChainBlockHash, err := conn.API().RPC.Chain.GetBlockHash(relayChainBlock)
 	if err != nil {
@@ -116,12 +110,37 @@ func ParachainHeadProofFn(cmd *cobra.Command, _ []string) error {
 		"paraId":              paraID,
 		"relayChainBlockHash": relayChainBlockHash.Hex(),
 	}).Info("parachain.CreateParachainMerkleProof")
-
-	merkleProofData, err := parachain.CreateParachainMerkleProof(paraHeadsAsSlice, paraID)
+	numParas := min(parachain.MaxParaHeads, len(paraHeadsAsSlice))
+	merkleProofData, err := parachain.CreateParachainMerkleProof(paraHeadsAsSlice[:numParas], paraID)
 	if err != nil {
 		log.WithError(err).Error("Cannot create merkle proof.")
 		return err
 	}
+
+	if merkleProofData.Root.Hex() != mmrProof.Leaf.ParachainHeads.Hex() {
+		log.WithFields(log.Fields{
+			"computedMmr": merkleProofData.Root.Hex(),
+			"mmr":         mmrProof.Leaf.ParachainHeads.Hex(),
+		}).Warn("MMR parachain merkle root does not match calculated merkle root. Filtering out parachain heads.")
+
+		paraHeadsAsSlice, err = conn.FilterParachainHeads(paraHeadsAsSlice, relayChainBlockHash)
+		if err != nil {
+			log.WithError(err).Fatal("Filtering out parachain heads failed.")
+		}
+
+		numParas = min(parachain.MaxParaHeads, len(paraHeadsAsSlice))
+		merkleProofData, err = parachain.CreateParachainMerkleProof(paraHeadsAsSlice[:numParas], paraID)
+		if err != nil {
+			log.WithError(err).Fatal("Filtering out parachain heads failed.")
+		}
+		if merkleProofData.Root.Hex() != mmrProof.Leaf.ParachainHeads.Hex() {
+			log.WithFields(log.Fields{
+				"computedMmr": merkleProofData.Root.Hex(),
+				"mmr":         mmrProof.Leaf.ParachainHeads.Hex(),
+			}).Fatal("MMR parachain merkle root does not match calculated merkle root.")
+		}
+	}
+
 	log.WithFields(log.Fields{
 		"paraHeadsAsSlice": paraHeadsAsSlice,
 		"paraId":           paraID,
@@ -132,7 +151,6 @@ func ParachainHeadProofFn(cmd *cobra.Command, _ []string) error {
 		"parachainId":           paraID,
 		"relaychainBlockHash":   relayChainBlockHash.Hex(),
 		"relaychainBlockNumber": relayChainBlock,
-		"parachainBlockNumber":  parachainBlock,
 		"paraHeads":             paraHeadsAsSlice,
 		"parachainHeader":       parachainHeader,
 	}).Info("Generated proof input for parachain block.")
diff --git a/relayer/relays/parachain/beefy-listener.go b/relayer/relays/parachain/beefy-listener.go
index d0f567351c..41d0e5e292 100644
--- a/relayer/relays/parachain/beefy-listener.go
+++ b/relayer/relays/parachain/beefy-listener.go
@@ -297,8 +297,8 @@ func (li *BeefyListener) generateAndValidateParasHeadsMerkleProof(input *ProofIn
 
 	// Try a filtering out parathreads
 	log.WithFields(log.Fields{
-		"beefyBlock": merkleProofData.Root.Hex(),
-		"leafIndex":  mmrProof.Leaf.ParachainHeads.Hex(),
+		"computedMmr": merkleProofData.Root.Hex(),
+		"mmr":         mmrProof.Leaf.ParachainHeads.Hex(),
 	}).Warn("MMR parachain merkle root does not match calculated merkle root. Trying to filtering out parathreads.")
 
 	paraHeads, err = li.relaychainConn.FilterParachainHeads(paraHeads, input.RelayBlockHash)
@@ -306,6 +306,7 @@ func (li *BeefyListener) generateAndValidateParasHeadsMerkleProof(input *ProofIn
 		return nil, paraHeads, fmt.Errorf("could not filter out parathreads: %w", err)
 	}
 
+	numParas = min(MaxParaHeads, len(paraHeads))
 	merkleProofData, err = CreateParachainMerkleProof(paraHeads[:numParas], input.ParaID)
 	if err != nil {
 		return nil, paraHeads, fmt.Errorf("create parachain header proof: %w", err)