From 9be537a5f01aec61eab2d126ca310f406727f951 Mon Sep 17 00:00:00 2001
From: Josh Rickmar <jrick@companyzero.com>
Date: Mon, 11 Nov 2024 18:32:58 +0000
Subject: [PATCH] mixpool: Do not return early for revealed secrets

When a RS (reveal secrets) message was published, the mixpool Receive method
would immediately return, and blame assignment was immediately run.  This
created a logic race (due to message propagation over the network) where not
all peers would run blame assignment with all currently-published messages,
and different peers would blame different sets of peers for missing messages.

Correct this by modifying the Receive method to not return ErrSecretsRevealed
until the total number of expected messages have been received first.
---
 mixing/mixpool/mixpool.go | 59 +++++++++++++++++++--------------------
 1 file changed, 28 insertions(+), 31 deletions(-)

diff --git a/mixing/mixpool/mixpool.go b/mixing/mixpool/mixpool.go
index 335dbb938..5c1637693 100644
--- a/mixing/mixpool/mixpool.go
+++ b/mixing/mixpool/mixpool.go
@@ -731,10 +731,13 @@ type Received struct {
 // such session has any messages currently accepted in the mixpool, the method
 // immediately errors.
 //
-// If any secrets messages are received for the described session, and r.RSs
-// is nil, Receive immediately returns ErrSecretsRevealed.  An additional call
-// to Receive with a non-nil RSs can be used to receive all of the secrets
-// after each peer publishes their own revealed secrets.
+// If r.RSs is nil and any secrets messages are received for the described
+// session, Receive will count these messages towards the total number of
+// expected messages based on the slice capacities (without overcounting if
+// the identity has also sent an expected non-RS message) and will return
+// ErrSecretsRevealed.  An additional call to Receive with a non-nil RSs can
+// be used to receive all of the secrets after each peer publishes their own
+// revealed secrets.
 func (p *Pool) Receive(ctx context.Context, r *Received) error {
 	sid := r.Sid
 	var bc *broadcast
@@ -784,42 +787,32 @@ Loop:
 	for {
 		// Pool is locked for reads.  Count if the total number of
 		// expected messages have been received.
-		received := 0
+		received := make(map[idPubKey]struct{})
+		countMsg := func(msg mixing.Message) {
+			received[*(*idPubKey)(msg.Pub())] = struct{}{}
+		}
 		for hash := range ses.hashes {
-			msgtype := p.pool[hash].msgtype
+			e := p.pool[hash]
+			msg := e.msg
+			msgtype := e.msgtype
 			switch {
 			case msgtype == msgtypeKE && r.KEs != nil:
-				received++
+				countMsg(msg)
 			case msgtype == msgtypeCT && r.CTs != nil:
-				received++
+				countMsg(msg)
 			case msgtype == msgtypeSR && r.SRs != nil:
-				received++
+				countMsg(msg)
 			case msgtype == msgtypeDC && r.DCs != nil:
-				received++
+				countMsg(msg)
 			case msgtype == msgtypeCM && r.CMs != nil:
-				received++
+				countMsg(msg)
 			case msgtype == msgtypeFP && r.FPs != nil:
-				received++
+				countMsg(msg)
 			case msgtype == msgtypeRS:
-				if r.RSs == nil {
-					// Since initial reporters of secrets
-					// need to take the blame for
-					// erroneous blame assignment if no
-					// issue was detected, we only trigger
-					// this for RS messages that do not
-					// reference any other previous RS.
-					rs := p.pool[hash].msg.(*wire.MsgMixSecrets)
-					prev := rs.PrevMsgs()
-					if len(prev) == 0 {
-						p.mtx.RUnlock()
-						return ErrSecretsRevealed
-					}
-				} else {
-					received++
-				}
+				countMsg(msg)
 			}
 		}
-		if received >= expectedMessages {
+		if len(received) >= expectedMessages {
 			break
 		}
 
@@ -836,6 +829,8 @@ Loop:
 		p.mtx.RLock()
 	}
 
+	var err error
+
 	// Pool is locked for reads.  Collect all of the messages.
 	for hash := range ses.hashes {
 		msg := p.pool[hash].msg
@@ -865,14 +860,16 @@ Loop:
 				r.FPs = append(r.FPs, msg)
 			}
 		case *wire.MsgMixSecrets:
-			if r.RSs != nil {
+			if r.RSs == nil {
+				err = ErrSecretsRevealed
+			} else {
 				r.RSs = append(r.RSs, msg)
 			}
 		}
 	}
 
 	p.mtx.RUnlock()
-	return nil
+	return err
 }
 
 var zeroHash chainhash.Hash