From 10cbf41f5e49ed1b9138add355de287c91b182d1 Mon Sep 17 00:00:00 2001
From: corverroos <corver.roos@gmail.com>
Date: Fri, 2 Dec 2022 12:22:07 +0200
Subject: [PATCH] core/qbft: including round when deduping

---
 core/qbft/qbft.go               | 23 ++++++------
 core/qbft/qbft_internal_test.go | 64 +++++++++++++++++++++++++++++++++
 2 files changed, 76 insertions(+), 11 deletions(-)

diff --git a/core/qbft/qbft.go b/core/qbft/qbft.go
index 150d1c06d..34cf49a31 100644
--- a/core/qbft/qbft.go
+++ b/core/qbft/qbft.go
@@ -159,6 +159,12 @@ var ruleLabels = map[UponRule]string{
 	UponRoundTimeout:             "round_timeout",
 }
 
+// dedupKey defines the key used to deduplicate upon rules.
+type dedupKey struct {
+	UponRule UponRule
+	Round    int64
+}
+
 // Run executes the consensus algorithm until the context closed.
 // The generic type I is the instance of consensus and can be anything.
 // The generic type V is the arbitrary data value being proposed; it only requires an Equal method.
@@ -188,7 +194,7 @@ func Run[I any, V comparable](ctx context.Context, d Definition[I, V], t Transpo
 		preparedJustification []Msg[I, V]
 		qCommit               []Msg[I, V]
 		buffer                = make(map[int64][]Msg[I, V])
-		dedupRules            = make(map[UponRule]int64) // map[UponRule]msg.Round()
+		dedupRules            = make(map[dedupKey]bool)
 		timerChan             <-chan time.Time
 		stopTimer             func()
 	)
@@ -219,17 +225,12 @@ func Run[I any, V comparable](ctx context.Context, d Definition[I, V], t Transpo
 
 	// isDuplicatedRule returns true if the rule has been already executed since last round change.
 	isDuplicatedRule := func(rule UponRule, msgRound int64) bool {
-		prevRound, ok := dedupRules[rule]
-		if !ok {
-			dedupRules[rule] = msgRound
+		key := dedupKey{UponRule: rule, Round: msgRound}
 
-			return false
-		}
+		if !dedupRules[key] {
+			dedupRules[key] = true
 
-		if prevRound != msgRound {
-			// Upon rules are either for the current round,
-			// or for a future round followed by a round change (which clears this map).
-			panic("bug: duplicate rule, but different round")
+			return false
 		}
 
 		return true
@@ -242,7 +243,7 @@ func Run[I any, V comparable](ctx context.Context, d Definition[I, V], t Transpo
 		}
 		d.LogRoundChange(ctx, instance, process, round, newRound, rule, extractRoundMsgs(buffer, round))
 		round = newRound
-		dedupRules = make(map[UponRule]int64)
+		dedupRules = make(map[dedupKey]bool)
 	}
 
 	// === Algorithm ===
diff --git a/core/qbft/qbft_internal_test.go b/core/qbft/qbft_internal_test.go
index d5381e175..53728c7c9 100644
--- a/core/qbft/qbft_internal_test.go
+++ b/core/qbft/qbft_internal_test.go
@@ -25,6 +25,8 @@ import (
 	"github.com/stretchr/testify/require"
 
 	"github.com/obolnetwork/charon/app/errors"
+	"github.com/obolnetwork/charon/app/log"
+	"github.com/obolnetwork/charon/app/z"
 )
 
 func TestQBFT(t *testing.T) {
@@ -535,3 +537,65 @@ func makeIsLeader(n int64) func(int64, int64, int64) bool {
 		return (instance+round)%n == process
 	}
 }
+
+// TestDuplicatePrePreparesRules tests that two pre-prepares for different rounds are not detected as duplicates.
+func TestDuplicatePrePreparesRules(t *testing.T) {
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+
+	const (
+		noLeader = 1
+		leader   = 2
+	)
+
+	newPreprepare := func(round int64) Msg[int64, int64] {
+		return msg{
+			msgType: MsgPrePrepare,
+			peerIdx: leader,
+			round:   round,
+			// Justification not required since nodes and quorum both 0.
+		}
+	}
+
+	def := noopDef
+	def.IsLeader = func(_ int64, _ int64, process int64) bool {
+		return process == leader
+	}
+	def.LogUponRule = func(ctx context.Context, instance int64, process, round int64, msg Msg[int64, int64], uponRule UponRule) {
+		log.Info(ctx, "UponRule", z.Str("rule", uponRule.String()), z.I64("round", msg.Round()))
+		require.Equal(t, uponRule, UponJustifiedPrePrepare)
+		if msg.Round() == 1 {
+			return
+		}
+		if msg.Round() == 2 {
+			cancel()
+			return
+		}
+		require.Fail(t, "unexpected round", "round=%d", round)
+	}
+
+	ch := make(chan Msg[int64, int64], 2)
+	ch <- newPreprepare(1)
+	ch <- newPreprepare(2)
+
+	transport := noopTransport
+	transport.Receive = ch
+
+	_ = Run(ctx, def, transport, 0, noLeader, 1)
+}
+
+// noopTransport is a transport that does nothing.
+var noopTransport = Transport[int64, int64]{
+	Broadcast: func(context.Context, MsgType, int64, int64, int64, int64, int64, int64, []Msg[int64, int64]) error {
+		return nil
+	},
+}
+
+// noopDef is a definition that does nothing.
+var noopDef = Definition[int64, int64]{
+	IsLeader:       func(int64, int64, int64) bool { return false },
+	NewTimer:       func(int64) (<-chan time.Time, func()) { return nil, func() {} },
+	LogUponRule:    func(context.Context, int64, int64, int64, Msg[int64, int64], UponRule) {},
+	LogRoundChange: func(context.Context, int64, int64, int64, int64, UponRule, []Msg[int64, int64]) {},
+	LogUnjust:      func(context.Context, int64, int64, Msg[int64, int64]) {},
+}