From 42783334c93e7e8912a0e05e25019ae5b53caa41 Mon Sep 17 00:00:00 2001
From: Sorin Stanculeanu <sorinstanculeanu26@gmail.com>
Date: Fri, 31 Jan 2025 15:02:45 +0200
Subject: [PATCH] fixes for multikey mode, send proof from random managed key

---
 consensus/spos/bls/v2/subroundEndRound.go | 38 +++++++++++++++++++++--
 1 file changed, 35 insertions(+), 3 deletions(-)

diff --git a/consensus/spos/bls/v2/subroundEndRound.go b/consensus/spos/bls/v2/subroundEndRound.go
index cf23daa3ae..ad8d027dd3 100644
--- a/consensus/spos/bls/v2/subroundEndRound.go
+++ b/consensus/spos/bls/v2/subroundEndRound.go
@@ -6,6 +6,7 @@ import (
 	"encoding/hex"
 	"fmt"
 	"math/bits"
+	"math/rand"
 	"sync"
 	"time"
 
@@ -100,7 +101,7 @@ func (sr *subroundEndRound) receivedProof(proof consensus.ProofHandler) {
 	sr.mutProcessingEndRound.Lock()
 	defer sr.mutProcessingEndRound.Unlock()
 
-	if sr.IsJobDone(sr.SelfPubKey(), sr.Current()) {
+	if sr.IsSelfJobDone(sr.Current()) {
 		return
 	}
 	if !sr.IsConsensusDataSet() {
@@ -555,18 +556,49 @@ func (sr *subroundEndRound) createAndBroadcastProof(signature []byte, bitmap []b
 		IsStartOfEpoch:      sr.GetHeader().IsStartOfEpochBlock(),
 	}
 
-	err := sr.BroadcastMessenger().BroadcastEquivalentProof(headerProof, []byte(sr.SelfPubKey()))
+	sender := sr.getEquivalentProofSender()
+	err := sr.BroadcastMessenger().BroadcastEquivalentProof(headerProof, []byte(sender))
 	if err != nil {
 		return err
 	}
 
 	log.Debug("step 3: block header proof has been sent",
 		"PubKeysBitmap", bitmap,
-		"AggregateSignature", signature)
+		"AggregateSignature", signature,
+		"proof sender", hex.EncodeToString([]byte(sender)))
 
 	return nil
 }
 
+func (sr *subroundEndRound) getEquivalentProofSender() string {
+	if sr.IsNodeInConsensusGroup(sr.SelfPubKey()) {
+		return sr.SelfPubKey() // single key mode
+	}
+
+	return sr.getRandomManagedKeyProofSender()
+}
+
+func (sr *subroundEndRound) getRandomManagedKeyProofSender() string {
+	// in multikey mode, we randomly select one managed key for the proof
+	consensusKeysManagedByCurrentNode := make([]string, 0)
+	for _, validator := range sr.ConsensusGroup() {
+		if !sr.IsKeyManagedBySelf([]byte(validator)) {
+			continue
+		}
+
+		consensusKeysManagedByCurrentNode = append(consensusKeysManagedByCurrentNode, validator)
+	}
+
+	if len(consensusKeysManagedByCurrentNode) == 0 {
+		return sr.SelfPubKey() // fallback return self pub key, should never happen
+	}
+
+	randIdx := rand.Intn(len(consensusKeysManagedByCurrentNode))
+	randManagedKey := consensusKeysManagedByCurrentNode[randIdx]
+
+	return randManagedKey
+}
+
 func (sr *subroundEndRound) createAndBroadcastInvalidSigners(invalidSigners []byte) {
 	if !sr.ShouldConsiderSelfKeyInConsensus() {
 		return